feat: add Java scalar UDF support by andygrove · Pull Request #46 · apache/datafusion-java · GitHub

Which issue does this PR close?

N/A

Rationale for this change

Add support for implement scalar UDFs in Java.

What changes are included in this PR?

Public Java API

ScalarUdf @FunctionalInterface with one method FieldVector evaluate(BufferAllocator allocator, List<FieldVector> args)
Volatility enum (IMMUTABLE / STABLE / VOLATILE)
SessionContext.registerUdf(name, udf, returnType, argTypes, volatility)

Internals

org.apache.datafusion.internal.JniBridge — per-call static trampoline that imports the input columns, calls user code, validates (non-null + row count), and exports the result via the Arrow C Data Interface
native/src/udf.rs — JavaScalarUdf implementing DataFusion's ScalarUDFImpl. Holds GlobalRefs to the user instance + bridge class and a cached JStaticMethodID; constructs an FFI struct-array view of the args, attaches the current thread to JNI, calls the bridge, translates any pending Java exception into a DataFusionError::Execution, imports the result via arrow::ffi::from_ffi, and validates the type matches the declared return type
JNI_OnLoad in native/src/lib.rs caches the JavaVM* in a OnceLock so DataFusion's worker threads can attach

Refinements from the design discussion

Args ride a single FFI struct pair (FFI_ArrowArray + FFI_ArrowSchema for a struct-array view of the columns), not an ArrowArrayStream — v1 always sees one batch per invoke and the simpler ABI avoids streaming overhead.
The per-call allocator is a single shared static RootAllocator on JniBridge rather than a fresh-per-call one — closing a per-call allocator while an FFI release callback still holds buffer references would throw.
StructArray::try_new_with_length(... , number_rows) is used for the args export so a zero-argument UDF doesn't panic.
number_rows is usize::try_into::<i32> checked before crossing JNI — a batch larger than i32::MAX would otherwise silently truncate and miscompare with Java's row-count check.

Docs + example

docs/source/user-guide/scalar-udf.md (linked from the user-guide toctree)
examples/src/main/java/org/apache/datafusion/examples/AddOneExample.java

How are these changes tested?

make test from a clean checkout. The new ScalarUdfTest has 12 tests covering:

Happy paths: add_one(Int32), concat(Utf8, Utf8), square(Float64), repeated invocations in one session, a 100-row VALUES scan
Contract violations: UDF returning null, wrong row count, wrong type, and a UDF that throws IllegalArgumentException — all surface as RuntimeExceptions with the expected message substrings (class name + user message preserved)
Lifecycle: two UDFs registered in the same session, register-after-close in a new session
Volatility: all three Volatility values round-trip through registration and execute correctly

Introduces JavaScalarUdf struct with its Signature, return_type, JNI references, and a ScalarUDFImpl impl whose invoke_with_args returns NotImplemented (placeholder for Task 6). Adds volatility_from_byte helper. Dead-code lints suppressed with comments referencing Task 5/6.

Add end-to-end UDF registration: SessionContext.registerUdf serialises the return/arg types as Arrow IPC, passes them to native via registerScalarUdf JNI, which decodes the schema, constructs a JavaScalarUdf with an exact Signature, and registers it on the DataFusion SessionContext. SQL planning now resolves registered UDFs by name; invocation still returns NotImplemented until Task 6.

Replace the NotImplemented stub with the real per-call data flow: materialise ColumnarValue args to arrays, pack them into a StructArray, export via to_ffi, call JniBridge.invokeScalarUdf over JNI with four FFI addresses, then import the result via from_ffi and return it as a ColumnarValue::Array. Update ScalarUdfTest to assert correct values from an AddOne UDF registered against a three-row constant table.

… and zero-arg case - Clear pending JNI exceptions in jthrowable_to_string fallback branches to avoid undefined behavior on thread detach - Add SAFETY comment on the from_ffi unsafe block explaining FFI initialization guarantees - Replace usize-to-i32 silent truncation with i32::try_from checked conversion for row count - Replace StructArray::new with StructArray::try_new_with_length to handle zero-arg UDFs

…ions

Thanks for letting this one linger so I could take a look!

