From 6087918323c9f90080b7600943d7bc47f222d4bc Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 20 Jul 2026 07:13:00 -0700 Subject: [PATCH 1/3] feat: add list min and max expressions Signed-off-by: Matt Katz --- vortex-array/Cargo.toml | 4 + vortex-array/benches/list_min_max.rs | 132 ++++ .../src/aggregate_fn/fns/min_max/grouped.rs | 649 ++++++++++++++++++ .../src/aggregate_fn/fns/min_max/mod.rs | 2 + vortex-array/src/aggregate_fn/session.rs | 11 + vortex-array/src/expr/exprs.rs | 32 + .../src/scalar_fn/fns/list_min_max/mod.rs | 241 +++++++ .../src/scalar_fn/fns/list_min_max/tests.rs | 513 ++++++++++++++ vortex-array/src/scalar_fn/fns/mod.rs | 1 + vortex-array/src/scalar_fn/session.rs | 4 + 10 files changed, 1589 insertions(+) create mode 100644 vortex-array/benches/list_min_max.rs create mode 100644 vortex-array/src/aggregate_fn/fns/min_max/grouped.rs create mode 100644 vortex-array/src/scalar_fn/fns/list_min_max/mod.rs create mode 100644 vortex-array/src/scalar_fn/fns/list_min_max/tests.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index a1aa5b047f8..5ddddc34d74 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -255,6 +255,10 @@ harness = false name = "list_sum" harness = false +[[bench]] +name = "list_min_max" +harness = false + [[bench]] name = "listview_rebuild" harness = false diff --git a/vortex-array/benches/list_min_max.rs b/vortex-array/benches/list_min_max.rs new file mode 100644 index 00000000000..06c28253bda --- /dev/null +++ b/vortex-array/benches/list_min_max.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for primitive `list_min` and `list_max` over `List`, `ListView`, and +//! `FixedSizeList` inputs. + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::expr::list_max; +use vortex_array::expr::list_min; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const BENCH_ARGS: &[(usize, usize)] = &[(8_192, 10), (8_192, 100), (8_192, 1_000)]; + +fn elements(num_lists: usize, list_size: usize) -> ArrayRef { + PrimitiveArray::from_option_iter( + (0..num_lists * list_size) + .map(|index| (index % 10 != 0).then_some(((index * 31) % 10_007) as i32)), + ) + .into_array() +} + +fn list_validity(num_lists: usize) -> Validity { + Validity::Array(BoolArray::from_iter((0..num_lists).map(|index| index % 10 != 0)).into_array()) +} + +fn make_list(num_lists: usize, list_size: usize) -> ArrayRef { + let offsets: Buffer = (0..=num_lists) + .map(|index| (index * list_size) as i32) + .collect(); + ListArray::try_new( + elements(num_lists, list_size), + offsets.into_array(), + list_validity(num_lists), + ) + .unwrap() + .into_array() +} + +fn make_listview(num_lists: usize, list_size: usize) -> ArrayRef { + let offsets: Buffer = (0..num_lists) + .map(|index| (((index * 17) % num_lists) * (list_size - 1)) as i32) + .collect(); + let sizes: Buffer = std::iter::repeat_n(list_size as i32, num_lists).collect(); + ListViewArray::new( + elements(num_lists, list_size), + offsets.into_array(), + sizes.into_array(), + list_validity(num_lists), + ) + .into_array() +} + +fn make_fixed_size_list(num_lists: usize, list_size: usize) -> ArrayRef { + FixedSizeListArray::new( + elements(num_lists, list_size), + list_size as u32, + list_validity(num_lists), + num_lists, + ) + .into_array() +} + +fn run_vortex(bencher: Bencher, array: ArrayRef, minimum: bool) { + let expression = if minimum { + list_min(root()) + } else { + list_max(root()) + }; + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + array + .clone() + .apply(&expression) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench(args = BENCH_ARGS)] +fn vortex_list_min(bencher: Bencher, (num_lists, list_size): (usize, usize)) { + run_vortex(bencher, make_list(num_lists, list_size), true); +} + +#[divan::bench(args = BENCH_ARGS)] +fn vortex_list_max(bencher: Bencher, (num_lists, list_size): (usize, usize)) { + run_vortex(bencher, make_list(num_lists, list_size), false); +} + +#[divan::bench(args = BENCH_ARGS)] +fn vortex_listview_min(bencher: Bencher, (num_lists, list_size): (usize, usize)) { + run_vortex(bencher, make_listview(num_lists, list_size), true); +} + +#[divan::bench(args = BENCH_ARGS)] +fn vortex_listview_max(bencher: Bencher, (num_lists, list_size): (usize, usize)) { + run_vortex(bencher, make_listview(num_lists, list_size), false); +} + +#[divan::bench(args = BENCH_ARGS)] +fn vortex_fixed_size_list_min(bencher: Bencher, (num_lists, list_size): (usize, usize)) { + run_vortex(bencher, make_fixed_size_list(num_lists, list_size), true); +} + +#[divan::bench(args = BENCH_ARGS)] +fn vortex_fixed_size_list_max(bencher: Bencher, (num_lists, list_size): (usize, usize)) { + run_vortex(bencher, make_fixed_size_list(num_lists, list_size), false); +} diff --git a/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs b/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs new file mode 100644 index 00000000000..b1c407eb81c --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs @@ -0,0 +1,649 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBufferMut; +use vortex_buffer::BitBufferView; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::GroupRanges; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::dtype::NativePType; +use crate::dtype::half::f16; +use crate::match_each_native_ptype; +use crate::validity::Validity; + +/// Encoding-specific grouped [`Min`] and [`Max`] kernel for primitive element arrays. +/// +/// Each outer list is one group over a shared primitive element buffer. The kernel consumes the +/// outer ranges and both validity levels in bulk, avoiding a per-list array slice and aggregate +/// state. Other element encodings may execute into [`Primitive`] before reaching this kernel. +#[derive(Debug)] +pub(crate) struct PrimitiveGroupedExtremaEncodingKernel; + +impl DynGroupedAggregateKernel for PrimitiveGroupedExtremaEncodingKernel { + fn grouped_aggregate( + &self, + aggregate_fn: &AggregateFnRef, + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if let Some(options) = aggregate_fn.as_opt::() { + return try_grouped_extreme(groups, ctx, options.skip_nans, Extreme::Min); + } + if let Some(options) = aggregate_fn.as_opt::() { + return try_grouped_extreme(groups, ctx, options.skip_nans, Extreme::Max); + } + Ok(None) + } +} + +#[derive(Clone, Copy)] +enum Extreme { + Min, + Max, +} + +/// Type-specific identities and combine operations used by the lane reducer. +/// +/// Starting every lane at an identity eliminates a seed search. A separate participation flag is +/// still required because the identity may also be a legitimate result. Keeping `CAN_NAN` and the +/// combine operation type-specific lets monomorphization remove float-only work for integers and +/// expose native min/max instructions to the optimizer. +trait ExtremaIdentity: NativePType { + const CAN_NAN: bool; + const MIN_IDENTITY: Self; + const MAX_IDENTITY: Self; + + fn minimum(candidate: Self, current: Self) -> Self; + fn maximum(candidate: Self, current: Self) -> Self; +} + +macro_rules! impl_integer_extrema_identity { + ($($t:ty),* $(,)?) => { + $( + impl ExtremaIdentity for $t { + const CAN_NAN: bool = false; + const MIN_IDENTITY: Self = Self::MAX; + const MAX_IDENTITY: Self = Self::MIN; + + #[inline(always)] + fn minimum(candidate: Self, current: Self) -> Self { + candidate.min(current) + } + + #[inline(always)] + fn maximum(candidate: Self, current: Self) -> Self { + candidate.max(current) + } + } + )* + }; +} + +macro_rules! impl_float_extrema_identity { + ($($t:ty),* $(,)?) => { + $( + impl ExtremaIdentity for $t { + const CAN_NAN: bool = true; + const MIN_IDENTITY: Self = Self::INFINITY; + const MAX_IDENTITY: Self = Self::NEG_INFINITY; + + #[inline(always)] + fn minimum(candidate: Self, current: Self) -> Self { + select(candidate.is_lt(current), candidate, current) + } + + #[inline(always)] + fn maximum(candidate: Self, current: Self) -> Self { + select(candidate.is_gt(current), candidate, current) + } + } + )* + }; +} + +impl_integer_extrema_identity!(u8, u16, u32, u64, i8, i16, i32, i64); +impl_float_extrema_identity!(f16, f32, f64); + +/// Fully specialized operations shared by the scalar and lane reducers for one invocation. +#[derive(Clone, Copy)] +struct ExtremaOps { + skip_nans: bool, + identity: T, + is_better: B, + combine: C, +} + +/// Element validity materialized once for the complete shared element array. +/// +/// The `Some` variant borrows the packed bitmap. Per-group reducers take zero-copy views or load a +/// word directly rather than constructing a new [`Mask`] and its cached valid-run representation. +#[derive(Clone, Copy)] +enum ElementValidity<'a> { + All, + None, + Some(BitBufferView<'a>), +} + +/// Run the encoding kernel when the shared elements are primitive. +/// +/// Returning `None` declines the encoding-specific dispatch and allows the grouped accumulator to +/// execute the elements further or use its generic per-group fallback. +fn try_grouped_extreme( + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + skip_nans: bool, + extreme: Extreme, +) -> VortexResult> { + if !groups.elements().is::() { + return Ok(None); + } + let elements = groups.elements().clone().downcast::(); + let group_ranges = groups.group_ranges(ctx)?; + let group_validity = groups.group_validity(ctx)?; + Ok(Some(grouped_extreme( + &elements, + &group_ranges, + &group_validity, + ctx, + skip_nans, + extreme, + )?)) +} + +/// Materialize element validity once, select the native physical type, and collect one result per +/// group. +/// +/// Type dispatch happens outside the group loop, so the inner reducer is monomorphized for both the +/// physical type and the requested extremum. +fn grouped_extreme( + elements: &PrimitiveArray, + group_ranges: &GroupRanges, + group_validity: &Mask, + ctx: &mut ExecutionCtx, + skip_nans: bool, + extreme: Extreme, +) -> VortexResult { + let element_validity = elements + .as_ref() + .validity()? + .execute_mask(elements.as_ref().len(), ctx)?; + + let result = match_each_native_ptype!(elements.ptype(), |T| { + let values = elements.as_slice::(); + match extreme { + Extreme::Min => collect_extrema( + values, + group_ranges, + group_validity, + &element_validity, + ExtremaOps { + skip_nans, + identity: T::MIN_IDENTITY, + is_better: |candidate: T, current| candidate.is_lt(current), + combine: ::minimum, + }, + ), + Extreme::Max => collect_extrema( + values, + group_ranges, + group_validity, + &element_validity, + ExtremaOps { + skip_nans, + identity: T::MAX_IDENTITY, + is_better: |candidate: T, current| candidate.is_gt(current), + combine: ::maximum, + }, + ), + } + }); + Ok(result.into_array()) +} + +/// Reduce every valid group directly into a final-length output buffer. +/// +/// Invalid outer lists are already represented by `group_validity`. `missing` records only valid +/// groups that have no participating element: empty lists, all-null lists, or float lists containing +/// only skipped NaNs. +fn collect_extrema( + values: &[T], + group_ranges: &GroupRanges, + group_validity: &Mask, + element_validity: &Mask, + ops: ExtremaOps, +) -> PrimitiveArray +where + T: ExtremaIdentity, + B: Fn(T, T) -> bool + Copy, + C: Fn(T, T) -> T + Copy, +{ + let element_validity = match element_validity.bit_buffer() { + AllOr::All => ElementValidity::All, + AllOr::None => ElementValidity::None, + AllOr::Some(validity) => ElementValidity::Some(validity.as_view()), + }; + let mut extrema = BufferMut::::zeroed(group_ranges.len()); + let mut missing = Vec::new(); + + for (index, ((offset, size), group_is_valid)) in + group_ranges.iter().zip(group_validity.iter()).enumerate() + { + if !group_is_valid { + continue; + } + match reduce_group(values, offset, size, element_validity, ops) { + Some(extreme) => extrema.as_mut_slice()[index] = extreme, + None => missing.push(index), + } + } + + PrimitiveArray::new(extrema.freeze(), output_validity(group_validity, &missing)) +} + +/// Combine outer-list validity with valid groups that did not produce an extremum. +/// +/// The common case reuses the original bitmap. A mutable copy is allocated only when at least one +/// otherwise-valid group is empty, all-null, or all-skipped-NaN. +fn output_validity(group_validity: &Mask, missing: &[usize]) -> Validity { + if missing.is_empty() { + return match group_validity.bit_buffer() { + AllOr::All => Validity::AllValid, + AllOr::None => Validity::AllInvalid, + AllOr::Some(validity) => Validity::from(validity.clone()), + }; + } + + let mut validity = match group_validity.bit_buffer() { + AllOr::All => BitBufferMut::new_set(group_validity.len()), + AllOr::None => BitBufferMut::new_unset(group_validity.len()), + AllOr::Some(validity) => BitBufferMut::copy_from(validity), + }; + for &index in missing { + validity.unset(index); + } + Validity::from(validity.freeze()) +} + +/// Reduce one `(offset, size)` range using a size-adaptive strategy. +/// +/// Groups shorter than 32 values use the scalar bitmap-word path. Groups from 32 through 511 use +/// one 16-byte native vector's worth of independent accumulators. Larger groups use four vectors' +/// worth to shorten the loop-carried dependency chain and amortize validity-mask handling. +/// +/// The range is independent of neighboring groups, which is required for overlapping and +/// out-of-order ListView values. +fn reduce_group( + values: &[T], + offset: usize, + size: usize, + element_validity: ElementValidity<'_>, + ops: ExtremaOps, +) -> Option +where + T: ExtremaIdentity, + B: Fn(T, T) -> bool + Copy, + C: Fn(T, T) -> T + Copy, +{ + if size < 32 { + return reduce_group_scalar( + values, + offset, + size, + element_validity, + ops.skip_nans, + ops.is_better, + ); + } + + // Long groups use four native SIMD vectors of independent accumulators to shorten the + // loop-carried dependency. Smaller groups use one vector to limit setup and reduction costs. + if size >= 512 { + return match size_of::() { + 1 => reduce_group_lanes::(values, offset, size, element_validity, ops), + 2 => reduce_group_lanes::(values, offset, size, element_validity, ops), + 4 => reduce_group_lanes::(values, offset, size, element_validity, ops), + 8 => reduce_group_lanes::(values, offset, size, element_validity, ops), + _ => reduce_group_lanes::(values, offset, size, element_validity, ops), + }; + } + + match size_of::() { + 1 => reduce_group_lanes::(values, offset, size, element_validity, ops), + 2 => reduce_group_lanes::(values, offset, size, element_validity, ops), + 4 => reduce_group_lanes::(values, offset, size, element_validity, ops), + 8 => reduce_group_lanes::(values, offset, size, element_validity, ops), + _ => reduce_group_lanes::(values, offset, size, element_validity, ops), + } +} + +/// Reduce a group through `LANES` independent accumulators and merge them into one result. +/// +/// `LANES` is selected to represent either 16 or 64 bytes of values. The element bitmap is sliced +/// as a borrowed [`BitBufferView`]; all-valid groups synthesize validity words without allocating a +/// bitmap. `found` distinguishes an empty reduction from a real result equal to `ops.identity`. +fn reduce_group_lanes( + values: &[T], + offset: usize, + size: usize, + element_validity: ElementValidity<'_>, + ops: ExtremaOps, +) -> Option +where + T: ExtremaIdentity, + B: Fn(T, T) -> bool + Copy, + C: Fn(T, T) -> T + Copy, +{ + debug_assert!(LANES > 0 && LANES.is_power_of_two() && 64 % LANES == 0); + let values = &values[offset..offset + size]; + let validity = match element_validity { + ElementValidity::All => None, + ElementValidity::None => return None, + ElementValidity::Some(validity) => Some(validity.slice(offset..offset + size)), + }; + let mut accumulators = [ops.identity; LANES]; + let mut found = false; + let poisoned = match validity { + Some(validity) => reduce_validity_words( + &mut accumulators, + values, + validity.chunks().iter_padded(), + ops.skip_nans, + &ops.combine, + &mut found, + ), + None => reduce_validity_words( + &mut accumulators, + values, + std::iter::repeat_n(u64::MAX, values.len().div_ceil(64)), + ops.skip_nans, + &ops.combine, + &mut found, + ), + }; + if let Some(nan) = poisoned { + return Some(nan); + } + + found.then(|| merge_accumulators(accumulators, &ops.is_better)) +} + +/// Reduce packed validity words without materializing per-group masks or valid-run vectors. +/// +/// Each input word describes at most 64 consecutive values. Tail bits are explicitly masked, so a +/// padded final word cannot make an empty lane appear to participate. A returned value is a NaN that +/// poisoned the group under `include_nans`; ordinary completion returns `None` and leaves the lane +/// results in `accumulators`. +fn reduce_validity_words( + accumulators: &mut [T; LANES], + values: &[T], + validity_words: impl Iterator, + skip_nans: bool, + combine: &impl Fn(T, T) -> T, + found: &mut bool, +) -> Option { + let mut base = 0; + for word in validity_words { + let word_len = values.len().saturating_sub(base).min(64); + let valid_mask = if word_len == 64 { + u64::MAX + } else { + (1u64 << word_len) - 1 + }; + if !T::CAN_NAN || !skip_nans { + *found |= word & valid_mask != 0; + } + let word_values = &values[base..base + word_len]; + let (chunks, remainder) = word_values.as_chunks::(); + let mut validity = word; + for lane_values in chunks { + if let Some(nan) = reduce_lane_chunk( + accumulators, + lane_values, + validity, + skip_nans, + combine, + found, + ) { + return Some(nan); + } + if LANES == 64 { + validity = 0; + } else { + validity >>= LANES; + } + } + + if !remainder.is_empty() + && let Some(nan) = + reduce_lane_remainder(accumulators, remainder, validity, skip_nans, combine, found) + { + return Some(nan); + } + base += word_len; + } + None +} + +/// Reduce a short group without constructing a bitmap slice or lane accumulators. +/// +/// Nullable groups load all relevant validity bits into one word. This path is restricted to fewer +/// than 32 values, which keeps the direct unaligned bitmap load simple even when the view begins at +/// a non-byte-aligned offset. +fn reduce_group_scalar( + values: &[T], + offset: usize, + size: usize, + element_validity: ElementValidity<'_>, + skip_nans: bool, + is_better: impl Fn(T, T) -> bool, +) -> Option { + let mut best = None; + let group_values = &values[offset..offset + size]; + match element_validity { + ElementValidity::All => { + reduce_scalar_run(&mut best, group_values, skip_nans, &is_better); + } + ElementValidity::None => {} + ElementValidity::Some(validity) => { + reduce_nullable_word_scalar( + &mut best, + group_values, + validity_word(validity, offset, size), + skip_nans, + &is_better, + ); + } + } + best +} + +/// Reduce the contiguous set-bit runs in one short group's validity word. +/// +/// The all-valid fast path scans the entire value slice. Otherwise, zero and one bit counts locate +/// valid runs without a per-element validity lookup. +fn reduce_nullable_word_scalar( + best: &mut Option, + values: &[T], + mut word: u64, + skip_nans: bool, + is_better: &impl Fn(T, T) -> bool, +) -> bool { + let valid_mask = if values.is_empty() { + 0 + } else { + (1u64 << values.len()) - 1 + }; + word &= valid_mask; + + if word == valid_mask { + return reduce_scalar_run(best, values, skip_nans, is_better); + } + + while word != 0 { + let start = word.trailing_zeros() as usize; + let run_len = (word >> start).trailing_ones() as usize; + if reduce_scalar_run(best, &values[start..start + run_len], skip_nans, is_better) { + return true; + } + word &= !(((1u64 << run_len) - 1) << start); + } + false +} + +#[inline] +/// Load and align the validity bits for a short group from the shared bitmap. +/// +/// The checked eight-byte chunk load compiles to an inline word load in the common case. The fold is +/// used only near the end of the backing buffer where eight bytes are not available. +fn validity_word(validity: BitBufferView<'_>, offset: usize, len: usize) -> u64 { + debug_assert!(len < 32); + debug_assert!(offset + len <= validity.len()); + if len == 0 { + return 0; + } + + let bit_offset = validity.offset() + offset; + let buffer = &validity.inner()[bit_offset / 8..]; + let word = buffer.first_chunk::<8>().map_or_else( + || { + buffer.iter().enumerate().fold(0, |word, (index, &byte)| { + word | u64::from(byte) << (index * 8) + }) + }, + |bytes| u64::from_le_bytes(*bytes), + ); + (word >> (bit_offset % 8)) & ((1u64 << len) - 1) +} + +/// Fold a contiguous run of valid values into `best`. +/// +/// Returns `true` only when an included NaN poisons the group, allowing callers to stop scanning +/// that group immediately. Skipped NaNs do not initialize or update `best`. +fn reduce_scalar_run( + best: &mut Option, + values: &[T], + skip_nans: bool, + is_better: &impl Fn(T, T) -> bool, +) -> bool { + for &candidate in values { + if candidate.is_nan() { + if skip_nans { + continue; + } + *best = Some(candidate); + return true; + } + match best { + Some(current) if is_better(candidate, *current) => *current = candidate, + None => *best = Some(candidate), + _ => {} + } + } + false +} + +#[inline(always)] +/// Combine one full lane chunk with its packed validity bits. +/// +/// The branchless combine-and-select form is intentional: after monomorphization, LLVM can lower +/// integer min/max plus validity selection to packed SIMD instructions. Float-only NaN logic is +/// removed entirely for integer types through [`ExtremaIdentity::CAN_NAN`]. +fn reduce_lane_chunk( + accumulators: &mut [T; LANES], + values: &[T; LANES], + validity: u64, + skip_nans: bool, + combine: &impl Fn(T, T) -> T, + found: &mut bool, +) -> Option { + for lane in 0..LANES { + let candidate = values[lane]; + let valid = validity & (1 << lane) != 0; + let current = accumulators[lane]; + if T::CAN_NAN { + let is_nan = candidate.is_nan(); + if skip_nans { + *found |= valid & !is_nan; + } + if valid & is_nan & !skip_nans { + return Some(candidate); + } + let reduced = combine(candidate, current); + accumulators[lane] = select(valid & !is_nan, reduced, current); + } else { + let reduced = combine(candidate, current); + accumulators[lane] = select(valid, reduced, current); + } + } + None +} + +/// Apply the same lane semantics to the final partial chunk of a validity word. +fn reduce_lane_remainder( + accumulators: &mut [T; LANES], + values: &[T], + validity: u64, + skip_nans: bool, + combine: &impl Fn(T, T) -> T, + found: &mut bool, +) -> Option { + for (lane, &candidate) in values.iter().enumerate() { + let valid = validity & (1 << lane) != 0; + let current = accumulators[lane]; + if T::CAN_NAN { + let is_nan = candidate.is_nan(); + if skip_nans { + *found |= valid & !is_nan; + } + if valid & is_nan & !skip_nans { + return Some(candidate); + } + let reduced = combine(candidate, current); + accumulators[lane] = select(valid & !is_nan, reduced, current); + } else { + let reduced = combine(candidate, current); + accumulators[lane] = select(valid, reduced, current); + } + } + None +} + +/// Tree-reduce independent lane accumulators into the group's scalar extremum. +/// +/// A tree keeps the dependency depth logarithmic in `LANES`, which matters most for the four-vector +/// path used by long groups. +fn merge_accumulators( + mut accumulators: [T; LANES], + is_better: &impl Fn(T, T) -> bool, +) -> T { + let mut len = LANES; + while len >= 2 { + let mid = len / 2; + for lane in 0..mid { + let candidate = accumulators[lane + mid]; + let current = accumulators[lane]; + accumulators[lane] = select(is_better(candidate, current), candidate, current); + } + len /= 2; + } + accumulators[0] +} + +#[inline(always)] +/// Select a value in a form that the lane loop can lower to a SIMD mask operation. +fn select(condition: bool, if_true: T, if_false: T) -> T { + if condition { if_true } else { if_false } +} diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 51601ce76fd..e6593241dfa 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -4,11 +4,13 @@ mod bool; mod decimal; mod extension; +mod grouped; mod primitive; mod varbin; use std::sync::LazyLock; +pub(crate) use grouped::PrimitiveGroupedExtremaEncodingKernel; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 72ef78a400c..ba86fc3bd74 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -28,6 +28,7 @@ use crate::aggregate_fn::fns::last::Last; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; +use crate::aggregate_fn::fns::min_max::PrimitiveGroupedExtremaEncodingKernel; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; @@ -117,6 +118,16 @@ impl Default for AggregateFnSession { Sum.id(), &PrimitiveGroupedSumEncodingKernel, ); + this.register_grouped_encoding_kernel( + Primitive.id(), + Min.id(), + &PrimitiveGroupedExtremaEncodingKernel, + ); + this.register_grouped_encoding_kernel( + Primitive.id(), + Max.id(), + &PrimitiveGroupedExtremaEncodingKernel, + ); this } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index da70427a0f7..42fb7ed3772 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -39,6 +39,8 @@ use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_length::ListLength; +use crate::scalar_fn::fns::list_min_max::ListMax; +use crate::scalar_fn::fns::list_min_max::ListMin; use crate::scalar_fn::fns::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; @@ -780,6 +782,36 @@ pub fn list_length(input: Expression) -> Expression { ListLength.new_expr(EmptyOptions, [input]) } +// ---- ListMin / ListMax ---- + +/// Creates an expression that returns the maximum participating element in each list for `List` and +/// `FixedSizeList` inputs. +/// +/// Null lists, empty lists, and lists whose elements are all null return null. Null elements and, +/// by default, float NaNs are ignored. See [`list_min_opts`] to include NaNs instead. +pub fn list_min(input: Expression) -> Expression { + ListMin.new_expr(NumericalAggregateOpts::default(), [input]) +} + +/// Creates a [`list_min`] expression with explicit NaN handling. +pub fn list_min_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { + ListMin.new_expr(options, [input]) +} + +/// Creates an expression that returns the minimum participating element in each list for `List` and +/// `FixedSizeList` inputs. +/// +/// Null lists, empty lists, and lists whose elements are all null return null. Null elements and, +/// by default, float NaNs are ignored. See [`list_max_opts`] to include NaNs instead. +pub fn list_max(input: Expression) -> Expression { + ListMax.new_expr(NumericalAggregateOpts::default(), [input]) +} + +/// Creates a [`list_max`] expression with explicit NaN handling. +pub fn list_max_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { + ListMax.new_expr(options, [input]) +} + // ---- ListSum ---- /// Creates an expression that sums the elements of each list for `List` and diff --git a/vortex-array/src/scalar_fn/fns/list_min_max/mod.rs b/vortex-array/src/scalar_fn/fns/list_min_max/mod.rs new file mode 100644 index 00000000000..2a121251736 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_min_max/mod.rs @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; + +/// Minimum of the non-null elements in each `List`, `ListView`, or `FixedSizeList` value. +/// +/// Null lists, empty lists, and lists without a participating element produce null. Null elements +/// are ignored. With the default [`NumericalAggregateOpts`], float NaNs are also ignored; with +/// [`NumericalAggregateOpts::include_nans`], any NaN poisons its list's result to NaN. +#[derive(Clone)] +pub struct ListMin; + +/// Maximum of the non-null elements in each `List`, `ListView`, or `FixedSizeList` value. +/// +/// Null lists, empty lists, and lists without a participating element produce null. Null elements +/// are ignored. With the default [`NumericalAggregateOpts`], float NaNs are also ignored; with +/// [`NumericalAggregateOpts::include_nans`], any NaN poisons its list's result to NaN. +#[derive(Clone)] +pub struct ListMax; + +impl ScalarFnVTable for ListMin { + type Options = NumericalAggregateOpts; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.min"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for list_min()"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + list_extreme_return_dtype("list_min", &arg_dtypes[0]) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + execute_list_extreme(Min, "list_min", options, args, ctx) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +impl ScalarFnVTable for ListMax { + type Options = NumericalAggregateOpts; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.max"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for list_max()"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + list_extreme_return_dtype("list_max", &arg_dtypes[0]) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + execute_list_extreme(Max, "list_max", options, args, ctx) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// Validate the list and element dtypes and derive the nullable scalar output dtype. +/// +/// The result is nullable even when both input levels are non-nullable because an empty list has no +/// minimum or maximum. +fn list_extreme_return_dtype(function: &str, input_dtype: &DType) -> VortexResult { + let element_dtype = list_element_dtype(function, input_dtype)?; + Ok(element_dtype.as_nullable()) +} + +/// Return the comparable element dtype accepted by the list-extrema implementation. +/// +/// Nested and extension values are rejected until their aggregate ordering and nested-null +/// semantics are explicitly defined. In particular, an extension's storage order is not +/// necessarily its logical order. +fn list_element_dtype(function: &str, input_dtype: &DType) -> VortexResult { + let element_dtype = match input_dtype { + DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { + element_dtype.as_ref() + } + other => vortex_bail!("{function}() requires List or FixedSizeList, got {other}"), + }; + + if !matches!( + element_dtype, + DType::Bool(_) + | DType::Primitive(..) + | DType::Decimal(..) + | DType::Utf8(_) + | DType::Binary(_) + ) { + vortex_bail!("{function}() cannot compare elements of type {element_dtype}") + } + + Ok(element_dtype.clone()) +} + +/// Execute a list extremum while preserving the constant-array fast path. +/// +/// Non-constant inputs are executed to a columnar representation and delegated to the grouped +/// aggregate machinery. A constant list is reduced once and its scalar result is broadcast to the +/// original row count, avoiding repeated work for both its elements and validity. +fn execute_list_extreme>( + aggregate: V, + function: &str, + options: &NumericalAggregateOpts, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let input = args.get(0)?; + let element_dtype = list_element_dtype(function, input.dtype())?; + + match input.execute::(ctx)? { + Columnar::Constant(constant) => { + // Evaluate a single list and broadcast its extremum rather than expanding all of the + // constant's repeated elements. + let one_row = ConstantArray::new(constant.scalar().clone(), 1) + .into_array() + .execute::(ctx)? + .into_array(); + let extreme = list_extreme_impl(aggregate, one_row, element_dtype, options, ctx)? + .execute_scalar(0, ctx)?; + Ok(ConstantArray::new(extreme, constant.len()).into_array()) + } + Columnar::Canonical(canonical) => list_extreme_impl( + aggregate, + canonical.into_array(), + element_dtype, + options, + ctx, + ), + } +} + +/// Treat every outer list value as one group and run the supplied `Min` or `Max` aggregate. +/// +/// The grouped accumulator preserves explicit ListView ranges, so repeated, overlapping, gapped, +/// and out-of-order views are reduced according to their logical rows rather than their physical +/// element order. +fn list_extreme_impl>( + aggregate: V, + canonical: ArrayRef, + element_dtype: DType, + options: &NumericalAggregateOpts, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut accumulator = GroupedAccumulator::try_new(aggregate, *options, element_dtype)?; + accumulator.accumulate_list(&canonical, ctx)?; + accumulator.finish() +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/list_min_max/tests.rs b/vortex-array/src/scalar_fn/fns/list_min_max/tests.rs new file mode 100644 index 00000000000..0382fb37a0e --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_min_max/tests.rs @@ -0,0 +1,513 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![allow(clippy::cast_possible_truncation, clippy::cognitive_complexity)] + +use std::sync::Arc; + +use prost::Message; +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_proto::expr as pb; + +use super::ListMax; +use super::ListMin; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::half::f16; +use crate::expr::Expression; +use crate::expr::list_max; +use crate::expr::list_max_opts; +use crate::expr::list_min; +use crate::expr::list_min_opts; +use crate::expr::proto::ExprSerializeProtoExt; +use crate::expr::root; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; +use crate::validity::Validity; + +fn create_list_elements() -> ArrayRef { + PrimitiveArray::from_option_iter::([ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + None, + ]) + .into_array() +} + +fn assert_primitive_extrema(list: &ArrayRef) -> VortexResult<()> { + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.clone().apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = + PrimitiveArray::from_option_iter::([Some(1), Some(3), None, Some(6)]); + let expected_max = + PrimitiveArray::from_option_iter::([Some(2), Some(5), None, Some(6)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +#[rstest] +#[case(buffer![0u32, 2, 5, 5, 7].into_array())] +#[case(buffer![0u64, 2, 5, 5, 7].into_array())] +fn list_extrema(#[case] offsets: ArrayRef) -> VortexResult<()> { + let list = + ListArray::try_new(create_list_elements(), offsets, Validity::NonNullable)?.into_array(); + assert_primitive_extrema(&list) +} + +#[test] +fn nullable_lists_and_elements() -> VortexResult<()> { + let list = ListArray::try_new( + create_list_elements(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()), + )? + .into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected = PrimitiveArray::from_option_iter::([Some(1), None, None, None]); + assert_arrays_eq!(minimum, expected, &mut ctx); + let expected = PrimitiveArray::from_option_iter::([Some(2), None, None, None]); + assert_arrays_eq!(maximum, expected, &mut ctx); + Ok(()) +} + +#[test] +fn listview_extrema() -> VortexResult<()> { + let list_view = ListViewArray::new( + create_list_elements(), + buffer![5u32, 0, 4, 1].into_array(), + buffer![2u32, 3, 0, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let minimum = list_view.clone().apply(&list_min(root()))?; + let maximum = list_view.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + // The views are [6, null], [1, 2, 3], [], and [2, 3]. + let expected_min = + PrimitiveArray::from_option_iter::([Some(6), Some(1), None, Some(2)]); + let expected_max = + PrimitiveArray::from_option_iter::([Some(6), Some(3), None, Some(3)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +#[test] +fn listview_extrema_across_validity_words() -> VortexResult<()> { + let elements = PrimitiveArray::from_option_iter::( + (0usize..160).map(|index| (!index.is_multiple_of(7)).then_some(index as i32)), + ) + .into_array(); + let list_view = ListViewArray::new( + elements, + buffer![63u32, 1, 65, 126, 0].into_array(), + buffer![70u32, 130, 64, 8, 0].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, true, true]).into_array()), + ) + .into_array(); + let minimum = list_view.clone().apply(&list_min(root()))?; + let maximum = list_view.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = + PrimitiveArray::from_option_iter::([Some(64), None, Some(65), Some(127), None]); + let expected_max = + PrimitiveArray::from_option_iter::([Some(132), None, Some(128), Some(132), None]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +#[test] +fn primitive_lane_widths() -> VortexResult<()> { + macro_rules! assert_extrema { + ($t:ty) => {{ + let elements = PrimitiveArray::from_option_iter::<$t, _>( + (0usize..70).map(|index| (!index.is_multiple_of(7)).then_some(index as $t)), + ) + .into_array(); + let list = ListArray::try_new( + elements, + buffer![0u32, 70].into_array(), + Validity::NonNullable, + )? + .into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = PrimitiveArray::from_option_iter::<$t, _>([Some(1 as $t)]); + let expected_max = PrimitiveArray::from_option_iter::<$t, _>([Some(69 as $t)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + }}; + } + + assert_extrema!(u8); + assert_extrema!(u16); + assert_extrema!(u32); + assert_extrema!(u64); + assert_extrema!(i8); + assert_extrema!(i16); + assert_extrema!(i32); + assert_extrema!(i64); + assert_extrema!(f32); + assert_extrema!(f64); + Ok(()) +} + +#[test] +fn primitive_wide_lane_widths() -> VortexResult<()> { + macro_rules! assert_extrema { + ($t:ty) => {{ + let elements = PrimitiveArray::from_option_iter::<$t, _>( + (0usize..520) + .map(|index| (!index.is_multiple_of(7)).then_some((index % 100) as $t)), + ) + .into_array(); + let list = ListArray::try_new( + elements, + buffer![0u32, 520].into_array(), + Validity::NonNullable, + )? + .into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = PrimitiveArray::from_option_iter::<$t, _>([Some(0 as $t)]); + let expected_max = PrimitiveArray::from_option_iter::<$t, _>([Some(99 as $t)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + }}; + } + + assert_extrema!(u8); + assert_extrema!(u16); + assert_extrema!(u32); + assert_extrema!(u64); + assert_extrema!(i8); + assert_extrema!(i16); + assert_extrema!(i32); + assert_extrema!(i64); + assert_extrema!(f32); + assert_extrema!(f64); + Ok(()) +} + +#[test] +fn f16_lane_width() -> VortexResult<()> { + let elements = PrimitiveArray::from_option_iter::( + (0usize..70).map(|index| (!index.is_multiple_of(7)).then_some(f16::from_f32(index as f32))), + ) + .into_array(); + let list = ListArray::try_new( + elements, + buffer![0u32, 70].into_array(), + Validity::NonNullable, + )? + .into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = PrimitiveArray::from_option_iter::([Some(f16::from_f32(1.0))]); + let expected_max = PrimitiveArray::from_option_iter::([Some(f16::from_f32(69.0))]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +#[test] +fn f16_wide_lane_width() -> VortexResult<()> { + let elements = + PrimitiveArray::from_option_iter::((0usize..520).map(|index| { + (!index.is_multiple_of(7)).then_some(f16::from_f32((index % 100) as f32)) + })) + .into_array(); + let list = ListArray::try_new( + elements, + buffer![0u32, 520].into_array(), + Validity::NonNullable, + )? + .into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = PrimitiveArray::from_option_iter::([Some(f16::from_f32(0.0))]); + let expected_max = PrimitiveArray::from_option_iter::([Some(f16::from_f32(99.0))]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +fn create_fixed_size_list(validity: Validity) -> ArrayRef { + let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array(); + FixedSizeListArray::new(elements, 2, validity, 4).into_array() +} + +#[test] +fn fixed_size_list_extrema() -> VortexResult<()> { + let list = create_fixed_size_list(Validity::Array( + BoolArray::from_iter([true, false, true, true]).into_array(), + )); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = + PrimitiveArray::from_option_iter::([Some(1), None, Some(5), Some(7)]); + let expected_max = + PrimitiveArray::from_option_iter::([Some(2), None, Some(6), Some(8)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +#[test] +fn zero_width_fixed_size_lists_are_null() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([0i32; 0]).into_array(); + let list = FixedSizeListArray::try_new(elements, 0, Validity::NonNullable, 3)?.into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected = PrimitiveArray::from_option_iter::([None, None, None]); + assert_arrays_eq!(minimum, expected, &mut ctx); + assert_arrays_eq!(maximum, expected, &mut ctx); + Ok(()) +} + +#[test] +fn string_extrema() -> VortexResult<()> { + let elements = VarBinArray::from_iter( + [Some("pear"), None, Some("apple"), Some("zebra")], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let list = ListArray::try_new( + elements, + buffer![0u32, 3, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.apply(&list_max(root()))?; + let mut ctx = array_session().create_execution_ctx(); + + let expected_min = VarBinArray::from_iter( + [Some("apple"), Some("zebra")], + DType::Utf8(Nullability::Nullable), + ); + let expected_max = VarBinArray::from_iter( + [Some("pear"), Some("zebra")], + DType::Utf8(Nullability::Nullable), + ); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + Ok(()) +} + +#[test] +fn nan_options() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0, f64::NAN, f64::NAN]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3, 5].into_array(), + Validity::NonNullable, + )? + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.clone().apply(&list_max(root()))?; + let expected_min = PrimitiveArray::from_option_iter::([Some(1.0), None]); + let expected_max = PrimitiveArray::from_option_iter::([Some(2.0), None]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + + for result in [ + list.clone().apply(&list_min_opts( + root(), + NumericalAggregateOpts::include_nans(), + ))?, + list.apply(&list_max_opts( + root(), + NumericalAggregateOpts::include_nans(), + ))?, + ] { + let result = result.execute::(&mut ctx)?; + assert!(result.as_slice::().iter().all(|value| value.is_nan())); + } + Ok(()) +} + +#[test] +fn nan_options_lane_path() -> VortexResult<()> { + let elements = PrimitiveArray::from_option_iter::((0usize..520).map(|index| { + if index == 411 { + Some(f64::NAN) + } else { + (!index.is_multiple_of(11)).then_some(index as f64) + } + })); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 520].into_array(), + Validity::NonNullable, + )? + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let minimum = list.clone().apply(&list_min(root()))?; + let maximum = list.clone().apply(&list_max(root()))?; + let expected_min = PrimitiveArray::from_option_iter::([Some(1.0)]); + let expected_max = PrimitiveArray::from_option_iter::([Some(519.0)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + + for result in [ + list.clone().apply(&list_min_opts( + root(), + NumericalAggregateOpts::include_nans(), + ))?, + list.apply(&list_max_opts( + root(), + NumericalAggregateOpts::include_nans(), + ))?, + ] { + let result = result.execute::(&mut ctx)?; + assert!(result.as_slice::()[0].is_nan()); + } + Ok(()) +} + +#[test] +fn constant_and_null_list_extrema() -> VortexResult<()> { + let list = ListArray::try_new( + create_list_elements(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + let scalar = list.execute_scalar(1, &mut ctx)?; + let constant = ConstantArray::new(scalar, 3).into_array(); + + let minimum = constant.clone().apply(&list_min(root()))?; + let maximum = constant.apply(&list_max(root()))?; + let expected_min = PrimitiveArray::from_option_iter::([Some(3), Some(3), Some(3)]); + let expected_max = PrimitiveArray::from_option_iter::([Some(5), Some(5), Some(5)]); + assert_arrays_eq!(minimum, expected_min, &mut ctx); + assert_arrays_eq!(maximum, expected_max, &mut ctx); + + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::Nullable, + ); + let nulls = ConstantArray::new(Scalar::null(dtype), 2).into_array(); + for result in [ + nulls.clone().apply(&list_min(root()))?, + nulls.apply(&list_max(root()))?, + ] { + assert_eq!(result.valid_count(&mut ctx)?, 0); + } + Ok(()) +} + +#[test] +fn dtype_validation() -> VortexResult<()> { + let opts = NumericalAggregateOpts::default(); + let non_list = DType::Primitive(PType::I32, Nullability::NonNullable); + assert!( + ListMin + .return_dtype(&opts, std::slice::from_ref(&non_list)) + .is_err() + ); + assert!(ListMax.return_dtype(&opts, &[non_list]).is_err()); + + let nested_element = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ); + let nested_list = DType::List(Arc::new(nested_element), Nullability::NonNullable); + assert!( + ListMin + .return_dtype(&opts, std::slice::from_ref(&nested_list)) + .is_err() + ); + assert!(ListMax.return_dtype(&opts, &[nested_list]).is_err()); + + let utf8_list = DType::List( + Arc::new(DType::Utf8(Nullability::NonNullable)), + Nullability::NonNullable, + ); + assert_eq!( + ListMin.return_dtype(&opts, std::slice::from_ref(&utf8_list))?, + DType::Utf8(Nullability::Nullable) + ); + assert_eq!( + ListMax.return_dtype(&opts, &[utf8_list])?, + DType::Utf8(Nullability::Nullable) + ); + Ok(()) +} + +#[test] +fn display() { + assert_eq!(list_min(root()).to_string(), "vortex.list.min($)"); + assert_eq!(list_max(root()).to_string(), "vortex.list.max($)"); + assert_eq!( + list_min_opts(root(), NumericalAggregateOpts::include_nans()).to_string(), + "vortex.list.min($, opts=skip_nans=false)" + ); + assert_eq!( + list_max_opts(root(), NumericalAggregateOpts::include_nans()).to_string(), + "vortex.list.max($, opts=skip_nans=false)" + ); +} + +#[test] +fn proto_round_trip() -> VortexResult<()> { + for expr in [ + list_min(root()), + list_min_opts(root(), NumericalAggregateOpts::include_nans()), + list_max(root()), + list_max_opts(root(), NumericalAggregateOpts::include_nans()), + ] { + let proto = expr.serialize_proto()?; + let buf = proto.encode_to_vec(); + let decoded = pb::Expr::decode(buf.as_slice())?; + let deserialized = Expression::from_proto(&decoded, &array_session())?; + assert_eq!(expr, deserialized); + } + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index f4618154329..a4fce50b70b 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -15,6 +15,7 @@ pub mod is_null; pub mod like; pub mod list_contains; pub mod list_length; +pub mod list_min_max; pub mod list_sum; pub mod literal; pub mod mask; diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index e937ac8d316..a3611b674f8 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -22,6 +22,8 @@ use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_length::ListLength; +use crate::scalar_fn::fns::list_min_max::ListMax; +use crate::scalar_fn::fns::list_min_max::ListMin; use crate::scalar_fn::fns::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::merge::Merge; @@ -71,6 +73,8 @@ impl Default for ScalarFnSession { this.register(Like); this.register(ListContains); this.register(ListLength); + this.register(ListMax); + this.register(ListMin); this.register(ListSum); this.register(Literal); this.register(Merge); From 8ec5fd180eff84920ac9e342c8dbb56b65ad3acb Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 20 Jul 2026 07:55:38 -0700 Subject: [PATCH 2/3] refactor: reuse AllOr for list validity Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/min_max/grouped.rs | 73 +++++++++---------- 1 file changed, 34 insertions(+), 39 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs b/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs index b1c407eb81c..1ab68313355 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs @@ -102,12 +102,20 @@ macro_rules! impl_float_extrema_identity { #[inline(always)] fn minimum(candidate: Self, current: Self) -> Self { - select(candidate.is_lt(current), candidate, current) + if candidate.is_lt(current) { + candidate + } else { + current + } } #[inline(always)] fn maximum(candidate: Self, current: Self) -> Self { - select(candidate.is_gt(current), candidate, current) + if candidate.is_gt(current) { + candidate + } else { + current + } } } )* @@ -126,17 +134,6 @@ struct ExtremaOps { combine: C, } -/// Element validity materialized once for the complete shared element array. -/// -/// The `Some` variant borrows the packed bitmap. Per-group reducers take zero-copy views or load a -/// word directly rather than constructing a new [`Mask`] and its cached valid-run representation. -#[derive(Clone, Copy)] -enum ElementValidity<'a> { - All, - None, - Some(BitBufferView<'a>), -} - /// Run the encoding kernel when the shared elements are primitive. /// /// Returning `None` declines the encoding-specific dispatch and allows the grouped accumulator to @@ -231,9 +228,9 @@ where C: Fn(T, T) -> T + Copy, { let element_validity = match element_validity.bit_buffer() { - AllOr::All => ElementValidity::All, - AllOr::None => ElementValidity::None, - AllOr::Some(validity) => ElementValidity::Some(validity.as_view()), + AllOr::All => AllOr::All, + AllOr::None => AllOr::None, + AllOr::Some(validity) => AllOr::Some(validity.as_view()), }; let mut extrema = BufferMut::::zeroed(group_ranges.len()); let mut missing = Vec::new(); @@ -244,7 +241,7 @@ where if !group_is_valid { continue; } - match reduce_group(values, offset, size, element_validity, ops) { + match reduce_group(values, offset, size, &element_validity, ops) { Some(extreme) => extrema.as_mut_slice()[index] = extreme, None => missing.push(index), } @@ -289,7 +286,7 @@ fn reduce_group( values: &[T], offset: usize, size: usize, - element_validity: ElementValidity<'_>, + element_validity: &AllOr>, ops: ExtremaOps, ) -> Option where @@ -338,7 +335,7 @@ fn reduce_group_lanes( values: &[T], offset: usize, size: usize, - element_validity: ElementValidity<'_>, + element_validity: &AllOr>, ops: ExtremaOps, ) -> Option where @@ -349,9 +346,9 @@ where debug_assert!(LANES > 0 && LANES.is_power_of_two() && 64 % LANES == 0); let values = &values[offset..offset + size]; let validity = match element_validity { - ElementValidity::All => None, - ElementValidity::None => return None, - ElementValidity::Some(validity) => Some(validity.slice(offset..offset + size)), + AllOr::All => None, + AllOr::None => return None, + AllOr::Some(validity) => Some(validity.slice(offset..offset + size)), }; let mut accumulators = [ops.identity; LANES]; let mut found = false; @@ -446,22 +443,22 @@ fn reduce_group_scalar( values: &[T], offset: usize, size: usize, - element_validity: ElementValidity<'_>, + element_validity: &AllOr>, skip_nans: bool, is_better: impl Fn(T, T) -> bool, ) -> Option { let mut best = None; let group_values = &values[offset..offset + size]; match element_validity { - ElementValidity::All => { + AllOr::All => { reduce_scalar_run(&mut best, group_values, skip_nans, &is_better); } - ElementValidity::None => {} - ElementValidity::Some(validity) => { + AllOr::None => {} + AllOr::Some(validity) => { reduce_nullable_word_scalar( &mut best, group_values, - validity_word(validity, offset, size), + validity_word(*validity, offset, size), skip_nans, &is_better, ); @@ -558,7 +555,7 @@ fn reduce_scalar_run( #[inline(always)] /// Combine one full lane chunk with its packed validity bits. /// -/// The branchless combine-and-select form is intentional: after monomorphization, LLVM can lower +/// The combine-and-select form is intentional: after monomorphization, LLVM can lower /// integer min/max plus validity selection to packed SIMD instructions. Float-only NaN logic is /// removed entirely for integer types through [`ExtremaIdentity::CAN_NAN`]. fn reduce_lane_chunk( @@ -582,10 +579,10 @@ fn reduce_lane_chunk( return Some(candidate); } let reduced = combine(candidate, current); - accumulators[lane] = select(valid & !is_nan, reduced, current); + accumulators[lane] = if valid & !is_nan { reduced } else { current }; } else { let reduced = combine(candidate, current); - accumulators[lane] = select(valid, reduced, current); + accumulators[lane] = if valid { reduced } else { current }; } } None @@ -612,10 +609,10 @@ fn reduce_lane_remainder( return Some(candidate); } let reduced = combine(candidate, current); - accumulators[lane] = select(valid & !is_nan, reduced, current); + accumulators[lane] = if valid & !is_nan { reduced } else { current }; } else { let reduced = combine(candidate, current); - accumulators[lane] = select(valid, reduced, current); + accumulators[lane] = if valid { reduced } else { current }; } } None @@ -635,15 +632,13 @@ fn merge_accumulators( for lane in 0..mid { let candidate = accumulators[lane + mid]; let current = accumulators[lane]; - accumulators[lane] = select(is_better(candidate, current), candidate, current); + accumulators[lane] = if is_better(candidate, current) { + candidate + } else { + current + }; } len /= 2; } accumulators[0] } - -#[inline(always)] -/// Select a value in a form that the lane loop can lower to a SIMD mask operation. -fn select(condition: bool, if_true: T, if_false: T) -> T { - if condition { if_true } else { if_false } -} From 716d31e5696c474cf55e66f2a15d39f483a8f835 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 20 Jul 2026 15:07:24 -0700 Subject: [PATCH 3/3] refactor: simplify grouped list extrema Signed-off-by: Matt Katz --- .../src/aggregate_fn/fns/min_max/grouped.rs | 30 +++++++++---------- vortex-array/src/expr/exprs.rs | 4 +-- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs b/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs index 1ab68313355..ebf1c0a590e 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/grouped.rs @@ -39,13 +39,15 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedExtremaEncodingKernel { groups: &GroupedArray, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(options) = aggregate_fn.as_opt::() { - return try_grouped_extreme(groups, ctx, options.skip_nans, Extreme::Min); - } - if let Some(options) = aggregate_fn.as_opt::() { - return try_grouped_extreme(groups, ctx, options.skip_nans, Extreme::Max); - } - Ok(None) + let (skip_nans, extreme) = if let Some(options) = aggregate_fn.as_opt::() { + (options.skip_nans, Extreme::Min) + } else if let Some(options) = aggregate_fn.as_opt::() { + (options.skip_nans, Extreme::Max) + } else { + return Ok(None); + }; + + try_grouped_extreme(groups, ctx, skip_nans, extreme) } } @@ -477,27 +479,23 @@ fn reduce_nullable_word_scalar( mut word: u64, skip_nans: bool, is_better: &impl Fn(T, T) -> bool, -) -> bool { - let valid_mask = if values.is_empty() { - 0 - } else { - (1u64 << values.len()) - 1 - }; +) { + let valid_mask = (1u64 << values.len()) - 1; word &= valid_mask; if word == valid_mask { - return reduce_scalar_run(best, values, skip_nans, is_better); + reduce_scalar_run(best, values, skip_nans, is_better); + return; } while word != 0 { let start = word.trailing_zeros() as usize; let run_len = (word >> start).trailing_ones() as usize; if reduce_scalar_run(best, &values[start..start + run_len], skip_nans, is_better) { - return true; + return; } word &= !(((1u64 << run_len) - 1) << start); } - false } #[inline] diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 42fb7ed3772..d465d35d2e9 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -784,7 +784,7 @@ pub fn list_length(input: Expression) -> Expression { // ---- ListMin / ListMax ---- -/// Creates an expression that returns the maximum participating element in each list for `List` and +/// Creates an expression that returns the minimum participating element in each list for `List` and /// `FixedSizeList` inputs. /// /// Null lists, empty lists, and lists whose elements are all null return null. Null elements and, @@ -798,7 +798,7 @@ pub fn list_min_opts(input: Expression, options: NumericalAggregateOpts) -> Expr ListMin.new_expr(options, [input]) } -/// Creates an expression that returns the minimum participating element in each list for `List` and +/// Creates an expression that returns the maximum participating element in each list for `List` and /// `FixedSizeList` inputs. /// /// Null lists, empty lists, and lists whose elements are all null return null. Null elements and,