diff --git a/Cargo.lock b/Cargo.lock index c05aae6923e..03093ddbcef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10544,6 +10544,27 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-scan-v2" +version = "0.1.0" +dependencies = [ + "futures", + "itertools 0.14.0", + "parking_lot", + "tracing", + "tracing-subscriber", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-file", + "vortex-io", + "vortex-layout", + "vortex-mask", + "vortex-scan", + "vortex-session", + "vortex-utils", +] + [[package]] name = "vortex-sequence" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 36aa5b2ac9e..14ebc128896 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ members = [ "vortex-btrblocks", "vortex-layout", "vortex-scan", + "vortex-scan-v2", "vortex-file", "vortex-ipc", "vortex", @@ -323,6 +324,7 @@ vortex-proto = { version = "0.1.0", path = "./vortex-proto", default-features = vortex-row = { version = "0.1.0", path = "./vortex-row", default-features = false } vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-features = false } vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = false } +vortex-scan-v2 = { version = "0.1.0", path = "./vortex-scan-v2", default-features = false } vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-features = false } vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false } vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false } 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..8787970d839 --- /dev/null +++ b/docs/developer-guide/internals/scan-planning.md @@ -0,0 +1,68 @@ +# 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. + +## Execution + +Each operator executes over a row range and selection mask. `SegmentScan` reads its segment, +structural operators combine their children, and `Eval` applies the remaining derived work. +`vortex-scan-v2` copies the existing scan orchestration around this API, so the original +`LayoutReader` scanner is untouched while the plan-native path is developed. + +## Future work + +Still to come: a plan registry and foreign operator placeholder so third-party operators survive +a round trip, and a serialization envelope. diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index e7c83ec2950..d4ce3912a40 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -297,7 +297,7 @@ fn row_idx_dtype() -> DType { } // Returns a SequenceArray representing the row indices for the given row range, -fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { +pub(crate) fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { Sequence::try_new( PValue::U64(row_offset + row_range.start), PValue::U64(1), diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..7a954155835 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -358,7 +358,11 @@ impl ZonedLayout { } impl ZonedData { - fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> { + pub(crate) fn zone_len(&self) -> usize { + self.zone_len + } + + pub(crate) fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> { match &self.zone_map_schema { ZoneMapSchema::LegacyStats(stats) => stats .iter() diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 2248aaed97c..d750dca73ae 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -18,6 +18,7 @@ use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::fns::sum::normalize_legacy_partial_array; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -92,7 +93,7 @@ impl ZoneMap { Ok(unsafe { Self::new_unchecked(column_dtype, array, aggregate_fns, zone_len, row_count) }) } - pub(super) unsafe fn new_unchecked( + pub(crate) unsafe fn new_unchecked( column_dtype: DType, array: StructArray, aggregate_fns: Arc<[AggregateFnRef]>, @@ -152,19 +153,32 @@ impl ZoneMap { session: &VortexSession, ) -> VortexResult { let mut ctx = session.create_execution_ctx(); - let num_zones = self.array.len(); - let predicate = self.lower_stats(predicate.clone())?; + self.applied_predicate(predicate)? + .null_as_false() + .execute(&mut ctx) + } - let array = self.array.clone().into_array(); - let applied = array.apply_bound(&predicate)?; + /// Evaluates a pruning predicate while preserving unknown (null) proof values. + pub(crate) fn evaluate( + &self, + predicate: &BoundExpression, + session: &VortexSession, + ) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + self.applied_predicate(predicate)? + .execute::(&mut ctx) + } + fn applied_predicate(&self, predicate: &BoundExpression) -> VortexResult { + let num_zones = self.array.len(); + let predicate = self.lower_stats(predicate.clone())?; + let applied = self.array.clone().into_array().apply_bound(&predicate)?; if !contains_row_count(&applied) { - return applied.null_as_false().execute(&mut ctx); + return Ok(applied); } let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?; - let substituted = substitute_row_count(applied, &row_count_array)?; - substituted.null_as_false().execute(&mut ctx) + substitute_row_count(applied, &row_count_array) } fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { 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/execution.rs b/vortex-layout/src/plan/execution.rs new file mode 100644 index 00000000000..5d9a3c5829f --- /dev/null +++ b/vortex-layout/src/plan/execution.rs @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::future::BoxFuture; +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::segments::SegmentSource; + +/// Future resolving to the array produced by a physical plan. +pub type PlanArrayFuture = BoxFuture<'static, VortexResult>; + +/// Runtime dependencies shared by every node in a plan execution. +#[derive(Clone)] +pub struct PlanExecutionContext { + segment_source: Arc, + session: VortexSession, +} + +impl PlanExecutionContext { + /// Creates an execution context over a segment source and Vortex session. + pub fn new(segment_source: Arc, session: VortexSession) -> Self { + Self { + segment_source, + session, + } + } + + /// Returns the segment source used to satisfy leaf reads. + pub fn segment_source(&self) -> &Arc { + &self.segment_source + } + + /// Returns the Vortex session used for array decoding and expression execution. + pub fn session(&self) -> &VortexSession { + &self.session + } +} diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs new file mode 100644 index 00000000000..675cb95e14a --- /dev/null +++ b/vortex-layout/src/plan/lower.rs @@ -0,0 +1,164 @@ +// 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 std::sync::Arc; + +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::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; +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; +use crate::plan::ZonedPlan; + +/// 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()); + } + if layout.is::() || layout.is::() { + return Ok(lower_zoned(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(), + layout.has_all_values_referenced(), + 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) + }) +} + +fn lower_zoned(layout: &LayoutRef) -> VortexResult { + // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones. + let metadata = if let Some(layout) = layout.as_opt::() { + layout.data() + } else if let Some(layout) = layout.as_opt::() { + layout.data() + } else { + vortex_bail!("Zoned plan requires a zoned layout") + }; + Ok(ZonedPlan::from_children( + layout.dtype().clone(), + layout.row_count(), + lazy_children(Arc::clone(layout), vec![0, 1]), + u64::try_from(metadata.zone_len())?, + metadata.aggregate_fns(), + )) +} diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs new file mode 100644 index 00000000000..c630b64cc2a --- /dev/null +++ b/vortex-layout/src/plan/mod.rs @@ -0,0 +1,87 @@ +// 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 execution; +mod lower; +mod optimize; +pub mod optimizer; +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 execution::PlanArrayFuture; +pub use execution::PlanExecutionContext; +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::RowIdxPartition; +pub use plans::RowIdxPartitionPlan; +pub use plans::RowIdxPlan; +pub use plans::RowIdxPlanMetadata; +pub use plans::RowIdxValues; +pub use plans::RowIdxValuesData; +pub use plans::RowIdxValuesPlan; +pub use plans::RowIdxValuesPlanMetadata; +pub use plans::SegmentScan; +pub use plans::SegmentScanData; +pub use plans::SegmentScanPlan; +pub use plans::Take; +pub use plans::TakeData; +pub use plans::TakePlan; +pub use plans::Zoned; +pub use plans::ZonedData; +pub use plans::ZonedPlan; +pub use plans::row_idx_dtype; +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..8efd498ca08 --- /dev/null +++ b/vortex-layout/src/plan/optimize.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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; + +use crate::plan::Eval; +use crate::plan::PlanRef; + +/// Optimizes `plan`, preserving its dtype and row domain. +pub fn optimize(plan: PlanRef) -> VortexResult { + if let Some(eval) = plan.as_opt::() { + return eval.optimize_top_down(None); + } + + 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 new file mode 100644 index 00000000000..301c7ffce61 --- /dev/null +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -0,0 +1,51 @@ +// 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; +use vortex_error::VortexResult; + +use super::Concat; +use super::Pack; +use super::PlanRef; +use super::RowIdx; +use super::Take; +use super::Zoned; +use super::plans::ExpressionConcatRule; +use super::plans::ExpressionPackRule; +use super::plans::ExpressionRowIdxRule; +use super::plans::ExpressionTakeRule; +use super::plans::ExpressionZonedRule; + +static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionConcatRule); +static EXPRESSION_TAKE_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionTakeRule); +static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule); +static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionPackRule); +static EXPRESSION_ZONED_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionZonedRule); + +static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ + &EXPRESSION_CONCAT_RULE, + &EXPRESSION_TAKE_RULE, + &EXPRESSION_ROW_IDX_RULE, + &EXPRESSION_PACK_RULE, + &EXPRESSION_ZONED_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/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) + } +} diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs new file mode 100644 index 00000000000..fe2e081fd42 --- /dev/null +++ b/vortex-layout/src/plan/plans/concat.rs @@ -0,0 +1,249 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::future; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::TryStreamExt; +use futures::stream::FuturesOrdered; +use vortex_array::Canonical; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::ChunkedArray; +use vortex_array::dtype::DType; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::label_bound_tree; +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::layouts::row_idx::RowIdx as RowIdxFn; +use crate::plan::Eval; +use crate::plan::EvalPlan; +use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; +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)] +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 execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "Concat row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + vortex_ensure!( + mask.len() == usize::try_from(row_range.end - row_range.start)?, + "Concat mask length mismatch" + ); + if row_range.is_empty() { + let empty = Canonical::empty(plan.dtype()).into_array(); + return Ok(future::ready(Ok(empty)).boxed()); + } + + let mut chunk_futures = Vec::new(); + for (chunk, &chunk_offset) in plan.children().iter().zip(plan.row_offsets()) { + let chunk = chunk?; + let chunk_end = chunk_offset + .checked_add(chunk.row_count()) + .ok_or_else(|| vortex_err!("Chunk row offset overflow"))?; + let start = row_range.start.max(chunk_offset); + let end = row_range.end.min(chunk_end); + if start < end { + let child_range = start - chunk_offset..end - chunk_offset; + let mask_range = usize::try_from(start - row_range.start)? + ..usize::try_from(end - row_range.start)?; + chunk_futures.push(chunk.execute(ctx, &child_range, mask.slice(mask_range))?); + } + } + + Ok(async move { + let chunks: Vec<_> = FuturesOrdered::from_iter(chunk_futures) + .try_collect() + .await?; + vortex_ensure!(!chunks.is_empty(), "Non-empty row range selected no chunks"); + if chunks.len() == 1 { + return Ok(chunks.into_iter().next().vortex_expect("one chunk")); + } + let dtype = chunks[0].dtype().clone(); + Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) + } + .boxed()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + 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 new file mode 100644 index 00000000000..a3e5523200d --- /dev/null +++ b/vortex-layout/src/plan/plans/eval.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::fmt; +use std::ops::Range; + +use futures::FutureExt; +use vortex_array::EmptyMetadata; +use vortex_array::MaskFuture; +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; + +use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; +use crate::plan::PlanId; +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)] +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 execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let child = plan.child_plan()?.execute(ctx, row_range, mask)?; + let expression = plan.expression().clone(); + Ok(async move { child.await?.apply_bound(&expression) }.boxed()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + if index == 0 { + Cow::Borrowed("child") + } else { + Cow::Owned(format!("child[{index}]")) + } + } +} + +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/list_pack.rs b/vortex-layout/src/plan/plans/list_pack.rs new file mode 100644 index 00000000000..904446049c6 --- /dev/null +++ b/vortex-layout/src/plan/plans/list_pack.rs @@ -0,0 +1,249 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::try_join; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +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::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; +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 execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "ListPack row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + let row_count = usize::try_from(row_range.end - row_range.start)?; + vortex_ensure!(mask.len() == row_count, "ListPack mask length mismatch"); + + let offsets_range = row_range.start + ..row_range + .end + .checked_add(1) + .ok_or_else(|| vortex_err!("List offsets range overflow"))?; + let offsets = plan.offsets()?.execute( + ctx, + &offsets_range, + MaskFuture::new_true(row_count.saturating_add(1)), + )?; + let validity = plan + .validity()? + .map(|validity| validity.execute(ctx, row_range, MaskFuture::new_true(row_count))) + .transpose()?; + let elements = plan.elements()?; + let execution = ctx.clone(); + let dtype = plan.dtype().clone(); + let nullability = dtype.nullability(); + + Ok(async move { + let (offsets, mask) = try_join!(offsets, mask)?; + if mask.all_false() { + return Ok(Canonical::empty(&dtype).into_array()); + } + + let elements_range = elements_range_from_offsets(&offsets, execution.session())?; + let elements_count = usize::try_from(elements_range.end - elements_range.start)?; + let elements = elements + .execute( + &execution, + &elements_range, + MaskFuture::new_true(elements_count), + )? + .await?; + let validity = match validity { + Some(validity) => Some(validity.await?), + None => None, + }; + let offsets = rebase_offsets(offsets, elements_range.start)?; + // SAFETY: lowering from a list layout guarantees compatible elements and monotonically + // increasing offsets. Rebasing preserves the represented list lengths. + let list = unsafe { + ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) + } + .into_array(); + if mask.all_true() { + Ok(list) + } else { + list.filter(mask) + } + } + .boxed()) + } + + 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}]")), + } + } +} + +fn elements_range_from_offsets( + offsets: &ArrayRef, + session: &vortex_session::VortexSession, +) -> VortexResult> { + if offsets.is_empty() { + return Ok(0..0); + } + let mut ctx = session.create_execution_ctx(); + let start = offsets + .execute_scalar(0, &mut ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value must fit in u64"); + let end = offsets + .execute_scalar(offsets.len() - 1, &mut ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value must fit in u64"); + Ok(start..end) +} + +fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { + if first == 0 { + return Ok(offsets); + } + let constant = ConstantArray::new(first, offsets.len()) + .into_array() + .cast(offsets.dtype().clone())?; + offsets.binary(constant, Operator::Sub) +} + +fn create_validity(validity: Option, nullability: Nullability) -> Validity { + match validity { + Some(validity) => Validity::Array(validity), + None if nullability.is_nullable() => Validity::AllValid, + None => Validity::NonNullable, + } +} diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs new file mode 100644 index 00000000000..d0af735685e --- /dev/null +++ b/vortex-layout/src/plan/plans/mod.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod concat; +pub(crate) mod eval; +mod list_pack; +mod pack; +mod row_idx; +mod row_idx_partition; +mod row_idx_values; +mod segment_scan; +mod take; +mod zoned; + +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; +pub(crate) use row_idx::ExpressionRowIdxRule; +pub use row_idx::RowIdx; +pub use row_idx::RowIdxData; +pub use row_idx::RowIdxPlan; +pub use row_idx::RowIdxPlanMetadata; +pub use row_idx_partition::RowIdxPartition; +pub use row_idx_partition::RowIdxPartitionPlan; +pub use row_idx_values::RowIdxValues; +pub use row_idx_values::RowIdxValuesData; +pub use row_idx_values::RowIdxValuesPlan; +pub use row_idx_values::RowIdxValuesPlanMetadata; +pub use row_idx_values::row_idx_dtype; +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::TakeData; +pub use take::TakePlan; +pub(crate) use zoned::ExpressionZonedRule; +pub use zoned::Zoned; +pub use zoned::ZonedData; +pub use zoned::ZonedPlan; diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs new file mode 100644 index 00000000000..939965bc3c6 --- /dev/null +++ b/vortex-layout/src/plan/plans/pack.rs @@ -0,0 +1,476 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::ops::Range; + +use futures::FutureExt; +use futures::try_join; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::StructArray; +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_array::validity::Validity; +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::PlanArrayFuture; +use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; +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)] +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 execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "Pack row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + vortex_ensure!( + mask.len() == usize::try_from(row_range.end - row_range.start)?, + "Pack mask length mismatch" + ); + let names = plan.fields().names().clone(); + let field_count = plan.nfields(); + let mut field_futures = Vec::with_capacity(field_count); + for index in 0..field_count { + let child = field_plan(plan, index)?; + field_futures.push(child.execute(ctx, row_range, mask.clone())?); + } + let validity = plan + .validity()? + .map(|validity| validity.execute(ctx, row_range, mask.clone())) + .transpose()?; + let output_mask = mask; + + Ok(async move { + let fields = futures::future::try_join_all(field_futures); + let validity = async move { + match validity { + Some(validity) => validity.await.map(Some), + None => Ok(None), + } + }; + let (fields, validity) = try_join!(fields, validity)?; + let len = output_mask.await?.true_count(); + let validity = validity.map_or(Validity::NonNullable, Validity::Array); + Ok(StructArray::try_new(names, fields, len, validity)?.into_array()) + } + .boxed()) + } + + 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}]")) + } +} + +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::