sunchao commented on code in PR #4744:
URL: https://github.com/apache/datafusion-comet/pull/4744#discussion_r4043804056
##########
native/core/src/execution/planner.rs:
##########
@@ -3670,6 +3689,134 @@ impl PhysicalPlanner {
}
}
+ fn create_high_order_function_expr(
+ &self,
+ expr: &HigherOrderFunc,
+ input_schema: SchemaRef,
+ ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+ let udf = create_comet_hof_func(&expr.func_name,
&self.session_ctx.state())?;
+
+ // 1. Plan value args.
+ let value_args: Vec<Arc<dyn PhysicalExpr>> = expr
+ .value_args
+ .iter()
+ .map(|e| self.create_expr(e, Arc::clone(&input_schema)))
+ .collect::<Result<_, _>>()?;
+
+ // 2. Resolve lambda param field types via the UDF (mirrors runtime).
+ let param_fields = Self::resolve_lambda_param_fields(
+ &udf,
+ &expr.func_name,
+ &value_args,
+ expr.lambdas.len(),
+ input_schema.as_ref(),
+ )?;
+
+ // 3. Plan lambdas with resolved param fields.
+ let lambdas: Vec<Arc<dyn PhysicalExpr>> = expr
+ .lambdas
+ .iter()
+ .zip(¶m_fields)
+ .map(|(l, fields)| self.create_lambda_expr(l, &input_schema,
fields))
+ .collect::<Result<_, _>>()?;
+
+ // 4. NOTE: assumes value args precede lambdas (holds for
array_filter).
+ let mut args = value_args;
+ args.extend(lambdas);
+
+ Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema(
+ udf,
+ args,
+ &input_schema,
+ Arc::new(ConfigOptions::default()),
+ )?))
Review Comment:
[P2] Skip native predicate evaluation when there are no array elements
Could the native filter return the empty/null arrays before evaluating its
lambda when the flattened element count is zero? With ANSI enabled, Spark
safely returns `[]` for a nonnull empty array in partition 0:
```sql
SELECT filter(a, x -> (1 DIV spark_partition_id()) > 0)
FROM t;
```
I verified this with normal Spark 4.1.3 optimization and Parquet input,
including a null-array row. The exact-head serializer admits it as a native
HOF. `SparkPartitionIdBuilder` lowers the partition ID to a scalar literal, so
partition 0 gives the native predicate a scalar division by zero.
DataFusion's `evaluate_single_list_lambda` returns early for an all-null
batch, but otherwise invokes the lambda even when the flattened values have
length zero. A component reproduction using Comet's actual
`decimal_integral_div` therefore raises `DIVIDE_BY_ZERO` on an empty array,
whereas Spark never evaluates the predicate. The component dependencies and
source-equivalence checks are described in the review summary; this was not a
full Comet/JNI query run.
Please preserve the row null mask in the empty-elements fast path and add a
regression with at least one nonnull empty array. Using a nonfoldable predicate
such as this one avoids an unrelated failure from Spark constant folding.
##########
spark/src/main/scala/org/apache/comet/serde/CometHighOrderFunction.scala:
##########
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.serde
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
HigherOrderFunction, LambdaFunction => SparkLambdaFunction, NamedLambdaVariable
=> SparkNamedLambdaVariable}
+
+import org.apache.comet.CometConf
+import org.apache.comet.serde.CometHighOrderFunction.{containsJvmDispatch,
namedLambdaVariable2Proto}
+import org.apache.comet.serde.ExprOuterClass.{HigherOrderFunc, LambdaFunction,
NamedLambdaVariable}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal,
serializeDataType}
+
+/**
+ * Serializer that converts Spark higher-order functions (e.g. `filter`,
`transform`, `exists`)
+ * into Comet's protobuf representation.
+ *
+ * Path selection happens in [[convert]] and has exactly two outcomes, chosen
in order:
+ *
+ * Native HOF proto - when
`spark.comet.exec.higherOrderFunction.native.enabled` is set and
+ * [[highOrderFunction2Proto]] serializes the whole expression natively. The
method returns `None`
+ * (and the HOF falls to the next path) if the HOF is structurally invalid
(lambda functions must
+ * be `LambdaFunction`, lambda arguments must be `NamedLambdaVariable`) or if
a lambda body
+ * contains a dispatch-only subexpression (regex, JSON, ...). Such
subexpressions are serialized
+ * as JVM codegen dispatch nodes, which cannot bind `NamedLambdaVariable`s, so
the lambda cannot
+ * be executed natively. A dispatch-only expression among the *value*
arguments is fine and stays
+ * native. JVM codegen dispatch - `CometScalaUDF.emitJvmCodegenDispatch` runs
the whole HOF
+ * (lambda included) on the JVM; `NamedLambdaVariable`s never cross this
boundary because the
+ * lambda is evaluated by Spark's own implementation. Returns `None` when the
dispatcher cannot
+ * handle the expression, which falls back to Spark entirely.
+ *
+ * Whether a lambda body requires dispatch is decided from the already-built
body proto by
+ * [[CometHighOrderFunction.containsJvmDispatch]] - the body is serialized
exactly once, and the
+ * verdict comes from the same proto that native execution would use.
+ */
+case class CometHighOrderFunction[T <: HigherOrderFunction](name: String)
+ extends CometExpressionSerde[T] {
+
+ def convert(expr: T, inputs: Seq[Attribute], binding: Boolean):
Option[ExprOuterClass.Expr] = {
+ if (!CometConf.COMET_EXEC_HIGHER_ORDER_FUNCTION_NATIVE_ENABLED.get()) {
+ return CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding)
+ }
+ highOrderFunction2Proto(expr, inputs, binding)
+ .orElse {
+ CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding)
+ }
+ }
+
+ private def highOrderFunction2Proto(
+ expr: T,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[ExprOuterClass.Expr] = {
+ val argumentsProto = expr.arguments.map(exprToProtoInternal(_, inputs,
binding))
+ val functionsProto = expr.functions
+ .map {
+ case slf: SparkLambdaFunction =>
+ exprToProtoInternal(slf.function, inputs, binding)
Review Comment:
[P2] Preserve conditional evaluation when serializing lambda bodies
Could this conversion preserve Spark's conditional evaluation, or decline
the native HOF cleanly when a speculative conversion would throw? With ANSI
enabled and a Parquet column `a` containing `[-1, 0]`, Spark 4.1.3 returns `[]`
for:
```sql
SELECT filter(a, x -> CASE WHEN x > 0
THEN CAST('bad' AS INT) > 0 ELSE false END)
FROM t;
```
Spark's normal optimizer retains the cast inside the guarded branch. The new
traversal of `slf.function` reaches `CometCast.convert`, which calls
`cast.eval()` for a literal child. That throws `CAST_INVALID_INPUT` during
serialization even though no element takes that branch. The exception escapes
before the fallback in `convert` can run.
I reproduced this through the exact-head public `QueryPlanSerde.exprToProto`
entry point using the optimized Spark expression. With
`spark.comet.exec.higherOrderFunction.native.enabled=false`, the same
expression emits JVM dispatch successfully. The base implementation also
dispatched the whole general filter without traversing its lambda. Please add
an ANSI regression with a guarded error and input that does not take the
failing branch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
