Dandandan commented on code in PR #21101:
URL: https://github.com/apache/datafusion/pull/21101#discussion_r2973028445
##########
datafusion/functions-nested/src/min_max.rs:
##########
@@ -198,20 +202,120 @@ impl ScalarUDFImpl for ArrayMin {
fn array_min_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
let [array] = take_function_args("array_min", args)?;
match array.data_type() {
- List(_) => array_min_max_helper(as_list_array(array)?, min_batch),
- LargeList(_) => array_min_max_helper(as_large_list_array(array)?,
min_batch),
+ List(_) => array_min_max_helper(as_list_array(array)?, true),
+ LargeList(_) => array_min_max_helper(as_large_list_array(array)?,
true),
arg_type => exec_err!("array_min does not support type: {arg_type}"),
}
}
fn array_min_max_helper<O: OffsetSizeTrait>(
array: &GenericListArray<O>,
- agg_fn: fn(&ArrayRef) -> Result<ScalarValue>,
+ is_min: bool,
) -> Result<ArrayRef> {
+ // Try the primitive fast path first
+ if let Some(result) = try_primitive_array_min_max(array, is_min) {
+ return result;
+ }
+
+ // Fallback: per-row ScalarValue path for non-primitive types
+ let agg_fn = if is_min { min_batch } else { max_batch };
let null_value = ScalarValue::try_from(array.value_type())?;
let result_vec: Vec<ScalarValue> = array
.iter()
.map(|arr| arr.as_ref().map_or_else(|| Ok(null_value.clone()), agg_fn))
.try_collect()?;
ScalarValue::iter_to_array(result_vec)
}
+
+/// Dispatches to a typed primitive min/max implementation, or returns `None`
if
+/// the element type is not a primitive.
+fn try_primitive_array_min_max<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ is_min: bool,
+) -> Option<Result<ArrayRef>> {
+ macro_rules! helper {
+ ($t:ty) => {
+ return Some(primitive_array_min_max::<O, $t>(list_array, is_min))
+ };
+ }
+ downcast_primitive! {
+ list_array.value_type() => (helper),
+ _ => {}
+ }
+ None
+}
+
+/// Threshold to switch from direct iteration to using `min` / `max` kernel
from
+/// `arrow::compute`. The latter has enough per-invocation overhead that direct
+/// iteration is faster for small lists.
+const ARROW_COMPUTE_THRESHOLD: usize = 32;
+
+/// Computes min or max for each row of a primitive ListArray.
+fn primitive_array_min_max<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
+ list_array: &GenericListArray<O>,
+ is_min: bool,
+) -> Result<ArrayRef> {
+ let values_array = list_array.values().as_primitive::<T>();
+ let values_slice = values_array.values();
+ let values_nulls = values_array.nulls();
+ let mut result_builder =
PrimitiveBuilder::<T>::with_capacity(list_array.len())
+ .with_data_type(values_array.data_type().clone());
+
+ for (row, w) in list_array.offsets().windows(2).enumerate() {
+ let row_result = if list_array.is_null(row) {
+ None
+ } else {
+ let start = w[0].as_usize();
+ let end = w[1].as_usize();
+ let len = end - start;
+
+ match len {
+ 0 => None,
+ _ if len < ARROW_COMPUTE_THRESHOLD => {
+ scalar_min_max::<T>(values_slice, values_nulls, start,
end, is_min)
+ }
+ _ => {
+ let slice = values_array.slice(start, len);
+ if is_min {
+ arrow::compute::min::<T>(&slice)
+ } else {
+ arrow::compute::max::<T>(&slice)
+ }
+ }
+ }
+ };
+
+ result_builder.append_option(row_result);
+ }
+
+ Ok(Arc::new(result_builder.finish()) as ArrayRef)
+}
+
+/// Computes min or max for a single list row by directly scanning a slice of
+/// the flat values buffer.
+#[inline]
+fn scalar_min_max<T: ArrowPrimitiveType>(
+ values_slice: &[T::Native],
+ values_nulls: Option<&arrow::buffer::NullBuffer>,
+ start: usize,
+ end: usize,
+ is_min: bool,
+) -> Option<T::Native> {
+ let mut best: Option<T::Native> = None;
+ for (i, &val) in values_slice[start..end].iter().enumerate() {
Review Comment:
I think for this small loop it makes sense to optimize for non-null case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
