[AURON #1862] Introduce FlinkAuronUDFWrapperContext for Flink UDF fallback by weiqingy · Pull Request #2490 · apache/auron

weiqingy · 2026-08-30T22:43:54Z

Which issue does this PR close?

Rationale for this change

A Flink Calc containing a call Auron cannot convert falls back in its entirety, so one unsupported user-defined function takes the whole projection back to Flink's generated code even when every other expression in it converts.

This adds a fallback for user-defined scalar functions that is scoped to the call rather than the Calc. The call becomes a native expression node whose evaluation is an upcall into the JVM; its arguments are ordinary native expressions, so anything that converts stays converted. Parameters cross the Arrow C Data Interface once per batch rather than once per row.

This is a coverage feature, not a performance one. A native implementation of a function is still the faster path and remains the one to prefer; what this buys is that reaching for the JVM costs one call instead of the surrounding query.

What changes are included in this PR?

The call is packaged into the existing UDF wrapper expression node, which is renamed to an engine-neutral name and dispatched through its interface rather than the Spark concrete class. Nothing about the node's shape or its protobuf field number changes.

The invocation itself is generated rather than reflected. At plan time the call's operands are rewritten to input references into the parameter row and handed to Flink's ExprCodeGenerator; the resulting source and its reference array travel in the node's opaque payload, and the task compiles it with CompileUtils against the user-code classloader. Overload selection, argument conversion and null handling therefore come from the same TypeInference that runs when Auron is not involved, rather than from logic here that would have to re-derive them.

One wrapper is retained per subtask, so a function's open runs once for that subtask and close runs at teardown. Without that, a busy subtask rebuilds the wrapper on each of roughly five drain cycles a second, and a function that acquires anything in open cannot work at all.

A JNI classloader fix is included. The lazy JNIEnv initialiser overwrote the calling thread's context classloader with a JVM-global captured at the first native call, so from the second job onward a Flink task thread carried a previous job's closed loader. It predates this feature and is shared with Spark.

Are there any user-facing changes?

A query whose Calc contains a user-defined scalar function call now keeps the rest of that Calc native instead of falling back wholesale. The generated call is the one Flink itself emits for the same expression, so results match for any function whose output depends only on its arguments.

Evaluation order is where the two diverge. Auron evaluates a call site across a whole batch; Flink's generated Calc evaluates every call site per row. Row order is preserved either way, so a function whose result depends only on its arguments returns the same rows under both.

Two shapes can differ. Where a Calc calls the same function at more than one site, Flink shares one instance across those sites while Auron gives each its own, so a function carrying mutable state sees separate instances and separate open/close pairs per subtask. And where one function reads state that another mutates outside the argument channel, a static or a singleton, the batch ordering changes what it reads. Nested calls are converted the same way and follow the same rule: f(g(x)) returns Flink's rows whenever f depends only on its arguments.

Sharing one instance would not reproduce Flink's answer either, since a shared instance under batch evaluation yields a third result rather than Flink's.

Calls are still declined, and fall back exactly as before, when an argument or return type falls outside the set Auron's Arrow boundary carries, when the function instance is not serializable or preparable, when eval takes no arguments, when it is varargs, or when an operand does not itself convert.

No configuration is added. No protobuf field is added or renumbered.

How was this patch tested?

Unit tests over the converter, the payload, the generator and the wrapper, and integration cases that run real queries through a TableEnvironment against the native library.

Every case that executes a query asserts both the fallback count and the row values, because a missing native arm converts cleanly and returns an empty result set, so neither assertion alone establishes that the native path ran.

Three properties are pinned by mutation rather than by assertion alone: removing the generated null short-circuit, reversing the operand order, and reinstating the deleted overload gate each fail the test written for them and no other.

Three integration cases compare against the same query with Auron off. Since the operator is activated by classpath shadowing and there is no switch, the comparison run appends a call whose return type declines the whole Calc and then drops that column. That this really disables the native path was checked by execution, with a probe function recording the stack reaching its own eval: the native run arrives through the generated invoker with no fallback recorded, and the comparison run arrives through Flink's StreamExecCalc with no Auron frame present.

Full build green, including the cross-version compile legs for Spark 3.0 through 4.1.

Was this patch authored or co-authored using generative AI tooling?

Yes
No

Generated-by: Claude Code (Opus 5)

…ne-neutral name The UDF wrapper expression is reused by any engine that needs to evaluate an unsupported expression on the JVM, so its Spark-specific naming no longer describes it. Rename the protobuf message to PhysicalUDFWrapperExprNode, the oneof field to udf_wrapper_expr, and the Rust module and struct to udf_wrapper and UDFWrapperExpr. Field number 10000 is unchanged, so the encoding is byte-identical. Dispatch still goes through the Spark concrete class, unchanged, so this commit has no behavioural content. SparkAuronUDFWrapperContext keeps its name in both the Scala class and the JNI binding: it is the Spark implementation of the shared context interface, and is genuinely Spark-specific.

…al interface Resolve eval through the AuronUDFWrapperContext interface rather than through the Spark concrete class. The object handle is already interface-typed, since JniBridge.getAuronUDFWrapperContext returns the interface, so only the method-ID lookup was engine-specific. This also removes a latent crash: the JNI bridge fills Spark-specific class bindings with a zeroed default on non-Spark engines, and JMethodID is repr(transparent) over a raw pointer, so the previous line would have passed a NULL jmethodID to CallVoidMethodA rather than failing gracefully. Add a scalar-subquery test, because this call site has a second consumer: the subquery wrapper reuses the same expression as its JVM-upcall vehicle, and it had no coverage. The test asserts the plan retains the subquery and that the result is non-empty, so it cannot pass without exercising the path.

Add FlinkAuronUDFWrapperContext, the Flink implementation of the shared AuronUDFWrapperContext interface. It imports parameters over the Arrow C Data Interface, invokes the user function once per row, and exports the results back, so one JNI crossing serves a whole batch. The eval method handle is bound with unreflect and asType rather than findVirtual, which requires an exact descriptor and would fail for every primitive-parameter function, since the descriptor is built from boxed conversion classes. Arrow roots are allocated from the root allocator rather than an operator-scoped child, because the exported result array is released only after eval returns. The argument array, output row and converters are reused across calls; the javadoc records the scheduling invariant that makes this safe and what would break it. FlinkUDFPayload carries the function instance, its resolved argument and return types, and the declared parameter type names, derived from the Method itself so the planner and runtime cannot spell them differently. Wire FlinkAuronAdaptor to return the context instead of throwing. Nothing emits the wrapper node yet, so no query behaviour changes.

…thread The lazy JNIEnv initialiser installed the JVM-global classloader on every thread that called in, including threads that arrived from Java and already had a context classloader of their own. On a Flink TaskManager the global is captured at the first native call and never refreshed, so from the second job onward a task thread had its own open user-code loader replaced by the previous job's, which Flink has already closed. Nothing restored it afterwards. Install the global only when the thread was not already attached, which is the case for threads native code creates. A thread arriving from Java keeps its own loader. The visible symptom is a closed-classloader failure when deserializing user code, but two quieter shapes matter more: a still-open loader belonging to a concurrently running job yields ClassNotFoundException, and the same class name from two different jars silently resolves the wrong version.

invokeWithArguments boxes into varargs and re-derives the handle type on every call, which puts the per-row invocation at roughly reflection speed. Bind with asSpreader and call invokeExact instead, which is the form the JIT can inline. Also correct the javadoc explaining why unreflect is used. It claimed findVirtual could not bind a primitive parameter because the descriptor is boxed, which is not true of this code: the declared parameter types come from the payload and resolve to primitive classes. The accurate reason is that unreflect takes the descriptor from the resolved Method itself.

Route an unconvertible user scalar function into the UDF wrapper node instead of failing the whole Calc back to Flink codegen. Detection keys on operator identity rather than SqlKind, since Flink's own IF, TRY_CAST and UNIX_TIMESTAMP are all OTHER_FUNCTION, and covers both BridgingSqlFunction and the deprecated ScalarSqlFunction path. Admission is decided at plan time: supported types, exactly one invokable eval overload, no open or close override, and a serializable instance. Rejecting at plan time matters because a type failure surviving to runtime surfaces as an empty result set rather than an error. Arguments that convert natively are emitted as parameters, so only the call itself runs on the JVM. Emit the short-circuit AND and OR nodes when a wrapper lands on a right operand. The native binary expression evaluates its right side on the full batch unless the left side is nearly all false, while Flink's generated code genuinely short-circuits, so without this a function could run on rows Flink would skip. Move the isSupported call inside the factory's try block so an exception from it no longer escapes conversion.

…subtask The wrapper was rebuilt on every drain cycle, so under a default watermark interval a user function was deserialized and reopened several times a second. Hold the wrappers in a per-subtask context that rides the channel the native runtime already uses to reach its worker threads, so a function is opened once per subtask and closed when the subtask finishes. Build the wrappers in the operator's open rather than lazily on first use. Lazily would run open on a worker thread that is discarded at the end of the first drain cycle, which would break a function that keeps state in a thread local, and would leave a subtask that receives no rows never opening its functions at all. With the lifetime fixed, drop the plan-time gate that rejected functions overriding open or close. The gate existed only because the wrapper could not offer them a sensible lifecycle. Carry a per-node ordinal in the payload so two call sites of the same function do not share a wrapper. Flink deep-copies a function per call site, and without a discriminator both sites produce identical payloads, which is invisible for a stateless function and wrong for one that initialises state in open.

…ing a method handle The wrapper bound a MethodHandle to the eval overload the planner selected and converted each argument through a DataStructureConverter it held itself. Selecting that overload meant re-deriving Java's most-specific rule, which the extraction helpers do not implement, so a function declaring two invokable overloads was declined rather than guessed at. Drive Flink's own generator instead. The planner rewrites the call's operands to input references into the parameter row, hands the result to ExprCodeGenerator, and emits a class implementing AuronGeneratedUDF; the payload carries that source with the reference array it is only valid against. The task compiles it with CompileUtils against the user-code classloader, which is the one loader that can see both the user function and the interface the generated class implements. Overload selection, argument conversion and null handling now come from the same TypeInference Flink uses for the query without Auron, so the ambiguity gate goes. The generated class declares its own getRuntimeContext(), which is what lets the emitted open statement build a real FunctionContext rather than the degraded form Flink's own hostless path settles for. Generation is validated by compiling at plan time, so a source Janino rejects declines to Flink's Calc instead of failing the job, and both the loading and evaluation paths name the user function in the failure that reaches the native side. The Arrow batch boundary, the native expression and the per-subtask wrapper retention are unchanged.

Three integration cases close the limitations the generated invocation shipped with. A function declaring two invokable eval overloads is admitted now that Flink resolves the call, but that was established only for the legacy registration path. The modern path resolves through TypeInference and reaches a different generator entry, so a case there pins the resolution by the value it returns. Rewriting the call's operands to input references loses the flag marking an operand as a literal. A call mixing a column and a literal runs natively, so the loss does not reach behaviour. A primitive parameter fed SQL NULL fails rather than returning a substituted default: Flink's generated call unboxes and raises before eval is entered. The native path fails the same way, inside the generated invoker rather than by declining. The two failures cannot be compared by identity, because the bridge carries the exception across as text and the function is not a stack frame on Flink's side, so the case asserts what is available and records why identity is not. There is no switch to disable Auron, so each comparison run appends a call whose return type declines the whole Calc and drops the companion column. That rests on declining being all-or-nothing per Calc, which the test class records, since a per-expression decline would leave the comparisons measuring Auron against itself. No lifecycle case is added; an existing one already forces several drain cycles and asserts a single open.

Copilot

Pull request overview

This PR adds Flink support for per-expression fallback of user-defined scalar functions by packaging the UDF call into an engine-neutral “UDF wrapper” native expression that upcalls into the JVM, allowing the rest of a Calc/projection to remain native. It also introduces Flink subtask-scoped wrapper retention (so UDF open/close run once per subtask), expands planner/runtime test coverage for this path, and fixes a JNI thread context classloader regression affecting multi-job JVM reuse.

Changes:

Renames the existing Spark-specific UDF wrapper node to an engine-neutral wrapper (protobuf tag unchanged) and wires it through the native planner/runtime.
Implements Flink UDF fallback: plan-time generation of a Flink codegen invoker, runtime compilation/instantiation, Arrow C Data Interface argument/result transfer, and subtask-scoped wrapper retention via FlinkAuronTaskContext.
Adds extensive unit/integration tests for wrapper admission rules, generated-code behavior, Arrow boundary correctness, lifecycle, and classloader/thread-context behavior; includes a JNI attach/classloader fix.

Reviewed changes

Copilot reviewed 34 out of 34 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala	Renames Spark-specific wrapper builder to engine-neutral wrapper node.
spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala	Adds Spark test asserting scalar subquery evaluation stays native.
native-engine/datafusion-ext-exprs/src/udf_wrapper.rs	Renames Spark wrapper expr to engine-neutral and routes JNI calls via `AuronUDFWrapperContext`.
native-engine/datafusion-ext-exprs/src/spark_scalar_subquery_wrapper.rs	Updates scalar subquery wrapper to use the engine-neutral UDF wrapper expr.
native-engine/datafusion-ext-exprs/src/lib.rs	Re-exports the renamed `udf_wrapper` module.
native-engine/auron-planner/src/planner.rs	Switches planner decoding to the renamed `UdfWrapperExpr` protobuf case.
native-engine/auron-planner/proto/auron.proto	Renames protobuf oneof field/message to `udf_wrapper_expr` (tag 10000 unchanged).
native-engine/auron-jni-bridge/src/jni_bridge.rs	Avoids overriding Java thread context classloader for already-attached JVM threads.
auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperatorTest.java	Adds operator-level tests for task-context publication, wrapper retention, and close semantics.
auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/jni/FlinkAuronAdaptorTest.java	Adds adaptor tests for thread context and wrapper-context creation semantics.
auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/functions/GeneratedUdfTestSupport.java	Test utility for building stand-in generated invoker payloads.
auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContextTest.java	Unit tests for Arrow-boundary UDF wrapper execution and error/leak behavior.
auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/functions/FlinkAuronTaskContextTest.java	Unit tests for subtask-scoped wrapper registry reuse/isolation and lifecycle.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/jni/FlinkAuronAdaptor.java	Implements Flink thread-context plumbing and UDF wrapper context retrieval.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java	Publishes/clears task context during runtime creation; eagerly builds wrappers in `open()`; closes context in `close()`.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/AuronPlanTreeRewriter.java	Adds generic protobuf traversal to collect UDF wrapper payloads for eager initialization.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkUDFPayload.java	Defines the serialized payload schema shared between planner and runtime.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java	Runtime wrapper executing generated invoker across Arrow C Data boundary.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronTaskContext.java	Subtask-scoped wrapper registry + thread-local publication channel for native workers.
auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/AuronGeneratedUDF.java	Interface implemented by planner-generated invoker classes.
auron-flink-extension/auron-flink-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecCalcTest.java	Adds planner test ensuring a Calc with a user scalar function stays native.
auron-flink-extension/auron-flink-planner/src/test/java/org/apache/auron/flink/table/runtime/AuronFlinkUDFITCase.java	End-to-end IT coverage for Flink UDF wrapper behavior, lifecycle, and fallbacks.
auron-flink-extension/auron-flink-planner/src/test/java/org/apache/auron/flink/table/planner/converter/RexCallConverterTest.java	Adds converter tests for UDF detection and short-circuit AND/OR emission.
auron-flink-extension/auron-flink-planner/src/test/java/org/apache/auron/flink/table/planner/converter/FlinkUDFFallbackBuilderTest.java	Unit tests for wrapper admission rules and payload properties.
auron-flink-extension/auron-flink-planner/src/test/java/org/apache/auron/flink/table/planner/converter/FlinkUDFCodeGeneratorTest.java	Tests generated-invoker compilation and correctness across Arrow boundary.
auron-flink-extension/auron-flink-planner/src/test/java/org/apache/auron/flink/table/planner/converter/FlinkNodeConverterFactoryTest.java	Ensures exceptions in `isSupported` are contained as “decline” rather than escaping.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/RexCallConverter.java	Adds UDF wrapper conversion and emits short-circuit AND/OR nodes when needed.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/NativePlanFusionBuilder.java	Documents classloader requirement for UDF invoker compilation.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkUDFFallbackBuilder.java	Implements UDF call admission + wrapper node construction + wrapper-subtree scanning.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkUDFCodeGenerator.java	Generates invoker via Flink codegen, validates it compiles, and returns source+refs.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkNodeConverterFactory.java	Treats `UnsupportedNodeException` as a normal decline (debug log) instead of warn.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkNodeConverter.java	Adds `UnsupportedNodeException` as an explicit decline signal.
auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/ConverterContext.java	Adds per-conversion UDF wrapper ordinal counter for payload uniqueness.
auron-core/src/main/java/org/apache/auron/jni/AuronAdaptor.java	Updates thread-context Javadoc to be engine-neutral task state, not just classloader.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

…fix a dangling javadoc Reading the wrapper payload through a duplicate leaves the buffer the native side passes at its original position and limit, so the entry does not depend on the caller treating the buffer as consumed. The input-schema javadoc had come to sit above a later method's own javadoc rather than above the accessor it describes, leaving that accessor undocumented.

…once per assertion The row count came from re-collecting the DataFrame the operator check had already collected. Spark 4 rebuilds query stages on every collect of the same adaptive plan, where 3.5 returns the finalized plan unchanged, so the second execution reached an assertion the first had not. Taking the count from its own execution leaves the checked plan executed once, and the plan assertion reads the executed plan without running it again.

Copilot

Pull request overview

Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/jni/FlinkAuronAdaptor.java:87

When setThreadContext is called with a ClassLoader (i.e., no FlinkAuronTaskContext published for this runtime), the existing ThreadLocal task context is left untouched. If the same worker thread previously ran with a published FlinkAuronTaskContext, subsequent getAuronUDFWrapperContext() calls on that thread could incorrectly resolve wrappers through stale per-subtask state.

        } else {
            Thread.currentThread().setContextClassLoader((ClassLoader) context);
        }

weiqingy · 2026-08-30T23:34:04Z

The Spark 4.0/4.1/4.2 failures were scalar subquery is evaluated natively, a test this PR adds. Fixed in 28e8976, and the diagnosis turned up a separate pre-existing bug, now #2491.

The test collected the same DataFrame twice: checkSparkAnswerAndOperator already collects it, and the row-count assertion collected it again. Spark 3.5 short-circuits the second collect, since getFinalPhysicalPlan() returns early when isFinalPlan is set. Spark 4 re-enters createQueryStages(..., firstRun = true) on every collect, and the comment on that branch names this case directly ("e.g, when we do df.collect multiple times"). The fix takes the count from its own execution so the checked plan runs exactly once.

Why the second collect fails rather than just redoing work is separate from this PR, and is what #2491 covers: AuronConverters builds fresh nodes without copying logicalLink, and Shims.setLogicalLink has had no callers since 441a1a2. Harmless on Spark 3; on Spark 4 a fully native, exchange-free plan carries no logical link and the assertion in setLogicalLinkForNewQueryStage fires. Neither of those files is touched here.

I kept that out of this PR deliberately. tryConvert is the single funnel for the 34 plan types convertSparkPlan handles, so restoring the call changes what AQE sees for every converted node on every supported Spark version, and it wants its own test matrix rather than riding along in a Flink UDF change. I will take a first pass at it on #2491.

Also addressed the two review comments in 0e1de00: reading the wrapper payload through byteBuffer.duplicate() so the caller's buffer is not consumed, and moving the input-schema javadoc back onto getInputType().

weiqingy · 2026-08-31T03:58:10Z

Hi @Tartarus0zm, could you please help review this PR when you get a chance? Thanks!

Tartarus0zm · 2026-09-04T03:51:00Z

hi @richox could you please help review this PR when you get a chance?
Because this involves a change to the UDF execution framework, and it has some relevance to Spark.

Tartarus0zm

hi @weiqingy thanks for your contribution!
I've left two minor comments.

Tartarus0zm · 2026-09-04T03:45:27Z

+     * @throws IllegalStateException if the payload cannot be turned into a wrapper
+     */
+    public AuronUDFWrapperContext getOrCreateWrapper(byte[] payload) {
+        return wrappers.computeIfAbsent(ByteBuffer.wrap(payload), key -> {


Is caching necessary here? Does Flink cache UDF objects? If a custom UDF maintains user-managed internal state, could reusing it here cause data correctness issues?

Good question. I read Flink 1.18.1's codegen, and Flink shares UDF instances more than I expected.

Flink names the generated field from the function's identity, not the call site, and collects the member, open and close statements in sets. So two calls of the same function in one Calc share a single instance, opened once and closed once per subtask.

That means we are stricter than Flink here, not looser. Two call sites produce different payload bytes, so they get two wrappers and two instances where Flink gives one. The generated className carries an incrementing counter, and that is what actually separates them. nodeOrdinal is belt and braces.

The cache turns out to matter for lifecycle rather than speed. FlinkAuronTaskContext.close() is the only path that reaches a user function's close(), since the native side only binds eval, and the runtime is rebuilt each drain cycle. Drop the cache and open() runs once per drain while close() never runs, so anything holding a resource leaks.

Your question did find a real bug though, just in the docs rather than the code. The javadoc on nodeOrdinal claimed the opposite of what Flink does. Fixed in dfb170d4.

Does that answer it, or is there a case you had in mind that I have missed?

Following up here, because my earlier answer was incomplete.

I said we end up stricter than Flink rather than looser. That is true, but it stops short of the part that matters: stricter still means a different answer. Two call sites of one function in a Calc get separate instances, and separate open/close pairs per subtask, so a function carrying state across invocations returns something different from vanilla Flink.

Sharing one instance would not fix it. Auron evaluates a call site across a whole batch where Flink evaluates both call sites per row, so a shared instance produces a third answer rather than Flink's. With a counting function: Flink gives (1,2) (3,4) (5,6), we give (1,1) (2,2) (3,3), a shared instance would give (1,4) (2,5) (3,6).

That leaves this divergence, or declining a Calc that calls one function at several sites, which would cost the stateless majority for something that only bites stateful functions.

I have qualified the PR description, and the two-call-site test now executes both sides and asserts they differ, so it is pinned rather than described.

Would you rather we kept the divergence, or declined the repeated-call-site case?

I think the behavior needs to be aligned with Flink — otherwise, inconsistent results would be disastrous for users.

Tartarus0zm · 2026-09-04T06:26:37Z

        }
+        // A user-defined scalar function is matched the same way, by operator identity. The
+        // admission checks run in convert, which may still decline and fall the whole Calc back.
+        if (FlinkUDFFallbackBuilder.userScalarFunctionOf(call).isPresent()) {


this conversion path is shared by both the standalone Calc and the Kafka source-Calc fusion. AuronOperatorFusionProcessor pushes the native plan containing the UDF wrapper down to AuronKafkaSourceFunction, but at runtime the Kafka source does not create, publish, or close a FlinkAuronTaskContext when it creates the AuronCallNativeWrapper. As a result, FlinkAuronAdaptor#getAuronUDFWrapperContext throws "no Flink task context is published on thread ..." directly when there is no task context on the current thread. This means queries that combine Kafka source-Calc fusion with UDFs fail at runtime, and this path also skips the UDF's open() / close() lifecycle — turning previously fallback-able existing queries into runtime exceptions. I'd suggest completing task context and UDF lifecycle management in the Kafka source runtime.

You're right. On master a UDF call made isSupported return false, so the whole Calc fell back to Flink and fusion never saw it. Here it gets admitted, fuses into the source, and dies on the first row.

I took your suggestion rather than gating the fusion, so a fused UDF now works. Fixed in fbcaa4f7.

The native runtime reads the thread context once when it is created and installs it on its workers, so the source publishes the context around that call. The context is built in open() along with its wrappers, which is what gives each function one open() and one close() per subtask. That was your lifecycle point, and it is the part that made this worth doing properly.

Teardown was the fiddly bit. It has to live in run() rather than close(), since closing the wrapper from the task thread while the source thread is waiting on a batch would free a receiver it is blocked on.

One note since it reaches past UDFs: the source publishes the context unconditionally, matching what FlinkAuronCalcOperator does for every native Calc, so its native workers get the user-code classloader whether or not a UDF is involved.

For tests there is an IT running a UDF over a fused source, plus lifecycle tests covering open and close once per subtask and the close path when the source never ran.

Does that match what you had in mind?

… and pin the node ordinal The javadoc on FlinkUDFPayload.nodeOrdinal stated Flink's behaviour backwards. It claimed Flink hands each call site its own copy of a user function, and that two calls of one function yield byte-equal payloads. Neither holds. Flink derives the generated field name from the function's identity rather than the call site and collects the member, open and close statements in sets, so two call sites of one function share a single instance. Payloads already differ per call site because the generated class name carries an incrementing counter, which makes the ordinal reinforcement rather than the operative mechanism. The class javadoc on FlinkAuronTaskContext rested on the same premise. Add a registry test holding the function, argument types and return type fixed so the ordinal is the only difference between two payloads. The existing test varies the user function at ordinal zero, so it passes whether or not the ordinal reaches the serialized form.

… source runtime A Calc fused into an Auron source carried its UDF wrapper into a runtime that published no FlinkAuronTaskContext, so the wrapper callback found no context on the thread and raised IllegalStateException on the first row. The native runtime reads the JVM thread context once when it is created and installs it on every worker of its pool, so the source publishes the context around the call that builds the runtime and clears it afterwards. The context itself is built in open(), where its wrappers are pre-built too, giving each user function one open() and one close() per subtask. Teardown belongs to run(), which is the only place where no evaluation can still be outstanding: closing the wrapper from the task thread while the source thread waits on a batch would free a receiver that thread is still blocked on. run() closes the wrapper and then the context, and close() closes the context only when run() never claimed it, so a user function built in open() is still closed when the source never ran. Covered by an integration test running a user function over a fused source, which asserts the fused plan carries no Calc operator alongside the row set, and by lifecycle tests pinning that a user function is opened and closed once per subtask and that the close path survives a source that never ran.

Copilot

🟡 Changes recommended

Two unresolved critical findings remain: registered scalar UDF definitions are rejected (3 votes), and repeated call sites do not share the UDF instance (1 vote).

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkUDFFallbackBuilder.java:219

Including a fresh nodeOrdinal in every payload makes repeated calls to the same registered ScalarFunction resolve to different entries in the byte-keyed task registry, so each call site gets its own function instance and its own open/close and mutable state. Flink's generated path instead names the reusable function field from the function identifier and reuses one instance per subtask; consequently stateful UDFs can return different results here (the added call-count test deliberately demonstrates that divergence), contradicting the stated unchanged semantics. Share the wrapper for the same Flink function identity, or explicitly treat this as a behavior change and document it.

                    udf.getClass().getName(),
                    context.nextUdfWrapperOrdinal()));

Files reviewed: 38/38 changed files
Comments generated: 2
Review effort level: Lite

…rgence and pin it by execution Two call sites of one function in a Calc resolve to separate wrappers, so each deserializes its own function instance and receives its own open and close per subtask. Flink's generated Calc shares a single instance across those sites. For a function whose result depends on state carried across invocations, that is a different answer, which the payload javadoc now says rather than stopping at the narrower observation that the wrappers do not share. The two-call-site test asserted only the Auron result, leaving the comparison with Flink to prose. It now runs the same query with the Calc declined and asserts both sides, so the divergence is established by execution and any future change on either side fails the test.

Copilot

🟡 Changes recommended

Unresolved critical and moderate findings block approval, with an additional message-quality nit.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (3)

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java:203

This evaluation boundary also catches OutOfMemoryError, StackOverflowError, and ThreadDeath and re-labels them as a normal UDF failure. Let fatal/control-flow errors propagate so Flink can apply its normal cancellation/recovery behavior, while wrapping only ordinary exceptions and the linkage errors that need the UDF name.

        } catch (Throwable t) {

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java:122

If the generated invoker's open acquires a resource and then throws, this constructor exits without retaining the invoker, so FlinkAuronTaskContext.close() can never call its close(). That leaks partial initialization on the exact failure path that leaves the operator's context unable to clean up the failed wrapper; close the generated invoker in a catch around open (preserving the original failure and suppressing any close failure) before rethrowing.

        generated.open(runtimeContext);

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/jni/FlinkAuronAdaptor.java:110

The exception text is grammatically broken ("a native runtime an Auron Flink operator created") and makes the missing-context condition harder to understand. Use wording that clearly identifies the required owner of the runtime.

                    + "; a UDF wrapper is only reachable from a native runtime an Auron Flink "
                    + "operator created");

