Skip to content

feat: centralizing higher-order list lambda evaluation helpers (DRAFT) #23911

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
saadtajwar wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

feat: centralizing higher-order list lambda evaluation helpers (DRAFT) #23911

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3228c33
feat: centralizing higher-order lis lambda evaluation helpers
saadtajwar Jul 26, 2026
071f404
replaced array_first with helper
saadtajwar Jul 26, 2026
d6883b5
removed imports
saadtajwar Jul 26, 2026
6024847
migrated array_filter and any_match to use helper
saadtajwar Jul 26, 2026
29ffb87
tests added
saadtajwar Jul 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 62 additions & 72 deletions datafusion/functions-nested/src/array_any_match.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -160,75 +156,25 @@ impl HigherOrderUDFImpl for ArrayAnyMatch {
}

fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue> {
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::<BooleanArray>()
.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::<i32>());
}
DataType::LargeList(_) => {
process_list!(list_array.as_list::<i64>());
}
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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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::<i64>::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::<BooleanArray>().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::<i32>::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::<BooleanArray>().unwrap(),
&BooleanArray::from(vec![Some(true), Some(true), Some(false)])
);
Ok(())
}
}
115 changes: 73 additions & 42 deletions datafusion/functions-nested/src/array_filter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
};

Expand Down Expand Up @@ -130,12 +127,9 @@ impl HigherOrderUDFImpl for ArrayFilter {
}

fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue> {
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() {
Expand All @@ -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::<BooleanArray>() 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::<i32>();
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::<i32>(),
)?;
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::<i64>();
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::<i64>(),
)?;
Arc::new(LargeListArray::new(
field,
new_offsets,
filtered_values,
large_list.nulls().cloned(),
evaluated.nulls().cloned(),
))
}
other => exec_err!("expected list, got {other}")?,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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::<i64>::from_lengths(vec![5]),
None,
);
let res = keep_greater_than_two(list).unwrap();
let actual = res.as_list::<i64>();
let expected = create_i32_large_list(
vec![3, 4, 5],
OffsetBuffer::<i64>::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::<i32>::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::<i32>();
let expected = create_i32_list(
vec![50, 50],
OffsetBuffer::<i32>::from_lengths(vec![1, 1, 0]),
None,
);
assert_eq!(actual, &expected);
}
}
Loading
Loading