Prune large IN sets against each zone's value range

java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java

@@ -175,7 +175,7 @@ private static Optional<Expression> 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<Expression> 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<Expression> 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<Expression> inBounds(Object[] values, DataType dataType, Expression field) {
-        Object low = null;
-        Object high = null;
-        Comparable<Object> lowKey = null;
-        Comparable<Object> highKey = null;
-        for (Object value : values) {
-            if (value == null) {
-                continue;
-            }
-            Optional<Comparable<Object>> 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<Expression> lower = literal(low, dataType);
-        Optional<Expression> 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<Comparable<Object>> orderKey(Object value, DataType dataType) {
-        Optional<? extends Comparable<?>> 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<Object>) comparable);
     }
 
     private static Optional<Expression> stringMatch(

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;
 

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<P: NonNanProof>(
     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::<Literal>()
-        .and_then(|literal| literal.as_list_opt())
-        .and_then(|list| list.elements());
-    let Some(elements) = elements else {
+        .or_else(|| list.as_opt::<PreparedSet>().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,21 +580,27 @@ fn list_contains_falsify<P: NonNanProof>(
         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())),
                 gt(value_min.clone(), lit(value.clone())),
             )
         }))
     };
+    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::<P>([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<VortexSession> = 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(())

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<Scalar>) -> Self {
+        debug_assert!(Self::supports(&dtype));
+        let mut values = elements
+            .into_iter()
+            .filter(|value| !value.is_null())
+            .collect::<Vec<_>>();
+        // 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<DType> {
+        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<ArrayRef> {
+        let mut probe = |idx| -> VortexResult<_> {
+            Ok(args
+                .get(idx)?
+                .execute::<RecursiveCanonical>(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::<VortexResult<Vec<_>>>()?;
+        Ok(BoolArray::from_iter(results).into_array())
+    }
+}