Files reviewed: 38/38 changed files
Comments generated: 2
Review effort level: Lite

… and let fatal errors pass A source whose close() ran before run() started had already taken and closed the task context, so the claim in run() returned null. The run published that null and built a native runtime anyway, and the first user function callback then failed with the adaptor's missing-context error. Abort the run instead: with no context there is nothing to serve a wrapper, and the teardown that took it has already finished. The two broad catches around generated-invoker loading and evaluation turned every Throwable into an IllegalStateException, including the JVM-fatal ones. Rethrow those first. The catches stay broad because loading a generated class reports its failures as LinkageError and ExceptionInInitializerError, which must still be wrapped with the function name.

Copilot

🟡 Changes recommended

A critical task-context cleanup issue and a moderate current-registration fallback issue remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (2)

auron-flink-extension/auron-flink-planner/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkUDFFallbackBuilder.java:148

BridgingSqlFunction.getDefinition() does not return the user function instance for the current registration path: Flink wraps it in ScalarFunctionDefinition, which is a FunctionDefinition rather than a ScalarFunction. Consequently this instanceof ScalarFunction check returns false for createTemporarySystemFunction calls, so RexCallConverter never invokes the fallback builder and the Calc still falls back wholesale; the advertised current-API integration case cannot exercise the new path. Extract getScalarFunction() from ScalarFunctionDefinition (while retaining the direct ScalarFunction case for inline definitions) before returning the optional.

        return definition instanceof ScalarFunction ? Optional.of((ScalarFunction) definition) : Optional.empty();

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/jni/FlinkAuronAdaptor.java:110

