feat: Add Aggregate UDF to FFI crate #14775

timsaucer · 2025-02-19T14:00:09Z

Which issue does this PR close?

This PR addresses part of #14562

Rationale for this change

This change allows for using user defined aggregate functions across FFI boundaries. It is useful for enabling shared libraries to pass functions back and forth. This feature will unlock:

Modules to provide DataFusion based FFI aggregates that can be reused in projects such as datafusion-python
Allows for use across different DataFusion versions between the function provider and consumer
Is a major step towards getting to the ability to export a function registry via FFI, a large blocker on sharing a session context

What changes are included in this PR?

This PR follows the same pattern as the previous FFI code. It exposes AggregateUDF and the Accumulators and AccumulatorArgs that go along with it.

Are these changes tested?

Included in this PR are both unit tests and two integration tests. The integration tests cover both grouping and non-grouping UDAFs. For the unit tests here is a tarpaulin report showing the coverage. The 12% not covered is almost entirely error handling.

Are there any user-facing changes?

There are no changes to existing API but new functions are exposed in the datafusion-ffi crate.

datafusion/ffi/src/tests/udf_udaf_udwf.rs

m09526

This is my first substantial review of someone's else code. Please let me know if I am not following proper steps in reviewing!

datafusion/ffi/src/arrow_wrappers.rs

datafusion/ffi/src/record_batch_stream.rs

datafusion/ffi/src/udaf/accumulator.rs

datafusion/ffi/src/udaf/mod.rs

datafusion/ffi/src/udaf/groups_accumulator.rs

timsaucer · 2025-03-02T11:29:01Z

@m09526 Thank you for the review! I've been out of country all last week, so I will try to get back to this soon now that I'm back.

alamb · 2025-03-12T15:51:36Z

Let me know when this is ready for review

timsaucer · 2025-03-12T18:58:53Z

Moved to draft until I get time to address the comments above. Additionally I discovered that Nullary type signatures do not make it properly through the type coercion, so that needs fixing before this can be widely used.

timsaucer · 2025-03-12T18:59:07Z

The above Nullary issue also impacts the existing Scalar UDFs.

timsaucer · 2025-03-31T14:01:44Z

Now that #15487 is resolved, this is unblocked.

alamb · 2025-04-07T12:55:08Z

Now that #15487 is resolved, this is unblocked.

Shall we mark it as "ready for review"?

crystalxyz · 2025-04-08T19:17:35Z

Now that #15487 is resolved, this is unblocked.

Shall we mark it as "ready for review"?

@alamb I'm helping Tim on this and I'm still addressing some comments. It shouldn't take too long. Will mark it ready for review when I'm done!

crystalxyz · 2025-04-09T02:03:29Z

@timsaucer @alamb Should be good for review now. Thanks!

timsaucer

Excellent work! Thank you so much for pulling this over the line.

I have a few suggestions, but the core functionality looks correct to me. As one of the authors I cannot approve the pull request.

timsaucer · 2025-04-10T11:05:04Z

datafusion/ffi/src/udaf/groups_accumulator.rs