The equivalent rust function takes a ScalarUDF struct that exposes many of these arguments, rather than passing them here at the top level. Is it worth emulating that structure here?

This is obviously a much broader question, it could be asked about nearly every public API we add. I waffled on it myself, and ended up deciding to strictly follow the rust API as closely as possible (my ScalarUDF is here - not saying we should emulate that; just for reference). My justification was that it would be easier to stay in line with the upstream rust as it evolved, and it would be clearer what Java supported.

Whatever we decide, I think we ought to document the strategy going forward, especially for the benefit of coding agents. I attempted this in my own bindings with Claude rules for defining the Public API as well as for documenting it, going so far as to ensure that Javadocs provide a working link to the rust equivalent. I'd be happy to do a version of this for these new bindings (wouldn't make sense to use those I've linked exactly).

I agree that documenting the strategy makes sense. Please go ahead and work on this when you have time. I also agree that we should stay as close to the Rust APIs as we can. I did some refactoring of this - let me know what you think.

Great API, great docs! Similar to but better than what I did.

I think it would be nice to propagate the stack trace (in full or in part), at least optionally. In my bindings I made it a session configurable option, though I could imagine many other strategies. The implementation would look quite different in the bridge.

Could make sense as follow on work for all upcalls, rather than in the scope of this PR.

That sounds like a good idea. I filed #55 to track this

…der class Mirror DataFusion's Rust shape: ScalarFunction (formerly the eval-only ScalarUdf interface) now carries name/argTypes/returnType/volatility and the per-batch evaluate body, mirroring ScalarUDFImpl. ScalarUdf becomes a holder that wraps a ScalarFunction and is what SessionContext.registerUdf takes — mirroring the ScalarUDF struct. The holder caches the four metadata values in its constructor so impls that allocate per call aren't double-evaluated at registration. The JNI bridge takes a ScalarFunction directly; the wrapper is bypassed on the hot path.

Looks great! I'll work on the rust/java API alignment documentation as a priority.

The other key documentation unlock that let me let coding agents loose implementing features without close supervision, was the patterns to follow in the rust bridge code. The version from my bindings is much less relevant here given the architectural differences, but I think something like it makes sense. I have much less conviction on that stuff too - my rust experience is as a hobbiest only (hundreds of hours of experience) vs my java experience (ten thousand+ hours) - but I'll try to do something at least.

Projects like these are so much more feasible these days with well-guided coding agents, which @timsaucer is clearly leaning into with the python bindings.

Looks great! I'll work on the rust/java API alignment documentation as a priority.

The other key documentation unlock that let me let coding agents loose implementing features without close supervision, was the patterns to follow in the rust bridge code. The version from my bindings is much less relevant here given the architectural differences, but I think something like it makes sense. I have much less conviction on that stuff too - my rust experience is as a hobbiest only (hundreds of hours of experience) vs my java experience (ten thousand+ hours) - but I'll try to do something at least.

Projects like these are so much more feasible these days with well-guided coding agents, which @timsaucer is clearly leaning into with the python bindings.

I look forward to having some LLM skills in this repo. I think this type of project is ideal for LLMs. Thanks for the review!

I think this type of project is ideal for LLMs.

Agreed! It was the realization of how good they'd gotten in the last six months that finally convinced me to take the project seriously. Even though it's a large surface area that requires a lot of code, adding features and keeping up with upstream changes should be a lot of implementing the same patterns over and over.

andygrove added 13 commits May 13, 2026 17:06

feat(native): cache JavaVM in OnceLock via JNI_OnLoad

7517f86

feat(udf): add ScalarUdf interface and Volatility enum

52a2589

feat(udf): add JniBridge trampoline for per-call UDF invocation

10c32bb

test(udf): cover concat, square, and repeated invocations

742629c

test(udf): cover null/wrong-rows/wrong-type/exception contract violat…

cf187a3

…ions

test(udf): cover multi-UDF, cross-session, and multi-row lifecycle

4ef01fd

test(udf): exercise all three Volatility kinds

086cba7

docs(udf): add user-guide page and runnable AddOne example

0d06aa5

refactor(udf): null-guard registerUdf arguments and drop dead assertion

296b3be

