From c96b022d4b92cd2c1acd8a9f8ebc3e5b4a361d31 Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Mon, 9 Nov 2026 16:05:00 +0000 Subject: [PATCH 1/2] Prune large IN sets against each zone's value range Signed-off-by: Christopher Wirt --- .../read/SparkFilterToVortexExpression.java | 2 +- vortex-expr/src/stats/rewrite.rs | 1 + vortex-expr/src/stats/rewrite/builtins.rs | 39 ++- vortex-expr/src/stats/rewrite/set_disjoint.rs | 146 +++++++++ vortex-layout/src/layouts/zoned/zone_map.rs | 304 ++++++++++++++++++ 5 files changed, 479 insertions(+), 13 deletions(-) create mode 100644 vortex-expr/src/stats/rewrite/set_disjoint.rs diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java index b4ca1b3be8d..4a4d44a1e7d 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java @@ -197,7 +197,7 @@ private static Optional convertIn(In in, StructType dataSchema) { Optional lookup = literal(null, dataType) .map(elementType -> Expression.in( field, Expression.literalList(elementType, false, elements.toArray(new Expression[0])))); - // Vortex derives no zone bounds from a set, so spell them out for pruning. + // Explicit bounds also support types whose set membership has no stats rewrite. return lookup.map(set -> inBounds(in.values(), dataType, field) .map(bounds -> Expression.and(bounds, set)) .orElse(set)); diff --git a/vortex-expr/src/stats/rewrite.rs b/vortex-expr/src/stats/rewrite.rs index 2c57bf8f0a7..3417749e0fa 100644 --- a/vortex-expr/src/stats/rewrite.rs +++ b/vortex-expr/src/stats/rewrite.rs @@ -24,6 +24,7 @@ mod dynamic_expression; #[cfg(test)] mod fn_tests; mod like_variant; +mod set_disjoint; pub(crate) use builtins::register_builtins; diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index bb90bf57106..fd7c691b732 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -42,6 +42,7 @@ use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; use vortex_array::scalar_fn::fns::is_null::IsNull; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::PreparedSet; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::lower::Lower; use vortex_array::scalar_fn::fns::operators::CompareOperator; @@ -59,6 +60,8 @@ use crate::stats::rewrite::dynamic_expression::StatsProof; use crate::stats::rewrite::dynamic_expression::dynamic_stats_proof; use crate::stats::rewrite::falsify; use crate::stats::rewrite::like_variant::LikeVariant; +use crate::stats::rewrite::set_disjoint::SetDisjoint; +use crate::stats::rewrite::set_disjoint::SortedValues; use crate::stats::session::StatsSession; /// Register built-in stats rewrite rules. @@ -524,9 +527,7 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { } } -/// Lists with more elements falsify `needle IN list` with the bounds of the list rather than -/// with a test per element, which keeps the pruning expression, and its evaluation per zone, -/// small for the large lists of dynamic join filters. +/// Larger lists use one binary range lookup per zone instead of a test per element. const MAX_LIST_ELEMENT_FALSIFICATIONS: usize = 16; /// The smallest and largest non-null elements, or `None` when there are none or two elements @@ -553,12 +554,17 @@ fn list_contains_falsify( let list = expr.child(0); let needle = expr.child(1); - let Some(list_scalar) = literal_stat(list, Stat::Min) else { + // Non-float proofs are emitted by one rule only. Check before cloning and sorting a + // potentially large list for the other rule, whose unguarded proof would be discarded. + if !P::EMIT_UNGUARDED_REWRITES && matches!(P::check(needle)?, NanCheck::NotNeeded) { return Ok(None); - }; - let elements = list_scalar + } + + let list_scalar = list .as_opt::() - .and_then(|literal| literal.as_list_opt()) + .or_else(|| list.as_opt::().map(|set| set.list())); + let elements = list_scalar + .and_then(|scalar| scalar.as_list_opt()) .and_then(|list| list.elements()); let Some(elements) = elements else { return Ok(None); @@ -575,6 +581,10 @@ fn list_contains_falsify( }; let value_predicate = if elements.len() > MAX_LIST_ELEMENT_FALSIFICATIONS + && let Some(values) = SortedValues::new(elements.clone()) + { + Some(SetDisjoint.try_new_bound_expr(values, [value_min, value_max])?) + } else if elements.len() > MAX_LIST_ELEMENT_FALSIFICATIONS && let Some((low, high)) = non_null_bounds(&elements) { Some(or( @@ -913,6 +923,8 @@ mod tests { use crate::stats::expr::StatOptions; use crate::stats::rewrite::BoundExpressionPruneExt; use crate::stats::rewrite::StatsRewriteRule; + use crate::stats::rewrite::set_disjoint::SetDisjoint; + use crate::stats::rewrite::set_disjoint::SortedValues; use crate::stats::session::StatsSessionExt; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -1232,7 +1244,7 @@ mod tests { } #[test] - fn rewrites_large_list_contains_falsifier_to_bounds() -> VortexResult<()> { + fn rewrites_large_list_contains_falsifier_to_range_lookup() -> VortexResult<()> { let list = Scalar::list( Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), (0..20) @@ -1249,10 +1261,13 @@ mod tests { assert_rewrite_eq!( falsify(&expr)?, - Some(or( - lt(stat(col("a"), Stat::Max), lit(Scalar::from(Some(0i32)))), - gt(stat(col("a"), Stat::Min), lit(Scalar::from(Some(19i32)))), - )) + Some( + SetDisjoint.new_expr( + SortedValues::new((0..20).map(|value| Scalar::from(Some(value))).collect()) + .expect("integers are ordered"), + [stat(col("a"), Stat::Min), stat(col("a"), Stat::Max)], + ) + ) ); Ok(()) } diff --git a/vortex-expr/src/stats/rewrite/set_disjoint.rs b/vortex-expr/src/stats/rewrite/set_disjoint.rs new file mode 100644 index 00000000000..fe60a438618 --- /dev/null +++ b/vortex-expr/src/stats/rewrite/set_disjoint.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A compact membership falsifier over the minimum and maximum of each zone. + +use std::cmp::Ordering; +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::arrays::BoolArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::registry::CachedId; + +/// Sorted non-null members, shared by clones of the pruning expression. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(super) struct SortedValues { + dtype: DType, + values: Arc<[Scalar]>, +} + +impl SortedValues { + pub(super) fn new(elements: Vec) -> Option { + let dtype = elements.first()?.dtype().clone(); + if !matches!( + dtype, + DType::Bool(_) + | DType::Primitive(..) + | DType::Decimal(..) + | DType::Utf8(_) + | DType::Binary(_) + ) || elements + .iter() + .any(|value| !value.dtype().eq_ignore_nullability(&dtype)) + { + return None; + } + let mut values = elements + .into_iter() + .filter(|value| !value.is_null()) + .collect::>(); + // These scalar types have the same total order as their min/max aggregates. In + // particular, floating-point order distinguishes signed zero and NaN bit patterns. + values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal)); + values.dedup(); + Some(Self { + dtype, + values: values.into(), + }) + } + + fn disjoint(&self, min: &Scalar, max: &Scalar) -> bool { + // Missing bounds prove nothing, including for all-null zones. Float callers also + // retain the existing non-NaN proof because min/max may omit matching NaN needles. + if min.is_null() || max.is_null() || min.partial_cmp(max).is_none() || min > max { + return false; + } + let first = self.values.partition_point(|value| value < min); + self.values + .get(first) + .is_none_or(|value| value.partial_cmp(max) == Some(Ordering::Greater)) + } +} + +impl Display for SortedValues { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{} members", self.values.len()) + } +} + +/// An internal stats expression: true only if no set member falls in `[min, max]`. +#[derive(Clone, Debug)] +pub(super) struct SetDisjoint; + +impl ScalarFnVTable for SetDisjoint { + type Options = SortedValues; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.stats.set_disjoint"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => "min".into(), + 1 => "max".into(), + _ => unreachable!(), + } + } + + fn return_dtype(&self, values: &Self::Options, args: &[DType]) -> VortexResult { + vortex_ensure!(args[0].eq_ignore_nullability(&args[1])); + vortex_ensure!( + values.dtype.eq_ignore_nullability(&args[0]), + "Set members must have the dtype of the zone bounds" + ); + Ok(DType::Bool(Nullability::NonNullable)) + } + + fn execute( + &self, + values: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Materialize once and retain probe state, rather than executing each bound per zone. + let mut min = args + .get(0)? + .execute::(ctx)? + .0 + .into_array() + .repeated_probe(); + let mut max = args + .get(1)? + .execute::(ctx)? + .0 + .into_array() + .repeated_probe(); + let results = (0..args.row_count()) + .map(|idx| { + Ok(values.disjoint( + &min.execute_scalar(idx, ctx)?, + &max.execute_scalar(idx, ctx)?, + )) + }) + .collect::>>()?; + Ok(BoolArray::from_iter(results).into_array()) + } +} diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 8d6eaf3a8d4..b3832a480b5 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -417,6 +417,7 @@ mod tests { use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -429,6 +430,7 @@ mod tests { use vortex_array::expr::dynamic; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; + use vortex_array::expr::in_list; use vortex_array::expr::is_not_null; use vortex_array::expr::is_null; use vortex_array::expr::lit; @@ -436,6 +438,7 @@ mod tests { use vortex_array::expr::not_eq; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::validity::Validity; @@ -469,6 +472,307 @@ mod tests { unsafe { NonZeroUsize::new_unchecked(64) } } + #[rstest] + #[case(16, false, false)] + #[case(17, false, false)] + #[case(80, false, false)] + #[case(80, true, false)] + #[case(80, false, true)] + #[case(80, true, true)] + fn in_list_prunes_gaps_between_members( + #[case] count: i32, + #[case] optimized: bool, + #[case] null_and_duplicate: bool, + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let list = Scalar::list( + Arc::new(dtype.clone()), + (0..count) + .rev() + .map(|value| Scalar::from(Some(value * 10))) + .chain( + null_and_duplicate + .then_some([Some(10i32).into(), Scalar::null(dtype.clone())]) + .into_iter() + .flatten(), + ) + .collect(), + Nullability::NonNullable, + ); + let last = (count - 1) * 10; + let min = Min.bind(NumericalAggregateOpts::skip_nans()); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let zone_map = ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&[ + ( + min.to_string(), + PrimitiveArray::from_option_iter([ + Some(-1i32), + Some(0), + Some(1), + Some(9), + Some(11), + Some(last), + Some(last + 1), + None, + ]) + .into_array(), + ), + ( + max.to_string(), + PrimitiveArray::from_option_iter([ + Some(-1i32), + Some(0), + Some(9), + Some(10), + Some(19), + Some(last), + Some(last + 10), + None, + ]) + .into_array(), + ), + ])?, + Arc::new([min, max]), + 4, + 32, + )?; + let expr = in_list(root(), lit(list)).bind(&dtype)?; + let expr = if optimized { + expr.optimize_recursive()? + } else { + expr + }; + let proof = expr + .falsify(&SESSION)? + .expect("constant IN list has a falsifier"); + assert_arrays_eq!( + zone_map.prune(&proof, &SESSION)?.into_array(), + BoolArray::from_iter([true, false, true, false, true, false, true, false]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[rstest] + fn large_string_in_list_prunes_gaps( + #[values(false, true)] optimized: bool, + ) -> VortexResult<()> { + let dtype = DType::Utf8(Nullability::Nullable); + let list = Scalar::list( + Arc::new(dtype.clone()), + (0..80) + .rev() + .map(|value| Scalar::utf8(format!("key-{value:03}"), Nullability::Nullable)) + .chain([Scalar::null(dtype.clone())]) + .collect(), + Nullability::NonNullable, + ); + let min = Min.bind(NumericalAggregateOpts::skip_nans()); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let zone_map = ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&[ + ( + min.to_string(), + VarBinViewArray::from_iter( + [ + Some("key-000"), + Some("key-010a"), + Some("key-078z"), + Some("key-080"), + None, + ], + dtype.clone(), + ) + .into_array(), + ), + ( + max.to_string(), + VarBinViewArray::from_iter( + [ + Some("key-000"), + Some("key-010z"), + Some("key-079"), + Some("key-099"), + None, + ], + dtype.clone(), + ) + .into_array(), + ), + ])?, + Arc::new([min, max]), + 4, + 20, + )?; + let expr = in_list(root(), lit(list)).bind(&dtype)?; + let expr = if optimized { + expr.optimize_recursive()? + } else { + expr + }; + let proof = expr + .falsify(&SESSION)? + .expect("constant IN list has a falsifier"); + assert_arrays_eq!( + zone_map.prune(&proof, &SESSION)?.into_array(), + BoolArray::from_iter([false, true, false, true, false]), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[rstest] + fn large_float_in_list_preserves_nan_and_signed_zero( + #[values(false, true)] optimized: bool, + #[values(false, true)] nan_stats: bool, + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::F64, Nullability::Nullable); + let list = Scalar::list( + Arc::new(dtype.clone()), + (10..30) + .map(|value| Scalar::from(Some(f64::from(value)))) + .chain([Some(-0.0f64).into(), Some(f64::NAN).into()]) + .collect(), + Nullability::NonNullable, + ); + let min = Min.bind(NumericalAggregateOpts::skip_nans()); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let mut fields = vec![ + ( + min.to_string(), + PrimitiveArray::from_option_iter([ + Some(5.0f64), + Some(5.0), + Some(-0.0), + Some(0.0), + None, + ]) + .into_array(), + ), + ( + max.to_string(), + PrimitiveArray::from_option_iter([ + Some(5.0f64), + Some(5.0), + Some(-0.0), + Some(0.0), + None, + ]) + .into_array(), + ), + ]; + let mut aggregates = vec![min, max]; + if nan_stats { + let nan_count = NanCount.bind(EmptyOptions); + fields.push(( + nan_count.to_string(), + PrimitiveArray::new(buffer![1u64, 0, 0, 0, 0], Validity::AllValid).into_array(), + )); + aggregates.push(nan_count); + } + let zone_map = ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&fields)?, + aggregates.into(), + 4, + 20, + )?; + let expr = in_list(root(), lit(list)).bind(&dtype)?; + let expr = if optimized { + expr.optimize_recursive()? + } else { + expr + }; + let proof = expr + .falsify(&SESSION)? + .expect("constant IN list has a falsifier"); + // Zone zero can contain a matching NaN although its finite extrema are outside the + // set. Signed zero membership follows the same total order as the min/max aggregates. + let mut ctx = SESSION.create_execution_ctx(); + let actual = PrimitiveArray::from_option_iter([ + Some(f64::NAN), + Some(5.0), + Some(-0.0), + Some(0.0), + None, + ]) + .into_array() + .apply_bound(&expr)?; + let matching = super::null_as_false_mask(actual.clone(), &mut ctx)?; + assert_arrays_eq!( + actual, + BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false), None,]), + &mut ctx + ); + let pruned = zone_map.prune(&proof, &SESSION)?; + for zone in 0..5 { + assert!( + !(matching.value(zone) && pruned.value(zone)), + "pruned matching zone {zone}" + ); + } + let expected = if nan_stats { + [false, true, false, true, false] + } else { + [false; 5] + }; + assert_arrays_eq!( + pruned.into_array(), + BoolArray::from_iter(expected), + &mut ctx + ); + Ok(()) + } + + #[test] + fn large_in_list_pruning_stress() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let list = Scalar::list( + Arc::new(dtype.clone()), + (0..16384) + .rev() + .map(|value| Scalar::from(value * 10)) + .collect(), + Nullability::NonNullable, + ); + let values = (0..8192) + .map(|zone| zone * 10 + i32::from(zone % 11 != 0)) + .collect::>(); + let needles = PrimitiveArray::new(values.clone(), Validity::NonNullable).into_array(); + let bounds = PrimitiveArray::new(values, Validity::AllValid).into_array(); + let min = Min.bind(NumericalAggregateOpts::skip_nans()); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let zone_map = ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&[ + (min.to_string(), bounds.clone()), + (max.to_string(), bounds), + ])?, + Arc::new([min, max]), + 4, + 32768, + )?; + let expr = in_list(root(), lit(list)) + .bind(&dtype)? + .optimize_recursive()?; + let proof = expr + .falsify(&SESSION)? + .expect("constant IN list has a falsifier"); + let pruned = zone_map.prune(&proof, &SESSION)?; + let mut ctx = SESSION.create_execution_ctx(); + // Every zone holds one distinct value. Its complement must agree with actual batch + // membership, including the thousands of zones between selected keys. + assert_arrays_eq!( + needles.apply_bound(&expr)?, + (!pruned).into_array(), + &mut ctx + ); + Ok(()) + } + #[test] fn test_zone_map_prunes() { // Construct a zone map with 3 zones: From 7d14d6923796c95d3bc251d68521464bf1001b40 Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Tue, 10 Nov 2026 10:40:37 +0000 Subject: [PATCH 2/2] refactor: simplify large IN pruning and drop redundant Spark bounds The Spark IN bounds only existed because the stats rewrite could not see a prepared set; it now can, so they duplicated the Rust pruning. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Christopher Wirt --- .../read/SparkFilterToVortexExpression.java | 68 +--- vortex-expr/src/stats/rewrite/builtins.rs | 54 +-- vortex-expr/src/stats/rewrite/set_disjoint.rs | 64 ++-- vortex-layout/src/layouts/zoned/zone_map.rs | 311 ++++++++---------- 4 files changed, 192 insertions(+), 305 deletions(-) diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java index 4a4d44a1e7d..ab23f77b3fd 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java @@ -175,7 +175,7 @@ private static Optional convertIn(In in, StructType dataSchema) { } Expression field = column.get().expression(); if (holdsNull) { - // A non-match must remain null beneath NOT. Range bounds would turn null into false outside the range. + // A non-match must remain null beneath NOT. return literal(null, dataType).map(elementType -> { elements.add(elementType); return Expression.in( @@ -194,73 +194,9 @@ private static Optional convertIn(In in, StructType dataSchema) { : Expression.or(comparisons.toArray(new Expression[0]))); } // A disjunction scans the column once per value, the set once in all. - Optional lookup = literal(null, dataType) + return literal(null, dataType) .map(elementType -> Expression.in( field, Expression.literalList(elementType, false, elements.toArray(new Expression[0])))); - // Explicit bounds also support types whose set membership has no stats rewrite. - return lookup.map(set -> inBounds(in.values(), dataType, field) - .map(bounds -> Expression.and(bounds, set)) - .orElse(set)); - } - - private static Optional inBounds(Object[] values, DataType dataType, Expression field) { - Object low = null; - Object high = null; - Comparable lowKey = null; - Comparable highKey = null; - for (Object value : values) { - if (value == null) { - continue; - } - Optional> key = orderKey(value, dataType); - if (key.isEmpty()) { - return Optional.empty(); - } - if (lowKey == null || key.get().compareTo(lowKey) < 0) { - low = value; - lowKey = key.get(); - } - if (highKey == null || key.get().compareTo(highKey) > 0) { - high = value; - highKey = key.get(); - } - } - Optional lower = literal(low, dataType); - Optional upper = literal(high, dataType); - if (lowKey == null || lower.isEmpty() || upper.isEmpty()) { - return Optional.empty(); - } - return Optional.of(Expression.and( - Expression.binary(BinaryOp.GTE, field, lower.get()), - Expression.binary(BinaryOp.LTE, field, upper.get()))); - } - - /** {@code value} as a key in the order Vortex compares {@code dataType} in, or empty for an unordered type. */ - @SuppressWarnings("unchecked") - private static Optional> orderKey(Object value, DataType dataType) { - Optional> key; - if (dataType instanceof ByteType - || dataType instanceof ShortType - || dataType instanceof IntegerType - || dataType instanceof LongType) { - key = value instanceof Number number ? Optional.of(number.longValue()) : Optional.empty(); - } else if (dataType instanceof DateType) { - key = dateDays(value); - } else if (dataType instanceof TimestampType || dataType instanceof TimestampNTZType) { - key = timestampMicros(value, dataType instanceof TimestampNTZType); - } else if (dataType instanceof DecimalType) { - key = value instanceof BigDecimal decimal - ? Optional.of(decimal) - : value instanceof Decimal sparkDecimal - ? Optional.of(sparkDecimal.toJavaBigDecimal()) - : Optional.empty(); - } else if (dataType instanceof StringType) { - // UTF8String orders by unsigned bytes, as Vortex does; String would order by UTF-16 units. - key = Optional.of(UTF8String.fromString(value.toString())); - } else { - key = Optional.empty(); - } - return key.map(comparable -> (Comparable) comparable); } private static Optional stringMatch( diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index fd7c691b732..0c4f7a27d21 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -527,7 +527,8 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { } } -/// Larger lists use one binary range lookup per zone instead of a test per element. +/// Larger lists use one binary range lookup per zone instead of a test per element, or the +/// list's overall bounds when its dtype is not supported by the lookup. const MAX_LIST_ELEMENT_FALSIFICATIONS: usize = 16; /// The smallest and largest non-null elements, or `None` when there are none or two elements @@ -554,8 +555,7 @@ fn list_contains_falsify( let list = expr.child(0); let needle = expr.child(1); - // Non-float proofs are emitted by one rule only. Check before cloning and sorting a - // potentially large list for the other rule, whose unguarded proof would be discarded. + // Only one rule emits non-float proofs, so skip the list work in the other. if !P::EMIT_UNGUARDED_REWRITES && matches!(P::check(needle)?, NanCheck::NotNeeded) { return Ok(None); } @@ -563,10 +563,10 @@ fn list_contains_falsify( let list_scalar = list .as_opt::() .or_else(|| list.as_opt::().map(|set| set.list())); - let elements = list_scalar - .and_then(|scalar| scalar.as_list_opt()) - .and_then(|list| list.elements()); - let Some(elements) = elements else { + let Some(list) = list_scalar.and_then(|scalar| scalar.as_list_opt()) else { + return Ok(None); + }; + let Some(elements) = list.elements() else { return Ok(None); }; if elements.is_empty() { @@ -580,18 +580,7 @@ fn list_contains_falsify( return Ok(None); }; - let value_predicate = if elements.len() > MAX_LIST_ELEMENT_FALSIFICATIONS - && let Some(values) = SortedValues::new(elements.clone()) - { - Some(SetDisjoint.try_new_bound_expr(values, [value_min, value_max])?) - } else if elements.len() > MAX_LIST_ELEMENT_FALSIFICATIONS - && let Some((low, high)) = non_null_bounds(&elements) - { - Some(or( - lt(value_max, lit(low.clone())), - gt(value_min, lit(high.clone())), - )) - } else { + let per_element = |elements: &[Scalar]| { and_collect(elements.iter().map(|value| { or( lt(value_max.clone(), lit(value.clone())), @@ -599,6 +588,19 @@ fn list_contains_falsify( ) })) }; + let value_predicate = if elements.len() <= MAX_LIST_ELEMENT_FALSIFICATIONS { + per_element(&elements) + } else if SortedValues::supports(list.element_dtype()) { + let values = SortedValues::new(list.element_dtype().clone(), elements); + Some(SetDisjoint.try_new_bound_expr(values, [value_min.clone(), value_max.clone()])?) + } else if let Some((low, high)) = non_null_bounds(&elements) { + Some(or( + lt(value_max.clone(), lit(low.clone())), + gt(value_min.clone(), lit(high.clone())), + )) + } else { + per_element(&elements) + }; value_predicate .map(|value_predicate| with_non_nan_guards::

