diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index 20bcae06e69..1fc8fdc26d1 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/docs/source/contributor-guide/adding_a_new_expression.md @@ -560,6 +560,14 @@ pub fn spark_ceil( } ``` +#### Returning a scalar or an array + +Return `ColumnarValue::Scalar` only when the result is the same for every row, for example when all +the inputs are scalars. That includes an empty batch: an expression whose result depends on a column +must return an empty array there, not a scalar. `IN` evaluates its candidates on an empty batch to +decide whether they are constants, and uses the value of a candidate that returns a scalar there for +every row. + #### Producing strings from arbitrary bytes Spark's `StringType` may contain bytes that are not valid UTF-8 (Spark stores them verbatim), but diff --git a/native/spark-expr/src/array_funcs/nested_comparison.rs b/native/spark-expr/src/array_funcs/nested_comparison.rs index 585c244c0eb..62f5c1eca58 100644 --- a/native/spark-expr/src/array_funcs/nested_comparison.rs +++ b/native/spark-expr/src/array_funcs/nested_comparison.rs @@ -29,12 +29,12 @@ use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer}; use arrow::compute::{not, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNode; use datafusion::common::{internal_err, DFSchema, Result, ScalarValue}; use datafusion::logical_expr::{ColumnarValue, Operator}; use datafusion::physical_expr::expressions::{in_list, BinaryExpr, Column, InListExpr, Literal}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_common::datum::apply_cmp; -use datafusion::physical_expr_common::physical_expr::is_volatile; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -653,6 +653,9 @@ pub fn spark_in_list( schema: &Schema, ) -> Result> { if !is_nested_with_float_leaf(&value.data_type(schema)?) || candidates.is_empty() { + // DataFusion takes a candidate as a constant when it returns a scalar for an empty + // batch, so an expression that reads a column must return an array there + // (apache/datafusion#26082) return in_list(value, candidates, &negated, schema); } validate_types(&value, &candidates, schema)?; @@ -660,7 +663,14 @@ pub fn spark_in_list( let constants = candidates .iter() .map(|child| { - if is_volatile(child) { + // Only a candidate built from literals is a constant. One that reads a column can + // still return a scalar for the empty batch. + let constant = !child + .exists(|e| { + Ok(e.is_volatile_node() || (e.children().is_empty() && !e.is::())) + }) + .unwrap_or(true); + if !constant { return None; } match child.evaluate(&empty).ok()? { @@ -1043,6 +1053,8 @@ mod tests { struct Probe { child: Arc, volatile: bool, + /// Returns a scalar NULL for an empty batch, whatever the child returns + scalar_on_empty: bool, calls: Arc, } impl Eq for Probe {} @@ -1074,6 +1086,11 @@ mod tests { } fn evaluate(&self, b: &RecordBatch) -> Result { self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if self.scalar_on_empty && b.num_rows() == 0 { + return Ok(ColumnarValue::Scalar(ScalarValue::try_from( + self.data_type(&b.schema())?, + )?)); + } self.child.evaluate(b) } fn is_volatile_node(&self) -> bool { @@ -1089,6 +1106,7 @@ mod tests { Ok(Arc::new(Self { child: Arc::clone(&children[0]), volatile: self.volatile, + scalar_on_empty: self.scalar_on_empty, calls: Arc::clone(&self.calls), })) } @@ -1107,6 +1125,7 @@ mod tests { let probe: Arc = Arc::new(Probe { child: literal(batch.column(0), 0), volatile, + scalar_on_empty: false, calls: Arc::clone(&calls), }); // Volatility must also be recognized through a nonvolatile parent node. @@ -1120,6 +1139,31 @@ mod tests { } Ok(()) } + #[test] + fn candidates_reading_a_column_remain_dynamic() -> Result<()> { + let l: ArrayRef = Arc::new(ListArray::from_iter_primitive::([ + Some(vec![Some(-0.0)]), + Some(vec![Some(1.0)]), + ])); + let r: ArrayRef = Arc::new(ListArray::from_iter_primitive::([ + Some(vec![Some(0.0)]), + None, + ])); + let batch = batch(l, r); + let (a, b) = columns(); + // An expression of b that returns a scalar for an empty batch + let candidate: Arc = Arc::new(Probe { + child: b, + volatile: false, + scalar_on_empty: true, + calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + }); + let expr = spark_in_list(a, vec![candidate], false, batch.schema().as_ref())?; + assert!(expr.as_ref().is::(), "{expr}"); + assert_eq!(results(&expr, &batch), vec![Some(true), None]); + Ok(()) + } + #[test] fn dynamic_input_is_evaluated_once_and_runtime_errors_propagate() -> Result<()> { use datafusion::physical_expr::expressions::CastExpr; @@ -1136,6 +1180,7 @@ mod tests { let value: Arc = Arc::new(Probe { child: Arc::clone(&a), volatile: false, + scalar_on_empty: false, calls: Arc::clone(&calls), }); let bad: Arc = Arc::new(CastExpr::new( @@ -1505,6 +1550,7 @@ mod tests { Arc::new(Probe { child: Arc::new(Literal::new(value)), volatile: false, + scalar_on_empty: false, calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), }) }; diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 192e6046f61..9676cad9505 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -17,7 +17,8 @@ use crate::{Cast, EvalMode, IfExpr, NormalizeNaNAndZero, SparkCastOptions, SparkComparison}; use arrow::array::{ - downcast_primitive, Array, ArrayRef, AsArray, BooleanArray, GenericByteArray, PrimitiveArray, + downcast_primitive, new_empty_array, Array, ArrayRef, AsArray, BooleanArray, GenericByteArray, + PrimitiveArray, }; use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::nullif; @@ -353,6 +354,13 @@ impl PhysicalExpr for CaseWhenExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { + // No row chooses a branch, which both the eager and the lazy evaluation answer with a + // scalar NULL. A scalar from an empty batch is taken to mean the expression is constant, + // as IN does to build its list once, so return an empty array instead. + if batch.num_rows() == 0 { + let data_type = self.data_type(&batch.schema())?; + return Ok(ColumnarValue::Array(new_empty_array(&data_type))); + } match self .eager .get_or_init(|| self.eager_result_type(&batch.schema())) @@ -1270,14 +1278,37 @@ mod tests { assert_eq!(scalar(expr), ScalarValue::Int64(None)); } + /// A scalar from an empty batch is taken to mean the expression is constant, so a CASE that + /// depends on a column has to return an empty array, whichever way it is evaluated. #[test] - fn empty_batch() { + fn empty_batch_returns_an_empty_array() { let batch = int_batch(vec![], vec![]); let schema = batch.schema(); let a = col("a", &schema).unwrap(); let b = col("b", &schema).unwrap(); - let when_then = vec![(binary(Arc::clone(&a), Operator::Lt, lit(0i64)), a)]; - check_against_case_expr(&batch, when_then, Some(b)); + let null = || lit(ScalarValue::Int64(None)); + let a_is_1 = || binary(Arc::clone(&a), Operator::Eq, lit(1i64)); + let a_div_b = binary(Arc::clone(&a), Operator::Divide, b); + let new_case = |when_then, else_expr| CaseWhenExpr::try_new(when_then, else_expr).unwrap(); + for (expr, eager) in [ + // IF(a = 1, NULL, a), as nullif(a, 1) is planned + ( + new_case(vec![(a_is_1(), null())], Some(Arc::clone(&a))), + true, + ), + // A branch that can fail is evaluated lazily, by CaseExpr, which also returns the + // NULL literal as a scalar for an empty batch + (new_case(vec![(a_is_1(), null())], Some(a_div_b)), false), + ] { + assert_eq!(expr.eager_result_type(&schema).is_some(), eager, "{expr}"); + match expr.evaluate(&batch).unwrap() { + ColumnarValue::Array(array) => { + assert_eq!(array.len(), 0, "{expr}"); + assert_eq!(array.data_type(), &DataType::Int64, "{expr}"); + } + other => panic!("{expr} returned {other:?} for an empty batch"), + } + } } #[test] diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/in_case_when_candidate.sql b/spark/src/test/resources/sql-tests/expressions/conditional/in_case_when_candidate.sql new file mode 100644 index 00000000000..6a00e49c238 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/conditional/in_case_when_candidate.sql @@ -0,0 +1,58 @@ +-- 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. + +-- IN builds its list once when every candidate is a constant, which it decides by evaluating the +-- candidates on an empty batch and checking for a scalar. A CASE or IF that depends on a column +-- must not return a scalar NULL there, or IN compares every row against NULL. + +-- Config: spark.comet.exec.range.enabled=true + +query +SELECT id, id IN (IF(id = 1, NULL, id)), id IN (nullif(id, 1), 5L) FROM range(0, 3) + +query +SELECT id, id IN (CASE WHEN id = 1 THEN NULL ELSE id END) FROM range(0, 3) + +-- coalesce is planned as a CASE +query +SELECT id, id IN (coalesce(nullif(id, 1), 5L)) FROM range(0, 3) + +-- no ELSE +query +SELECT id, id IN (CASE WHEN id <> 1 THEN id END), id NOT IN (CASE WHEN id <> 1 THEN id END) +FROM range(0, 3) + +query +SELECT id, id IN (CASE WHEN id = 0 THEN id WHEN id = 2 THEN id END) FROM range(0, 3) + +-- A CASE that reads no column but that Spark cannot fold, because Comet plans +-- spark_partition_id() as a literal +query +SELECT id, id IN (CASE WHEN spark_partition_id() = 0 THEN 1L END, 5L) FROM range(0, 3, 1, 1) + +-- nested operands +query +SELECT id, named_struct('a', id) IN (named_struct('a', IF(id = 1, NULL, id))) FROM range(0, 3) + +query +SELECT id, named_struct('a', CAST(id AS DOUBLE)) + IN (named_struct('a', IF(id = 1, NULL, CAST(id AS DOUBLE)))) +FROM range(0, 3) + +query +SELECT id, array(CAST(id AS DOUBLE)) IN (array(IF(id = 1, NULL, CAST(id AS DOUBLE)))) +FROM range(0, 3)