kazantsev-maksim commented on code in PR #4744:
URL: https://github.com/apache/datafusion-comet/pull/4744#discussion_r4052952151
##########
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:
Thanks for the thorough reviews and guidance, @sunchao!
Both runtime P2 issues have been resolved, and all corresponding ANSI
regression tests are now passing:
### 1. Zero-row lambda guard (Empty & mixed arrays)
- Implemented `EmptyBatchGuardExpr`, a lightweight physical expression
adapter in `native/core/src/execution/lambda.rs`, and wrapped `body_expr`
before calling `LambdaExpr::try_new`.
- When `batch.num_rows() == 0`, it short-circuits evaluation and returns
`arrow::array::new_empty_array` directly, avoiding scalar runtime errors (such
as `1 DIV spark_partition_id()`). DataFusion retains responsibility for
reconstructing the output array offsets and row null masks.
- The adapter properly delegates `children()`, `with_new_children()`,
`fmt_sql()`, and satisfies `DynEq`/`DynHash` via `dyn_eq`/`dyn_hash`.
### 2. Per-element short-circuiting in conditional branches (Guarded AND /
OR / CASE / IF)
- Added `hasGuardedFallibleBranch` and `isFallibleExpr` in
`CometHighOrderFunction.scala`.
- Because DataFusion evaluates vectorized branches without masking when more
than 20% of batch rows require evaluation (which easily triggers on small array
batches like `[0, 1]`), native execution cannot guarantee Spark's per-element
short-circuiting in ANSI mode.
- When conditional expressions (`And`, `Or`, `CaseWhen`, `If`, `Coalesce`)
contain potentially fallible operations (integer division, non-try casts,
arithmetic overflow, out-of-bounds array indexing) in guarded branches under
ANSI mode, the native path is declined and execution cleanly falls back to JVM
codegen dispatch.
### 3. Regression SQL tests added
Added test queries under ANSI mode covering:
- Non-null empty arrays and mixed `[[], NULL]` using `1 DIV
spark_partition_id()`.
- Guarded AND with `x <> 0 AND (1 DIV x) > 0` on `[0, 1]`.
- Guarded OR with `x = 0 OR (1 DIV x) > 0` on `[0, 1]`.
- Guarded CASE WHEN with `CAST('bad' AS INT)` on `[-1, 0]`.
Could you please take another look when you have time?
--
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]
