andygrove commented on code in PR #6716: URL: https://github.com/apache/datafusion-comet/pull/6716#discussion_r4219593282
########## native/spark-expr/src/conditional_funcs/null_short_circuit.rs: ########## @@ -0,0 +1,674 @@ +// 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::NullBuffer; +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use datafusion::common::tree_node::{Transformed, TreeNode}; +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::utils::collect_columns; +use datafusion::physical_expr::PhysicalExpr; +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. An argument that can be evaluated for any row, such as a column, is evaluated for all of +/// them. +/// +/// 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>, + 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 masked = inner + .children() + .iter() + .enumerate() + .map(|(i, arg)| i > 0 && !is_infallible(arg, input_schema)) + .collect::<Vec<_>>(); + if !masked.contains(&true) { + return Ok(inner); + } + let plans = inner + .children() + .into_iter() + .zip(masked) + .map(|(arg, masked)| { + if arg.is::<Literal>() { + return Ok((None, masked)); + } + // A masked argument is NULL on the rows it is skipped for + let field = arg.return_field(input_schema)?; + Ok(( + Some(Arc::new(field.as_ref().clone().with_nullable(true))), + masked, + )) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Arc::new(Self::try_new(inner, plans)?)) + } + + /// `plans` holds each argument's [`Arg::field`] and whether it is masked. + fn try_new(inner: Arc<dyn PhysicalExpr>, plans: Vec<(Option<FieldRef>, bool)>) -> Result<Self> { + let children = inner.children(); + if children.len() != plans.len() { + return internal_err!( + "NullShortCircuit has {} plans for {} arguments", + plans.len(), + children.len() + ); + } + let args = children + .into_iter() + .zip(plans) + .map(|(expr, (field, masked))| Arg::try_new(expr, field, masked)) + .collect::<Result<Vec<_>>>()?; + let mut body_children = Vec::with_capacity(args.len()); + let mut column = 0; + for arg in &args { + if let Some(field) = &arg.field { + body_children.push(Arc::new(Column::new(field.name(), column)) as _); + column += 1; + } else { + body_children.push(Arc::clone(&arg.expr)); + } + } + let body = Arc::clone(&inner).with_new_children(body_children)?; + Ok(Self { inner, body, 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 an argument evaluated so far is NULL, or `None` while there are none + let mut nulls: Option<NullBuffer> = None; + let mut fields = vec![]; + let mut columns = vec![]; + let last = self.args.len() - 1; + for (i, arg) in self.args.iter().enumerate() { + let value = arg.evaluate(batch, nulls.as_ref())?; + if i < last { + let value_nulls = match &value { + ColumnarValue::Array(array) if array.len() != num_rows => { + return internal_err!( + "{} returned {} rows for a batch of {num_rows}", + arg.expr, + array.len() + ); + } + ColumnarValue::Array(array) => array.logical_nulls(), + ColumnarValue::Scalar(scalar) if scalar.is_null() => { + Some(NullBuffer::new_null(num_rows)) + } + ColumnarValue::Scalar(_) => None, + }; + nulls = NullBuffer::union(nulls.as_ref(), value_nulls.as_ref()) + .filter(|n| n.null_count() > 0); + if nulls.as_ref().is_some_and(|n| n.null_count() == num_rows) { + // Every row is NULL, and no later argument is evaluated for any of them + let data_type = self.inner.data_type(batch.schema_ref())?; + return Ok(ColumnarValue::Array(new_null_array(&data_type, num_rows))); + } + } + if let Some(field) = &arg.field { + let column = value.into_array(num_rows)?; + // The batch is checked against the types the arguments have at run time, which + // can differ from their planned types in the names and nullability of nested fields + fields.push(if field.data_type() == column.data_type() { + Arc::clone(field) + } else { + Arc::new( + field + .as_ref() + .clone() + .with_data_type(column.data_type().clone()), + ) + }); + columns.push(column); + } + } + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + let args = + RecordBatch::try_new_with_options(Arc::new(Schema::new(fields)), columns, &options)?; + self.body.evaluate(&args) + } + + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + self.inner.children() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn PhysicalExpr>>, + ) -> Result<Arc<dyn PhysicalExpr>> { + let inner = Arc::clone(&self.inner).with_new_children(children)?; + let plans = self + .args + .iter() + .map(|arg| (arg.field.clone(), arg.masked.is_some())) + .collect(); + Ok(Arc::new(Self::try_new(inner, plans)?)) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.inner.fmt_sql(f) + } + + fn is_volatile_node(&self) -> bool { + self.inner.is_volatile_node() + } +} + +/// An argument, and how it is evaluated. +#[derive(Debug)] +struct Arg { + expr: Arc<dyn PhysicalExpr>, + /// The field of the column that `body` reads the argument from, or `None` for a literal, + /// which `body` keeps + field: Option<FieldRef>, + /// For an argument evaluated for only the rows where no earlier argument is NULL, the input + /// columns it reads, and the argument reading them from a batch of just those columns, so + /// that skipping rows filters no other column + masked: Option<(Vec<usize>, Arc<dyn PhysicalExpr>)>, +} + +impl Arg { + fn try_new( + expr: &Arc<dyn PhysicalExpr>, + field: Option<FieldRef>, + masked: bool, + ) -> Result<Self> { + let masked = masked.then(|| project(expr)).transpose()?; + Ok(Self { + expr: Arc::clone(expr), + field, + masked, + }) + } + + /// Evaluates the argument, for just the rows outside `nulls` if it is masked, returning a + /// value for every row of the batch. + fn evaluate(&self, batch: &RecordBatch, nulls: Option<&NullBuffer>) -> Result<ColumnarValue> { + match (&self.masked, nulls) { + (Some((projection, projected)), Some(nulls)) => { + let rows = BooleanArray::new(nulls.inner().clone(), None); + projected.evaluate_selection(&batch.project(projection)?, &rows) + } + _ => self.expr.evaluate(batch), Review Comment: Fixed in 9c9ff57832. A masked argument is now evaluated through its rebuilt copy for every batch, including one with no NULL array, so `shuffle` keeps one generator however the batches alternate. The original expression is evaluated only for an argument that isn't masked. Your case is now a unit test, `stateful_argument_keeps_its_state_across_batches`. It evaluates `array_union(a, shuffle(b))` over a batch with no NULL array, then one with a NULL, then one with none again, and compares each with the unwrapped function. Its second batch comes out as `[NULL,[1,4,3,2],[3,2,1,4]]`, the sequence you got from Spark 3.5.9. There's also a Spark 4.0+ fixture, `array_null_short_circuit_shuffle.sql`, which runs `array_union(a, shuffle(b, 42))` over 8200 rows in one file so that only the scan's second batch has NULL arrays. Both fail with the previous evaluation, and the fixture passes on 4.0, 4.1 and 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]
