From bcc32898beaced6b6758179e3881a615859eb77d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 23 Jul 2026 10:37:57 -0400 Subject: [PATCH] Use list length bounds for zone pruning Signed-off-by: Matt Katz --- .../aggregate_fn/fns/list_length_min_max.rs | 6 ++ vortex-array/src/stats/rewrite/builtins.rs | 83 +++++++++++++++++++ vortex-layout/src/layouts/zoned/writer.rs | 18 ++++ vortex-layout/src/layouts/zoned/zone_map.rs | 39 +++++++++ 4 files changed, 146 insertions(+) diff --git a/vortex-array/src/aggregate_fn/fns/list_length_min_max.rs b/vortex-array/src/aggregate_fn/fns/list_length_min_max.rs index d0a238b9821..1efae2102ed 100644 --- a/vortex-array/src/aggregate_fn/fns/list_length_min_max.rs +++ b/vortex-array/src/aggregate_fn/fns/list_length_min_max.rs @@ -14,7 +14,9 @@ use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMaxResult; @@ -114,6 +116,10 @@ impl AggregateFnVTable for ListLengthMinMax { .then(list_length_min_max_dtype) } + fn zone_stat_default(&self, input_dtype: &DType) -> Option { + matches!(input_dtype, DType::List(..)).then(|| self.bind(EmptyOptions)) + } + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { self.return_dtype(options, input_dtype) } diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index ea656f24546..0629e2d32d2 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -11,12 +11,16 @@ use crate::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; +use crate::aggregate_fn::fns::list_length_min_max::ListLengthMinMax; +use crate::aggregate_fn::fns::list_length_min_max::MAX_LENGTH; +use crate::aggregate_fn::fns::list_length_min_max::MIN_LENGTH; use crate::dtype::DType; use crate::expr::Expression; use crate::expr::and; use crate::expr::and_collect; use crate::expr::cast; use crate::expr::eq; +use crate::expr::get_item; use crate::expr::gt; use crate::expr::gt_eq; use crate::expr::lit; @@ -40,6 +44,7 @@ use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeVariant; use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_length::ListLength; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -616,6 +621,10 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option return cast_stat(expr.child(0), dtype, stat, ctx); } + if expr.is::() { + return list_length_stat_expr(expr.child(0), stat, ctx); + } + let aggregate_fn = stat.aggregate_fn()?; // The aggregate may not support the expression's dtype, e.g. min/max over structs, // even when the predicate itself is well-typed. Such stats cannot be lowered later, @@ -627,6 +636,47 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option .then(|| stat_fn(expr.clone(), aggregate_fn)) } +fn list_length_stat_expr( + list: &Expression, + stat: Stat, + ctx: &StatsRewriteCtx<'_>, +) -> Option { + match ctx.return_dtype(list).ok()? { + DType::List(..) => match stat { + Stat::Min | Stat::Max => { + let field = match stat { + Stat::Min => MIN_LENGTH, + Stat::Max => MAX_LENGTH, + _ => unreachable!(), + }; + Some(get_item( + field, + stat_fn(list.clone(), ListLengthMinMax.bind(AggregateEmptyOptions)), + )) + } + // list_length carries the outer list validity unchanged. + Stat::NullCount => stat_expr(list, Stat::NullCount, ctx), + Stat::IsConstant + | Stat::IsSorted + | Stat::IsStrictSorted + | Stat::Sum + | Stat::UncompressedSizeInBytes + | Stat::NaNCount => None, + }, + DType::FixedSizeList(_, size, _) => match stat { + Stat::Min | Stat::Max => Some(lit(size as u64)), + Stat::NullCount => stat_expr(list, Stat::NullCount, ctx), + Stat::IsConstant + | Stat::IsSorted + | Stat::IsStrictSorted + | Stat::Sum + | Stat::UncompressedSizeInBytes + | Stat::NaNCount => None, + }, + _ => None, + } +} + fn with_non_nan_guards<'a, P: NonNanProof>( ctx: &StatsRewriteCtx<'_>, exprs: impl IntoIterator, @@ -706,6 +756,9 @@ mod tests { use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; + use crate::aggregate_fn::fns::list_length_min_max::ListLengthMinMax; + use crate::aggregate_fn::fns::list_length_min_max::MAX_LENGTH; + use crate::aggregate_fn::fns::list_length_min_max::MIN_LENGTH; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -717,12 +770,14 @@ mod tests { use crate::expr::col; use crate::expr::dynamic; use crate::expr::eq; + use crate::expr::get_item; use crate::expr::gt; use crate::expr::gt_eq; use crate::expr::is_not_null; use crate::expr::is_null; use crate::expr::like; use crate::expr::list_contains; + use crate::expr::list_length; use crate::expr::lit; use crate::expr::lt; use crate::expr::lt_eq; @@ -757,6 +812,13 @@ mod tests { ("f", DType::Primitive(PType::F32, Nullability::NonNullable)), ("s", DType::Utf8(Nullability::NonNullable)), ("t", DType::Utf8(Nullability::NonNullable)), + ( + "l", + DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable, + ), + ), ("n", nested_struct_dtype()), ]), Nullability::NonNullable, @@ -832,6 +894,27 @@ mod tests { Ok(()) } + #[test] + fn rewrites_list_length_comparison_falsifier() -> VortexResult<()> { + let length_min_max = stat_fn(col("l"), ListLengthMinMax.bind(AggregateEmptyOptions)); + + assert_eq!( + falsify(>(list_length(col("l")), lit(2u64)))?, + Some(lt_eq( + get_item(MAX_LENGTH, length_min_max.clone()), + lit(2u64) + )) + ); + assert_eq!( + falsify(&eq(list_length(col("l")), lit(2u64)))?, + Some(or( + gt(get_item(MIN_LENGTH, length_min_max.clone()), lit(2u64)), + gt(lit(2u64), get_item(MAX_LENGTH, length_min_max)), + )) + ); + Ok(()) + } + #[test] fn rewrites_between_falsifier() -> VortexResult<()> { let expr = between( diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 048905dade6..b4a5d1dc049 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -230,8 +230,11 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A #[cfg(test)] mod tests { + use std::sync::Arc; + use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; + use vortex_array::aggregate_fn::fns::list_length_min_max::ListLengthMinMax; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::sum::Sum; @@ -282,4 +285,19 @@ mod tests { .all(|aggregate_fn| !aggregate_fn.is::()) ); } + + #[test] + fn default_aggregates_include_list_length_min_max() { + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable, + ); + let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); + + assert!( + aggregate_fns + .iter() + .any(|aggregate_fn| aggregate_fn.is::()) + ); + } } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index e5f2af494de..bf374eed759 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -371,6 +371,9 @@ mod tests { use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; + use vortex_array::aggregate_fn::fns::list_length_min_max::ListLengthMinMax; + use vortex_array::aggregate_fn::fns::list_length_min_max::MAX_LENGTH; + use vortex_array::aggregate_fn::fns::list_length_min_max::MIN_LENGTH; use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; @@ -390,6 +393,7 @@ mod tests { use vortex_array::expr::gt_eq; use vortex_array::expr::is_not_null; use vortex_array::expr::is_null; + use vortex_array::expr::list_length; use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::not_eq; @@ -401,6 +405,7 @@ mod tests { use vortex_array::stats::all_null; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; use crate::layouts::zoned::zone_map::ZoneMap; use crate::test::SESSION; @@ -483,6 +488,40 @@ mod tests { ); } + #[test] + fn list_length_min_max_prunes_list_length() -> VortexResult<()> { + let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)); + let column_dtype = DType::List(element_dtype, Nullability::Nullable); + let list_length_min_max = ListLengthMinMax.bind(EmptyOptions); + let length_min_max_values = StructArray::try_new( + [MIN_LENGTH, MAX_LENGTH].into(), + vec![ + PrimitiveArray::new(buffer![0u64, 3, 0], Validity::NonNullable).into_array(), + PrimitiveArray::new(buffer![2u64, 5, 0], Validity::NonNullable).into_array(), + ], + 3, + Validity::from_iter([true, true, false]), + )? + .into_array(); + let zone_map = ZoneMap::try_new( + column_dtype.clone(), + StructArray::from_fields(&[(list_length_min_max.to_string(), length_min_max_values)])?, + Arc::new([list_length_min_max]), + 4, + 12, + )?; + + let predicate = gt(list_length(root()), lit(2u64)); + let pruning_expr = falsify(&predicate, column_dtype); + let mask = zone_map.prune(&pruning_expr, &SESSION)?; + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([true, false, false]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + #[test] fn bounded_display_names_satisfy_min_max_rewrites() { let bounded_max = BoundedMax.bind(BoundedMaxOptions {