diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index 52d4e6a545a..18fd199cc6a 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -103,6 +103,7 @@ jobs: env: VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + VORTEX_USE_PLAN_V2: "1" # Makes python output nicer COLUMNS: 120 strategy: diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..6c106617511 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9980,6 +9980,7 @@ dependencies = [ "url", "vortex", "vortex-arrow", + "vortex-scan-v2", "vortex-utils", ] @@ -10544,6 +10545,29 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-scan-v2" +version = "0.1.0" +dependencies = [ + "bit-vec", + "futures", + "itertools 0.14.0", + "parking_lot", + "rstest", + "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 44e154eadc3..594d48e570e 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/internals/scan-planning.md b/docs/developer-guide/internals/scan-planning.md index e03131924ac..8787970d839 100644 --- a/docs/developer-guide/internals/scan-planning.md +++ b/docs/developer-guide/internals/scan-planning.md @@ -55,9 +55,14 @@ ID, dtype, row count, and lazy children. Only the unsized tail containing the vt 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 -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. +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-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 49fb22d4f59..56afe335f85 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -38,6 +38,7 @@ tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-arrow = { workspace = true } +vortex-scan-v2 = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 89abfe68d18..ccb06a6fecd 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -3,6 +3,7 @@ use std::ops::Range; use std::sync::Arc; +use std::sync::Once; use std::sync::Weak; use arrow_array::RecordBatchOptions; @@ -42,6 +43,10 @@ use tracing::Instrument; use vortex::array::VortexSessionExecute; use vortex::error::VortexError; use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::expr::BoundExpression; +use vortex::expr::Expression; +use vortex::expr::root; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; @@ -50,6 +55,8 @@ use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; +use vortex_scan_v2::FilterMode; +use vortex_scan_v2::ScanBuilder as PlanScanBuilder; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; @@ -327,40 +334,6 @@ impl FileOpener for VortexOpener { .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; let projector = leftover_projection.make_projector(&stream_schema)?; - // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once. - let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { - Entry::Occupied(mut occupied_entry) => { - if let Some(reader) = occupied_entry.get().upgrade() { - tracing::trace!("reusing layout reader for {}", occupied_entry.key()); - reader - } else { - tracing::trace!("creating layout reader for {}", occupied_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!( - "Failed to create layout reader: {e}" - )) - })?; - occupied_entry.insert(Arc::downgrade(&reader)); - reader - } - } - Entry::Vacant(vacant_entry) => { - tracing::trace!("creating layout reader for {}", vacant_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!("Failed to create layout reader: {e}")) - })?; - vacant_entry.insert(Arc::downgrade(&reader)); - - reader - } - }; - - let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); - - if let Some(vortex_plan) = file.extensions.get::() { - scan_builder = vortex_plan.apply_to_builder(scan_builder); - } - let filter = filter .and_then(|f| { // Verify that all filters we've accepted from DataFusion get pushed down. @@ -396,76 +369,220 @@ impl FileOpener for VortexOpener { .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) .transpose() .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; + let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); + let stream = if std::env::var("VORTEX_USE_PLAN_V2").as_deref() == Ok("1") { + if file.extensions.get::().is_some() { + return Err(exec_datafusion_err!( + "plan-v2 scans do not support VortexAccessPlan" + )); + } + let filter_mode = + if std::env::var("VORTEX_PLAN_V2_FILTER_MODE").as_deref() == Ok("adaptive") { + FilterMode::Adaptive + } else { + FilterMode::Parallel + }; + let mut scan_builder = PlanScanBuilder::try_new( + vxf.footer().layout(), + vxf.segment_source(), + session.clone(), + ) + .map_err(|error| DataFusionError::External(Box::new(error)))? + .with_projection( + unbind(&scan_projection) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ) + .with_some_filter( + filter + .as_ref() + .map(unbind) + .transpose() + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ) + .with_filter_mode(filter_mode); + if let Some(limit) = limit + && filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + if let Some(concurrency) = scan_concurrency { + scan_builder = scan_builder.with_concurrency(concurrency); + } - if let Some(limit) = limit - && filter.is_none() - { - scan_builder = scan_builder.with_limit(limit); - } - - if let Some(concurrency) = scan_concurrency { - scan_builder = scan_builder.with_concurrency(concurrency); - } + announce_plan_v2(filter_mode); + + // DataFusion hands each partition a byte range of the file. Translate it to the + // row range whose natural splits that byte range owns, so partitions of one file + // cover it exactly once. + if let Some(file_range) = file.range.as_ref() { + let byte_range = Range { + start: u64::try_from(file_range.start).map_err(|_| { + exec_datafusion_err!("Vortex file range start is negative") + })?, + end: u64::try_from(file_range.end).map_err(|_| { + exec_datafusion_err!("Vortex file range end is negative") + })?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + let file_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + file.object_meta.size, + || { + scan_builder.full_file_splits().map_err(|e| { + exec_datafusion_err!( + "Failed to compute plan-v2 natural splits: {e}" + ) + }) + }, + )?; + let Some(row_range) = + split_aligned_row_range(byte_range.clone(), file_splits.as_ref()) + else { + tracing::debug!( + ?byte_range, + "plan-v2 byte range owns no natural split" + ); + return Ok(stream::empty().boxed()); + }; + tracing::debug!( + ?byte_range, + ?row_range, + "plan-v2 scanning a partial file range" + ); + scan_builder = scan_builder.with_row_range(row_range); + } + } - // Set before the byte-range translation below, which computes natural splits for - // the fields the scan's projection and filter reference. - scan_builder = scan_builder - .with_projection(scan_projection) - .with_some_filter(filter); - - if let Some(file_range) = file.range { - let byte_range = Range { - start: u64::try_from(file_range.start) - .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, + let location = file.object_meta.location.clone(); + let session = session.clone(); + scan_builder + .with_ordered(has_output_ordering) + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + .into_stream() + .map_err(|e| exec_datafusion_err!("Failed to create plan-v2 stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new(e.with_context(format!( + "Failed to read Vortex file with plan v2: {location}" + )))) + }) + .boxed() + } else { + // Share layout readers between partitions so each file's layout is read once. + let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { + Entry::Occupied(mut occupied_entry) => { + if let Some(reader) = occupied_entry.get().upgrade() { + tracing::trace!("reusing layout reader for {}", occupied_entry.key()); + reader + } else { + tracing::trace!("creating layout reader for {}", occupied_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + occupied_entry.insert(Arc::downgrade(&reader)); + reader + } + } + Entry::Vacant(vacant_entry) => { + tracing::trace!("creating layout reader for {}", vacant_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + vacant_entry.insert(Arc::downgrade(&reader)); + reader + } }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_splits = natural_splits_for_file( - natural_splits.as_ref(), - &file.object_meta.location, - &scan_builder, - file.object_meta.size, - )?; - - let Some(row_range) = - split_aligned_row_range(byte_range, natural_splits.as_ref()) - else { - return Ok(stream::empty().boxed()); - }; - scan_builder = scan_builder - .with_row_range(row_range) - // Hand the shared full-file boundaries back to the scan so prepare() - // skips its own layout walk. - .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + let mut scan_builder = + ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); + if let Some(vortex_plan) = file.extensions.get::() { + scan_builder = vortex_plan.apply_to_builder(scan_builder); + } + if let Some(limit) = limit + && filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + if let Some(concurrency) = scan_concurrency { + scan_builder = scan_builder.with_concurrency(concurrency); } - } - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); - let stream = scan_builder - .with_metrics_registry(metrics_registry) - .with_ordered(has_output_ordering) - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) - .into_stream() - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - file.object_meta.location - )))) - }) + // Set before translating byte ranges because natural splits depend on the fields + // referenced by the projection and filter. + scan_builder = scan_builder + .with_projection(scan_projection) + .with_some_filter(filter); + if let Some(file_range) = file.range { + let byte_range = Range { + start: u64::try_from(file_range.start).map_err(|_| { + exec_datafusion_err!("Vortex file range start is negative") + })?, + end: u64::try_from(file_range.end).map_err(|_| { + exec_datafusion_err!("Vortex file range end is negative") + })?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + let natural_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + file.object_meta.size, + || { + scan_builder.full_file_splits().map_err(|e| { + exec_datafusion_err!( + "Failed to compute Vortex natural splits: {e}" + ) + }) + }, + )?; + let Some(row_range) = + split_aligned_row_range(byte_range, natural_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + scan_builder = scan_builder + .with_row_range(row_range) + .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + } + } + + let location = file.object_meta.location.clone(); + let session = session.clone(); + scan_builder + .with_metrics_registry(metrics_registry) + .with_ordered(has_output_ordering) + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + .into_stream() + .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new( + e.with_context(format!("Failed to read Vortex file: {location}")), + )) + }) + .boxed() + }; + let stream = stream .map(move |batch| { let batch = if projector.projection().as_ref().is_empty() { batch @@ -494,6 +611,23 @@ impl FileOpener for VortexOpener { } } +fn unbind(expression: &BoundExpression) -> VortexResult { + if expression.is_root() { + return Ok(root()); + } + Expression::try_new( + expression + .as_scalar() + .vortex_expect("non-root bound expression must have a scalar function") + .clone(), + expression + .children() + .iter() + .map(unbind) + .collect::>>()?, + ) +} + /// A file's natural split boundaries plus the precomputed byte each split is assigned to, /// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row /// boundaries with a binary search instead of re-projecting every split per partition. @@ -549,11 +683,11 @@ impl NaturalSplits { } /// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. -fn natural_splits_for_file( +fn natural_splits_for_file( natural_splits: &DashMap>, path: &Path, - scan_builder: &ScanBuilder, total_size: u64, + row_boundaries: impl FnOnce() -> DFResult>, ) -> DFResult> { if let Some(splits) = natural_splits.get(path) { return Ok(Arc::clone(splits.value())); @@ -565,27 +699,23 @@ fn natural_splits_for_file( match natural_splits.entry(path.clone()) { Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), Entry::Vacant(entry) => { - let splits = compute_natural_splits(scan_builder, total_size)?; + let splits = Arc::new(NaturalSplits::new(row_boundaries()?.into(), total_size)); entry.insert(Arc::clone(&splits)); Ok(splits) } } } -/// Walk the layout tree to compute the file's full natural split boundaries for the fields -/// referenced by the scan's projection and filter. -fn compute_natural_splits( - scan_builder: &ScanBuilder, - total_size: u64, -) -> DFResult> { - let row_boundaries = scan_builder - .full_file_splits() - .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; - - Ok(Arc::new(NaturalSplits::new( - row_boundaries.into(), - total_size, - ))) +/// Logs once per process that scans are running through plan-v2, so a benchmark or CI job log +/// records which scan path produced its timings. +fn announce_plan_v2(filter_mode: FilterMode) { + static ANNOUNCED: Once = Once::new(); + ANNOUNCED.call_once(|| { + tracing::info!( + ?filter_mode, + "Vortex file scans are using the plan-v2 scan path (VORTEX_USE_PLAN_V2=1)" + ); + }); } /// Translate a DataFusion byte range to the contiguous natural split ranges it owns. diff --git a/vortex-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 c84c0b443dd..157f121d4c9 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -16,6 +16,7 @@ use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; use vortex_array::aggregate_fn::fns::all_null::AllNull; use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -84,7 +85,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]>, @@ -144,19 +145,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/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 index e8911734549..1e388de220f 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -6,6 +6,8 @@ //! 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; @@ -24,6 +26,8 @@ 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; @@ -31,6 +35,7 @@ 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. /// @@ -52,6 +57,9 @@ pub fn lower(layout: &LayoutRef) -> VortexResult { 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() @@ -118,6 +126,7 @@ fn lower_dict(layout: &DictLayout) -> VortexResult { TakePlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), + layout.has_all_values_referenced(), lazy_children(layout.to_layout(), vec![1, 0]), ) }) @@ -151,3 +160,21 @@ fn lazy_children(layout: LayoutRef, slots: Vec) -> PlanChildren { 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 index 06f7acd0031..c630b64cc2a 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -9,8 +9,10 @@ mod children; mod display; +mod execution; mod lower; mod optimize; +pub mod optimizer; mod plans; mod typed; mod vtable; @@ -21,6 +23,8 @@ 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; @@ -37,13 +41,24 @@ 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; 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 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 index 2e2893fde5f..c0a38abace2 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -2,20 +2,40 @@ // 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)] @@ -136,7 +156,100 @@ impl PlanVTable for Concat { 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::try_new(expression.clone(), chunk?)?.into_plan())) + .collect::>>()?; + Ok(Some( + ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(), + )) + } +} diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index 70a645da687..01df14a9531 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -3,20 +3,34 @@ 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_error::vortex_bail; 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)] @@ -105,6 +119,17 @@ impl PlanVTable for Eval { 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") @@ -123,3 +148,94 @@ fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> V } Ok(()) } + +impl EvalPlan { + /// Optimizes this plan top-down, applying parent-reduction rules as they become applicable. + /// + /// `blocked_child_type` suppresses one rule re-firing on its own residual output, which would + /// otherwise loop when a rewrite leaves an expression above the same child kind. + pub(crate) fn optimize_top_down( + &self, + blocked_child_type: Option, + ) -> VortexResult { + if self.expression().is_root() { + return optimize(self.child_plan()?); + } + + let child = self.child_plan()?; + let child_type = child.id(); + let parent = EvalPlan::try_new(self.expression().clone(), child.clone())?.into_plan(); + if blocked_child_type != Some(child_type) + && let Some(rewritten) = reduce_parent(&parent, 0)? + { + return Self::optimize_rewrite(rewritten, child_type); + } + + let child = optimize(child)?; + + let child_type = child.id(); + let parent = EvalPlan::try_new(self.expression().clone(), child)?.into_plan(); + if blocked_child_type != Some(child_type) + && let Some(rewritten) = reduce_parent(&parent, 0)? + { + return Self::optimize_rewrite(rewritten, child_type); + } + Ok(parent) + } + + fn optimize_rewrite(rewritten: PlanRef, previous_child_type: PlanId) -> VortexResult { + let Some(eval) = rewritten.as_opt::() else { + return optimize(rewritten); + }; + // A residual expression may remain above the same child kind after a successful rewrite. + // Do not immediately apply that rule again; recursively optimize only the retained child. + let child_type = eval.child_plan()?.id(); + let blocked = (child_type == previous_child_type).then_some(previous_child_type); + eval.optimize_top_down(blocked) + } +} + +/// Rewrites partition accessors in `expression` to read from a partitioned root. +pub(crate) fn rewrite_partition_root( + expression: BoundExpression, + root_dtype: DType, + collapsed: &[(FieldName, FieldName)], +) -> VortexResult { + Ok(expression + .transform_down(|node| { + if let Some(value_name) = node + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + { + let partition_access = &node.children()[0]; + if let Some(partition_name) = partition_access + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + && partition_access.children()[0].is_root() + && collapsed.iter().any(|(partition, value)| { + partition == partition_name && value == value_name + }) + { + return Ok(Transformed { + value: BoundExpression::try_new( + GetItem.bind(partition_name.clone()), + [BoundExpression::new_root(root_dtype.clone())], + )?, + changed: true, + order: TraversalOrder::Skip, + }); + } + } + + if node.is_root() { + Ok(Transformed { + value: BoundExpression::new_root(root_dtype.clone()), + changed: true, + order: TraversalOrder::Skip, + }) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} diff --git a/vortex-layout/src/plan/plans/list_pack.rs b/vortex-layout/src/plan/plans/list_pack.rs index d9a23de818b..1c2992e3d74 100644 --- a/vortex-layout/src/plan/plans/list_pack.rs +++ b/vortex-layout/src/plan/plans/list_pack.rs @@ -2,18 +2,35 @@ // 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; @@ -116,6 +133,75 @@ impl PlanVTable for ListPack { validate_children(plan.dtype(), plan.row_count(), children) } + 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"), @@ -188,3 +274,42 @@ fn validate_children(dtype: &DType, row_count: u64, children: &PlanChildren) -> } Ok(()) } + +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 index 2e4a6dbad5e..d0af735685e 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -2,31 +2,50 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod concat; -mod eval; +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 index 9cf47ebe4b7..005901bdd32 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -2,23 +2,52 @@ // 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)] @@ -174,6 +203,51 @@ impl PlanVTable for Pack { 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> { assert!( index < plan.children().len(), @@ -224,3 +298,260 @@ fn validate_validity_child(expected_row_count: u64, child: &PlanRef) -> VortexRe } Ok(()) } + +impl PackPlan { + /// Rebuilds this plan with only `fields`, which must be a subset of the current fields. + /// + /// Pruning is only sound for a non-nullable struct: dropping a field of a nullable struct + /// would drop the validity child that the remaining fields depend on. + pub(crate) fn with_pruned_fields( + &self, + fields: Vec<(FieldName, PlanRef)>, + ) -> VortexResult { + vortex_ensure!( + !self.dtype().is_nullable(), + "Cannot prune fields from a nullable Pack" + ); + let struct_fields = StructFields::from_iter( + fields + .iter() + .map(|(name, plan)| (name.clone(), plan.dtype().clone())), + ); + let field_plans = fields.into_iter().map(|(_, plan)| plan).collect::>(); + PackPlan::try_new( + struct_fields, + Nullability::NonNullable, + self.row_count(), + field_plans, + None, + ) + } +} + +/// Pushes an expression into the referenced fields of a [`Pack`], pruning the rest. +#[derive(Debug)] +pub(crate) struct ExpressionPackRule; + +impl PlanParentReduceRule for ExpressionPackRule { + type Parent = Eval; + + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + _child_idx: usize, + ) -> VortexResult> { + if child.dtype().is_nullable() { + return Ok(None); + } + + let expression = parent.expression(); + let fields = child.fields(); + let referenced_fields = + descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) + .get(&ExactBoundExpr(expression.clone())) + .vortex_expect("Bound expression missing free-field annotations") + .clone(); + let expanded = expand_struct_root(expression.clone(), fields)?; + let partitioned = + partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?; + + if partitioned.partition_names.is_empty() { + let selected_indices = fields + .names() + .iter() + .enumerate() + .filter_map(|(index, name)| referenced_fields.contains(name).then_some(index)) + .collect::>(); + if selected_indices.len() == fields.nfields() { + return Ok(None); + } + + let pruned_fields = selected_indices + .into_iter() + .map(|field_index| { + Ok(( + field_name(fields, field_index)?, + field_plan(child, field_index)?, + )) + }) + .collect::>>()?; + let rewritten = child.with_pruned_fields(pruned_fields)?.into_plan(); + return Ok(Some( + EvalPlan::try_new(expression.clone(), rewritten)?.into_plan(), + )); + } + + if partitioned.partition_names.len() == 1 { + let name = partitioned + .partition_names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition has no field"))?; + let index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, index)?; + let lowered = step_into_struct_field(expanded, name, field.dtype().clone())?; + return Ok(Some(EvalPlan::try_new(lowered, field)?.into_plan())); + } + + let residual = partitioned.root; + let mut collapsed = Vec::with_capacity(partitioned.partitions.len()); + let mut field_expressions = vec![None; fields.nfields()]; + for index in 0..partitioned.partitions.len() { + let name = &partitioned.partition_names[index]; + let partition = &partitioned.partitions[index]; + let field_index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, field_index)?; + let lowered = if let Some(pack) = partition + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + && partition.children().len() == 1 + { + let value_name = pack + .names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; + collapsed.push((name.clone(), value_name.clone())); + partition.children()[0].clone() + } else { + partition.clone() + }; + let lowered = step_into_struct_field(lowered, name, field.dtype().clone())?; + field_expressions[field_index] = Some(lowered); + } + + let mut pruned_fields = Vec::with_capacity(partitioned.partition_names.len()); + for (field_index, expression) in field_expressions.into_iter().enumerate() { + let Some(expression) = expression else { + continue; + }; + let field = field_plan(child, field_index)?; + pruned_fields.push(( + field_name(fields, field_index)?, + EvalPlan::try_new(expression, field)?.into_plan(), + )); + } + let rewritten = child.with_pruned_fields(pruned_fields)?.into_plan(); + let residual = rewrite_partition_root(residual, rewritten.dtype().clone(), &collapsed)?; + + Ok(Some(EvalPlan::try_new(residual, rewritten)?.into_plan())) + } +} + +fn field_name(fields: &StructFields, index: usize) -> VortexResult { + Ok(fields + .field_name(index) + .ok_or_else(|| vortex_err!("Struct field {index} has no name"))? + .clone()) +} + +fn field_plan(plan: &Plan, index: usize) -> VortexResult { + plan.child(index)? + .ok_or_else(|| vortex_err!("Struct field {index} has no plan")) +} + +fn expanded_struct_root( + root_dtype: &DType, + fields: &StructFields, +) -> VortexResult { + let root = BoundExpression::new_root(root_dtype.clone()); + let children = fields + .names() + .iter() + .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])) + .collect::>>()?; + bound_pack(fields.names().clone(), children) +} + +fn expand_struct_root( + expression: BoundExpression, + fields: &StructFields, +) -> VortexResult { + Ok(expression + .transform_down(|node| { + if node.is_root() { + return Ok(Transformed { + value: expanded_struct_root(node.dtype(), fields)?, + 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 scalar_fn.is::() { + return Ok(Transformed { + value: node, + changed: false, + order: TraversalOrder::Skip, + }); + } + + if let Some(selection) = scalar_fn.as_opt::