sunchao commented on code in PR #6716: URL: https://github.com/apache/datafusion-comet/pull/6716#discussion_r4196438196
########## native/spark-expr/src/conditional_funcs/null_short_circuit.rs: ########## @@ -0,0 +1,632 @@ +// 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. + +use super::case_when::is_infallible; +use arrow::array::{new_null_array, Array, BooleanArray}; +use arrow::buffer::BooleanBuffer; +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::{internal_datafusion_err, internal_err, Result}; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::{Column, Literal}; +use datafusion::physical_expr::PhysicalExpr; +use std::collections::BTreeSet; +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +/// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant +/// `BinaryExpression` or `TernaryExpression`: the first for every row, and each later one only for +/// the rows where no earlier argument is NULL. +/// +/// Spark returns NULL as soon as an argument is NULL, without evaluating the ones after it. +/// DataFusion evaluates every argument over the whole batch before it calls a function, so an +/// argument that can fail, such as an ANSI cast of a malformed string, would fail on a row that +/// Spark answers with NULL, and a nondeterministic argument would see rows that it does not see in +/// Spark. +/// +/// The function still sees every row, with NULL for each argument skipped there, so it has to +/// return NULL wherever an argument other than the last is NULL. Its arguments have to be its +/// children, in the order Spark evaluates them, and it has to compute its result from them alone. +/// `ArrayInsert` evaluates its own arguments the same way. +#[derive(Debug)] +pub struct NullShortCircuit { + /// The function over its arguments + inner: Arc<dyn PhysicalExpr>, + /// The function over the evaluated arguments, reading each one that is not a literal from a + /// column of the batch that holds them + body: Arc<dyn PhysicalExpr>, + /// The fields of those columns + fields: Vec<FieldRef>, + args: Vec<Arg>, +} + +impl NullShortCircuit { + /// Wraps `inner` if evaluating one of its arguments after the first for a row that Spark skips + /// could change the outcome, and returns it unchanged otherwise. + pub fn wrap( + inner: Arc<dyn PhysicalExpr>, + input_schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>> { + let children = inner.children(); + if children + .iter() + .skip(1) + .all(|arg| is_infallible(arg, input_schema)) + { + return Ok(inner); + } + // An argument after the first is NULL on the rows it is skipped for + let fields = children + .iter() + .filter(|arg| !arg.is::<Literal>()) + .map(|arg| { + let field = arg.return_field(input_schema)?; + Ok(Arc::new(field.as_ref().clone().with_nullable(true)) as FieldRef) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Arc::new(Self::try_new(inner, fields)?)) + } + + fn try_new(inner: Arc<dyn PhysicalExpr>, fields: Vec<FieldRef>) -> Result<Self> { + let args = inner + .children() + .into_iter() + .map(Arg::try_new) + .collect::<Result<Vec<_>>>()?; + let mut columns = fields.iter().enumerate(); + let body_children = args + .iter() + .map(|arg| { + if arg.is_literal { + return Ok(Arc::clone(&arg.expr)); + } + let (index, field) = columns + .next() + .ok_or_else(|| internal_datafusion_err!("NullShortCircuit lacks a field"))?; + Ok(Arc::new(Column::new(field.name(), index)) as Arc<dyn PhysicalExpr>) + }) + .collect::<Result<Vec<_>>>()?; + if columns.next().is_some() { + return internal_err!("NullShortCircuit has a field for no argument"); + } + let body = Arc::clone(&inner).with_new_children(body_children)?; + Ok(Self { + inner, + body, + fields, + args, + }) + } +} + +impl Hash for NullShortCircuit { + fn hash<H: Hasher>(&self, state: &mut H) { + self.inner.hash(state); + } +} + +impl PartialEq for NullShortCircuit { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +impl Eq for NullShortCircuit {} + +impl Display for NullShortCircuit { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "NullShortCircuit({})", self.inner) + } +} + +impl PhysicalExpr for NullShortCircuit { + fn data_type(&self, input_schema: &Schema) -> Result<DataType> { + self.inner.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result<bool> { + self.inner.nullable(input_schema) + } + + fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { + self.inner.return_field(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + let num_rows = batch.num_rows(); + // The rows where no argument evaluated so far is NULL, or `None` while that is all of them + let mut rows: Option<BooleanArray> = None; + let mut columns = Vec::with_capacity(self.fields.len()); + let last = self.args.len() - 1; + for (i, arg) in self.args.iter().enumerate() { + let value = match &rows { + Some(rows) if !arg.is_literal => arg.evaluate_rows(batch, rows)?, Review Comment: [P2] Preserve Spark’s errors for shared fallible arguments. With ANSI enabled and the default `spark.sql.subexpressionElimination.enabled=true`, `SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)) FROM t` over Parquet rows `(NULL, 'bad')` and `(array(1), '1')` raises `CAST_INVALID_INPUT` in Spark. Spark’s projection evaluates the shared cast before either array expression. Masking each call independently here instead returns `(NULL, NULL)` and `(true, 1)`, silently allowing a query that Spark and the base native path reject. Could affected projections preserve Spark’s shared evaluation or fall back, with a regression test covering the default configuration? Evidence: Executed the query on Spark 3.5.9 with ANSI enabled: subexpression elimination enabled raised CAST_INVALID_INPUT, while disabling it returned [(NULL,NULL),(true,1)]. An exact-head Rust reproduction using the native array_contains/array_position functions and Cast showed both unwrapped expressions raising CAST_INVALID_INPUT and both NullShortCircuit-wrapped expressions returning [NULL,true] and [NULL,1]. Spark 4.1.3 source corroborates the trigger: ProjectExec.doConsume emits shared-expression evaluation before result expressions, and EquivalentExpressions recurses into these binary expressions’ children. Reproduction artifacts: /tmp/6716-spark-cse-repro.log, /tmp/6716-native-repro.rs, /tmp/6716-native-repro.log. -- 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]