+pub struct FFI_GroupsAccumulator {
+    pub update_batch: unsafe extern "C" fn(
+        accumulator: &mut Self,
+        values: RVec<WrappedArray>,
+        group_indices: RVec<usize>,
+        opt_filter: ROption<WrappedArray>,
+        total_num_groups: usize,
+    ) -> RResult<(), RString>,


We need to add documentation for the struct, but it should be very similar to the others except pointing to GroupsAccumulator

I would personally suggest not implementing the GroupsAccumulator in the first FFI implementation -- it is hard to implement correctly as is.

Perhaps we can avoid this for now 🤔

Though I suppose GroupsAccumulator is basically required for high performance aggregation 🤔

datafusion/ffi/src/udaf/groups_accumulator.rs

datafusion/ffi/tests/ffi_integration.rs

timsaucer · 2025-04-10T11:13:21Z

@robtandy As someone keenly interested in the FFI work, would you be able to do a review?

alamb

THanks @timsaucer -- this looks like it is pretty close. It has a few unresolved comments, but then would be ready to go it seems

datafusion/core/src/datasource/file_format/parquet.rs

datafusion/core/src/lib.rs

datafusion/ffi/src/arrow_wrappers.rs

alamb · 2025-05-07T21:00:15Z

datafusion/ffi/src/udaf/groups_accumulator.rs

+pub struct FFI_GroupsAccumulator {
+    pub update_batch: unsafe extern "C" fn(
+        accumulator: &mut Self,
+        values: RVec<WrappedArray>,
+        group_indices: RVec<usize>,
+        opt_filter: ROption<WrappedArray>,
+        total_num_groups: usize,
+    ) -> RResult<(), RString>,


I would personally suggest not implementing the GroupsAccumulator in the first FFI implementation -- it is hard to implement correctly as is.

Perhaps we can avoid this for now 🤔

Though I suppose GroupsAccumulator is basically required for high performance aggregation 🤔

timsaucer · 2025-06-04T21:31:14Z

@alamb I believe this one is ready.

alamb

Thanks @timsaucer -- I think this looks good to me. I didn't have time to review all the APIs in detail, but I think we need to get the big initial PR in and then work on fixing any blemishes over time

Requested changes were made

alamb · 2025-06-05T18:17:27Z

Thank you very much @m09526 for the review

alamb · 2025-06-05T18:24:15Z

I merged up one more time to make sure the CI tests pass but assuming they do I plan to merge this one in

alamb · 2025-06-05T19:08:03Z

It seems github is experiencing issues. I will close/reopen this PR to restart the checks

https://www.githubstatus.com/

alamb · 2025-06-05T19:58:55Z

gogogo

* Work in progress adding user defined aggregate function FFI support * Intermediate work. Going through groups accumulator * MVP for aggregate udf via FFI * Clean up after rebase * Add unit test for FFI Accumulator Args * Adding unit tests and fixing memory errors in aggregate ffi udf * Working through additional unit and integration tests for UDAF ffi * Switch to a accumulator that supports convert to state to get a little better coverage * Set feature so we do not get an error warning in stable rustc * Add more options to test * Add unit test for FFI RecordBatchStream * Add a few more args to ffi accumulator test fn * Adding more unit tests on ffi aggregate udaf * taplo format * Update code comment * Correct function name * Temp fix record batch test dependencies * Address some comments * Revise comments and address PR comments * Remove commented code * Refactor GroupsAccumulator * Add documentation * Split integration tests * Address comments to refactor error handling for opt filter * Fix linting errors * Fix linting and add deref * Remove extra tests and unnecessary code * Adjustments to FFI aggregate functions after rebase on main * cargo fmt * cargo clippy * Re-implement cleaned up code that was removed in last push * Minor review comments --------- Co-authored-by: Crystal Zhou <crystal.zhouxiaoyue@hotmail.com> Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

timsaucer force-pushed the feat/aggregate-udf-ffi branch 2 times, most recently from 9aae315 to 77fd002 Compare February 27, 2025 07:00

timsaucer marked this pull request as ready for review February 27, 2025 07:16

timsaucer changed the title ~~feat: Add Aggregate UDF to FFI interface~~ feat: Add Aggregate UDF to FFI crate Feb 27, 2025

Omega359 reviewed Feb 27, 2025

View reviewed changes

datafusion/ffi/src/tests/udf_udaf_udwf.rs Show resolved Hide resolved

m09526 previously requested changes Feb 28, 2025

View reviewed changes

alamb mentioned this pull request Mar 4, 2025

Weekly Plan (Andrew Lamb) March 3, 2025 #14978

Closed

12 tasks

timsaucer marked this pull request as draft March 12, 2025 18:57

alamb mentioned this pull request Mar 17, 2025

[Epic] A collection of FFI related tasks #15283

Open

10 tasks

github-actions bot added logical-expr Logical plan and expressions core Core DataFusion crate ffi Changes to the ffi crate labels Apr 2, 2025

timsaucer commented Apr 10, 2025

View reviewed changes

timsaucer marked this pull request as ready for review April 10, 2025 15:01

This was referenced Apr 28, 2025

Weekly Plan: Andrew Lamb 2025-04-28 #15880

Closed

Weekly Plan: Andrew Lamb 2025-05-05 #15943

Closed

alamb reviewed May 7, 2025

View reviewed changes

alamb mentioned this pull request May 11, 2025

Weekly Plan: Andrew Lamb 2025-05-12 #16022

Closed

24 tasks

alamb mentioned this pull request May 19, 2025

Weekly Plan: Andrew Lamb 2025-05-19 #16101

Closed

18 tasks

alamb mentioned this pull request May 27, 2025

Weekly Plan: Andrew Lamb 2025-05-26 #16198

Closed

11 tasks

timsaucer mentioned this pull request Jun 4, 2025

Release DataFusion 48.0.0 (June 2025) #15771

Open

29 tasks

crystalxyz and others added 14 commits June 4, 2025 14:34

Temp fix record batch test dependencies

ec05091

Address some comments

fc40bc0

Revise comments and address PR comments

4d164cc

Remove commented code

173d7c0

Refactor GroupsAccumulator

45ea283

Add documentation

b6da0e9

Split integration tests

a2fbda8

Address comments to refactor error handling for opt filter

0b4a8f5

Fix linting errors

4b3c533

Fix linting and add deref

1b85dd9

Remove extra tests and unnecessary code

8a4de4a

Adjustments to FFI aggregate functions after rebase on main

5e02c72

cargo fmt

d128b85

cargo clippy

4282c2a

timsaucer force-pushed the feat/aggregate-udf-ffi branch from ec4dda9 to 4282c2a Compare June 4, 2025 19:22

timsaucer added 2 commits June 4, 2025 15:27

Re-implement cleaned up code that was removed in last push

bb84d08

Minor review comments

1c3fad9

timsaucer mentioned this pull request Jun 5, 2025

feat: Add Window UDFs to FFI Crate #16261

Merged

alamb approved these changes Jun 5, 2025

View reviewed changes

Merge branch 'main' into feat/aggregate-udf-ffi

e63af82

alamb closed this Jun 5, 2025

alamb reopened this Jun 5, 2025

alamb merged commit 5d3ed9c into apache:main Jun 5, 2025
45 of 58 checks passed

feat: Add Aggregate UDF to FFI crate #14775

feat: Add Aggregate UDF to FFI crate #14775

Uh oh!

Conversation

timsaucer commented Feb 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Uh oh!

m09526 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timsaucer commented Mar 2, 2025

Uh oh!

alamb commented Mar 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

timsaucer commented Mar 12, 2025

Uh oh!

timsaucer commented Mar 12, 2025

Uh oh!

timsaucer commented Mar 31, 2025

Uh oh!

alamb commented Apr 7, 2025

Uh oh!

crystalxyz commented Apr 8, 2025

Uh oh!

crystalxyz commented Apr 9, 2025

Uh oh!

timsaucer left a comment

Choose a reason for hiding this comment

Uh oh!

timsaucer Apr 10, 2025

Choose a reason for hiding this comment

Uh oh!

alamb May 7, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timsaucer commented Apr 10, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alamb May 7, 2025

Choose a reason for hiding this comment

Uh oh!

timsaucer commented Jun 4, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

alamb commented Jun 5, 2025

Uh oh!

alamb commented Jun 5, 2025

Uh oh!

alamb commented Jun 5, 2025

Uh oh!

Uh oh!

alamb commented Jun 5, 2025

Uh oh!

timsaucer commented Feb 19, 2025 •

edited

Loading

alamb commented Mar 12, 2025 •

edited

Loading