fix: return an empty array from CASE WHEN for an empty batch

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

native/spark-expr/src/array_funcs/nested_comparison.rs

@@ -27,11 +27,11 @@ use arrow::buffer::{BooleanBuffer, 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::physical_expr::is_volatile;
 use std::fmt::{Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
@@ -369,14 +369,24 @@ pub fn spark_in_list(
     schema: &Schema,
 ) -> Result<Arc<dyn PhysicalExpr>> {
     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)?;
     let empty = RecordBatch::new_empty(Arc::new(schema.clone()));
     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::<Literal>()))
+                })
+                .unwrap_or(true);
+            if !constant {
                 return None;
             }
             match child.evaluate(&empty).ok()? {
@@ -759,6 +769,8 @@ mod tests {
     struct Probe {
         child: Arc<dyn PhysicalExpr>,
         volatile: bool,
+        /// Returns a scalar NULL for an empty batch, whatever the child returns
+        scalar_on_empty: bool,
         calls: Arc<std::sync::atomic::AtomicUsize>,
     }
     impl Eq for Probe {}
@@ -790,6 +802,11 @@ mod tests {
         }
         fn evaluate(&self, b: &RecordBatch) -> Result<ColumnarValue> {
             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 {
@@ -805,6 +822,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),
             }))
         }
@@ -823,6 +841,7 @@ mod tests {
             let probe: Arc<dyn PhysicalExpr> = 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.
@@ -836,6 +855,31 @@ mod tests {
         }
         Ok(())
     }
+    #[test]
+    fn candidates_reading_a_column_remain_dynamic() -> Result<()> {
+        let l: ArrayRef = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>([
+            Some(vec![Some(-0.0)]),
+            Some(vec![Some(1.0)]),
+        ]));
+        let r: ArrayRef = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>([
+            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<dyn PhysicalExpr> = 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::<NestedPredicate>(), "{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;
@@ -852,6 +896,7 @@ mod tests {
         let value: Arc<dyn PhysicalExpr> = Arc::new(Probe {
             child: Arc::clone(&a),
             volatile: false,
+            scalar_on_empty: false,
             calls: Arc::clone(&calls),
         });
         let bad: Arc<dyn PhysicalExpr> = Arc::new(CastExpr::new(

native/spark-expr/src/conditional_funcs/case_when.rs

@@ -17,7 +17,8 @@
 
 use crate::{Cast, EvalMode, IfExpr, NormalizeNaNAndZero, SparkCastOptions};
 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;
@@ -349,6 +350,13 @@ impl PhysicalExpr for CaseWhenExpr {
     }
 
     fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        // 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()))
@@ -1238,14 +1246,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]

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)