fix: check native UDF types at registration and explain planning-time evaluation

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

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)

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.")

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.

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