diff --git a/Cargo.lock b/Cargo.lock index a1ece6159e3..a701966688a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10023,6 +10023,7 @@ dependencies = [ "async-stream", "async-trait", "bit-vec", + "codspeed-criterion-compat-walltime", "flatbuffers", "futures", "insta", diff --git a/docs/developer-guide/index.md b/docs/developer-guide/index.md index 0bc908693d5..9afff8877aa 100644 --- a/docs/developer-guide/index.md +++ b/docs/developer-guide/index.md @@ -23,6 +23,7 @@ internals/session internals/async-runtime internals/vtables internals/execution +internals/scan-planning internals/stats-pruning internals/io internals/serialization diff --git a/docs/developer-guide/internals/scan-planning.md b/docs/developer-guide/internals/scan-planning.md new file mode 100644 index 00000000000..2fa352c733d --- /dev/null +++ b/docs/developer-guide/internals/scan-planning.md @@ -0,0 +1,193 @@ +# Expression Pushdown in Scan Plans + +:::{note} +This is a provisional design, not a description of the current scan planner. +::: + +## Expression scan plan + +Add a physical plan that applies an expression to the output of another scan plan: + +```rust +pub struct ExpressionScanPlan { + expression: Expression, + child: ScanPlanRef, +} +``` + +`ExpressionScanPlan` inherits its row domain from `child` and derives its output dtype from +`expression`. Any part of the expression that cannot be pushed down remains in this plan and is +evaluated over the child's result. + +Scalar functions remain nodes in `Expression`; they do not each become physical scan-plan nodes. +Instead, pushdown behavior is supplied by pluggable kernels registered for a scalar function (or +other expression node) and a concrete scan-plan type. + +Composite scan plans expose ordered logical children to this optimizer: + +```text +StructScanPlan.children = [field(0), ..., field(n - 1), validity?] +DictScanPlan.children = [codes, values] +``` + +## Pushdown + +Optimizing `ExpressionScanPlan(expression, child)` has three phases. + +1. **Annotate dependencies.** Walk the expression and annotate every node with the indices of the + immediate scan-plan children needed to evaluate it. A registered + `(expression node, scan plan) -> [child index]` kernel provides plan-specific dependencies; + ordinary scalar functions can otherwise take the union of their expression children's + annotations. For a struct, `get_item($, "a")` needs field `a` and, when present, the struct + validity child. `is_not_null($)` and `is_null($)` need only the top-level validity child. + +2. **Partition and group.** Only a subexpression annotated with exactly one scan-plan child is + eligible for pushdown. Cut maximal eligible subexpressions, group all cuts for the same child, + and build one packed expression for that child. Replace the cuts in the remaining expression + with references to the group results; this remaining expression is the combination expression. + Expressions needing zero or multiple children stay above the current plan. This follows the + grouping model of the existing expression partitioner and ensures each child group is evaluated + once. + +3. **Lower into each child.** Rewrite every root reference in a group from the current plan's scope + into the selected child's scope. This uses a second pluggable + `(expression node, scan plan, child index) -> expression` kernel. For a non-nullable struct, + `get_item($, "a")` becomes `$` when lowering into field `a`; targeting another field rejects that + pushdown. When lowering into a struct's validity child, `is_not_null($)` becomes `$` and + `is_null($)` becomes `not($)`. The lowered group is installed as an `ExpressionScanPlan` over + that child, and pushdown then continues recursively. Every rule must prove equivalence across + the plan boundary; lowering is not a generic replacement of `$`. + +The combination expression is retained above the grouped child plans. Missing kernels or failed +lowering leave the affected expression at the current level, preserving the generic execution +fallback. + +## Examples + +`@name` denotes a reference from the combination expression to a grouped child result. + +### Struct without validity + +```text +plan = StructScanPlan.children = [a, b] +expr = (get_item($, "a") + 1) * get_item($, "b") + +annotations: + get_item($, "a") + 1 -> {a} + get_item($, "b") -> {b} + expr -> {a, b} + +groups before lowering: + @a = get_item($, "a") + 1 + @b = get_item($, "b") + +groups after lowering: + @a = $ + 1 + @b = $ + +combine = @a * @b +``` + +The current non-nullable struct shape can therefore install `$ + 1` over child `a`, read child +`b` directly, and evaluate only the multiplication above the grouped results. The `+` needs no +struct-specific rule: it is rebuilt after its `get_item` child is lowered. + +### Struct with validity + +```text +plan = StructScanPlan.children = [a, b, validity] +expr = is_not_null($) && (get_item($, "a") > 0) + +annotations: + is_not_null($) -> {validity} + get_item($, "a") -> {a, validity} + get_item($, "a") > 0 -> {a, validity} + expr -> {a, validity} + +groups before lowering: + @valid = is_not_null($) + +groups after lowering: + @valid = $ + +combine = @valid && (get_item($, "a") > 0) +``` + +The direct `is_not_null($) -> $` lowering is valid because `$` on the validity child is exactly the +top-level struct validity. The analogous direct lowering of `get_item` is not valid: + +```text +get_item(nullable_struct, "a") = mask(a, validity) + +invalid lowering: + get_item($, "a") -> $ + +counterexample: + a = 7, validity = false + get_item(nullable_struct, "a") = null + lowered result = 7 +``` + +The singleton rule prevents this mistake: `get_item($, "a")` is annotated with `{a, validity}` and +therefore remains above the struct. A stronger validity-aware rule may lower a strict function by +factoring out the mask: + +```text +(get_item($, "a") + 1) + = mask(a, validity) + 1 + = mask(a + 1, validity) + +groups after lowering: + @a = $ + 1 + @valid = $ + +combine = mask(@a, @valid) +``` + +This requires `+` to be strict. It must also be infallible, or execute under `@valid`, so that +evaluating values hidden by the mask cannot introduce a new error. Non-strict functions require +different combination expressions: + +```text +is_not_null(get_item($, "a")) + = @valid && is_not_null(@a) + +is_null(get_item($, "a")) + = !@valid || is_null(@a) +``` + +Lowering through validity is therefore a proven factorization into child expressions and a +residual combination expression, not just scope substitution. + +### Dictionary + +```text +plan = DictScanPlan.children = [codes, values] +expr = byte_length($) + +annotations: + byte_length($) -> {values} + +groups before lowering: + @values = byte_length($) + +groups after lowering: + @values = byte_length($) + +combine = $ + +result = DictScanPlan.children = [ + codes, + ExpressionScanPlan(byte_length($), values), +] +``` + +This is valid for the same strict, infallible, negative-cost functions accepted by the current +dictionary pushdown. The dictionary plan reuses `codes` and applies the function once to +`values`. + +## Future work + +The same optimizer may eventually support `ScanPlan x ScanPlan -> ScanPlan` transforms. Those +plan-to-plan rewrites are outside the scope of this proposal; this design covers only expression +and scalar-function pushdown through scan plans. diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 25f9003fd31..1d1f5c74546 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -119,7 +119,6 @@ pub static ALLOWED_ENCODINGS: LazyLock> = LazyLock::new(|| { allowed.insert(ZigZag.id()); // Experimental encodings - if use_experimental_patches() { allowed.insert(Patched.id()); } diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..05507c7f7ee 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -55,6 +55,7 @@ vortex-session = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] +criterion = { workspace = true } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } rstest = { workspace = true } @@ -63,6 +64,10 @@ tokio = { workspace = true, features = ["rt", "macros"] } vortex-array = { path = "../vortex-array", features = ["_test-harness"] } vortex-io = { path = "../vortex-io", features = ["tokio"] } +[[bench]] +name = "scan_plan" +harness = false + [features] _test-harness = [] tokio = ["dep:tokio", "vortex-error/tokio"] diff --git a/vortex-layout/benches/scan_plan.rs b/vortex-layout/benches/scan_plan.rs new file mode 100644 index 00000000000..416bfd3692e --- /dev/null +++ b/vortex-layout/benches/scan_plan.rs @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use criterion::BenchmarkId; +use criterion::Criterion; +use criterion::black_box; +use criterion::criterion_group; +use criterion::criterion_main; +use futures::FutureExt; +use vortex_array::MaskFuture; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; +use vortex_array::expr::get_item; +use vortex_array::expr::root; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_layout::LayoutChildren; +use vortex_layout::LayoutParts; +use vortex_layout::LayoutReaderContext; +use vortex_layout::LayoutReaderRef; +use vortex_layout::LayoutRef; +use vortex_layout::layouts::flat::FlatLayout; +use vortex_layout::layouts::struct_::Struct; +use vortex_layout::layouts::struct_::StructLayout; +use vortex_layout::scan::plan_v2::LayoutReaderScanPlanV2; +use vortex_layout::scan::plan_v2::ScanPlanRef; +use vortex_layout::scan::plan_v2::StructScanPlan; +use vortex_layout::segments::SegmentFuture; +use vortex_layout::segments::SegmentId; +use vortex_layout::segments::SegmentSource; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +#[derive(Clone)] +struct CountingLayoutChildren { + children: Arc<[LayoutRef]>, + materializations: Arc, +} + +impl LayoutChildren for CountingLayoutChildren { + fn to_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn child(&self, idx: usize, dtype: &DType) -> VortexResult { + let child = self + .children + .get(idx) + .vortex_expect("benchmark child index must be valid"); + vortex_ensure!( + child.dtype() == dtype, + "benchmark child dtype mismatch: {} != {dtype}", + child.dtype() + ); + self.materializations.fetch_add(1, Ordering::Relaxed); + Ok(Arc::clone(child)) + } + + fn child_row_count(&self, idx: usize) -> u64 { + self.children + .get(idx) + .vortex_expect("benchmark child index must be valid") + .row_count() + } + + fn nchildren(&self) -> usize { + self.children.len() + } +} + +struct NoSegments; + +impl SegmentSource for NoSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + async move { vortex_bail!("benchmark must not poll segment {id}") }.boxed() + } +} + +struct ColdGetItemFixture { + layout: StructLayout, + segment_source: Arc, + session: VortexSession, + reader_context: LayoutReaderContext, + get_item: Expression, + child_materializations: Arc, +} + +impl ColdGetItemFixture { + fn new(width: usize) -> Self { + let child_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let fields = StructFields::from_iter( + (0..width).map(|idx| (format!("field_{idx}"), child_dtype.clone())), + ); + let dtype = DType::Struct(fields, Nullability::NonNullable); + let children = (0..width) + .map(|idx| { + FlatLayout::new( + 1, + child_dtype.clone(), + SegmentId::try_from(idx).vortex_expect("benchmark width must fit in SegmentId"), + ReadContext::new([]), + ) + .into_layout() + }) + .collect::>(); + let child_materializations = Arc::new(AtomicUsize::new(0)); + let counting_children = CountingLayoutChildren { + children: children.into(), + materializations: Arc::clone(&child_materializations), + }; + let layout = LayoutParts::new( + Struct, + dtype, + 1, + Vec::new(), + Arc::new(counting_children), + (), + ) + .into_typed(); + + Self { + layout, + segment_source: Arc::new(NoSegments), + session: VortexSession::empty(), + reader_context: LayoutReaderContext::default(), + get_item: get_item(format!("field_{}", width - 1), root()), + child_materializations, + } + } + + fn fresh_reader(&self) -> LayoutReaderRef { + self.layout + .new_reader( + Arc::from("benchmark"), + Arc::clone(&self.segment_source), + &self.session, + &self.reader_context, + ) + .vortex_expect("benchmark struct reader must be constructed") + } + + fn child_materializations(&self) -> usize { + self.child_materializations.load(Ordering::Relaxed) + } + + // A fresh reader gives every timed iteration a new LazyReaderChildren cache. The counter + // deltas below prove that construction is lazy and first access materializes one child. + fn layout_reader(&self) { + let before = self.child_materializations(); + let reader = self.fresh_reader(); + assert_eq!( + self.child_materializations(), + before, + "constructing a fresh StructReader must not materialize children" + ); + + let future = reader + .projection_evaluation( + &(0..1), + &self.get_item, + MaskFuture::ready(Mask::new_true(1)), + ) + .vortex_expect("benchmark projection must be planned"); + assert_eq!( + self.child_materializations(), + before + 1, + "a cold GetItem must materialize exactly one child reader" + ); + drop(black_box(future)); + } + + fn struct_scan_plan(&self) { + let before = self.child_materializations(); + let reader = self.fresh_reader(); + let source: ScanPlanRef = Arc::new(LayoutReaderScanPlanV2::new(reader)); + let struct_plan: ScanPlanRef = Arc::new( + StructScanPlan::try_new(source) + .vortex_expect("benchmark source must construct a struct scan plan"), + ); + let field = Arc::clone(&struct_plan) + .apply_expr(self.get_item.clone()) + .vortex_expect("benchmark GetItem must reduce") + .optimize() + .vortex_expect("benchmark field plan must optimize"); + assert_eq!( + self.child_materializations(), + before, + "constructing and reducing a StructScanPlan must not materialize children" + ); + + let future = field + .projection_evaluation(&(0..1), MaskFuture::ready(Mask::new_true(1))) + .vortex_expect("benchmark projection must be planned"); + assert_eq!( + self.child_materializations(), + before + 1, + "a cold reduced GetItem must materialize exactly one child reader" + ); + drop(black_box(future)); + } +} + +fn bench_cold_get_item(c: &mut Criterion) { + let mut group = c.benchmark_group("cold_get_item"); + for width in [10, 100, 1_000] { + let fixture = ColdGetItemFixture::new(width); + group.bench_with_input(BenchmarkId::new("layout_reader", width), &width, |b, _| { + b.iter(|| fixture.layout_reader()); + }); + group.bench_with_input( + BenchmarkId::new("struct_scan_plan", width), + &width, + |b, _| { + b.iter(|| fixture.struct_scan_plan()); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_cold_get_item); +criterion_main!(benches); diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index 98fd1918a42..6b2c07a5b83 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -5,11 +5,12 @@ pub mod arrow; mod filter; pub mod layout; pub mod multi; +mod plan; +pub mod plan_v2; pub mod repeated_scan; pub mod scan_builder; pub mod split_by; mod splits; -mod tasks; #[cfg(test)] mod test; diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/plan.rs similarity index 85% rename from vortex-layout/src/scan/tasks.rs rename to vortex-layout/src/scan/plan.rs index a86546e15ef..439cbee1be3 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/plan.rs @@ -16,11 +16,46 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; -use crate::LayoutReader; +use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; pub type TaskFuture = BoxFuture<'static, VortexResult>; +pub(crate) struct Plan { + layout_reader: LayoutReaderRef, + projection: Expression, + filter: Option, +} + +impl Plan { + pub(crate) fn new( + layout_reader: LayoutReaderRef, + projection: Expression, + filter: Option, + ) -> Self { + Self { + layout_reader, + projection, + filter, + } + } + + pub(crate) fn task_context( + &self, + mapper: Arc VortexResult + Send + Sync>, + ) -> Arc> { + Arc::new(TaskContext { + filter: self + .filter + .clone() + .map(|filter| Arc::new(FilterExpr::new(filter))), + reader: Arc::clone(&self.layout_reader), + projection: self.projection.clone(), + mapper, + }) + } +} + /// Logic for executing a single split reading task. /// N.B. read_mask should be evaluated against all_false() before calling this /// method to avoid creating an empty TaskFuture. @@ -153,13 +188,9 @@ pub fn split_exec( /// Information needed to execute a single split task. /// /// Row selection is evaluated before creating a split task so it's not included -pub struct TaskContext { - /// The shared filter expression. - pub filter: Option>, - /// The layout reader. - pub reader: Arc, - /// The projection expression to apply to gather the scanned rows. - pub projection: Expression, - /// Function that maps into an A. - pub mapper: Arc VortexResult + Send + Sync>, +pub(crate) struct TaskContext { + filter: Option>, + reader: LayoutReaderRef, + projection: Expression, + mapper: Arc VortexResult + Send + Sync>, } diff --git a/vortex-layout/src/scan/plan_v2.rs b/vortex-layout/src/scan/plan_v2.rs new file mode 100644 index 00000000000..b5751e31115 --- /dev/null +++ b/vortex-layout/src/scan/plan_v2.rs @@ -0,0 +1,751 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +// See https://github.com/vortex-data/vortex/issues/9062 + +use std::any::Any; +use std::ops::BitAnd; +use std::ops::Range; +use std::sync::Arc; + +use bit_vec::BitVec; +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex_array::ArrayRef; +use vortex_array::MaskFuture; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::Expression; +use vortex_array::expr::get_item; +use vortex_array::expr::root; +use vortex_array::expr::transform::replace; +use vortex_array::scalar_fn::ReduceCtx; +use vortex_array::scalar_fn::ReduceNode; +use vortex_array::scalar_fn::ReduceNodeRef; +use vortex_array::scalar_fn::ScalarFnRef; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::get_item::GetItem; +use vortex_array::scalar_fn::fns::pack::Pack; +use vortex_array::scalar_fn::fns::pack::PackOptions; +use vortex_array::scalar_fn::fns::root::Root; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; +use vortex_scan::row_mask::RowMask; + +use crate::ArrayFuture; +use crate::LayoutReaderRef; +use crate::scan::filter::FilterExpr; + +pub(crate) struct PlanV2 { + projection: ScanPlanRef, + predicates: Vec, + filter: Option, +} + +impl PlanV2 { + pub(crate) fn new( + projection: ScanPlanRef, + predicates: Vec, + filter: Option, + ) -> Self { + Self { + projection, + predicates, + filter, + } + } + + pub(crate) fn task_context( + &self, + mapper: Arc VortexResult + Send + Sync>, + ) -> Arc> { + Arc::new(TaskContext { + filter: self + .filter + .clone() + .map(|filter| Arc::new(FilterExpr::new(filter))), + predicates: self.predicates.clone(), + projection: Arc::clone(&self.projection), + mapper, + }) + } +} + +/// Shared handle to a heap-allocated V2 physical scan plan. +pub type ScanPlanRef = Arc; + +/// A heap-allocated physical scan plan. +/// +/// A source plan represents an instantiated layout. [`apply_expr`](Self::apply_expr) derives +/// another plan whose root value is the applied expression, and [`optimize`](Self::optimize) +/// rewrites that derived plan before execution. Execution therefore selects an already-bound plan +/// and supplies only its row range and mask. +pub trait ScanPlan: 'static + Send + Sync { + /// Return this plan as [`Any`] for plan-specific optimization rules. + fn as_any(&self) -> &dyn Any; + + /// Apply `expr` to this plan's root value and return the resulting plan. + fn apply_expr(self: Arc, expr: Expression) -> VortexResult; + + /// Optimize this plan and return the resulting plan. + fn optimize(self: Arc) -> VortexResult; + + /// Returns the name of the underlying layout reader for debugging. + fn name(&self) -> &Arc; + + /// Returns the dtype produced by this plan. + fn dtype(&self) -> &DType; + + /// Returns the number of rows in this plan's row domain. + fn row_count(&self) -> u64; + + /// Returns a mask where all false values are proven false for this plan. + fn pruning_evaluation(&self, row_range: &Range, mask: Mask) -> VortexResult; + + /// Evaluates this boolean plan and intersects it with `mask`. + fn filter_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult; + + /// Evaluates this plan over the selected rows. + fn projection_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult; +} + +/// A non-nullable struct plan whose fields can be selected independently. +/// +/// The struct presents itself to scalar-function reduction as a [`Pack`] over its field plans. +/// This lets rules such as `GetItem(Pack(...))` reduce directly to a field plan while the +/// compatibility source remains available for expressions that cannot yet be reduced physically. +/// Field plans are created only when a reduction selects them, so planning one field does not +/// construct and discard plans for its siblings. +pub struct StructScanPlan { + source: ScanPlanRef, + pack: ScalarFnRef, +} + +impl StructScanPlan { + /// Wrap a non-nullable struct source with independently selectable field plans. + pub fn try_new(source: ScanPlanRef) -> VortexResult { + let DType::Struct(struct_fields, Nullability::NonNullable) = source.dtype() else { + vortex_bail!( + "StructScanPlan requires a non-nullable struct source, got {}", + source.dtype() + ); + }; + + let names = struct_fields.names().clone(); + let pack = Pack.bind(PackOptions { + names, + nullability: Nullability::NonNullable, + }); + + Ok(Self { source, pack }) + } + + fn field(&self, idx: usize) -> ScanPlanRef { + let struct_fields = self + .source + .dtype() + .as_struct_fields_opt() + .vortex_expect("StructScanPlan source must have a struct dtype"); + let name = struct_fields + .field_name(idx) + .vortex_expect("Pack child must have a matching struct field") + .clone(); + let dtype = struct_fields + .field_by_index(idx) + .vortex_expect("Pack child must have a matching struct field"); + + Arc::new(StructFieldScanPlan { + source: Arc::clone(&self.source), + name, + dtype, + }) + } +} + +impl ScanPlan for StructScanPlan { + fn as_any(&self) -> &dyn Any { + self + } + + fn apply_expr(self: Arc, expr: Expression) -> VortexResult { + if expr.is::() && expr.child(0).is::() { + let root = Arc::clone(&self); + let root: ScanPlanRef = root; + if let Some(reduced) = reduce_expr(&root, &expr)? { + return Ok(reduced); + } + } + + Arc::clone(&self.source).apply_expr(expr) + } + + fn optimize(self: Arc) -> VortexResult { + let source = Arc::clone(&self.source).optimize()?; + Ok(Arc::new(Self::try_new(source)?)) + } + + fn name(&self) -> &Arc { + self.source.name() + } + + fn dtype(&self) -> &DType { + self.source.dtype() + } + + fn row_count(&self) -> u64 { + self.source.row_count() + } + + fn pruning_evaluation(&self, row_range: &Range, mask: Mask) -> VortexResult { + self.source.pruning_evaluation(row_range, mask) + } + + fn filter_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.source.filter_evaluation(row_range, mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.source.projection_evaluation(row_range, mask) + } +} + +struct StructFieldScanPlan { + source: ScanPlanRef, + name: vortex_array::dtype::FieldName, + dtype: DType, +} + +impl StructFieldScanPlan { + fn expr(&self) -> Expression { + get_item(self.name.clone(), root()) + } + + fn applied(&self) -> VortexResult { + Arc::clone(&self.source).apply_expr(self.expr()) + } +} + +impl ScanPlan for StructFieldScanPlan { + fn as_any(&self) -> &dyn Any { + self + } + + fn apply_expr(self: Arc, expr: Expression) -> VortexResult { + let expr = replace(expr, &root(), self.expr()); + Arc::clone(&self.source).apply_expr(expr) + } + + fn optimize(self: Arc) -> VortexResult { + self.applied()?.optimize() + } + + fn name(&self) -> &Arc { + self.source.name() + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.source.row_count() + } + + fn pruning_evaluation(&self, row_range: &Range, mask: Mask) -> VortexResult { + self.applied()?.pruning_evaluation(row_range, mask) + } + + fn filter_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.applied()?.filter_evaluation(row_range, mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.applied()?.projection_evaluation(row_range, mask) + } +} + +struct ScanPlanReduceNode { + plan: ScanPlanRef, +} + +impl ReduceNode for ScanPlanReduceNode { + fn as_any(&self) -> &dyn Any { + self + } + + fn node_dtype(&self) -> VortexResult { + Ok(self.plan.dtype().clone()) + } + + fn scalar_fn(&self) -> Option<&ScalarFnRef> { + self.plan + .as_any() + .downcast_ref::() + .map(|plan| &plan.pack) + } + + fn child(&self, idx: usize) -> ReduceNodeRef { + let plan = self + .plan + .as_any() + .downcast_ref::() + .vortex_expect("only StructScanPlan nodes expose reduction children"); + Arc::new(Self { + plan: plan.field(idx), + }) + } + + fn child_count(&self) -> usize { + self.plan + .as_any() + .downcast_ref::() + .map_or(0, |plan| plan.source.dtype().as_struct_fields().nfields()) + } +} + +struct AppliedExprReduceNode { + scalar_fn: ScalarFnRef, + dtype: DType, + child: ReduceNodeRef, +} + +impl ReduceNode for AppliedExprReduceNode { + fn as_any(&self) -> &dyn Any { + self + } + + fn node_dtype(&self) -> VortexResult { + Ok(self.dtype.clone()) + } + + fn scalar_fn(&self) -> Option<&ScalarFnRef> { + Some(&self.scalar_fn) + } + + fn child(&self, idx: usize) -> ReduceNodeRef { + (idx == 0) + .then(|| Arc::clone(&self.child)) + .vortex_expect("GetItem reduction has exactly one child") + } + + fn child_count(&self) -> usize { + 1 + } +} + +struct ScanPlanReduceCtx; + +impl ReduceCtx for ScanPlanReduceCtx { + fn new_node( + &self, + _scalar_fn: ScalarFnRef, + _children: &[ReduceNodeRef], + ) -> VortexResult { + vortex_bail!("scan-plan reduction cannot yet create scalar-function plans") + } +} + +fn reduce_expr(plan: &ScanPlanRef, expr: &Expression) -> VortexResult> { + let root_node: ReduceNodeRef = Arc::new(ScanPlanReduceNode { + plan: Arc::clone(plan), + }); + let node = AppliedExprReduceNode { + scalar_fn: expr.scalar_fn().clone(), + dtype: expr.return_dtype(plan.dtype())?, + child: root_node, + }; + + let Some(reduced) = expr.scalar_fn().reduce(&node, &ScanPlanReduceCtx)? else { + return Ok(None); + }; + let Some(reduced) = reduced.as_any().downcast_ref::() else { + vortex_bail!("scan-plan reduction returned a non-plan node") + }; + + Ok(Some(Arc::clone(&reduced.plan))) +} + +/// Compatibility V2 source and expression plan backed by a layout reader. +/// +/// Applying an expression and optimizing it produce new heap-allocated plans. Execution delegates +/// the resulting expression to the established reader implementation. Layout-specific source plans +/// can replace this compatibility node without changing the split execution loop. +pub struct LayoutReaderScanPlanV2 { + reader: LayoutReaderRef, + expr: Expression, + dtype: DType, +} + +impl LayoutReaderScanPlanV2 { + /// Create a V2 source plan for `reader`. + pub fn new(reader: LayoutReaderRef) -> Self { + let dtype = reader.dtype().clone(); + Self { + reader, + expr: root(), + dtype, + } + } + + fn try_new(reader: LayoutReaderRef, expr: Expression) -> VortexResult { + let dtype = expr.return_dtype(reader.dtype())?; + Ok(Self { + reader, + expr, + dtype, + }) + } +} + +impl ScanPlan for LayoutReaderScanPlanV2 { + fn as_any(&self) -> &dyn Any { + self + } + + fn apply_expr(self: Arc, expr: Expression) -> VortexResult { + let expr = replace(expr, &root(), self.expr.clone()); + Ok(Arc::new(Self::try_new(Arc::clone(&self.reader), expr)?)) + } + + fn optimize(self: Arc) -> VortexResult { + let expr = self.expr.optimize_recursive(self.reader.dtype())?; + Ok(Arc::new(Self::try_new(Arc::clone(&self.reader), expr)?)) + } + + fn name(&self) -> &Arc { + self.reader.name() + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.reader.row_count() + } + + fn pruning_evaluation(&self, row_range: &Range, mask: Mask) -> VortexResult { + self.reader.pruning_evaluation(row_range, &self.expr, mask) + } + + fn filter_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.reader.filter_evaluation(row_range, &self.expr, mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.reader + .projection_evaluation(row_range, &self.expr, mask) + } +} + +/// Environment variable selecting the scan planning implementation. +pub const SCAN_IMPL_ENV: &str = "VORTEX_SCAN_IMPL"; + +/// Returns whether V2 heap-allocated planning is enabled for this process. +/// +/// The existing `plan` path remains the default on this extraction branch. Set +/// `VORTEX_SCAN_IMPL=planv2` to exercise the V2 path with the same execution implementation. +pub fn plan_v2_enabled() -> VortexResult { + match std::env::var(SCAN_IMPL_ENV) { + Ok(value) => parse_scan_impl(&value), + Err(std::env::VarError::NotPresent) => Ok(false), + Err(std::env::VarError::NotUnicode(value)) => { + vortex_bail!("{SCAN_IMPL_ENV} must be valid unicode, got {value:?}") + } + } +} + +fn parse_scan_impl(value: &str) -> VortexResult { + match value { + "" | "plan" | "v1" | "legacy" | "layout-reader" => Ok(false), + "planv2" | "plan-v2" | "v2" | "planned" | "scan-plan" => Ok(true), + other => vortex_bail!( + "{SCAN_IMPL_ENV} must be one of plan, v1, legacy, layout-reader, planv2, plan-v2, v2, planned, or scan-plan, got {other:?}" + ), + } +} + +/// Execute one split using a V2 physical scan plan. +/// +/// The execution order intentionally mirrors [`crate::scan::plan::split_exec`]. Expressions were +/// consumed during planning, so execution selects a predicate or projection plan without passing +/// an expression. +pub(crate) fn split_exec( + ctx: Arc>, + read_mask: RowMask, + limit: Option<&mut u64>, +) -> VortexResult>>> { + let row_range = read_mask.row_range(); + let row_mask = read_mask.mask().clone(); + + let filter_mask = match ctx.filter.as_ref() { + None => { + let row_mask = match limit { + Some(l) if *l == 0 => Mask::new_false(row_mask.len()), + Some(l) => { + let true_count = row_mask.true_count(); + let mask_limit = usize::try_from(*l) + .map(|l| l.min(true_count)) + .unwrap_or(true_count); + let row_mask = row_mask.limit(mask_limit); + *l -= mask_limit as u64; + row_mask + } + None => row_mask, + }; + + MaskFuture::ready(row_mask) + } + Some(filter) => { + if filter.conjuncts().len() != ctx.predicates.len() { + vortex_bail!( + "physical predicate count {} does not match conjunct count {}", + ctx.predicates.len(), + filter.conjuncts().len() + ); + } + + let ctx = Arc::clone(&ctx); + let filter = Arc::clone(filter); + let row_range = row_range.clone(); + + MaskFuture::new(row_mask.len(), async move { + let mut mask = row_mask; + let mut dynamic_versions = vec![None; filter.conjuncts().len()]; + + for (idx, predicate) in ctx.predicates.iter().enumerate() { + if mask.all_false() { + return Ok(mask); + } + + dynamic_versions[idx] = filter.dynamic_updates(idx).map(|du| du.version()); + let conjunct_mask = predicate + .pruning_evaluation(&row_range, mask.clone())? + .await?; + mask = mask.bitand(&conjunct_mask); + } + + let mut remaining = BitVec::from_elem(filter.conjuncts().len(), true); + while let Some(idx) = filter.next_conjunct(&remaining) { + remaining.set(idx, false); + if mask.all_false() { + return Ok(mask); + } + + let current_version = filter.dynamic_updates(idx).map(|du| du.version()); + if let Some(version) = current_version + && dynamic_versions[idx].is_none_or(|old| old < version) + { + dynamic_versions[idx] = Some(version); + let conjunct_mask = ctx.predicates[idx] + .pruning_evaluation(&row_range, mask.clone())? + .await?; + mask = mask.bitand(&conjunct_mask); + } + if mask.all_false() { + return Ok(mask); + } + + let conjunct_mask = ctx.predicates[idx] + .filter_evaluation(&row_range, MaskFuture::ready(mask))? + .await?; + filter.report_selectivity(idx, conjunct_mask.density()); + mask = conjunct_mask; + } + + Ok(mask) + }) + } + }; + + let projection_future = ctx + .projection + .projection_evaluation(&row_range, filter_mask.clone())?; + + let mapper = Arc::clone(&ctx.mapper); + let array_fut = async move { + let mask = filter_mask.await?; + if mask.all_false() { + return Ok(None); + } + + let array = projection_future.await?; + mapper(array).map(Some) + }; + + Ok(array_fut.boxed()) +} + +/// Information needed to execute one split from a V2 physical scan plan. +pub(crate) struct TaskContext { + filter: Option>, + predicates: Vec, + projection: ScanPlanRef, + mapper: Arc VortexResult + Send + Sync>, +} + +#[cfg(test)] +mod tests { + use std::any::Any; + + use vortex_array::dtype::FieldMask; + use vortex_array::dtype::FieldName; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::expr::eq; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + + use super::*; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; + + #[test] + fn scan_impl_accepts_v1_and_v2_values() -> VortexResult<()> { + for value in ["", "plan", "v1", "legacy", "layout-reader"] { + assert!(!parse_scan_impl(value)?); + } + for value in ["planv2", "plan-v2", "v2", "planned", "scan-plan"] { + assert!(parse_scan_impl(value)?); + } + Ok(()) + } + + #[test] + fn scan_impl_rejects_unknown_value() { + assert!(parse_scan_impl("unknown").is_err()); + } + + struct TestLayoutReader { + name: Arc, + dtype: DType, + } + + impl TestLayoutReader { + fn new() -> Self { + Self { + name: Arc::from("test"), + dtype: DType::Struct( + StructFields::from_iter([( + FieldName::from("a"), + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ), + } + } + } + + impl LayoutReader for TestLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 1 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + _split_range: &SplitRange, + _splits: &mut RowSplits, + ) -> VortexResult<()> { + unimplemented!("not needed for scan-plan construction") + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + _mask: Mask, + ) -> VortexResult { + unimplemented!("not needed for scan-plan construction") + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + _mask: MaskFuture, + ) -> VortexResult { + unimplemented!("not needed for scan-plan construction") + } + + fn projection_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + _mask: MaskFuture, + ) -> VortexResult { + unimplemented!("not needed for scan-plan construction") + } + } + + #[test] + fn get_item_reduces_to_a_struct_field_scan_plan() -> VortexResult<()> { + let reader: LayoutReaderRef = Arc::new(TestLayoutReader::new()); + let reader_plan: ScanPlanRef = Arc::new(LayoutReaderScanPlanV2::new(reader)); + let struct_plan = Arc::new(StructScanPlan::try_new(reader_plan)?); + let source: ScanPlanRef = struct_plan; + + let field = Arc::clone(&source).apply_expr(get_item("a", root()))?; + assert!(field.as_any().is::()); + assert_eq!( + field.dtype(), + &DType::Primitive(PType::I32, Nullability::NonNullable) + ); + + let predicate = field + .optimize()? + .apply_expr(eq(root(), lit(1_i32)))? + .optimize()?; + assert_eq!(predicate.dtype(), &DType::Bool(Nullability::NonNullable)); + + Ok(()) + } +} diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 681f33639bc..117fa276168 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -26,10 +26,10 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; -use crate::scan::filter::FilterExpr; +use crate::scan::plan; +use crate::scan::plan_v2; +use crate::scan::plan_v2::ScanPlanRef; use crate::scan::splits::Splits; -use crate::scan::tasks::TaskContext; -use crate::scan::tasks::split_exec; /// A projected subset (by indices, range, and filter) of rows from a Vortex data source. /// @@ -37,9 +37,7 @@ use crate::scan::tasks::split_exec; /// data source. pub struct RepeatedScan { session: VortexSession, - layout_reader: LayoutReaderRef, - projection: Expression, - filter: Option, + execution: ExecutionPlan, ordered: bool, /// Optionally read a subset of the rows in the file. row_range: Option>, @@ -57,6 +55,16 @@ pub struct RepeatedScan { dtype: DType, } +enum ExecutionPlan { + Plan(plan::Plan), + PlanV2(plan_v2::PlanV2), +} + +enum ExecutionTaskContext { + Plan(Arc>), + PlanV2(Arc>), +} + impl RepeatedScan { pub fn dtype(&self) -> &DType { &self.dtype @@ -89,7 +97,7 @@ impl RepeatedScan { clippy::too_many_arguments, reason = "all arguments are needed for scan construction" )] - pub fn new( + pub fn new_plan( session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -105,9 +113,40 @@ impl RepeatedScan { ) -> Self { Self { session, - layout_reader, - projection, - filter, + execution: ExecutionPlan::Plan(plan::Plan::new(layout_reader, projection, filter)), + ordered, + row_range, + selection, + splits, + concurrency, + map_fn, + limit, + dtype, + } + } + + /// Construct a repeated scan from a prepared heap-allocated physical plan. + #[expect( + clippy::too_many_arguments, + reason = "all arguments are needed for scan construction" + )] + pub fn new_plan_v2( + session: VortexSession, + projection: ScanPlanRef, + predicates: Vec, + filter: Option, + ordered: bool, + row_range: Option>, + selection: Selection, + splits: Splits, + concurrency: usize, + map_fn: Arc VortexResult + Send + Sync>, + limit: Option, + dtype: DType, + ) -> Self { + Self { + session, + execution: ExecutionPlan::PlanV2(plan_v2::PlanV2::new(projection, predicates, filter)), ordered, row_range, selection, @@ -173,12 +212,14 @@ impl RepeatedScan { let mut limit = self.limit; let mut tasks = Vec::new(); - let ctx = Arc::new(TaskContext { - filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), - reader: Arc::clone(&self.layout_reader), - projection: self.projection.clone(), - mapper: Arc::clone(&self.map_fn), - }); + let ctx = match &self.execution { + ExecutionPlan::Plan(plan) => { + ExecutionTaskContext::Plan(plan.task_context(Arc::clone(&self.map_fn))) + } + ExecutionPlan::PlanV2(plan_v2) => { + ExecutionTaskContext::PlanV2(plan_v2.task_context(Arc::clone(&self.map_fn))) + } + }; for range in ranges { let row_mask = self.selection.row_mask(&range); @@ -186,7 +227,14 @@ impl RepeatedScan { continue; } - tasks.push(split_exec(Arc::clone(&ctx), row_mask, limit.as_mut())?); + tasks.push(match &ctx { + ExecutionTaskContext::Plan(ctx) => { + plan::split_exec(Arc::clone(ctx), row_mask, limit.as_mut())? + } + ExecutionTaskContext::PlanV2(ctx) => { + plan_v2::split_exec(Arc::clone(ctx), row_mask, limit.as_mut())? + } + }); if limit.is_some_and(|l| l == 0) { break; } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 11fd5c7b882..60eff4ed422 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -16,8 +16,10 @@ use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; +use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; use vortex_array::expr::analysis::referenced_field_paths; +use vortex_array::expr::forms::conjuncts; use vortex_array::expr::root; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; @@ -40,6 +42,10 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReader; use crate::LayoutReaderRef; use crate::layouts::row_idx::RowIdxLayoutReader; +use crate::scan::plan_v2::LayoutReaderScanPlanV2; +use crate::scan::plan_v2::ScanPlanRef; +use crate::scan::plan_v2::StructScanPlan; +use crate::scan::plan_v2::plan_v2_enabled; use crate::scan::repeated_scan::RepeatedScan; use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; @@ -309,7 +315,40 @@ impl ScanBuilder { )?) }; - Ok(RepeatedScan::new( + if plan_v2_enabled()? { + let source: ScanPlanRef = + Arc::new(LayoutReaderScanPlanV2::new(Arc::clone(&layout_reader))); + let source: ScanPlanRef = match source.dtype() { + DType::Struct(_, Nullability::NonNullable) => { + Arc::new(StructScanPlan::try_new(source)?) + } + _ => source, + }; + let projection_plan = Arc::clone(&source).apply_expr(projection)?.optimize()?; + let predicate_plans = filter + .as_ref() + .map(conjuncts) + .unwrap_or_default() + .into_iter() + .map(|expr| Arc::clone(&source).apply_expr(expr)?.optimize()) + .collect::>>()?; + return Ok(RepeatedScan::new_plan_v2( + self.session.clone(), + projection_plan, + predicate_plans, + filter, + self.ordered, + self.row_range, + self.selection, + splits, + self.concurrency, + self.map_fn, + self.limit, + dtype, + )); + } + + Ok(RepeatedScan::new_plan( self.session.clone(), layout_reader, projection,