Skip to content

Use list length bounds for zone pruning #8924

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
mhk197 wants to merge 1 commit into
base: mk/list-shape-stats
Choose a base branch
from
Draft

Use list length bounds for zone pruning #8924

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 6 additions & 0 deletions vortex-array/src/aggregate_fn/fns/list_length_min_max.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -114,6 +116,10 @@ impl AggregateFnVTable for ListLengthMinMax {
.then(list_length_min_max_dtype)
}

fn zone_stat_default(&self, input_dtype: &DType) -> Option<AggregateFnRef> {
matches!(input_dtype, DType::List(..)).then(|| self.bind(EmptyOptions))
}

fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
self.return_dtype(options, input_dtype)
}
Expand Down
83 changes: 83 additions & 0 deletions vortex-array/src/stats/rewrite/builtins.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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::<ListLength>() {
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,
Expand All @@ -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<Expression> {
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<Item = &'a Expression>,
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(&gt(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(
Expand Down
18 changes: 18 additions & 0 deletions vortex-layout/src/layouts/zoned/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -282,4 +285,19 @@ mod tests {
.all(|aggregate_fn| !aggregate_fn.is::<Sum>())
);
}

#[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::<ListLengthMinMax>())
);
}
}
39 changes: 39 additions & 0 deletions vortex-layout/src/layouts/zoned/zone_map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down
Loading