diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index ea656f24546..ee6e1706ec3 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -83,6 +83,14 @@ impl StatsRewriteRule for BinaryNanCountStatsRewrite { ) -> VortexResult> { binary_falsify::(expr, ctx) } + + fn satisfy( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + binary_satisfy::(expr, ctx) + } } #[derive(Debug)] @@ -100,6 +108,14 @@ impl StatsRewriteRule for BinaryAllNonNanStatsRewrite { ) -> VortexResult> { binary_falsify::(expr, ctx) } + + fn satisfy( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + binary_satisfy::(expr, ctx) + } } fn binary_falsify( @@ -168,6 +184,82 @@ fn binary_falsify( }) } +fn binary_satisfy( + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + let operator = expr.as_::(); + let lhs = expr.child(0); + let rhs = expr.child(1); + + Ok(match operator { + Operator::Eq => min(lhs, ctx) + .zip(max(rhs, ctx)) + .zip(max(lhs, ctx).zip(min(rhs, ctx))) + .map(|((min_lhs, max_rhs), (max_lhs, min_rhs))| { + with_comparison_guards::

( + ctx, + [lhs, rhs], + and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)), + ) + }) + .transpose()? + .flatten(), + Operator::NotEq => min(lhs, ctx) + .zip(max(rhs, ctx)) + .zip(min(rhs, ctx).zip(max(lhs, ctx))) + .map(|((min_lhs, max_rhs), (min_rhs, max_lhs))| { + with_comparison_guards::

( + ctx, + [lhs, rhs], + or(gt(min_lhs, max_rhs), gt(min_rhs, max_lhs)), + ) + }) + .transpose()? + .flatten(), + Operator::Gt => min(lhs, ctx) + .zip(max(rhs, ctx)) + .map(|(a, b)| with_comparison_guards::

(ctx, [lhs, rhs], gt(a, b))) + .transpose()? + .flatten(), + Operator::Gte => min(lhs, ctx) + .zip(max(rhs, ctx)) + .map(|(a, b)| with_comparison_guards::

(ctx, [lhs, rhs], gt_eq(a, b))) + .transpose()? + .flatten(), + Operator::Lt => max(lhs, ctx) + .zip(min(rhs, ctx)) + .map(|(a, b)| with_comparison_guards::

(ctx, [lhs, rhs], lt(a, b))) + .transpose()? + .flatten(), + Operator::Lte => max(lhs, ctx) + .zip(min(rhs, ctx)) + .map(|(a, b)| with_comparison_guards::

(ctx, [lhs, rhs], lt_eq(a, b))) + .transpose()? + .flatten(), + Operator::And => { + if !P::EMIT_UNGUARDED_REWRITES { + return Ok(None); + } + + match (ctx.satisfy(lhs)?, ctx.satisfy(rhs)?) { + (Some(lhs), Some(rhs)) => Some(and(lhs, rhs)), + _ => None, + } + } + Operator::Or => { + if !P::EMIT_UNGUARDED_REWRITES { + return Ok(None); + } + + let lhs_satisfier = ctx.satisfy(lhs)?; + let rhs_satisfier = ctx.satisfy(rhs)?; + or_collect(lhs_satisfier.into_iter().chain(rhs_satisfier)) + } + Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None, + }) +} + #[derive(Debug)] struct BetweenStatsRewrite; @@ -190,6 +282,21 @@ impl StatsRewriteRule for BetweenStatsRewrite { let rhs = Binary.new_expr(options.upper_strict.to_operator(), [arr, upper]); ctx.falsify(&and(lhs, rhs)) } + + fn satisfy( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + let options = expr.as_::(); + let arr = expr.child(0).clone(); + let lower = expr.child(1).clone(); + let upper = expr.child(2).clone(); + + let lhs = Binary.new_expr(options.lower_strict.to_operator(), [lower, arr.clone()]); + let rhs = Binary.new_expr(options.upper_strict.to_operator(), [arr, upper]); + ctx.satisfy(&and(lhs, rhs)) + } } #[derive(Debug)] @@ -529,6 +636,12 @@ enum NanCheck { Unavailable, } +enum NullCheck { + NotNeeded, + Check(Expression), + Unavailable, +} + trait NonNanProof { const EMIT_UNGUARDED_REWRITES: bool; @@ -652,6 +765,42 @@ fn with_non_nan_guards<'a, P: NonNanProof>( }) } +fn with_comparison_guards<'a, P: NonNanProof>( + ctx: &StatsRewriteCtx<'_>, + exprs: impl IntoIterator, + value_predicate: Expression, +) -> VortexResult> { + let exprs = exprs.into_iter().collect::>(); + let mut null_checks = Vec::new(); + for expr in &exprs { + match null_check(ctx, expr)? { + NullCheck::NotNeeded => {} + NullCheck::Check(check) => null_checks.push(check), + NullCheck::Unavailable => return Ok(None), + } + } + + // `AllNonNull` is not a `Stat`, so FileStatsBinder cannot bind it. Keep this + // in `NullCount == 0` form so the same satisfier works for zone and file stats. + let value_predicate = if let Some(null_checks) = and_collect(null_checks) { + and(null_checks, value_predicate) + } else { + value_predicate + }; + with_non_nan_guards::

