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..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])))); - // Vortex derives no zone bounds from a set, so spell them out for pruning. - 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.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..0c4f7a27d21 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,8 @@ 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, 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 @@ -553,14 +555,18 @@ fn list_contains_falsify( let list = expr.child(0); let needle = expr.child(1); - let Some(list_scalar) = literal_stat(list, Stat::Min) else { + // 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); - }; - let elements = list_scalar + } + + let list_scalar = list .as_opt::() - .and_then(|literal| literal.as_list_opt()) - .and_then(|list| list.elements()); - let Some(elements) = elements else { + .or_else(|| list.as_opt::().map(|set| set.list())); + 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() { @@ -574,14 +580,7 @@ fn list_contains_falsify( return Ok(None); }; - let value_predicate = 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())), @@ -589,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() @@ -913,6 +925,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 +1246,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,9 +1263,12 @@ 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( + 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 new file mode 100644 index 00000000000..db6a25dca3d --- /dev/null +++ b/vortex-expr/src/stats/rewrite/set_disjoint.rs @@ -0,0 +1,148 @@ +// 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 { + /// 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(_) + ) + } + + /// 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()) + .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(); + Self { + dtype, + values: values.into(), + } + } + + fn disjoint(&self, min: &Scalar, max: &Scalar) -> bool { + // 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 > max) + } +} + +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]), + "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" + ); + Ok(DType::Bool(Nullability::NonNullable)) + } + + fn execute( + &self, + values: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + 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( + &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..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; @@ -417,6 +419,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 +432,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 +440,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 +474,252 @@ 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)] + #[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 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 (_, 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]), + &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 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 (_, 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]), + &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 bounds = PrimitiveArray::from_option_iter([ + Some(5.0f64), + Some(5.0), + Some(-0.0), + Some(0.0), + None, + ]) + .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!( + 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 expected = if nan_stats { + [false, true, false, true, false] + } else { + [false; 5] + }; + assert_arrays_eq!( + zone_map.prune(&proof, &SESSION)?.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 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)?; + // 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!( + (!pruned).into_array(), + needles.apply_bound(&expr)?, + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + #[test] fn test_zone_map_prunes() { // Construct a zone map with 3 zones: