andygrove commented on code in PR #6716: URL: https://github.com/apache/datafusion-comet/pull/6716#discussion_r4197302246
########## 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: Fixed in ed51460cc. Before an operator is serialized, Comet now asks Spark's own `EquivalentExpressions` which subexpressions its expressions share, and a call whose later argument holds one stays unmasked, so the shared cast is evaluated for every row as in Spark's generated code. It isn't gated on `spark.sql.subexpressionElimination.enabled`, because Spark's aggregate path eliminates subexpressions regardless of that setting. The array fixture now has your query and an aggregation version, both expecting `CAST_INVALID_INPUT` under the default configuration, plus a control where the whole call is shared and Spark still skips the cast. With the shared check removed, both error queries fail. All three pass on Spark 3.4 through 4.2. -- 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]
