diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md index ba98cd6acbb..147a2069c94 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 5c2a9608307..0feb9605922 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 eeda2447b1f..44e0554653d 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 22c71039a2d..d1f444cc4ba 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 9f36d18ff78..4bad30e2c96 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