When a wrapper callback arrives without a published task context, this is the diagnostic users receive, but “a native runtime an Auron Flink operator created” is grammatically ambiguous and makes the ownership condition hard to understand. Rephrase it as “a native runtime created by an Auron Flink operator.”

                    + "; a UDF wrapper is only reachable from a native runtime an Auron Flink "
                    + "operator created");

Files reviewed: 38/38 changed files
Comments generated: 1
Review effort level: Lite

…ing what they held Each resource in the operator's close() was released by closing it and then nulling the field. A close that threw skipped the assignment, so the field kept its reference. For the task context that retains the user-code classloader and the runtime context on the operator, and a user function's close() can throw: the context clears its registry in a finally but rethrows the first failure. Take each field into a local and clear it before closing, so the reference is gone whichever way the close goes. The ordering and its nesting are unchanged; only the point at which each field is cleared moved. A second close of an already-closed resource is no longer reachable either.

Copilot

🟡 Changes recommended

The critical nested/stateful UDF evaluation-order issue remains unresolved; the diagnostic nit also remains.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java:123

generated.open(runtimeContext) is outside the diagnostic wrapping used for compilation and evaluation. If a UDF's open(FunctionContext) throws, the constructor propagates the raw exception and FlinkAuronTaskContext replaces it with the generic error creating Flink UDF wrapper context, so the native-side error does not identify the failing UDF despite this class's constructor contract. Wrap this call with an exception naming udfClassName.

        generated.open(runtimeContext);

