sunchao commented on code in PR #6716: URL: https://github.com/apache/datafusion-comet/pull/6716#discussion_r4196865164
########## 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: [P2] Keep one stateful argument instance across masked and unmasked batches. For `array_union(a, shuffle(b))`, with columns ordered `(a, b)`, projecting `b` from index 1 to 0 reconstructs `ShuffleExpr` with a separate generator. This branch evaluates the original expression for all-valid batches, while line 270 evaluates the projected expression for mixed-null batches. Switching between them restarts the random sequence and changes array order relative to Spark. This also regresses the base path when `b` is NULL exactly where `a` is NULL, because the original shuffle already skips those rows. Could both paths use the same argument instance, with a regression spanning both batch shapes? Evidence: At the reviewed SHA, a native test evaluated `array_union(a, ShuffleExpr(b, 42))` over batch 1 containing two rows `([], [1,2,3,4])`, then batch 2 containing `(NULL,NULL)` and two more identical non-null rows. Spark 3.5.9 with resolved seed 42 and the unwrapped native expression produce batch 2 `[NULL,[1,4,3,2],[3,2,1,4]]`. The wrapper produces `[NULL,[4,3,2,1],[4,1,2,3]]`, repeating batch 1's permutations. The equality assertion failed reproducibly. Source confirms `ShuffleExpr::with_new_children` creates fresh generator state. Reproduction source and logs: `/tmp/6716-2fd9-review-tests.rs`, `/tmp/6716-2fd9-repro.log`, `/tmp/6716-2fd9-spark.py`, `/tmp/6716-2fd9-spark.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]