pgwhalen reviewed May 17, 2026

View reviewed changes

andygrove mentioned this pull request May 17, 2026

Propagate Java stack traces from JVM upcalls into DataFusionError #55

Open

andygrove marked this pull request as ready for review May 17, 2026 19:57

andygrove requested a review from pgwhalen May 17, 2026 20:07

pgwhalen approved these changes May 18, 2026

View reviewed changes

andygrove merged commit dd5d6fc into apache:main May 18, 2026
2 of 3 checks passed

andygrove deleted the feat/scalar-udf branch May 18, 2026 03:20

LantaoJin mentioned this pull request May 18, 2026

fix(udf): pass batch row count to ScalarFunction.evaluate #57

Closed

+                public void registerUdf(
+                    String name,
+                    ScalarUdf udf,
+                    ArrowType returnType,
+                    List<ArrowType> argTypes,
+                    Volatility volatility) {

+              ## Implement
+              Implement the `ScalarUdf` interface:
+              ```java
+              import java.util.List;
+              import org.apache.arrow.memory.BufferAllocator;
+              import org.apache.arrow.vector.FieldVector;
+              import org.apache.arrow.vector.IntVector;
+              import org.apache.datafusion.ScalarUdf;
+              public final class AddOne implements ScalarUdf {
+                  @Override
+                  public FieldVector evaluate(BufferAllocator allocator, List<FieldVector> args) {
+                      IntVector in = (IntVector) args.get(0);
+                      IntVector out = new IntVector("add_one", allocator);
+                      out.allocateNew(in.getValueCount());
+                      for (int i = 0; i < in.getValueCount(); i++) {
+                          if (in.isNull(i)) out.setNull(i);
+                          else out.set(i, in.get(i) + 1);
+                      }
+                      out.setValueCount(in.getValueCount());
+                      return out;
+                  }
+              }
+              ```
+              Allocate any new vectors — including the result — from the supplied
+              `BufferAllocator`. The input vectors are read-only views; do not close them.
+              Ownership of the returned vector transfers to the framework on return.

feat: add Java scalar UDF support by andygrove · Pull Request #46 · apache/datafusion-java · GitHub

Repository navigation

feat: add Java scalar UDF support - #46

feat: add Java scalar UDF support#46
andygrove merged 14 commits into
apache:mainfrom
andygrove:feat/scalar-udf

andygrove commented May 14, 2026 •
edited
Loading

Uh oh!

Uh oh!

pgwhalen left a comment

Uh oh!

pgwhalen May 17, 2026

Uh oh!

andygrove May 17, 2026

Uh oh!

pgwhalen May 17, 2026

Uh oh!

pgwhalen May 17, 2026

Uh oh!

andygrove May 17, 2026

Uh oh!

pgwhalen left a comment

Uh oh!

andygrove commented May 18, 2026

Uh oh!

Uh oh!

pgwhalen commented May 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

+                  }
+              }
+              /// Best-effort: extract class name and getMessage() from a Java throwable.

feat: add Java scalar UDF support by andygrove · Pull Request #46 · apache/datafusion-java · GitHub

Repository navigation

feat: add Java scalar UDF support - #46

feat: add Java scalar UDF support#46andygrove merged 14 commits intoapache:mainapache/datafusion-java:mainfrom andygrove:feat/scalar-udfandygrove/datafusion-java:feat/scalar-udfCopy head branch name to clipboard

Conversation

andygrove commented May 14, 2026 • 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?

How are these changes tested?

Uh oh!

pgwhalen left a comment

Choose a reason for hiding this comment

Uh oh!

pgwhalen May 17, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove May 17, 2026

Choose a reason for hiding this comment

Uh oh!

pgwhalen May 17, 2026

Choose a reason for hiding this comment

Uh oh!

pgwhalen May 17, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove May 17, 2026

Choose a reason for hiding this comment

Uh oh!

pgwhalen left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove commented May 18, 2026

Uh oh!

Uh oh!

pgwhalen commented May 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

feat: add Java scalar UDF support#46
andygrove merged 14 commits into
apache:mainfrom
andygrove:feat/scalar-udf

andygrove commented May 14, 2026 •
edited
Loading