Files reviewed: 38/38 changed files
Comments generated: 1
Review effort level: Lite

An operand of a user function call is converted the same way as any other, so a nested call becomes a wrapper inside a wrapper and recurses without a depth limit. That shape had no coverage. Auron evaluates each parameter across the batch before invoking the outer function per row, where Flink interleaves the two per row. Row order is preserved on both sides, so the inner call's output sequence is the same and an outer function that depends only on its arguments returns the same rows. The two new cases pin that for a stateful call nested inside a pure one and for the reverse, each running the query a second time with the Calc declined and asserting the two agree.

Copilot

🟡 Changes recommended

One critical and one moderate finding remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java:253

open() registers this exporter in the JVM-global JniBridge resource map, but a no-row subtask never starts the native FFI reader, so nothing removes that entry. Closing the Java exporter here does not remove it; repeated empty subtasks therefore retain closed exporters (and their allocator/user-code references) in the static map. Remove resourceId during operator cleanup (a no-op when the native reader already consumed it).

                    FlinkArrowFFIExporter batchExporter = exporter;
                    exporter = null;
                    try {
                        if (batchExporter != null) {
                            batchExporter.close();

Files reviewed: 38/38 changed files
Comments generated: 1
Review effort level: Lite

…, not owned The two ArrowArray handles in eval wrap addresses the native caller owns. Wrapping an address neither takes ownership of it nor releases it on close, and the import moves the argument buffers into the params root before the block ends, so the declaration order of these resources carries no lifetime meaning. Reading the order as ownership is the wrong conclusion to reach, and nothing at the call sites says otherwise.

Copilot

🔵 Needs a closer look

One or more issues must be addressed before approval.

Review details

Suppressed comments (3)

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunction.java:456

Cancellation can still race this startup: after claimTaskContextForRun() marks nativeRuntimeStarted, close() deliberately stops taking the context, but run() does not check isRunning before constructing AuronCallNativeWrapper. If cancellation/close occurs in this window, the source has already closed its Kafka resources while run() starts and drives a native plan anyway (callbacks merely discard rows because isRunning is false). Guard the claim/start transition with the same lock, or otherwise make close() wait for or abort startup, so no runtime is created after cancellation.

            FlinkAuronTaskContext.setCurrent(runTaskContext);
            try {
                wrapper = new AuronCallNativeWrapper(

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java:197

The Arrow C Data wrappers must outlive the imported/exported roots: the native C array structs are still needed while Data.importIntoVectorSchemaRoot releases paramsRoot, and while the exported outputRoot is transferred. Because try-with-resources closes in reverse declaration order, this currently closes exportArray/importArray before the roots; declare the two ArrowArray resources before the roots (as the wrapper tests do) to keep the C structs alive through root cleanup.

        try (VectorSchemaRoot paramsRoot = VectorSchemaRoot.create(paramsArrowSchema, allocator);
                VectorSchemaRoot outputRoot = VectorSchemaRoot.create(outputArrowSchema, allocator);
                ArrowArray importArray = ArrowArray.wrap(importFFIArrayPtr);
                ArrowArray exportArray = ArrowArray.wrap(exportFFIArrayPtr)) {

native-engine/auron-jni-bridge/src/jni_bridge.rs:56

This avoids installing the stale loader on Java-created threads, but the fallback still keeps JavaClasses.classloader as a global JNI reference in the process-wide JNI_JAVA_CLASSES OnceCell. Since that reference is captured from the first job's context classloader, every later-job classloader cleanup can still be defeated by the first user-code loader being retained for the TaskManager lifetime. Please remove this process-wide strong reference (or otherwise make the native-thread bootstrap use a non-retaining/current loader) now that runtime startup propagates the current loader explicitly.

Files reviewed: 38/38 changed files
Comments generated: 0 new
Review effort level: Lite

weiqingy · 2026-09-15T04:04:26Z

@Tartarus0zm thanks for the review. Here is what moved since, as the PR has grown a fair bit.

The Kafka source-Calc fusion gap was a real regression this PR introduced, and I took your suggestion rather than gating the fusion. The source runtime now builds, publishes and closes a FlinkAuronTaskContext, so a fused UDF runs natively (fbcaa4f7). Teardown sits in run() rather than close(), since closing the wrapper from the task thread while the source thread waits on a batch would free a receiver it is blocked on.

On caching I owe you a correction. My first answer said we end up stricter than Flink rather than looser. True, but it stops short: stricter still means a different answer for a function carrying state across invocations. The PR description now states the divergence as the batch-versus-row mechanism, and a test runs both sides and asserts they differ rather than describing it in prose (dfb170d4, 6d9df4db).

Copilot raised five more. Three were real and are fixed: a close()/run() race leaving a null task context, catch (Throwable) swallowing JVM-fatal errors at the native boundary (both a12a97c7), and a throwing UDF close() pinning the user-code classloader (314564b6). Two did not survive checking, with the evidence in each thread: nested UDF calls return Flink's rows whenever the outer function depends only on its arguments (measured, tests in fdad6c18), and a wrapped ArrowArray is not released on close, so the C-Data declaration order carries no lifetime meaning (767a594b).

One question is still open: where a Calc calls the same function at more than one site, is the divergence worth keeping, or should that shape decline instead? Sharing an instance does not reach parity, since Auron evaluates a call site across a batch where Flink evaluates both sites per row, so a shared instance gives a third answer rather than Flink's. I kept the divergence and documented it, since declining would cost the stateless majority.

Tartarus0zm · 2026-09-20T11:48:52Z

hi @weiqingy At the moment, there are two issues:

When the same UDF has multiple call sites, the results are inconsistent with Flink.
In nested UDF scenarios, the invocation order is inconsistent with Flink.

It looks like these two issues need a redesigned implementation, so we'll need further discussion.

weiqingy added 9 commits August 30, 2026 15:04

Copilot AI lite review requested due to automatic review settings August 30, 2026 22:43

github-actions Bot added spark native core flink labels Aug 30, 2026

Copilot started reviewing on behalf of weiqingy August 30, 2026 22:44 View session

Copilot AI reviewed Aug 30, 2026

View reviewed changes

Comment thread ...link-extension/auron-flink-runtime/src/main/java/org/apache/auron/jni/FlinkAuronAdaptor.java Outdated

Comment thread ...k-planner/src/main/java/org/apache/auron/flink/table/planner/converter/ConverterContext.java

weiqingy added 2 commits August 30, 2026 16:14

Copilot AI review requested due to automatic review settings August 30, 2026 23:23

Copilot started reviewing on behalf of weiqingy August 30, 2026 23:23 View session

Copilot AI reviewed Aug 30, 2026

View reviewed changes

weiqingy mentioned this pull request Aug 30, 2026

Converted native plans drop logicalLink, breaking repeated DataFrame execution on Spark 4 #2491

Open

Tartarus0zm requested review from Tartarus0zm and richox September 2, 2026 04:18

Tartarus0zm requested changes Sep 4, 2026

View reviewed changes

weiqingy added 2 commits September 13, 2026 16:58

Copilot AI review requested due to automatic review settings September 14, 2026 02:33

Copilot started reviewing on behalf of weiqingy September 14, 2026 02:33 View session

Copilot AI reviewed Sep 14, 2026

View reviewed changes

Comment thread ...er/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkUDFFallbackBuilder.java

Comment thread ...sion/auron-flink-runtime/src/main/java/org/apache/auron/flink/functions/FlinkUDFPayload.java

Copilot AI review requested due to automatic review settings September 14, 2026 05:32

Copilot started reviewing on behalf of weiqingy September 14, 2026 05:32 View session

Copilot AI reviewed Sep 14, 2026

View reviewed changes

Comment thread ...k-runtime/src/main/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunction.java

Comment thread ...link-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java

Copilot AI review requested due to automatic review settings September 14, 2026 16:33

Copilot started reviewing on behalf of weiqingy September 14, 2026 16:44 View session

Copilot AI reviewed Sep 14, 2026

View reviewed changes

Comment thread ...nk-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java Outdated

Copilot AI review requested due to automatic review settings September 15, 2026 01:37

Copilot started reviewing on behalf of weiqingy September 15, 2026 01:37 View session

Copilot AI reviewed Sep 15, 2026

View reviewed changes

Comment thread ...er/src/main/java/org/apache/auron/flink/table/planner/converter/FlinkUDFFallbackBuilder.java

Copilot AI review requested due to automatic review settings September 15, 2026 02:25

Copilot started reviewing on behalf of weiqingy September 15, 2026 02:26 View session

Copilot AI reviewed Sep 15, 2026

View reviewed changes

Comment thread ...link-runtime/src/main/java/org/apache/auron/flink/functions/FlinkAuronUDFWrapperContext.java

Copilot AI review requested due to automatic review settings September 15, 2026 02:51

Copilot started reviewing on behalf of weiqingy September 15, 2026 02:51 View session

Copilot AI reviewed Sep 15, 2026

View reviewed changes

Conversation

weiqingy commented Aug 30, 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?

Are there any user-facing changes?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

weiqingy commented Aug 30, 2026

Uh oh!

weiqingy commented Aug 31, 2026

Uh oh!

Tartarus0zm commented Sep 4, 2026

Uh oh!

Tartarus0zm left a comment

Choose a reason for hiding this comment

Uh oh!

Tartarus0zm Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

weiqingy Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

weiqingy Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

Tartarus0zm Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Tartarus0zm Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

weiqingy Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (3)

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (2)

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

weiqingy commented Aug 30, 2026 •

edited

Loading