(ctx, exprs, value_predicate) +} + +fn null_check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + if !ctx.return_dtype(expr)?.is_nullable() { + return Ok(NullCheck::NotNeeded); + } + + Ok(match null_count(expr, ctx) { + Some(null_count) => NullCheck::Check(eq(null_count, lit(0u64))), + None => NullCheck::Unavailable, + }) +} + fn literal_stat(expr: &Expression, stat: Stat) -> Option { let scalar = expr.as_opt::()?; match stat { @@ -726,6 +875,7 @@ mod tests { use crate::expr::lit; use crate::expr::lt; use crate::expr::lt_eq; + use crate::expr::not_eq; use crate::expr::or; use crate::expr::stats::Stat; use crate::scalar::Scalar; @@ -757,6 +907,10 @@ mod tests { ("f", DType::Primitive(PType::F32, Nullability::NonNullable)), ("s", DType::Utf8(Nullability::NonNullable)), ("t", DType::Utf8(Nullability::NonNullable)), + ( + "nullable", + DType::Primitive(PType::I32, Nullability::Nullable), + ), ("n", nested_struct_dtype()), ]), Nullability::NonNullable, @@ -819,6 +973,33 @@ mod tests { Ok(()) } + #[test] + fn rewrites_comparison_satisfier() -> VortexResult<()> { + assert_eq!( + satisfy(>(col("a"), lit(10)))?, + Some(gt(stat(col("a"), Stat::Min), lit(10))) + ); + assert_eq!( + satisfy(<(col("a"), lit(10)))?, + Some(lt(stat(col("a"), Stat::Max), lit(10))) + ); + assert_eq!( + satisfy(&eq(col("a"), col("b")))?, + Some(and( + eq(stat(col("a"), Stat::Min), stat(col("b"), Stat::Max)), + eq(stat(col("a"), Stat::Max), stat(col("b"), Stat::Min)), + )) + ); + assert_eq!( + satisfy(¬_eq(col("a"), col("b")))?, + Some(or( + gt(stat(col("a"), Stat::Min), stat(col("b"), Stat::Max)), + gt(stat(col("b"), Stat::Min), stat(col("a"), Stat::Max)), + )) + ); + Ok(()) + } + #[test] fn rewrites_boolean_falsifiers() -> VortexResult<()> { let expr = and(gt(col("a"), lit(10)), lt(col("a"), lit(50))); @@ -832,6 +1013,29 @@ mod tests { Ok(()) } + #[test] + fn rewrites_boolean_satisfiers() -> VortexResult<()> { + let satisfiable = gt(col("a"), lit(10)); + let unavailable = like(col("s"), lit("prefix%")); + + assert_eq!( + satisfy(&and(satisfiable.clone(), unavailable.clone()))?, + None + ); + assert_eq!( + satisfy(&or(satisfiable.clone(), unavailable))?, + Some(gt(stat(col("a"), Stat::Min), lit(10))) + ); + assert_eq!( + satisfy(&and(satisfiable, lt(col("a"), lit(50))))?, + Some(and( + gt(stat(col("a"), Stat::Min), lit(10)), + lt(stat(col("a"), Stat::Max), lit(50)), + )) + ); + Ok(()) + } + #[test] fn rewrites_between_falsifier() -> VortexResult<()> { let expr = between( @@ -854,6 +1058,67 @@ mod tests { Ok(()) } + #[test] + fn rewrites_between_satisfier() -> VortexResult<()> { + let expr = between( + col("a"), + lit(10), + lit(50), + BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::NonStrict, + }, + ); + + assert_eq!( + satisfy(&expr)?, + Some(and( + lt_eq(lit(10), stat(col("a"), Stat::Min)), + lt_eq(stat(col("a"), Stat::Max), lit(50)), + )) + ); + Ok(()) + } + + #[test] + fn satisfier_guards_nullable_operand() -> VortexResult<()> { + assert_eq!( + satisfy(>(col("nullable"), lit(10)))?, + Some(and( + eq(stat(col("nullable"), Stat::NullCount), lit(0u64)), + gt(stat(col("nullable"), Stat::Min), lit(10)), + )) + ); + Ok(()) + } + + #[test] + fn satisfier_nan_guards_float_operand() -> VortexResult<()> { + assert_eq!( + satisfy(>(col("f"), lit(1.0f32)))?, + Some(nan_guarded( + col("f"), + gt(stat(col("f"), Stat::Min), lit(1.0f32)), + )) + ); + Ok(()) + } + + #[test] + fn no_satisfier_for_dynamic_and_like() -> VortexResult<()> { + let dynamic = dynamic( + CompareOperator::Gt, + || Some(10i32.into()), + DType::Primitive(PType::I32, Nullability::NonNullable), + true, + col("a"), + ); + + assert_eq!(satisfy(&dynamic)?, None); + assert_eq!(satisfy(&like(col("s"), lit("prefix%")))?, None); + Ok(()) + } + #[test] fn rewrites_null_falsifiers() -> VortexResult<()> { assert_eq!( diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index c9a71f85c55..30fa9eab1e3 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -27,6 +27,7 @@ use vortex_session::VortexSession; use crate::FileStatistics; use crate::footer::Footer; use crate::pruning::can_prune_file_stats; +use crate::pruning::can_satisfy_file_stats; use crate::v2::FileStatsLayoutReader; /// Represents a Vortex file, providing access to its metadata and content. @@ -206,6 +207,30 @@ impl VortexFile { ) } + /// Returns `true` if file-level statistics prove the expression matches + /// every row in this file. + /// + /// Row-count-aware satisfier predicates are evaluated with the file's total + /// row count as their scope. + pub fn can_match_all(&self, filter: &Expression) -> VortexResult { + let Some((stats, fields)) = self + .footer + .statistics() + .zip(self.footer.dtype().as_struct_fields_opt()) + else { + return Ok(false); + }; + + can_satisfy_file_stats( + filter, + self.footer.dtype(), + self.footer.row_count(), + stats, + fields, + &self.session, + ) + } + /// Return the file's natural row splits as root-coordinate ranges. /// /// These are the ranges that [`SplitBy::Layout`] would use for an all-fields scan. diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index 559df97d2d0..3bde611392d 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -34,7 +34,46 @@ pub(crate) fn can_prune_file_stats( struct_fields: &StructFields, session: &VortexSession, ) -> VortexResult { - let Some(pruning_expr) = expr.falsify(dtype, session)? else { + file_stats_proof( + expr, + dtype, + row_count, + file_stats, + struct_fields, + session, + Expression::falsify, + ) +} + +pub(crate) fn can_satisfy_file_stats( + expr: &Expression, + dtype: &DType, + row_count: u64, + file_stats: &FileStatistics, + struct_fields: &StructFields, + session: &VortexSession, +) -> VortexResult { + file_stats_proof( + expr, + dtype, + row_count, + file_stats, + struct_fields, + session, + Expression::satisfy, + ) +} + +fn file_stats_proof( + expr: &Expression, + dtype: &DType, + row_count: u64, + file_stats: &FileStatistics, + struct_fields: &StructFields, + session: &VortexSession, + rewrite: fn(&Expression, &DType, &VortexSession) -> VortexResult>, +) -> VortexResult { + let Some(proof_expr) = rewrite(expr, dtype, session)? else { return Ok(false); }; @@ -43,19 +82,19 @@ pub(crate) fn can_prune_file_stats( file_stats, struct_fields, }; - let pruning_expr = bind_stats(pruning_expr, &binder)?; + let proof_expr = bind_stats(proof_expr, &binder)?; - let simplified = pruning_expr.optimize_recursive(&DType::Null)?; + let simplified = proof_expr.optimize_recursive(&DType::Null)?; if let Some(result) = simplified.as_opt::() { return Ok(result.as_bool().value() == Some(true)); } - let pruning = NullArray::new(1).into_array().apply(&pruning_expr)?; - let row_count_replacement = ConstantArray::new(row_count, pruning.len()).into_array(); - let pruning = substitute_row_count(pruning, &row_count_replacement)?; + let proof = NullArray::new(1).into_array().apply(&proof_expr)?; + let row_count_replacement = ConstantArray::new(row_count, proof.len()).into_array(); + let proof = substitute_row_count(proof, &row_count_replacement)?; let mut ctx = session.create_execution_ctx(); - let result = pruning + let result = proof .execute::(&mut ctx)? .into_bool() .into_array() diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index e665ab18d50..678385f2178 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2282,6 +2282,11 @@ async fn test_can_prune_composite_predicates() -> VortexResult<()> { assert!(!file.can_prune(&eq(col("age"), lit(18)))?); assert!(!file.can_prune(&and(gt(col("age"), lit(20)), gt(col("price"), lit(100))))?); + // Satisfaction is the dual one-way proof: it only returns true when every row matches. + assert!(file.can_match_all(>(col("age"), lit(10)))?); + assert!(file.can_match_all(&and(gt(col("age"), lit(10)), gt(col("price"), lit(100)),))?); + assert!(!file.can_match_all(>(col("age"), lit(20)))?); + Ok(()) } diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index fc14eaf57d7..aa98df1f169 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -5,7 +5,9 @@ //! //! If file-level statistics prove that a filter expression cannot match any rows in the file, //! [`FileStatsLayoutReader`] short-circuits [`pruning_evaluation`](LayoutReader::pruning_evaluation) -//! by returning an all-false mask — avoiding all downstream I/O. +//! by returning an all-false mask. If they prove it matches every row, it short-circuits +//! [`filter_evaluation`](LayoutReader::filter_evaluation) by returning the input mask. Both avoid +//! downstream I/O. use std::ops::Range; use std::sync::Arc; @@ -27,14 +29,16 @@ use vortex_utils::aliases::dash_map::DashMap; use crate::FileStatistics; use crate::pruning::can_prune_file_stats; +use crate::pruning::can_satisfy_file_stats; -/// A [`LayoutReader`] decorator that prunes entire files based on file-level statistics. +/// A [`LayoutReader`] decorator that short-circuits file-wide filter proofs from file statistics. /// -/// This reader wraps an inner `LayoutReader` and intercepts `pruning_evaluation` calls. +/// This reader wraps an inner `LayoutReader` and intercepts pruning and filter evaluation calls. /// When file-level stats prove that a filter expression is false for the entire file, -/// it returns an all-false mask immediately — avoiding all downstream I/O. +/// it returns an all-false mask immediately. When they prove it true for every row, it returns +/// the filter's input mask. Both cases avoid all downstream I/O. /// -/// Pruning results are cached per-expression since file-level stats are global +/// File-level proof results are cached per expression since they are global /// (the result is the same regardless of which row range is requested). pub struct FileStatsLayoutReader { child: LayoutReaderRef, @@ -42,6 +46,7 @@ pub struct FileStatsLayoutReader { struct_fields: StructFields, session: VortexSession, prune_cache: DashMap, + satisfy_cache: DashMap, } impl FileStatsLayoutReader { @@ -64,6 +69,7 @@ impl FileStatsLayoutReader { struct_fields, session, prune_cache: Default::default(), + satisfy_cache: Default::default(), } } @@ -82,6 +88,18 @@ impl FileStatsLayoutReader { ) } + /// Evaluates whether file-level statistics prove `expr` matches every row. + fn evaluate_satisfy_file_stats(&self, expr: &Expression) -> VortexResult { + can_satisfy_file_stats( + expr, + self.child.dtype(), + self.child.row_count(), + &self.file_stats, + &self.struct_fields, + &self.session, + ) + } + /// Returns the file-level statistics used by this reader. pub fn file_stats(&self) -> &FileStatistics { &self.file_stats @@ -141,7 +159,21 @@ impl LayoutReader for FileStatsLayoutReader { expr: &Expression, mask: MaskFuture, ) -> VortexResult { - self.child.filter_evaluation(row_range, expr, mask) + if let Some(satisfied) = self.satisfy_cache.get(expr) { + if *satisfied { + return Ok(mask); + } + return self.child.filter_evaluation(row_range, expr, mask); + } + + let satisfied = self.evaluate_satisfy_file_stats(expr)?; + self.satisfy_cache.insert(expr.clone(), satisfied); + + if satisfied { + Ok(mask) + } else { + self.child.filter_evaluation(row_range, expr, mask) + } } fn projection_evaluation( @@ -162,8 +194,11 @@ impl LayoutReader for FileStatsLayoutReader { mod tests { use std::sync::Arc; use std::sync::LazyLock; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use vortex_array::ArrayContext; + use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -192,7 +227,10 @@ mod tests { use vortex_layout::LayoutStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::table::TableStrategy; + use vortex_layout::segments::SegmentFuture; + use vortex_layout::segments::SegmentId; use vortex_layout::segments::SegmentSink; + use vortex_layout::segments::SegmentSource; use vortex_layout::segments::TestSegments; use vortex_layout::sequence::SequenceId; use vortex_layout::sequence::SequentialArrayStreamExt; @@ -208,6 +246,18 @@ mod tests { .with::() }); + struct CountingSegmentSource { + inner: Arc, + request_count: Arc, + } + + impl SegmentSource for CountingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.request_count.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + } + fn test_file_stats(min: i32, max: i32) -> FileStatistics { let mut stats = StatsSet::default(); stats.set(Stat::Min, Precision::exact(ScalarValue::from(min))); @@ -230,6 +280,43 @@ mod tests { ) } + fn test_nullable_file_stats(min: i32, max: i32, null_count: u64) -> FileStatistics { + let mut stats = StatsSet::default(); + stats.set(Stat::Min, Precision::exact(ScalarValue::from(min))); + stats.set(Stat::Max, Precision::exact(ScalarValue::from(max))); + stats.set( + Stat::NullCount, + Precision::exact(ScalarValue::from(null_count)), + ); + FileStatistics::new( + Arc::from([stats]), + Arc::from([DType::Primitive(PType::I32, Nullability::Nullable)]), + ) + } + + async fn write_single_column_table( + session: &VortexSession, + segments: Arc, + values: ArrayRef, + ) -> VortexResult { + let ctx = ArrayContext::empty(); + let (ptr, eof) = SequenceId::root().split(); + let struct_array = StructArray::from_fields([("col", values)].as_slice())?; + let strategy = TableStrategy::new( + Arc::new(FlatLayoutStrategy::default()), + Arc::new(FlatLayoutStrategy::default()), + ); + strategy + .write_stream( + ctx, + segments, + struct_array.into_array().to_array_stream().sequenced(ptr), + eof, + session, + ) + .await + } + #[test] fn pruning_when_filter_out_of_range() -> VortexResult<()> { block_on(|handle| async { @@ -309,6 +396,86 @@ mod tests { }) } + #[test] + fn filter_skipped_when_file_satisfies() -> VortexResult<()> { + block_on(|handle| async { + let session = SESSION.clone().with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let layout = write_single_column_table( + &session, + Arc::clone(&segments), + buffer![60i32, 70, 80, 90, 100].into_array(), + ) + .await?; + let request_count = Arc::new(AtomicUsize::new(0)); + let source: Arc = Arc::new(CountingSegmentSource { + inner: segments, + request_count: Arc::clone(&request_count), + }); + let child = layout.new_reader("".into(), source, &session, &Default::default())?; + let reader = FileStatsLayoutReader::new(child, test_file_stats(60, 100), session); + let expr = gt(get_item("col", root()), lit(50i32)); + let input = Mask::from_iter([true, false, true, false, true]); + + let result = reader + .filter_evaluation(&(0..5), &expr, MaskFuture::ready(input.clone()))? + .await?; + + assert_eq!(result, input); + assert_eq!(request_count.load(Ordering::Relaxed), 0); + Ok(()) + }) + } + + #[test] + fn filter_not_skipped_when_partial() -> VortexResult<()> { + block_on(|handle| async { + let session = SESSION.clone().with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let layout = write_single_column_table( + &session, + Arc::clone(&segments), + buffer![0i32, 25, 50, 75, 100].into_array(), + ) + .await?; + let child = layout.new_reader("".into(), segments, &session, &Default::default())?; + let reader = FileStatsLayoutReader::new(child, test_file_stats(0, 100), session); + let expr = gt(get_item("col", root()), lit(50i32)); + + let result = reader + .filter_evaluation(&(0..5), &expr, MaskFuture::new_true(5))? + .await?; + + assert_eq!(result, Mask::from_iter([false, false, false, true, true])); + Ok(()) + }) + } + + #[test] + fn filter_nullable_not_satisfied() -> VortexResult<()> { + block_on(|handle| async { + let session = SESSION.clone().with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let layout = write_single_column_table( + &session, + Arc::clone(&segments), + PrimitiveArray::from_option_iter([None::, Some(60), Some(100)]).into_array(), + ) + .await?; + let child = layout.new_reader("".into(), segments, &session, &Default::default())?; + let reader = + FileStatsLayoutReader::new(child, test_nullable_file_stats(60, 100, 1), session); + let expr = gt(get_item("col", root()), lit(50i32)); + + let result = reader + .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? + .await?; + + assert_eq!(result, Mask::from_iter([false, true, true])); + Ok(()) + }) + } + #[test] fn no_pruning_for_computed_expression_stats() -> VortexResult<()> { block_on(|handle| async { diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index a9870229792..56382ce1873 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -269,9 +269,17 @@ impl LayoutReader for ChunkedReader { let mut chunk_evals = vec![]; for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) { + let chunk_mask = mask.slice(mask_range); + if let Some(Ok(chunk_mask)) = chunk_mask.clone().now_or_never() + && chunk_mask.all_false() + { + chunk_evals.push(MaskFuture::ready(chunk_mask)); + continue; + } + let chunk_reader = self.chunk_reader(chunk_idx)?; let chunk_eval = chunk_reader - .filter_evaluation(&chunk_range, expr, mask.slice(mask_range)) + .filter_evaluation(&chunk_range, expr, chunk_mask) .map_err(|err| { err.with_context(format!("While evaluating filter on chunk {chunk_idx}")) })?; @@ -346,6 +354,8 @@ impl LayoutReader for ChunkedReader { #[cfg(test)] mod test { use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use futures::stream; use rstest::fixture; @@ -359,10 +369,14 @@ mod test { use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; + use vortex_array::expr::gt; + use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; + use vortex_mask::Mask; use vortex_session::registry::ReadContext; use crate::IntoLayout; @@ -374,6 +388,7 @@ mod test { use crate::layouts::flat::FlatLayout; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::scan::split_by::SplitBy; + use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::TestSegments; @@ -382,6 +397,18 @@ mod test { use crate::sequence::SequentialStreamExt as _; use crate::test::SESSION; + struct CountingSegmentSource { + inner: Arc, + request_count: Arc, + } + + impl SegmentSource for CountingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.request_count.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + } + #[fixture] /// Create a chunked layout with three chunks of primitive arrays. fn chunked_layout() -> (Arc, LayoutRef) { @@ -489,4 +516,32 @@ mod test { assert_arrays_eq!(result, expected, &mut ctx); }) } + + #[rstest] + fn all_false_chunk_slice_skips_child( + #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + block_on(|handle| async { + let request_count = Arc::new(AtomicUsize::new(0)); + let source = Arc::new(CountingSegmentSource { + inner: segments, + request_count: Arc::clone(&request_count), + }); + let session = SESSION.clone().with_handle(handle); + let reader = layout.new_reader("".into(), source, &session, &Default::default())?; + let mask = Mask::from_iter([false, false, false, true, true, true, true, true, true]); + + let result = reader + .filter_evaluation( + &(0..layout.row_count()), + >(root(), lit(0i32)), + MaskFuture::ready(mask.clone()), + )? + .await?; + + assert_eq!(result, mask); + assert_eq!(request_count.load(Ordering::Relaxed), 2); + Ok(()) + }) + } } diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index b517df985c9..c2fed3f5935 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -35,6 +35,7 @@ use crate::layouts::zoned::zone_map::ZoneMap; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = Shared>>>; +pub(super) type SharedSatisfyMask = Shared>>; type PredicateCache = Arc>>; pub(super) struct PruningState { @@ -47,8 +48,10 @@ pub(super) struct PruningState { session: VortexSession, pruning_result: LazyLock>>, + satisfy_result: LazyLock>>, zone_map: OnceLock, pruning_predicates: LazyLock>>, + satisfy_predicates: LazyLock>>, } impl PruningState { @@ -67,8 +70,10 @@ impl PruningState { lazy_children, session, pruning_result: Default::default(), + satisfy_result: Default::default(), zone_map: Default::default(), pruning_predicates: Default::default(), + satisfy_predicates: Default::default(), } } @@ -116,6 +121,42 @@ impl PruningState { .clone() } + pub(super) fn satisfy_mask_future(&self, expr: Expression) -> Option { + if let Some(result) = self.satisfy_result.get(&expr) { + return result.value().clone(); + } + + self.satisfy_result + .entry(expr.clone()) + .or_insert_with(|| match self.satisfy_predicate(expr.clone()) { + None => { + trace!(%expr, "no satisfy predicate"); + None + } + Some(predicate) => { + trace!(%expr, ?predicate, "constructed satisfy predicate"); + let zone_map = self.zone_map(); + let session = self.session.clone(); + + Some( + async move { + let zone_map = zone_map.await?; + zone_map.prune(&predicate, &session).map_err(|err| { + err.with_context(format!( + "While evaluating satisfy predicate {} (derived from {})", + predicate, expr + )) + }) + } + .map_err(Arc::new) + .boxed() + .shared(), + ) + } + }) + .clone() + } + fn pruning_predicate(&self, expr: Expression) -> Option { self.pruning_predicates .entry(expr.clone()) @@ -130,6 +171,25 @@ impl PruningState { .clone() } + fn satisfy_predicate(&self, expr: Expression) -> Option { + if DynamicExprUpdates::new(&expr).is_some() { + trace!(%expr, "no satisfy predicate for dynamic expression"); + return None; + } + + self.satisfy_predicates + .entry(expr.clone()) + .or_default() + .get_or_init(move || match expr.satisfy(&self.dtype, &self.session) { + Ok(predicate) => predicate, + Err(error) => { + trace!(%expr, %error, "failed to construct stats rewrite predicate"); + None + } + }) + .clone() + } + fn zone_map(&self) -> SharedZoneMap { self.zone_map .get_or_init(move || { diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 162742b6f22..668858df650 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -2,9 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::BitAnd; +use std::ops::BitOr; use std::ops::Range; use std::sync::Arc; +use futures::FutureExt; use futures::future::BoxFuture; use itertools::Itertools; use tracing::trace; @@ -88,6 +90,41 @@ impl ZonedReader { .saturating_mul(self.layout.zone_len as u64) .min(self.layout.row_count()) } + + fn zone_row_lengths(&self, row_range: &Range) -> VortexResult<(Range, Vec)> { + let row_count = row_range.end - row_range.start; + let zone_range = self.zone_range(row_range); + let zone_lengths = zone_range + .clone() + .map(|zone_idx| { + let start = usize::try_from( + self.first_row_offset(zone_idx) + .saturating_sub(row_range.start), + )?; + let end = usize::try_from( + self.first_row_offset(zone_idx + 1) + .saturating_sub(row_range.start) + .min(row_count), + )?; + Ok::<_, VortexError>(end - start) + }) + .try_collect()?; + + Ok((zone_range, zone_lengths)) + } + + fn expand_zone_mask(&self, row_range: &Range, zone_mask: &Mask) -> VortexResult { + let (zone_range, zone_lengths) = self.zone_row_lengths(row_range)?; + let row_mask_len = usize::try_from(row_range.end - row_range.start)?; + let mut builder = BitBufferMut::with_capacity(row_mask_len); + for (zone_idx, &zone_length) in zone_range.zip_eq(&zone_lengths) { + builder.append_n(zone_mask.value(usize::try_from(zone_idx)?), zone_length); + } + + let row_mask = Mask::from(builder.freeze()); + assert_eq!(row_mask.len(), row_mask_len, "Mask length mismatch"); + Ok(row_mask) + } } impl LayoutReader for ZonedReader { @@ -138,24 +175,7 @@ impl LayoutReader for ZonedReader { return Ok(data_eval); }; - let row_count = row_range.end - row_range.start; - let zone_range = self.zone_range(row_range); - let zone_lengths: Vec<_> = zone_range - .clone() - .map(|zone_idx| { - // Figure out the range in the mask that corresponds to the zone - let start = usize::try_from( - self.first_row_offset(zone_idx) - .saturating_sub(row_range.start), - )?; - let end = usize::try_from( - self.first_row_offset(zone_idx + 1) - .saturating_sub(row_range.start) - .min(row_count), - )?; - Ok::<_, VortexError>(end - start) - }) - .try_collect()?; + let (zone_range, zone_lengths) = self.zone_row_lengths(row_range)?; let name = Arc::clone(&self.name); let expr = expr.clone(); @@ -201,7 +221,39 @@ impl LayoutReader for ZonedReader { expr: &Expression, mask: MaskFuture, ) -> VortexResult { - self.data_child()?.filter_evaluation(row_range, expr, mask) + if self.layout.zone_len == 0 { + return self.data_child()?.filter_evaluation(row_range, expr, mask); + } + + let Some(satisfy_mask_future) = self.pruning.satisfy_mask_future(expr.clone()) else { + return self.data_child()?.filter_evaluation(row_range, expr, mask); + }; + + // The zone map is normally resolved by pruning before filter evaluation. If it or the + // input mask is not ready, delegate to preserve eager child prefetching. + let Some(Ok(satisfied_zones)) = satisfy_mask_future.clone().now_or_never() else { + return self.data_child()?.filter_evaluation(row_range, expr, mask); + }; + let Some(Ok(input)) = mask.clone().now_or_never() else { + return self.data_child()?.filter_evaluation(row_range, expr, mask); + }; + + let satisfied_rows = self.expand_zone_mask(row_range, &satisfied_zones)?; + let child_input = input.clone().bitand_not(&satisfied_rows); + if child_input.all_false() { + return Ok(MaskFuture::ready(input)); + } + + let passthrough = input.bitand(&satisfied_rows); + let child_result = self.data_child()?.filter_evaluation( + row_range, + expr, + MaskFuture::ready(child_input), + )?; + + Ok(MaskFuture::new(passthrough.len(), async move { + Ok(child_result.await?.bitor(&passthrough)) + })) } fn projection_evaluation( @@ -221,6 +273,8 @@ impl LayoutReader for ZonedReader { mod test { use std::num::NonZeroUsize; use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use rstest::fixture; use rstest::rstest; @@ -262,12 +316,37 @@ mod test { use crate::layouts::zoned::Zoned; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::session::LayoutSession; + struct CountingSegmentSource { + inner: Arc, + request_count: Arc, + } + + impl SegmentSource for CountingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.request_count.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + } + + fn counting_source( + inner: Arc, + ) -> (Arc, Arc) { + let request_count = Arc::new(AtomicUsize::new(0)); + let source: Arc = Arc::new(CountingSegmentSource { + inner, + request_count: Arc::clone(&request_count), + }); + (source, request_count) + } + fn session_with_handle(handle: Handle) -> VortexSession { array_session() .with::() @@ -308,6 +387,37 @@ mod test { (segments, layout) } + #[fixture] + fn nullable_stats_layout() -> (Arc, LayoutRef) { + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let strategy = ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(3).vortex_expect("non zero"), + ..Default::default() + }, + ); + let array_stream = PrimitiveArray::new( + buffer![0i32, 4, 5], + Validity::from_iter([false, true, true]), + ) + .into_array() + .to_array_stream() + .sequenced(ptr); + let segments2 = Arc::::clone(&segments); + let layout = block_on(|handle| async move { + let session = session_with_handle(handle); + strategy + .write_stream(ctx, segments2, array_stream, eof, &session) + .await + }) + .unwrap(); + (segments, layout) + } + #[rstest] fn test_stats_evaluator( #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), @@ -363,6 +473,129 @@ mod test { }) } + #[rstest] + fn satisfied_zones_skip_data_child( + #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + block_on(|handle| async { + let row_count = layout.row_count(); + let (source, request_count) = counting_source(segments); + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), source, &session, &Default::default())?; + let input = Mask::new_true(usize::try_from(row_count)?); + + reader + .pruning_evaluation(&(0..row_count), >(root(), lit(10)), input.clone())? + .await?; + let requests_before_filter = request_count.load(Ordering::Relaxed); + + let result = reader + .filter_evaluation( + &(0..row_count), + >(root(), lit(0)), + MaskFuture::ready(input.clone()), + )? + .await?; + + assert_eq!(result, input); + assert_eq!( + request_count.load(Ordering::Relaxed), + requests_before_filter, + "a fully satisfied filter must not read the data child" + ); + Ok(()) + }) + } + + #[rstest] + fn partial_satisfaction_stitches( + #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + block_on(|handle| async { + let row_count = layout.row_count(); + let (source, request_count) = counting_source(segments); + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), source, &session, &Default::default())?; + let input = Mask::new_true(usize::try_from(row_count)?); + + reader + .pruning_evaluation(&(0..row_count), >(root(), lit(10)), input.clone())? + .await?; + let requests_before_filter = request_count.load(Ordering::Relaxed); + + let result = reader + .filter_evaluation( + &(0..row_count), + >(root(), lit(3)), + MaskFuture::ready(input), + )? + .await?; + + assert_eq!( + result, + Mask::from_iter([false, false, false, true, true, true, true, true, true]) + ); + assert_eq!( + request_count.load(Ordering::Relaxed), + requests_before_filter + 1, + "only the unsatisfied zone should read a data segment" + ); + Ok(()) + }) + } + + #[rstest] + fn nullable_zone_with_nulls_not_skipped( + #[from(nullable_stats_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + block_on(|handle| async { + let row_count = layout.row_count(); + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; + let input = Mask::new_true(usize::try_from(row_count)?); + + reader + .pruning_evaluation(&(0..row_count), >(root(), lit(10)), input.clone())? + .await?; + + let result = reader + .filter_evaluation( + &(0..row_count), + >(root(), lit(3)), + MaskFuture::ready(input), + )? + .await?; + + assert_eq!(result, Mask::from_iter([false, true, true])); + Ok(()) + }) + } + + #[rstest] + fn unwarmed_zone_map_keeps_filter_correct( + #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + block_on(|handle| async { + let row_count = layout.row_count(); + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; + + let result = reader + .filter_evaluation( + &(0..row_count), + >(root(), lit(3)), + MaskFuture::new_true(usize::try_from(row_count)?), + )? + .await?; + + assert_eq!( + result, + Mask::from_iter([false, false, false, true, true, true, true, true, true]) + ); + Ok(()) + }) + } + #[test] fn test_default_zoned_null_count_pruning_mask() { let ctx = ArrayContext::empty(); @@ -428,7 +661,7 @@ mod test { } #[rstest] - fn test_legacy_zero_zone_len_skips_zoned_pruning( + fn legacy_zero_zone_len_delegates_pruning_and_filter( #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), ) -> VortexResult<()> { let zoned_layout = layout.as_::(); @@ -469,6 +702,18 @@ mod test { .await?; assert!(result.all_true()); + + let result = reader + .filter_evaluation( + &(0..row_count), + >(root(), lit(7)), + MaskFuture::new_true(usize::try_from(row_count)?), + )? + .await?; + assert_eq!( + result, + Mask::from_iter([false, false, false, false, false, false, false, true, true]) + ); Ok(()) }) }