diff --git a/datafusion/functions-nested/src/array_any_match.rs b/datafusion/functions-nested/src/array_any_match.rs index b83c56e9e227f..0f620f18bd8f2 100644 --- a/datafusion/functions-nested/src/array_any_match.rs +++ b/datafusion/functions-nested/src/array_any_match.rs @@ -18,17 +18,11 @@ //! [`datafusion_expr::HigherOrderUDF`] definitions for array_any_match function. use arrow::{ - array::{Array, AsArray, BooleanArray, BooleanBuilder, new_null_array}, + array::{Array, BooleanArray, BooleanBuilder}, buffer::NullBuffer, - compute::take_arrays, - datatypes::{ArrowNativeType, DataType, Field, FieldRef}, -}; -use datafusion_common::{ - Result, exec_datafusion_err, exec_err, plan_err, - utils::{ - adjust_offsets_for_slice, list_values, list_values_row_number, take_function_args, - }, + datatypes::{DataType, Field, FieldRef}, }; +use datafusion_common::{Result, plan_err, utils::take_function_args}; use datafusion_expr::{ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, @@ -37,7 +31,9 @@ use datafusion_expr::{ use datafusion_macros::user_doc; use std::{fmt::Debug, sync::Arc}; -use crate::lambda_utils::coerce_single_list_arg; +use crate::lambda_utils::{ + SingleListLambdaResult, coerce_single_list_arg, evaluate_single_list_predicate, +}; make_higher_order_function_expr_and_func!( ArrayAnyMatch, @@ -160,75 +156,25 @@ impl HigherOrderUDFImpl for ArrayAnyMatch { } fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { - let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)] = - take_function_args(self.name(), &args.args)? - else { - return exec_err!("{} expects a value followed by a lambda", self.name()); + let evaluated = match evaluate_single_list_predicate(self.name(), &args)? { + SingleListLambdaResult::EarlyReturn(v) => return Ok(v), + SingleListLambdaResult::Ready(v) => v, }; - let list_array = list.to_array(args.number_rows)?; + let predicate = evaluated.boolean_predicate(self.name())?; - // fast path: fully null input — also required for FixedSizeList which can't be - // handled by clear_null_values when fully null - if list_array.null_count() == list_array.len() { - return Ok(ColumnarValue::Array(new_null_array( - args.return_type(), - list_array.len(), - ))); - } - - let list_values = list_values(&list_array)?; - - let values_param = || Ok(Arc::clone(&list_values)); - - let predicate_results = lambda - .evaluate(&[&values_param], |arrays| { - let indices = list_values_row_number(&list_array)?; - Ok(take_arrays(arrays, &indices, None)?) - })? - .into_array(list_values.len())?; - - let predicate_bool = predicate_results - .as_any() - .downcast_ref::() - .ok_or_else(|| { - exec_datafusion_err!( - "{} predicate must return boolean array", - self.name() - ) - })?; - - let mut values = BooleanBuilder::with_capacity(list_array.len()); - - // Maps predicate results (flat over all elements) back to one Boolean per row. - // Uses adjusted offsets so sliced lists index correctly into the predicate array. - macro_rules! process_list { - ($list_typed:expr) => {{ - let offsets = adjust_offsets_for_slice($list_typed); - for i in 0..$list_typed.len() { - let start = offsets[i].as_usize(); - let end = offsets[i + 1].as_usize(); - // any_match_for_range returns None when nulls poison the result; - // null rows produce an empty range and return Some(false), but their - // null bit is preserved by attaching the original null bitmap below. - values.append_option(any_match_for_range(predicate_bool, start, end)); - } - }}; - } - - match list_array.data_type() { - DataType::List(_) => { - process_list!(list_array.as_list::()); - } - DataType::LargeList(_) => { - process_list!(list_array.as_list::()); - } - other => return exec_err!("expected list, got {other}"), + let mut values = BooleanBuilder::with_capacity(evaluated.len()); + for i in 0..evaluated.len() { + let (start, end) = evaluated.row_range(i); + // any_match_for_range returns None when nulls poison the result; + // null rows produce an empty range and return Some(false), but their + // null bit is preserved by attaching the original null bitmap below. + values.append_option(any_match_for_range(&predicate, start, end)); } let (boolean_buffer, predicate_nulls) = values.finish().into_parts(); // Merge: a row is null if the input list row was null or the predicate returned null. - let nulls = NullBuffer::union(list_array.nulls(), predicate_nulls.as_ref()); + let nulls = NullBuffer::union(evaluated.nulls(), predicate_nulls.as_ref()); Ok(ColumnarValue::Array(Arc::new(BooleanArray::new( boolean_buffer, nulls, @@ -260,6 +206,10 @@ mod tests { use datafusion_physical_expr::create_physical_expr; use crate::array_any_match::{ArrayAnyMatch, array_any_match_higher_order_function}; + use crate::lambda_utils::test_utils::{ + create_i32_large_list, create_i32_list, eval_hof_on_i32_list, + eval_hof_on_i32_list_with_outer, v, + }; fn run_any_match( list: impl arrow::array::Array + Clone + 'static, @@ -500,4 +450,44 @@ mod tests { ); Ok(()) } + + #[test] + fn test_any_match_large_list_parity() -> Result<()> { + let list = create_i32_large_list( + vec![1, 2, 3], + OffsetBuffer::::from_lengths(vec![3]), + None, + ); + let result = eval_hof_on_i32_list( + array_any_match_higher_order_function(), + list, + v().gt(lit(2i32)), + )?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &BooleanArray::from(vec![Some(true)]) + ); + Ok(()) + } + + #[test] + fn test_any_match_captured_outer_column() -> Result<()> { + let list = create_i32_list( + vec![1, 50, 4, 50, 7, 50], + OffsetBuffer::::from_lengths(vec![2, 2, 2]), + None, + ); + let number = Int32Array::from(vec![10, 40, 60]); + let result = eval_hof_on_i32_list_with_outer( + array_any_match_higher_order_function(), + list, + number, + v().gt(col("number")), + )?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &BooleanArray::from(vec![Some(true), Some(true), Some(false)]) + ); + Ok(()) + } } diff --git a/datafusion/functions-nested/src/array_filter.rs b/datafusion/functions-nested/src/array_filter.rs index 7dd7230ae9e06..3439699433272 100644 --- a/datafusion/functions-nested/src/array_filter.rs +++ b/datafusion/functions-nested/src/array_filter.rs @@ -23,13 +23,10 @@ use arrow::{ OffsetSizeTrait, new_empty_array, }, buffer::{OffsetBuffer, ScalarBuffer}, - compute::{filter as arrow_filter, take_arrays}, + compute::filter as arrow_filter, datatypes::{DataType, Field, FieldRef}, }; -use datafusion_common::{ - Result, ScalarValue, exec_err, - utils::{adjust_offsets_for_slice, list_values_row_number}, -}; +use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, @@ -39,7 +36,7 @@ use datafusion_macros::user_doc; use std::sync::Arc; use crate::lambda_utils::{ - ListValuesResult, coerce_single_list_arg, extract_list_values, + SingleListLambdaResult, coerce_single_list_arg, evaluate_single_list_predicate, single_list_lambda_parameters, value_lambda_pair, }; @@ -130,12 +127,9 @@ impl HigherOrderUDFImpl for ArrayFilter { } fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { - let (list, lambda) = value_lambda_pair(self.name(), &args.args)?; - let list_array = list.to_array(args.number_rows)?; - - let list_values = match extract_list_values(&list_array, args.return_type())? { - ListValuesResult::EarlyReturn(v) => return Ok(v), - ListValuesResult::Values(v) => v, + let evaluated = match evaluate_single_list_predicate(self.name(), &args)? { + SingleListLambdaResult::EarlyReturn(v) => return Ok(v), + SingleListLambdaResult::Ready(v) => v, }; let field = match args.return_field.data_type() { @@ -149,56 +143,47 @@ impl HigherOrderUDFImpl for ArrayFilter { } }; - let values_param = || Ok(Arc::clone(&list_values)); - let predicate_output = lambda.evaluate(&[&values_param], |arrays| { - let indices = list_values_row_number(&list_array)?; - Ok(take_arrays(arrays, &indices, None)?) - })?; - // Scalar predicate short-circuit: x -> true or x -> false/null - if let ColumnarValue::Scalar(ScalarValue::Boolean(b)) = &predicate_output { + if let ColumnarValue::Scalar(ScalarValue::Boolean(b)) = + &evaluated.evaluated_result + { return match b { - Some(true) => Ok(ColumnarValue::Array(list_array)), + Some(true) => Ok(ColumnarValue::Array(evaluated.original_list)), _ => Ok(ColumnarValue::Array(empty_filtered_list( - &list_array, + &evaluated.original_list, field, )?)), }; } - let predicate = predicate_output.into_array(list_values.len())?; - let Some(predicate) = predicate.as_any().downcast_ref::() else { - return exec_err!( - "{} lambda must return boolean, got {}", - self.name(), - predicate.data_type() - ); - }; + let predicate = evaluated.boolean_predicate(self.name())?; // ListView and LargeListView are coerced to List/LargeList by coerce_value_types. - let filtered_list = match list_array.data_type() { + let filtered_list = match evaluated.original_list.data_type() { DataType::List(_) => { - let list = list_array.as_list::(); - let adjusted_offsets = adjust_offsets_for_slice(list); - let (filtered_values, new_offsets) = - filter_list_values(&list_values, predicate, &adjusted_offsets)?; + let (filtered_values, new_offsets) = filter_list_values( + &evaluated.flattened_values, + &predicate, + &evaluated.adjusted_offsets::(), + )?; Arc::new(ListArray::new( field, new_offsets, filtered_values, - list.nulls().cloned(), + evaluated.nulls().cloned(), )) as ArrayRef } DataType::LargeList(_) => { - let large_list = list_array.as_list::(); - let adjusted_offsets = adjust_offsets_for_slice(large_list); - let (filtered_values, new_offsets) = - filter_list_values(&list_values, predicate, &adjusted_offsets)?; + let (filtered_values, new_offsets) = filter_list_values( + &evaluated.flattened_values, + &predicate, + &evaluated.adjusted_offsets::(), + )?; Arc::new(LargeListArray::new( field, new_offsets, filtered_values, - large_list.nulls().cloned(), + evaluated.nulls().cloned(), )) } other => exec_err!("expected list, got {other}")?, @@ -284,9 +269,14 @@ mod tests { buffer::{NullBuffer, OffsetBuffer}, }; + use arrow::array::Int32Array; + use crate::array_filter::array_filter_higher_order_function; - use crate::lambda_utils::test_utils::{create_i32_list, eval_hof_on_i32_list, v}; - use datafusion_expr::lit; + use crate::lambda_utils::test_utils::{ + create_i32_large_list, create_i32_list, eval_hof_on_i32_list, + eval_hof_on_i32_list_with_outer, v, + }; + use datafusion_expr::{col, lit}; fn keep_greater_than_two( list: impl Array + Clone + 'static, @@ -456,4 +446,45 @@ mod tests { ); assert_eq!(actual, &expected); } + + #[test] + fn filter_large_list_parity() { + let list = create_i32_large_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![5]), + None, + ); + let res = keep_greater_than_two(list).unwrap(); + let actual = res.as_list::(); + let expected = create_i32_large_list( + vec![3, 4, 5], + OffsetBuffer::::from_lengths(vec![3]), + None, + ); + assert_eq!(actual, &expected); + } + + #[test] + fn filter_captured_outer_column() { + let list = create_i32_list( + vec![1, 50, 4, 50, 7, 50], + OffsetBuffer::::from_lengths(vec![2, 2, 2]), + None, + ); + let number = Int32Array::from(vec![10, 40, 60]); + let res = eval_hof_on_i32_list_with_outer( + array_filter_higher_order_function(), + list, + number, + v().gt(col("number")), + ) + .unwrap(); + let actual = res.as_list::(); + let expected = create_i32_list( + vec![50, 50], + OffsetBuffer::::from_lengths(vec![1, 1, 0]), + None, + ); + assert_eq!(actual, &expected); + } } diff --git a/datafusion/functions-nested/src/array_first.rs b/datafusion/functions-nested/src/array_first.rs index 07154a2db74b7..615dc47394379 100644 --- a/datafusion/functions-nested/src/array_first.rs +++ b/datafusion/functions-nested/src/array_first.rs @@ -18,17 +18,11 @@ //! [`datafusion_expr::HigherOrderUDF`] definitions for array_first function. use arrow::{ - array::{ - Array, AsArray, BooleanArray, GenericListArray, OffsetSizeTrait, UInt64Array, - UInt64Builder, new_null_array, - }, - compute::{take, take_arrays}, + array::{Array, BooleanArray, UInt64Array, UInt64Builder}, + compute::take, datatypes::{DataType, FieldRef}, }; -use datafusion_common::{ - Result, exec_datafusion_err, exec_err, plan_err, - utils::{adjust_offsets_for_slice, list_values, list_values_row_number}, -}; +use datafusion_common::{Result, exec_err, plan_err}; use datafusion_expr::{ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, @@ -38,7 +32,8 @@ use datafusion_macros::user_doc; use std::sync::Arc; use crate::lambda_utils::{ - coerce_single_list_arg, single_list_lambda_parameters, value_lambda_pair, + EvaluatedListLambda, SingleListLambdaResult, coerce_single_list_arg, + evaluate_single_list_predicate, single_list_lambda_parameters, value_lambda_pair, }; make_higher_order_function_expr_and_func!( @@ -147,58 +142,20 @@ impl HigherOrderUDFImpl for ArrayFirst { } fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { - let (list, lambda) = value_lambda_pair(self.name(), &args.args)?; - - let list_array = list.to_array(args.number_rows)?; - - // Fast path: fully null input. Also required for FixedSizeList which - // can't be handled by clear_null_values when fully null. - if list_array.null_count() == list_array.len() { - return Ok(ColumnarValue::Array(new_null_array( - args.return_type(), - list_array.len(), - ))); - } + let evaluated = match evaluate_single_list_predicate(self.name(), &args)? { + SingleListLambdaResult::EarlyReturn(v) => return Ok(v), + SingleListLambdaResult::Ready(v) => v, + }; - let list_values = list_values(&list_array)?; - - // Evaluate the predicate over every flat element. Captured columns are - // spread to align with the flattened values via list_values_row_number. - let values_param = || Ok(Arc::clone(&list_values)); - - let predicate_results = lambda - .evaluate(&[&values_param], |arrays| { - let indices = list_values_row_number(&list_array)?; - Ok(take_arrays(arrays, &indices, None)?) - })? - .into_array(list_values.len())?; - - let predicate_bool = predicate_results - .as_any() - .downcast_ref::() - .ok_or_else(|| { - exec_datafusion_err!( - "{} predicate must return boolean array, got {}", - self.name(), - predicate_results.data_type() - ) - })?; - - // For each row, find the flat index of the first element whose predicate - // is true. Rows with no match, including empty rows and null rows that - // clear_null_values truncated to empty, map to a null index, producing - // a null result via `take`. - let indices = match list_array.data_type() { - DataType::List(_) => { - first_match_indices(list_array.as_list::(), predicate_bool) - } - DataType::LargeList(_) => { - first_match_indices(list_array.as_list::(), predicate_bool) + let predicate = evaluated.boolean_predicate(self.name())?; + let indices = match evaluated.original_list.data_type() { + DataType::List(_) | DataType::LargeList(_) => { + first_match_indices(&evaluated, &predicate) } other => return exec_err!("expected list, got {other}"), }; - let result = take(list_values.as_ref(), &indices, None)?; + let result = take(evaluated.flattened_values.as_ref(), &indices, None)?; Ok(ColumnarValue::Array(result)) } @@ -213,18 +170,14 @@ impl HigherOrderUDFImpl for ArrayFirst { /// /// A null predicate value is treated as not matching. The matched element itself /// may be null and is still returned. -fn first_match_indices( - list: &GenericListArray, +fn first_match_indices( + evaluated: &EvaluatedListLambda, predicate: &BooleanArray, ) -> UInt64Array { - // Offsets are adjusted so that sliced lists index correctly into the - // predicate / values arrays returned by list_values. - let offsets = adjust_offsets_for_slice(list); - let mut builder = UInt64Builder::with_capacity(list.len()); + let mut builder = UInt64Builder::with_capacity(evaluated.len()); - for i in 0..list.len() { - let start = offsets[i].as_usize(); - let end = offsets[i + 1].as_usize(); + for i in 0..evaluated.len() { + let (start, end) = evaluated.row_range(i); match (start..end).find(|&j| predicate.is_valid(j) && predicate.value(j)) { Some(j) => builder.append_value(j as u64), @@ -244,9 +197,12 @@ mod tests { }; use crate::array_first::array_first_higher_order_function; - use crate::lambda_utils::test_utils::{create_i32_list, eval_hof_on_i32_list, v}; + use crate::lambda_utils::test_utils::{ + create_i32_large_list, create_i32_list, eval_hof_on_i32_list, + eval_hof_on_i32_list_with_outer, v, + }; use datafusion_common::Result; - use datafusion_expr::lit; + use datafusion_expr::{col, lit}; fn first_greater_than_two( list: impl Array + Clone + 'static, @@ -411,6 +367,42 @@ mod tests { ); } + #[test] + fn test_first_large_list_parity() -> Result<()> { + let list = create_i32_large_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![5]), + None, + ); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(3)]) + ); + Ok(()) + } + + #[test] + fn test_first_captured_outer_column() -> Result<()> { + let list = create_i32_list( + vec![1, 50, 4, 50, 7, 50], + OffsetBuffer::::from_lengths(vec![2, 2, 2]), + None, + ); + let number = Int32Array::from(vec![10, 40, 60]); + let res = eval_hof_on_i32_list_with_outer( + array_first_higher_order_function(), + list, + number, + v().gt(col("number")), + )?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(50), Some(50), None]) + ); + Ok(()) + } + #[test] fn test_first_string_elements() -> Result<()> { use arrow::array::ListArray; diff --git a/datafusion/functions-nested/src/lambda_utils.rs b/datafusion/functions-nested/src/lambda_utils.rs index 927ca5a51461c..1b620456a2493 100644 --- a/datafusion/functions-nested/src/lambda_utils.rs +++ b/datafusion/functions-nested/src/lambda_utils.rs @@ -17,13 +17,19 @@ //! Shared utilities for `(array, lambda)` style higher-order functions. -use arrow::array::ArrayRef; -use arrow::datatypes::{DataType, FieldRef}; +use arrow::array::{ArrayRef, AsArray, BooleanArray, OffsetSizeTrait, new_null_array}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::compute::take_arrays; +use arrow::datatypes::{ArrowNativeType, DataType, FieldRef}; +use datafusion_common::utils::{adjust_offsets_for_slice, list_values_row_number}; use datafusion_common::{ Result, ScalarValue, plan_err, utils::{list_values, take_function_args}, }; -use datafusion_expr::{ColumnarValue, LambdaParametersProgress, ValueOrLambda}; +use datafusion_common::{exec_datafusion_err, exec_err}; +use datafusion_expr::{ + ColumnarValue, HigherOrderFunctionArgs, LambdaParametersProgress, ValueOrLambda, +}; use std::sync::Arc; /// Extracts a `(value, lambda)` pair from a [`ValueOrLambda`] slice. @@ -126,12 +132,225 @@ pub(crate) fn extract_list_values( Ok(ListValuesResult::Values(values)) } +pub(crate) enum SingleListLambdaResult { + EarlyReturn(ColumnarValue), + Ready(EvaluatedListLambda), +} + +pub(crate) struct EvaluatedListLambda { + pub original_list: ArrayRef, + pub flattened_values: ArrayRef, + pub evaluated_result: ColumnarValue, + row_offsets: Vec, +} + +impl EvaluatedListLambda { + pub(crate) fn len(&self) -> usize { + self.original_list.len() + } + + pub(crate) fn nulls(&self) -> Option<&NullBuffer> { + self.original_list.nulls() + } + + pub(crate) fn row_range(&self, i: usize) -> (usize, usize) { + (self.row_offsets[i], self.row_offsets[i + 1]) + } + + pub(crate) fn adjusted_offsets(&self) -> OffsetBuffer { + OffsetBuffer::from_lengths(self.row_offsets.windows(2).map(|w| w[1] - w[0])) + } + + pub(crate) fn boolean_predicate(&self, name: &str) -> Result { + let arr = self + .evaluated_result + .clone() + .into_array(self.flattened_values.len())?; + + let predicate = arr.as_any().downcast_ref::().ok_or_else(|| { + exec_datafusion_err!("{} predicate must return boolean array", name) + })?; + + Ok(predicate.clone()) + } +} + +fn adjusted_row_offsets(list: &ArrayRef) -> Result> { + Ok(match list.data_type() { + DataType::List(_) => adjust_offsets_for_slice(list.as_list::()) + .iter() + .map(|o| o.as_usize()) + .collect(), + DataType::LargeList(_) => adjust_offsets_for_slice(list.as_list::()) + .iter() + .map(|o| o.as_usize()) + .collect(), + other => return exec_err!("expected list, got {other}"), + }) +} + +fn evaluate_single_list_lambda( + name: &str, + args: &HigherOrderFunctionArgs, +) -> Result { + let (original_list, lambda) = value_lambda_pair(name, &args.args)?; + let original_list = original_list.to_array(args.number_rows)?; + + if original_list.null_count() == original_list.len() { + return Ok(SingleListLambdaResult::EarlyReturn(ColumnarValue::Array( + new_null_array(args.return_type(), original_list.len()), + ))); + } + + let flattened_values = list_values(&original_list)?; + let values_param = || Ok(Arc::clone(&flattened_values)); + + let evaluated_result = lambda.evaluate(&[&values_param], |arrays| { + let indices = list_values_row_number(&original_list)?; + Ok(take_arrays(arrays, &indices, None)?) + })?; + + let row_offsets = adjusted_row_offsets(&original_list)?; + + Ok(SingleListLambdaResult::Ready(EvaluatedListLambda { + original_list, + flattened_values, + evaluated_result, + row_offsets, + })) +} + +pub(crate) fn evaluate_single_list_predicate( + name: &str, + args: &HigherOrderFunctionArgs, +) -> Result { + let result = evaluate_single_list_lambda(name, args)?; + let SingleListLambdaResult::Ready(evaluated_list_lambda) = &result else { + return Ok(result); + }; + + match &evaluated_list_lambda.evaluated_result { + ColumnarValue::Scalar(ScalarValue::Boolean(_)) => Ok(result), + ColumnarValue::Scalar(scalar) => exec_err!( + "{name} lambda must return boolean, got {}", + scalar.data_type() + ), + ColumnarValue::Array(array) if array.as_any().is::() => Ok(result), + ColumnarValue::Array(array) => exec_err!( + "{name} lambda must return boolean, got {}", + array.data_type() + ), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::{ + array::{Array, ArrayRef}, + buffer::{NullBuffer, OffsetBuffer}, + datatypes::{DataType, Field}, + }; + use datafusion_common::Result; + + use super::{adjusted_row_offsets, coerce_single_list_arg}; + use crate::lambda_utils::test_utils::{create_i32_large_list, create_i32_list}; + + #[test] + fn adjusted_row_offsets_matches_list_lengths() -> Result<()> { + let list = create_i32_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![2, 0, 3]), + None, + ); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 2, 2, 5]); + Ok(()) + } + + #[test] + fn adjusted_row_offsets_on_sliced_list() -> Result<()> { + let list = create_i32_list( + vec![10, 1, 2, 3, 4], + OffsetBuffer::::from_lengths(vec![1, 2, 2]), + None, + ) + .slice(1, 2); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 2, 4]); + Ok(()) + } + + #[test] + fn adjusted_row_offsets_null_rows_keep_backing_lengths() -> Result<()> { + let list = create_i32_list( + vec![1, 99, 100, 2], + OffsetBuffer::::from_lengths(vec![1, 2, 1]), + Some(NullBuffer::from(vec![true, false, true])), + ); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 1, 3, 4]); + Ok(()) + } + + #[test] + fn adjusted_row_offsets_large_list_parity() -> Result<()> { + let list = create_i32_large_list( + vec![1, 2, 3, 4], + OffsetBuffer::::from_lengths(vec![1, 3]), + None, + ); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 1, 4]); + Ok(()) + } + + #[test] + fn coerce_single_list_arg_supports_advertised_list_likes() -> Result<()> { + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + assert_eq!( + coerce_single_list_arg("test", &[DataType::List(Arc::clone(&field))])?, + vec![DataType::List(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg("test", &[DataType::LargeList(Arc::clone(&field))])?, + vec![DataType::LargeList(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg( + "test", + &[DataType::FixedSizeList(Arc::clone(&field), 3)] + )?, + vec![DataType::List(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg("test", &[DataType::ListView(Arc::clone(&field))])?, + vec![DataType::List(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg( + "test", + &[DataType::LargeListView(Arc::clone(&field))] + )?, + vec![DataType::LargeList(field)] + ); + Ok(()) + } + + #[test] + fn coerce_single_list_arg_rejects_non_list() { + let err = coerce_single_list_arg("test", &[DataType::Int32]).unwrap_err(); + assert!(err.to_string().contains("expected a list")); + } +} + #[cfg(test)] pub(crate) mod test_utils { use std::{collections::HashMap, sync::Arc}; use arrow::{ - array::{Array, ArrayRef, Int32Array, ListArray, RecordBatch}, + array::{Array, ArrayRef, Int32Array, LargeListArray, ListArray, RecordBatch}, buffer::{NullBuffer, OffsetBuffer}, datatypes::{DataType, Field}, }; @@ -154,6 +373,15 @@ pub(crate) mod test_utils { ListArray::new(list_field, offsets, Arc::new(values.into()), nulls) } + pub(crate) fn create_i32_large_list( + values: impl Into, + offsets: OffsetBuffer, + nulls: Option, + ) -> LargeListArray { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + LargeListArray::new(list_field, offsets, Arc::new(values.into()), nulls) + } + pub(crate) fn eval_hof_on_i32_list( func: Arc, list: impl Array + Clone + 'static, @@ -185,6 +413,39 @@ pub(crate) mod test_utils { .into_array(list.len()) } + /// Evaluates a HOF whose lambda body may capture an outer `number` column. + pub(crate) fn eval_hof_on_i32_list_with_outer( + func: Arc, + list: impl Array + Clone + 'static, + number: Int32Array, + lambda_body: Expr, + ) -> Result { + assert_eq!(list.len(), number.len()); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("list", list.data_type().clone(), list.is_nullable()), + Field::new("number", DataType::Int32, true), + ] + .into(), + HashMap::new(), + )?; + + create_physical_expr( + &Expr::HigherOrderFunction(HigherOrderFunction::new( + func, + vec![col("list"), lambda(["v"], lambda_body)], + )), + &schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )? + .evaluate(&RecordBatch::try_new( + Arc::clone(schema.inner()), + vec![Arc::new(list.clone()), Arc::new(number)], + )?)? + .into_array(list.len()) + } + pub(crate) fn v() -> Expr { Expr::LambdaVariable(LambdaVariable::new( "v".to_string(), diff --git a/datafusion/sqllogictest/test_files/array/array_any_match.slt b/datafusion/sqllogictest/test_files/array/array_any_match.slt index 37aa47c55adcf..82133054e118a 100644 --- a/datafusion/sqllogictest/test_files/array/array_any_match.slt +++ b/datafusion/sqllogictest/test_files/array/array_any_match.slt @@ -109,6 +109,29 @@ SELECT array_any_match(NULL, x -> x > 2); ---- NULL +# predicate can reference an outer column +query B +SELECT array_any_match(list, x -> x > number) FROM t; +---- +true +true +false + +# large list works +query B +SELECT array_any_match(arrow_cast([1, 2, 3], 'LargeList(Int32)'), x -> x > 2); +---- +true + +# other list representations are coerced during planning +query BBB +SELECT + array_any_match(arrow_cast([1, 2, 3], 'FixedSizeList(3, Int32)'), x -> x > 2), + array_any_match(arrow_cast([1, 2, 3], 'ListView(Int32)'), x -> x > 2), + array_any_match(arrow_cast([1, 2, 3], 'LargeListView(Int32)'), x -> x > 2); +---- +true true true + statement ok drop table t; diff --git a/datafusion/sqllogictest/test_files/array/array_filter.slt b/datafusion/sqllogictest/test_files/array/array_filter.slt index 9b564c5061205..b6d73fbe7d09d 100644 --- a/datafusion/sqllogictest/test_files/array/array_filter.slt +++ b/datafusion/sqllogictest/test_files/array/array_filter.slt @@ -120,6 +120,20 @@ SELECT array_filter(arrow_cast(list, 'ListView(Int32)'), v -> v > 2) from t; [4, 50] [7, 50] +# large list works +query ? +SELECT array_filter(arrow_cast([1, 2, 3, 4, 5], 'LargeList(Int32)'), v -> v > 2); +---- +[3, 4, 5] + +# FixedSizeList / LargeListView coercions during planning +query ?? +SELECT + array_filter(arrow_cast([1, 2, 3, 4], 'FixedSizeList(4, Int32)'), v -> v > 2), + array_filter(arrow_cast([1, 2, 3, 4], 'LargeListView(Int32)'), v -> v > 2); +---- +[3, 4] [3, 4] + # null array argument returns null query ? SELECT array_filter(arrow_cast(NULL, 'List(Int32)'), v -> v > 0);