sunchao commented on code in PR #4744:
URL: https://github.com/apache/datafusion-comet/pull/4744#discussion_r4049685570
##########
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()),
+ )?))
+ }
+
+ fn resolve_lambda_param_fields(
+ udf: &HigherOrderUDF,
+ func_name: &str,
+ value_args: &[Arc<dyn PhysicalExpr>],
+ lambda_count: usize,
+ schema: &Schema,
+ ) -> Result<Vec<Vec<FieldRef>>, ExecutionError> {
+ let mut planning_fields: Vec<ValueOrLambda<FieldRef,
Option<FieldRef>>> = value_args
+ .iter()
+ .map(|e| Ok(ValueOrLambda::Value(e.return_field(schema)?)))
+ .collect::<Result<_, DataFusionError>>()?;
+ planning_fields.extend(std::iter::repeat_n(
+ ValueOrLambda::Lambda(None),
+ lambda_count,
+ ));
+
+ match udf.lambda_parameters(0, &planning_fields)? {
+ LambdaParametersProgress::Complete(items) if items.len() >=
lambda_count => Ok(items),
+ LambdaParametersProgress::Complete(items) =>
Err(GeneralError(format!(
+ "{func_name}: expected parameter fields for {lambda_count}
lambdas, got {}",
+ items.len()
+ ))),
+ LambdaParametersProgress::Partial(_) => Err(GeneralError(format!(
+ "{func_name}: multi-step lambda resolution is not supported
yet"
+ ))),
+ }
+ }
+
+ fn create_lambda_expr(
+ &self,
+ lambda: &LambdaFunction,
+ input_schema: &SchemaRef,
+ param_fields: &[FieldRef],
+ ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+ if param_fields.len() < lambda.args.len() {
+ return Err(GeneralError(format!(
+ "lambda declares {} params but the function resolved only {}",
+ lambda.args.len(),
+ param_fields.len()
+ )));
+ }
+
+ // Build extended schema = input schema ++ lambda params, and the
scope.
+ let mut body_fields: Vec<FieldRef> =
input_schema.fields().iter().map(Arc::clone).collect();
+ let mut scope = LambdaScope::with_capacity(lambda.args.len());
+ let mut scope_entries: Vec<(usize, FieldRef)> =
Vec::with_capacity(lambda.args.len());
+ let mut param_names: Vec<String> =
Vec::with_capacity(lambda.args.len());
+
+ for (arg, resolved) in lambda.args.iter().zip(param_fields) {
+ // Runtime uses `param.renamed(name)` — do the same here.
+ let field: FieldRef =
Arc::new(resolved.as_ref().clone().with_name(&arg.name));
+ let idx = body_fields.len();
+ if scope
+ .insert(arg.expr_id, (idx, Arc::clone(&field)))
+ .is_some()
+ {
+ return Err(GeneralError(format!(
+ "duplicate lambda variable exprId {} ('{}')",
+ arg.expr_id, arg.name
+ )));
+ }
+ scope_entries.push((idx, Arc::clone(&field)));
+ param_names.push(arg.name.clone());
+ body_fields.push(field);
+ }
+
+ let body_schema = Arc::new(Schema::new(
+ body_fields
+ .iter()
+ .map(|f| f.as_ref().clone())
+ .collect::<Vec<_>>(),
+ ));
+ let lambda_body = lambda
+ .body
+ .as_ref()
+ .ok_or_else(|| GeneralError("lambda has no body".to_string()))?;
+
+ // Plan the body under this scope; the guard pops on any `?` / drop.
+ let body_expr = self
+ .lambda_scopes
+ .with_scope(scope, || self.create_expr(lambda_body, body_schema))?;
+
+ Ok(Arc::new(LambdaExpr::try_new(param_names, body_expr)?))
Review Comment:
[P2] Preserve per-element short-circuiting in native lambda bodies
Could we preserve Spark's per-element AND/OR evaluation before admitting
these lambda bodies to the native path? With ANSI enabled and a Parquet array
column `a` containing `[0, 1]`:
```sql
SELECT filter(a, x -> x <> 0 AND 1 DIV x > 0) FROM t;
-- Spark: [1]
SELECT filter(a, x -> x = 0 OR 1 DIV x > 0) FROM t;
-- Spark: [0, 1]
```
At `5c4ba0a`, both expressions serialize to native HOFs, including with JVM
codegen dispatch disabled. I verified that Spark 4.1.3 retains the guard on the
left in the optimized expression and returns the results above. The
corresponding native component probes raise `DIVIDE_BY_ZERO` for both
expressions.
`AndBuilder` and `OrBuilder` construct DataFusion `BinaryExpr`. In
DataFusion 55.1.0, mixed boolean batches only mask the right operand when at
most 20% of rows need it. With `[0, 1]`, division therefore runs on the zero
element despite the guard. The control `[0, 0, 0, 0, 1]` succeeds. The base
implementation dispatched the whole general filter to Spark, and the new
serialization `NonFatal` catch cannot intercept this runtime error.
Could we preserve the evaluation mask, or fall back for predicates whose
skipped branches can raise, and add ANSI regressions for both guarded AND and
OR?
Validation boundary: these are native component reproductions using the
current Comet division kernel and DataFusion components whose relevant source
files match public 55.1.0 byte-for-byte, paired with exact-head serializer
checks and Spark reference executions. Full Comet/JNI query execution remains
unrun.
##########
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:
Rechecked at `5c4ba0aab36c33370dfd339eba53920621efd25b`: the empty-elements
issue is still present in this head. With ANSI enabled, the partition-zero
predicate from the original comment still serializes natively. Spark returns
`[]` and `null`, while the native component probe raises `DIVIDE_BY_ZERO` for
both a nonnull empty array and mixed empty/null input. The all-null control
returns null correctly.
The proposed zero-row guard at the lambda-body boundary looks suitable. I
tested that guard in the component harness: it produces `[]`, preserves `[[],
null]` for mixed input, and leaves the all-null result unchanged. DataFusion
can retain responsibility for rebuilding the output arrays and their null mask.
There is no lambda wrapper left in this head, so this needs an actual
physical-expression adapter around `body_expr` before `LambdaExpr::try_new`.
Its `children()` should expose the inner body and `with_new_children()` should
rebuild the adapter so projection rewriting preserves the guard. Could you add
that implementation and an ANSI regression with at least one nonnull empty
array, including mixed empty/null input?
This validation used component probes plus exact-head serializer checks and
Spark 4.1.3 reference execution. I have not run the full Comet/JNI query.
--
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]
