From 7392c29183adc7daf26de7201923ce903c0a2f96 Mon Sep 17 00:00:00 2001 From: Han-Wen Tsao Date: Sun, 25 Oct 2026 12:13:33 +0000 Subject: [PATCH] fix: compare every row when arrays_overlap gets a scalar beside a column (#6743) * fix: compare every row when arrays_overlap gets a scalar beside a column When one argument was a scalar, SparkArraysOverlap turned it into a one-row list while the column kept its full length, and returned row 0's result as a scalar. A constant array beside an array column therefore answered every row with the first row's result, or panicked with an index out of bounds when the constant was the second argument. Keep the scalar as a one-row list that every output row reads, and return an array result when the other argument is an array. Repeating the scalar to the batch length instead could overflow the list's 32-bit offsets for a long constant. Two scalar arguments still produce a scalar result. Assisted-by: Claude Code (claude-opus-5-5) * test: reuse the arrays_overlap fixture tables for the scalar case Write test_arrays_overlap and test_overlap_struct as one file each, so their rows share a batch. The existing constant-second queries then cover the scalar-beside-column path, and a constant-first query with answers that change by row, plus a struct query in both orders, complete it. This replaces the separate test_overlap_batch table and the flat and struct Rust tests it duplicated. The long-constant and two-scalar Rust tests stay, since a fixture cannot reach them. Assisted-by: Claude Code (claude-opus-5-5) (cherry picked from commit c39eb711e3e5a11aec73e1f685db2d67d10ed15b) --- .../src/array_funcs/arrays_overlap.rs | 175 ++++++++++++++++-- .../expressions/array/arrays_overlap.sql | 14 +- 2 files changed, 170 insertions(+), 19 deletions(-) diff --git a/native/spark-expr/src/array_funcs/arrays_overlap.rs b/native/spark-expr/src/array_funcs/arrays_overlap.rs index c761a9a51bd..ce30e96c142 100644 --- a/native/spark-expr/src/array_funcs/arrays_overlap.rs +++ b/native/spark-expr/src/array_funcs/arrays_overlap.rs @@ -156,26 +156,43 @@ impl ScalarUDFImpl for SparkArraysOverlap { Ok(ColumnarValue::Array(result)) } (left, right) => { - // Handle scalar inputs by converting to arrays + // At least one side is a scalar. It stays a one-row list that every output row + // reads, rather than being repeated to the batch length: repeating a long list + // could overflow the list's offsets. Only two scalars produce a scalar result. + let left_scalar = matches!(left, ColumnarValue::Scalar(_)); + let right_scalar = matches!(right, ColumnarValue::Scalar(_)); + let both_scalar = left_scalar && right_scalar; + let rows = RowMapping { + len: if both_scalar { 1 } else { args.number_rows }, + left_scalar, + right_scalar, + }; let left_arr = left.to_array(1)?; let right_arr = right.to_array(1)?; let result = match (left_arr.data_type(), right_arr.data_type()) { - (DataType::List(_), DataType::List(_)) => arrays_overlap_list::( + (DataType::List(_), DataType::List(_)) => arrays_overlap_rows::( left_arr.as_any().downcast_ref().unwrap(), right_arr.as_any().downcast_ref().unwrap(), self.normalize_floats, + rows, )?, - (DataType::LargeList(_), DataType::LargeList(_)) => arrays_overlap_list::( + (DataType::LargeList(_), DataType::LargeList(_)) => arrays_overlap_rows::( left_arr.as_any().downcast_ref().unwrap(), right_arr.as_any().downcast_ref().unwrap(), self.normalize_floats, + rows, )?, (l, r) => { return exec_err!("{} does not support types '{l}' and '{r}'", self.name()) } }; - let scalar = ScalarValue::try_from_array(&result, 0)?; - Ok(ColumnarValue::Scalar(scalar)) + if both_scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(result)) + } } } } @@ -194,19 +211,69 @@ fn arrays_overlap_list( left: &GenericListArray, right: &GenericListArray, normalize_floats: bool, +) -> Result { + arrays_overlap_rows( + left, + right, + normalize_floats, + RowMapping::aligned(left.len()), + ) +} + +/// Which input row each output row reads. A scalar input is a one-row list that every output row +/// reads; otherwise output row `i` reads input row `i`. +#[derive(Clone, Copy)] +struct RowMapping { + len: usize, + left_scalar: bool, + right_scalar: bool, +} + +impl RowMapping { + fn aligned(len: usize) -> Self { + Self { + len, + left_scalar: false, + right_scalar: false, + } + } + + fn left(&self, i: usize) -> usize { + if self.left_scalar { + 0 + } else { + i + } + } + + fn right(&self, i: usize) -> usize { + if self.right_scalar { + 0 + } else { + i + } + } +} + +/// `arrays_overlap_list` over `rows.len` output rows, reading the input rows `rows` maps them to. +fn arrays_overlap_rows( + left: &GenericListArray, + right: &GenericListArray, + normalize_floats: bool, + rows: RowMapping, ) -> Result { let left_values = left.values(); let right_values = right.values(); if left_values.data_type() != right_values.data_type() { - return arrays_overlap_list_generic(left, right); + return arrays_overlap_list_generic(left, right, rows); } // Fast paths for flat element types: probe the flat value buffers directly instead of // slicing each row and running an Arrow compare kernel once per probe element. macro_rules! flat_fast_path { ($l:expr, $r:expr) => { - return Ok(overlap_rows(left, right, flat_row_overlap($l, $r))) + return Ok(overlap_rows(left, right, rows, flat_row_overlap($l, $r))) }; } macro_rules! primitive_fast_path { @@ -263,7 +330,7 @@ fn arrays_overlap_list( left_values.as_string::(), right_values.as_string::() ), - _ => arrays_overlap_list_generic(left, right), + _ => arrays_overlap_list_generic(left, right, rows), } } @@ -273,9 +340,10 @@ fn arrays_overlap_list( fn overlap_rows( left: &GenericListArray, right: &GenericListArray, + rows: RowMapping, mut row_overlap: impl FnMut(Range, Range) -> bool, ) -> ArrayRef { - let len = left.len(); + let len = rows.len; let left_offsets = left.offsets(); let right_offsets = right.offsets(); let left_element_nulls = left.values().nulls(); @@ -284,13 +352,14 @@ fn overlap_rows( let mut builder = BooleanArray::builder(len); for i in 0..len { - if left.is_null(i) || right.is_null(i) { + let (li, ri) = (rows.left(i), rows.right(i)); + if left.is_null(li) || right.is_null(ri) { builder.append_null(); continue; } - let left_range = left_offsets[i].as_usize()..left_offsets[i + 1].as_usize(); - let right_range = right_offsets[i].as_usize()..right_offsets[i + 1].as_usize(); + let left_range = left_offsets[li].as_usize()..left_offsets[li + 1].as_usize(); + let right_range = right_offsets[ri].as_usize()..right_offsets[ri + 1].as_usize(); if left_range.is_empty() || right_range.is_empty() { builder.append_value(false); @@ -476,6 +545,7 @@ fn normalize_list_element_floats( fn arrays_overlap_list_generic( left: &GenericListArray, right: &GenericListArray, + rows: RowMapping, ) -> Result { let left_owned = has_float_leaf(left.values().data_type()).then(|| normalize_list_element_floats(left)); @@ -484,17 +554,18 @@ fn arrays_overlap_list_generic( has_float_leaf(right.values().data_type()).then(|| normalize_list_element_floats(right)); let right: &GenericListArray = right_owned.as_ref().unwrap_or(right); - let len = left.len(); + let len = rows.len; let mut builder = BooleanArray::builder(len); for i in 0..len { - if left.is_null(i) || right.is_null(i) { + let (li, ri) = (rows.left(i), rows.right(i)); + if left.is_null(li) || right.is_null(ri) { builder.append_null(); continue; } - let left_values = left.value(i); - let right_values = right.value(i); + let left_values = left.value(li); + let right_values = right.value(ri); if left_values.is_empty() || right_values.is_empty() { builder.append_value(false); @@ -602,7 +673,7 @@ mod tests { use super::*; use arrow::array::{ Float64Array, Float64Builder, Int32Array, Int32Builder, LargeListArray, ListArray, - ListBuilder, StructArray, StructBuilder, + ListBuilder, NullArray, StructArray, StructBuilder, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{Field, Fields}; @@ -1372,4 +1443,74 @@ mod tests { assert!(result.is_null(0)); Ok(()) } + + fn invoke_columnar( + left: ColumnarValue, + right: ColumnarValue, + rows: usize, + ) -> Result { + SparkArraysOverlap::new().invoke_with_args(ScalarFunctionArgs { + args: vec![left, right], + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("result", DataType::Boolean, true)), + config_options: Arc::new(datafusion::config::ConfigOptions::default()), + }) + } + + fn scalar_list(list: ListArray) -> ColumnarValue { + ColumnarValue::Scalar(ScalarValue::List(Arc::new(list))) + } + + fn booleans(result: ColumnarValue) -> Vec> { + match result { + ColumnarValue::Array(array) => array + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect(), + ColumnarValue::Scalar(_) => panic!("expected an array result"), + } + } + + /// A long constant list is read in place for every row rather than repeated to the batch + /// length: 262,144 elements repeated over 8,192 rows would overflow a list's 32-bit offsets. + #[test] + fn long_scalar_list_is_not_repeated() -> Result<()> { + let field = Arc::new(Field::new("item", DataType::Null, true)); + let constant = ListArray::new( + Arc::clone(&field), + OffsetBuffer::new(vec![0, 262_144].into()), + Arc::new(NullArray::new(262_144)), + None, + ); + let rows = 8_192; + let column = ListArray::new( + field, + OffsetBuffer::new(vec![0; rows + 1].into()), + Arc::new(NullArray::new(0)), + None, + ); + + let result = invoke_columnar( + scalar_list(constant), + ColumnarValue::Array(Arc::new(column)), + rows, + )?; + assert_eq!(booleans(result), vec![Some(false); rows]); + Ok(()) + } + + /// Two constants still produce a constant. + #[test] + fn two_scalars_return_a_scalar() -> Result<()> { + let one = make_list_array(&Int32Array::from(vec![1]), &[0, 1], None); + let result = invoke_columnar(scalar_list(one.clone()), scalar_list(one), 4)?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) + )); + Ok(()) + } } diff --git a/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql index 1ddfe14dbcd..7afb001dcce 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql @@ -18,8 +18,11 @@ statement CREATE TABLE test_arrays_overlap(a array, b array) USING parquet +-- COALESCE(1) writes the table as one file, so the rows share a batch: a constant array reaches +-- native as a scalar, and beside a column it has to be compared with every row, not only the +-- first. statement -INSERT INTO test_arrays_overlap VALUES (array(1, 2, 3), array(3, 4, 5)), (array(1, 2), array(3, 4)), (array(), array(1)), (NULL, array(1)), (array(1, NULL), array(NULL, 2)) +INSERT INTO test_arrays_overlap SELECT /*+ COALESCE(1) */ * FROM VALUES (array(1, 2, 3), array(3, 4, 5)), (array(1, 2), array(3, 4)), (array(), array(1)), (NULL, array(1)), (array(1, NULL), array(NULL, 2)) AS t(a, b) query SELECT arrays_overlap(a, b) FROM test_arrays_overlap @@ -32,6 +35,9 @@ SELECT arrays_overlap(a, array(3, 4, 5)) FROM test_arrays_overlap query SELECT arrays_overlap(array(1, 2, 3), b) FROM test_arrays_overlap +query +SELECT arrays_overlap(array(3, 4, 5), a) FROM test_arrays_overlap + -- literal + literal query SELECT arrays_overlap(array(1, 2, 3), array(3, 4, 5)), arrays_overlap(array(1, 2), array(3, 4)), arrays_overlap(array(), array(1)), arrays_overlap(cast(NULL as array), array(1)) @@ -269,12 +275,16 @@ SELECT a, b, arrays_overlap(a, b) FROM test_overlap_struct_dbl statement CREATE TABLE test_overlap_struct(a array>, b array>) USING parquet +-- One file, so the rows share a batch for the struct constants below. statement -INSERT INTO test_overlap_struct VALUES (array(named_struct('x', 1, 'y', 2)), array(named_struct('x', 1, 'y', 2))), (array(named_struct('x', 1, 'y', 2)), array(named_struct('x', 3, 'y', 4))), (array(named_struct('x', 1, 'y', cast(NULL as int))), array(named_struct('x', 1, 'y', cast(NULL as int)))), (array(cast(NULL as struct)), array(cast(NULL as struct))) +INSERT INTO test_overlap_struct SELECT /*+ COALESCE(1) */ * FROM VALUES (array(named_struct('x', 1, 'y', 2)), array(named_struct('x', 1, 'y', 2))), (array(named_struct('x', 1, 'y', 2)), array(named_struct('x', 3, 'y', 4))), (array(named_struct('x', 1, 'y', cast(NULL as int))), array(named_struct('x', 1, 'y', cast(NULL as int)))), (array(cast(NULL as struct)), array(cast(NULL as struct))) AS t(a, b) query SELECT a, b, arrays_overlap(a, b) FROM test_overlap_struct +query +SELECT b, arrays_overlap(b, array(named_struct('x', 3, 'y', 4))), arrays_overlap(array(named_struct('x', 3, 'y', 4)), b) FROM test_overlap_struct + -- mixed column and literal with NULL elements query SELECT arrays_overlap(a, array(99, NULL)) FROM test_arrays_overlap