([needle], value_predicate)) .transpose() @@ -1261,13 +1263,13 @@ mod tests { assert_rewrite_eq!( falsify(&expr)?, - Some( - SetDisjoint.new_expr( - SortedValues::new((0..20).map(|value| Scalar::from(Some(value))).collect()) - .expect("integers are ordered"), - [stat(col("a"), Stat::Min), stat(col("a"), Stat::Max)], - ) - ) + Some(SetDisjoint.new_expr( + SortedValues::new( + DType::Primitive(PType::I32, Nullability::Nullable), + (0..20).map(|value| Scalar::from(Some(value))).collect(), + ), + [stat(col("a"), Stat::Min), stat(col("a"), Stat::Max)], + )) ); Ok(()) } diff --git a/vortex-expr/src/stats/rewrite/set_disjoint.rs b/vortex-expr/src/stats/rewrite/set_disjoint.rs index fe60a438618..db6a25dca3d 100644 --- a/vortex-expr/src/stats/rewrite/set_disjoint.rs +++ b/vortex-expr/src/stats/rewrite/set_disjoint.rs @@ -33,21 +33,21 @@ pub(super) struct SortedValues { } impl SortedValues { - pub(super) fn new(elements: Vec) -> Option { - let dtype = elements.first()?.dtype().clone(); - if !matches!( + /// Whether members of `dtype` order as their min/max aggregates do. + pub(super) fn supports(dtype: &DType) -> bool { + matches!( dtype, DType::Bool(_) | DType::Primitive(..) | DType::Decimal(..) | DType::Utf8(_) | DType::Binary(_) - ) || elements - .iter() - .any(|value| !value.dtype().eq_ignore_nullability(&dtype)) - { - return None; - } + ) + } + + /// The members of a list of `dtype` elements, which must be [supported](Self::supports). + pub(super) fn new(dtype: DType, elements: Vec) -> Self { + debug_assert!(Self::supports(&dtype)); let mut values = elements .into_iter() .filter(|value| !value.is_null()) @@ -56,22 +56,22 @@ impl SortedValues { // particular, floating-point order distinguishes signed zero and NaN bit patterns. values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal)); values.dedup(); - Some(Self { + Self { dtype, values: values.into(), - }) + } } fn disjoint(&self, min: &Scalar, max: &Scalar) -> bool { - // Missing bounds prove nothing, including for all-null zones. Float callers also - // retain the existing non-NaN proof because min/max may omit matching NaN needles. - if min.is_null() || max.is_null() || min.partial_cmp(max).is_none() || min > max { + // Missing bounds prove nothing, including for all-null zones. + if min.is_null() + || max.is_null() + || matches!(min.partial_cmp(max), None | Some(Ordering::Greater)) + { return false; } let first = self.values.partition_point(|value| value < min); - self.values - .get(first) - .is_none_or(|value| value.partial_cmp(max) == Some(Ordering::Greater)) + self.values.get(first).is_none_or(|value| value > max) } } @@ -106,7 +106,12 @@ impl ScalarFnVTable for SetDisjoint { } fn return_dtype(&self, values: &Self::Options, args: &[DType]) -> VortexResult { - vortex_ensure!(args[0].eq_ignore_nullability(&args[1])); + vortex_ensure!( + args[0].eq_ignore_nullability(&args[1]), + "Zone bounds must share a dtype, got {} and {}", + args[0], + args[1] + ); vortex_ensure!( values.dtype.eq_ignore_nullability(&args[0]), "Set members must have the dtype of the zone bounds" @@ -120,19 +125,16 @@ impl ScalarFnVTable for SetDisjoint { args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Materialize once and retain probe state, rather than executing each bound per zone. - let mut min = args - .get(0)? - .execute::(ctx)? - .0 - .into_array() - .repeated_probe(); - let mut max = args - .get(1)? - .execute::(ctx)? - .0 - .into_array() - .repeated_probe(); + let mut probe = |idx| -> VortexResult<_> { + Ok(args + .get(idx)? + .execute::(ctx)? + .0 + .into_array() + .repeated_probe()) + }; + let mut min = probe(0)?; + let mut max = probe(1)?; let results = (0..args.row_count()) .map(|idx| { Ok(values.disjoint( diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index b3832a480b5..84f8d5d38a8 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -397,8 +397,10 @@ mod tests { use std::sync::atomic::Ordering; use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; @@ -472,6 +474,50 @@ mod tests { unsafe { NonZeroUsize::new_unchecked(64) } } + /// A zone map holding the skip-NaN min and max of each zone, plus any `extra` aggregates. + fn min_max_zone_map( + dtype: &DType, + mins: ArrayRef, + maxs: ArrayRef, + extra: impl IntoIterator, + row_count: u64, + ) -> VortexResult { + let min = Min.bind(NumericalAggregateOpts::skip_nans()); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let (aggregates, fields): (Vec<_>, Vec<_>) = [(min, mins), (max, maxs)] + .into_iter() + .chain(extra) + .map(|(aggregate, array)| { + let name = aggregate.to_string(); + (aggregate, (name, array)) + }) + .unzip(); + ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&fields)?, + aggregates.into(), + 4, + row_count, + ) + } + + /// The bound `root() IN list` and its falsifier, optionally after optimization prepares the + /// list as a set. + fn in_list_falsifier( + dtype: &DType, + list: Scalar, + optimized: bool, + ) -> VortexResult<(BoundExpression, BoundExpression)> { + let mut expr = in_list(root(), lit(list)).bind(dtype)?; + if optimized { + expr = expr.optimize_recursive()?; + } + let proof = expr + .falsify(&SESSION)? + .expect("constant IN list has a falsifier"); + Ok((expr, proof)) + } + #[rstest] #[case(16, false, false)] #[case(17, false, false)] @@ -500,53 +546,34 @@ mod tests { Nullability::NonNullable, ); let last = (count - 1) * 10; - let min = Min.bind(NumericalAggregateOpts::skip_nans()); - let max = Max.bind(NumericalAggregateOpts::skip_nans()); - let zone_map = ZoneMap::try_new( - dtype.clone(), - StructArray::from_fields(&[ - ( - min.to_string(), - PrimitiveArray::from_option_iter([ - Some(-1i32), - Some(0), - Some(1), - Some(9), - Some(11), - Some(last), - Some(last + 1), - None, - ]) - .into_array(), - ), - ( - max.to_string(), - PrimitiveArray::from_option_iter([ - Some(-1i32), - Some(0), - Some(9), - Some(10), - Some(19), - Some(last), - Some(last + 10), - None, - ]) - .into_array(), - ), - ])?, - Arc::new([min, max]), - 4, + let zone_map = min_max_zone_map( + &dtype, + PrimitiveArray::from_option_iter([ + Some(-1i32), + Some(0), + Some(1), + Some(9), + Some(11), + Some(last), + Some(last + 1), + None, + ]) + .into_array(), + PrimitiveArray::from_option_iter([ + Some(-1i32), + Some(0), + Some(9), + Some(10), + Some(19), + Some(last), + Some(last + 10), + None, + ]) + .into_array(), + [], 32, )?; - let expr = in_list(root(), lit(list)).bind(&dtype)?; - let expr = if optimized { - expr.optimize_recursive()? - } else { - expr - }; - let proof = expr - .falsify(&SESSION)? - .expect("constant IN list has a falsifier"); + let (_, proof) = in_list_falsifier(&dtype, list, optimized)?; assert_arrays_eq!( zone_map.prune(&proof, &SESSION)?.into_array(), BoolArray::from_iter([true, false, true, false, true, false, true, false]), @@ -569,53 +596,34 @@ mod tests { .collect(), Nullability::NonNullable, ); - let min = Min.bind(NumericalAggregateOpts::skip_nans()); - let max = Max.bind(NumericalAggregateOpts::skip_nans()); - let zone_map = ZoneMap::try_new( - dtype.clone(), - StructArray::from_fields(&[ - ( - min.to_string(), - VarBinViewArray::from_iter( - [ - Some("key-000"), - Some("key-010a"), - Some("key-078z"), - Some("key-080"), - None, - ], - dtype.clone(), - ) - .into_array(), - ), - ( - max.to_string(), - VarBinViewArray::from_iter( - [ - Some("key-000"), - Some("key-010z"), - Some("key-079"), - Some("key-099"), - None, - ], - dtype.clone(), - ) - .into_array(), - ), - ])?, - Arc::new([min, max]), - 4, + let zone_map = min_max_zone_map( + &dtype, + VarBinViewArray::from_iter( + [ + Some("key-000"), + Some("key-010a"), + Some("key-078z"), + Some("key-080"), + None, + ], + dtype.clone(), + ) + .into_array(), + VarBinViewArray::from_iter( + [ + Some("key-000"), + Some("key-010z"), + Some("key-079"), + Some("key-099"), + None, + ], + dtype.clone(), + ) + .into_array(), + [], 20, )?; - let expr = in_list(root(), lit(list)).bind(&dtype)?; - let expr = if optimized { - expr.optimize_recursive()? - } else { - expr - }; - let proof = expr - .falsify(&SESSION)? - .expect("constant IN list has a falsifier"); + let (_, proof) = in_list_falsifier(&dtype, list, optimized)?; assert_arrays_eq!( zone_map.prune(&proof, &SESSION)?.into_array(), BoolArray::from_iter([false, true, false, true, false]), @@ -638,89 +646,45 @@ mod tests { .collect(), Nullability::NonNullable, ); - let min = Min.bind(NumericalAggregateOpts::skip_nans()); - let max = Max.bind(NumericalAggregateOpts::skip_nans()); - let mut fields = vec![ - ( - min.to_string(), - PrimitiveArray::from_option_iter([ - Some(5.0f64), - Some(5.0), - Some(-0.0), - Some(0.0), - None, - ]) - .into_array(), - ), - ( - max.to_string(), - PrimitiveArray::from_option_iter([ - Some(5.0f64), - Some(5.0), - Some(-0.0), - Some(0.0), - None, - ]) - .into_array(), - ), - ]; - let mut aggregates = vec![min, max]; - if nan_stats { - let nan_count = NanCount.bind(EmptyOptions); - fields.push(( - nan_count.to_string(), - PrimitiveArray::new(buffer![1u64, 0, 0, 0, 0], Validity::AllValid).into_array(), - )); - aggregates.push(nan_count); - } - let zone_map = ZoneMap::try_new( - dtype.clone(), - StructArray::from_fields(&fields)?, - aggregates.into(), - 4, - 20, - )?; - let expr = in_list(root(), lit(list)).bind(&dtype)?; - let expr = if optimized { - expr.optimize_recursive()? - } else { - expr - }; - let proof = expr - .falsify(&SESSION)? - .expect("constant IN list has a falsifier"); - // Zone zero can contain a matching NaN although its finite extrema are outside the - // set. Signed zero membership follows the same total order as the min/max aggregates. - let mut ctx = SESSION.create_execution_ctx(); - let actual = PrimitiveArray::from_option_iter([ - Some(f64::NAN), + let bounds = PrimitiveArray::from_option_iter([ + Some(5.0f64), Some(5.0), Some(-0.0), Some(0.0), None, ]) - .into_array() - .apply_bound(&expr)?; - let matching = super::null_as_false_mask(actual.clone(), &mut ctx)?; + .into_array(); + let nan_count = nan_stats.then(|| { + ( + NanCount.bind(EmptyOptions), + PrimitiveArray::new(buffer![1u64, 0, 0, 0, 0], Validity::AllValid).into_array(), + ) + }); + let zone_map = min_max_zone_map(&dtype, bounds.clone(), bounds, nan_count, 20)?; + let (expr, proof) = in_list_falsifier(&dtype, list, optimized)?; + // Zone zero holds a matching NaN although its finite extrema are outside the set, so + // only NaN stats can prune it. + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( - actual, - BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false), None,]), + PrimitiveArray::from_option_iter([ + Some(f64::NAN), + Some(5.0), + Some(-0.0), + Some(0.0), + None, + ]) + .into_array() + .apply_bound(&expr)?, + BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false), None]), &mut ctx ); - let pruned = zone_map.prune(&proof, &SESSION)?; - for zone in 0..5 { - assert!( - !(matching.value(zone) && pruned.value(zone)), - "pruned matching zone {zone}" - ); - } let expected = if nan_stats { [false, true, false, true, false] } else { [false; 5] }; assert_arrays_eq!( - pruned.into_array(), + zone_map.prune(&proof, &SESSION)?.into_array(), BoolArray::from_iter(expected), &mut ctx ); @@ -743,32 +707,15 @@ mod tests { .collect::>(); let needles = PrimitiveArray::new(values.clone(), Validity::NonNullable).into_array(); let bounds = PrimitiveArray::new(values, Validity::AllValid).into_array(); - let min = Min.bind(NumericalAggregateOpts::skip_nans()); - let max = Max.bind(NumericalAggregateOpts::skip_nans()); - let zone_map = ZoneMap::try_new( - dtype.clone(), - StructArray::from_fields(&[ - (min.to_string(), bounds.clone()), - (max.to_string(), bounds), - ])?, - Arc::new([min, max]), - 4, - 32768, - )?; - let expr = in_list(root(), lit(list)) - .bind(&dtype)? - .optimize_recursive()?; - let proof = expr - .falsify(&SESSION)? - .expect("constant IN list has a falsifier"); + let zone_map = min_max_zone_map(&dtype, bounds.clone(), bounds, [], 32768)?; + let (expr, proof) = in_list_falsifier(&dtype, list, true)?; let pruned = zone_map.prune(&proof, &SESSION)?; - let mut ctx = SESSION.create_execution_ctx(); - // Every zone holds one distinct value. Its complement must agree with actual batch - // membership, including the thousands of zones between selected keys. + // Every zone holds one distinct value, so the zones kept must be exactly those whose + // value is a member, including among the thousands of zones between selected keys. assert_arrays_eq!( - needles.apply_bound(&expr)?, (!pruned).into_array(), - &mut ctx + needles.apply_bound(&expr)?, + &mut SESSION.create_execution_ctx() ); Ok(()) }