From d7e911d61cbf578ec67d470f956ffef84c808519 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Fri, 6 Nov 2026 16:17:59 +0800 Subject: [PATCH 1/4] perf: avoid repeated row copies in literal CASE branches Signed-off-by: Dennis Zhuang --- datafusion/physical-expr/benches/case_when.rs | 70 +++++++++- .../physical-expr/src/expressions/case.rs | 121 ++++++++++++++++-- 2 files changed, 177 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-expr/benches/case_when.rs b/datafusion/physical-expr/benches/case_when.rs index 44534a41aea73..0add92dea070b 100644 --- a/datafusion/physical-expr/benches/case_when.rs +++ b/datafusion/physical-expr/benches/case_when.rs @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use arrow::array::{Array, ArrayRef, Int32Array, Int32Builder, StringArray}; +use arrow::array::{ + Array, ArrayRef, BooleanArray, Int32Array, Int32Builder, StringArray, +}; use arrow::datatypes::{ArrowNativeTypeOp, Field, Schema}; use arrow::record_batch::RecordBatch; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; @@ -87,6 +89,72 @@ fn criterion_benchmark(c: &mut Criterion) { benchmark_lookup_table_case_when(c, 8192); benchmark_divide_by_zero_protection(c, 8192); + benchmark_literal_when_selection(c, 8192); +} + +fn benchmark_literal_when_selection(c: &mut Criterion, row_count: usize) { + let mut group = c.benchmark_group("literal_when_selection"); + for branches in [1, 4, 16] { + for first_only in [false, true] { + if branches == 1 && first_only { + continue; + } + let batch = RecordBatch::try_from_iter(vec![ + ( + "partition", + Arc::new(Int32Array::from_iter_values( + (0..row_count) + .map(|i| if first_only { 0 } else { (i % branches) as i32 }), + )) as ArrayRef, + ), + ( + "left_key", + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|i| format!("{i:032x}")), + )) as ArrayRef, + ), + ( + "right_key", + Arc::new(StringArray::from_iter_values((0..row_count).map(|i| { + let value = if i % 3 == 0 { i } else { i + 1 }; + format!("{value:032x}") + }))) as ArrayRef, + ), + ]) + .unwrap(); + let schema = batch.schema(); + let then: Arc = Arc::new(BinaryExpr::new( + col("left_key", &schema).unwrap(), + Operator::Eq, + col("right_key", &schema).unwrap(), + )); + let expr = case( + Some(col("partition", &schema).unwrap()), + (0..branches) + .map(|i| (lit(i as i32), Arc::clone(&then))) + .collect(), + None, + ) + .unwrap(); + let expected = BooleanArray::from( + (0..row_count).map(|i| i % 3 == 0).collect::>(), + ); + let result = expr + .evaluate(&batch) + .unwrap() + .into_array(row_count) + .unwrap(); + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &expected + ); + let distribution = if first_only { "first_only" } else { "uniform" }; + group.bench_function(BenchmarkId::new(distribution, branches), |b| { + b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) + }); + } + } + group.finish(); } fn run_benchmarks(c: &mut Criterion, batch: &RecordBatch) { diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 91f282f3287f6..ef5393ecce6d9 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -385,6 +385,19 @@ fn filter_record_batch( } } +fn take_case_rows<'a>( + batch: &'a RecordBatch, + indices: &dyn Array, +) -> std::result::Result, ArrowError> { + // These indices are an ordered subset of the original batch's rows. + if indices.len() == batch.num_rows() { + return Ok(Cow::Borrowed(batch)); + } + Ok(Cow::Owned(arrow::compute::take_record_batch( + batch, indices, + )?)) +} + // This function exists purely to be able to use the same call style // for `filter_record_batch` and `filter_array` at the point of use. // When https://github.com/apache/arrow-rs/pull/8693 is available, replace @@ -771,11 +784,17 @@ impl CaseBody { return_type: &DataType, ) -> Result { let mut result_builder = ResultBuilder::new(return_type, batch.num_rows()); + // Literal WHEN expressions do not read rows. Keep indices between branches + // instead of repeatedly copying the unselected columns. + let defer_remainder = self + .when_then_expr + .iter() + .all(|(when, _)| when.is::()); // `remainder_rows` contains the indices of the rows that need to be evaluated let mut remainder_rows: ArrayRef = Arc::new(UInt32Array::from_iter_values(0..batch.num_rows() as u32)); - // `remainder_batch` contains the rows themselves that need to be evaluated + // Non-literal WHEN expressions need the remaining rows materialized. let mut remainder_batch = Cow::Borrowed(batch); // evaluate the base expression @@ -823,10 +842,12 @@ impl CaseBody { return result_builder.finish(); } - // Remove the null rows from the remainder batch + // Remove null rows from the remaining selection. let not_null_filter = create_filter(&base_not_nulls, true); - remainder_batch = - Cow::Owned(filter_record_batch(&remainder_batch, ¬_null_filter)?); + if !defer_remainder { + remainder_batch = + Cow::Owned(filter_record_batch(&remainder_batch, ¬_null_filter)?); + } remainder_rows = filter_array(&remainder_rows, ¬_null_filter)?; base_values = filter_array(&base_values, ¬_null_filter)?; } @@ -855,22 +876,29 @@ impl CaseBody { continue; } - // If the 'when' predicate matched all remaining rows, there is no need to filter + // No later branch can match if this branch selects every remaining row. if when_value.null_count() == 0 && !when_value.has_false() { let then_expression = &self.when_then_expr[i].1; - let then_value = then_expression.evaluate(&remainder_batch)?; + let then_batch = if defer_remainder { + take_case_rows(batch, &remainder_rows)? + } else { + Cow::Borrowed(remainder_batch.as_ref()) + }; + let then_value = then_expression.evaluate(&then_batch)?; result_builder.add_branch_result(&remainder_rows, then_value)?; return result_builder.finish(); } - // Filter the remainder batch based on the 'when' value - // This results in a batch containing only the rows that need to be evaluated - // for the current branch + // Evaluate THEN on the selected rows. // Still no need to call `prep_null_mask_filter` since `create_filter` will already do // this unconditionally. let then_filter = create_filter(&when_value, true); - let then_batch = filter_record_batch(&remainder_batch, &then_filter)?; let then_rows = filter_array(&remainder_rows, &then_filter)?; + let then_batch = if defer_remainder { + take_case_rows(batch, &then_rows)? + } else { + Cow::Owned(filter_record_batch(&remainder_batch, &then_filter)?) + }; let then_expression = &self.when_then_expr[i].1; let then_value = then_expression.evaluate(&then_batch)?; @@ -892,8 +920,10 @@ impl CaseBody { } }?; let next_filter = create_filter(&next_selection, true); - remainder_batch = - Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?); + if !defer_remainder { + remainder_batch = + Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?); + } remainder_rows = filter_array(&remainder_rows, &next_filter)?; base_values = filter_array(&base_values, &next_filter)?; } @@ -903,7 +933,12 @@ impl CaseBody { if let Some(e) = &self.else_expr { // keep `else_expr`'s data type and return type consistent let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?; - let else_value = expr.evaluate(&remainder_batch)?; + let else_batch = if defer_remainder { + take_case_rows(batch, &remainder_rows)? + } else { + Cow::Borrowed(remainder_batch.as_ref()) + }; + let else_value = expr.evaluate(&else_batch)?; result_builder.add_branch_result(&remainder_rows, else_value)?; } @@ -1644,6 +1679,66 @@ mod tests { } } + #[test] + fn literal_when_preserves_row_indices_and_short_circuits() -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![ + ( + "key", + Arc::new(Int32Array::from(vec![ + Some(2), + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(0), + ])) as ArrayRef, + ), + ( + "value", + Arc::new(Int32Array::from(vec![20, 99, 10, 21, 30, 40, 0])) as ArrayRef, + ), + ])?; + let schema = batch.schema(); + let invalid = binary(lit(1i32), Operator::Divide, lit(0i32), &schema)?; + for with_else in [false, true] { + let expr = case( + Some(col("key", &schema)?), + vec![ + (lit(99i32), Arc::clone(&invalid)), + (lit(ScalarValue::Int32(None)), Arc::clone(&invalid)), + (lit(2i32), col("value", &schema)?), + (lit(1i32), col("value", &schema)?), + (lit(2i32), Arc::clone(&invalid)), + (lit(3i32), lit(300i32)), + ], + with_else.then(|| col("value", &schema)).transpose()?, + )?; + let expected = if with_else { + vec![ + Some(20), + Some(99), + Some(10), + Some(21), + Some(300), + Some(40), + Some(0), + ] + } else { + vec![Some(20), None, Some(10), Some(21), Some(300), None, None] + }; + // Four rows also exercise the branch that matches every remaining row. + for count in [4, 7] { + let result = expr.evaluate(&batch.slice(0, count))?.into_array(count)?; + assert_eq!( + as_int32_array(&result)?, + &Int32Array::from(expected[..count].to_vec()) + ); + } + } + Ok(()) + } + #[test] fn case_with_expr() -> Result<()> { let batch = case_test_batch()?; From 0747e4ceec47e0a168c984b46fae507ed47e9909 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Fri, 6 Nov 2026 16:25:18 +0800 Subject: [PATCH 2/4] bench: cover partial literal CASE matches with ELSE Signed-off-by: Dennis Zhuang --- datafusion/physical-expr/benches/case_when.rs | 122 +++++++++--------- 1 file changed, 64 insertions(+), 58 deletions(-) diff --git a/datafusion/physical-expr/benches/case_when.rs b/datafusion/physical-expr/benches/case_when.rs index 0add92dea070b..1c65c1c4eed64 100644 --- a/datafusion/physical-expr/benches/case_when.rs +++ b/datafusion/physical-expr/benches/case_when.rs @@ -94,65 +94,71 @@ fn criterion_benchmark(c: &mut Criterion) { fn benchmark_literal_when_selection(c: &mut Criterion, row_count: usize) { let mut group = c.benchmark_group("literal_when_selection"); - for branches in [1, 4, 16] { - for first_only in [false, true] { - if branches == 1 && first_only { - continue; - } - let batch = RecordBatch::try_from_iter(vec![ - ( - "partition", - Arc::new(Int32Array::from_iter_values( - (0..row_count) - .map(|i| if first_only { 0 } else { (i % branches) as i32 }), - )) as ArrayRef, - ), - ( - "left_key", - Arc::new(StringArray::from_iter_values( - (0..row_count).map(|i| format!("{i:032x}")), - )) as ArrayRef, - ), - ( - "right_key", - Arc::new(StringArray::from_iter_values((0..row_count).map(|i| { - let value = if i % 3 == 0 { i } else { i + 1 }; - format!("{value:032x}") - }))) as ArrayRef, - ), - ]) - .unwrap(); - let schema = batch.schema(); - let then: Arc = Arc::new(BinaryExpr::new( - col("left_key", &schema).unwrap(), - Operator::Eq, - col("right_key", &schema).unwrap(), - )); - let expr = case( - Some(col("partition", &schema).unwrap()), - (0..branches) - .map(|i| (lit(i as i32), Arc::clone(&then))) - .collect(), - None, - ) + for (distribution, branches) in [ + ("uniform", 1), + ("uniform", 4), + ("first_only", 4), + ("uniform", 16), + ("first_only", 16), + ("half_match_with_else", 1), + ] { + let first_only = distribution == "first_only"; + let with_else = distribution == "half_match_with_else"; + let partitions = if with_else { 2 } else { branches }; + let batch = RecordBatch::try_from_iter(vec![ + ( + "partition", + Arc::new(Int32Array::from_iter_values((0..row_count).map(|i| { + if first_only { + 0 + } else { + (i % partitions) as i32 + } + }))) as ArrayRef, + ), + ( + "left_key", + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|i| format!("{i:032x}")), + )) as ArrayRef, + ), + ( + "right_key", + Arc::new(StringArray::from_iter_values((0..row_count).map(|i| { + let value = if i % 3 == 0 { i } else { i + 1 }; + format!("{value:032x}") + }))) as ArrayRef, + ), + ]) + .unwrap(); + let schema = batch.schema(); + let then: Arc = Arc::new(BinaryExpr::new( + col("left_key", &schema).unwrap(), + Operator::Eq, + col("right_key", &schema).unwrap(), + )); + let expr = case( + Some(col("partition", &schema).unwrap()), + (0..branches) + .map(|i| (lit(i as i32), Arc::clone(&then))) + .collect(), + with_else.then(|| Arc::clone(&then)), + ) + .unwrap(); + let expected = + BooleanArray::from((0..row_count).map(|i| i % 3 == 0).collect::>()); + let result = expr + .evaluate(&batch) + .unwrap() + .into_array(row_count) .unwrap(); - let expected = BooleanArray::from( - (0..row_count).map(|i| i % 3 == 0).collect::>(), - ); - let result = expr - .evaluate(&batch) - .unwrap() - .into_array(row_count) - .unwrap(); - assert_eq!( - result.as_any().downcast_ref::().unwrap(), - &expected - ); - let distribution = if first_only { "first_only" } else { "uniform" }; - group.bench_function(BenchmarkId::new(distribution, branches), |b| { - b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) - }); - } + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &expected + ); + group.bench_function(BenchmarkId::new(distribution, branches), |b| { + b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) + }); } group.finish(); } From e6dfc8f8f95687277a66b2e79e6cb7b6fc747e81 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Fri, 6 Nov 2026 16:55:07 +0800 Subject: [PATCH 3/4] fix(physical-expr): preserve nested field names in CASE selection Signed-off-by: Dennis Zhuang --- .../physical-expr/src/expressions/case.rs | 61 ++++++++++++++++++- 1 file changed, 59 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index ef5393ecce6d9..32193fca9b3b7 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -393,8 +393,13 @@ fn take_case_rows<'a>( if indices.len() == batch.num_rows() { return Ok(Cow::Borrowed(batch)); } - Ok(Cow::Owned(arrow::compute::take_record_batch( - batch, indices, + let columns = arrow::compute::take_arrays(batch.columns(), indices, None)?; + Ok(Cow::Owned(RecordBatch::try_new_with_options( + batch.schema(), + columns, + &RecordBatchOptions::new() + .with_match_field_names(false) + .with_row_count(Some(indices.len())), )?)) } @@ -1739,6 +1744,58 @@ mod tests { Ok(()) } + #[test] + fn taking_case_rows_preserves_zero_column_batch_length() -> Result<()> { + let batch = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(5)), + )?; + let selected = take_case_rows(&batch, &UInt32Array::from(vec![1, 3]))?; + assert_eq!(selected.num_rows(), 2); + Ok(()) + } + + #[test] + fn literal_when_preserves_nested_field_names() -> Result<()> { + let mut payload = ListBuilder::new(Int32Builder::new()).with_field(Field::new( + "value", + DataType::Int32, + true, + )); + for value in [10, 20] { + payload.values().append_value(value); + payload.append(true); + } + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new( + "payload", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + false, + ), + ])); + let batch = RecordBatch::try_new_with_options( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(payload.finish()), + ], + &RecordBatchOptions::new().with_match_field_names(false), + )?; + let expr = case( + Some(col("key", &schema)?), + vec![(lit(1i32), expressions::is_null(col("payload", &schema)?)?)], + Some(lit(false)), + )?; + let result = expr.evaluate(&batch)?.into_array(2)?; + assert_eq!( + as_boolean_array(&result)?, + &BooleanArray::from(vec![false, false]) + ); + Ok(()) + } + #[test] fn case_with_expr() -> Result<()> { let batch = case_test_batch()?; From 0e1e39466aef7a7e6b9e4502586ca87e69c69adc Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sun, 8 Nov 2026 22:38:04 -0700 Subject: [PATCH 4/4] fix: reject CASE inputs whose nested field names differ from the schema Signed-off-by: Dennis Zhuang --- .../physical-expr/src/expressions/case.rs | 49 +------------------ 1 file changed, 2 insertions(+), 47 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 188368a088dd8..07d7b1e449b54 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -402,13 +402,8 @@ fn take_case_rows<'a>( ); return Ok(Cow::Borrowed(batch)); } - let columns = arrow::compute::take_arrays(batch.columns(), indices, None)?; - Ok(Cow::Owned(RecordBatch::try_new_with_options( - batch.schema(), - columns, - &RecordBatchOptions::new() - .with_match_field_names(false) - .with_row_count(Some(indices.len())), + Ok(Cow::Owned(arrow::compute::take_record_batch( + batch, indices, )?)) } @@ -1898,46 +1893,6 @@ mod tests { Ok(()) } - #[test] - fn literal_when_preserves_nested_field_names() -> Result<()> { - let mut payload = ListBuilder::new(Int32Builder::new()).with_field(Field::new( - "value", - DataType::Int32, - true, - )); - for value in [10, 20] { - payload.values().append_value(value); - payload.append(true); - } - let schema = Arc::new(Schema::new(vec![ - Field::new("key", DataType::Int32, false), - Field::new( - "payload", - DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), - false, - ), - ])); - let batch = RecordBatch::try_new_with_options( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(payload.finish()), - ], - &RecordBatchOptions::new().with_match_field_names(false), - )?; - let expr = case( - Some(col("key", &schema)?), - vec![(lit(1i32), expressions::is_null(col("payload", &schema)?)?)], - Some(lit(false)), - )?; - let result = expr.evaluate(&batch)?.into_array(2)?; - assert_eq!( - as_boolean_array(&result)?, - &BooleanArray::from(vec![false, false]) - ); - Ok(()) - } - #[test] fn case_with_expr() -> Result<()> { let batch = case_test_batch()?;