You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/11 03:46:35 UTC
(datafusion-comet) branch main updated: fix: check native UDF types at registration and explain planning-time evaluation (#6777)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 78b0eb0df3 fix: check native UDF types at registration and explain planning-time evaluation (#6777)
78b0eb0df3 is described below
commit 78b0eb0df342e631a619c4951fd9b531f0375a3e
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sun Oct 11 02:38:56 2026 +0000
fix: check native UDF types at registration and explain planning-time evaluation (#6777)
Registering a native UDF with a type Comet has no native representation
for, such as a user-defined type, succeeded, and every call then fell
back to Spark, which cannot evaluate it. CometNativeUDF.register now
refuses such a signature before it loads the library.
Spark also evaluates a call itself while planning a query: over local
data, in a filter on partition columns, and to sample the keys of a
global sort. The error a native UDF raises there said that Comet had not
taken the operator holding the call, which sent users to the explain
output for a fallback that never happened. The error, the scaladoc and
the Rust UDF guide now list those cases, and the guide says to sort on a
column holding a UDF's result rather than on the call.
Split out of #6697.
---
docs/source/user-guide/latest/rust_udfs.md | 17 +++++++++++
.../org/apache/comet/udf/CometNativeUDF.scala | 20 +++++++++----
.../org/apache/comet/udf/CometUdfExceptions.scala | 11 ++++---
.../scala/org/apache/comet/udf/NativeUdfCall.scala | 7 +++--
.../org/apache/comet/CometNativeUdfSuite.scala | 35 ++++++++++++++++++++++
5 files changed, 76 insertions(+), 14 deletions(-)
diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md
index ba98cd6acb..147a2069c9 100644
--- a/docs/source/user-guide/latest/rust_udfs.md
+++ b/docs/source/user-guide/latest/rust_udfs.md
@@ -158,6 +158,7 @@ CometNativeUDF.register(
Registration loads the library on the driver and verifies that a function with that name exists, so
a bad path or a missing function fails immediately with a clear error rather than at execution time.
+It also checks that Comet can carry the argument and return types natively.
The function is then callable from SQL or the DataFrame API like any other:
@@ -181,6 +182,19 @@ way it would for any other shared object, through `LD_LIBRARY_PATH` on Linux and
paths may be loaded, so it makes no difference to the trust decision described under
[Limitations](#limitations).
+### When Spark evaluates the call
+
+A call runs only inside Comet's native execution. Spark cannot run the function, so wherever Spark
+has to evaluate the call itself, the query fails with `CometUdfNotEvaluatedException`. That happens:
+
+- when Comet does not take the operator holding the call, for example because another expression
+ in it is not supported, or because the call is an argument of an ordinary Scala or Java UDF. The
+ query's extended explain output gives the reason.
+- while Spark plans the query: over local data such as `VALUES`, which the optimizer evaluates
+ eagerly, in a filter on partition columns, which Spark evaluates to prune partitions, and in the
+ sort keys of a global sort, which Spark evaluates on a sample of rows to choose range bounds. To
+ sort on a UDF's result, select it as a column and sort on that column instead.
+
## Return types
A UDF does not have one fixed return type. Its `return_field` is called with the actual argument
@@ -277,6 +291,9 @@ the first panic, and the SDK cannot prevent it.
This feature is at an early stage. The current limitations are:
- **Scalar functions only.** Aggregate, window, and table functions are not supported.
+- **Spark cannot evaluate a call.** A query fails wherever Spark would have to run the function
+ itself, which includes a few places while it plans the query. See
+ [When Spark evaluates the call](#when-spark-evaluates-the-call).
- **Immutable functions only.** A UDF must return the same output for the same input. Comet plans
every Rust UDF with DataFusion's `Volatility::Immutable`, so a function that reads a clock, draws
from an RNG, or carries state across batches may be folded at plan time, evaluated once and
diff --git a/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala
index 5c2a960830..0feb960592 100644
--- a/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala
+++ b/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala
@@ -23,6 +23,7 @@ import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.comet.CometUdfErrors
import org.apache.spark.sql.types.DataType
+import org.apache.comet.serde.QueryPlanSerde
import org.apache.comet.shims.ShimSessionFunctionRegistry
/**
@@ -46,9 +47,10 @@ object CometNativeUDF {
/**
* Register a single native UDF with an explicit signature.
*
- * Validates the library on the driver (loads it, confirms a UDF named `name` exists), then
- * installs `name` as a temporary function of the session, as `spark.udf.register` would: other
- * sessions do not see it, and registering another function under the same name replaces it.
+ * Checks that Comet can carry the argument and return types natively and validates the library
+ * on the driver (loads it, confirms a UDF named `name` exists), then installs `name` as a
+ * temporary function of the session, as `spark.udf.register` would: other sessions do not see
+ * it, and registering another function under the same name replaces it.
*
* Executors need no registration: the library path travels with the plan in the
* `NativeScalarUdf` proto, and each executor loads the library itself on first use. The path
@@ -58,8 +60,9 @@ object CometNativeUDF {
* argument types differ, other than in nullability, and inserts no casts, so cast the arguments
* in the query instead.
*
- * Spark cannot evaluate the call itself: if Comet does not take the operator holding it, the
- * query fails.
+ * Spark cannot evaluate the call itself, so a query fails where Spark would have to: in an
+ * operator Comet does not take, and while planning over local data, in a filter on partition
+ * columns, or to sample the keys of a global sort.
*
* `deterministic` must be `true`. Comet plans every imported kernel as immutable, so a
* nondeterministic UDF cannot yet be expressed; passing `false` fails here rather than silently
@@ -84,6 +87,13 @@ object CometNativeUDF {
"eliminated as a common subexpression. " +
"See https://github.com/apache/datafusion-comet/issues/5249")
}
+ // With a type Comet cannot represent natively, every call would fall back to Spark, which
+ // cannot evaluate it.
+ (inputTypes :+ returnType).find(QueryPlanSerde.serializeDataType(_).isEmpty).foreach { t =>
+ throw new IllegalArgumentException(
+ s"native UDF '$name': Comet has no native representation for type " +
+ s"${t.catalogString}, so no call to it could run.")
+ }
validateLibrary(libraryPath, name)
ShimSessionFunctionRegistry
.functionRegistry(spark)
diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala b/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala
index eeda2447b1..44e0554653 100644
--- a/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala
+++ b/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala
@@ -21,12 +21,11 @@ package org.apache.comet.udf
import org.apache.comet.CometRuntimeException
-/**
- * Thrown when Spark itself evaluates a call to a UDF registered with Comet, which means Comet did
- * not take the operator holding it.
- */
+/** Thrown when Spark itself evaluates a call to a native UDF, which only Comet can evaluate. */
class CometUdfNotEvaluatedException(name: String)
extends CometRuntimeException(
s"UDF '$name' is registered with Comet and runs only inside Comet's native execution, but " +
- "Spark evaluated it, which means Comet did not take the operator holding the call. The " +
- "query's extended explain output gives the reason that operator fell back to Spark.")
+ "Spark evaluated it. Either Comet did not take the operator holding the call, and the " +
+ "query's extended explain output gives the reason, or Spark evaluated the call itself " +
+ "while planning: over local data, in a filter on partition columns, or to sample the " +
+ "keys of a global sort.")
diff --git a/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala b/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala
index 22c71039a2..d1f444cc4b 100644
--- a/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala
+++ b/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala
@@ -29,9 +29,10 @@ import org.apache.spark.sql.types.DataType
* in the session builds one for each call it resolves.
*
* Comet replaces it with a `NativeScalarUdf` that runs the function named `name` in the library
- * at `libraryPath`. Spark has no way to evaluate it, so `eval` throws. That happens only when
- * Comet does not take the operator holding the call, and failing there makes the fallback visible
- * rather than silent.
+ * at `libraryPath`. Spark has no way to evaluate it, so `eval` throws, which makes a fallback
+ * visible rather than silent. Spark evaluates a call where Comet does not take the operator
+ * holding it, and while planning in a few places: over local data, in a filter on partition
+ * columns, and to sample the keys of a global sort.
*
* `argumentTypes` is the registered signature, one type per child. Spark's analyzer checks each
* call against it, disregarding nullability, and inserts no casts.
diff --git a/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala
index 9f36d18ff7..4bad30e2c9 100644
--- a/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala
@@ -446,6 +446,41 @@ class CometNativeUdfSuite extends CometTestBase {
}
}
+ test("Spark evaluating a native UDF while planning fails the query") {
+ // Spark's optimizer evaluates a projection over local data itself, so this fails even though
+ // Comet would take every operator.
+ CometNativeUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType)
+ val e = intercept[Exception] {
+ sql("SELECT add_one_c(x) FROM VALUES (1L), (2L) AS t(x)").collect()
+ }
+ assert(
+ stackTraceContains(e, "Spark evaluated the call itself while planning"),
+ s"unhelpful error: $e")
+ }
+
+ test("a global sort on a native UDF's result runs when the sort is on a column") {
+ // Spark evaluates a global sort's keys on a sample of rows to choose range bounds, so the
+ // sort has to be on a column holding the result rather than on the call itself.
+ CometNativeUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType)
+ val df = spark.range(0, 20, 1, 4).selectExpr("add_one_c(19 - id) AS y").orderBy("y")
+ assert(df.collect().map(_.getLong(0)).toSeq == (1L to 20L))
+ }
+
+ test("a type Comet cannot carry natively is refused at registration") {
+ // Earlier tests register `echo_c`, so drop it first to see that the refusal installs nothing.
+ sql("DROP TEMPORARY FUNCTION IF EXISTS echo_c")
+ val e = intercept[IllegalArgumentException] {
+ CometNativeUDF.register(
+ spark,
+ "echo_c",
+ libPath,
+ Seq(ObjectType(classOf[String])),
+ LongType)
+ }
+ assert(e.getMessage.contains("Comet has no native representation"), e.getMessage)
+ assert(!spark.catalog.functionExists("echo_c"))
+ }
+
// A native UDF call has to reach its serde wherever Comet converts expressions, not just in a
// projection. Filters, join conditions, grouping keys and window partitioning each route
// through a different Comet operator, and a regression in any one of them would show up only as
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org