From 6d8c163dfd822bad10ae056218a72aba54cffdb1 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 10 Aug 2026 11:04:26 +0100 Subject: [PATCH 01/13] Add layout scan physical plan model Signed-off-by: Joe Isaacs --- docs/developer-guide/index.md | 1 + .../internals/scan-planning.md | 63 +++ vortex-layout/src/lib.rs | 1 + vortex-layout/src/plan/children.rs | 136 +++++ vortex-layout/src/plan/display.rs | 131 +++++ vortex-layout/src/plan/lower.rs | 137 +++++ vortex-layout/src/plan/mod.rs | 72 +++ vortex-layout/src/plan/optimize.rs | 35 ++ vortex-layout/src/plan/plans/concat.rs | 136 +++++ vortex-layout/src/plan/plans/eval.rs | 101 ++++ vortex-layout/src/plan/plans/list_pack.rs | 123 +++++ vortex-layout/src/plan/plans/mod.rs | 32 ++ vortex-layout/src/plan/plans/pack.rs | 137 +++++ vortex-layout/src/plan/plans/row_idx.rs | 106 ++++ vortex-layout/src/plan/plans/segment_scan.rs | 95 ++++ vortex-layout/src/plan/plans/take.rs | 104 ++++ vortex-layout/src/plan/tests.rs | 478 ++++++++++++++++++ vortex-layout/src/plan/typed.rs | 376 ++++++++++++++ vortex-layout/src/plan/vtable.rs | 71 +++ 19 files changed, 2335 insertions(+) create mode 100644 docs/developer-guide/internals/scan-planning.md create mode 100644 vortex-layout/src/plan/children.rs create mode 100644 vortex-layout/src/plan/display.rs create mode 100644 vortex-layout/src/plan/lower.rs create mode 100644 vortex-layout/src/plan/mod.rs create mode 100644 vortex-layout/src/plan/optimize.rs create mode 100644 vortex-layout/src/plan/plans/concat.rs create mode 100644 vortex-layout/src/plan/plans/eval.rs create mode 100644 vortex-layout/src/plan/plans/list_pack.rs create mode 100644 vortex-layout/src/plan/plans/mod.rs create mode 100644 vortex-layout/src/plan/plans/pack.rs create mode 100644 vortex-layout/src/plan/plans/row_idx.rs create mode 100644 vortex-layout/src/plan/plans/segment_scan.rs create mode 100644 vortex-layout/src/plan/plans/take.rs create mode 100644 vortex-layout/src/plan/tests.rs create mode 100644 vortex-layout/src/plan/typed.rs create mode 100644 vortex-layout/src/plan/vtable.rs 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..e03131924ac --- /dev/null +++ b/docs/developer-guide/internals/scan-planning.md @@ -0,0 +1,63 @@ +# Scan Plans + +A scan plan is the physical plan for satisfying one scan query. It is a tree of physical operators +over a row domain, describing the reads and derived work needed to produce that query's result. + +## Operators, not layout mirrors + +Plan operators describe *what work happens*, not *which layout produced it*. Their identity and +operator-specific state are independent of the source layout kind. The complete plan node is not: +its common lazy-child container can own hidden source state used to materialize individual children +on demand. + +| Operator | Work | +| --- | --- | +| `SegmentScan` | read one segment and decode it to an array | +| `Concat` | concatenate its children row-wise | +| `Pack` | assemble a struct from one child per field, plus optional validity | +| `Take` | index `values` by `codes` | +| `ListPack` | assemble a list from elements and offsets, plus optional validity | +| `Eval` | apply an expression to its child | +| `RowIdx` | offset row numbers into the file's row domain | + +Naming operators for what they compute is what lets one rule cover every case. `Concat` of +`Concat` flattens on shape alone, and `Take` over `SegmentScan` is the dictionary pushdown, +regardless of the source layout. + +The stored layout tree describes all physical data in a file. A plan is query-specific: it is built +from that tree for one projection, filter, and row domain. Different queries over the same file can +therefore produce different plans. + +## Optimization + +Child replacement is implemented by the common plan container rather than by every operator. It +replaces the external child container, clones `PlanData`, then invokes the operator's +`PlanVTable::with_children` callback to validate the new children and refresh derived caches such +as `Concat` row offsets. Rules therefore rewrite the generic tree without reconstructing common +plan fields inside each operator. + +Optimization rewrites the initial tree so that each expression is evaluated as close as possible to +the physical data that can satisfy it. Every rewrite must preserve the query result, including its +dtype, row domain, row order, row identity, null behavior, and observable errors. + +Planning does not read segment data. It constructs and optimizes a description of the work that a +later execution stage will perform. + +## Vtables + +Each operator is a small vtable type implementing `PlanVTable`, paired with a `Plan` container +over a shared `PlanRef`. `PlanRef` points to one allocation whose ordinary fields hold the operator +ID, dtype, row count, and lazy children. Only the unsized tail containing the vtable and +`V::PlanData` is erased behind `dyn DynPlan`, so common-field reads do not use dynamic dispatch. +`Plan` provides typed access to that operator data through `Deref`. + +`PlanVTable` also carries `id` and a `Metadata` codec. Operators with no unrecoverable state +already serialize their metadata; the ones holding a read context or a bound expression return +`None` until those codecs exist. + +## Future work + +Plans currently stop at construction and optimization. Still to come: a plan registry and foreign +operator placeholder so third-party operators survive a round trip, a serialization envelope, and +an execution stage that walks an optimized plan, reads the referenced segments, and returns the +query result. diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 03cd832a280..0dd7527ba28 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -15,6 +15,7 @@ //! optional bound filter, optional row range, [`Selection`](vortex_scan::selection::Selection), //! split strategy, and task concurrency settings, then produces array streams or iterators. pub mod layouts; +pub mod plan; pub use children::*; pub use encoding::*; diff --git a/vortex-layout/src/plan/children.rs b/vortex-layout/src/plan/children.rs new file mode 100644 index 00000000000..4b639a57fd8 --- /dev/null +++ b/vortex-layout/src/plan/children.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::sync::Arc; + +use once_cell::sync::OnceCell; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::plan::PlanRef; + +type ChildInitializer = dyn Fn(usize) -> VortexResult + 'static + Send + Sync; + +/// Ordered plan children that may be initialized one slot at a time. +/// +/// Eagerly constructed operators store already-filled slots. Layout lowering instead installs an +/// initializer that owns the source layout and lowers each child on first access. +#[derive(Clone)] +pub struct PlanChildren { + initializer: Option>, + cache: Arc<[OnceCell]>, +} + +impl PlanChildren { + /// Creates lazy child slots backed by `initializer`. + pub(crate) fn lazy( + len: usize, + initializer: impl Fn(usize) -> VortexResult + 'static + Send + Sync, + ) -> Self { + Self { + initializer: Some(Arc::new(initializer)), + cache: (0..len).map(|_| OnceCell::new()).collect::>().into(), + } + } + + /// Returns the number of children without initializing any slot. + pub fn len(&self) -> usize { + self.cache.len() + } + + /// Returns whether there are no children. + pub fn is_empty(&self) -> bool { + self.cache.is_empty() + } + + /// Returns a child, initializing and caching its slot on first access. + pub fn get(&self, index: usize) -> VortexResult> { + let Some(cell) = self.cache.get(index) else { + return Ok(None); + }; + if let Some(child) = cell.get() { + return Ok(Some(child.clone())); + } + + let initializer = self + .initializer + .as_ref() + .ok_or_else(|| vortex_err!("Plan child {index} was not initialized"))?; + Ok(Some(cell.get_or_try_init(|| initializer(index))?.clone())) + } + + /// Iterates over the children in logical order, initializing slots as they are visited. + pub fn iter(&self) -> impl ExactSizeIterator> + '_ { + (0..self.len()).map(|index| { + self.get(index)? + .ok_or_else(|| vortex_err!("Plan child {index} is absent")) + }) + } + + /// Materializes all children into an eager vector. + pub fn to_vec(&self) -> VortexResult> { + self.iter().collect() + } + + /// Returns a child collection with one slot replaced. + pub fn with_child(&self, index: usize, child: PlanRef) -> VortexResult { + if index >= self.len() { + vortex_bail!("Plan child index out of bounds: {index} of {}", self.len()); + } + + let source = self.clone(); + Ok(Self::lazy(source.len(), move |child_index| { + if child_index == index { + return Ok(child.clone()); + } + source + .get(child_index)? + .ok_or_else(|| vortex_err!("Plan child {child_index} is absent")) + })) + } +} + +impl From> for PlanChildren { + fn from(children: Vec) -> Self { + let cache = children + .into_iter() + .map(OnceCell::with_value) + .collect::>() + .into(); + Self { + initializer: None, + cache, + } + } +} + +impl From<[PlanRef; N]> for PlanChildren { + fn from(children: [PlanRef; N]) -> Self { + Vec::from(children).into() + } +} + +impl Default for PlanChildren { + fn default() -> Self { + Vec::new().into() + } +} + +impl fmt::Debug for PlanChildren { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PlanChildren") + .field("len", &self.len()) + .field( + "initialized", + &self + .cache + .iter() + .filter(|slot| slot.get().is_some()) + .count(), + ) + .finish() + } +} diff --git a/vortex-layout/src/plan/display.rs b/vortex-layout/src/plan/display.rs new file mode 100644 index 00000000000..8a05f184761 --- /dev/null +++ b/vortex-layout/src/plan/display.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; + +pub use vortex_utils::tree::DepthContext as PlanTreeContext; +pub use vortex_utils::tree::IndentedFormatter as PlanIndentedFormatter; +use vortex_utils::tree::TreeDisplayAdapter; +pub use vortex_utils::tree::TreeDisplayExtractor as PlanTreeExtractor; +use vortex_utils::tree::write_indented_tree; + +use super::PlanRef; + +/// Adds the plan's display representation to a tree node's header. +pub struct PlanSummaryExtractor; + +impl PlanSummaryExtractor { + /// Writes a plan directly to `formatter`. + pub fn write(plan: &PlanRef, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "{plan}") + } +} + +impl PlanTreeExtractor for PlanSummaryExtractor { + fn write_header( + &self, + plan: &PlanRef, + _context: &PlanTreeContext, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + write!(formatter, " ")?; + Self::write(plan, formatter) + } +} + +/// Composable display builder for a physical plan tree. +/// +/// Call `plan.tree_display()` for the default extractors. Use `plan.tree_display_builder()` to +/// start with only node and child names, then add extractors with [`Self::with`]. +pub struct PlanTreeDisplay<'a> { + plan: &'a PlanRef, + extractors: Vec>>, +} + +impl<'a> PlanTreeDisplay<'a> { + /// Creates a tree display for `plan` with no extractors. + pub fn new(plan: &'a PlanRef) -> Self { + Self { + plan, + extractors: Vec::new(), + } + } + + /// Creates a tree display using each plan's display representation. + pub fn default_display(plan: &'a PlanRef) -> Self { + Self::new(plan).with(PlanSummaryExtractor) + } + + /// Adds an extractor to the display pipeline. + pub fn with + 'static>( + mut self, + extractor: E, + ) -> Self { + self.extractors.push(Box::new(extractor)); + self + } + + /// Adds a pre-boxed extractor to the display pipeline. + pub fn with_boxed( + mut self, + extractor: Box>, + ) -> Self { + self.extractors.push(extractor); + self + } +} + +impl TreeDisplayAdapter for PlanTreeDisplay<'_> { + type Context = PlanTreeContext; + type Node = PlanRef; + + fn write_node( + &self, + plan: &PlanRef, + context: &PlanTreeContext, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + for extractor in &self.extractors { + extractor.write_header(plan, context, formatter)?; + } + Ok(()) + } + + fn write_details( + &self, + plan: &PlanRef, + context: &PlanTreeContext, + formatter: &mut PlanIndentedFormatter<'_, '_>, + ) -> fmt::Result { + for extractor in &self.extractors { + extractor.write_details(plan, context, formatter)?; + } + Ok(()) + } + + fn visit_children( + &self, + plan: &PlanRef, + visit: &mut dyn FnMut(&str, &PlanRef, bool) -> fmt::Result, + ) -> fmt::Result { + let children = plan.children(); + for index in 0..children.len() { + let child = plan.child_required(index).map_err(|_| fmt::Error)?; + let child_name = plan.child_name(index); + visit(child_name.as_ref(), &child, index + 1 == children.len())?; + } + Ok(()) + } +} + +impl fmt::Display for PlanTreeDisplay<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write_indented_tree( + self, + "root", + self.plan, + &mut PlanTreeContext::default(), + formatter, + ) + } +} diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs new file mode 100644 index 00000000000..82c36b584e5 --- /dev/null +++ b/vortex-layout/src/plan/lower.rs @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Test support for constructing physical plans from stored layout trees. +//! +//! This module is only used to build physical-plan fixtures for tests. It is not a production +//! planning API. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::LayoutRef; +use crate::layouts::chunked::Chunked; +use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::dict::Dict; +use crate::layouts::dict::DictLayout; +use crate::layouts::flat::Flat; +use crate::layouts::flat::FlatLayout; +use crate::layouts::list::ELEMENTS_CHILD_INDEX; +use crate::layouts::list::List; +use crate::layouts::list::ListLayout; +use crate::layouts::list::OFFSETS_CHILD_INDEX; +use crate::layouts::list::VALIDITY_CHILD_INDEX; +use crate::layouts::struct_::Struct; +use crate::layouts::struct_::StructLayout; +use crate::plan::ConcatPlan; +use crate::plan::ListPackPlan; +use crate::plan::PackPlan; +use crate::plan::PlanChildren; +use crate::plan::PlanRef; +use crate::plan::SegmentScanPlan; +use crate::plan::TakePlan; + +/// Constructs a physical-plan fixture from `layout` for tests. +/// +/// The root operator is built immediately. Its child container owns a hidden clone of the source +/// layout and lowers each child independently on first access. +pub fn lower(layout: &LayoutRef) -> VortexResult { + if let Some(layout) = layout.as_opt::() { + return Ok(lower_flat(layout).into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_chunked(layout)?.into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_struct(layout)?.into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_dict(layout)?.into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_list(layout)?.into_plan()); + } + vortex_bail!( + "No physical plan implementation for layout '{}'", + layout.encoding_id() + ) +} + +fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan { + SegmentScanPlan::new( + layout.dtype().clone(), + layout.row_count(), + layout.segment_id(), + layout.array_ctx().clone(), + layout.array_tree().cloned(), + ) +} + +fn lower_chunked(layout: &ChunkedLayout) -> VortexResult { + let mut row_offsets = Vec::with_capacity(layout.nchildren()); + let mut row_count = 0u64; + for index in 0..layout.nchildren() { + row_offsets.push(row_count); + row_count = row_count + .checked_add(layout.child_row_count(index)) + .ok_or_else(|| vortex_err!("Chunked row count overflow"))?; + } + Ok(ConcatPlan::from_children( + layout.dtype().clone(), + layout.row_count(), + row_offsets.into(), + lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), + )) +} + +fn lower_struct(layout: &StructLayout) -> VortexResult { + // Struct layout slot 0 is validity and field i is slot i + 1. The plan puts validity last so + // field indices are identical to their plan-child indices. + let fields = layout.struct_fields().clone(); + let mut slots = (1..=fields.nfields()).collect::>(); + if layout.dtype().is_nullable() { + slots.push(0); + } + Ok(PackPlan::from_children( + fields, + layout.dtype().nullability(), + layout.row_count(), + lazy_children(layout.to_layout(), slots), + )) +} + +fn lower_dict(layout: &DictLayout) -> VortexResult { + // Dict serialization stores values before codes; the plan order is deliberately codes, + // values because that is the optimizer-facing logical shape. + Ok(TakePlan::from_children( + layout.dtype().clone(), + layout.row_count(), + lazy_children(layout.to_layout(), vec![1, 0]), + )) +} + +fn lower_list(layout: &ListLayout) -> VortexResult { + let mut slots = vec![ELEMENTS_CHILD_INDEX, OFFSETS_CHILD_INDEX]; + if layout.dtype().is_nullable() { + slots.push(VALIDITY_CHILD_INDEX); + } + Ok(ListPackPlan::from_children( + layout.dtype().clone(), + layout.row_count(), + lazy_children(layout.to_layout(), slots), + )) +} + +fn lazy_children(layout: LayoutRef, slots: Vec) -> PlanChildren { + PlanChildren::lazy(slots.len(), move |index| { + let slot = slots + .get(index) + .copied() + .ok_or_else(|| vortex_err!("Missing plan child slot {index}"))?; + let child = layout + .slot(slot)? + .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?; + lower(&child) + }) +} diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs new file mode 100644 index 00000000000..06f7acd0031 --- /dev/null +++ b/vortex-layout/src/plan/mod.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Physical plans for scans. +//! +//! A plan is a tree of physical operators over a row domain. Operator identity and +//! operator-specific state do not depend on the source layout kind, so rewrites can reason about a +//! plan's shape alone. The common child container can initialize individual slots lazily. + +mod children; +mod display; +mod lower; +mod optimize; +mod plans; +mod typed; +mod vtable; + +pub use children::PlanChildren; +pub use display::PlanIndentedFormatter; +pub use display::PlanSummaryExtractor; +pub use display::PlanTreeContext; +pub use display::PlanTreeDisplay; +pub use display::PlanTreeExtractor; +pub use lower::lower; +pub use optimize::optimize; +pub use plans::Concat; +pub use plans::ConcatData; +pub use plans::ConcatPlan; +pub use plans::Eval; +pub use plans::EvalData; +pub use plans::EvalPlan; +pub use plans::ListPack; +pub use plans::ListPackData; +pub use plans::ListPackPlan; +pub use plans::Pack; +pub use plans::PackData; +pub use plans::PackPlan; +pub use plans::RowIdx; +pub use plans::RowIdxData; +pub use plans::RowIdxPlan; +pub use plans::RowIdxPlanMetadata; +pub use plans::SegmentScan; +pub use plans::SegmentScanData; +pub use plans::SegmentScanPlan; +pub use plans::Take; +pub use plans::TakePlan; +pub use typed::DynPlan; +pub use typed::Plan; +pub use typed::PlanParts; +pub use typed::PlanRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +pub use vtable::PlanId; +pub use vtable::PlanVTable; + +/// Returns an error when `children` does not have exactly `expected` entries. +pub(crate) fn check_child_count( + name: &str, + children: &PlanChildren, + expected: usize, +) -> VortexResult<()> { + if children.len() != expected { + vortex_bail!( + "{name} expects {expected} children but got {}", + children.len() + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/plan/optimize.rs b/vortex-layout/src/plan/optimize.rs new file mode 100644 index 00000000000..d78798cdb65 --- /dev/null +++ b/vortex-layout/src/plan/optimize.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Generic bottom-up optimization over physical plans. + +use vortex_error::VortexResult; + +use crate::plan::Eval; +use crate::plan::PlanRef; + +/// Optimizes `plan`, preserving its dtype and row domain. +pub fn optimize(plan: PlanRef) -> VortexResult { + let mut children = Vec::with_capacity(plan.child_count()); + let mut changed = false; + for child in plan.children().iter() { + let child = child?; + let optimized = optimize(child.clone())?; + changed |= !PlanRef::ptr_eq(&child, &optimized); + children.push(optimized); + } + + let plan = if changed { + plan.with_children(children)? + } else { + plan + }; + + let Some(eval) = plan.as_opt::() else { + return Ok(plan); + }; + if eval.expression().is_root() { + return eval.child_plan(); + } + Ok(plan) +} diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs new file mode 100644 index 00000000000..414a685ec5d --- /dev/null +++ b/vortex-layout/src/plan/plans/concat.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::sync::Arc; + +use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; + +/// Concatenates its children row-wise. +#[derive(Clone, Debug)] +pub struct Concat; + +/// Row offsets of each concatenated child. +#[derive(Clone, Debug)] +pub struct ConcatData { + row_offsets: Arc<[u64]>, +} + +/// A plan that concatenates its children row-wise. +pub type ConcatPlan = Plan; + +impl ConcatPlan { + pub(crate) fn from_children( + dtype: DType, + row_count: u64, + row_offsets: Arc<[u64]>, + children: PlanChildren, + ) -> Self { + PlanParts { + vtable: Concat, + dtype, + row_count, + children, + data: ConcatData { row_offsets }, + } + .into_typed() + } + + /// Creates a concatenation over `children`. + /// + /// Every child must produce `dtype`, and the row domain is the sum of the child row counts. + pub fn try_new(dtype: DType, children: Vec) -> VortexResult { + let mut row_offsets = Vec::with_capacity(children.len()); + let mut row_count = 0u64; + for child in &children { + if child.dtype() != &dtype { + vortex_bail!( + "Concat child dtype {} does not match {dtype}", + child.dtype() + ); + } + row_offsets.push(row_count); + row_count += child.row_count(); + } + Ok(Self::from_children( + dtype, + row_count, + row_offsets.into(), + children.into(), + )) + } + + /// Returns the first row of each child within this plan's row domain. + pub fn row_offsets(&self) -> &[u64] { + &self.data().row_offsets + } +} + +impl PlanVTable for Concat { + type PlanData = ConcatData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.concat"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + // Row offsets are derived from the children, so nothing needs storing. + Some(EmptyMetadata) + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + data: &mut Self::PlanData, + ) -> VortexResult<()> { + if children.len() != plan.children().len() { + vortex_bail!( + "Concat expects {} children but got {}", + plan.children().len(), + children.len() + ); + } + + let mut row_offsets = Vec::with_capacity(children.len()); + let mut row_count = 0u64; + for child in children.iter() { + let child = child?; + if child.dtype() != plan.dtype() { + vortex_bail!( + "Concat child dtype {} does not match {}", + child.dtype(), + plan.dtype() + ); + } + row_offsets.push(row_count); + row_count = row_count + .checked_add(child.row_count()) + .ok_or_else(|| vortex_error::vortex_err!("Concat row count overflow"))?; + } + if row_count != plan.row_count() { + vortex_bail!( + "Concat children have {row_count} rows but the plan has {}", + plan.row_count() + ); + } + data.row_offsets = row_offsets.into(); + Ok(()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + Cow::Owned(format!("chunks[{index}]")) + } +} diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs new file mode 100644 index 00000000000..05421eeda74 --- /dev/null +++ b/vortex-layout/src/plan/plans/eval.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::fmt; + +use vortex_array::EmptyMetadata; +use vortex_array::expr::BoundExpression; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::plan::check_child_count; + +/// Applies an expression to the output of its child. +#[derive(Clone, Debug)] +pub struct Eval; + +/// The expression evaluated by an [`Eval`]. +#[derive(Clone, Debug)] +pub struct EvalData { + expression: BoundExpression, +} + +/// A plan that applies an expression to its child. +pub type EvalPlan = Plan; + +impl EvalPlan { + /// Creates an evaluation of `expression`, which must be bound to the child's dtype. + pub fn new(expression: BoundExpression, child: PlanRef) -> Self { + PlanParts { + vtable: Eval, + dtype: expression.dtype().clone(), + row_count: child.row_count(), + children: vec![child].into(), + data: EvalData { expression }, + } + .into_typed() + } + + /// Returns the expression evaluated by this plan. + pub fn expression(&self) -> &BoundExpression { + &self.data().expression + } + + /// Returns the child plan supplying the expression root. + pub fn child_plan(&self) -> VortexResult { + self.child_required(0) + } +} + +impl PlanVTable for Eval { + type PlanData = EvalData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.eval"); + *ID + } + + fn fmt(plan: &Plan, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, " expr={}", plan.expression()) + } + + fn metadata(_plan: &Plan) -> Option { + // Expressions serialize through `vortex.expr` protobuf, which is not wired up here yet. + None + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + check_child_count("Eval", children, 1)?; + let child = children + .get(0)? + .ok_or_else(|| vortex_error::vortex_err!("Eval child is absent"))?; + if child.row_count() != plan.row_count() { + vortex_error::vortex_bail!( + "Eval child has {} rows but the plan has {}", + child.row_count(), + plan.row_count() + ); + } + Ok(()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + if index == 0 { + Cow::Borrowed("child") + } else { + Cow::Owned(format!("child[{index}]")) + } + } +} diff --git a/vortex-layout/src/plan/plans/list_pack.rs b/vortex-layout/src/plan/plans/list_pack.rs new file mode 100644 index 00000000000..f62f8c1994b --- /dev/null +++ b/vortex-layout/src/plan/plans/list_pack.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::sync::Arc; + +use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; + +const ELEMENTS: usize = 0; +const OFFSETS: usize = 1; +const VALIDITY: usize = 2; + +/// Assembles a list from elements and offsets, plus an optional trailing validity child. +#[derive(Clone, Debug)] +pub struct ListPack; + +/// Operator-specific list assembly data. +#[derive(Clone, Debug)] +pub struct ListPackData; + +/// A plan that assembles a list from its children. +pub type ListPackPlan = Plan; + +impl ListPackPlan { + pub(crate) fn from_children(dtype: DType, row_count: u64, children: PlanChildren) -> Self { + PlanParts { + vtable: ListPack, + dtype, + row_count, + children, + data: ListPackData, + } + .into_typed() + } + + /// Creates a list assembly from `elements` and `offsets`. + /// + /// `validity` is required exactly when `nullability` is [`Nullability::Nullable`]. The row + /// domain is one fewer than the number of offsets. + pub fn try_new( + nullability: Nullability, + row_count: u64, + elements: PlanRef, + offsets: PlanRef, + validity: Option, + ) -> VortexResult { + if validity.is_some() != (nullability == Nullability::Nullable) { + vortex_bail!( + "ListPack validity child must be present exactly when the list is nullable" + ); + } + let dtype = DType::List(Arc::new(elements.dtype().clone()), nullability); + let mut children = vec![elements, offsets]; + children.extend(validity); + Ok(Self::from_children(dtype, row_count, children.into())) + } + + /// Returns the plan producing list elements. + pub fn elements(&self) -> VortexResult { + self.child_required(ELEMENTS) + } + + /// Returns the plan producing list offsets. + pub fn offsets(&self) -> VortexResult { + self.child_required(OFFSETS) + } + + /// Returns the plan producing list validity, if the list is nullable. + pub fn validity(&self) -> VortexResult> { + self.child(VALIDITY) + } +} + +impl PlanVTable for ListPack { + type PlanData = ListPackData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.list_pack"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + // Nullability is recoverable from the plan dtype. + Some(EmptyMetadata) + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + if children.len() != plan.children().len() { + vortex_bail!( + "ListPack expects {} children but got {}", + plan.children().len(), + children.len() + ); + } + Ok(()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + match index { + ELEMENTS => Cow::Borrowed("elements"), + OFFSETS => Cow::Borrowed("offsets"), + VALIDITY => Cow::Borrowed("validity"), + _ => Cow::Owned(format!("child[{index}]")), + } + } +} diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs new file mode 100644 index 00000000000..2e4a6dbad5e --- /dev/null +++ b/vortex-layout/src/plan/plans/mod.rs @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod concat; +mod eval; +mod list_pack; +mod pack; +mod row_idx; +mod segment_scan; +mod take; + +pub use concat::Concat; +pub use concat::ConcatData; +pub use concat::ConcatPlan; +pub use eval::Eval; +pub use eval::EvalData; +pub use eval::EvalPlan; +pub use list_pack::ListPack; +pub use list_pack::ListPackData; +pub use list_pack::ListPackPlan; +pub use pack::Pack; +pub use pack::PackData; +pub use pack::PackPlan; +pub use row_idx::RowIdx; +pub use row_idx::RowIdxData; +pub use row_idx::RowIdxPlan; +pub use row_idx::RowIdxPlanMetadata; +pub use segment_scan::SegmentScan; +pub use segment_scan::SegmentScanData; +pub use segment_scan::SegmentScanPlan; +pub use take::Take; +pub use take::TakePlan; diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs new file mode 100644 index 00000000000..156fcd5af4c --- /dev/null +++ b/vortex-layout/src/plan/plans/pack.rs @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; + +use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; + +/// Assembles a struct from one child per field, plus an optional trailing validity child. +#[derive(Clone, Debug)] +pub struct Pack; + +/// Operator-specific struct assembly data. +#[derive(Clone, Debug)] +pub struct PackData; + +/// A plan that assembles a struct from its children. +pub type PackPlan = Plan; + +impl PackPlan { + pub(crate) fn from_children( + fields: StructFields, + nullability: Nullability, + row_count: u64, + children: PlanChildren, + ) -> Self { + PlanParts { + vtable: Pack, + dtype: DType::Struct(fields, nullability), + row_count, + children, + data: PackData, + } + .into_typed() + } + + /// Creates a struct assembly from `fields` and one child per field. + /// + /// `validity` is required exactly when `nullability` is [`Nullability::Nullable`]. + pub fn try_new( + fields: StructFields, + nullability: Nullability, + row_count: u64, + field_plans: Vec, + validity: Option, + ) -> VortexResult { + if field_plans.len() != fields.nfields() { + vortex_bail!( + "Pack expects {} field children but got {}", + fields.nfields(), + field_plans.len() + ); + } + if validity.is_some() != (nullability == Nullability::Nullable) { + vortex_bail!("Pack validity child must be present exactly when the struct is nullable"); + } + + let mut children = field_plans; + children.extend(validity); + Ok(Self::from_children( + fields, + nullability, + row_count, + children.into(), + )) + } + + /// Returns the struct fields assembled by this plan. + pub fn fields(&self) -> &StructFields { + self.dtype() + .as_struct_fields_opt() + .vortex_expect("Pack dtype must be a struct") + } + + /// Returns the number of struct fields, excluding any validity child. + pub fn nfields(&self) -> usize { + self.fields().nfields() + } + + /// Returns the plan producing struct validity, if the struct is nullable. + pub fn validity(&self) -> VortexResult> { + self.child(self.nfields()) + } +} + +impl PlanVTable for Pack { + type PlanData = PackData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.pack"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + // The struct fields are recoverable from the plan dtype. + Some(EmptyMetadata) + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + if children.len() != plan.children().len() { + vortex_bail!( + "Pack expects {} children but got {}", + plan.children().len(), + children.len() + ); + } + Ok(()) + } + + fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { + if let Some(name) = plan.fields().field_name(index) { + return Cow::Borrowed(name.as_ref()); + } + if index == plan.fields().nfields() { + return Cow::Borrowed("validity"); + } + Cow::Owned(format!("child[{index}]")) + } +} diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs new file mode 100644 index 00000000000..99d8f2f963d --- /dev/null +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; + +use vortex_array::ProstMetadata; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::plan::check_child_count; + +/// Adds row-index support to its child, offsetting row numbers into the file's row domain. +#[derive(Clone, Debug)] +pub struct RowIdx; + +/// The row offset applied to the child domain. +#[derive(Clone, Debug)] +pub struct RowIdxData { + row_offset: u64, +} + +/// A plan that adds row-index support to its child. +pub type RowIdxPlan = Plan; + +impl RowIdxPlan { + /// Creates a row-index plan with `row_offset` applied to its child domain. + pub fn new(row_offset: u64, child: PlanRef) -> Self { + PlanParts { + vtable: RowIdx, + dtype: child.dtype().clone(), + row_count: child.row_count(), + children: vec![child].into(), + data: RowIdxData { row_offset }, + } + .into_typed() + } + + /// Returns the row offset applied to the child domain. + pub fn row_offset(&self) -> u64 { + self.data().row_offset + } + + /// Returns the child plan. + pub fn child_plan(&self) -> VortexResult { + self.child_required(0) + } +} + +impl PlanVTable for RowIdx { + type PlanData = RowIdxData; + type Metadata = ProstMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.row_idx"); + *ID + } + + fn metadata(plan: &Plan) -> Option { + Some(ProstMetadata(RowIdxPlanMetadata { + row_offset: plan.data().row_offset, + })) + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + check_child_count("RowIdx", children, 1)?; + let child = children + .get(0)? + .ok_or_else(|| vortex_error::vortex_err!("RowIdx child is absent"))?; + if child.dtype() != plan.dtype() || child.row_count() != plan.row_count() { + vortex_error::vortex_bail!( + "RowIdx child shape changed from ({}, {}) to ({}, {})", + plan.dtype(), + plan.row_count(), + child.dtype(), + child.row_count() + ); + } + Ok(()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + if index == 0 { + Cow::Borrowed("child") + } else { + Cow::Owned(format!("child[{index}]")) + } + } +} + +/// Serialized metadata for a [`RowIdx`] plan. +#[derive(Clone, PartialEq, Eq, ::prost::Message)] +pub struct RowIdxPlanMetadata { + /// The row offset applied to the child domain. + #[prost(uint64, tag = "1")] + pub row_offset: u64, +} diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/plans/segment_scan.rs new file mode 100644 index 00000000000..d2b20df89ad --- /dev/null +++ b/vortex-layout/src/plan/plans/segment_scan.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use vortex_session::registry::ReadContext; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanVTable; +use crate::plan::check_child_count; +use crate::segments::SegmentId; + +/// Reads one serialized array segment. +#[derive(Clone, Debug)] +pub struct SegmentScan; + +/// Data needed to read and decode a single segment. +#[derive(Clone, Debug)] +pub struct SegmentScanData { + segment_id: SegmentId, + array_ctx: ReadContext, + array_tree: Option, +} + +/// A plan that reads one serialized array segment. +pub type SegmentScanPlan = Plan; + +impl SegmentScanPlan { + /// Creates a segment scan over `segment_id`. + pub fn new( + dtype: DType, + row_count: u64, + segment_id: SegmentId, + array_ctx: ReadContext, + array_tree: Option, + ) -> Self { + PlanParts { + vtable: SegmentScan, + dtype, + row_count, + children: PlanChildren::default(), + data: SegmentScanData { + segment_id, + array_ctx, + array_tree, + }, + } + .into_typed() + } + + /// Returns the segment this plan reads. + pub fn segment_id(&self) -> SegmentId { + self.data().segment_id + } + + /// Returns the read context for the serialized array. + pub fn array_ctx(&self) -> &ReadContext { + &self.data().array_ctx + } + + /// Returns the serialized array encoding tree, when it is stored out of line. + pub fn array_tree(&self) -> Option<&ByteBuffer> { + self.data().array_tree.as_ref() + } +} + +impl PlanVTable for SegmentScan { + type PlanData = SegmentScanData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.segment_scan"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + // The segment ID and read context are not yet covered by a metadata codec. + None + } + + fn with_children( + _plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + check_child_count("SegmentScan", children, 0)?; + Ok(()) + } +} diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/plans/take.rs new file mode 100644 index 00000000000..1184db10e65 --- /dev/null +++ b/vortex-layout/src/plan/plans/take.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; + +use vortex_array::EmptyMetadata; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::plan::check_child_count; + +const CODES: usize = 0; +const VALUES: usize = 1; + +/// Indexes `values` by `codes`, with children ordered as `[codes, values]`. +#[derive(Clone, Debug)] +pub struct Take; + +/// A plan that indexes one child by another. +pub type TakePlan = Plan; + +impl TakePlan { + pub(crate) fn from_children( + dtype: vortex_array::dtype::DType, + row_count: u64, + children: PlanChildren, + ) -> Self { + PlanParts { + vtable: Take, + dtype, + row_count, + children, + data: (), + } + .into_typed() + } + + /// Creates a take of `values` at `codes`. + /// + /// The row domain is that of `codes`, and the output dtype is that of `values`. + pub fn new(codes: PlanRef, values: PlanRef) -> Self { + Self::from_children( + values.dtype().clone(), + codes.row_count(), + vec![codes, values].into(), + ) + } + + /// Returns the plan producing indices. + pub fn codes(&self) -> VortexResult { + self.child_required(CODES) + } + + /// Returns the plan producing the values being indexed. + pub fn values(&self) -> VortexResult { + self.child_required(VALUES) + } +} + +impl PlanVTable for Take { + type PlanData = (); + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.take"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + Some(EmptyMetadata) + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + check_child_count("Take", children, 2)?; + let codes = children + .get(CODES)? + .ok_or_else(|| vortex_error::vortex_err!("Take codes child is absent"))?; + let values = children + .get(VALUES)? + .ok_or_else(|| vortex_error::vortex_err!("Take values child is absent"))?; + if codes.row_count() != plan.row_count() || values.dtype() != plan.dtype() { + vortex_error::vortex_bail!("Take child shape does not match the plan output"); + } + Ok(()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + match index { + CODES => Cow::Borrowed("codes"), + VALUES => Cow::Borrowed("values"), + _ => Cow::Owned(format!("child[{index}]")), + } + } +} diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs new file mode 100644 index 00000000000..9c478317ee8 --- /dev/null +++ b/vortex-layout/src/plan/tests.rs @@ -0,0 +1,478 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::sync::Arc; + +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::get_item; +use vortex_array::expr::root; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::LayoutRef; +use crate::OwnedLayoutChildren; +use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::dict::DictLayout; +use crate::layouts::flat::FlatLayout; +use crate::layouts::foreign::new_foreign_layout; +use crate::layouts::list::ListLayout; +use crate::layouts::row_idx::row_idx; +use crate::layouts::struct_::StructLayout; +use crate::segments::SegmentId; + +fn primitive(ptype: PType, nullability: Nullability) -> DType { + DType::Primitive(ptype, nullability) +} + +fn flat(row_count: u64, dtype: DType, segment: u32) -> LayoutRef { + FlatLayout::new( + row_count, + dtype, + SegmentId::from(segment), + ReadContext::new([]), + ) + .into_layout() +} + +fn unsupported(row_count: u64, dtype: DType) -> LayoutRef { + static ID: CachedId = CachedId::new("vortex.test.unsupported"); + new_foreign_layout(*ID, dtype, row_count, Vec::new(), Vec::new(), Vec::new()) +} + +fn make_plan(layout: LayoutRef) -> VortexResult { + lower(&layout) +} + +fn child_of(plan: &PlanRef, index: usize) -> VortexResult { + plan.child(index)? + .ok_or_else(|| vortex_err!("missing child {index}")) +} + +fn assert_unsupported(error: vortex_error::VortexError) { + assert!( + error + .to_string() + .contains("No physical plan implementation for layout 'vortex.test.unsupported'"), + "unexpected error: {error}" + ); +} + +#[test] +fn unsupported_layout_has_no_plan() -> VortexResult<()> { + let layout = unsupported(3, DType::Null); + + assert_unsupported( + lower(&layout) + .err() + .ok_or_else(|| vortex_err!("unsupported layout unexpectedly produced a plan"))?, + ); + Ok(()) +} + +#[test] +fn flat_plan_has_no_children() -> VortexResult<()> { + let plan = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; + + assert!(plan.is::()); + assert_eq!(plan.child_count(), 0); + assert!(plan.child(0)?.is_none()); + Ok(()) +} + +#[test] +fn chunked_plan_exposes_chunks() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 3, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + ) + .into_layout(); + let plan = make_plan(layout)?; + + assert!(plan.is::()); + assert_eq!(plan.child_count(), 2); + assert_eq!(child_of(&plan, 0)?.row_count(), 2); + assert_eq!(child_of(&plan, 1)?.row_count(), 1); + Ok(()) +} + +#[test] +fn chunked_plan_lowers_each_chunk_on_access() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 2, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![ + flat(1, dtype.clone(), 0), + unsupported(1, dtype), + ]), + ) + .into_layout(); + + let plan = make_plan(layout)?; + assert_eq!(child_of(&plan, 0)?.row_count(), 1); + assert_unsupported( + child_of(&plan, 1) + .err() + .ok_or_else(|| vortex_err!("unsupported chunk unexpectedly produced a plan"))?, + ); + Ok(()) +} + +#[test] +fn struct_plan_lowers_each_field_on_access() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 1, + DType::Struct( + StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), + Nullability::NonNullable, + ), + vec![flat(1, field_dtype.clone(), 0), unsupported(1, field_dtype)], + ) + .into_layout(); + + let plan = make_plan(layout)?; + assert_eq!(child_of(&plan, 0)?.row_count(), 1); + assert_unsupported( + child_of(&plan, 1) + .err() + .ok_or_else(|| vortex_err!("unsupported field unexpectedly produced a plan"))?, + ); + Ok(()) +} + +#[test] +fn dict_plan_orders_codes_before_values() -> VortexResult<()> { + let values_dtype = primitive(PType::I32, Nullability::NonNullable); + let codes_dtype = primitive(PType::U8, Nullability::NonNullable); + let layout = DictLayout::new( + flat(2, values_dtype.clone(), 0), + flat(3, codes_dtype.clone(), 1), + ) + .into_layout(); + let plan = make_plan(layout)?; + + assert!(plan.is::()); + assert_eq!(plan.child_count(), 2); + assert_eq!(child_of(&plan, 0)?.dtype(), &codes_dtype); + assert_eq!(child_of(&plan, 1)?.dtype(), &values_dtype); + Ok(()) +} + +#[test] +fn list_plan_appends_validity_when_nullable() -> VortexResult<()> { + let element_dtype = primitive(PType::I32, Nullability::NonNullable); + let offsets_dtype = primitive(PType::U32, Nullability::NonNullable); + let non_nullable = ListLayout::new( + DType::List(Arc::new(element_dtype.clone()), Nullability::NonNullable), + flat(3, element_dtype.clone(), 0), + flat(3, offsets_dtype.clone(), 1), + None, + ) + .into_layout(); + let plan = make_plan(non_nullable)?; + + assert!(plan.is::()); + assert_eq!(plan.child_count(), 2); + assert_eq!(child_of(&plan, 0)?.dtype(), &element_dtype); + assert_eq!(child_of(&plan, 1)?.dtype(), &offsets_dtype); + + let nullable = ListLayout::new( + DType::List(Arc::new(element_dtype.clone()), Nullability::Nullable), + flat(3, element_dtype, 2), + flat(3, offsets_dtype, 3), + Some(flat(2, DType::Bool(Nullability::NonNullable), 4)), + ) + .into_layout(); + let nullable_plan = make_plan(nullable)?; + assert_eq!(nullable_plan.child_count(), 3); + assert_eq!( + child_of(&nullable_plan, 2)?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + Ok(()) +} + +#[test] +fn struct_plan_appends_validity_when_nullable() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let fields = StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]); + let non_nullable = StructLayout::new( + 3, + DType::Struct(fields.clone(), Nullability::NonNullable), + vec![ + flat(3, field_dtype.clone(), 0), + flat(3, field_dtype.clone(), 1), + ], + ) + .into_layout(); + let plan = make_plan(non_nullable)?; + + assert!(plan.is::()); + assert_eq!(plan.child_count(), 2); + assert_eq!(child_of(&plan, 0)?.dtype(), &field_dtype); + assert_eq!(child_of(&plan, 1)?.dtype(), &field_dtype); + + let nullable = StructLayout::new( + 3, + DType::Struct(fields, Nullability::Nullable), + vec![ + flat(3, DType::Bool(Nullability::NonNullable), 2), + flat(3, field_dtype.clone(), 3), + flat(3, field_dtype, 4), + ], + ) + .into_layout(); + let nullable_plan = make_plan(nullable)?; + assert_eq!(nullable_plan.child_count(), 3); + assert_eq!( + child_of(&nullable_plan, 2)?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + Ok(()) +} + +#[test] +fn with_children_rejects_mismatched_arity() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 3, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + ) + .into_layout(); + let plan = make_plan(layout)?; + + let error = plan + .with_children(vec![child_of(&plan, 0)?]) + .err() + .ok_or_else(|| vortex_err!("mismatched arity unexpectedly succeeded"))?; + assert!( + error + .to_string() + .contains("Concat expects 2 children but got 1"), + "unexpected error: {error}" + ); + Ok(()) +} + +#[test] +fn with_children_replaces_children_in_order() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 3, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + ) + .into_layout(); + let plan = make_plan(layout)?; + + let swapped = plan.with_children(vec![child_of(&plan, 1)?, child_of(&plan, 0)?])?; + assert_eq!(child_of(&swapped, 0)?.row_count(), 1); + assert_eq!(child_of(&swapped, 1)?.row_count(), 2); + assert_eq!(swapped.as_::().row_offsets(), &[0, 1]); + Ok(()) +} + +#[test] +fn optimize_drops_identity_expressions() -> VortexResult<()> { + let child = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; + let expression = root().bind(child.dtype())?; + let plan: PlanRef = EvalPlan::new(expression, child).into_plan(); + + assert!(optimize(plan)?.is::()); + Ok(()) +} + +#[test] +fn optimize_rewrites_nested_children() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 3, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + ) + .into_layout(); + let chunked = make_plan(layout)?; + + // Wrap the first chunk in an identity expression, which optimization should remove without + // any chunked-specific rule. + let chunk = child_of(&chunked, 0)?; + let identity: PlanRef = EvalPlan::new(root().bind(chunk.dtype())?, chunk).into_plan(); + let wrapped = chunked.with_children(vec![identity, child_of(&chunked, 1)?])?; + + let optimized = optimize(wrapped)?; + assert!(optimized.is::()); + assert!(child_of(&optimized, 0)?.is::()); + assert!(child_of(&optimized, 1)?.is::()); + Ok(()) +} + +#[test] +fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), + Nullability::NonNullable, + ), + vec![flat(3, field_dtype.clone(), 0), flat(3, field_dtype, 1)], + ) + .into_layout(); + let child = make_plan(layout)?; + let expression = get_item("a", root()).bind(child.dtype())?; + let plan = EvalPlan::new(expression, child); + + assert_eq!(plan.to_string(), "vortex.plan.eval(i32, rows=3) expr=$.a"); + let plan: PlanRef = plan.into_plan(); + + assert_eq!(plan.to_string(), "vortex.plan.eval(i32, rows=3) expr=$.a"); + insta::assert_snapshot!(plan.display_tree(), @r" + root: vortex.plan.eval(i32, rows=3) expr=$.a + child: vortex.plan.pack({a=i32, b=i32}, rows=3) + a: vortex.plan.segment_scan(i32, rows=3) + b: vortex.plan.segment_scan(i32, rows=3) + "); + + struct DepthExtractor; + + impl PlanTreeExtractor for DepthExtractor { + fn write_header( + &self, + _plan: &PlanRef, + context: &PlanTreeContext, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + write!(formatter, " depth={}", context.depth()) + } + } + + insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @r" + root: depth=0 + child: depth=1 + a: depth=2 + b: depth=2 + "); + + let nullable_fields = StructFields::from_iter([ + ("a", primitive(PType::I32, Nullability::NonNullable)), + ("b", primitive(PType::I32, Nullability::NonNullable)), + ]); + let nullable_layout = StructLayout::new( + 3, + DType::Struct(nullable_fields, Nullability::Nullable), + vec![ + flat(3, DType::Bool(Nullability::NonNullable), 2), + flat(3, primitive(PType::I32, Nullability::NonNullable), 3), + flat(3, primitive(PType::I32, Nullability::NonNullable), 4), + ], + ) + .into_layout(); + let nullable = make_plan(nullable_layout)?; + insta::assert_snapshot!(nullable.tree_display_builder(), @r" + root: + a: + b: + validity: + "); + Ok(()) +} + +#[test] +fn chunked_plan_display_names_chunks() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 3, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + ) + .into_layout(); + let plan = make_plan(layout)?; + + insta::assert_snapshot!(plan.display_tree(), @r" + root: vortex.plan.concat(i32, rows=3) + chunks[0]: vortex.plan.segment_scan(i32, rows=2) + chunks[1]: vortex.plan.segment_scan(i32, rows=1) + "); + Ok(()) +} + +#[test] +fn dict_plan_display_names_logical_children() -> VortexResult<()> { + let layout = DictLayout::new( + flat(2, primitive(PType::I32, Nullability::NonNullable), 0), + flat(3, primitive(PType::U8, Nullability::NonNullable), 1), + ) + .into_layout(); + let plan = make_plan(layout)?; + + insta::assert_snapshot!(plan.display_tree(), @r" + root: vortex.plan.take(i32, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.segment_scan(i32, rows=2) + "); + Ok(()) +} + +#[test] +fn list_plan_display_handles_optional_validity() -> VortexResult<()> { + let element_dtype = primitive(PType::I32, Nullability::NonNullable); + let offsets_dtype = primitive(PType::U32, Nullability::NonNullable); + let non_nullable_layout = ListLayout::new( + DType::List(Arc::new(element_dtype.clone()), Nullability::NonNullable), + flat(4, element_dtype.clone(), 0), + flat(3, offsets_dtype.clone(), 1), + None, + ) + .into_layout(); + let non_nullable = make_plan(non_nullable_layout)?; + + insta::assert_snapshot!(non_nullable.display_tree(), @r" + root: vortex.plan.list_pack(list(i32), rows=2) + elements: vortex.plan.segment_scan(i32, rows=4) + offsets: vortex.plan.segment_scan(u32, rows=3) + "); + + let nullable_layout = ListLayout::new( + DType::List(Arc::new(element_dtype.clone()), Nullability::Nullable), + flat(4, element_dtype, 2), + flat(3, offsets_dtype, 3), + Some(flat(2, DType::Bool(Nullability::NonNullable), 4)), + ) + .into_layout(); + let nullable = make_plan(nullable_layout)?; + + insta::assert_snapshot!(nullable.display_tree(), @r" + root: vortex.plan.list_pack(list(i32)?, rows=2) + elements: vortex.plan.segment_scan(i32, rows=4) + offsets: vortex.plan.segment_scan(u32, rows=3) + validity: vortex.plan.segment_scan(bool, rows=2) + "); + Ok(()) +} + +#[test] +fn row_idx_plan_preserves_row_index_expressions() -> VortexResult<()> { + let layout = flat(3, primitive(PType::I32, Nullability::NonNullable), 0); + let plan = RowIdxPlan::new(10, make_plan(layout)?).into_plan(); + let bound_expression = row_idx().bind(plan.dtype())?; + let plan = optimize(EvalPlan::new(bound_expression.clone(), plan).into_plan())?; + let expression = plan + .as_opt::() + .ok_or_else(|| vortex_err!("optimized plan is not an expression plan"))?; + + assert_eq!(expression.expression(), &bound_expression); + assert!(expression.child_plan()?.is::()); + assert_eq!(expression.row_count(), 3); + Ok(()) +} diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/typed.rs new file mode 100644 index 00000000000..133d2ff0d42 --- /dev/null +++ b/vortex-layout/src/plan/typed.rs @@ -0,0 +1,376 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::borrow::Cow; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; +use std::fmt::Formatter; +use std::marker::PhantomData; +use std::ops::Deref; +use std::sync::Arc; + +use vortex_array::SerializeMetadata; +use vortex_array::dtype::DType; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanVTable; +use crate::plan::display::PlanTreeDisplay; + +/// The combined allocation behind [`PlanRef`]. +/// +/// Common plan state is stored before the unsized `data` tail, so reading the operator ID, dtype, +/// row count, or children does not dispatch through the operator vtable. Only `PlanData` is +/// erased to [`DynPlan`]. +struct PlanInner { + id: PlanId, + dtype: DType, + row_count: u64, + children: PlanChildren, + data: D, // must be last for unsized coercion +} + +/// Shared, erased handle to a plan operator. +#[derive(Clone)] +pub struct PlanRef(Arc>); + +impl PlanRef { + fn from_inner(inner: Arc>>) -> Self { + let inner: Arc> = inner; + Self(inner) + } + + fn dyn_plan(&self) -> &dyn DynPlan { + &self.0.data + } + + /// Returns whether two references point at the same plan. + pub fn ptr_eq(lhs: &Self, rhs: &Self) -> bool { + Arc::ptr_eq(&lhs.0, &rhs.0) + } + + /// Returns the operator ID. + pub fn id(&self) -> PlanId { + self.0.id + } + + /// Returns the dtype produced by this plan. + pub fn dtype(&self) -> &DType { + &self.0.dtype + } + + /// Returns the number of rows in this plan's row domain. + pub fn row_count(&self) -> u64 { + self.0.row_count + } + + /// Returns the common child container without initializing any child. + pub fn children(&self) -> &PlanChildren { + &self.0.children + } + + /// Returns the number of children without initializing any child. + pub fn child_count(&self) -> usize { + self.0.children.len() + } + + /// Returns the child at `index`, initializing it on first access. + pub fn child(&self, index: usize) -> VortexResult> { + self.0.children.get(index) + } + + /// Returns the child at `index`, or an error when the index is out of bounds. + pub fn child_required(&self, index: usize) -> VortexResult { + self.child(index)? + .ok_or_else(|| vortex_err!("Missing plan child {index}")) + } + + /// Rebuilds this plan with `children` stored outside its erased operator data. + pub fn with_children(&self, children: impl Into) -> VortexResult { + self.dyn_plan().dyn_with_children(self, children.into()) + } + + /// Rebuilds this plan with one child replaced, preserving laziness in all other slots. + pub fn with_child(&self, index: usize, child: PlanRef) -> VortexResult { + self.with_children(self.children().with_child(index, child)?) + } + + /// Returns the display name of the child at `index`. + pub fn child_name(&self, index: usize) -> Cow<'_, str> { + self.dyn_plan().dyn_child_name(self, index) + } + + /// Serializes operator-specific metadata, or `None` when the operator is not serializable. + pub fn metadata(&self) -> Option> { + self.dyn_plan().dyn_metadata(self) + } + + /// Returns whether this plan uses vtable `V`. + pub fn is(&self) -> bool { + self.dyn_plan().as_any().is::>() + } + + /// Downcasts this plan to vtable `V`. + pub fn as_(&self) -> &Plan { + self.as_opt::().vortex_expect("Failed to downcast") + } + + /// Attempts to borrow this plan as a typed handle for vtable `V`. + pub fn as_opt(&self) -> Option<&Plan> { + if !self.is::() { + return None; + } + + // SAFETY: Plan is transparent over PlanRef, and the type check above proves that its + // erased tail contains PlanData. + Some(unsafe { &*(std::ptr::from_ref(self).cast::>()) }) + } + + /// Displays this plan and its descendants with the default plan extractors. + pub fn display_tree(&self) -> PlanTreeDisplay<'_> { + PlanTreeDisplay::default_display(self) + } + + /// Creates a composable tree display with no extractors. + pub fn tree_display_builder(&self) -> PlanTreeDisplay<'_> { + PlanTreeDisplay::new(self) + } +} + +impl Display for PlanRef { + fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result { + write!( + formatter, + "{}({}, rows={})", + self.id(), + self.dtype(), + self.row_count() + )?; + self.dyn_plan().dyn_fmt(self, formatter) + } +} + +impl Debug for PlanRef { + fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("Plan") + .field("id", &self.0.id) + .field("dtype", &self.0.dtype) + .field("row_count", &self.0.row_count) + .field("children", &self.0.children) + .field("data", &&self.0.data) + .finish() + } +} + +/// Pieces used to construct a typed plan. +pub struct PlanParts { + /// The vtable identifying the operator. + pub vtable: V, + /// Logical dtype produced by the operator. + pub dtype: DType, + /// Number of rows in the operator's row domain. + pub row_count: u64, + /// Child operators, in stable logical order. + pub children: PlanChildren, + /// Operator-specific, non-child data. + pub data: V::PlanData, +} + +impl PlanParts { + /// Converts these parts into a typed plan. + pub fn into_typed(self) -> Plan { + Plan::from_parts(self) + } + + /// Erases these parts into a plan reference. + pub fn into_plan(self) -> PlanRef { + self.into_typed().into_plan() + } +} + +/// A typed, shared handle to a plan operator. +#[repr(transparent)] +pub struct Plan { + inner: PlanRef, + _vtable: PhantomData, +} + +impl Plan { + /// Constructs a plan from explicit parts. + pub fn from_parts(parts: PlanParts) -> Self { + let inner = Arc::new(PlanInner { + id: parts.vtable.id(), + dtype: parts.dtype, + row_count: parts.row_count, + children: parts.children, + data: PlanData { + vtable: parts.vtable, + data: parts.data, + }, + }); + Self { + inner: PlanRef::from_inner(inner), + _vtable: PhantomData, + } + } + + fn typed_data(&self) -> &PlanData { + self.inner + .dyn_plan() + .as_any() + .downcast_ref::>() + .vortex_expect("Typed plan contains the wrong vtable") + } + + /// Returns the vtable. + pub fn vtable(&self) -> &V { + &self.typed_data().vtable + } + + /// Returns operator-specific data. + pub fn data(&self) -> &V::PlanData { + &self.typed_data().data + } + + /// Returns the dtype produced by this plan. + pub fn dtype(&self) -> &DType { + self.inner.dtype() + } + + /// Returns the number of rows in this plan's row domain. + pub fn row_count(&self) -> u64 { + self.inner.row_count() + } + + /// Returns the common child container without initializing any child. + pub fn children(&self) -> &PlanChildren { + self.inner.children() + } + + /// Returns a child, initializing it on first access. + pub fn child(&self, index: usize) -> VortexResult> { + self.inner.child(index) + } + + /// Returns the child at `index`, or an error when the index is out of bounds. + pub fn child_required(&self, index: usize) -> VortexResult { + self.inner.child_required(index) + } + + /// Erases this typed plan into a shared reference. + pub fn to_plan(&self) -> PlanRef { + self.inner.clone() + } + + /// Erases this typed plan into a shared reference. + pub fn into_plan(self) -> PlanRef { + self.inner + } +} + +impl Clone for Plan { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + _vtable: PhantomData, + } + } +} + +impl Debug for Plan { + fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result { + Debug::fmt(&self.inner, formatter) + } +} + +impl Display for Plan { + fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(&self.inner, formatter) + } +} + +impl Deref for Plan { + type Target = V::PlanData; + + fn deref(&self) -> &Self::Target { + self.data() + } +} + +impl From> for PlanRef { + fn from(value: Plan) -> Self { + value.into_plan() + } +} + +/// A vtable value paired with its operator-specific plan data. +struct PlanData { + vtable: V, + data: V::PlanData, +} + +impl Debug for PlanData { + fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PlanData") + .field("vtable", &self.vtable) + .field("data", &self.data) + .finish() + } +} + +/// Erased operator-specific behavior stored in the unsized tail of a [`PlanRef`]. +#[doc(hidden)] +pub trait DynPlan: 'static + Send + Sync + Debug { + /// Returns this operator data as [`Any`] for downcasting. + fn as_any(&self) -> &dyn Any; + + /// Formats operator-specific fields. + fn dyn_fmt(&self, plan: &PlanRef, formatter: &mut Formatter<'_>) -> fmt::Result; + + /// Clones operator data, runs its child-replacement callback, and rebuilds the common node. + fn dyn_with_children(&self, plan: &PlanRef, children: PlanChildren) -> VortexResult; + + /// Returns the display name of the child at `index`. + fn dyn_child_name<'a>(&'a self, plan: &'a PlanRef, index: usize) -> Cow<'a, str>; + + /// Serializes operator-specific metadata, or `None` when the operator is not serializable. + fn dyn_metadata(&self, plan: &PlanRef) -> Option>; +} + +impl DynPlan for PlanData { + fn as_any(&self) -> &dyn Any { + self + } + + fn dyn_fmt(&self, plan: &PlanRef, formatter: &mut Formatter<'_>) -> fmt::Result { + ::fmt(plan.as_::(), formatter) + } + + fn dyn_with_children(&self, plan: &PlanRef, children: PlanChildren) -> VortexResult { + let mut data = self.data.clone(); + V::with_children(plan.as_::(), &children, &mut data)?; + Ok(PlanParts { + vtable: self.vtable.clone(), + dtype: plan.dtype().clone(), + row_count: plan.row_count(), + children, + data, + } + .into_plan()) + } + + fn dyn_child_name<'a>(&'a self, plan: &'a PlanRef, index: usize) -> Cow<'a, str> { + V::child_name(plan.as_::(), index) + } + + fn dyn_metadata(&self, plan: &PlanRef) -> Option> { + V::metadata(plan.as_::()).map(SerializeMetadata::serialize) + } +} diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/vtable.rs new file mode 100644 index 00000000000..caf004383b5 --- /dev/null +++ b/vortex-layout/src/plan/vtable.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::fmt; +use std::fmt::Debug; + +use vortex_array::DeserializeMetadata; +use vortex_array::SerializeMetadata; +use vortex_error::VortexResult; +use vortex_session::registry::Id; + +use crate::plan::PlanChildren; +use crate::plan::typed::Plan; + +/// A unique identifier for a plan operator. +pub type PlanId = Id; + +/// Operator-specific behavior for a typed [`Plan`]. +/// +/// Common fields — dtype, row count, and children — are stored outside the erased operator data. +/// Implementations own only their operator-specific data, its metadata codec, and a callback for +/// refreshing cached data after generic child replacement. +/// +/// Operators describe physical work over a row domain. Their identity and operator-specific data +/// do not depend on the source layout kind. The common lazy-child storage may nevertheless own a +/// hidden source-layout reference used for on-demand lowering. +pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { + /// Operator-specific data, excluding children. + /// + /// Children belong in [`PlanParts::children`](crate::plan::PlanParts::children) so that + /// traversal and rewriting can discover them generically. + type PlanData: 'static + Send + Sync + Clone + Debug; + + /// Serialized form of [`PlanData`](Self::PlanData). + type Metadata: SerializeMetadata + DeserializeMetadata + Debug; + + /// Returns the globally unique operator ID. + fn id(&self) -> PlanId; + + /// Writes operator-specific fields after the plan's standard display summary. + fn fmt(plan: &Plan, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let _ = (plan, formatter); + Ok(()) + } + + /// Returns the serializable metadata for this operator. + /// + /// Returns `None` when the operator holds state that cannot be serialized. + fn metadata(plan: &Plan) -> Option; + + /// Refreshes cloned operator data after the common child container is replaced. + /// + /// The plan layer clones [`PlanData`](Self::PlanData), replaces the children externally, and + /// invokes this callback. Implementations validate the new children and update any derived + /// values in `data`; they do not rebuild the plan itself. + fn with_children( + plan: &Plan, + children: &PlanChildren, + data: &mut Self::PlanData, + ) -> VortexResult<()> { + let _ = (plan, children, data); + Ok(()) + } + + /// Returns the display name of the child at `index`. + fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { + let _ = plan; + Cow::Owned(format!("child[{index}]")) + } +} From bc0495fc2f8ff382ffcf505ce9f5166e87074fab Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 11 Aug 2026 17:31:50 +0100 Subject: [PATCH 02/13] fix up Signed-off-by: Joe Isaacs --- vortex-array/src/expr/bound_expression.rs | 30 ++++++ vortex-layout/src/plan/children.rs | 18 ---- vortex-layout/src/plan/lower.rs | 60 +++++++---- vortex-layout/src/plan/plans/concat.rs | 22 ++-- vortex-layout/src/plan/plans/eval.rs | 26 ++++- vortex-layout/src/plan/plans/list_pack.rs | 97 +++++++++++++++--- vortex-layout/src/plan/plans/pack.rs | 119 +++++++++++++++++++--- vortex-layout/src/plan/plans/take.rs | 27 +++-- vortex-layout/src/plan/tests.rs | 25 ++++- vortex-layout/src/plan/typed.rs | 5 - 10 files changed, 333 insertions(+), 96 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index bc74633608a..45040ab2cd1 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -226,6 +226,22 @@ impl BoundExpression { matches!(self.kind, BoundKind::Root) } + /// Return whether every scope root in this expression has `dtype`. + /// + /// Expressions without a scope root, such as literals, match every dtype. + pub fn is_root_bound_to(&self, dtype: &DType) -> bool { + let mut is_bound_to = true; + pre_order_visit_down(self, |node| { + if node.is_root() && node.dtype() != dtype { + is_bound_to = false; + return Ok(TraversalOrder::Stop); + } + Ok(TraversalOrder::Continue) + }) + .vortex_expect("bound expression traversal cannot not fail"); + is_bound_to + } + /// Return an expression that proves this predicate is definitely false from statistics. pub fn falsify(&self, session: &VortexSession) -> VortexResult> { StatsRewriteCtx::new(session).falsify(self) @@ -363,6 +379,20 @@ mod tests { Ok(()) } + #[test] + fn bound_to_checks_every_root() -> VortexResult<()> { + let dtype = struct_dtype(); + let bound = eq(col("a"), col("a")).bind(&dtype)?; + assert!(bound.is_root_bound_to(&dtype)); + assert!(!bound.is_root_bound_to(&DType::Bool(Nullability::NonNullable))); + assert!( + lit(true) + .bind(&dtype)? + .is_root_bound_to(&DType::Bool(Nullability::NonNullable)) + ); + Ok(()) + } + #[test] fn bound_display_matches_unbound() -> VortexResult<()> { for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { diff --git a/vortex-layout/src/plan/children.rs b/vortex-layout/src/plan/children.rs index 4b639a57fd8..b045c1be8d2 100644 --- a/vortex-layout/src/plan/children.rs +++ b/vortex-layout/src/plan/children.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use once_cell::sync::OnceCell; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_err; use crate::plan::PlanRef; @@ -73,23 +72,6 @@ impl PlanChildren { pub fn to_vec(&self) -> VortexResult> { self.iter().collect() } - - /// Returns a child collection with one slot replaced. - pub fn with_child(&self, index: usize, child: PlanRef) -> VortexResult { - if index >= self.len() { - vortex_bail!("Plan child index out of bounds: {index} of {}", self.len()); - } - - let source = self.clone(); - Ok(Self::lazy(source.len(), move |child_index| { - if child_index == index { - return Ok(child.clone()); - } - source - .get(child_index)? - .ok_or_else(|| vortex_err!("Plan child {child_index} is absent")) - })) - } } impl From> for PlanChildren { diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index 82c36b584e5..e8911734549 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -77,12 +77,16 @@ fn lower_chunked(layout: &ChunkedLayout) -> VortexResult { .checked_add(layout.child_row_count(index)) .ok_or_else(|| vortex_err!("Chunked row count overflow"))?; } - Ok(ConcatPlan::from_children( - layout.dtype().clone(), - layout.row_count(), - row_offsets.into(), - lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), - )) + // SAFETY: Chunked layout construction validates that every child has the parent dtype and + // that their row counts sum to the parent; offsets were computed with checked addition above. + Ok(unsafe { + ConcatPlan::from_children_unchecked( + layout.dtype().clone(), + layout.row_count(), + row_offsets.into(), + lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), + ) + }) } fn lower_struct(layout: &StructLayout) -> VortexResult { @@ -93,22 +97,30 @@ fn lower_struct(layout: &StructLayout) -> VortexResult { if layout.dtype().is_nullable() { slots.push(0); } - Ok(PackPlan::from_children( - fields, - layout.dtype().nullability(), - layout.row_count(), - lazy_children(layout.to_layout(), slots), - )) + // SAFETY: Struct layout construction validates child count and row counts, and its logical + // slot schema defines every field dtype plus the optional non-nullable boolean validity dtype. + Ok(unsafe { + PackPlan::from_children_unchecked( + fields, + layout.dtype().nullability(), + layout.row_count(), + lazy_children(layout.to_layout(), slots), + ) + }) } fn lower_dict(layout: &DictLayout) -> VortexResult { // Dict serialization stores values before codes; the plan order is deliberately codes, // values because that is the optimizer-facing logical shape. - Ok(TakePlan::from_children( - layout.dtype().clone(), - layout.row_count(), - lazy_children(layout.to_layout(), vec![1, 0]), - )) + // SAFETY: Dict layout construction validates its values and codes slots. The plan reorders + // those slots to [codes, values] while preserving the derived output dtype and row domain. + Ok(unsafe { + TakePlan::from_children_unchecked( + layout.dtype().clone(), + layout.row_count(), + lazy_children(layout.to_layout(), vec![1, 0]), + ) + }) } fn lower_list(layout: &ListLayout) -> VortexResult { @@ -116,11 +128,15 @@ fn lower_list(layout: &ListLayout) -> VortexResult { if layout.dtype().is_nullable() { slots.push(VALIDITY_CHILD_INDEX); } - Ok(ListPackPlan::from_children( - layout.dtype().clone(), - layout.row_count(), - lazy_children(layout.to_layout(), slots), - )) + // SAFETY: List layout construction validates its element, offsets, and optional validity + // child shapes before this plan preserves them in the same logical order. + Ok(unsafe { + ListPackPlan::from_children_unchecked( + layout.dtype().clone(), + layout.row_count(), + lazy_children(layout.to_layout(), slots), + ) + }) } fn lazy_children(layout: LayoutRef, slots: Vec) -> PlanChildren { diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 414a685ec5d..2e2893fde5f 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -31,7 +31,13 @@ pub struct ConcatData { pub type ConcatPlan = Plan; impl ConcatPlan { - pub(crate) fn from_children( + /// Creates a concatenation from potentially unresolved children without validation. + /// + /// # Safety + /// + /// Every child must have `dtype`; `row_offsets` must contain the cumulative row offset of + /// every child; and the sum of all child row counts must equal `row_count` without overflow. + pub(crate) unsafe fn from_children_unchecked( dtype: DType, row_count: u64, row_offsets: Arc<[u64]>, @@ -61,14 +67,14 @@ impl ConcatPlan { ); } row_offsets.push(row_count); - row_count += child.row_count(); + row_count = row_count + .checked_add(child.row_count()) + .ok_or_else(|| vortex_error::vortex_err!("Concat row count overflow"))?; } - Ok(Self::from_children( - dtype, - row_count, - row_offsets.into(), - children.into(), - )) + // SAFETY: Child dtypes and the checked cumulative row metadata were validated above. + Ok(unsafe { + Self::from_children_unchecked(dtype, row_count, row_offsets.into(), children.into()) + }) } /// Returns the first row of each child within this plan's row domain. diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index 05421eeda74..70a645da687 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -7,6 +7,7 @@ use std::fmt; use vortex_array::EmptyMetadata; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_session::registry::CachedId; use crate::plan::Plan; @@ -32,7 +33,19 @@ pub type EvalPlan = Plan; impl EvalPlan { /// Creates an evaluation of `expression`, which must be bound to the child's dtype. - pub fn new(expression: BoundExpression, child: PlanRef) -> Self { + pub fn try_new(expression: BoundExpression, child: PlanRef) -> VortexResult { + validate_expression_child(&expression, &child)?; + + // SAFETY: The expression root dtype was validated against the child dtype above. + Ok(unsafe { Self::new_unchecked(expression, child) }) + } + + /// Creates an evaluation without validating the expression's root dtype. + /// + /// # Safety + /// + /// Every scope root in `expression` must have the same dtype as `child`. + pub unsafe fn new_unchecked(expression: BoundExpression, child: PlanRef) -> Self { PlanParts { vtable: Eval, dtype: expression.dtype().clone(), @@ -81,6 +94,7 @@ impl PlanVTable for Eval { let child = children .get(0)? .ok_or_else(|| vortex_error::vortex_err!("Eval child is absent"))?; + validate_expression_child(plan.expression(), &child)?; if child.row_count() != plan.row_count() { vortex_error::vortex_bail!( "Eval child has {} rows but the plan has {}", @@ -99,3 +113,13 @@ impl PlanVTable for Eval { } } } + +fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> { + if !expression.is_root_bound_to(child.dtype()) { + vortex_bail!( + "Eval expression is not bound to child dtype {}", + child.dtype() + ); + } + Ok(()) +} diff --git a/vortex-layout/src/plan/plans/list_pack.rs b/vortex-layout/src/plan/plans/list_pack.rs index f62f8c1994b..d9a23de818b 100644 --- a/vortex-layout/src/plan/plans/list_pack.rs +++ b/vortex-layout/src/plan/plans/list_pack.rs @@ -9,6 +9,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::plan::Plan; @@ -34,7 +35,18 @@ pub struct ListPackData; pub type ListPackPlan = Plan; impl ListPackPlan { - pub(crate) fn from_children(dtype: DType, row_count: u64, children: PlanChildren) -> Self { + /// Creates a list assembly from potentially unresolved children without validation. + /// + /// # Safety + /// + /// `dtype` must be a list whose element dtype matches the elements child. The offsets child + /// must be a non-nullable integer with `row_count + 1` rows. A non-nullable boolean validity + /// child with `row_count` rows must be present exactly when `dtype` is nullable. + pub(crate) unsafe fn from_children_unchecked( + dtype: DType, + row_count: u64, + children: PlanChildren, + ) -> Self { PlanParts { vtable: ListPack, dtype, @@ -56,15 +68,14 @@ impl ListPackPlan { offsets: PlanRef, validity: Option, ) -> VortexResult { - if validity.is_some() != (nullability == Nullability::Nullable) { - vortex_bail!( - "ListPack validity child must be present exactly when the list is nullable" - ); - } let dtype = DType::List(Arc::new(elements.dtype().clone()), nullability); let mut children = vec![elements, offsets]; children.extend(validity); - Ok(Self::from_children(dtype, row_count, children.into())) + let children = PlanChildren::from(children); + validate_children(&dtype, row_count, &children)?; + + // SAFETY: All child shape invariants were validated above. + Ok(unsafe { Self::from_children_unchecked(dtype, row_count, children) }) } /// Returns the plan producing list elements. @@ -102,14 +113,7 @@ impl PlanVTable for ListPack { children: &PlanChildren, _data: &mut Self::PlanData, ) -> VortexResult<()> { - if children.len() != plan.children().len() { - vortex_bail!( - "ListPack expects {} children but got {}", - plan.children().len(), - children.len() - ); - } - Ok(()) + validate_children(plan.dtype(), plan.row_count(), children) } fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { @@ -121,3 +125,66 @@ impl PlanVTable for ListPack { } } } + +fn validate_children(dtype: &DType, row_count: u64, children: &PlanChildren) -> VortexResult<()> { + let elements_dtype = dtype + .as_list_element_opt() + .ok_or_else(|| vortex_err!("ListPack output dtype must be a list, got {dtype}"))?; + let expected_children = 2 + usize::from(dtype.is_nullable()); + if children.len() != expected_children { + vortex_bail!( + "ListPack expects {expected_children} children but got {}", + children.len() + ); + } + + let elements = children + .get(ELEMENTS)? + .ok_or_else(|| vortex_err!("ListPack elements child is absent"))?; + if elements.dtype() != elements_dtype.as_ref() { + vortex_bail!( + "ListPack elements child has dtype {} but the list element dtype is {}", + elements.dtype(), + elements_dtype + ); + } + + let offsets = children + .get(OFFSETS)? + .ok_or_else(|| vortex_err!("ListPack offsets child is absent"))?; + if !offsets.dtype().is_int() || offsets.dtype().is_nullable() { + vortex_bail!( + "ListPack offsets child must have a non-nullable integer dtype, got {}", + offsets.dtype() + ); + } + let offsets_row_count = row_count + .checked_add(1) + .ok_or_else(|| vortex_err!("ListPack offsets row count overflow"))?; + if offsets.row_count() != offsets_row_count { + vortex_bail!( + "ListPack offsets child has {} rows but must have {offsets_row_count}", + offsets.row_count() + ); + } + + if dtype.is_nullable() { + let validity = children + .get(VALIDITY)? + .ok_or_else(|| vortex_err!("ListPack validity child is absent"))?; + let validity_dtype = DType::Bool(Nullability::NonNullable); + if validity.dtype() != &validity_dtype { + vortex_bail!( + "ListPack validity child has dtype {} but must have dtype {validity_dtype}", + validity.dtype() + ); + } + if validity.row_count() != row_count { + vortex_bail!( + "ListPack validity child has {} rows but the plan has {row_count}", + validity.row_count() + ); + } + } + Ok(()) +} diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 156fcd5af4c..9cf47ebe4b7 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -10,6 +10,7 @@ use vortex_array::dtype::StructFields; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::plan::Plan; @@ -31,7 +32,14 @@ pub struct PackData; pub type PackPlan = Plan; impl PackPlan { - pub(crate) fn from_children( + /// Creates a struct assembly from potentially unresolved children without validation. + /// + /// # Safety + /// + /// `children` must contain one child per field, followed by a non-nullable boolean validity + /// child exactly when `nullability` is nullable. Every child must have `row_count` rows, and + /// every field child must have its corresponding field dtype. + pub(crate) unsafe fn from_children_unchecked( fields: StructFields, nullability: Nullability, row_count: u64, @@ -49,7 +57,9 @@ impl PackPlan { /// Creates a struct assembly from `fields` and one child per field. /// - /// `validity` is required exactly when `nullability` is [`Nullability::Nullable`]. + /// Each field child must have the corresponding field dtype and `row_count` rows. `validity` + /// is required exactly when `nullability` is [`Nullability::Nullable`] and must produce + /// non-nullable booleans with `row_count` rows. pub fn try_new( fields: StructFields, nullability: Nullability, @@ -68,14 +78,40 @@ impl PackPlan { vortex_bail!("Pack validity child must be present exactly when the struct is nullable"); } + for (index, (field_dtype, field_plan)) in + fields.fields().zip(field_plans.iter()).enumerate() + { + validate_field_child(index, &field_dtype, row_count, field_plan)?; + } + if let Some(validity) = validity.as_ref() { + validate_validity_child(row_count, validity)?; + } + + // SAFETY: The child count, presence, dtypes, and row counts were validated above. + Ok(unsafe { Self::new_unchecked(fields, nullability, row_count, field_plans, validity) }) + } + + /// Creates a struct assembly without validating its children. + /// + /// # Safety + /// + /// The caller must ensure that: + /// + /// - `field_plans` contains exactly one child per field, in field order; + /// - every field child has the corresponding field dtype and `row_count` rows; and + /// - `validity` is present exactly when the struct is nullable and, when present, produces + /// non-nullable booleans with `row_count` rows. + pub unsafe fn new_unchecked( + fields: StructFields, + nullability: Nullability, + row_count: u64, + field_plans: Vec, + validity: Option, + ) -> Self { let mut children = field_plans; children.extend(validity); - Ok(Self::from_children( - fields, - nullability, - row_count, - children.into(), - )) + // SAFETY: The caller guarantees the same child invariants required by this constructor. + unsafe { Self::from_children_unchecked(fields, nullability, row_count, children.into()) } } /// Returns the struct fields assembled by this plan. @@ -115,23 +151,76 @@ impl PlanVTable for Pack { children: &PlanChildren, _data: &mut Self::PlanData, ) -> VortexResult<()> { - if children.len() != plan.children().len() { + let expected_children = plan.nfields() + usize::from(plan.dtype().is_nullable()); + if children.len() != expected_children { vortex_bail!( - "Pack expects {} children but got {}", - plan.children().len(), + "Pack expects {expected_children} children but got {}", children.len() ); } + + for (index, field_dtype) in plan.fields().fields().enumerate() { + let child = children + .get(index)? + .ok_or_else(|| vortex_err!("Pack field child {index} is absent"))?; + validate_field_child(index, &field_dtype, plan.row_count(), &child)?; + } + if plan.dtype().is_nullable() { + let validity = children + .get(plan.nfields())? + .ok_or_else(|| vortex_err!("Pack validity child is absent"))?; + validate_validity_child(plan.row_count(), &validity)?; + } Ok(()) } fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { + assert!( + index < plan.children().len(), + "Pack child index out of bounds: {index} of {}", + plan.children().len() + ); if let Some(name) = plan.fields().field_name(index) { return Cow::Borrowed(name.as_ref()); } - if index == plan.fields().nfields() { - return Cow::Borrowed("validity"); - } - Cow::Owned(format!("child[{index}]")) + Cow::Borrowed("validity") + } +} + +fn validate_field_child( + index: usize, + expected_dtype: &DType, + expected_row_count: u64, + child: &PlanRef, +) -> VortexResult<()> { + if child.dtype() != expected_dtype { + vortex_bail!( + "Pack field child {index} has dtype {} but the field has dtype {expected_dtype}", + child.dtype() + ); + } + if child.row_count() != expected_row_count { + vortex_bail!( + "Pack field child {index} has {} rows but the plan has {expected_row_count}", + child.row_count() + ); + } + Ok(()) +} + +fn validate_validity_child(expected_row_count: u64, child: &PlanRef) -> VortexResult<()> { + let expected_dtype = DType::Bool(Nullability::NonNullable); + if child.dtype() != &expected_dtype { + vortex_bail!( + "Pack validity child has dtype {} but must have dtype {expected_dtype}", + child.dtype() + ); + } + if child.row_count() != expected_row_count { + vortex_bail!( + "Pack validity child has {} rows but the plan has {expected_row_count}", + child.row_count() + ); } + Ok(()) } diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/plans/take.rs index 1184db10e65..1e68281029b 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/plans/take.rs @@ -4,6 +4,7 @@ use std::borrow::Cow; use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -26,8 +27,14 @@ pub struct Take; pub type TakePlan = Plan; impl TakePlan { - pub(crate) fn from_children( - dtype: vortex_array::dtype::DType, + /// Creates a take from potentially unresolved children without validation. + /// + /// # Safety + /// + /// `children` must be `[codes, values]`; `codes` must have `row_count` rows; and `dtype` must + /// be the values dtype unioned with the codes nullability. + pub(crate) unsafe fn from_children_unchecked( + dtype: DType, row_count: u64, children: PlanChildren, ) -> Self { @@ -45,11 +52,12 @@ impl TakePlan { /// /// The row domain is that of `codes`, and the output dtype is that of `values`. pub fn new(codes: PlanRef, values: PlanRef) -> Self { - Self::from_children( - values.dtype().clone(), - codes.row_count(), - vec![codes, values].into(), - ) + let dtype = values + .dtype() + .union_nullability(codes.dtype().nullability()); + let row_count = codes.row_count(); + // SAFETY: Parent metadata is derived from the ordered children immediately above. + unsafe { Self::from_children_unchecked(dtype, row_count, vec![codes, values].into()) } } /// Returns the plan producing indices. @@ -88,7 +96,10 @@ impl PlanVTable for Take { let values = children .get(VALUES)? .ok_or_else(|| vortex_error::vortex_err!("Take values child is absent"))?; - if codes.row_count() != plan.row_count() || values.dtype() != plan.dtype() { + let dtype = values + .dtype() + .union_nullability(codes.dtype().nullability()); + if codes.row_count() != plan.row_count() || &dtype != plan.dtype() { vortex_error::vortex_bail!("Take child shape does not match the plan output"); } Ok(()) diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 9c478317ee8..019ebbaf1ba 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -283,11 +283,28 @@ fn with_children_replaces_children_in_order() -> VortexResult<()> { Ok(()) } +#[test] +fn eval_try_new_validates_expression_root_dtype() -> VortexResult<()> { + let expression = root().bind(&primitive(PType::I32, Nullability::NonNullable))?; + let child = make_plan(flat(3, primitive(PType::I64, Nullability::NonNullable), 0))?; + + let error = EvalPlan::try_new(expression, child) + .err() + .ok_or_else(|| vortex_err!("mismatched Eval root dtype unexpectedly succeeded"))?; + assert!( + error + .to_string() + .contains("Eval expression is not bound to child dtype i64"), + "unexpected error: {error}" + ); + Ok(()) +} + #[test] fn optimize_drops_identity_expressions() -> VortexResult<()> { let child = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; let expression = root().bind(child.dtype())?; - let plan: PlanRef = EvalPlan::new(expression, child).into_plan(); + let plan: PlanRef = EvalPlan::try_new(expression, child)?.into_plan(); assert!(optimize(plan)?.is::()); Ok(()) @@ -307,7 +324,7 @@ fn optimize_rewrites_nested_children() -> VortexResult<()> { // Wrap the first chunk in an identity expression, which optimization should remove without // any chunked-specific rule. let chunk = child_of(&chunked, 0)?; - let identity: PlanRef = EvalPlan::new(root().bind(chunk.dtype())?, chunk).into_plan(); + let identity: PlanRef = EvalPlan::try_new(root().bind(chunk.dtype())?, chunk)?.into_plan(); let wrapped = chunked.with_children(vec![identity, child_of(&chunked, 1)?])?; let optimized = optimize(wrapped)?; @@ -331,7 +348,7 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { .into_layout(); let child = make_plan(layout)?; let expression = get_item("a", root()).bind(child.dtype())?; - let plan = EvalPlan::new(expression, child); + let plan = EvalPlan::try_new(expression, child)?; assert_eq!(plan.to_string(), "vortex.plan.eval(i32, rows=3) expr=$.a"); let plan: PlanRef = plan.into_plan(); @@ -466,7 +483,7 @@ fn row_idx_plan_preserves_row_index_expressions() -> VortexResult<()> { let layout = flat(3, primitive(PType::I32, Nullability::NonNullable), 0); let plan = RowIdxPlan::new(10, make_plan(layout)?).into_plan(); let bound_expression = row_idx().bind(plan.dtype())?; - let plan = optimize(EvalPlan::new(bound_expression.clone(), plan).into_plan())?; + let plan = optimize(EvalPlan::try_new(bound_expression.clone(), plan)?.into_plan())?; let expression = plan .as_opt::() .ok_or_else(|| vortex_err!("optimized plan is not an expression plan"))?; diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/typed.rs index 133d2ff0d42..bf544801263 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/typed.rs @@ -95,11 +95,6 @@ impl PlanRef { self.dyn_plan().dyn_with_children(self, children.into()) } - /// Rebuilds this plan with one child replaced, preserving laziness in all other slots. - pub fn with_child(&self, index: usize, child: PlanRef) -> VortexResult { - self.with_children(self.children().with_child(index, child)?) - } - /// Returns the display name of the child at `index`. pub fn child_name(&self, index: usize) -> Cow<'_, str> { self.dyn_plan().dyn_child_name(self, index) From 2eeee0760c571071583bf7f58c91f54f049378f0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 10 Aug 2026 11:04:52 +0100 Subject: [PATCH 03/13] Add plan parent-reduction rule API Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/mod.rs | 1 + vortex-layout/src/plan/optimizer/mod.rs | 11 ++ vortex-layout/src/plan/optimizer/rules.rs | 149 ++++++++++++++++++++++ 3 files changed, 161 insertions(+) create mode 100644 vortex-layout/src/plan/optimizer/mod.rs create mode 100644 vortex-layout/src/plan/optimizer/rules.rs diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 06f7acd0031..11d94805365 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -11,6 +11,7 @@ mod children; mod display; mod lower; mod optimize; +pub mod optimizer; mod plans; mod typed; mod vtable; diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/optimizer/mod.rs new file mode 100644 index 00000000000..c9df9fb547e --- /dev/null +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Static parent-child rewrite rules for physical plans. + +mod rules; + +pub use rules::DynPlanParentReduceRule; +pub use rules::PlanParentReduceRule; +pub use rules::PlanParentReduceRuleAdapter; +pub use rules::PlanParentRuleSet; diff --git a/vortex-layout/src/plan/optimizer/rules.rs b/vortex-layout/src/plan/optimizer/rules.rs new file mode 100644 index 00000000000..0e4d1a70bd2 --- /dev/null +++ b/vortex-layout/src/plan/optimizer/rules.rs @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Typed and type-erased interfaces for parent-child plan rewrites. + +use std::any::type_name; +use std::fmt::Debug; +use std::marker::PhantomData; + +use vortex_error::VortexResult; + +use crate::plan::Plan; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; + +/// A metadata-only rewrite where a child plan rewrites its parent plan. +/// +/// Rules return one rewrite without recursively optimizing the replacement. The plan optimizer +/// owns traversal and drives further rewrites. +pub trait PlanParentReduceRule: Debug + Send + Sync + 'static { + /// The concrete parent operator matched by this rule. + type Parent: PlanVTable; + + /// Attempts to replace `parent` based on its child at `child_idx`. + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + child_idx: usize, + ) -> VortexResult>; +} + +/// Type-erased interface used by [`PlanParentRuleSet`]. +pub trait DynPlanParentReduceRule: Debug + Send + Sync + 'static { + /// Returns whether this rule supports the concrete child and parent operators. + fn matches(&self, child: &PlanRef, parent: &PlanRef) -> bool; + + /// Attempts to replace `parent` based on `child` at `child_idx`. + fn reduce_parent( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult>; +} + +/// Bridges a typed [`PlanParentReduceRule`] to a type-erased static registry. +pub struct PlanParentReduceRuleAdapter { + rule: R, + _child: PhantomData C>, +} + +impl PlanParentReduceRuleAdapter { + /// Creates an adapter for a typed parent-child rule. + pub const fn new(rule: R) -> Self { + Self { + rule, + _child: PhantomData, + } + } +} + +impl Debug for PlanParentReduceRuleAdapter +where + C: PlanVTable, + R: PlanParentReduceRule, +{ + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PlanParentReduceRuleAdapter") + .field("parent", &type_name::()) + .field("child", &type_name::()) + .field("rule", &self.rule) + .finish() + } +} + +impl DynPlanParentReduceRule for PlanParentReduceRuleAdapter +where + C: PlanVTable, + R: PlanParentReduceRule, +{ + fn matches(&self, child: &PlanRef, parent: &PlanRef) -> bool { + child.is::() && parent.is::() + } + + fn reduce_parent( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult> { + let Some(child) = child.as_opt::() else { + return Ok(None); + }; + let Some(parent) = parent.as_opt::() else { + return Ok(None); + }; + self.rule.reduce_parent(child, parent, child_idx) + } +} + +/// An ordered static collection of parent-child plan rewrite rules. +pub struct PlanParentRuleSet { + rules: &'static [&'static dyn DynPlanParentReduceRule], +} + +impl PlanParentRuleSet { + /// Creates a rule set whose first successful rewrite wins. + pub const fn new(rules: &'static [&'static dyn DynPlanParentReduceRule]) -> Self { + Self { rules } + } + + /// Evaluates rules registered for the concrete `(parent, child)` pair. + pub fn evaluate( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult> { + for rule in self.rules { + if !rule.matches(child, parent) { + continue; + } + let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? else { + continue; + }; + + #[cfg(debug_assertions)] + { + vortex_error::vortex_ensure!( + reduced.row_count() == parent.row_count(), + "Plan rewrite from {rule:?} changed row count from {} to {}", + parent.row_count(), + reduced.row_count() + ); + vortex_error::vortex_ensure!( + reduced.dtype() == parent.dtype(), + "Plan rewrite from {rule:?} changed dtype from {} to {}", + parent.dtype(), + reduced.dtype() + ); + } + + return Ok(Some(reduced)); + } + Ok(None) + } +} From dab471074d70f8d09d5cffec87c50076ee9d65e1 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 10 Aug 2026 11:05:18 +0100 Subject: [PATCH 04/13] Push expressions through plan operators Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/optimize.rs | 34 ++- vortex-layout/src/plan/optimizer/mod.rs | 30 +++ vortex-layout/src/plan/plans/concat.rs | 48 ++++ vortex-layout/src/plan/plans/eval.rs | 100 +++++++++ vortex-layout/src/plan/plans/mod.rs | 5 +- vortex-layout/src/plan/plans/pack.rs | 284 ++++++++++++++++++++++++ vortex-layout/src/plan/plans/take.rs | 49 ++++ vortex-layout/src/plan/tests.rs | 274 +++++++++++++++++++++++ 8 files changed, 802 insertions(+), 22 deletions(-) diff --git a/vortex-layout/src/plan/optimize.rs b/vortex-layout/src/plan/optimize.rs index d78798cdb65..8efd498ca08 100644 --- a/vortex-layout/src/plan/optimize.rs +++ b/vortex-layout/src/plan/optimize.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Generic bottom-up optimization over physical plans. +//! Plan optimization. +//! +//! Optimization is driven top-down from [`Eval`] nodes, which apply the static parent-reduction +//! rules in [`crate::plan::optimizer`] as they become applicable. Operators without a rule simply +//! optimize their children. use vortex_error::VortexResult; @@ -10,26 +14,14 @@ use crate::plan::PlanRef; /// Optimizes `plan`, preserving its dtype and row domain. pub fn optimize(plan: PlanRef) -> VortexResult { - let mut children = Vec::with_capacity(plan.child_count()); - let mut changed = false; - for child in plan.children().iter() { - let child = child?; - let optimized = optimize(child.clone())?; - changed |= !PlanRef::ptr_eq(&child, &optimized); - children.push(optimized); + if let Some(eval) = plan.as_opt::() { + return eval.optimize_top_down(None); } - let plan = if changed { - plan.with_children(children)? - } else { - plan - }; - - let Some(eval) = plan.as_opt::() else { - return Ok(plan); - }; - if eval.expression().is_root() { - return eval.child_plan(); - } - Ok(plan) + let children = plan + .children() + .iter() + .map(|child| optimize(child?)) + .collect::>>()?; + plan.with_children(children) } diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/optimizer/mod.rs index c9df9fb547e..a2c2d359804 100644 --- a/vortex-layout/src/plan/optimizer/mod.rs +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -9,3 +9,33 @@ pub use rules::DynPlanParentReduceRule; pub use rules::PlanParentReduceRule; pub use rules::PlanParentReduceRuleAdapter; pub use rules::PlanParentRuleSet; +use vortex_error::VortexResult; + +use super::Concat; +use super::Pack; +use super::PlanRef; +use super::Take; +use super::plans::ExpressionConcatRule; +use super::plans::ExpressionPackRule; +use super::plans::ExpressionTakeRule; + +static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionConcatRule); +static EXPRESSION_TAKE_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionTakeRule); +static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionPackRule); + +static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ + &EXPRESSION_CONCAT_RULE, + &EXPRESSION_TAKE_RULE, + &EXPRESSION_PACK_RULE, +]); + +/// Attempts a static rewrite for `parent` and its child at `child_idx`. +pub(crate) fn reduce_parent(parent: &PlanRef, child_idx: usize) -> VortexResult> { + let Some(child) = parent.child(child_idx)? else { + return Ok(None); + }; + PARENT_RULES.evaluate(&child, parent, child_idx) +} diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 2e2893fde5f..697409690ba 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -6,16 +6,22 @@ use std::sync::Arc; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::label_bound_tree; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; +use crate::layouts::row_idx::RowIdx as RowIdxFn; +use crate::plan::Eval; +use crate::plan::EvalPlan; use crate::plan::Plan; use crate::plan::PlanChildren; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::optimizer::PlanParentReduceRule; /// Concatenates its children row-wise. #[derive(Clone, Debug)] @@ -140,3 +146,45 @@ impl PlanVTable for Concat { Cow::Owned(format!("chunks[{index}]")) } } + +/// Pushes an expression into every chunk of a [`Concat`]. +#[derive(Debug)] +pub(crate) struct ExpressionConcatRule; + +impl PlanParentReduceRule for ExpressionConcatRule { + type Parent = Eval; + + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + _child_idx: usize, + ) -> VortexResult> { + let expression = parent.expression(); + // Row-index expressions are relative to the whole row domain, so they cannot be evaluated + // chunk by chunk. + let references_row_idx = label_bound_tree( + expression, + |node| { + node.as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + }, + |acc, &child| acc | child, + ) + .get(&ExactBoundExpr(expression.clone())) + .copied() + .unwrap_or(false); + if references_row_idx { + return Ok(None); + } + + let chunks = child + .children() + .iter() + .map(|chunk| Ok(EvalPlan::try_new(expression.clone(), chunk?)?.into_plan())) + .collect::>>()?; + Ok(Some( + ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(), + )) + } +} diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index 70a645da687..bc7255d83ff 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -5,7 +5,14 @@ use std::borrow::Cow; use std::fmt; use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; use vortex_array::expr::BoundExpression; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; @@ -17,6 +24,8 @@ use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; use crate::plan::check_child_count; +use crate::plan::optimize; +use crate::plan::optimizer::reduce_parent; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] @@ -123,3 +132,94 @@ fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> V } Ok(()) } + +impl EvalPlan { + /// Optimizes this plan top-down, applying parent-reduction rules as they become applicable. + /// + /// `blocked_child_type` suppresses one rule re-firing on its own residual output, which would + /// otherwise loop when a rewrite leaves an expression above the same child kind. + pub(crate) fn optimize_top_down( + &self, + blocked_child_type: Option, + ) -> VortexResult { + if self.expression().is_root() { + return optimize(self.child_plan()?); + } + + let child = self.child_plan()?; + let child_type = child.id(); + let parent = EvalPlan::try_new(self.expression().clone(), child.clone())?.into_plan(); + if blocked_child_type != Some(child_type) + && let Some(rewritten) = reduce_parent(&parent, 0)? + { + return Self::optimize_rewrite(rewritten, child_type); + } + + let child = optimize(child)?; + + let child_type = child.id(); + let parent = EvalPlan::try_new(self.expression().clone(), child)?.into_plan(); + if blocked_child_type != Some(child_type) + && let Some(rewritten) = reduce_parent(&parent, 0)? + { + return Self::optimize_rewrite(rewritten, child_type); + } + Ok(parent) + } + + fn optimize_rewrite(rewritten: PlanRef, previous_child_type: PlanId) -> VortexResult { + let Some(eval) = rewritten.as_opt::() else { + return optimize(rewritten); + }; + // A residual expression may remain above the same child kind after a successful rewrite. + // Do not immediately apply that rule again; recursively optimize only the retained child. + let child_type = eval.child_plan()?.id(); + let blocked = (child_type == previous_child_type).then_some(previous_child_type); + eval.optimize_top_down(blocked) + } +} + +/// Rewrites partition accessors in `expression` to read from a partitioned root. +pub(crate) fn rewrite_partition_root( + expression: BoundExpression, + root_dtype: DType, + collapsed: &[(FieldName, FieldName)], +) -> VortexResult { + Ok(expression + .transform_down(|node| { + if let Some(value_name) = node + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + { + let partition_access = &node.children()[0]; + if let Some(partition_name) = partition_access + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + && partition_access.children()[0].is_root() + && collapsed.iter().any(|(partition, value)| { + partition == partition_name && value == value_name + }) + { + return Ok(Transformed { + value: BoundExpression::try_new( + GetItem.bind(partition_name.clone()), + [BoundExpression::new_root(root_dtype.clone())], + )?, + changed: true, + order: TraversalOrder::Skip, + }); + } + } + + if node.is_root() { + Ok(Transformed { + value: BoundExpression::new_root(root_dtype.clone()), + changed: true, + order: TraversalOrder::Skip, + }) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index 2e4a6dbad5e..e795b7bd5c5 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod concat; -mod eval; +pub(crate) mod eval; mod list_pack; mod pack; mod row_idx; @@ -12,12 +12,14 @@ mod take; pub use concat::Concat; pub use concat::ConcatData; pub use concat::ConcatPlan; +pub(crate) use concat::ExpressionConcatRule; pub use eval::Eval; pub use eval::EvalData; pub use eval::EvalPlan; pub use list_pack::ListPack; pub use list_pack::ListPackData; pub use list_pack::ListPackPlan; +pub(crate) use pack::ExpressionPackRule; pub use pack::Pack; pub use pack::PackData; pub use pack::PackPlan; @@ -28,5 +30,6 @@ pub use row_idx::RowIdxPlanMetadata; pub use segment_scan::SegmentScan; pub use segment_scan::SegmentScanData; pub use segment_scan::SegmentScanPlan; +pub(crate) use take::ExpressionTakeRule; pub use take::Take; pub use take::TakePlan; diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 9cf47ebe4b7..678dba62dc3 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -5,20 +5,40 @@ use std::borrow::Cow; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::descendent_bound_annotations; +use vortex_array::expr::make_bound_free_field_annotator; +use vortex_array::expr::transform::partition_bound; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::get_item::GetItem; +use vortex_array::scalar_fn::fns::pack::Pack as PackFn; +use vortex_array::scalar_fn::fns::pack::PackOptions; +use vortex_array::scalar_fn::fns::select::Select; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::registry::CachedId; +use crate::plan::Eval; +use crate::plan::EvalPlan; use crate::plan::Plan; use crate::plan::PlanChildren; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::optimizer::PlanParentReduceRule; +use crate::plan::plans::eval::rewrite_partition_root; /// Assembles a struct from one child per field, plus an optional trailing validity child. #[derive(Clone, Debug)] @@ -224,3 +244,267 @@ fn validate_validity_child(expected_row_count: u64, child: &PlanRef) -> VortexRe } Ok(()) } + +impl PackPlan { + /// Rebuilds this plan with only `fields`, which must be a subset of the current fields. + /// + /// Pruning is only sound for a non-nullable struct: dropping a field of a nullable struct + /// would drop the validity child that the remaining fields depend on. + pub(crate) fn with_pruned_fields( + &self, + fields: Vec<(FieldName, PlanRef)>, + ) -> VortexResult { + vortex_ensure!( + !self.dtype().is_nullable(), + "Cannot prune fields from a nullable Pack" + ); + let struct_fields = StructFields::from_iter( + fields + .iter() + .map(|(name, plan)| (name.clone(), plan.dtype().clone())), + ); + let field_plans = fields.into_iter().map(|(_, plan)| plan).collect::>(); + PackPlan::try_new( + struct_fields, + Nullability::NonNullable, + self.row_count(), + field_plans, + None, + ) + } +} + +/// Pushes an expression into the referenced fields of a [`Pack`], pruning the rest. +#[derive(Debug)] +pub(crate) struct ExpressionPackRule; + +impl PlanParentReduceRule for ExpressionPackRule { + type Parent = Eval; + + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + _child_idx: usize, + ) -> VortexResult> { + if child.dtype().is_nullable() { + return Ok(None); + } + + let expression = parent.expression(); + let fields = child.fields(); + let referenced_fields = + descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) + .get(&ExactBoundExpr(expression.clone())) + .vortex_expect("Bound expression missing free-field annotations") + .clone(); + let expanded_root = expanded_struct_root(child.dtype(), fields)?; + let expanded = expand_struct_root(expression.clone(), &expanded_root, fields)?; + let partitioned = + partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?; + + if partitioned.partition_names.is_empty() { + let selected_indices = fields + .names() + .iter() + .enumerate() + .filter_map(|(index, name)| referenced_fields.contains(name).then_some(index)) + .collect::>(); + if selected_indices.len() == fields.nfields() { + return Ok(None); + } + + let pruned_fields = selected_indices + .into_iter() + .map(|field_index| { + Ok(( + field_name(fields, field_index)?, + field_plan(child, field_index)?, + )) + }) + .collect::>>()?; + let rewritten = child.with_pruned_fields(pruned_fields)?.into_plan(); + return Ok(Some( + EvalPlan::try_new(expression.clone(), rewritten)?.into_plan(), + )); + } + + if partitioned.partition_names.len() == 1 { + let name = partitioned + .partition_names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition has no field"))?; + let index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, index)?; + let lowered = step_into_struct_field(expanded, name, field.dtype().clone())?; + return Ok(Some(EvalPlan::try_new(lowered, field)?.into_plan())); + } + + let residual = partitioned.root; + let mut collapsed = Vec::with_capacity(partitioned.partitions.len()); + let mut field_expressions = vec![None; fields.nfields()]; + for index in 0..partitioned.partitions.len() { + let name = &partitioned.partition_names[index]; + let partition = &partitioned.partitions[index]; + let field_index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, field_index)?; + let lowered = if let Some(pack) = partition + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + && partition.children().len() == 1 + { + let value_name = pack + .names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; + collapsed.push((name.clone(), value_name.clone())); + partition.children()[0].clone() + } else { + partition.clone() + }; + let lowered = step_into_struct_field(lowered, name, field.dtype().clone())?; + field_expressions[field_index] = Some(lowered); + } + + let mut pruned_fields = Vec::with_capacity(partitioned.partition_names.len()); + for (field_index, expression) in field_expressions.into_iter().enumerate() { + let Some(expression) = expression else { + continue; + }; + let field = field_plan(child, field_index)?; + pruned_fields.push(( + field_name(fields, field_index)?, + EvalPlan::try_new(expression, field)?.into_plan(), + )); + } + let rewritten = child.with_pruned_fields(pruned_fields)?.into_plan(); + let residual = rewrite_partition_root(residual, rewritten.dtype().clone(), &collapsed)?; + + Ok(Some(EvalPlan::try_new(residual, rewritten)?.into_plan())) + } +} + +fn field_name(fields: &StructFields, index: usize) -> VortexResult { + Ok(fields + .field_name(index) + .ok_or_else(|| vortex_err!("Struct field {index} has no name"))? + .clone()) +} + +fn field_plan(plan: &Plan, index: usize) -> VortexResult { + plan.child(index)? + .ok_or_else(|| vortex_err!("Struct field {index} has no plan")) +} + +fn expanded_struct_root( + root_dtype: &DType, + fields: &StructFields, +) -> VortexResult { + let root = BoundExpression::new_root(root_dtype.clone()); + let children = fields + .names() + .iter() + .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])) + .collect::>>()?; + bound_pack(fields.names().clone(), children) +} + +fn expand_struct_root( + expression: BoundExpression, + expanded_root: &BoundExpression, + fields: &StructFields, +) -> VortexResult { + Ok(expression + .transform_down(|node| { + if node.is_root() { + return Ok(Transformed { + value: expanded_root.clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + let Some(scalar_fn) = node.as_scalar() else { + return Ok(Transformed::no(node)); + }; + if !node + .children() + .first() + .is_some_and(BoundExpression::is_root) + { + return Ok(Transformed::no(node)); + } + + if let Some(field_name) = scalar_fn.as_opt::() { + let index = fields.find(field_name).ok_or_else(|| { + vortex_err!("Field {field_name} not found while expanding struct root") + })?; + return Ok(Transformed { + value: expanded_root.children()[index].clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + if let Some(selection) = scalar_fn.as_opt::() { let names = selection.normalize_to_included_fields(fields.names())?; + let root = node.children()[0].clone(); let children = names .iter() .map(|name| { - let index = fields - .find(name) - .vortex_expect("normalized selection fields must exist in the root"); - expanded_root.children()[index].clone() + BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()]) }) - .collect(); + .collect::>>()?; return Ok(Transformed { value: bound_pack(names, children)?, changed: true, diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 8b117665531..c52f2f3dd64 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -29,6 +29,8 @@ use vortex_array::expr::gt; use vortex_array::expr::is_null; use vortex_array::expr::lit; use vortex_array::expr::root; +use vortex_array::expr::select; +use vortex_array::expr::select_exclude; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_io::runtime::single::block_on; @@ -278,6 +280,72 @@ fn struct_plan_appends_validity_when_nullable() -> VortexResult<()> { Ok(()) } +#[test] +fn struct_select_preserves_struct_output() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), + Nullability::NonNullable, + ), + vec![ + flat(3, field_dtype.clone(), 0), + flat(3, field_dtype.clone(), 1), + ], + ) + .into_layout(); + + let plan = make_eval(select(["a"], root()), make_plan(layout)?)?.into_plan(); + let optimized = optimize(plan)?; + + assert_eq!( + optimized.dtype(), + &DType::Struct( + StructFields::from_iter([("a", field_dtype)]), + Nullability::NonNullable, + ) + ); + insta::assert_snapshot!(optimized.display_tree(), @r" + root: vortex.plan.eval({a=i32}, rows=3) expr=pack(a: $) + child: vortex.plan.segment_scan(i32, rows=3) + "); + Ok(()) +} + +#[test] +fn struct_select_exclude_preserves_struct_output() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), + Nullability::NonNullable, + ), + vec![ + flat(3, field_dtype.clone(), 0), + flat(3, field_dtype.clone(), 1), + ], + ) + .into_layout(); + + let plan = make_eval(select_exclude(["b"], root()), make_plan(layout)?)?.into_plan(); + let optimized = optimize(plan)?; + + assert_eq!( + optimized.dtype(), + &DType::Struct( + StructFields::from_iter([("a", field_dtype)]), + Nullability::NonNullable, + ) + ); + insta::assert_snapshot!(optimized.display_tree(), @r" + root: vortex.plan.eval({a=i32}, rows=3) expr=pack(a: $) + child: vortex.plan.segment_scan(i32, rows=3) + "); + Ok(()) +} + #[test] fn with_children_rejects_mismatched_arity() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); From aae072ee4dbaf25e5bc918f9ea3536bd70b060bf Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 11 Aug 2026 11:35:35 +0100 Subject: [PATCH 09/13] Run SQL benchmarks with plan v2 Signed-off-by: Joe Isaacs --- .github/workflows/sql-bench-matrix.yml | 2 + Cargo.lock | 1 + vortex-datafusion/Cargo.toml | 1 + vortex-datafusion/src/persistent/opener.rs | 282 ++++++++++++++------- 4 files changed, 188 insertions(+), 98 deletions(-) diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index 52d4e6a545a..98550988a47 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -103,6 +103,8 @@ jobs: env: VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + VORTEX_USE_PLAN_V2: "1" + DATAFUSION_OPTIMIZER_REPARTITION_FILE_SCANS: "false" # Makes python output nicer COLUMNS: 120 strategy: diff --git a/Cargo.lock b/Cargo.lock index e1e76d70772..6c4ddc6fd29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9980,6 +9980,7 @@ dependencies = [ "url", "vortex", "vortex-arrow", + "vortex-scan-v2", "vortex-utils", ] diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 49fb22d4f59..56afe335f85 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -38,6 +38,7 @@ tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-arrow = { workspace = true } +vortex-scan-v2 = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index f5fa2147100..c48e7d0a929 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -42,6 +42,10 @@ use tracing::Instrument; use vortex::array::VortexSessionExecute; use vortex::error::VortexError; use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::expr::BoundExpression; +use vortex::expr::Expression; +use vortex::expr::root; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; @@ -50,6 +54,7 @@ use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; +use vortex_scan_v2::ScanBuilder as PlanScanBuilder; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; @@ -327,40 +332,6 @@ impl FileOpener for VortexOpener { .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; let projector = leftover_projection.make_projector(&stream_schema)?; - // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once. - let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { - Entry::Occupied(mut occupied_entry) => { - if let Some(reader) = occupied_entry.get().upgrade() { - tracing::trace!("reusing layout reader for {}", occupied_entry.key()); - reader - } else { - tracing::trace!("creating layout reader for {}", occupied_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!( - "Failed to create layout reader: {e}" - )) - })?; - occupied_entry.insert(Arc::downgrade(&reader)); - reader - } - } - Entry::Vacant(vacant_entry) => { - tracing::trace!("creating layout reader for {}", vacant_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!("Failed to create layout reader: {e}")) - })?; - vacant_entry.insert(Arc::downgrade(&reader)); - - reader - } - }; - - let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); - - if let Some(vortex_plan) = file.extensions.get::() { - scan_builder = vortex_plan.apply_to_builder(scan_builder); - } - let filter = filter .and_then(|f| { // Verify that all filters we've accepted from DataFusion get pushed down. @@ -396,76 +367,174 @@ impl FileOpener for VortexOpener { .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) .transpose() .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; + let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); + let stream = if std::env::var("VORTEX_USE_PLAN_V2").as_deref() == Ok("1") { + if file.extensions.get::().is_some() { + return Err(exec_datafusion_err!( + "plan-v2 scans do not support VortexAccessPlan" + )); + } + if let Some(file_range) = file.range.as_ref() { + let start = u64::try_from(file_range.start) + .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?; + let end = u64::try_from(file_range.end) + .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?; + if start != 0 || end != file.object_meta.size { + return Err(exec_datafusion_err!( + "plan-v2 scans require DataFusion file-scan repartitioning to be disabled" + )); + } + } - if let Some(limit) = limit - && filter.is_none() - { - scan_builder = scan_builder.with_limit(limit); - } - - if let Some(concurrency) = scan_concurrency { - scan_builder = scan_builder.with_concurrency(concurrency); - } + let mut scan_builder = PlanScanBuilder::try_new( + vxf.footer().layout(), + vxf.segment_source(), + session.clone(), + ) + .map_err(|error| DataFusionError::External(Box::new(error)))? + .with_projection( + unbind(&scan_projection) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ) + .with_some_filter( + filter + .as_ref() + .map(unbind) + .transpose() + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + if let Some(limit) = limit + && filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + if let Some(concurrency) = scan_concurrency { + scan_builder = scan_builder.with_concurrency(concurrency); + } - // Set before the byte-range translation below, which computes natural splits for - // the fields the scan's projection and filter reference. - scan_builder = scan_builder - .with_projection(scan_projection) - .with_some_filter(filter); - - if let Some(file_range) = file.range { - let byte_range = Range { - start: u64::try_from(file_range.start) - .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, + let location = file.object_meta.location.clone(); + let session = session.clone(); + scan_builder + .with_ordered(has_output_ordering) + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + .into_stream() + .map_err(|e| exec_datafusion_err!("Failed to create plan-v2 stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new(e.with_context(format!( + "Failed to read Vortex file with plan v2: {location}" + )))) + }) + .boxed() + } else { + // Share layout readers between partitions so each file's layout is read once. + let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { + Entry::Occupied(mut occupied_entry) => { + if let Some(reader) = occupied_entry.get().upgrade() { + tracing::trace!("reusing layout reader for {}", occupied_entry.key()); + reader + } else { + tracing::trace!("creating layout reader for {}", occupied_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + occupied_entry.insert(Arc::downgrade(&reader)); + reader + } + } + Entry::Vacant(vacant_entry) => { + tracing::trace!("creating layout reader for {}", vacant_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + vacant_entry.insert(Arc::downgrade(&reader)); + reader + } }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_splits = natural_splits_for_file( - natural_splits.as_ref(), - &file.object_meta.location, - &scan_builder, - file.object_meta.size, - )?; - - let Some(row_range) = - split_aligned_row_range(byte_range, natural_splits.as_ref()) - else { - return Ok(stream::empty().boxed()); - }; - scan_builder = scan_builder - .with_row_range(row_range) - // Hand the shared full-file boundaries back to the scan so prepare() - // skips its own layout walk. - .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + let mut scan_builder = + ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); + if let Some(vortex_plan) = file.extensions.get::() { + scan_builder = vortex_plan.apply_to_builder(scan_builder); + } + if let Some(limit) = limit + && filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + if let Some(concurrency) = scan_concurrency { + scan_builder = scan_builder.with_concurrency(concurrency); } - } - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); - let stream = scan_builder - .with_metrics_registry(metrics_registry) - .with_ordered(has_output_ordering) - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) - .into_stream() - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - file.object_meta.location - )))) - }) + // Set before translating byte ranges because natural splits depend on the fields + // referenced by the projection and filter. + scan_builder = scan_builder + .with_projection(scan_projection) + .with_some_filter(filter); + if let Some(file_range) = file.range { + let byte_range = Range { + start: u64::try_from(file_range.start).map_err(|_| { + exec_datafusion_err!("Vortex file range start is negative") + })?, + end: u64::try_from(file_range.end).map_err(|_| { + exec_datafusion_err!("Vortex file range end is negative") + })?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + let natural_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + &scan_builder, + file.object_meta.size, + )?; + let Some(row_range) = + split_aligned_row_range(byte_range, natural_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + scan_builder = scan_builder + .with_row_range(row_range) + .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + } + } + + let location = file.object_meta.location.clone(); + let session = session.clone(); + scan_builder + .with_metrics_registry(metrics_registry) + .with_ordered(has_output_ordering) + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + .into_stream() + .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new(e.with_context(format!( + "Failed to read Vortex file: {location}" + )))) + }) + .boxed() + }; + let stream = stream .map(move |batch| { let batch = if projector.projection().as_ref().is_empty() { batch @@ -494,6 +563,23 @@ impl FileOpener for VortexOpener { } } +fn unbind(expression: &BoundExpression) -> VortexResult { + if expression.is_root() { + return Ok(root()); + } + Expression::try_new( + expression + .as_scalar() + .vortex_expect("non-root bound expression must have a scalar function") + .clone(), + expression + .children() + .iter() + .map(unbind) + .collect::>>()?, + ) +} + /// A file's natural split boundaries plus the precomputed byte each split is assigned to, /// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row /// boundaries with a binary search instead of re-projecting every split per partition. From c657a5af7b8be93c742660ac5e418f30eb08617a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 11 Aug 2026 12:06:34 +0100 Subject: [PATCH 10/13] Add adaptive plan filter execution Signed-off-by: Joe Isaacs --- Cargo.lock | 1 + vortex-datafusion/src/persistent/opener.rs | 11 +- vortex-layout/src/plan/plans/zoned.rs | 35 ++++++- vortex-layout/src/plan/tests.rs | 31 ++++++ vortex-layout/src/scan/mod.rs | 2 + vortex-scan-v2/Cargo.toml | 1 + vortex-scan-v2/src/filter.rs | 111 +++++++++++++++++++++ vortex-scan-v2/src/lib.rs | 2 + vortex-scan-v2/src/repeated_scan.rs | 5 +- vortex-scan-v2/src/scan_builder.rs | 73 +++++++++++--- vortex-scan-v2/src/tasks.rs | 17 +--- vortex-scan-v2/src/tests.rs | 43 ++++++++ 12 files changed, 297 insertions(+), 35 deletions(-) create mode 100644 vortex-scan-v2/src/filter.rs diff --git a/Cargo.lock b/Cargo.lock index 6c4ddc6fd29..9fa86b76314 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10549,6 +10549,7 @@ dependencies = [ name = "vortex-scan-v2" version = "0.1.0" dependencies = [ + "bit-vec", "futures", "itertools 0.14.0", "parking_lot", diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index c48e7d0a929..36cafa4519c 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -54,6 +54,7 @@ use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; +use vortex_scan_v2::FilterMode; use vortex_scan_v2::ScanBuilder as PlanScanBuilder; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; @@ -386,6 +387,13 @@ impl FileOpener for VortexOpener { } } + let filter_mode = if std::env::var("VORTEX_PLAN_V2_FILTER_MODE").as_deref() + == Ok("adaptive") + { + FilterMode::Adaptive + } else { + FilterMode::Parallel + }; let mut scan_builder = PlanScanBuilder::try_new( vxf.footer().layout(), vxf.segment_source(), @@ -402,7 +410,8 @@ impl FileOpener for VortexOpener { .map(unbind) .transpose() .map_err(|error| DataFusionError::External(Box::new(error)))?, - ); + ) + .with_filter_mode(filter_mode); if let Some(limit) = limit && filter.is_none() { diff --git a/vortex-layout/src/plan/plans/zoned.rs b/vortex-layout/src/plan/plans/zoned.rs index 040db545c25..33d17d2b591 100644 --- a/vortex-layout/src/plan/plans/zoned.rs +++ b/vortex-layout/src/plan/plans/zoned.rs @@ -35,6 +35,7 @@ use vortex_session::registry::CachedId; use crate::layouts::zoned::zone_map::ZoneMap; use crate::plan::Eval; +use crate::plan::EvalPlan; use crate::plan::Plan; use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; @@ -223,6 +224,26 @@ impl ZonedPlan { )) } + fn with_data_expression(&self, expression: BoundExpression) -> VortexResult> { + let Some(data_plan) = self.data_plan()? else { + return Ok(None); + }; + Ok(Some( + PlanParts { + vtable: Zoned, + dtype: expression.dtype().clone(), + row_count: self.row_count(), + children: vec![ + EvalPlan::try_new(expression, data_plan)?.into_plan(), + self.zones_plan()?, + ] + .into(), + data: self.data().clone(), + } + .into_typed(), + )) + } + fn execute_pruning( &self, ctx: &PlanExecutionContext, @@ -380,7 +401,6 @@ impl PlanVTable for Zoned { if data_plan.dtype() != plan.dtype() || data_plan.row_count() != plan.row_count() { vortex_error::vortex_bail!("Zoned data child shape does not match the plan output"); } - data.column_dtype = plan.dtype().clone(); Ok(()) } @@ -414,7 +434,7 @@ impl PlanVTable for Zoned { } } -/// Rewrites an abstract statistic expression over a zoned plan into its pruning state. +/// Pushes data expressions through a zoned plan and rewrites statistic expressions into pruning. #[derive(Debug)] pub(crate) struct ExpressionZonedRule; @@ -444,12 +464,17 @@ impl PlanParentReduceRule for ExpressionZonedRule { contains_root |= expression.is_root(); Ok(Transformed::no(expression)) })?; - if !parent.dtype().is_boolean() || !contains_stat || contains_root { - return Ok(None); + if contains_stat { + if !parent.dtype().is_boolean() || contains_root { + return Ok(None); + } + return Ok(child + .with_pruning(parent.expression().clone())? + .map(Plan::into_plan)); } Ok(child - .with_pruning(parent.expression().clone())? + .with_data_expression(parent.expression().clone())? .map(Plan::into_plan)) } } diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index c52f2f3dd64..caaa6b1ed70 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -966,6 +966,37 @@ fn zoned_plan_exposes_data_and_zones() -> VortexResult<()> { Ok(()) } +#[test] +fn data_expression_pushes_through_zoned_and_preserves_zones() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let zones_dtype = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let aggregate_fns: Arc<[AggregateFnRef]> = Vec::new().into(); + let layout = ZonedLayout::try_new( + flat(5, dtype.clone(), 0), + flat(2, zones_dtype, 1), + zone_len, + aggregate_fns, + )? + .into_layout(); + let source = make_plan(layout)?; + let expression = gt(root(), lit(5_i32)).bind(&dtype)?; + + let optimized = optimize(EvalPlan::try_new(expression, source)?.into_plan())?; + insta::assert_snapshot!(optimized.display_tree(), @r" + root: vortex.plan.zoned(bool, rows=5) + data: vortex.plan.eval(bool, rows=5) expr=($ > 5i32) + child: vortex.plan.segment_scan(i32, rows=5) + zones: vortex.plan.segment_scan({}, rows=2) + "); + let optimized_zoned = optimized + .as_opt::() + .ok_or_else(|| vortex_err!("optimized plan is not zoned"))?; + assert!(!optimized_zoned.is_pruning()); + assert_eq!(optimized_zoned.zones_plan()?.row_count(), 2); + Ok(()) +} + #[test] fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index 98fd1918a42..89556736689 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -13,6 +13,8 @@ mod tasks; #[cfg(test)] mod test; +pub use filter::FilterExpr; + /// A heuristic for an ideal split size. /// /// We don't actually know if this is right, but it is probably a good estimate. diff --git a/vortex-scan-v2/Cargo.toml b/vortex-scan-v2/Cargo.toml index cece44f9730..604e55fdfda 100644 --- a/vortex-scan-v2/Cargo.toml +++ b/vortex-scan-v2/Cargo.toml @@ -14,6 +14,7 @@ rust-version = { workspace = true } version = { workspace = true } [dependencies] +bit-vec = { workspace = true } futures = { workspace = true, features = ["alloc", "async-await"] } itertools = { workspace = true } tracing = { workspace = true } diff --git a/vortex-scan-v2/src/filter.rs b/vortex-scan-v2/src/filter.rs new file mode 100644 index 00000000000..be9f9e095ef --- /dev/null +++ b/vortex-scan-v2/src/filter.rs @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use bit_vec::BitVec; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_error::VortexResult; +use vortex_layout::plan::PlanExecutionContext; +use vortex_layout::plan::PlanRef; +use vortex_layout::scan::FilterExpr; +use vortex_mask::Mask; + +/// Controls how a scan executes top-level filter conjunctions. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum FilterMode { + /// Optimizes the complete predicate as one plan, allowing independent branches to run in + /// parallel. + #[default] + Parallel, + /// Splits top-level conjunctions and executes them as an adaptively ordered mask chain. + Adaptive, +} + +#[derive(Clone)] +pub(crate) enum FilterPlan { + Parallel(PlanRef), + Adaptive { + filter: Arc, + plans: Arc<[PlanRef]>, + }, +} + +impl FilterPlan { + pub(crate) fn parallel(plan: PlanRef) -> Self { + Self::Parallel(plan) + } + + pub(crate) fn adaptive(filter: FilterExpr, plans: Vec) -> Self { + Self::Adaptive { + filter: Arc::new(filter), + plans: plans.into(), + } + } + + pub(crate) fn plans(&self) -> Vec<&PlanRef> { + match self { + Self::Parallel(plan) => vec![plan], + Self::Adaptive { plans, .. } => plans.iter().collect(), + } + } + + pub(crate) fn execute( + &self, + execution: &PlanExecutionContext, + row_range: &Range, + row_mask: Mask, + ) -> VortexResult { + match self { + Self::Parallel(filter) => { + let predicate = + filter.execute(execution, row_range, MaskFuture::ready(row_mask.clone()))?; + let session = execution.session().clone(); + Ok(MaskFuture::new(row_mask.len(), async move { + let predicate = predicate.await?; + let mut execution = session.create_execution_ctx(); + let predicate: Mask = predicate.null_as_false().execute(&mut execution)?; + Ok(row_mask.intersect_by_rank(&predicate)) + })) + } + Self::Adaptive { filter, plans } => { + let execution = execution.clone(); + let row_range = row_range.clone(); + let filter = Arc::clone(filter); + let plans = Arc::clone(plans); + Ok(MaskFuture::new(row_mask.len(), async move { + let mut row_mask = row_mask; + let mut remaining = BitVec::from_elem(plans.len(), true); + while let Some(index) = filter.next_conjunct(&remaining) { + remaining.set(index, false); + if row_mask.all_false() { + break; + } + + let input_rows = row_mask.true_count(); + let predicate = plans[index].execute( + &execution, + &row_range, + MaskFuture::ready(row_mask.clone()), + )?; + let predicate = predicate.await?; + let mut ctx = execution.session().create_execution_ctx(); + let predicate: Mask = predicate.null_as_false().execute(&mut ctx)?; + row_mask = row_mask.intersect_by_rank(&predicate); + filter.report_selectivity(index, row_mask.density()); + tracing::trace!( + target: "vortex_scan_v2::execution", + conjunct = index, + input_rows, + output_rows = row_mask.true_count(), + "applied an adaptive plan filter conjunct" + ); + } + Ok(row_mask) + })) + } + } + } +} diff --git a/vortex-scan-v2/src/lib.rs b/vortex-scan-v2/src/lib.rs index 52249d77cee..8477484e7ae 100644 --- a/vortex-scan-v2/src/lib.rs +++ b/vortex-scan-v2/src/lib.rs @@ -10,6 +10,7 @@ //! Set `RUST_LOG=vortex_scan_v2=debug` to log source and optimized plan trees and selected scan //! splits. Use `trace` to also log execution of each split. +mod filter; mod repeated_scan; mod scan_builder; mod splits; @@ -18,6 +19,7 @@ mod tasks; #[cfg(test)] mod tests; +pub use filter::FilterMode; pub use repeated_scan::RepeatedScan; pub use scan_builder::ScanBuilder; pub use splits::SplitBy; diff --git a/vortex-scan-v2/src/repeated_scan.rs b/vortex-scan-v2/src/repeated_scan.rs index 50016a1d98d..ac47c6eaf51 100644 --- a/vortex-scan-v2/src/repeated_scan.rs +++ b/vortex-scan-v2/src/repeated_scan.rs @@ -25,6 +25,7 @@ use vortex_layout::plan::PlanRef; use vortex_scan::selection::Selection; use vortex_utils::parallelism::get_available_parallelism; +use crate::filter::FilterPlan; use crate::splits::Splits; use crate::tasks::TaskContext; use crate::tasks::split_exec; @@ -34,7 +35,7 @@ pub struct RepeatedScan { execution: PlanExecutionContext, projection: PlanRef, pruning: Option, - filter: Option, + filter: Option, ordered: bool, row_range: Option>, selection: Selection, @@ -82,7 +83,7 @@ impl RepeatedScan { execution: PlanExecutionContext, projection: PlanRef, pruning: Option, - filter: Option, + filter: Option, ordered: bool, row_range: Option>, selection: Selection, diff --git a/vortex-scan-v2/src/scan_builder.rs b/vortex-scan-v2/src/scan_builder.rs index 6c44b131fd3..269f12ec26b 100644 --- a/vortex-scan-v2/src/scan_builder.rs +++ b/vortex-scan-v2/src/scan_builder.rs @@ -38,6 +38,7 @@ use vortex_layout::plan::RowIdxValues; use vortex_layout::plan::Zoned; use vortex_layout::plan::lower; use vortex_layout::plan::optimize; +use vortex_layout::scan::FilterExpr; use vortex_layout::segments::SegmentSource; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -45,6 +46,8 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::RepeatedScan; +use crate::filter::FilterMode; +use crate::filter::FilterPlan; use crate::splits::SplitBy; use crate::splits::Splits; use crate::splits::attempt_split_ranges; @@ -55,6 +58,7 @@ pub struct ScanBuilder { base_plan: PlanRef, projection: Expression, filter: Option, + filter_mode: FilterMode, ordered: bool, row_range: Option>, selection: Selection, @@ -96,6 +100,7 @@ impl ScanBuilder { base_plan, projection: root(), filter: None, + filter_mode: FilterMode::default(), ordered: true, row_range: None, selection: Selection::default(), @@ -140,6 +145,12 @@ impl ScanBuilder { self } + /// Configures whether filter conjunctions execute in parallel or as an adaptive chain. + pub fn with_filter_mode(mut self, filter_mode: FilterMode) -> Self { + self.filter_mode = filter_mode; + self + } + /// Sets the projection expression. pub fn with_projection(mut self, projection: Expression) -> Self { self.projection = projection; @@ -232,6 +243,7 @@ impl ScanBuilder { base_plan: self.base_plan, projection: self.projection, filter: self.filter, + filter_mode: self.filter_mode, ordered: self.ordered, row_range: self.row_range, selection: self.selection, @@ -263,7 +275,7 @@ impl ScanBuilder { &source, self.execution.session(), )?; - let filter = optimize_filter_plan(filter_expression.as_ref(), &source)?; + let filter = optimize_filter_plan(filter_expression.as_ref(), &source, self.filter_mode)?; let splits = if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) { @@ -274,7 +286,9 @@ impl ScanBuilder { .clone() .unwrap_or_else(|| 0..self.base_plan.row_count()); let mut plans = vec![&projection]; - plans.extend(filter.as_ref()); + if let Some(filter) = &filter { + plans.extend(filter.plans()); + } plans.extend(pruning.as_ref()); Splits::Natural(self.split_by.splits(&plans, &row_range)?) }; @@ -424,21 +438,52 @@ fn build_pruning_plan( fn optimize_filter_plan( filter: Option<&BoundExpression>, source: &PlanRef, -) -> VortexResult> { + mode: FilterMode, +) -> VortexResult> { let Some(expression) = filter else { return Ok(None); }; - let filter = optimize(EvalPlan::try_new(expression.clone(), source.clone())?.into_plan())?; - vortex_ensure!( - filter.dtype().is_boolean(), - "Filter plan must produce booleans" - ); - tracing::debug!( - target: "vortex_scan_v2::planner", - plan = %filter.display_tree(), - "optimized the filter physical plan" - ); - Ok(Some(filter)) + match mode { + FilterMode::Parallel => { + let filter = + optimize(EvalPlan::try_new(expression.clone(), source.clone())?.into_plan())?; + vortex_ensure!( + filter.dtype().is_boolean(), + "Filter plan must produce booleans" + ); + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %filter.display_tree(), + "optimized the parallel filter physical plan" + ); + Ok(Some(FilterPlan::parallel(filter))) + } + FilterMode::Adaptive => { + let filter = FilterExpr::new(expression.clone()); + let plans = filter + .conjuncts() + .iter() + .enumerate() + .map(|(index, expression)| { + let plan = optimize( + EvalPlan::try_new(expression.clone(), source.clone())?.into_plan(), + )?; + vortex_ensure!( + plan.dtype().is_boolean(), + "Filter conjunct plan must produce booleans" + ); + tracing::debug!( + target: "vortex_scan_v2::planner", + conjunct = index, + plan = %plan.display_tree(), + "optimized an adaptive filter conjunct plan" + ); + Ok(plan) + }) + .collect::>>()?; + Ok(Some(FilterPlan::adaptive(filter, plans))) + } + } } fn uses_only_pruning_sources(plan: &PlanRef) -> VortexResult { diff --git a/vortex-scan-v2/src/tasks.rs b/vortex-scan-v2/src/tasks.rs index 3b8dd701875..c9576d7d8ff 100644 --- a/vortex-scan-v2/src/tasks.rs +++ b/vortex-scan-v2/src/tasks.rs @@ -14,6 +14,8 @@ use vortex_layout::plan::PlanRef; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; +use crate::filter::FilterPlan; + pub(crate) type TaskFuture = BoxFuture<'static, VortexResult>; pub(crate) fn split_exec( @@ -78,18 +80,7 @@ pub(crate) fn split_exec( } let filter_mask = if let Some(filter) = &ctx.filter { - let predicate = filter.execute( - &ctx.execution, - &row_range, - MaskFuture::ready(row_mask.clone()), - )?; - let session = ctx.execution.session().clone(); - MaskFuture::new(row_mask.len(), async move { - let predicate = predicate.await?; - let mut execution = session.create_execution_ctx(); - let predicate: Mask = predicate.null_as_false().execute(&mut execution)?; - Ok(row_mask.intersect_by_rank(&predicate)) - }) + filter.execute(&ctx.execution, &row_range, row_mask)? } else { MaskFuture::ready(row_mask) }; @@ -126,7 +117,7 @@ pub(crate) fn split_exec( pub(crate) struct TaskContext { pub(crate) execution: PlanExecutionContext, pub(crate) pruning: Option, - pub(crate) filter: Option, + pub(crate) filter: Option, pub(crate) projection: PlanRef, pub(crate) mapper: Arc VortexResult + Send + Sync>, } diff --git a/vortex-scan-v2/src/tests.rs b/vortex-scan-v2/src/tests.rs index e058727c0a6..e7544b8997b 100644 --- a/vortex-scan-v2/src/tests.rs +++ b/vortex-scan-v2/src/tests.rs @@ -19,6 +19,7 @@ use vortex_array::expr::checked_add; use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::lit; +use vortex_array::expr::lt; use vortex_array::expr::root; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; @@ -46,6 +47,7 @@ use vortex_layout::sequence::SequentialArrayStreamExt; use vortex_layout::session::LayoutSession; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; +use crate::FilterMode; use crate::ScanBuilder; use crate::SplitBy; @@ -109,6 +111,47 @@ fn scans_layout_through_optimized_plans() -> VortexResult<()> { }) } +#[test] +fn supports_parallel_and_adaptive_filter_modes() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = PrimitiveArray::from_iter(0_i32..10).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + let filter = and(gt(root(), lit(2_i32)), lt(root(), lit(8_i32))); + let expected = PrimitiveArray::from_iter(3_i32..8).into_array(); + + for mode in [FilterMode::Parallel, FilterMode::Adaptive] { + let actual = ScanBuilder::try_new( + &layout, + Arc::::clone(&segments), + session.clone(), + )? + .with_filter(filter.clone()) + .with_filter_mode(mode) + .with_split_by(SplitBy::RowCount(3)) + .into_array_stream()? + .read_all() + .await?; + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + } + Ok(()) + }) +} + #[test] fn filter_and_projection_share_flat_segment_request() -> VortexResult<()> { block_on(|handle| async { From eb4af24c070d17c3c8b8c6ac7d3ceaa592a70aa6 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 11 Aug 2026 18:47:31 +0100 Subject: [PATCH 11/13] Avoid reading fields for empty projections Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/plans/row_idx.rs | 7 ++++ vortex-layout/src/plan/tests.rs | 37 +++++++++++++++++ vortex-scan-v2/src/splits.rs | 6 +++ vortex-scan-v2/src/tests.rs | 55 +++++++++++++++++++++++++ 4 files changed, 105 insertions(+) diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index 5a8c026e7df..6961ebd505b 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -164,6 +164,13 @@ impl PlanParentReduceRule for ExpressionRowIdxRule { } })?; + // A root-independent expression does not need either side of RowIdx. + if partitioned.partition_annotations.is_empty() { + return Ok(Some( + EvalPlan::try_new(expression.clone(), child.child_plan()?)?.into_plan(), + )); + } + if partitioned.partition_annotations.len() == 1 { return match partitioned.partition_annotations[0] { RowIdxExpressionPartition::RowIdx => { diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index caaa6b1ed70..964bbaef36d 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -28,6 +28,7 @@ use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::is_null; use vortex_array::expr::lit; +use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::expr::select; use vortex_array::expr::select_exclude; @@ -607,6 +608,42 @@ fn row_idx_only_expression_uses_generated_values_plan() -> VortexResult<()> { Ok(()) } +#[test] +fn empty_projection_prunes_row_idx_child_fields() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ), + vec![flat(3, value_dtype.clone(), 0), flat(3, value_dtype, 1)], + ) + .into_layout(); + let projection = pack( + std::iter::empty::<(&str, Expression)>(), + Nullability::NonNullable, + ); + let plan = make_eval( + projection, + RowIdxPlan::new(0, make_plan(layout)?).into_plan(), + )? + .into_plan(); + + let optimized = optimize(plan)?; + let projection = optimized + .as_opt::() + .ok_or_else(|| vortex_err!("optimized plan has no projection expression"))?; + let child = projection.child_plan()?; + let empty_struct = child + .as_opt::() + .ok_or_else(|| vortex_err!("empty projection did not prune the RowIdx child"))?; + + assert_eq!(empty_struct.nfields(), 0); + assert_eq!(empty_struct.children().len(), 0); + Ok(()) +} + #[test] fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { let value_dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-scan-v2/src/splits.rs b/vortex-scan-v2/src/splits.rs index a26a223e95f..34ff2dd04d9 100644 --- a/vortex-scan-v2/src/splits.rs +++ b/vortex-scan-v2/src/splits.rs @@ -86,6 +86,12 @@ fn collect_plan_splits( return Ok(()); } + // A childless Pack preserves row count even though it exposes no leaf boundaries. + if plan.is::() && plan.child_count() == 0 { + boundaries.push(row_offset + row_range.end); + return Ok(()); + } + if plan.is::() || plan.is::() { for index in 0..plan.child_count() { if let Some(child) = plan.child(index)? diff --git a/vortex-scan-v2/src/tests.rs b/vortex-scan-v2/src/tests.rs index e7544b8997b..37e1a7be2d7 100644 --- a/vortex-scan-v2/src/tests.rs +++ b/vortex-scan-v2/src/tests.rs @@ -14,12 +14,17 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; use vortex_array::expr::and; use vortex_array::expr::checked_add; use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::lit; use vortex_array::expr::lt; +use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; @@ -196,6 +201,56 @@ fn filter_and_projection_share_flat_segment_request() -> VortexResult<()> { }) } +#[test] +fn empty_projection_reads_no_struct_fields() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = StructArray::from_fields( + [ + ("a", buffer![1_i32, 2, 3, 4, 5, 6].into_array()), + ("b", buffer![7_i32, 8, 9, 10, 11, 12].into_array()), + ] + .as_slice(), + )? + .into_array(); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let layout = TableStrategy::new(Arc::clone(&flat), flat) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + let tracking = TrackingSource::new(segments); + let source: Arc = Arc::new(tracking.clone()); + let projection = pack( + std::iter::empty::<(&str, Expression)>(), + Nullability::NonNullable, + ); + + let actual = ScanBuilder::try_new(&layout, source, session)? + .with_projection(projection) + .into_array_stream()? + .read_all() + .await?; + + assert_eq!(actual.len(), 6); + assert_eq!( + actual.dtype(), + &DType::Struct(StructFields::empty(), Nullability::NonNullable) + ); + assert!(tracking.requests().is_empty()); + Ok(()) + }) +} + #[test] fn zoned_pruning_skips_a_falsified_data_chunk() -> VortexResult<()> { block_on(|handle| async { From 93a1d505390221bfaabcda0701d4504d8924c5ab Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 11 Aug 2026 13:36:11 +0000 Subject: [PATCH 12/13] Support DataFusion file-scan repartitioning in plan-v2 scans Translate each partition's byte range to the row range whose natural splits it owns, mirroring the LayoutReader path, so plan-v2 scans no longer require repartitioning to be disabled. Benchmarks can now run against the same DataFusion configuration as the stored baseline. Log the plan-v2 scan path once per process so benchmark and CI logs record which scan path produced their timings. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MjXw22CLhx6NgA8mxjJCmQ --- .github/workflows/sql-bench-matrix.yml | 1 - Cargo.lock | 1 + vortex-datafusion/src/persistent/opener.rs | 87 +++++++++++++------- vortex-scan-v2/Cargo.toml | 1 + vortex-scan-v2/src/scan_builder.rs | 84 ++++++++++++++----- vortex-scan-v2/src/tests.rs | 96 ++++++++++++++++++++++ 6 files changed, 218 insertions(+), 52 deletions(-) diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index 98550988a47..18fd199cc6a 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -104,7 +104,6 @@ jobs: VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" VORTEX_USE_PLAN_V2: "1" - DATAFUSION_OPTIMIZER_REPARTITION_FILE_SCANS: "false" # Makes python output nicer COLUMNS: 120 strategy: diff --git a/Cargo.lock b/Cargo.lock index 9fa86b76314..6c106617511 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10553,6 +10553,7 @@ dependencies = [ "futures", "itertools 0.14.0", "parking_lot", + "rstest", "tracing", "tracing-subscriber", "vortex-array", diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 36cafa4519c..64ebcd5d2aa 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -3,6 +3,7 @@ use std::ops::Range; use std::sync::Arc; +use std::sync::Once; use std::sync::Weak; use arrow_array::RecordBatchOptions; @@ -375,18 +376,6 @@ impl FileOpener for VortexOpener { "plan-v2 scans do not support VortexAccessPlan" )); } - if let Some(file_range) = file.range.as_ref() { - let start = u64::try_from(file_range.start) - .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?; - let end = u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?; - if start != 0 || end != file.object_meta.size { - return Err(exec_datafusion_err!( - "plan-v2 scans require DataFusion file-scan repartitioning to be disabled" - )); - } - } - let filter_mode = if std::env::var("VORTEX_PLAN_V2_FILTER_MODE").as_deref() == Ok("adaptive") { @@ -421,6 +410,42 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } + announce_plan_v2(filter_mode); + + // DataFusion hands each partition a byte range of the file. Translate it to the + // row range whose natural splits that byte range owns, so partitions of one file + // cover it exactly once. + if let Some(file_range) = file.range.as_ref() { + let byte_range = Range { + start: u64::try_from(file_range.start).map_err(|_| { + exec_datafusion_err!("Vortex file range start is negative") + })?, + end: u64::try_from(file_range.end).map_err(|_| { + exec_datafusion_err!("Vortex file range end is negative") + })?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + let file_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + file.object_meta.size, + || { + scan_builder.full_file_splits().map_err(|e| { + exec_datafusion_err!( + "Failed to compute plan-v2 natural splits: {e}" + ) + }) + }, + )?; + let Some(row_range) = + split_aligned_row_range(byte_range, file_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + scan_builder = scan_builder.with_row_range(row_range); + } + } + let location = file.object_meta.location.clone(); let session = session.clone(); scan_builder @@ -505,8 +530,14 @@ impl FileOpener for VortexOpener { let natural_splits = natural_splits_for_file( natural_splits.as_ref(), &file.object_meta.location, - &scan_builder, file.object_meta.size, + || { + scan_builder.full_file_splits().map_err(|e| { + exec_datafusion_err!( + "Failed to compute Vortex natural splits: {e}" + ) + }) + }, )?; let Some(row_range) = split_aligned_row_range(byte_range, natural_splits.as_ref()) @@ -644,11 +675,11 @@ impl NaturalSplits { } /// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. -fn natural_splits_for_file( +fn natural_splits_for_file( natural_splits: &DashMap>, path: &Path, - scan_builder: &ScanBuilder, total_size: u64, + row_boundaries: impl FnOnce() -> DFResult>, ) -> DFResult> { if let Some(splits) = natural_splits.get(path) { return Ok(Arc::clone(splits.value())); @@ -660,27 +691,23 @@ fn natural_splits_for_file( match natural_splits.entry(path.clone()) { Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), Entry::Vacant(entry) => { - let splits = compute_natural_splits(scan_builder, total_size)?; + let splits = Arc::new(NaturalSplits::new(row_boundaries()?.into(), total_size)); entry.insert(Arc::clone(&splits)); Ok(splits) } } } -/// Walk the layout tree to compute the file's full natural split boundaries for the fields -/// referenced by the scan's projection and filter. -fn compute_natural_splits( - scan_builder: &ScanBuilder, - total_size: u64, -) -> DFResult> { - let row_boundaries = scan_builder - .full_file_splits() - .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; - - Ok(Arc::new(NaturalSplits::new( - row_boundaries.into(), - total_size, - ))) +/// Logs once per process that scans are running through plan-v2, so a benchmark or CI job log +/// records which scan path produced its timings. +fn announce_plan_v2(filter_mode: FilterMode) { + static ANNOUNCED: Once = Once::new(); + ANNOUNCED.call_once(|| { + tracing::info!( + ?filter_mode, + "Vortex file scans are using the plan-v2 scan path (VORTEX_USE_PLAN_V2=1)" + ); + }); } /// Translate a DataFusion byte range to the contiguous natural split ranges it owns. diff --git a/vortex-scan-v2/Cargo.toml b/vortex-scan-v2/Cargo.toml index 604e55fdfda..28d615be182 100644 --- a/vortex-scan-v2/Cargo.toml +++ b/vortex-scan-v2/Cargo.toml @@ -30,6 +30,7 @@ vortex-utils = { workspace = true } [dev-dependencies] parking_lot = { workspace = true } +rstest = { workspace = true } tracing-subscriber = { workspace = true, features = ["env-filter"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-file = { workspace = true, features = ["tokio"] } diff --git a/vortex-scan-v2/src/scan_builder.rs b/vortex-scan-v2/src/scan_builder.rs index 269f12ec26b..23bd6ff7324 100644 --- a/vortex-scan-v2/src/scan_builder.rs +++ b/vortex-scan-v2/src/scan_builder.rs @@ -222,6 +222,42 @@ impl ScanBuilder { self } + /// Optimizes the projection, pruning, and filter plans over the row-index-aware source. + fn optimized_plans(&self) -> VortexResult { + let source = RowIdxPlan::new(self.row_offset, self.base_plan.clone()).into_plan(); + tracing::debug!( + target: "vortex_scan_v2::planner", + row_offset = self.row_offset, + plan = %source.display_tree(), + "planning expressions over the row-index-aware source" + ); + let projection = optimize_projection_plan(self.projection.clone(), &source)?; + let filter_expression = optimize_filter_expression(self.filter.clone(), &source)?; + let pruning = optimize_pruning_plan( + filter_expression.as_ref(), + &source, + self.execution.session(), + )?; + let filter = optimize_filter_plan(filter_expression.as_ref(), &source, self.filter_mode)?; + Ok(OptimizedPlans { + projection, + pruning, + filter, + }) + } + + /// Returns the split boundaries the scan would use over the whole file, ignoring any + /// configured row range or selection. + /// + /// Callers that partition one file across several scans use these boundaries to divide the + /// file into disjoint row ranges, so they must be computed from the full row count rather + /// than from an already-restricted range. + pub fn full_file_splits(&self) -> VortexResult> { + let plans = self.optimized_plans()?; + self.split_by + .splits(&plans.split_sources(), &(0..self.base_plan.row_count())) + } + /// Returns the dtype produced by the projection expression. pub fn dtype(&self) -> VortexResult { self.projection.return_dtype(self.base_plan.dtype()) @@ -261,21 +297,7 @@ impl ScanBuilder { vortex_bail!("Vortex doesn't support scans with both a filter and a limit") } - let source = RowIdxPlan::new(self.row_offset, self.base_plan.clone()).into_plan(); - tracing::debug!( - target: "vortex_scan_v2::planner", - row_offset = self.row_offset, - plan = %source.display_tree(), - "planning expressions over the row-index-aware source" - ); - let projection = optimize_projection_plan(self.projection, &source)?; - let filter_expression = optimize_filter_expression(self.filter, &source)?; - let pruning = optimize_pruning_plan( - filter_expression.as_ref(), - &source, - self.execution.session(), - )?; - let filter = optimize_filter_plan(filter_expression.as_ref(), &source, self.filter_mode)?; + let plans = self.optimized_plans()?; let splits = if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) { @@ -285,12 +307,7 @@ impl ScanBuilder { .row_range .clone() .unwrap_or_else(|| 0..self.base_plan.row_count()); - let mut plans = vec![&projection]; - if let Some(filter) = &filter { - plans.extend(filter.plans()); - } - plans.extend(pruning.as_ref()); - Splits::Natural(self.split_by.splits(&plans, &row_range)?) + Splits::Natural(self.split_by.splits(&plans.split_sources(), &row_range)?) }; match &splits { Splits::Natural(boundaries) => tracing::debug!( @@ -307,6 +324,11 @@ impl ScanBuilder { ), } + let OptimizedPlans { + projection, + pruning, + filter, + } = plans; Ok(RepeatedScan::new( self.execution, projection, @@ -346,6 +368,26 @@ impl ScanBuilder { } } +/// The optimized plans a scan executes: the projection, the optional pruning plan, and the +/// optional filter plan. +struct OptimizedPlans { + projection: PlanRef, + pruning: Option, + filter: Option, +} + +impl OptimizedPlans { + /// The plans whose layout boundaries contribute to the scan's natural splits. + fn split_sources(&self) -> Vec<&PlanRef> { + let mut plans = vec![&self.projection]; + if let Some(filter) = &self.filter { + plans.extend(filter.plans()); + } + plans.extend(self.pruning.as_ref()); + plans + } +} + fn optimize_projection_plan(expression: Expression, source: &PlanRef) -> VortexResult { tracing::debug!( target: "vortex_scan_v2::planner", diff --git a/vortex-scan-v2/src/tests.rs b/vortex-scan-v2/src/tests.rs index 37e1a7be2d7..bcf933840cb 100644 --- a/vortex-scan-v2/src/tests.rs +++ b/vortex-scan-v2/src/tests.rs @@ -2,9 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::num::NonZeroUsize; +use std::ops::Range; use std::sync::Arc; use parking_lot::Mutex; +use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -408,3 +410,97 @@ fn scans_selected_rows_from_a_list_plan() -> VortexResult<()> { Ok(()) }) } + +#[test] +fn full_file_splits_ignore_the_configured_row_range() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = ChunkedArray::from_iter([ + buffer![1_i32, 2, 3].into_array(), + buffer![4_i32, 5, 6].into_array(), + buffer![7_i32, 8, 9].into_array(), + ]) + .into_array(); + let layout = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + + let whole_file = ScanBuilder::try_new( + &layout, + Arc::::clone(&segments) as Arc, + session.clone(), + )? + .full_file_splits()?; + let restricted = ScanBuilder::try_new(&layout, segments, session.clone())? + .with_row_range(3..5) + .full_file_splits()?; + + assert_eq!(whole_file, restricted); + assert_eq!(whole_file.first(), Some(&0)); + assert_eq!(whole_file.last(), Some(&9)); + Ok(()) + }) +} + +#[rstest] +#[case(vec![0..9])] +#[case(vec![0..3, 3..9])] +#[case(vec![0..4, 4..7, 7..9])] +fn row_ranges_partition_the_file_exactly_once( + #[case] row_ranges: Vec>, +) -> VortexResult<()> { + block_on(|handle| async move { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = ChunkedArray::from_iter([ + buffer![1_i32, 2, 3].into_array(), + buffer![4_i32, 5, 6].into_array(), + buffer![7_i32, 8, 9].into_array(), + ]) + .into_array(); + let layout = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + + let mut scanned = Vec::new(); + for row_range in row_ranges { + let chunk = ScanBuilder::try_new( + &layout, + Arc::::clone(&segments) as Arc, + session.clone(), + )? + .with_filter(gt(root(), lit(2_i32))) + .with_row_range(row_range) + .into_array_stream()? + .read_all() + .await?; + scanned.push(chunk); + } + + let actual = ChunkedArray::from_iter(scanned).into_array(); + let expected = PrimitiveArray::from_iter(3_i32..10).into_array(); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +} From daa0791bf42694c29d54fda3d1eb15617e0289c6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 11 Aug 2026 14:15:38 +0000 Subject: [PATCH 13/13] Trace plan-v2 partial file ranges Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MjXw22CLhx6NgA8mxjJCmQ --- vortex-datafusion/src/persistent/opener.rs | 30 ++++++++++++++-------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 64ebcd5d2aa..7acd14086b0 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -376,13 +376,12 @@ impl FileOpener for VortexOpener { "plan-v2 scans do not support VortexAccessPlan" )); } - let filter_mode = if std::env::var("VORTEX_PLAN_V2_FILTER_MODE").as_deref() - == Ok("adaptive") - { - FilterMode::Adaptive - } else { - FilterMode::Parallel - }; + let filter_mode = + if std::env::var("VORTEX_PLAN_V2_FILTER_MODE").as_deref() == Ok("adaptive") { + FilterMode::Adaptive + } else { + FilterMode::Parallel + }; let mut scan_builder = PlanScanBuilder::try_new( vxf.footer().layout(), vxf.segment_source(), @@ -438,10 +437,19 @@ impl FileOpener for VortexOpener { }, )?; let Some(row_range) = - split_aligned_row_range(byte_range, file_splits.as_ref()) + split_aligned_row_range(byte_range.clone(), file_splits.as_ref()) else { + tracing::debug!( + ?byte_range, + "plan-v2 byte range owns no natural split" + ); return Ok(stream::empty().boxed()); }; + tracing::debug!( + ?byte_range, + ?row_range, + "plan-v2 scanning a partial file range" + ); scan_builder = scan_builder.with_row_range(row_range); } } @@ -568,9 +576,9 @@ impl FileOpener for VortexOpener { .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {location}" - )))) + DataFusionError::External(Box::new( + e.with_context(format!("Failed to read Vortex file: {location}")), + )) }) .boxed() };