From 991458035d89e90a1a37c0252c07d9beafd51c91 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 10 Aug 2026 11:04:26 +0100 Subject: [PATCH 1/6] 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 9b8ffd722ce33c84fe6c1c78b6eda0968b858ec0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 10 Aug 2026 11:04:52 +0100 Subject: [PATCH 2/6] 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 7e3aae3986575745d773f7354da1527005a84a06 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 10 Aug 2026 11:05:18 +0100 Subject: [PATCH 3/6] 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 | 285 ++++++++++++++++++++++++ vortex-layout/src/plan/plans/take.rs | 49 ++++ vortex-layout/src/plan/tests.rs | 274 +++++++++++++++++++++++ 8 files changed, 803 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 414a685ec5d..2f65329d23c 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)] @@ -134,3 +140,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::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 05421eeda74..326d29081a3 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_session::registry::CachedId; @@ -16,6 +23,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)] @@ -99,3 +108,94 @@ impl PlanVTable for Eval { } } } + +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::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::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 156fcd5af4c..590d1aafd91 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -5,19 +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)] @@ -135,3 +156,267 @@ impl PlanVTable for Pack { Cow::Owned(format!("child[{index}]")) } } + +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::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::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::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::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::