From 29f589dc5af64656f9b052ced3a8417cb28c96a2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 13:42:41 +0000 Subject: [PATCH] feat(layout): planner-driven V2 scan over physical plans An opt-in scan (`VORTEX_SCAN_V2=1`) that lowers a layout once into a physical plan and runs each split through announce, pruning, filter (adaptive conjunct order) and projection planners. Each stage of a split runs as one pipeline `Scan`, and a split's scans share a `ValueCache` dropped with the split. Adds the V2 `ScanBuilder`/`RepeatedScanV2`, plan-derived splits, early IO announcements and file pruning. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hai26WGsNERxRHXpA3JboE Signed-off-by: Claude --- vortex-layout/src/layouts/chunked/mod.rs | 5 + vortex-layout/src/plan/physical/typed.rs | 8 + vortex-layout/src/scan/filter.rs | 6 +- vortex-layout/src/scan/mod.rs | 2 + vortex-layout/src/scan/planning/announce.rs | 409 +++++ vortex-layout/src/scan/planning/filter.rs | 511 ++++++ .../src/scan/planning/filter/tests.rs | 190 +++ vortex-layout/src/scan/planning/graph.rs | 257 +++ vortex-layout/src/scan/planning/mod.rs | 171 ++ vortex-layout/src/scan/planning/projection.rs | 568 +++++++ vortex-layout/src/scan/scan_builder.rs | 44 + vortex-layout/src/scan/v2/conjuncts.rs | 226 +++ vortex-layout/src/scan/v2/file.rs | 102 ++ vortex-layout/src/scan/v2/io.rs | 321 ++++ vortex-layout/src/scan/v2/legacy_tests.rs | 121 ++ vortex-layout/src/scan/v2/list_tests.rs | 56 + vortex-layout/src/scan/v2/lower.rs | 221 +++ vortex-layout/src/scan/v2/mod.rs | 113 ++ vortex-layout/src/scan/v2/prefetch.rs | 145 ++ vortex-layout/src/scan/v2/pruning.rs | 275 ++++ vortex-layout/src/scan/v2/repeated_scan.rs | 606 +++++++ vortex-layout/src/scan/v2/scan_builder.rs | 333 ++++ vortex-layout/src/scan/v2/share.rs | 148 ++ vortex-layout/src/scan/v2/split.rs | 90 ++ vortex-layout/src/scan/v2/splits.rs | 343 ++++ vortex-layout/src/scan/v2/stream.rs | 141 ++ vortex-layout/src/scan/v2/tests.rs | 1401 +++++++++++++++++ 27 files changed, 6812 insertions(+), 1 deletion(-) create mode 100644 vortex-layout/src/scan/planning/announce.rs create mode 100644 vortex-layout/src/scan/planning/filter.rs create mode 100644 vortex-layout/src/scan/planning/filter/tests.rs create mode 100644 vortex-layout/src/scan/planning/graph.rs create mode 100644 vortex-layout/src/scan/planning/mod.rs create mode 100644 vortex-layout/src/scan/planning/projection.rs create mode 100644 vortex-layout/src/scan/v2/conjuncts.rs create mode 100644 vortex-layout/src/scan/v2/file.rs create mode 100644 vortex-layout/src/scan/v2/io.rs create mode 100644 vortex-layout/src/scan/v2/legacy_tests.rs create mode 100644 vortex-layout/src/scan/v2/list_tests.rs create mode 100644 vortex-layout/src/scan/v2/lower.rs create mode 100644 vortex-layout/src/scan/v2/mod.rs create mode 100644 vortex-layout/src/scan/v2/prefetch.rs create mode 100644 vortex-layout/src/scan/v2/pruning.rs create mode 100644 vortex-layout/src/scan/v2/repeated_scan.rs create mode 100644 vortex-layout/src/scan/v2/scan_builder.rs create mode 100644 vortex-layout/src/scan/v2/share.rs create mode 100644 vortex-layout/src/scan/v2/split.rs create mode 100644 vortex-layout/src/scan/v2/splits.rs create mode 100644 vortex-layout/src/scan/v2/stream.rs create mode 100644 vortex-layout/src/scan/v2/tests.rs diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 459429046c2..2c8bd006a9d 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -96,6 +96,11 @@ impl VTable for Chunked { } impl Layout { + /// Cached chunk boundaries, including the final row count. + pub(crate) fn chunk_offsets(&self) -> &[u64] { + &self.chunk_offsets + } + /// Construct a chunked layout. pub fn new(row_count: u64, dtype: DType, children: Arc) -> Self { let offsets = chunk_offsets(children.as_ref()).vortex_expect("chunk row counts overflow"); diff --git a/vortex-layout/src/plan/physical/typed.rs b/vortex-layout/src/plan/physical/typed.rs index 1c6d99fd93c..2c96cd68b44 100644 --- a/vortex-layout/src/plan/physical/typed.rs +++ b/vortex-layout/src/plan/physical/typed.rs @@ -99,6 +99,14 @@ impl PlanRef { .ok_or_else(|| vortex_err!("Missing plan child {index}")) } + /// Borrows the child at `index`, or returns an error if it is out of bounds. + pub(crate) fn child_ref_required(&self, index: usize) -> VortexResult<&PlanRef> { + self.0 + .children + .get_ref(index)? + .ok_or_else(|| vortex_err!("Missing plan child {index}")) + } + /// Rebuilds this plan with `children` stored outside its erased operator data. pub fn with_children(&self, children: impl Into) -> VortexResult { self.dyn_plan().dyn_with_children(self, children.into()) diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 2cb36068913..f0c81b26348 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -55,7 +55,11 @@ fn bound_conjuncts(expr: &BoundExpression) -> Vec { impl FilterExpr { pub fn new(expr: BoundExpression) -> Self { - let conjuncts = bound_conjuncts(&expr); + Self::from_conjuncts(bound_conjuncts(&expr)) + } + + /// Builds a filter from conjuncts that are already split, without splitting them further. + pub(crate) fn from_conjuncts(conjuncts: Vec) -> Self { let num_conjuncts = conjuncts.len(); let dynamic_conjuncts = conjuncts.iter().map(DynamicExprUpdates::new).collect_vec(); diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index e08e4cb3f31..21c93583eea 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -6,6 +6,7 @@ mod filter; pub mod layout; mod limit; pub mod multi; +pub mod planning; pub mod repeated_scan; pub mod scan_builder; pub mod split_by; @@ -13,6 +14,7 @@ mod splits; mod tasks; #[cfg(test)] mod test; +pub mod v2; /// A heuristic for an ideal split size. /// diff --git a/vortex-layout/src/scan/planning/announce.rs b/vortex-layout/src/scan/planning/announce.rs new file mode 100644 index 00000000000..e6b7a640918 --- /dev/null +++ b/vortex-layout/src/scan/planning/announce.rs @@ -0,0 +1,409 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::LazyLock; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_scan::planning::next::Next; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; + +use crate::plan::physical::PlanRef; +use crate::scan::planning::FilterPlans; +use crate::scan::planning::ProjectionPlanner; +use crate::scan::planning::ScanPlans; +use crate::scan::planning::SelectedRows; +use crate::scan::planning::plan_selected; +use crate::scan::v2::prefetch::plan_segments; +use crate::scan::v2::splits::projection_splits; + +/// Announces the segments a split is likely to read, then hands the split to `next`. +/// +/// An announcement reads nothing itself, but the IO service can coalesce every announced segment +/// near one that it does read, as the layout reader's splits do by registering their reads when +/// their futures are built. +pub struct AnnouncePlanner { + plans: Option, + filter: Option, + /// The split, until it is handed to `next`. + selected: Option, + announced: bool, + next: Option, +} + +enum Continuation { + Next(Next), + Split(Option), + SparseProjection { rows: Option> }, +} + +impl AnnouncePlanner { + /// Creates a planner announcing the reads of `filter` and the projection over `selected`. + pub fn new( + plans: ScanPlans, + filter: Option, + selected: SelectedRows, + next: Next, + ) -> Self { + Self { + plans: Some(plans), + filter, + selected: Some(selected), + announced: false, + next: Some(Continuation::Next(next)), + } + } + + pub(super) fn for_split( + plans: ScanPlans, + pruning: Option, + filter: Option, + selected: SelectedRows, + ) -> Self { + Self { + plans: Some(plans), + filter, + selected: Some(selected), + announced: false, + next: Some(Continuation::Split(pruning)), + } + } + + pub(super) fn for_sparse_projection(plans: ScanPlans, selected: SelectedRows) -> Self { + // Input index masks retain their indices unless limited. Keep contiguous ranges on + // the projection path that preserves their chunk slices. + let has_gaps = selected + .mask + .first() + .zip(selected.mask.last()) + .is_some_and(|(first, last)| last - first + 1 > selected.mask.true_count()); + if !has_gaps { + return Self::for_split(plans, None, None, selected); + } + Self { + next: Some(Continuation::SparseProjection { rows: None }), + ..Self::for_split(plans, None, None, selected) + } + } + + fn announce( + &self, + selected: &SelectedRows, + project_before_filter: bool, + ) -> VortexResult<(IoBatch, Option>)> { + let plans = self + .plans + .as_ref() + .ok_or_else(|| vortex_err!("AnnouncePlanner has no plans"))?; + let rows = &selected.scope.rows; + let mut ids = Vec::new(); + let mut consumer = Vec::new(); + let mut announce = |plan: &PlanRef, rows| -> VortexResult<()> { + consumer.clear(); + plan_segments(plan, rows, &mut consumer)?; + consumer.sort_unstable(); + consumer.dedup(); + ids.extend(consumer.iter().copied()); + Ok(()) + }; + for filter in self.filter.iter().flat_map(FilterPlans::plans) { + announce(filter, rows.clone())?; + } + // V1 constructs projection futures before filtering, so a shared segment read remains + // live for each pending projection. Keep one registration per potential consumer. + let mut projection_rows: Option> = None; + if self.filter.is_none() || project_before_filter { + for rows in projection_splits(&plans.projection_starts, rows.clone()) { + // The input selection is already known, even before predicates run. Empty chunks + // cannot produce a projection consumer and must not keep read registrations alive. + let start = usize::try_from(rows.start - selected.scope.rows.start)?; + let end = usize::try_from(rows.end - selected.scope.rows.start)?; + if selected.mask.slice(start..end).all_false() { + continue; + } + projection_rows.get_or_insert_with(|| rows.clone()).end = rows.end; + announce(&plans.projection, rows)?; + } + } + let batch = ids + .into_iter() + .enumerate() + .map(|(index, id)| { + let location = plans + .locations + .get(*id as usize) + .ok_or_else(|| vortex_err!("segment {id} has no known location"))?; + Ok(IoRequest { + intent: IoIntent::Announce, + request: IoRequestId(u32::try_from(index)?), + target: location.target(), + }) + }) + .collect::>()?; + Ok((batch, projection_rows)) + } +} + +impl IoConsumer for AnnouncePlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for AnnouncePlanner { + fn state(&self) -> State { + if self.selected.is_some() { + State::NeedsCompute + } else { + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + static PROJECT_BEFORE_FILTER: LazyLock = LazyLock::new(|| { + !std::env::var("VORTEX_SCAN_PROJECT_ANNOUNCE").is_ok_and(|value| value == "0") + }); + if !self.announced { + self.announced = true; + let (batch, rows) = match &self.selected { + Some(selected) => self.announce(selected, *PROJECT_BEFORE_FILTER)?, + None => (Vec::new(), None), + }; + if let Some(Continuation::SparseProjection { rows: selected }) = &mut self.next { + *selected = rows; + } + if !batch.is_empty() { + return Ok(PlannerOutput::NeedsIO(batch)); + } + } + let Some(selected) = self.selected.take() else { + vortex_bail!("AnnouncePlanner: compute called after Done"); + }; + let scope = selected.scope.clone(); + let next = match self.next.take() { + Some(Continuation::Next(next)) => next(selected)?, + Some(Continuation::Split(pruning)) => { + let plans = self + .plans + .take() + .ok_or_else(|| vortex_err!("AnnouncePlanner has no plans"))?; + plan_selected(plans, pruning, self.filter.take(), selected)? + } + Some(Continuation::SparseProjection { rows }) => { + let mut plans = self + .plans + .take() + .ok_or_else(|| vortex_err!("AnnouncePlanner has no plans"))?; + let mut selected = selected; + if let Some(rows) = rows { + // Reuse the populated cuts found by announcements to trim empty edges, + // especially when only one cut contains rows. + let offset = selected.scope.rows.start; + let start = usize::try_from(rows.start - offset)?; + let end = usize::try_from(rows.end - offset)?; + selected.mask = selected.mask.slice(start..end); + selected.scope.rows = rows; + } + if selected.mask.true_count() <= selected.mask.len() / 8 { + // The graph follows each column's own chunk boundaries. Keep the original + // boundaries for announcements, which must still skip unselected chunks. + plans.projection_starts = Default::default(); + } + Box::new(ProjectionPlanner::new(plans, selected)) + } + None => vortex_bail!("AnnouncePlanner has no continuation"), + }; + Ok(PlannerOutput::Planner(scope, next)) + } +} + +#[cfg(test)] +#[allow( + clippy::single_range_in_vec_init, + reason = "cases list the ranges announced, and some announce one" +)] +mod tests { + use std::ops::Range; + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_buffer::Alignment; + use vortex_mask::Mask; + use vortex_scan::planning::planner::WorkScope; + use vortex_session::registry::ReadContext; + + use super::*; + use crate::plan::physical::ConcatPlan; + use crate::plan::physical::FilterPlan; + use crate::plan::physical::SegmentScanPlan; + use crate::scan::planning::SegmentLocation; + use crate::segments::SegmentId; + use crate::test::new_session; + + #[rstest] + #[case::default(false, &[150, 2800], None)] + #[case::one_cut(true, &[150], Some(vec![100..1000]))] + #[case::two_rows_in_one_cut(true, &[150, 160], Some(vec![100..1000]))] + #[case::two_cuts(true, &[150, 2800], Some(vec![100..2900]))] + #[case::contiguous_across_cuts(true, &[999, 1000], Some(vec![100..1000, 1000..2000]))] + fn announcements_skip_chunks_without_selected_rows( + #[case] sparse: bool, + #[case] indices: &[usize], + #[case] morsel_rows: Option>>, + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, NonNullable); + let chunks = (0..3) + .map(|segment| { + FilterPlan::new( + SegmentScanPlan::new( + dtype.clone(), + 1000, + SegmentId::from(segment), + ReadContext::new([]), + None, + ) + .into_plan(), + ) + .into_plan() + }) + .collect(); + let locations: Arc<[_]> = (0..3) + .map(|offset| SegmentLocation { + offset, + length: 1, + alignment: Alignment::none(), + }) + .collect(); + let selected = SelectedRows { + scope: WorkScope { + file_ordinal: 0, + rows: 100..2900, + }, + mask: Mask::from_indices(2800, indices.iter().map(|row| row - 100)), + }; + let plans = ScanPlans { + session: new_session(), + locations: Arc::clone(&locations), + projection: ConcatPlan::try_new(dtype, chunks)?.into_plan(), + projection_starts: Arc::from([0, 1000, 2000]), + row_offset: 0, + values: Default::default(), + }; + let mut planner = if sparse { + AnnouncePlanner::for_sparse_projection(plans, selected) + } else { + AnnouncePlanner::for_split(plans, None, None, selected) + }; + let PlannerOutput::NeedsIO(requests) = planner.compute()? else { + vortex_bail!("expected announcements"); + }; + assert!( + requests + .iter() + .all(|request| request.intent == IoIntent::Announce) + ); + let targets: Vec<_> = requests.iter().map(|request| request.target).collect(); + let mut expected: Vec<_> = indices + .iter() + .map(|row| locations[row / 1000].target()) + .collect(); + expected.dedup(); + assert_eq!(targets, expected); + if let Some(morsel_rows) = morsel_rows { + let PlannerOutput::Planner(_, mut projection) = planner.compute()? else { + vortex_bail!("expected projection planner"); + }; + for rows in morsel_rows { + let mut output = projection.compute()?; + if matches!(output, PlannerOutput::NeedsIO(_)) { + output = projection.compute()?; + } + let PlannerOutput::Morsel(scope, _) = output else { + vortex_bail!("expected a projection morsel"); + }; + assert_eq!(scope.rows, rows); + } + assert_eq!(projection.state(), State::Done); + } + Ok(()) + } + + #[rstest] + #[case::early_projection(true, true, &[0, 0, 1])] + #[case::deferred_projection(true, false, &[0])] + #[case::unfiltered(false, false, &[0, 1])] + fn announcements_keep_filter_segments_and_unfiltered_projection( + #[case] filtered: bool, + #[case] early_projection: bool, + #[case] expected: &[usize], + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, NonNullable); + let chunks: Vec<_> = (0..2) + .map(|segment| { + FilterPlan::new( + SegmentScanPlan::new( + dtype.clone(), + 1000, + SegmentId::from(segment), + ReadContext::new([]), + None, + ) + .into_plan(), + ) + .into_plan() + }) + .collect(); + let filter = filtered.then(|| FilterPlans::single(chunks[0].clone())); + let locations: Arc<[_]> = (0..2) + .map(|offset| SegmentLocation { + offset, + length: 1, + alignment: Alignment::none(), + }) + .collect(); + let selected = SelectedRows { + scope: WorkScope { + file_ordinal: 0, + rows: 0..2000, + }, + mask: Mask::new_true(2000), + }; + let planner = AnnouncePlanner::for_split( + ScanPlans { + session: new_session(), + locations: Arc::clone(&locations), + projection: ConcatPlan::try_new(dtype, chunks)?.into_plan(), + projection_starts: Arc::from([0, 1000]), + row_offset: 0, + values: Default::default(), + }, + None, + filter, + selected, + ); + let (batch, _) = planner.announce(planner.selected.as_ref().unwrap(), early_projection)?; + let actual: Vec<_> = batch.iter().map(|request| request.target).collect(); + let expected: Vec<_> = expected + .iter() + .map(|&index| locations[index].target()) + .collect(); + assert_eq!(actual, expected); + assert!( + batch + .iter() + .all(|request| request.intent == IoIntent::Announce) + ); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/planning/filter.rs b/vortex-layout/src/scan/planning/filter.rs new file mode 100644 index 00000000000..cf9e01c5d93 --- /dev/null +++ b/vortex-layout/src/scan/planning/filter.rs @@ -0,0 +1,511 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::type_name; +use std::env; +use std::sync::Arc; +use std::sync::LazyLock; + +use bit_vec::BitVec; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::expr::BoundExpression; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::between::Between; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::fill_null::FillNull; +use vortex_array::scalar_fn::fns::get_item::GetItem; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +use vortex_scan::planning::next::Next; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; + +use crate::plan::physical::PlanRef; +use crate::plan::physical::Take; +use crate::scan::filter::FilterExpr; +use crate::scan::planning::ScanPlans; +use crate::scan::planning::graph::GraphStep; +use crate::scan::planning::graph::Piece; +use crate::scan::planning::graph::ProtocolGraph; +use crate::scan::planning::graph::ScanGraph; +use crate::scan::v2::prefetch::plan_segments; +use crate::scan::v2::prefetch::selected_ranges; +use crate::scan::v2::splits::projection_splits; + +/// Which rows a [`FilterPlanner`] keeps. +#[derive(Clone, Copy)] +enum Keep { + /// Rows whose value is true. + True, + /// Rows whose value is not true. + False, +} + +/// The rows of a split that survived its filter. +pub struct SelectedRows { + /// The split the rows belong to. + pub scope: WorkScope, + /// Which rows of `scope.rows` are selected. + pub mask: Mask, +} + +/// The plans a [`FilterPlanner`] evaluates, one per conjunct of the filter, and how it orders +/// them. +/// +/// Every row must satisfy every conjunct, so the planner may evaluate them in any order, each over +/// the rows the previous ones kept. With a filter expression, the order adapts to each conjunct's +/// selectivity as splits report it, shared across the whole scan. +#[derive(Clone)] +pub struct FilterPlans { + plans: Arc<[PlanRef]>, + order: Option>, + /// Whether the predicate can safely evaluate rows outside the selection. + infallible: Arc<[bool]>, + /// Predicates mapped through dictionary codes are cheaper before filtering those codes. + dictionary: Arc<[bool]>, + /// Numeric disjunctions avoid filtering compressed data separately for each comparison. + numeric_disjunction: Arc<[bool]>, +} + +impl FilterPlans { + /// A single plan, evaluated as is. + pub fn single(plan: PlanRef) -> Self { + Self { + plans: Arc::from([plan]), + order: None, + infallible: Arc::from([false]), + dictionary: Arc::from([false]), + numeric_disjunction: Arc::from([false]), + } + } + + /// One plan per conjunct of `filter`, in the order of [`FilterExpr::conjuncts`], evaluated in + /// the order `filter` prefers. + pub(crate) fn conjuncts(filter: Arc, plans: Vec) -> VortexResult { + debug_assert_eq!(filter.conjuncts().len(), plans.len()); + let infallible = filter.conjuncts().iter().map(is_infallible).collect(); + let numeric_disjunction = filter + .conjuncts() + .iter() + .map(|expression| { + matches!(expression.as_opt::(), Some(Operator::Or)) + && is_numeric_predicate(expression) + }) + .collect(); + let dictionary = plans + .iter() + .map(has_dictionary_predicate) + .collect::>()?; + Ok(Self { + infallible, + dictionary, + numeric_disjunction, + plans: plans.into(), + order: Some(filter), + }) + } + + /// Every plan, for callers that walk them all. + pub fn plans(&self) -> &[PlanRef] { + &self.plans + } + + /// The next plan to evaluate among those still `remaining`. + fn next(&self, remaining: &BitVec) -> Option { + match &self.order { + Some(filter) => filter.next_conjunct(remaining), + None => remaining.iter().position(|pending| pending), + } + } + + /// Records that plan `index` kept `output` of `input` rows. + fn report(&self, index: usize, input: usize, output: usize) { + if let Some(filter) = &self.order + && input > 0 + { + filter.report_selectivity(index, output as f64 / input as f64); + } + } +} + +fn has_dictionary_predicate(plan: &PlanRef) -> VortexResult { + if plan.is::() && plan.dtype().is_boolean() { + return Ok(true); + } + for child in plan.children().iter_refs() { + if has_dictionary_predicate(child?)? { + return Ok(true); + } + } + Ok(false) +} + +fn is_infallible(expression: &BoundExpression) -> bool { + expression + .as_scalar() + .is_none_or(|scalar| scalar.signature().is_infallible()) + && expression.children().iter().all(is_infallible) +} + +fn is_numeric_predicate(expression: &BoundExpression) -> bool { + if expression.is_root() || expression.as_opt::().is_some() { + return expression.dtype().is_primitive() || expression.dtype().is_decimal(); + } + if expression.as_opt::().is_some() { + return true; + } + let simple = expression.as_opt::().is_some_and(|operator| { + operator.is_comparison() || matches!(operator, Operator::And | Operator::Or) + }) || expression.as_opt::().is_some() + || expression.as_opt::().is_some(); + simple && expression.children().iter().all(is_numeric_predicate) +} + +/// Evaluates a split's filter to a selection, then hands the selection to `next`. +/// +/// The filter is split into conjuncts, and each runs as its own plan over the rows the previous +/// ones kept, most selective first as measured so far. Infallible predicates over dense selections +/// or dictionary codes run before applying the selection; other predicates run only over selected +/// rows. The planner stops as soon as none remain, without creating a child. +/// +/// The first conjunct fetches its segments, and with that request the planner prefetches the +/// segments of every other conjunct over the same rows, so they arrive while the first is being +/// evaluated. A fetch is served ahead of a prefetch. The projection's segments are not asked for +/// here: the projection planner does that once the filter has kept rows. +/// +/// The same stage prunes: built with [`FilterPlanner::pruning`], it runs a pruning plan whose +/// value is true where zone statistics prove the filter false, and keeps the other rows. +pub struct FilterPlanner { + plans: ScanPlans, + filters: FilterPlans, + keep: Keep, + scope: WorkScope, + mask: Mask, + next: Next, + /// Plans not yet evaluated. + remaining: BitVec, + /// The plan being evaluated, the number of rows selected when it started, and its graph. + running: Option<(usize, usize, ProtocolGraph)>, + pieces: Vec, + /// The protocol id the next plan's first request gets; ids never repeat within the planner. + next_io_id: u32, + /// Whether the conjuncts after the first have been prefetched. + prefetched: bool, + done: bool, + /// The current predicate returns every row, to be intersected with the input mask. + evaluate_all: bool, + /// Rows excluded before this filter must remain excluded from its segment reads. + initial_mask: Mask, + selection_revision: usize, +} + +impl FilterPlanner { + /// Creates a planner that filters the rows of `scope` selected by `mask`. + pub fn new( + plans: ScanPlans, + filters: FilterPlans, + scope: WorkScope, + mask: Mask, + next: Next, + ) -> Self { + Self::with_keep(plans, filters, Keep::True, scope, mask, next) + } + + /// Creates a planner that prunes the rows of `scope` selected by `mask`, dropping the rows for + /// which `pruning` is true. + pub fn pruning( + plans: ScanPlans, + pruning: PlanRef, + scope: WorkScope, + mask: Mask, + next: Next, + ) -> Self { + Self::with_keep( + plans, + FilterPlans::single(pruning), + Keep::False, + scope, + mask, + next, + ) + } + + fn with_keep( + plans: ScanPlans, + filters: FilterPlans, + keep: Keep, + scope: WorkScope, + mask: Mask, + next: Next, + ) -> Self { + let remaining = BitVec::from_elem(filters.plans.len(), true); + let initial_mask = mask.clone(); + Self { + plans, + filters, + keep, + scope, + mask, + next, + remaining, + running: None, + pieces: Vec::new(), + next_io_id: 0, + prefetched: false, + done: false, + evaluate_all: false, + initial_mask, + selection_revision: 0, + } + } + + /// Prefetches remaining filter segments and, within `projection_budget`, projection segments + /// over the selected rows. Their ids count down from the top, clear of the ids the plans' + /// graphs count up from. + fn prefetch_others(&self, running: usize, projection_budget: usize) -> VortexResult { + let mut ids = Vec::new(); + for range in selected_ranges(&self.scope.rows, &self.mask) { + for (index, plan) in self.filters.plans.iter().enumerate() { + if index != running && self.remaining[index] { + plan_segments(plan, range.clone(), &mut ids)?; + } + } + } + ids.sort_unstable(); + ids.dedup(); + // Projection reads can overlap predicate evaluation, but would defeat zone pruning. + if projection_budget > 0 && matches!(self.keep, Keep::True) { + let mut projection = Vec::new(); + for rows in projection_splits(&self.plans.projection_starts, self.scope.rows.clone()) { + if !self.mask.all_true() { + let start = usize::try_from(rows.start - self.scope.rows.start)?; + let end = usize::try_from(rows.end - self.scope.rows.start)?; + if self.mask.slice(start..end).all_false() { + continue; + } + } + plan_segments(&self.plans.projection, rows, &mut projection)?; + } + projection.sort_unstable(); + projection.dedup(); + let mut budget = projection_budget; + let mut selected = Vec::new(); + for id in projection { + if ids.binary_search(&id).is_ok() { + continue; + } + let location = self + .plans + .locations + .get(*id as usize) + .ok_or_else(|| vortex_err!("segment {id} has no known location"))?; + let length = location.length as usize; + if length <= budget { + budget -= length; + selected.push(id); + } + } + ids.extend(selected); + } + ids.into_iter() + .enumerate() + .map(|(index, id)| { + let location = self + .plans + .locations + .get(*id as usize) + .ok_or_else(|| vortex_err!("segment {id} has no known location"))?; + Ok(IoRequest { + intent: IoIntent::Prefetch, + request: IoRequestId(u32::MAX - u32::try_from(index)?), + target: location.target(), + }) + }) + .collect() + } + + /// Starts the next plan over the rows still selected, or finishes: without a child when no + /// row is left, or by handing the selection to `next` when every plan has run. + fn start_next(&mut self) -> VortexResult { + let next = if self.mask.all_false() { + None + } else { + self.filters.next(&self.remaining) + }; + let Some(index) = next else { + self.done = true; + if self.mask.all_false() { + return Ok(PlannerOutput::Done); + } + let scope = self.scope.clone(); + let child = (self.next)(SelectedRows { + scope: scope.clone(), + mask: self.mask.clone(), + })?; + return Ok(PlannerOutput::Planner(scope, child)); + }; + // Filtering compressed dictionary codes rebuilds their encoding before a cheap boolean + // lookup. Evaluate that lookup first, as V1 does. Dense masks likewise cost less to AND + // with a full predicate result than to compact and then scatter back by rank. + self.evaluate_all = self.filters.infallible[index] + && (self.mask.density() >= 0.2 + || self.filters.dictionary[index] + || self.filters.numeric_disjunction[index]); + let plan = self.filters.plans[index].clone(); + // Pipelines apply the mask where they read, so a predicate over a partial selection only + // produces the selected rows. + if !self.initial_mask.all_true() { + self.evaluate_all = false; + } + let evaluation_mask = if self.evaluate_all { + self.initial_mask.clone() + } else { + self.mask.clone() + }; + let graph = ScanGraph::try_new( + self.plans.session.clone(), + &plan, + self.scope.rows.clone(), + evaluation_mask, + self.plans.row_offset, + self.plans.values.clone(), + )?; + self.running = Some(( + index, + self.mask.true_count(), + ProtocolGraph::new(graph, Arc::clone(&self.plans.locations), self.next_io_id), + )); + Ok(PlannerOutput::Continue) + } + + /// Narrows the mask to the rows plan `index` keeps. + fn narrow(&mut self, index: usize) -> VortexResult<()> { + self.pieces.sort_by_key(|piece| piece.rows.start); + let mut values: Vec = std::mem::take(&mut self.pieces) + .into_iter() + .map(|piece| piece.array) + .collect(); + // A lone piece executes through its own kernels; wrapped in a chunked array it would go + // through the generic builder instead. + let values = if values.len() == 1 { + values.remove(0) + } else { + let dtype = self.filters.plans[index].dtype().clone(); + ChunkedArray::try_new(values, dtype)?.into_array() + }; + let mut ctx = self.plans.session.create_execution_ctx(); + let values: Mask = values.fill_null(false)?.execute(&mut ctx)?; + let keep = match self.keep { + Keep::True => values, + Keep::False => !values, + }; + self.mask = if self.evaluate_all { + &self.mask & &keep + } else { + self.mask.intersect_by_rank(&keep) + }; + Ok(()) + } +} + +impl IoConsumer for FilterPlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some((_, _, graph)) = self.running.as_mut() { + graph.set_io_result(request, result); + } + } +} + +impl Planner for FilterPlanner { + fn trace_name(&self) -> &'static str { + match self.keep { + Keep::True => type_name::(), + Keep::False => "vortex_layout::scan::planning::filter::PruningPlanner", + } + } + + fn trace_selection(&self) -> Option<(usize, usize)> { + Some((self.selection_revision, self.mask.true_count())) + } + + fn state(&self) -> State { + match &self.running { + _ if self.done => State::Done, + None => State::NeedsCompute, + Some((_, _, graph)) => match graph.state() { + State::Done => State::NeedsCompute, + state => state, + }, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + vortex_bail!("FilterPlanner: compute called after Done"); + } + let Some((index, input, graph)) = self.running.as_mut() else { + return self.start_next(); + }; + if graph.state() == State::Done { + let (index, input) = (*index, *input); + self.next_io_id = graph.next_id(); + self.running = None; + self.narrow(index)?; + self.remaining.set(index, false); + self.selection_revision += 1; + self.filters.report(index, input, self.mask.true_count()); + return Ok(PlannerOutput::Continue); + } + let running = *index; + Ok(match graph.compute()? { + GraphStep::Yield => PlannerOutput::Continue, + GraphStep::NeedsIO(mut batch) => { + // The running conjunct's fetches come first, then the other conjuncts' prefetches. + if !self.prefetched { + static PROJECTION_BUDGET: LazyLock = LazyLock::new(|| { + env::var("VORTEX_SCAN_PROJECTION_PREFETCH_BYTES") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(0) + }); + self.prefetched = true; + batch.extend(self.prefetch_others(running, *PROJECTION_BUDGET)?); + } + PlannerOutput::NeedsIO(batch) + } + GraphStep::Piece(mut piece) => { + // A care hint lets concat skip whole chunks. Dense predicates still need a + // result for their row positions so narrowing can use a bitmap intersection. + if self.evaluate_all && piece.array.is_empty() { + piece.array = ConstantArray::new( + Scalar::bool(false, self.filters.plans[running].dtype().nullability()), + usize::try_from(piece.rows.end - piece.rows.start)?, + ) + .into_array(); + } + self.pieces.push(piece); + PlannerOutput::Continue + } + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/scan/planning/filter/tests.rs b/vortex-layout/src/scan/planning/filter/tests.rs new file mode 100644 index 00000000000..24048b4e023 --- /dev/null +++ b/vortex-layout/src/scan/planning/filter/tests.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use rstest::rstest; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::expr::checked_add; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::or; +use vortex_array::expr::root; +use vortex_buffer::Alignment; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; +use vortex_scan::planning::next::Next; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; +use vortex_session::registry::ReadContext; + +use super::FilterPlanner; +use super::FilterPlans; +use super::Keep; +use super::SelectedRows; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::RowIdxPlan; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::TakePlan; +use crate::scan::filter::FilterExpr; +use crate::scan::planning::ScanPlans; +use crate::scan::planning::SegmentLocation; +use crate::segments::SegmentId; +use crate::test::new_session; + +#[rstest] +#[case::disabled(Keep::True, 0, false, &[])] +#[case::one_segment(Keep::True, 4, false, &[0])] +#[case::skips_large_segment(Keep::True, 11, false, &[0, 2])] +#[case::two_segments(Keep::True, 16, false, &[0, 1])] +#[case::all_segments(Keep::True, 23, false, &[0, 1, 2])] +#[case::fragmented_excluded_chunk(Keep::True, 23, true, &[0, 2])] +#[case::zone_pruning(Keep::False, 23, false, &[])] +fn projection_prefetch_budget( + #[case] keep: Keep, + #[case] budget: usize, + #[case] exclude_middle: bool, + #[case] expected: &[usize], +) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, NonNullable); + let chunks = (0..3) + .map(|segment| { + SegmentScanPlan::new( + dtype.clone(), + 1000, + SegmentId::from(segment), + ReadContext::new([]), + None, + ) + .into_plan() + }) + .collect(); + let locations: Arc<[_]> = [4, 12, 7, 1] + .into_iter() + .enumerate() + .map(|(offset, length)| SegmentLocation { + offset: offset as u64 * 100, + length, + alignment: Alignment::none(), + }) + .collect(); + let filter = SegmentScanPlan::new( + DType::Bool(NonNullable), + 3000, + SegmentId::from(3), + ReadContext::new([]), + None, + ) + .into_plan(); + let next: Next = Arc::new(|_| vortex_bail!("test stops before handoff")); + let planner = FilterPlanner::with_keep( + ScanPlans { + session: new_session(), + locations: Arc::clone(&locations), + projection: ConcatPlan::try_new(dtype, chunks)?.into_plan(), + projection_starts: Arc::from([0, 1000, 2000]), + row_offset: 0, + values: Default::default(), + }, + FilterPlans::single(filter), + keep, + WorkScope { + file_ordinal: 0, + rows: 0..3000, + }, + // More than 64 selected runs must still exclude the empty middle chunk. + Mask::from_iter( + (0..3000).map(|row| !exclude_middle || (row % 2 == 0 && !(1000..2000).contains(&row))), + ), + next, + ); + let actual: Vec<_> = planner + .prefetch_others(0, budget)? + .iter() + .map(|request| request.target) + .collect(); + let expected: Vec<_> = expected.iter().map(|&id| locations[id].target()).collect(); + assert_eq!(actual, expected); + Ok(()) +} + +#[rstest] +#[case::dense(&[0, 2, 4, 7], false, false, false, true, &[4, 7])] +#[case::sparse(&[7], false, false, false, false, &[7])] +#[case::sparse_dictionary(&[7], true, false, false, true, &[7])] +#[case::sparse_disjunction(&[7], false, false, true, true, &[7])] +#[case::fallible_dense(&[0, 1, 2, 3], false, true, false, false, &[3])] +fn predicate_selection( + #[case] selected: &[usize], + #[case] dictionary: bool, + #[case] fallible: bool, + #[case] disjunction: bool, + #[case] evaluate_all: bool, + #[case] expected: &[usize], + #[values(false, true)] initial_selection: bool, +) -> VortexResult<()> { + let source = RowIdxPlan::new(8).into_plan(); + let expression = if fallible { + // Rows 4..8 overflow, so evaluating outside the selection would fail. + gt(checked_add(root(), lit(u64::MAX - 3)), lit(u64::MAX - 1)) + } else if disjunction { + or(gt(root(), lit(3_u64)), lt(root(), lit(1_u64))) + } else { + gt(root(), lit(3_u64)) + } + .bind(source.dtype())?; + let predicate = EvalPlan::try_new(expression.clone(), source.clone())?.into_plan(); + let predicate = if dictionary { + TakePlan::new(source.clone(), predicate).into_plan() + } else { + predicate + }; + let filters = FilterPlans::conjuncts( + Arc::new(FilterExpr::from_conjuncts(vec![expression])), + vec![predicate], + )?; + let next: Next = Arc::new(|_| vortex_bail!("test stops before handoff")); + let selection = Mask::from_iter((0..8).map(|row| selected.contains(&row))); + let mut planner = FilterPlanner::new( + ScanPlans { + session: new_session(), + locations: Arc::from([]), + projection: source, + projection_starts: Arc::from([]), + row_offset: 0, + values: Default::default(), + }, + filters, + WorkScope { + file_ordinal: 0, + rows: 0..8, + }, + if initial_selection { + selection.clone() + } else { + Mask::new_true(8) + }, + next, + ); + // A preceding conjunct can narrow the selection without excluding segments at scan setup. + planner.mask = selection; + planner.compute()?; + assert_eq!(planner.evaluate_all, evaluate_all && !initial_selection); + for _ in 0..100 { + if planner.remaining.iter().all(|pending| !pending) { + let actual: Vec<_> = (0..8).filter(|&row| planner.mask.value(row)).collect(); + assert_eq!(actual, expected); + assert_eq!(planner.trace_selection(), Some((1, expected.len()))); + return Ok(()); + } + assert_eq!(planner.state(), State::NeedsCompute); + planner.compute()?; + } + vortex_bail!("predicate did not finish") +} diff --git a/vortex-layout/src/scan/planning/graph.rs b/vortex-layout/src/scan/planning/graph.rs new file mode 100644 index 00000000000..efa253dec8e --- /dev/null +++ b/vortex-layout/src/scan/planning/graph.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use rustc_hash::FxHashMap; +use vortex_array::ArrayRef; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_io::request; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +use vortex_scan::planning::planner::State; +use vortex_session::VortexSession; + +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::ReadId; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::Turn; +use crate::plan::physical::pipeline::ValueCache; +use crate::scan::planning::SegmentLocation; +use crate::segments::SegmentId; + +/// Rows a running plan produced, in row order. +#[derive(Clone)] +pub(crate) struct Piece { + /// Orders the piece among the pieces of one graph. + pub(crate) rows: Range, + /// The selected rows the piece holds. + pub(crate) array: ArrayRef, +} + +/// A plan running over one row range on the pipeline [`Scan`]. +/// +/// The scan produces the selected rows in row order and leaves its reads to its owner. +pub(crate) struct ScanGraph { + scan: Scan, + /// A step the scan already took that the owner has not seen. + pending: Option, + done: bool, + /// Orders the pieces handed out. + next_piece: u64, +} + +/// What a graph wants from its owner. +enum GraphState { + NeedsCompute, + Waiting, + Done, +} + +impl ScanGraph { + /// Compiles nothing yet: the scan compiles `plan` over `rows`, restricted to `mask`, on its + /// first step. `row_offset` is the global row index of the plan's first row, and `values` the + /// whole values shared with the split's other scans. + pub(crate) fn try_new( + session: VortexSession, + plan: &PlanRef, + rows: Range, + mask: Mask, + row_offset: u64, + values: ValueCache, + ) -> VortexResult { + let scan = Scan::try_new(session, plan.clone(), vec![Split { rows, mask }])? + .with_row_offset(row_offset) + .with_value_cache(values); + Ok(Self { + scan, + pending: None, + done: false, + next_piece: 0, + }) + } + + fn state(&self) -> GraphState { + match &self.pending { + _ if self.done => GraphState::Done, + Some(Turn::Waiting) => GraphState::Waiting, + _ => GraphState::NeedsCompute, + } + } + + fn step(&mut self) -> VortexResult { + let turn = match self.pending.take() { + Some(Turn::Waiting) | None => self.scan.step()?, + Some(turn) => turn, + }; + match &turn { + Turn::Done => self.done = true, + Turn::Waiting => self.pending = Some(Turn::Waiting), + Turn::Read(_) | Turn::Output(..) => {} + } + Ok(turn) + } + + fn piece(&mut self, array: ArrayRef) -> Piece { + let start = self.next_piece; + self.next_piece += 1; + Piece { + rows: start..start + 1, + array, + } + } + + fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { + self.scan.deliver(read, bytes)?; + if matches!(self.pending, Some(Turn::Waiting)) { + self.pending = None; + } + Ok(()) + } +} + +/// What one [`ProtocolGraph::compute`] produced. +pub(crate) enum GraphStep { + /// Control returns to the owner with CPU work still ready. + Yield, + /// Requests to publish, each returned once. + NeedsIO(IoBatch), + /// Rows the plan produced. + Piece(Piece), +} + +/// A [`ScanGraph`] whose segment reads are published as protocol requests. +/// +/// The scan names segments by id; the protocol names byte ranges. Each scan read becomes a +/// `Fetch` of the segment's range at the segment's alignment, and each delivery is handed back to +/// the scan under its own id. +pub(crate) struct ProtocolGraph { + graph: ScanGraph, + locations: Arc<[SegmentLocation]>, + /// Published and undelivered requests, with the scan read and segment each one answers. + outstanding: FxHashMap, + /// A delivery the scan could not take, reported by the next compute. + failed: Option, + /// The protocol id the next published request gets. + next_id: u32, +} + +impl ProtocolGraph { + /// Wraps `graph`, numbering its requests from `first_id`. + /// + /// Request ids must be unique over an owner's lifetime, so an owner that runs several graphs + /// starts each where the previous one's [`next_id`](Self::next_id) left off. + pub(crate) fn new(graph: ScanGraph, locations: Arc<[SegmentLocation]>, first_id: u32) -> Self { + Self { + graph, + locations, + outstanding: FxHashMap::default(), + failed: None, + next_id: first_id, + } + } + + /// The protocol id the next request would get: the first id a following graph of the same + /// owner may use. + pub(crate) fn next_id(&self) -> u32 { + self.next_id + } + + /// The owner's protocol state while the graph is running. + pub(crate) fn state(&self) -> State { + if self.failed.is_some() { + return State::NeedsCompute; + } + match self.graph.state() { + GraphState::Done => State::Done, + GraphState::NeedsCompute => State::NeedsCompute, + GraphState::Waiting => State::Waiting, + } + } + + pub(crate) fn compute(&mut self) -> VortexResult { + if let Some(error) = self.failed.take() { + return Err(error); + } + Ok(match self.graph.step()? { + Turn::Waiting | Turn::Done => GraphStep::Yield, + Turn::Output(_, array) => GraphStep::Piece(self.graph.piece(array)), + Turn::Read(read) => { + // Publish every read the scan has ready together, so the owner can coalesce them. + let mut batch = vec![self.request(read.id, read.segment_id)?]; + loop { + match self.graph.step()? { + Turn::Read(read) => batch.push(self.request(read.id, read.segment_id)?), + turn => { + if !matches!(turn, Turn::Waiting) { + self.graph.pending = Some(turn); + self.graph.done = false; + } + break; + } + } + } + GraphStep::NeedsIO(batch) + } + }) + } + + fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { + let location = self.location(segment_id)?; + let id = self.next_id; + self.next_id = self + .next_id + .checked_add(1) + .ok_or_else(|| vortex_err!("ProtocolGraph ran out of request ids"))?; + self.outstanding.insert(id, (read, segment_id)); + Ok(request::IoRequest { + intent: IoIntent::Fetch, + request: request::IoRequestId(id), + target: location.target(), + }) + } + + /// Hands a delivered read to the scan. The driver only delivers requests this graph + /// published, so anything else is a driver bug. Bytes the scan cannot take fail the owner's + /// next compute. + pub(crate) fn set_io_result(&mut self, id: request::IoRequestId, result: IoResult) { + let Some((read, segment_id)) = self.outstanding.remove(&id.0) else { + vortex_panic!("ProtocolGraph: delivery of {id:?}, which is not outstanding"); + }; + let IoResult::Bytes(bytes) = result else { + vortex_panic!("ProtocolGraph: segment {segment_id} answered with a size"); + }; + if let Err(error) = self.deliver(read, segment_id, bytes) + && self.failed.is_none() + { + self.failed = Some(error.with_context(format!("delivering segment {segment_id}"))); + } + } + + fn deliver( + &mut self, + read: ReadId, + segment_id: SegmentId, + bytes: BufferHandle, + ) -> VortexResult<()> { + // A source that honours the target's alignment makes this a check, not a copy. + let alignment = self.location(segment_id)?.alignment; + let bytes = BufferHandle::new_host(bytes.try_into_host_sync()?.aligned(alignment)); + self.graph.deliver(read, bytes) + } + + fn location(&self, id: SegmentId) -> VortexResult { + self.locations + .get(*id as usize) + .copied() + .ok_or_else(|| vortex_err!("segment {id} has no known location")) + } +} diff --git a/vortex-layout/src/scan/planning/mod.rs b/vortex-layout/src/scan/planning/mod.rs new file mode 100644 index 00000000000..ca368c10134 --- /dev/null +++ b/vortex-layout/src/scan/planning/mod.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Layout-level planners and morsels for the planning protocol. +//! +//! A split is evaluated over physical plans: an [`AnnouncePlanner`] tells the IO service which +//! segments the split is likely to read, a pruning [`FilterPlanner`] drops rows whose zone +//! statistics prove the filter false, a [`FilterPlanner`] runs the filter plan over the rows that +//! remain, and a [`ProjectionPlanner`] emits a [`ProjectionMorsel`] that runs the projection plan +//! over the rows that survive. Each is a separate stage, and the evaluating ones execute through +//! the plan exec graph; [`plan_split`] composes them for one split. +//! +//! Planners that decide which splits to read, and that know where each segment lives in a file, +//! belong to the crates that own those sources. + +mod announce; +mod filter; +pub(crate) mod graph; +mod projection; + +use std::sync::Arc; + +pub use announce::AnnouncePlanner; +pub use filter::FilterPlanner; +pub use filter::FilterPlans; +pub use filter::SelectedRows; +pub use projection::ProjectionMorsel; +pub use projection::ProjectionPlanner; +use vortex_buffer::Alignment; +use vortex_error::VortexResult; +use vortex_io::request::IoTarget; +use vortex_mask::Mask; +use vortex_scan::planning::next::Next; +use vortex_scan::planning::next::next_fn; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::WorkScope; +use vortex_session::VortexSession; + +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::ValueCache; + +/// Where a segment's bytes live in its source. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SegmentLocation { + /// Byte offset of the segment from the start of the source. + pub offset: u64, + /// Length of the segment in bytes. + pub length: u32, + /// Alignment the segment's bytes must have once delivered. + pub alignment: Alignment, +} + +impl SegmentLocation { + /// The byte range a read of the segment asks for. + pub fn target(&self) -> IoTarget { + IoTarget::Range { + offset: self.offset, + len: self.length as usize, + alignment: self.alignment, + } + } +} + +/// What every split of a scan shares: the projection plan and where its segments live. +#[derive(Clone)] +pub struct ScanPlans { + /// The session used for decoding and expression evaluation. + pub session: VortexSession, + /// The location of every segment, indexed by segment id. + pub locations: Arc<[SegmentLocation]>, + /// The optimized projection plan, over the scan's root row domain. + pub projection: PlanRef, + /// Where the chunks the projection reads start, over the root row domain, sorted. A split's + /// projection is cut at each of these, so every projection split reads at most one chunk of + /// every column. Bounded sparse identity scans may omit these cuts after IO announcements. + pub projection_starts: Arc<[u64]>, + /// The global row index of the root row domain's first row. + pub row_offset: u64, + /// Whole values, such as dictionary values, kept between the scans of one split. Each split + /// starts with its own, dropped when the split finishes. + pub values: ValueCache, +} + +/// The work for one split: announce its likely reads, prune the rows of `scope` selected by +/// `mask` with zone statistics, filter the rows that remain, then project the rows that survive. +/// The pruning and filter stages are skipped when their plan is absent. +pub fn plan_split( + plans: ScanPlans, + pruning: Option, + filter: Option, + scope: WorkScope, + mask: Mask, +) -> VortexResult> { + Ok(Box::new(AnnouncePlanner::for_split( + ScanPlans { + values: ValueCache::default(), + ..plans + }, + pruning, + filter, + SelectedRows { scope, mask }, + ))) +} + +/// Plans a bounded row-index selection with an identity projection and no filter or pruning. +/// One graph per split avoids decoding a column again at another column's chunk boundary. +pub(crate) fn plan_sparse_projection( + plans: ScanPlans, + scope: WorkScope, + mask: Mask, +) -> Box { + Box::new(AnnouncePlanner::for_sparse_projection( + ScanPlans { + values: ValueCache::default(), + ..plans + }, + SelectedRows { scope, mask }, + )) +} + +/// Builds execution stages only after the announcement. Split construction and early IO +/// announcements happen before tasks run; allocating continuations here +/// lets the split's execution task do that work. +fn plan_selected( + plans: ScanPlans, + pruning: Option, + filter: Option, + selected: SelectedRows, +) -> VortexResult> { + if pruning.is_none() && filter.is_none() { + return Ok(Box::new(ProjectionPlanner::new(plans, selected))); + } + let project: Next = { + let plans = plans.clone(); + next_fn(move |selected| Ok(ProjectionPlanner::new(plans.clone(), selected))) + }; + let Some(pruning) = pruning else { + return Ok(match filter { + Some(filter) => Box::new(FilterPlanner::new( + plans, + filter, + selected.scope, + selected.mask, + project, + )), + None => Box::new(ProjectionPlanner::new(plans, selected)), + }); + }; + let next = match filter { + None => project, + Some(filter) => { + let plans = plans.clone(); + next_fn(move |selected: SelectedRows| { + Ok(FilterPlanner::new( + plans.clone(), + filter.clone(), + selected.scope, + selected.mask, + Arc::clone(&project), + )) + }) + } + }; + Ok(Box::new(FilterPlanner::pruning( + plans, + pruning, + selected.scope, + selected.mask, + next, + ))) +} diff --git a/vortex-layout/src/scan/planning/projection.rs b/vortex-layout/src/scan/planning/projection.rs new file mode 100644 index 00000000000..3bfec039ebf --- /dev/null +++ b/vortex-layout/src/scan/planning/projection.rs @@ -0,0 +1,568 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoTarget; +use vortex_scan::planning::morsel::Morsel; +use vortex_scan::planning::morsel::MorselOutput; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::scan::planning::ScanPlans; +use crate::scan::planning::SelectedRows; +use crate::scan::planning::graph::GraphStep; +use crate::scan::planning::graph::Piece; +use crate::scan::planning::graph::ProtocolGraph; +use crate::scan::planning::graph::ScanGraph; +use crate::scan::v2::prefetch::plan_segments; +use crate::scan::v2::splits::projection_splits; + +/// Turns a filter split's selected rows into the morsels that project them. +/// +/// The rows are cut into projection splits where the projection's chunks start, and each split +/// with a selected row becomes one [`ProjectionMorsel`], so a morsel reads at most one chunk of +/// every column. The planner finishes once it has handed out every morsel. +/// +/// When there are multiple morsels, the planner first prefetches every segment the projection +/// reads over the selected rows, so their reads start together. Nothing asks for projection +/// segments earlier by default. Opt-in projection read-ahead can start a bounded set during +/// filtering instead. +/// Early announcements can share neighboring reads; wide projections withdraw interests +/// belonging only to empty cuts once the final mask is known. +pub struct ProjectionPlanner { + plans: ScanPlans, + /// The filter split's selected rows, until the first compute cuts them. + selected: Option, + /// The projection splits not yet handed out, last first. + pending: Vec, + pruned_cuts: bool, +} + +impl ProjectionPlanner { + /// Creates a planner for `selected`. + pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + plans, + selected: Some(selected), + pending: Vec::new(), + pruned_cuts: false, + } + } + + /// Prefetches the nonempty projection chunks, which are already cut for execution. + fn prefetch(&self) -> VortexResult { + if self.pending.len() == 1 { + // One morsel publishes its own reads before waiting; there are no later morsels + // whose reads need to start ahead of it. + return Ok(Vec::new()); + } + let mut ids = Vec::new(); + for selected in &self.pending { + plan_segments( + &self.plans.projection, + selected.scope.rows.clone(), + &mut ids, + )?; + } + ids.sort_unstable(); + ids.dedup(); + ids.into_iter() + .enumerate() + .map(|(index, id)| { + let location = self + .plans + .locations + .get(*id as usize) + .ok_or_else(|| vortex_err!("segment {id} has no known location"))?; + Ok(IoRequest { + intent: IoIntent::Prefetch, + request: IoRequestId(u32::try_from(index)?), + target: location.target(), + }) + }) + .collect() + } + + /// The projection splits of `selected` that have a selected row, last first. + fn cut(&mut self, selected: &SelectedRows) -> Vec { + let SelectedRows { scope, mask } = selected; + let start = scope.rows.start; + let index = |row: u64| usize::try_from(row - start).vortex_expect("split row fits usize"); + let mut pruned_cuts = false; + let mut splits = projection_splits(&self.plans.projection_starts, scope.rows.clone()) + .into_iter() + .filter_map(|rows| { + let mask = mask.slice(index(rows.start)..index(rows.end)); + if mask.all_false() { + pruned_cuts = true; + return None; + } + Some(SelectedRows { + scope: WorkScope { + file_ordinal: scope.file_ordinal, + rows, + }, + mask, + }) + }) + .collect::>(); + splits.reverse(); + self.pruned_cuts = pruned_cuts; + splits + } + + fn forget_pruned( + &self, + selected: &SelectedRows, + mut prefetch: IoBatch, + ) -> VortexResult { + if !self.pruned_cuts { + return Ok(prefetch); + } + let mut kept: HashSet<_> = prefetch.iter().map(|request| request.target).collect(); + if let [selected] = self.pending.as_slice() { + // A single morsel submits its own reads, so no prefetch batch protects its aliases. + let mut ids = Vec::new(); + plan_segments( + &self.plans.projection, + selected.scope.rows.clone(), + &mut ids, + )?; + for id in ids { + kept.insert( + self.plans + .locations + .get(*id as usize) + .ok_or_else(|| vortex_err!("segment {id} has no known location"))? + .target(), + ); + } + } + let mut seen: HashSet<_> = HashSet::default(); + let mut ids = Vec::new(); + plan_segments( + &self.plans.projection, + selected.scope.rows.clone(), + &mut ids, + )?; + let mut forget = Vec::new(); + for id in ids { + let target = self + .plans + .locations + .get(*id as usize) + .ok_or_else(|| vortex_err!("segment {id} has no known location"))? + .target(); + // Different segment IDs may name the same range. Retain every selected alias. + if kept.contains(&target) || !seen.insert(target) { + continue; + } + forget.push(IoRequest { + intent: IoIntent::Forget, + request: IoRequestId(u32::try_from(prefetch.len() + forget.len())?), + target, + }); + } + // Withdraw stale neighbors before the selected projection reads become eligible. + forget.append(&mut prefetch); + Ok(forget) + } +} + +impl IoConsumer for ProjectionPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for ProjectionPlanner { + fn state(&self) -> State { + if self.selected.is_none() && self.pending.is_empty() { + State::Done + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + static FORGET_PRUNED: LazyLock> = LazyLock::new(|| { + std::env::var("VORTEX_SCAN_IO_FORGET_PRUNED_PROJECTION") + .ok() + .map(|value| value == "1") + }); + if let Some(selected) = self.selected.take() { + tracing::debug!( + target: "vortex_layout::filter_stages", + row_start = selected.scope.rows.start, + row_end = selected.scope.rows.end, + selected_rows = selected.mask.true_count(), + "projection mask ready" + ); + self.pending = self.cut(&selected); + // No projection split has a selected row. + if self.pending.is_empty() { + return Ok(PlannerOutput::Done); + } + let wide = self + .plans + .projection + .dtype() + .as_struct_fields_opt() + .is_some_and(|fields| fields.nfields() >= 8); + let forget = (*FORGET_PRUNED).unwrap_or(wide); + let prefetch = self.prefetch()?; + let prefetch = if forget { + self.forget_pruned(&selected, prefetch)? + } else { + prefetch + }; + if !prefetch.is_empty() { + return Ok(PlannerOutput::NeedsIO(prefetch)); + } + } + let Some(selected) = self.pending.pop() else { + vortex_bail!("ProjectionPlanner: compute called after Done"); + }; + let scope = selected.scope.clone(); + Ok(PlannerOutput::Morsel( + scope, + Box::new(ProjectionMorsel::new(self.plans.clone(), selected)), + )) + } +} + +/// Runs the projection plan over a split's selected rows and returns them as one batch, in row +/// order. A selection that projects to no rows finishes without a batch. +pub struct ProjectionMorsel { + plans: ScanPlans, + selected: SelectedRows, + graph: Option, + pieces: Vec, + done: bool, +} + +impl ProjectionMorsel { + /// Creates a morsel projecting `selected`. + pub fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + plans, + selected, + graph: None, + pieces: Vec::new(), + done: false, + } + } + + fn finish(&mut self) -> VortexResult { + self.done = true; + self.pieces.sort_by_key(|piece| piece.rows.start); + let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) + .into_iter() + .map(|piece| piece.array) + .filter(|array| !array.is_empty()) + .collect(); + let array = match arrays.len() { + 0 => return Ok(MorselOutput::Done), + 1 => arrays.remove(0), + _ => ChunkedArray::try_new(arrays, self.plans.projection.dtype().clone())?.into_array(), + }; + Ok(MorselOutput::Batch(array)) + } +} + +impl IoConsumer for ProjectionMorsel { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some(graph) = self.graph.as_mut() { + graph.set_io_result(request, result); + } + } +} + +impl Morsel for ProjectionMorsel { + fn state(&self) -> State { + match &self.graph { + _ if self.done => State::Done, + None => State::NeedsCompute, + Some(graph) => match graph.state() { + State::Done => State::NeedsCompute, + state => state, + }, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + vortex_bail!("ProjectionMorsel: compute called after Done"); + } + let Some(graph) = self.graph.as_mut() else { + let graph = ScanGraph::try_new( + self.plans.session.clone(), + &self.plans.projection, + self.selected.scope.rows.clone(), + self.selected.mask.clone(), + self.plans.row_offset, + self.plans.values.clone(), + )?; + self.graph = Some(ProtocolGraph::new( + graph, + Arc::clone(&self.plans.locations), + 0, + )); + return Ok(MorselOutput::Continue); + }; + if graph.state() == State::Done { + return self.finish(); + } + Ok(match graph.compute()? { + GraphStep::Yield => MorselOutput::Continue, + GraphStep::NeedsIO(batch) => { + if tracing::enabled!(target: "vortex_layout::filter_stages", tracing::Level::DEBUG) + { + for request in &batch { + if request.intent == IoIntent::Fetch + && let IoTarget::Range { offset, len, .. } = request.target + { + tracing::debug!( + target: "vortex_layout::filter_stages", + row_start = self.selected.scope.rows.start, + row_end = self.selected.scope.rows.end, + request = request.request.0, + offset, + length = len, + "projection input requested" + ); + } + } + } + MorselOutput::NeedsIO(batch) + } + GraphStep::Piece(piece) => { + self.pieces.push(piece); + MorselOutput::Continue + } + }) + } +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_buffer::Alignment; + use vortex_mask::Mask; + use vortex_session::registry::ReadContext; + + use super::*; + use crate::plan::physical::ConcatPlan; + use crate::plan::physical::FilterPlan; + use crate::plan::physical::PackPlan; + use crate::plan::physical::SegmentScanPlan; + use crate::scan::planning::SegmentLocation; + use crate::segments::SegmentId; + use crate::test::new_session; + + #[rstest::rstest] + fn single_morsel_withdrawal_preserves_selected_aliases( + #[values(false, true)] aliased: bool, + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, NonNullable); + let chunks = (0..2) + .map(|segment| { + FilterPlan::new( + SegmentScanPlan::new( + dtype.clone(), + 1000, + SegmentId::from(segment), + ReadContext::new([]), + None, + ) + .into_plan(), + ) + .into_plan() + }) + .collect(); + let locations: Arc<[_]> = (0..2) + .map(|offset| SegmentLocation { + offset: if aliased { 0 } else { offset }, + length: 1, + alignment: Alignment::none(), + }) + .collect(); + let selected = SelectedRows { + scope: WorkScope { + file_ordinal: 0, + rows: 0..2000, + }, + mask: Mask::from_iter((0..2000).map(|row| row < 1000)), + }; + let mut planner = ProjectionPlanner::new( + ScanPlans { + session: new_session(), + locations: Arc::clone(&locations), + projection: ConcatPlan::try_new(dtype, chunks)?.into_plan(), + projection_starts: Arc::from([0, 1000]), + row_offset: 0, + values: Default::default(), + }, + SelectedRows { + scope: selected.scope.clone(), + mask: selected.mask.clone(), + }, + ); + planner.pending = planner.cut(&selected); + assert_eq!(planner.pending.len(), 1); + let prefetch = planner.prefetch()?; + assert!(prefetch.is_empty()); + let withdrawal = planner.forget_pruned(&selected, prefetch)?; + let targets: Vec<_> = withdrawal.iter().map(|request| request.target).collect(); + assert_eq!( + targets, + if aliased { + vec![] + } else { + vec![locations[1].target()] + } + ); + Ok(()) + } + + #[rstest::rstest] + fn prefetch_skips_empty_chunks_without_enumerating_row_runs( + #[values(false, true)] aliased: bool, + #[values(0, 2, 8)] fields: usize, + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, NonNullable); + let chunks = (0..3) + .map(|segment| { + FilterPlan::new( + SegmentScanPlan::new( + dtype.clone(), + 1000, + SegmentId::from(segment), + ReadContext::new([]), + None, + ) + .into_plan(), + ) + .into_plan() + }) + .collect(); + let locations: Arc<[_]> = (0..3) + .map(|offset| SegmentLocation { + offset: if aliased && offset == 1 { 0 } else { offset }, + length: 1, + alignment: Alignment::none(), + }) + .collect(); + // More than 64 runs used to fall back to prefetching the entire bounding range, + // including the empty middle chunk, after enumerating every run. + let mask = + Mask::from_iter((0..3000).map(|row| row % 2 == 0 && !(1000..2000).contains(&row))); + let projection = ConcatPlan::try_new(dtype.clone(), chunks)?.into_plan(); + let projection = if fields == 0 { + projection + } else { + PackPlan::try_new( + StructFields::from_iter( + (0..fields).map(|index| (format!("f{index}"), dtype.clone())), + ), + NonNullable, + 3000, + vec![projection; fields], + None, + )? + .into_plan() + }; + let mut planner = ProjectionPlanner::new( + ScanPlans { + session: new_session(), + locations: Arc::clone(&locations), + projection, + projection_starts: Arc::from([0, 1000, 2000]), + row_offset: 0, + values: Default::default(), + }, + SelectedRows { + scope: WorkScope { + file_ordinal: 0, + rows: 0..3000, + }, + mask: mask.clone(), + }, + ); + let PlannerOutput::NeedsIO(requests) = planner.compute()? else { + vortex_bail!("expected projection prefetch"); + }; + let forgotten = std::env::var("VORTEX_SCAN_IO_FORGET_PRUNED_PROJECTION") + .ok() + .map(|value| value == "1") + .unwrap_or(fields >= 8); + assert_eq!(planner.pending.len(), 2); + assert_eq!( + requests + .iter() + .filter(|r| r.intent == IoIntent::Forget) + .count(), + usize::from(forgotten && !aliased) + ); + let prefetch: IoBatch = requests + .into_iter() + .filter(|request| request.intent == IoIntent::Prefetch) + .collect(); + let targets: Vec<_> = prefetch.iter().map(|request| request.target).collect(); + assert_eq!(targets, vec![locations[0].target(), locations[2].target()]); + let pruned = planner.forget_pruned( + &SelectedRows { + scope: WorkScope { + file_ordinal: 0, + rows: 0..3000, + }, + mask: mask.clone(), + }, + prefetch, + )?; + let forgotten: Vec<_> = pruned + .iter() + .filter(|request| request.intent == IoIntent::Forget) + .map(|request| request.target) + .collect(); + assert_eq!( + forgotten, + if aliased { + vec![] + } else { + vec![locations[1].target()] + } + ); + assert_eq!( + pruned.first().map(|request| request.intent), + Some(if aliased { + IoIntent::Prefetch + } else { + IoIntent::Forget + }) + ); + let Mask::Values(values) = mask else { + vortex_bail!("expected a partial mask"); + }; + assert!(values.cached_slices().is_none()); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 814c0759453..2c039a6fbf4 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -266,6 +266,11 @@ impl ScanBuilder { self } + /// Whether the scan has both a filter and a row limit. + pub(crate) fn has_filter_and_limit(&self) -> bool { + self.filter.is_some() && self.limit.is_some() + } + /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { Ok(self.projection.dtype().clone()) @@ -300,6 +305,27 @@ impl ScanBuilder { } } + /// Splits the builder into its configuration, for [`v2::ScanBuilder`](crate::scan::v2::ScanBuilder) + /// to copy. + pub(crate) fn into_parts(self) -> ScanParts { + ScanParts { + session: self.session, + layout_reader: self.layout_reader, + projection: self.projection, + filter: self.filter, + ordered: self.ordered, + row_range: self.row_range, + selection: self.selection, + split_by: self.split_by, + natural_splits: self.natural_splits, + concurrency: self.concurrency, + map_fn: self.map_fn, + metrics_registry: self.metrics_registry, + limit: self.limit, + row_offset: self.row_offset, + } + } + /// Rewrite range comparisons in the filter into betweens, compute split ranges, and return an /// executable repeated scan. pub fn prepare(self) -> VortexResult> { @@ -391,6 +417,24 @@ impl ScanBuilder { } } +/// Everything a [`ScanBuilder`] holds. +pub(crate) struct ScanParts { + pub(crate) session: VortexSession, + pub(crate) layout_reader: LayoutReaderRef, + pub(crate) projection: BoundExpression, + pub(crate) filter: Option, + pub(crate) ordered: bool, + pub(crate) row_range: Option>, + pub(crate) selection: Selection, + pub(crate) split_by: SplitBy, + pub(crate) natural_splits: Option>, + pub(crate) concurrency: usize, + pub(crate) map_fn: Arc VortexResult + Send + Sync>, + pub(crate) metrics_registry: Option>, + pub(crate) limit: Option, + pub(crate) row_offset: u64, +} + enum LazyScanState { Builder(Option>>), Preparing(Task>>>), diff --git a/vortex-layout/src/scan/v2/conjuncts.rs b/vortex-layout/src/scan/v2/conjuncts.rs new file mode 100644 index 00000000000..f7fa6246ee9 --- /dev/null +++ b/vortex-layout/src/scan/v2/conjuncts.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use vortex_array::dtype::FieldPath; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::analysis::referenced_field_paths; +use vortex_array::expr::bound::and as bound_and; +use vortex_error::VortexResult; + +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Filter; +use crate::plan::physical::FilterPlan; +use crate::plan::physical::PlanRef; + +/// Merges conjuncts that read a common field back into one conjunct each. +/// +/// Evaluating two conjuncts separately reads and decodes their shared columns twice, so the +/// conjuncts are grouped by the fields they share, transitively: in `a > 5 AND a < b AND b = 2` +/// all three read a field another one reads, so they become one conjunct. Groups keep the order of +/// their first member, and members keep their order within a group. +pub(super) fn group_conjuncts(conjuncts: &[BoundExpression]) -> VortexResult> { + if conjuncts.len() < 2 { + return Ok(conjuncts.to_vec()); + } + let fields = conjuncts + .iter() + .map(|conjunct| Ok(referenced_field_paths(conjunct)?.into_iter().collect_vec())) + .collect::>>>()?; + + let mut sets = DisjointSets::new(conjuncts.len()); + for (index, paths) in fields.iter().enumerate() { + for (other, other_paths) in fields[..index].iter().enumerate() { + let shared = paths + .iter() + .any(|path| other_paths.iter().any(|other| path.overlap(other))); + if shared { + sets.union(index, other); + } + } + } + + // Every set is rooted at its lowest member, so filling groups in index order keeps both the + // groups and their members in input order. + let mut groups = vec![Vec::new(); conjuncts.len()]; + for (index, conjunct) in conjuncts.iter().enumerate() { + groups[sets.find(index)].push(conjunct.clone()); + } + Ok(groups + .into_iter() + .filter_map(|group| group.into_iter().reduce(bound_and)) + .collect()) +} + +/// Moves each filter above the predicate evaluated over it, so the predicate runs over every row +/// of a segment and its result is filtered, rather than filtering the segment first. +/// +/// The default scan evaluates filter predicates the same way: filtering decoded data first would +/// decompress it, while a predicate over the whole encoded segment can run on its encoding and +/// hand the filter a lazy result. Rows outside the selection hold real data, so the predicate sees +/// only values the column contains. Only used for filter conjuncts; projections still filter +/// first, since a projection may be invalid on rows the filter removed. +pub(super) fn filter_after_eval(plan: PlanRef) -> VortexResult { + let plan = map_children(plan, filter_after_eval)?; + let Some(eval) = plan.as_opt::() else { + return Ok(plan); + }; + let child = eval.child_plan()?; + let Some(filter) = child.as_opt::() else { + return Ok(plan); + }; + let evaluated = EvalPlan::try_new(eval.expression().clone(), filter.child_plan()?)?; + Ok(FilterPlan::new(evaluated.into_plan()).into_plan()) +} + +/// Rebuilds `plan` over its children mapped by `f`, keeping `plan` itself when no child changed, +/// so plans shared between the filter and the projection stay shared. +pub(super) fn map_children( + plan: PlanRef, + mut f: impl FnMut(PlanRef) -> VortexResult, +) -> VortexResult { + let mut children: Option> = None; + for (index, child) in plan.children().iter_refs().enumerate() { + let child = child?; + let mapped = f(child.clone())?; + if let Some(children) = &mut children { + children.push(mapped); + } else if !PlanRef::ptr_eq(child, &mapped) { + let mut mapped_children = Vec::with_capacity(plan.child_count()); + for previous in plan.children().iter_refs().take(index) { + mapped_children.push(previous?.clone()); + } + mapped_children.push(mapped); + children = Some(mapped_children); + } + } + match children { + Some(children) => plan.with_children(children), + None => Ok(plan), + } +} + +/// Union-find over conjunct indices, where each set is rooted at its lowest member. +struct DisjointSets(Vec); + +impl DisjointSets { + fn new(len: usize) -> Self { + Self((0..len).collect()) + } + + fn find(&mut self, mut index: usize) -> usize { + while self.0[index] != index { + self.0[index] = self.0[self.0[index]]; + index = self.0[index]; + } + index + } + + fn union(&mut self, lhs: usize, rhs: usize) { + let (lhs, rhs) = (self.find(lhs), self.find(rhs)); + if lhs < rhs { + self.0[rhs] = lhs; + } else { + self.0[lhs] = rhs; + } + } +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::expr::Expression; + use vortex_array::expr::and; + use vortex_array::expr::col; + use vortex_array::expr::eq; + use vortex_array::expr::gt; + use vortex_array::expr::lit; + use vortex_array::expr::lt; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::filter_after_eval; + use super::group_conjuncts; + use crate::plan::physical::Eval; + use crate::plan::physical::EvalPlan; + use crate::plan::physical::Filter; + use crate::plan::physical::FilterPlan; + use crate::plan::physical::SegmentScan; + use crate::plan::physical::SegmentScanPlan; + use crate::scan::filter::FilterExpr; + use crate::segments::SegmentId; + + fn dtype() -> DType { + let i32 = DType::Primitive(PType::I32, NonNullable); + DType::Struct( + StructFields::from_iter([("a", i32.clone()), ("b", i32.clone()), ("c", i32)]), + NonNullable, + ) + } + + fn grouped(filter: Expression) -> VortexResult> { + let conjuncts = FilterExpr::new(filter.bind(&dtype())?).conjuncts().to_vec(); + Ok(group_conjuncts(&conjuncts)? + .iter() + .map(ToString::to_string) + .collect()) + } + + #[test] + fn groups_conjuncts_sharing_a_field_transitively() -> VortexResult<()> { + let filter = and( + and(gt(col("a"), lit(5)), lt(col("a"), col("b"))), + and(eq(col("c"), lit(1)), eq(col("b"), lit(2))), + ); + let groups = grouped(filter)?; + assert_eq!(groups.len(), 2, "{groups:?}"); + assert!( + groups[0].contains("$.a") && groups[0].contains("$.b"), + "{groups:?}" + ); + assert!( + groups[1].contains("$.c") && !groups[1].contains("$.a"), + "{groups:?}" + ); + Ok(()) + } + + #[test] + fn moves_filters_above_the_predicate() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, NonNullable); + let scan = SegmentScanPlan::new( + dtype.clone(), + 10, + SegmentId::from(0), + ReadContext::new([]), + None, + ) + .into_plan(); + let filtered = FilterPlan::new(scan).into_plan(); + let predicate = gt(root(), lit(5)).bind(&dtype)?; + let plan = EvalPlan::try_new(predicate, filtered)?.into_plan(); + + let rewritten = filter_after_eval(plan)?; + assert!(rewritten.is::(), "{}", rewritten.display_tree()); + let eval = rewritten.child_required(0)?; + assert!(eval.is::(), "{}", rewritten.display_tree()); + assert!( + eval.child_required(0)?.is::(), + "{}", + rewritten.display_tree() + ); + Ok(()) + } + + #[test] + fn keeps_disjoint_conjuncts_apart() -> VortexResult<()> { + let groups = grouped(and(gt(col("a"), lit(5)), lt(col("b"), lit(3))))?; + assert_eq!(groups.len(), 2, "{groups:?}"); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/v2/file.rs b/vortex-layout/src/scan/v2/file.rs new file mode 100644 index 00000000000..bfa85885829 --- /dev/null +++ b/vortex-layout/src/scan/v2/file.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; +use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::Weak; + +use once_cell::sync::OnceCell; +use parking_lot::Mutex; +use vortex_error::VortexResult; + +use crate::LayoutReader; +use crate::plan::physical::PlanRef; +use crate::scan::v2::ScanFile; +use crate::scan::v2::io::SegmentRanges; +use crate::scan::v2::io::segment_ranges; +use crate::scan::v2::lower::lower; +use crate::scan::v2::lower::lower_with_zones; + +/// Layout-derived state with the same lifetime as an explicitly cached layout reader. +/// +/// Holds only the lowered plans and segment ranges. Decoded values, such as dictionary values, +/// live no longer than the split that read them, and zone statistics no longer than the query +/// that proves over them. +#[derive(Default)] +pub struct FilePlans { + root: OnceCell, + zones: OnceCell, + ranges: OnceCell, +} + +/// Lowered plans shared by active scans over one layout reader. +/// +/// The registry holds only weak references, so it never extends their lifetime. An explicitly +/// cached reader owns its [`FilePlans`] independently, as V1 retains state on its reader tree. +pub(super) struct SharedFile { + pub(super) file: ScanFile, + /// The file's layout lowered to a plan. + pub(super) root: PlanRef, + /// Layout-derived state shared by active scans or retained by an explicitly cached reader. + plans: Arc, + _reader: Arc, +} + +impl SharedFile { + pub(super) fn zones(&self) -> VortexResult { + self.plans + .zones + .get_or_try_init(|| lower_with_zones(&self.file.layout)) + .cloned() + } + + pub(super) fn segment_ranges(&self) -> SegmentRanges { + Arc::clone( + self.plans + .ranges + .get_or_init(|| segment_ranges(&self.file.locations)), + ) + } +} + +/// A reader and a weak reference to its active scans' shared file. +type Entry = (Weak, Arc>>); + +static FILES: LazyLock>> = LazyLock::new(Default::default); + +/// The shared file of `reader`, created over `file` for its first scan. +pub(super) fn shared_file( + reader: &Arc, + file: ScanFile, +) -> VortexResult> { + let cell = { + let mut files = FILES.lock(); + files.retain(|(reader, _)| reader.strong_count() > 0); + if let Some((_, cell)) = files + .iter() + .find(|(known, _)| ptr::addr_eq(known.as_ptr(), Arc::as_ptr(reader))) + { + Arc::clone(cell) + } else { + let cell = Arc::default(); + files.push((Arc::downgrade(reader), Arc::clone(&cell))); + cell + } + }; + // Partitions of one file share initialization, while different files can lower in parallel. + let mut known = cell.lock(); + if let Some(shared) = known.upgrade() { + return Ok(shared); + } + let plans = file.plans.clone().unwrap_or_default(); + let root = plans.root.get_or_try_init(|| lower(&file.layout))?.clone(); + let shared = Arc::new(SharedFile { + root, + plans, + _reader: Arc::clone(reader), + file, + }); + *known = Arc::downgrade(&shared); + Ok(shared) +} diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs new file mode 100644 index 00000000000..f8524a2bc30 --- /dev/null +++ b/vortex-layout/src/scan/v2/io.rs @@ -0,0 +1,321 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::VecDeque; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use futures::FutureExt; +use futures::StreamExt; +use futures::TryFutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use futures::stream::FuturesUnordered; +use parking_lot::Mutex; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoResult; +use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::scan::planning::SegmentLocation; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// Byte ranges of a file's segments, keyed by `(offset, length)`, back to their segment ids. +pub(super) type SegmentRanges = Arc>; + +pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { + let ranges = locations.iter().enumerate().map(|(id, location)| { + let id = SegmentId::from(u32::try_from(id).unwrap_or(u32::MAX)); + ((location.offset, location.length as usize), id) + }); + Arc::new(ranges.collect()) +} + +/// Serves splits from any [`SegmentSource`], by segment id. +/// +/// A file opened for reading provides an [`IoService`] that drives the file's coalescing reads +/// directly. Any other segment source is served by this one. +pub(super) struct SegmentScanIo { + segments: Arc, + ranges: SegmentRanges, +} + +impl SegmentScanIo { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self { segments, ranges } + } +} + +impl IoService for SegmentScanIo { + fn session(&self) -> Arc { + Arc::new(SegmentIoSource::new( + Arc::clone(&self.segments), + Arc::clone(&self.ranges), + )) + } +} + +/// Serves a split's range requests from a segment source, without blocking. +/// +/// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once, +/// [`poll`](IoSource::poll) takes a finished one if there is any, and the split's future awaits +/// the next with [`poll_completion`](IoSource::poll_completion) when the driver has nothing else +/// to run. [`wait`](IoSource::wait) is never called. +/// +/// An announcement requests the segment and holds the request, never polled, until a prefetch or +/// fetch consumes it, so a source that coalesces requests can fold it into a nearby read. +/// A read started ahead of the fetch that needs it. +type Prefetched = Shared>>>; + +pub(super) struct SegmentIoSource { + segments: Arc, + ranges: SegmentRanges, + reads: Mutex>>, + /// Reads started by prefetches, consumed by the next fetch of their segment. + prefetched: Mutex>, + /// Announced segments, requested and never polled. + announced: Mutex>>, +} + +impl SegmentIoSource { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self { + segments, + ranges, + reads: Mutex::new(FuturesUnordered::new()), + prefetched: Mutex::default(), + announced: Mutex::default(), + } + } +} + +impl IoSource for SegmentIoSource { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + let reads = self.reads.lock(); + let mut prefetched = self.prefetched.lock(); + for request in batch { + if request.intent == IoIntent::Forget { + continue; + } + let IoTarget::Range { offset, len, .. } = request.target else { + vortex_bail!( + "SegmentIoSource only serves byte ranges, not {:?}", + request.target + ); + }; + let segment = *self + .ranges + .get(&(offset, len)) + .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; + if request.intent == IoIntent::Announce { + self.announced + .lock() + .entry(segment) + .or_default() + .push_back(self.segments.request(segment)); + continue; + } + if request.intent == IoIntent::Prefetch { + prefetched.entry(segment).or_insert_with(|| { + let read = self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared(); + // Polling registers the read as wanted, so the source starts it now. + drop( + read.clone() + .poll_unpin(&mut Context::from_waker(Waker::noop())), + ); + read + }); + continue; + } + let bytes = match prefetched.remove(&segment) { + Some(read) => read + .map_err(move |err| { + vortex_err!("prefetched read of segment {segment} failed: {err}") + }) + .boxed(), + None => self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)), + }; + let id = request.request; + reads.push( + async move { + Completion { + owner, + request: id, + result: bytes.await.map(IoResult::Bytes), + } + } + .boxed(), + ); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + // The split's future polls again with its own waker before it waits, so a read that + // finishes after this check is not missed. + let mut cx = Context::from_waker(Waker::noop()); + match self.reads.lock().poll_next_unpin(&mut cx) { + Poll::Ready(completion) => Ok(completion), + Poll::Pending => Ok(None), + } + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + match self.reads.lock().poll_next_unpin(cx) { + Poll::Ready(Some(completion)) => Poll::Ready(Ok(completion)), + Poll::Ready(None) => Poll::Ready(Err(vortex_err!( + "The split's driver is waiting with no reads in flight" + ))), + Poll::Pending => Poll::Pending, + } + } + + fn wait(&self) -> VortexResult { + vortex_bail!("SegmentIoSource is awaited through poll_completion, not waited on") + } + + fn release(&self, _owner: IoOwnerId) {} + + fn clear(&self) { + self.reads.lock().clear(); + self.prefetched.lock().clear(); + self.announced.lock().clear(); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use futures::FutureExt; + use futures::future; + use vortex_array::buffer::BufferHandle; + use vortex_buffer::Alignment; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use vortex_io::request::IoIntent; + use vortex_io::request::IoOwnerId; + use vortex_io::request::IoRequest; + use vortex_io::request::IoRequestId; + use vortex_io::request::IoSource; + use vortex_io::request::IoTarget; + + use super::SegmentIoSource; + use super::segment_ranges; + use crate::scan::planning::SegmentLocation; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; + use crate::segments::SegmentSource; + + /// Serves four bytes per segment and counts the reads it is asked for. + #[derive(Default)] + struct CountingSegments(AtomicUsize); + + impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.0.fetch_add(1, Ordering::Relaxed); + let bytes = ByteBuffer::from(vec![u8::try_from(*id).unwrap_or(u8::MAX); 4]); + future::ready(Ok(BufferHandle::new_host(bytes))).boxed() + } + } + + fn request(intent: IoIntent, id: u32) -> IoRequest { + IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::range(0, 4), + } + } + + /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, + /// and only the fetch is delivered. + #[tokio::test] + async fn fetch_reuses_a_prefetched_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 1)])?; + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(1)); + assert!(completion.result.is_ok()); + assert!(io.poll()?.is_none()); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2)])?; + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(2)); + assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[tokio::test] + async fn announcements_keep_independent_consumers() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + io.submit( + IoOwnerId(0), + vec![ + request(IoIntent::Announce, 0), + request(IoIntent::Announce, 1), + ], + )?; + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + for id in 2..5 { + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); + assert_eq!( + segments.0.load(Ordering::Relaxed), + if id < 4 { 2 } else { 3 } + ); + } + Ok(()) + } +} diff --git a/vortex-layout/src/scan/v2/legacy_tests.rs b/vortex-layout/src/scan/v2/legacy_tests.rs new file mode 100644 index 00000000000..61e953922de --- /dev/null +++ b/vortex-layout/src/scan/v2/legacy_tests.rs @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::TryStreamExt; +use futures::stream; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::DeserializeMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::expr::stats::Stat; +use vortex_array::stats::as_stat_bitset_bytes; +use vortex_error::VortexResult; +use vortex_io::session::RuntimeSessionExt; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use super::scan_file; +use crate::LayoutBuildContext; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::OwnedLayoutChildren; +use crate::VTable; +use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::LegacyStatsMetadata; +use crate::scan::v2::ScanBuilder; +use crate::segments::SegmentSource; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; +use crate::test::new_session; + +async fn flat( + array: ArrayRef, + segments: Arc, + session: &VortexSession, +) -> VortexResult { + let (mut sequence, eof) = SequenceId::root().split(); + FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + segments, + SequentialStreamAdapter::new( + array.dtype().clone(), + stream::iter([Ok((sequence.advance(), array))]), + ) + .sendable(), + eof, + session, + ) + .await +} + +#[tokio::test] +async fn legacy_string_zones_scan() -> VortexResult<()> { + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let data = VarBinViewArray::from_iter_str(["a", "b", "y", "z"]).into_array(); + let zones = StructArray::from_fields(&[ + ( + "max", + VarBinViewArray::from_iter([Some("b"), Some("z")], data.dtype().as_nullable()) + .into_array(), + ), + ( + "max_is_truncated", + BoolArray::from_iter([false, false]).into_array(), + ), + ])? + .into_array(); + let children = OwnedLayoutChildren::layout_children(vec![ + flat(data.clone(), Arc::clone(&segments), &session).await?, + flat(zones, Arc::clone(&segments), &session).await?, + ]); + let read_ctx = ReadContext::new([]); + let mut metadata = 2u32.to_le_bytes().to_vec(); + metadata.extend(as_stat_bitset_bytes(&[Stat::Max])); + let layout = ::build( + &LegacyStats, + data.dtype(), + 4, + &LegacyStatsMetadata::deserialize(&metadata)?, + vec![], + children.as_ref(), + &LayoutBuildContext { + session: &session, + array_read_ctx: &read_ctx, + }, + )? + .into_layout(); + let segments: Arc = segments; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let output = ScanBuilder::new(session.clone(), reader, scan_file(&segments, &layout)?) + .with_filter(gt(root(), lit("m")).bind(data.dtype())?) + .into_stream()? + .try_collect::>() + .await?; + assert_arrays_eq!( + ChunkedArray::try_new(output, data.dtype().clone())?, + VarBinViewArray::from_iter_str(["y", "z"]), + &mut session.create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-layout/src/scan/v2/list_tests.rs b/vortex-layout/src/scan/v2/list_tests.rs new file mode 100644 index 00000000000..40a66af289e --- /dev/null +++ b/vortex-layout/src/scan/v2/list_tests.rs @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::stream; +use vortex_array::ArrayContext; +use vortex_array::IntoArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_io::session::RuntimeSessionExt; + +use super::scan_file; +use crate::LayoutStrategy; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; +use crate::layouts::list::writer::ListLayoutStrategy; +use crate::scan::v2::supports; +use crate::segments::SegmentSource; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; +use crate::test::new_session; + +/// Lists are scanned by the default executor: their element ranges are only known once their +/// offsets are read, after a pipeline would have been compiled. +#[tokio::test] +async fn lists_are_not_supported() -> VortexResult<()> { + let session = new_session().with_tokio(); + let chunk = ListArray::try_new( + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2), Some(1), Some(0)]) + .into_array(), + buffer![0u64, 2, 2, 3, 5].into_array(), + Validity::from_iter([true, true, false, true]), + )? + .into_array(); + let dtype = chunk.dtype().clone(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let chunks = vec![Ok((sequence.advance(), chunk))]; + let layout = ChunkedLayoutStrategy::new(ListLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(chunks)).sendable(), + eof, + &session, + ) + .await?; + let segments: Arc = segments; + assert!(!supports(&scan_file(&segments, &layout)?)); + Ok(()) +} diff --git a/vortex-layout/src/scan/v2/lower.rs b/vortex-layout/src/scan/v2/lower.rs new file mode 100644 index 00000000000..aedd612713c --- /dev/null +++ b/vortex-layout/src/scan/v2/lower.rs @@ -0,0 +1,221 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lowers a file's layout tree into the physical plan the V2 executor runs. +//! +//! A copy of the test-only [`plan::lower`](crate::plan::physical::lower) with two modes for zoned layouts: +//! data plans lower them to their data child, and pruning plans keep them as a +//! [`ZonedPlan`] so pruning proofs can reach the zone statistics. + +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::physical::ConcatPlan; +use crate::plan::physical::FilterPlan; +use crate::plan::physical::LegacyStatsPlan; +use crate::plan::physical::ListPackPlan; +use crate::plan::physical::PackPlan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::SharePlan; +use crate::plan::physical::TakePlan; +use crate::plan::physical::ZonedPlan; + +/// Lowers `layout` into a physical plan. +/// +/// 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. +/// +/// A flat layout lowers to a [`Filter`](crate::plan::physical::Filter) over its segment scan, so the scan +/// returns only the rows it is executed with. +pub(crate) fn lower(layout: &LayoutRef) -> VortexResult { + lower_with(layout, Zones::Drop) +} + +/// Lowers `layout` like [`lower`], but keeps each zoned layout as a [`ZonedPlan`] over its data +/// and zone statistics, so pruning proofs can be pushed down to the zones. +pub(crate) fn lower_with_zones(layout: &LayoutRef) -> VortexResult { + lower_with(layout, Zones::Keep) +} + +/// What a zoned layout lowers to. +#[derive(Clone, Copy)] +enum Zones { + /// Its data child: data plans do not read zone statistics. + Drop, + /// A [`ZonedPlan`] over its data and zones. + Keep, +} + +fn lower_with(layout: &LayoutRef, zones: Zones) -> VortexResult { + if layout.is::() || layout.is::() { + return match zones { + Zones::Drop => { + let data = layout + .slot(0)? + .ok_or_else(|| vortex_err!("Zoned layout has no data child"))?; + lower_with(&data, zones) + } + Zones::Keep => lower_zoned(layout, zones), + }; + } + if let Some(layout) = layout.as_opt::() { + return Ok(FilterPlan::new(lower_flat(layout).into_plan()).into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_chunked(layout, zones)?.into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_struct(layout, zones)?.into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_dict(layout, zones)?.into_plan()); + } + if let Some(layout) = layout.as_opt::() { + return Ok(lower_list(layout, zones)?.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, zones: Zones) -> VortexResult { + let row_offsets = &layout.chunk_offsets()[..layout.nchildren()]; + // SAFETY: Chunked layout construction validates that every child has the parent dtype and + // that their row counts sum to the parent; its cached offsets use checked addition. + Ok(unsafe { + ConcatPlan::from_children_unchecked( + layout.dtype().clone(), + layout.row_count(), + row_offsets.into(), + lazy_children(layout.to_layout(), (0..layout.nchildren()).collect(), zones), + ) + }) +} + +fn lower_struct(layout: &StructLayout, zones: Zones) -> 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); + } + // SAFETY: Struct layout construction validates child count and row counts, and its logical + // slot schema defines every field dtype plus the optional non-nullable boolean validity dtype. + Ok(unsafe { + PackPlan::from_children_unchecked( + fields, + layout.dtype().nullability(), + layout.row_count(), + lazy_children(layout.to_layout(), slots, zones), + ) + }) +} + +fn lower_dict(layout: &DictLayout, zones: Zones) -> VortexResult { + // The values are shared, so an expression pushed onto them runs over values canonicalized + // once, and the projection of the same dictionary reuses them, as the V1 dictionary reader does. + // Dict serialization stores values before codes; the plan order is codes, values. + let layout = layout.to_layout(); + let slot = |slot| { + layout + .slot(slot)? + .ok_or_else(|| vortex_err!("Dict layout slot {slot} is absent")) + }; + let codes = lower_with(&slot(1)?, zones)?; + let values = SharePlan::new(lower_with(&slot(0)?, zones)?).into_plan(); + Ok(TakePlan::new(codes, values)) +} + +fn lower_list(layout: &ListLayout, zones: Zones) -> VortexResult { + let mut slots = vec![ELEMENTS_CHILD_INDEX, OFFSETS_CHILD_INDEX]; + if layout.dtype().is_nullable() { + slots.push(VALIDITY_CHILD_INDEX); + } + // SAFETY: List layout construction validates its element, offsets, and optional validity + // child shapes before this plan preserves them in the same logical order. + Ok(unsafe { + ListPackPlan::from_children_unchecked( + layout.dtype().clone(), + layout.row_count(), + lazy_children(layout.to_layout(), slots, zones), + ) + }) +} + +fn lower_zoned(layout: &LayoutRef, zones: Zones) -> 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") + }; + let children = lazy_children(Arc::clone(layout), vec![0, 1], zones); + let zone_len = u64::try_from(metadata.zone_len())?; + // A legacy table is converted to the aggregate schema when the plan is optimized. + Ok(match metadata.legacy_stats() { + Some(stats) => LegacyStatsPlan::from_children( + layout.dtype().clone(), + layout.row_count(), + children, + zone_len, + Arc::clone(stats), + ) + .into_plan(), + None => ZonedPlan::from_children( + layout.dtype().clone(), + layout.row_count(), + children, + zone_len, + metadata.aggregate_fns(), + ) + .into_plan(), + }) +} + +fn lazy_children(layout: LayoutRef, slots: Vec, zones: Zones) -> 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_with(&child, zones) + }) +} diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs new file mode 100644 index 00000000000..ea477a62ae5 --- /dev/null +++ b/vortex-layout/src/scan/v2/mod.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An alternative scan executor, selected by callers when [`enabled`] returns true. +//! +//! The executor has its own [`ScanBuilder`], a copy of the default one that also takes the +//! [`ScanFile`] the reader was opened over. +//! +//! Callers that configure a default [`ScanBuilder`](crate::scan::scan_builder::ScanBuilder) can +//! instead swap only the final call: [`prepare`] for its `prepare`, and [`into_stream`] for its +//! `into_stream`. Both copy the default builder into this executor's and return the same types as +//! the calls they replace. +//! +//! Each split runs as a [`FilterPlanner`](crate::scan::planning::FilterPlanner), which hands the +//! rows that survive to a [`ProjectionPlanner`](crate::scan::planning::ProjectionPlanner) and its +//! morsel. Both execute physical plans lowered from the file's layout, on the planning driver, and +//! the file's own segment source serves their reads. +//! +//! The driver runs inside the split's future, on whichever thread polls it, and the future awaits +//! the reads it starts from the file's segment source. Every split announces the segments its +//! plans are likely to read when it starts, so the source can coalesce them. +//! +//! A caller that drives the planning protocol itself takes the splits as [`SplitPlan`]s from +//! [`RepeatedScanV2::split_plans`] and admits them to its own run, reading through +//! [`RepeatedScanV2::io`]. + +mod conjuncts; +mod file; +pub(crate) mod io; +mod lower; +pub(crate) mod prefetch; +mod pruning; +mod repeated_scan; +mod scan_builder; +mod share; +mod split; +pub(crate) mod splits; +mod stream; + +use std::env; +use std::sync::Arc; +use std::sync::LazyLock; + +pub use file::FilePlans; +pub use repeated_scan::RepeatedScanV2; +pub use repeated_scan::prepare; +pub use scan_builder::ScanBuilder; +pub use split::SplitPlan; +pub use stream::into_stream; +use vortex_array::dtype::DType; +use vortex_io::request::IoService; + +use crate::LayoutRef; +use crate::scan::planning::SegmentLocation; +use crate::scan::scan_builder::ScanBuilder as DefaultScanBuilder; +use crate::segments::SegmentSource; + +/// The environment variable that selects this executor when set to `1`. +pub const ENV_VAR: &str = "VORTEX_SCAN_V2"; + +static ENABLED: LazyLock = LazyLock::new(|| env::var(ENV_VAR).is_ok_and(|v| v == "1")); + +/// Whether [`ENV_VAR`] selects this executor. Read once per process. +pub fn enabled() -> bool { + *ENABLED +} + +/// Whether this executor can scan `file`. +/// +/// Physical plans run on pipelines, which are compiled before any read, so a list's element +/// range, known only once its offsets are read, cannot be compiled: files with list columns are +/// scanned by the default executor. +pub fn supports(file: &ScanFile) -> bool { + !has_list(file.layout.dtype()) +} + +/// Whether this executor can run `builder`'s scan over `file`: it supports the file, and the scan +/// does not combine a filter with a row limit. +pub fn supports_scan(builder: &DefaultScanBuilder, file: &ScanFile) -> bool { + supports(file) && !builder.has_filter_and_limit() +} + +fn has_list(dtype: &DType) -> bool { + match dtype { + DType::List(..) => true, + DType::FixedSizeList(element, ..) => has_list(element), + DType::Struct(fields, _) => fields.fields().any(|field| has_list(&field)), + DType::Extension(ext) => has_list(ext.storage_dtype()), + _ => false, + } +} + +/// The file a scan reads: its root layout, where each segment lives, and the source that fetches +/// segments by id. +#[derive(Clone)] +pub struct ScanFile { + /// The root layout the scan builder's reader was opened over. + pub layout: LayoutRef, + /// The location of every segment, indexed by segment id. + pub locations: Arc<[SegmentLocation]>, + /// Fetches segment bytes by id. + pub segments: Arc, + /// Serves the splits' reads, when the file provides a service for them; otherwise they are + /// served from `segments`. + pub io: Option>, + /// Retains the lowered layout plans for an explicitly cached reader. No decoded data is + /// retained here. + /// Only share this state between scans over the same layout and segment source. + pub plans: Option>, +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/scan/v2/prefetch.rs b/vortex-layout/src/scan/v2/prefetch.rs new file mode 100644 index 00000000000..86fd469e33e --- /dev/null +++ b/vortex-layout/src/scan/v2/prefetch.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Registers a split's reads with the segment source before the split runs. +//! +//! A registered read does no IO by itself, but a file's segment source coalesces it into any +//! nearby read that is issued, so reads registered for every split up front are served in a few +//! large IOs rather than many small ones. The layout reader registers its reads the same way. + +use std::cmp; +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::plan::physical::Concat; +use crate::plan::physical::Eval; +use crate::plan::physical::Filter; +use crate::plan::physical::Pack; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::Share; +use crate::plan::physical::Take; +use crate::segments::SegmentId; + +/// The most row ranges reads are asked for ahead of time one by one; a more scattered selection +/// is asked for as the range spanning it. +const MAX_PREFETCH_RANGES: usize = 64; + +/// The row ranges to ask for ahead of time when `mask` selects rows of `rows`. +pub(crate) fn selected_ranges(rows: &Range, mask: &Mask) -> Vec> { + let start = rows.start; + match mask.slices() { + AllOr::All => vec![rows.clone()], + AllOr::None => Vec::new(), + AllOr::Some(slices) if slices.len() <= MAX_PREFETCH_RANGES => slices + .iter() + .map(|&(begin, end)| start + begin as u64..start + end as u64) + .collect(), + AllOr::Some(slices) => { + let first = slices.first().map_or(0, |slice| slice.0); + let last = slices.last().map_or(0, |slice| slice.1); + let range = start + first as u64..start + last as u64; + vec![range] + } + } +} + +/// Adds the segments `plan` is likely to read over `rows` to `segments`. +/// +/// This is a hint, not a promise: it follows plans whose children cover the same rows, or a known +/// part of them, and stops at any other plan, such as a list whose elements cover other rows. +pub(crate) fn plan_segments( + plan: &PlanRef, + rows: Range, + segments: &mut Vec, +) -> VortexResult<()> { + if rows.is_empty() { + return Ok(()); + } + if let Some(scan) = plan.as_opt::() { + segments.push(scan.segment_id()); + } else if let Some(concat) = plan.as_opt::() { + let offsets = concat.row_offsets(); + let first = offsets.partition_point(|&offset| offset <= rows.start); + for (index, &offset) in offsets + .iter() + .enumerate() + .skip(first.saturating_sub(1)) + .take_while(|&(_, &offset)| offset < rows.end) + { + let child = plan.child_ref_required(index)?; + let start = cmp::max(rows.start, offset); + let end = cmp::min(rows.end, offset + child.row_count()); + plan_segments(child, start - offset..end - offset, segments)?; + } + } else if plan.is::() { + plan_segments(plan.child_ref_required(0)?, rows, segments)?; + let values = plan.child_ref_required(1)?; + let len = values.row_count(); + plan_segments(values, 0..len, segments)?; + } else if plan.is::() { + plan_segments(plan.child_ref_required(0)?, rows, segments)?; + } else if plan.is::() || plan.is::() || plan.is::() { + for child in plan.children().iter_refs() { + plan_segments(child?, rows.clone(), segments)?; + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::*; + use crate::plan::physical::ConcatPlan; + use crate::plan::physical::FilterPlan; + use crate::plan::physical::SegmentScanPlan; + + const DTYPE: DType = DType::Primitive(PType::I32, NonNullable); + + /// Four flat chunks of 1000 rows, in segments 0 to 3. + fn chunked() -> VortexResult { + let chunks = (0..4) + .map(|segment| { + let scan = SegmentScanPlan::new( + DTYPE, + 1000, + SegmentId::from(segment), + ReadContext::new(ArrayContext::empty().to_ids()), + None, + ); + FilterPlan::new(scan.into_plan()).into_plan() + }) + .collect(); + Ok(ConcatPlan::try_new(DTYPE, chunks)?.into_plan()) + } + + #[rstest] + #[case::everything(0..4000, &[0, 1, 2, 3])] + #[case::inside_one_chunk(1200..1800, &[1])] + #[case::whole_chunk(1000..2000, &[1])] + #[case::across_a_boundary(999..1001, &[0, 1])] + #[case::several_chunks(500..2500, &[0, 1, 2])] + #[case::last_row(3999..4000, &[3])] + #[case::empty(1500..1500, &[])] + fn concat_reads_only_the_chunks_it_overlaps( + #[case] rows: Range, + #[case] expected: &[u32], + ) -> VortexResult<()> { + let mut segments = Vec::new(); + plan_segments(&chunked()?, rows, &mut segments)?; + let expected: Vec = expected.iter().copied().map(SegmentId::from).collect(); + assert_eq!(segments, expected); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/v2/pruning.rs b/vortex-layout/src/scan/v2/pruning.rs new file mode 100644 index 00000000000..ec0f6096e5b --- /dev/null +++ b/vortex-layout/src/scan/v2/pruning.rs @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::env; +use std::future; +use std::ops::Range; +use std::sync::Arc; +use std::sync::LazyLock; + +use futures::channel::oneshot; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::StructArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; +use vortex_array::validity::Validity; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoService; +use vortex_mask::Mask; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; + +use crate::plan::physical::Eval; +use crate::plan::physical::Pack; +use crate::plan::physical::PlanRef; +use crate::plan::physical::Zoned; +use crate::plan::physical::pipeline::ValueCache; +use crate::scan::planning::ScanPlans; +use crate::scan::planning::graph::GraphStep; +use crate::scan::planning::graph::ProtocolGraph; +use crate::scan::planning::graph::ScanGraph; + +pub(super) fn file_pruning_enabled() -> bool { + static ENABLED: LazyLock = + LazyLock::new(|| env::var("VORTEX_SCAN_FILE_PRUNING").is_ok_and(|value| value == "1")); + *ENABLED +} + +pub(super) struct PrunedFile { + pub ranges: Vec>, + pub dynamic: bool, +} + +impl PrunedFile { + /// Keeps holes inside a filter split as a mask, so disjoint surviving zones do not create + /// separate tasks that decode the same compressed chunk independently. + pub(super) fn select(&self, rows: Range) -> VortexResult, Mask)>> { + let lo = self.ranges.partition_point(|range| range.end <= rows.start); + let hi = self.ranges.partition_point(|range| range.start < rows.end); + if lo >= hi { + return Ok(None); + } + let selected = &self.ranges[lo..hi]; + let rows = + rows.start.max(selected[0].start)..rows.end.min(selected[selected.len() - 1].end); + let len = usize::try_from(rows.end - rows.start)?; + if selected.len() == 1 { + return Ok(Some((rows, Mask::new_true(len)))); + } + let mut bits = BitBufferMut::with_capacity(len); + let mut cursor = rows.start; + for range in selected { + let start = rows.start.max(range.start); + let end = rows.end.min(range.end); + bits.append_n(false, usize::try_from(start - cursor)?); + bits.append_n(true, usize::try_from(end - start)?); + cursor = end; + } + Ok(Some((rows, Mask::from(bits.freeze())))) + } +} + +/// Loads zone proofs through the planning protocol before any data splits are announced. +pub(super) async fn prune_file( + plans: &ScanPlans, + proof: PlanRef, + io: &Arc, +) -> VortexResult { + if proof.row_count() == 0 { + return Ok(PrunedFile { + ranges: Vec::new(), + dynamic: false, + }); + } + let mut boundaries = vec![0, proof.row_count()]; + let mut dynamic = false; + zone_boundaries(&proof, &mut boundaries, &mut dynamic)?; + boundaries.sort_unstable(); + boundaries.dedup(); + // A single representative row loads every complete zone table and initializes each proof's + // cache. Expanding the proof over the file here would allocate one boolean per data row. + let graph = ScanGraph::try_new( + plans.session.clone(), + &proof, + 0..1, + Mask::new_true(1), + plans.row_offset, + ValueCache::default(), + )?; + let scope = WorkScope { + file_ordinal: 0, + rows: 0..proof.row_count(), + }; + let (sender, receiver) = oneshot::channel(); + let planner = FilePruningPlanner { + graph: ProtocolGraph::new(graph, Arc::clone(&plans.locations), 0), + proof, + plans: plans.clone(), + boundaries, + result: Some(sender), + }; + let mut run = Run::new(); + run.admit(Box::new(planner), scope, io.session()); + loop { + match run.advance()? { + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + Progress::Batch(_) => vortex_bail!("file pruning emitted a data batch"), + } + } + Ok(PrunedFile { + ranges: receiver + .await + .map_err(|_| vortex_err!("file pruning did not finish"))?, + dynamic, + }) +} + +struct FilePruningPlanner { + graph: ProtocolGraph, + proof: PlanRef, + plans: ScanPlans, + boundaries: Vec, + result: Option>>>, +} + +impl IoConsumer for FilePruningPlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.graph.set_io_result(request, result); + } +} + +impl Planner for FilePruningPlanner { + fn state(&self) -> State { + if self.result.is_none() { + State::Done + } else if self.graph.state() == State::Done { + State::NeedsCompute + } else { + self.graph.state() + } + } + + fn compute(&mut self) -> VortexResult { + if self.graph.state() != State::Done { + return Ok(match self.graph.compute()? { + GraphStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + GraphStep::Yield | GraphStep::Piece(_) => PlannerOutput::Continue, + }); + } + let starts = &self.boundaries[..self.boundaries.len() - 1]; + let values = compact_proof(&self.proof, starts, &self.plans)?; + let pruned: Mask = values + .fill_null(false)? + .execute(&mut self.plans.session.create_execution_ctx())?; + let mut ranges: Vec> = Vec::new(); + for (index, pair) in self.boundaries.windows(2).enumerate() { + if pruned.value(index) { + continue; + } + if let Some(last) = ranges.last_mut() + && last.end == pair[0] + { + last.end = pair[1]; + } else { + ranges.push(pair[0]..pair[1]); + } + } + tracing::debug!(target: "vortex_layout::scan::v2::pruning", + file_rows = self.proof.row_count(), intervals = starts.len(), + surviving_rows = ranges.iter().map(|r| r.end - r.start).sum::(), + surviving_ranges = ranges.len(), "file pruning finished"); + if let Some(result) = self.result.take() { + drop(result.send(ranges)); + } + Ok(PlannerOutput::Done) + } +} + +fn zone_boundaries( + plan: &PlanRef, + boundaries: &mut Vec, + dynamic: &mut bool, +) -> VortexResult<()> { + if let Some(zoned) = plan.as_opt::() { + let expression = zoned + .pruning_expression() + .ok_or_else(|| vortex_err!("missing zone proof"))?; + *dynamic |= DynamicExprUpdates::new(expression).is_some(); + let mut row = zoned.zone_len(); + while row < plan.row_count() { + boundaries.push(row); + row = row.saturating_add(zoned.zone_len()); + } + } else { + if let Some(eval) = plan.as_opt::() { + *dynamic |= DynamicExprUpdates::new(eval.expression()).is_some(); + } + for child in plan.children().iter() { + zone_boundaries(&child?, boundaries, dynamic)?; + } + } + Ok(()) +} + +/// Evaluates one representative per interval in the union of all column zone boundaries. This +/// preserves compound proofs even when columns have different zone sizes; memory scales with +/// the zone count, rather than the file's row count. +fn compact_proof(plan: &PlanRef, starts: &[u64], plans: &ScanPlans) -> VortexResult { + if let Some(zoned) = plan.as_opt::() { + let expression = zoned + .pruning_expression() + .ok_or_else(|| vortex_err!("missing zone proof"))?; + let mask = zoned + .pruned_zones() + .ok_or_else(|| vortex_err!("missing zone cache"))? + .get(expression, &plans.session)? + .ok_or_else(|| vortex_err!("zone proof was not loaded"))?; + let validity = if plan.dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }; + return Ok(BoolArray::try_new( + starts + .iter() + .map(|row| Ok(mask.value(usize::try_from(row / zoned.zone_len())?))) + .collect::>()?, + validity, + )? + .into_array()); + } + if let Some(eval) = plan.as_opt::() { + return compact_proof(&eval.child_plan()?, starts, plans)?.apply_bound(eval.expression()); + } + if let Some(pack) = plan.as_opt::() { + let fields = (0..pack.nfields()) + .map(|index| compact_proof(&pack.child_required(index)?, starts, plans)) + .collect::>>()?; + let validity = match pack.validity()? { + Some(validity) => Validity::Array(compact_proof(&validity, starts, plans)?), + None => Validity::NonNullable, + }; + return Ok(StructArray::try_new( + pack.fields().names().clone(), + fields, + starts.len(), + validity, + )? + .into_array()); + } + vortex_bail!("unsupported file pruning plan: {}", plan.id()) +} diff --git a/vortex-layout/src/scan/v2/repeated_scan.rs b/vortex-layout/src/scan/v2/repeated_scan.rs new file mode 100644 index 00000000000..1318e303dd9 --- /dev/null +++ b/vortex-layout/src/scan/v2/repeated_scan.rs @@ -0,0 +1,606 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cmp; +use std::iter; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use itertools::Either; +use itertools::Itertools; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::BoundExpression; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::request::IoService; +use vortex_mask::Mask; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; +use vortex_session::VortexSession; +use vortex_utils::parallelism::get_available_parallelism; + +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Pack; +use crate::plan::physical::PlanRef; +use crate::plan::physical::Zoned; +use crate::plan::physical::optimize; +use crate::plan::physical::plan_row_idx_expression; +use crate::scan::filter::FilterExpr; +use crate::scan::planning::FilterPlans; +use crate::scan::planning::ScanPlans; +use crate::scan::planning::plan_sparse_projection; +use crate::scan::planning::plan_split; +use crate::scan::scan_builder; +use crate::scan::splits::Splits; +use crate::scan::splits::attempt_split_ranges; +use crate::scan::v2::ScanBuilder; +use crate::scan::v2::ScanFile; +use crate::scan::v2::conjuncts::filter_after_eval; +use crate::scan::v2::conjuncts::group_conjuncts; +use crate::scan::v2::file::SharedFile; +use crate::scan::v2::file::shared_file; +use crate::scan::v2::io::SegmentScanIo; +use crate::scan::v2::pruning::PrunedFile; +use crate::scan::v2::pruning::file_pruning_enabled; +use crate::scan::v2::pruning::prune_file; +use crate::scan::v2::share::unshare_unread; +use crate::scan::v2::split::SplitPlan; +use crate::scan::v2::splits::MAX_SPARSE_ROWS; +use crate::scan::v2::splits::chunk_starts; +use crate::scan::v2::splits::filter_split_boundaries; +use crate::scan::v2::splits::layout_chunk_starts; +use crate::scan::v2::splits::max_split_rows; +use crate::scan::v2::splits::sparse_filter_split_boundaries; + +/// Computes split ranges for `builder` and returns an executable scan over `file`, the file the +/// builder's reader was opened over. +/// +/// The file's layout is lowered once into a physical plan, and the filter and projection are +/// planned and optimized over it. Every split then runs those plans through the planning +/// protocol, and never calls the layout reader's evaluation methods. The builder's reader is only +/// used to choose the splits. +/// +/// The replacement for the default [`ScanBuilder::prepare`](scan_builder::ScanBuilder::prepare). +/// Equivalent to copying `builder` into a [`ScanBuilder`] and preparing that. +pub fn prepare( + builder: scan_builder::ScanBuilder, + file: ScanFile, +) -> VortexResult> { + ScanBuilder::from_default(builder, file).prepare() +} + +/// Plans `builder`'s scan over its file. See [`prepare`]. +pub(super) fn prepare_scan( + builder: ScanBuilder, +) -> VortexResult> { + let dtype = builder.dtype()?; + + if builder.filter.is_some() && builder.limit.is_some() { + vortex_bail!("Vortex doesn't support scans with both a filter and a limit") + } + + let shared = shared_file(&builder.layout_reader, builder.file)?; + let prepared = PreparedPlans::new( + &shared, + &builder.projection, + builder.filter.as_ref(), + &builder.selection, + &builder.session, + )?; + let plans = ScanPlans { + session: builder.session, + locations: Arc::clone(&shared.file.locations), + projection: prepared.projection.clone(), + projection_starts: Arc::clone(&prepared.projection_starts), + row_offset: builder.row_offset, + values: Default::default(), + }; + + let sparse_projection = builder.projection.is_root() + && builder.filter.is_none() + && matches!(&builder.selection, Selection::IncludeByIndex(indices) if indices.len() <= MAX_SPARSE_ROWS); + let splits = match attempt_split_ranges(&builder.selection, builder.row_range.as_ref()) { + Some(ranges) => Splits::Ranges(ranges), + None => { + let threads = get_available_parallelism().unwrap_or(1); + let boundaries = if sparse_projection + && builder.limit.is_none() + && let Selection::IncludeByIndex(indices) = &builder.selection + { + let indices = indices.as_slice(); + let indices = match &builder.row_range { + Some(range) => { + let start = indices.partition_point(|row| *row < range.start); + let end = indices.partition_point(|row| *row < range.end); + &indices[start..end] + } + None => indices, + }; + sparse_filter_split_boundaries( + &prepared.filter_starts, + &prepared.all_starts, + shared.root.row_count(), + indices, + threads, + ) + } else { + filter_split_boundaries( + &prepared.filter_starts, + &prepared.all_starts, + 0..shared.root.row_count(), + max_split_rows(shared.root.row_count(), threads), + ) + }; + Splits::Natural(boundaries.into()) + } + }; + + Ok(RepeatedScanV2 { + pruning: prepared.pruning.clone(), + plans, + filter: prepared.filter, + io: shared.file.io.clone().unwrap_or_else(|| { + Arc::new(SegmentScanIo::new( + Arc::clone(&shared.file.segments), + shared.segment_ranges(), + )) + }), + shared, + sparse_projection, + row_range: builder.row_range, + selection: builder.selection, + splits, + map_fn: builder.map_fn, + limit: builder.limit, + dtype, + }) +} + +/// Expression plans and chunk boundaries for a scan. +struct PreparedPlans { + projection: PlanRef, + filter: Option, + pruning: Option, + projection_starts: Arc<[u64]>, + filter_starts: Arc<[u64]>, + all_starts: Arc<[u64]>, +} + +impl PreparedPlans { + fn new( + shared: &SharedFile, + projection: &BoundExpression, + filter: Option<&BoundExpression>, + selection: &Selection, + session: &VortexSession, + ) -> VortexResult { + if projection.is_root() && filter.is_none() { + // Identity scans need no expression rewrites or changes to dictionary value forms. + // Read chunk metadata without lowering the unselected parts of the layout into plans. + let indices = match selection { + Selection::IncludeByIndex(indices) + if indices.len() as u64 <= shared.root.row_count() / 8 => + { + Some(indices.as_slice()) + } + _ => None, + }; + let starts: Arc<[u64]> = layout_chunk_starts(&shared.file.layout, indices)?.into(); + return Ok(Self { + projection: shared.root.clone(), + filter: None, + pruning: None, + projection_starts: Arc::clone(&starts), + filter_starts: Arc::clone(&starts), + all_starts: starts, + }); + } + let root = shared.root.clone(); + let plan = |expression| optimize(plan_row_idx_expression(expression, root.clone())?); + let filter_expression = filter; + let filter = filter + .cloned() + .map(|filter| { + let conjuncts = group_conjuncts(FilterExpr::new(filter).conjuncts())?; + VortexResult::Ok(Arc::new(FilterExpr::from_conjuncts(conjuncts))) + }) + .transpose()?; + // The projection comes first, then one plan per conjunct. + let mut all_plans = vec![plan(projection.clone())?]; + for conjunct in filter.iter().flat_map(|filter| filter.conjuncts()) { + all_plans.push(filter_after_eval(plan(conjunct.clone())?)?); + } + let mut all_plans = unshare_unread(all_plans)?; + let conjunct_plans = all_plans.split_off(1); + let projection = all_plans.remove(0); + // Filter splits are sized by the chunks of the columns the filter reads, or of those the + // projection reads when there is no filter, and each is cut into projection splits. + let projection_starts: Arc<[u64]> = chunk_starts([&projection])?.into(); + let filter_starts = if conjunct_plans.is_empty() { + projection_starts.to_vec() + } else { + chunk_starts(&conjunct_plans)? + }; + let mut all_starts = filter_starts.clone(); + all_starts.extend_from_slice(&projection_starts); + all_starts.sort_unstable(); + all_starts.dedup(); + let filter = filter + .map(|filter| FilterPlans::conjuncts(filter, conjunct_plans)) + .transpose()?; + let pruning = filter_expression + .map(|filter| pruning_plan(filter, &shared.zones()?, session)) + .transpose()? + .flatten(); + Ok(Self { + projection, + filter, + pruning, + projection_starts, + filter_starts: filter_starts.into(), + all_starts: all_starts.into(), + }) + } +} + +/// Joins a filter split's batches into one array without copying them: a struct whose fields are +/// chunked over the batches when every batch is a struct of the scan's non-nullable struct dtype, +/// and a chunked array otherwise. +fn join_batches(mut batches: Vec, dtype: &DType) -> VortexResult> { + if batches.len() <= 1 { + return Ok(batches.pop()); + } + if let DType::Struct(fields, Nullability::NonNullable) = dtype + && let Some(structs) = batches + .iter() + .map(|batch| batch.as_opt::().map(|batch| batch.into_owned())) + .collect::>>() + { + let len = structs.iter().map(|batch| batch.len()).sum(); + let columns = (0..fields.nfields()) + .map(|index| { + let chunks = structs + .iter() + .map(|batch| batch.unmasked_field(index).clone()) + .collect::>(); + let dtype = chunks[0].dtype().clone(); + Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) + }) + .collect::>>()?; + let array = + StructArray::try_new(fields.names().clone(), columns, len, Validity::NonNullable)?; + return Ok(Some(array.into_array())); + } + Ok(Some( + ChunkedArray::try_new(batches, dtype.clone())?.into_array(), + )) +} + +/// A prepared scan that turns row ranges into one task per filter split. +/// +/// The replacement for [`RepeatedScan`](crate::scan::repeated_scan::RepeatedScan). +pub struct RepeatedScanV2 { + sparse_projection: bool, + shared: Arc, + /// Proves rows can't match the filter from zone statistics, when the filter allows it. + pruning: Option, + plans: ScanPlans, + filter: Option, + /// Serves the splits' reads. + io: Arc, + row_range: Option>, + selection: Selection, + splits: Splits, + map_fn: Arc VortexResult + Send + Sync>, + limit: Option, + dtype: DType, +} + +impl RepeatedScanV2 { + /// The dtype of the projected arrays. + pub fn dtype(&self) -> &DType { + &self.dtype + } + + /// The service the scan's splits read through: one session per planning run root. + pub fn io(&self) -> &Arc { + &self.io + } + + /// Converts a projected array into the scan's output. + pub fn map(&self, array: ArrayRef) -> VortexResult { + (self.map_fn)(array) + } + + /// Joins the batches of one filter split, in row order, into one array without copying them. + /// Returns `None` when there are none. + pub fn join(&self, batches: Vec) -> VortexResult> { + join_batches(batches, self.plans.projection.dtype()) + } + + /// Returns one task per filter split of `row_range` that has selected rows. + /// + /// A task returns its projection splits' batches as one array, whose fields are chunked over + /// them, so nothing is copied. + pub fn execute( + &self, + row_range: Option>, + ) -> VortexResult>>>> { + self.execute_splits(self.split_plans(row_range)?) + } + + /// Like [`execute`](Self::execute), with file pruning before split creation when + /// `VORTEX_SCAN_FILE_PRUNING=1`. Statistics IO completes before this returns the data tasks. + pub async fn execute_pruned( + &self, + row_range: Option>, + ) -> VortexResult>>>> { + self.execute_splits(self.execution_splits(row_range).await?) + } + + fn execute_splits( + &self, + splits: Vec, + ) -> VortexResult>>>> { + let dtype = self.plans.projection.dtype().clone(); + Ok(splits + .into_iter() + .map(|split| { + let io = Arc::clone(&self.io); + let map_fn = Arc::clone(&self.map_fn); + let dtype = dtype.clone(); + let run = split.run(io); + let shared = Arc::clone(&self.shared); + async move { + let _shared = shared; + join_batches(run.await?, &dtype)? + .map(|array| map_fn(array)) + .transpose() + } + .boxed() + }) + .collect()) + } + + /// Returns one task per filter split of `row_range` that has selected rows, each returning + /// one batch per projection split with selected rows, in row order. + pub fn execute_batches( + &self, + row_range: Option>, + ) -> VortexResult>>>> { + self.execute_batch_splits(self.split_plans(row_range)?) + } + + pub(super) async fn execute_batches_pruned( + &self, + row_range: Option>, + ) -> VortexResult>>>> { + self.execute_batch_splits(self.execution_splits(row_range).await?) + } + + fn execute_batch_splits( + &self, + splits: Vec, + ) -> VortexResult>>>> { + Ok(splits + .into_iter() + .map(|split| { + let io = Arc::clone(&self.io); + let map_fn = Arc::clone(&self.map_fn); + let run = split.run(io); + let shared = Arc::clone(&self.shared); + async move { + let _shared = shared; + run.await?.into_iter().map(|array| map_fn(array)).collect() + } + .boxed() + }) + .collect()) + } + + async fn execution_splits( + &self, + row_range: Option>, + ) -> VortexResult> { + if file_pruning_enabled() { + self.pruned_split_plans(row_range).await + } else { + self.split_plans(row_range) + } + } + + /// Prunes the whole file's zone statistics before constructing or announcing data splits. + /// Surviving splits retain pruning only when dynamic bounds can change during execution. + pub async fn pruned_split_plans( + &self, + row_range: Option>, + ) -> VortexResult> { + let pruned = match &self.pruning { + Some(proof) => Some(prune_file(&self.plans, proof.clone(), &self.io).await?), + None => None, + }; + self.split_plans_with_pruning(row_range, pruned.as_ref()) + } + + /// The filter splits of `row_range` that have selected rows, in row order, each ready to be + /// admitted to a planning run that reads through a session of [`io`](Self::io). + pub fn split_plans(&self, row_range: Option>) -> VortexResult> { + self.split_plans_with_pruning(row_range, None) + } + + fn split_plans_with_pruning( + &self, + row_range: Option>, + pruned: Option<&PrunedFile>, + ) -> VortexResult> { + let selection_range: Option> = match &self.selection { + Selection::IncludeByIndex(buf) if !buf.is_empty() => { + Some(buf[0]..buf[buf.len() - 1] + 1) + } + Selection::IncludeRoaring(roaring) if !roaring.is_empty() => { + Some(roaring.min().vortex_expect("empty")..roaring.max().vortex_expect("empty") + 1) + } + _ => None, + }; + let row_range = intersect_ranges(self.row_range.as_ref(), row_range); + let row_range = intersect_ranges(row_range.as_ref(), selection_range); + + let ranges = match &self.splits { + Splits::Natural(vec) => { + let splits_iter = match row_range { + None => Either::Left(vec.iter().copied()), + Some(range) => { + if range.is_empty() { + return Ok(Vec::new()); + } + let lo = vec.partition_point(|&x| x <= range.start); + let hi = vec.partition_point(|&x| x < range.end); + Either::Right( + iter::once(range.start) + .chain(vec[lo..hi].iter().copied()) + .chain(iter::once(range.end)), + ) + } + }; + Either::Left(splits_iter.tuple_windows().map(|(start, end)| start..end)) + } + Splits::Ranges(ranges) => Either::Right(match row_range { + None => Either::Left(ranges.iter().cloned()), + Some(range) => { + if range.is_empty() { + return Ok(Vec::new()); + } + Either::Right(ranges.iter().filter_map(move |r| { + let start = cmp::max(r.start, range.start); + let end = cmp::min(r.end, range.end); + (start < end).then_some(start..end) + })) + } + }), + }; + + let pruning = if pruned.is_some_and(|pruned| !pruned.dynamic) { + None + } else { + self.pruning.clone() + }; + let mut limit = self.limit; + let mut splits = Vec::new(); + for range in ranges { + let (range, surviving) = match pruned { + Some(pruned) => match pruned.select(range)? { + Some((range, mask)) => (range, Some(mask)), + None => continue, + }, + None => (range, None), + }; + let row_mask = self.selection.row_mask(&range); + let selected = match surviving { + Some(surviving) => row_mask.mask() & &surviving, + None => row_mask.mask().clone(), + }; + if selected.all_false() { + continue; + } + let mask = match (&self.filter, limit.as_mut()) { + (None, Some(0)) => Mask::new_false(selected.len()), + (None, Some(l)) => { + let true_count = selected.true_count(); + let mask_limit = usize::try_from(*l) + .map(|l| l.min(true_count)) + .unwrap_or(true_count); + *l -= mask_limit as u64; + selected.limit(mask_limit) + } + _ => selected, + }; + if !mask.all_false() { + let scope = WorkScope { + file_ordinal: 0, + rows: row_mask.row_range(), + }; + splits.push(SplitPlan { + root: if self.sparse_projection { + plan_sparse_projection(self.plans.clone(), scope.clone(), mask) + } else { + plan_split( + self.plans.clone(), + pruning.clone(), + self.filter.clone(), + scope.clone(), + mask, + )? + }, + scope, + }); + } + if limit.is_some_and(|l| l == 0) { + break; + } + } + + tracing::debug!( + target: "vortex_layout::scan::v2::pruning", + splits = splits.len(), + file_pruned = pruned.is_some(), + "filter splits prepared" + ); + Ok(splits) + } +} + +/// Plans the zone-statistics proof that `filter` is false, over a copy of the file's plan that keeps +/// its zones. +/// +/// The filter is falsified into a predicate over statistics, which the optimizer pushes down to +/// the zoned columns it reads and rewrites into pruning plans over their zone tables. Returns +/// `None` when the filter cannot be falsified, or when any part of the proof would still read +/// column data rather than zone statistics. +fn pruning_plan( + filter: &BoundExpression, + zones: &PlanRef, + session: &VortexSession, +) -> VortexResult> { + let Some(predicate) = filter.falsify(session)? else { + return Ok(None); + }; + let plan = optimize(EvalPlan::try_new(predicate, zones.clone())?.into_plan())?; + Ok(reads_only_zones(&plan)?.then_some(plan)) +} + +/// Whether every leaf of `plan` is a pruning plan over zone statistics. +fn reads_only_zones(plan: &PlanRef) -> VortexResult { + if let Some(zoned) = plan.as_opt::() { + return Ok(zoned.is_pruning()); + } + if !plan.is::() && !plan.is::() { + return Ok(false); + } + for child in plan.children().iter() { + if !reads_only_zones(&child?)? { + return Ok(false); + } + } + Ok(true) +} + +fn intersect_ranges(left: Option<&Range>, right: Option>) -> Option> { + match (left, right) { + (None, None) => None, + (None, Some(r)) => Some(r), + (Some(l), None) => Some(l.clone()), + (Some(l), Some(r)) => Some(cmp::max(l.start, r.start)..cmp::min(l.end, r.end)), + } +} diff --git a/vortex-layout/src/scan/v2/scan_builder.rs b/vortex-layout/src/scan/v2/scan_builder.rs new file mode 100644 index 00000000000..b620b5a070d --- /dev/null +++ b/vortex-layout/src/scan/v2/scan_builder.rs @@ -0,0 +1,333 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use futures::Stream; +use futures::future::BoxFuture; +use vortex_array::ArrayRef; +use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; +use vortex_array::iter::ArrayIterator; +use vortex_array::iter::ArrayIteratorAdapter; +use vortex_array::stream::ArrayStream; +use vortex_array::stream::ArrayStreamAdapter; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_io::runtime::BlockingRuntime; +use vortex_metrics::MetricsRegistry; +use vortex_scan::selection::Selection; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_session::VortexSession; + +use crate::LayoutReader; +use crate::LayoutReaderRef; +use crate::scan::limit::ScanLimit; +use crate::scan::scan_builder; +use crate::scan::scan_builder::referenced_field_masks; +use crate::scan::split_by::SplitBy; +use crate::scan::v2::RepeatedScanV2; +use crate::scan::v2::ScanFile; +use crate::scan::v2::repeated_scan::prepare_scan; +use crate::scan::v2::stream::LazyScanStream; + +/// Builder for scanning a file with this executor into arrays, streams, iterators, or mapped +/// outputs. +/// +/// A copy of the default [`ScanBuilder`](scan_builder::ScanBuilder), which it takes every option +/// from, plus the [`ScanFile`] the reader was opened over. It prepares and runs the scan through +/// the planning protocol rather than through the layout reader. +/// +/// A scan has three independent row restriction mechanisms: +/// +/// - [`with_row_range`](Self::with_row_range) selects a contiguous range before scanning. +/// - [`with_selection`](Self::with_selection) applies a [`Selection`] inside that range. +/// - [`with_filter`](Self::with_filter) evaluates an expression predicate during execution. +/// +/// Projection and filter expressions must be bound against the reader dtype. +/// +/// This executor sizes its own splits from the chunks of the columns the scan reads, so it does +/// not yet honour [`with_split_by`](Self::with_split_by) or +/// [`with_natural_splits`](Self::with_natural_splits), and it records no metrics. The layout +/// reader is only used to share the file's plans between scans of one reader. +pub struct ScanBuilder { + pub(super) session: VortexSession, + pub(super) layout_reader: LayoutReaderRef, + /// The file the layout reader was opened over. + pub(super) file: ScanFile, + pub(super) projection: BoundExpression, + pub(super) filter: Option, + /// Whether the scan needs to return splits in the order they appear in the file. + pub(super) ordered: bool, + /// Optionally read a subset of the rows in the file. + pub(super) row_range: Option>, + /// The selection mask to apply to the selected row range. + pub(super) selection: Selection, + /// How to split the file for concurrent processing. + split_by: SplitBy, + /// Precomputed full-file natural split boundaries. + natural_splits: Option>, + /// The number of splits to make progress on concurrently **per-thread**. + pub(super) concurrency: usize, + /// Function to apply to each [`ArrayRef`] within the split tasks. + pub(super) map_fn: Arc VortexResult + Send + Sync>, + metrics_registry: Option>, + /// Maximal number of rows to read (after filtering) + pub(super) limit: Option, + /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, + /// but not by the scan [`Selection`] which remains relative. + pub(super) row_offset: u64, +} + +impl ScanBuilder { + /// Create a scan builder over `layout_reader`, opened over `file`, using `session` for + /// runtime and execution state. + pub fn new( + session: VortexSession, + layout_reader: Arc, + file: ScanFile, + ) -> Self { + Self::from_default(scan_builder::ScanBuilder::new(session, layout_reader), file) + } + + /// Returns an [`ArrayStream`] using the session's runtime handle. + /// + /// See [`ScanBuilder::into_stream`] for more details. + pub fn into_array_stream(self) -> VortexResult { + let dtype = self.dtype()?; + let stream = self.into_stream()?; + Ok(ArrayStreamAdapter::new(dtype, stream)) + } + + /// Returns an [`ArrayIterator`] using the given blocking runtime. + pub fn into_array_iter( + self, + runtime: &B, + ) -> VortexResult { + let stream = self.into_array_stream()?; + let dtype = stream.dtype().clone(); + Ok(ArrayIteratorAdapter::new( + dtype, + runtime.block_on_stream(stream), + )) + } +} + +impl ScanBuilder { + /// Copies every option of a default builder, for a scan of `file`, the file the builder's + /// reader was opened over. + pub fn from_default(builder: scan_builder::ScanBuilder, file: ScanFile) -> Self { + let parts = builder.into_parts(); + Self { + session: parts.session, + layout_reader: parts.layout_reader, + file, + projection: parts.projection, + filter: parts.filter, + ordered: parts.ordered, + row_range: parts.row_range, + selection: parts.selection, + split_by: parts.split_by, + natural_splits: parts.natural_splits, + concurrency: parts.concurrency, + map_fn: parts.map_fn, + metrics_registry: parts.metrics_registry, + limit: parts.limit.map(|limit| match limit { + ScanLimit::PerExecution(limit) => limit, + // Only the default executor's layout and multi-file scans share a budget. + ScanLimit::Shared(_) => { + vortex_panic!("The V2 scan does not support a limit shared between scans") + } + }), + row_offset: parts.row_offset, + } + } + + /// Add a filter expression bound against the reader dtype. + pub fn with_filter(mut self, filter: BoundExpression) -> Self { + self.filter = Some(filter); + self + } + + /// Add or clear a filter expression bound against the reader dtype. + pub fn with_some_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Set a projection expression bound against the reader dtype. + pub fn with_projection(mut self, projection: BoundExpression) -> Self { + self.projection = projection; + self + } + + /// Returns whether output chunks are yielded in file order. + pub fn ordered(&self) -> bool { + self.ordered + } + + /// Configure whether output chunks must be yielded in file order. + pub fn with_ordered(mut self, ordered: bool) -> Self { + self.ordered = ordered; + self + } + + /// Restrict scanning to a contiguous row range. + pub fn with_row_range(mut self, row_range: Range) -> Self { + self.row_range = Some(row_range); + self + } + + /// Apply a row selection to the selected row range. + pub fn with_selection(mut self, selection: Selection) -> Self { + self.selection = selection; + self + } + + /// Select rows by strictly sorted absolute indices relative to the scan input. + pub fn with_row_indices(mut self, row_indices: StrictSortedBuffer) -> Self { + self.selection = Selection::IncludeByIndex(row_indices); + self + } + + /// Set the root row offset used by row-index expressions. + pub fn with_row_offset(mut self, row_offset: u64) -> Self { + self.row_offset = row_offset; + self + } + + /// Configure how natural scan work is split for concurrency. + pub fn with_split_by(mut self, split_by: SplitBy) -> Self { + self.split_by = split_by; + self + } + + /// Supply precomputed full-file natural split boundaries (see + /// [`full_file_splits`](Self::full_file_splits)). Boundaries must be strictly increasing. + pub fn with_natural_splits(mut self, boundaries: Arc<[u64]>) -> Self { + debug_assert!( + boundaries.windows(2).all(|w| w[0] < w[1]), + "natural split boundaries must be strictly increasing" + ); + self.natural_splits = Some(boundaries); + self + } + + /// Compute the full-file natural split boundaries for the fields referenced by this scan's + /// projection and filter, ignoring any configured row range. + pub fn full_file_splits(&self) -> VortexResult> { + let field_mask = referenced_field_masks(&self.projection, self.filter.as_ref())?; + self.split_by.splits( + self.layout_reader.as_ref(), + &(0..self.layout_reader.row_count()), + &field_mask, + ) + } + + /// Returns the per-worker row-split concurrency. + pub fn concurrency(&self) -> usize { + self.concurrency + } + + /// The number of row splits to make progress on concurrently per-thread, must + /// be greater than 0. + pub fn with_concurrency(mut self, concurrency: usize) -> Self { + assert!(concurrency > 0); + self.concurrency = concurrency; + self + } + + /// Add or clear the metrics registry used by scan execution. + pub fn with_some_metrics_registry(mut self, metrics: Option>) -> Self { + self.metrics_registry = metrics; + self + } + + /// Set the metrics registry used by scan execution. + pub fn with_metrics_registry(mut self, metrics: Arc) -> Self { + self.metrics_registry = Some(metrics); + self + } + + /// Add or clear the maximum number of rows returned after filtering. + pub fn with_some_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + /// Set the maximum number of rows returned after filtering. + pub fn with_limit(mut self, limit: u64) -> Self { + self.limit = Some(limit); + self + } + + /// The [`DType`] returned by the scan, after applying the projection. + pub fn dtype(&self) -> VortexResult { + Ok(self.projection.dtype().clone()) + } + + /// The session used by the scan. + pub fn session(&self) -> &VortexSession { + &self.session + } + + /// Map each split of the scan. The function runs within the split task. + pub fn map( + self, + map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, + ) -> ScanBuilder { + let old_map_fn = self.map_fn; + ScanBuilder { + session: self.session, + layout_reader: self.layout_reader, + file: self.file, + projection: self.projection, + filter: self.filter, + ordered: self.ordered, + row_range: self.row_range, + selection: self.selection, + split_by: self.split_by, + natural_splits: self.natural_splits, + concurrency: self.concurrency, + metrics_registry: self.metrics_registry, + limit: self.limit, + row_offset: self.row_offset, + map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)), + } + } + + /// Plan and optimize the expressions over the file, compute split ranges, and return an + /// executable repeated scan. + pub fn prepare(self) -> VortexResult> { + prepare_scan(self) + } + + /// Constructs a task per filter split of the scan, returned as a vector of futures. + pub fn build(self) -> VortexResult>>>> { + // The ultimate short circuit + if self.limit.is_some_and(|l| l == 0) { + return Ok(vec![]); + } + + self.prepare()?.execute(None) + } + + /// Returns a [`Stream`] using the session's runtime handle. + /// + /// Parallel splits are spawned onto the runtime; a single split is polled directly. + pub fn into_stream( + self, + ) -> VortexResult> + Send + 'static + use> { + Ok(LazyScanStream::new(self)) + } + + /// Returns an [`Iterator`] using the session's runtime. + pub fn into_iter( + self, + runtime: &B, + ) -> VortexResult> + 'static> { + let stream = self.into_stream()?; + Ok(runtime.block_on_stream(stream)) + } +} diff --git a/vortex-layout/src/scan/v2/share.rs b/vortex-layout/src/scan/v2/share.rs new file mode 100644 index 00000000000..9293f09ee50 --- /dev/null +++ b/vortex-layout/src/scan/v2/share.rs @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::scalar_fn::fns::dynamic::DynamicComparison; +use vortex_error::VortexResult; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::plan::physical::Eval; +use crate::plan::physical::PlanByPtr; +use crate::plan::physical::PlanRef; +use crate::plan::physical::Share; +use crate::plan::physical::Take; +use crate::scan::v2::conjuncts::map_children; + +/// Removes each [`Share`] of dictionary values that no plan of the scan reads whole. +/// +/// Values read whole are canonicalized anyway, so an expression pushed onto the same values +/// reuses that work by running over the share, as the default scan's dictionary reader does once +/// it holds the canonical values. Without such a reader, the expression runs over the encoded +/// values instead, where kernels such as a compare over FSST avoid decompressing them. +pub(super) fn unshare_unread(plans: Vec) -> VortexResult> { + let mut read = HashSet::default(); + let mut has_share = false; + for plan in &plans { + has_share |= collect_read(plan, &mut read)?; + } + if !has_share { + return Ok(plans); + } + plans.into_iter().map(|plan| unshare(plan, &read)).collect() +} + +/// Records the shares that a take reads whole in `plan`, or evaluates a dynamic comparison over: +/// a take evaluates that again on every execution, so the values under it stay shared. +fn collect_read(plan: &PlanRef, read: &mut HashSet) -> VortexResult { + if let Some(take) = plan.as_opt::() { + let mut values = take.values()?; + if let Some(eval) = values.as_opt::() + && eval.expression().contains::()? + { + values = eval.child_plan()?; + } + if values.is::() { + read.insert(PlanByPtr(values)); + } + } + let mut has_share = plan.is::(); + for child in plan.children().iter_refs() { + has_share |= collect_read(child?, read)?; + } + Ok(has_share) +} + +fn unshare(plan: PlanRef, read: &HashSet) -> VortexResult { + if let Some(share) = plan.as_opt::() + && !read.contains(&PlanByPtr(plan.clone())) + { + return unshare(share.child_plan()?, read); + } + map_children(plan, |child| unshare(child, read)) +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::eq; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_session::registry::ReadContext; + + use super::unshare_unread; + use crate::LayoutRef; + use crate::layouts::dict::DictLayout; + use crate::layouts::flat::FlatLayout; + use crate::plan::physical::Eval; + use crate::plan::physical::PlanRef; + use crate::plan::physical::Share; + use crate::plan::physical::Take; + use crate::plan::physical::optimize; + use crate::plan::physical::plan_row_idx_expression; + use crate::scan::v2::lower::lower; + use crate::segments::SegmentId; + + fn dictionary() -> LayoutRef { + let flat = |row_count, dtype, segment: u32| { + FlatLayout::new( + row_count, + dtype, + SegmentId::from(segment), + ReadContext::new([]), + ) + .into_layout() + }; + DictLayout::new( + flat(2, DType::Utf8(NonNullable), 0), + flat(3, DType::Primitive(PType::U8, NonNullable), 1), + ) + .into_layout() + } + + fn plan(root_plan: &PlanRef, expression: Expression) -> VortexResult { + let expression = expression.bind(root_plan.dtype())?; + optimize(plan_row_idx_expression(expression, root_plan.clone())?) + } + + /// The values under the pushed predicate: the child of the `Eval` on the take's values. + fn predicate_input(filter: &PlanRef) -> VortexResult { + let values = filter + .as_opt::() + .ok_or_else(|| vortex_err!("filter is not a take: {filter}"))? + .values()?; + values + .as_opt::() + .ok_or_else(|| vortex_err!("predicate was not pushed onto the values: {values}"))? + .child_plan() + } + + #[test] + fn share_is_kept_where_values_are_read_whole() -> VortexResult<()> { + let root_plan = lower(&dictionary())?; + let plans = unshare_unread(vec![ + plan(&root_plan, root())?, + plan(&root_plan, eq(root(), lit("a")))?, + ])?; + + let read_whole = plans[0] + .as_opt::() + .ok_or_else(|| vortex_err!("projection is not a take"))? + .values()?; + assert!(read_whole.is::()); + let under_predicate = predicate_input(&plans[1])?; + assert!(PlanRef::ptr_eq(&read_whole, &under_predicate)); + Ok(()) + } + + #[test] + fn share_is_removed_where_values_are_not_read_whole() -> VortexResult<()> { + let root_plan = lower(&dictionary())?; + let plans = unshare_unread(vec![plan(&root_plan, eq(root(), lit("a")))?])?; + assert!(!predicate_input(&plans[0])?.is::()); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs new file mode 100644 index 00000000000..f9be150e354 --- /dev/null +++ b/vortex-layout/src/scan/v2/split.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::env; +use std::future; +use std::sync::Arc; +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::WorkScope; + +/// One filter split of a scan, ready to be admitted to a planning run. +/// +/// The split's planners announce, prune, filter, and project the rows of `scope`, and its morsels +/// emit one batch per projection split with selected rows, in whatever order their reads arrive. +pub struct SplitPlan { + /// The rows of the scan's root row domain the split covers. + pub scope: WorkScope, + /// The split's first planner. + pub root: Box, +} + +impl SplitPlan { + /// Runs the split on a planning run of its own, reading through a session of `io`, and + /// returns its batches in row order. + /// + /// The run advances inside this future, on whichever thread polls it, and the future awaits + /// its reads between steps. Nothing is handed to another thread. + pub(super) fn run( + mut self, + io: Arc, + ) -> impl Future>> + Send + 'static { + static EARLY_ANNOUNCE: LazyLock = LazyLock::new(|| { + !env::var("VORTEX_SCAN_EARLY_ANNOUNCE").is_ok_and(|value| value == "0") + }); + // Register neighboring ranges before any split starts fetching, so the driver can + // coalesce them. Keep the session alive for fetches and cancellation. Setting + // VORTEX_SCAN_EARLY_ANNOUNCE=0 restores late announcements for comparisons. + let prepared: VortexResult>> = (|| { + if !*EARLY_ANNOUNCE { + return Ok(None); + } + let source = io.session(); + match self.root.compute()? { + PlannerOutput::NeedsIO(batch) => { + vortex_ensure!( + batch + .iter() + .all(|request| request.intent == IoIntent::Announce), + "early announcement encountered an active read" + ); + source.submit(IoOwnerId(0), batch)?; + } + PlannerOutput::Planner(scope, root) => { + self.scope = scope; + self.root = root; + } + output => vortex_bail!("unexpected initial announcement output: {output:?}"), + } + Ok(Some(source)) + })(); + async move { + let source = prepared?.unwrap_or_else(|| io.session()); + let mut run = Run::new(); + run.admit(self.root, self.scope, source); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + // Morsels finish in whatever order their reads arrive. + batches.sort_by_key(|batch| batch.scope.rows.start); + Ok(batches.into_iter().map(|batch| batch.array).collect()) + } + } +} diff --git a/vortex-layout/src/scan/v2/splits.rs b/vortex-layout/src/scan/v2/splits.rs new file mode 100644 index 00000000000..366d21186a4 --- /dev/null +++ b/vortex-layout/src/scan/v2/splits.rs @@ -0,0 +1,343 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splits a scan from its plans rather than from the layout reader. +//! +//! A scan splits twice. The filter runs over filter splits, cut where the chunks of the columns the +//! filter reads start, or of those the projection reads when the scan has no filter, and cut +//! further at other columns' chunk starts when too large to give every thread several splits. Under each filter split, the projection runs +//! over one or more projection splits, cut where the chunks of the columns the projection reads +//! start, so each projection split reads at most one chunk of every column. + +use std::iter; +use std::ops::Range; + +use itertools::Itertools; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::LayoutRef; +use crate::layouts::chunked::Chunked; +use crate::layouts::dict::Dict; +use crate::layouts::struct_::Struct; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; +use crate::plan::physical::Concat; +use crate::plan::physical::Eval; +use crate::plan::physical::Filter; +use crate::plan::physical::Pack; +use crate::plan::physical::PlanRef; +use crate::plan::physical::Take; + +/// Bounds the selected output retained by a sparse projection graph. +pub(crate) const MAX_SPARSE_ROWS: usize = 1024; + +/// The most rows a filter split has, which amortises a split's fixed cost over a large scan. +const MAX_SPLIT_ROWS: u64 = 1 << 16; + +/// The fewest rows the cap on a filter split goes down to for a small scan. +const MIN_SPLIT_ROWS: u64 = 1 << 13; + +/// How many filter splits a scan aims to give each thread, so threads that finish early find more +/// work. +const SPLITS_PER_THREAD: u64 = 4; + +/// The most rows a filter split of a scan over `rows` rows has when `threads` threads run it. +pub(super) fn max_split_rows(rows: u64, threads: usize) -> u64 { + let splits = SPLITS_PER_THREAD * u64::try_from(threads.max(1)).unwrap_or(u64::MAX); + (rows / splits).clamp(MIN_SPLIT_ROWS, MAX_SPLIT_ROWS) +} + +/// Identity scans retain natural chunk boundaries without materializing the full physical plan. +/// Dictionary values and list elements are separate row domains, as in [`chunk_starts`]. +/// Sparse indices restrict the walk to selected chunks; their end boundaries separate empty gaps. +pub(super) fn layout_chunk_starts( + layout: &LayoutRef, + indices: Option<&[u64]>, +) -> VortexResult> { + let mut starts = Vec::new(); + collect_layout_starts(layout, 0, indices, &mut starts)?; + starts.sort_unstable(); + starts.dedup(); + Ok(starts) +} + +fn collect_layout_starts( + layout: &LayoutRef, + offset: u64, + indices: Option<&[u64]>, + starts: &mut Vec, +) -> VortexResult<()> { + if let Some(chunked) = layout.as_opt::() { + if let Some(mut indices) = indices { + let offsets = chunked.chunk_offsets(); + while let Some(&row) = indices.first() { + let index = offsets + .partition_point(|&start| start <= row.saturating_sub(offset)) + .saturating_sub(1); + if index >= chunked.nchildren() { + break; + } + let start = offset + .checked_add(offsets[index]) + .ok_or_else(|| vortex_err!("Chunked row offset overflow"))?; + let end = offset + .checked_add(offsets[index + 1]) + .ok_or_else(|| vortex_err!("Chunked row offset overflow"))?; + let count = indices.partition_point(|&row| row < end); + // Both edges are needed: the gap before the next selected chunk must remain + // an empty projection split, so IO planning does not announce its segments. + starts.extend([start, end]); + if !chunked.children().child_is_indivisible(index) { + let child = layout + .slot(index)? + .ok_or_else(|| vortex_err!("Missing chunk {index}"))?; + collect_layout_starts(&child, start, Some(&indices[..count]), starts)?; + } + indices = &indices[count..]; + } + return Ok(()); + } + for (index, &child_offset) in chunked.chunk_offsets()[..chunked.nchildren()] + .iter() + .enumerate() + { + let offset = offset + .checked_add(child_offset) + .ok_or_else(|| vortex_err!("Chunked row offset overflow"))?; + starts.push(offset); + if !chunked.children().child_is_indivisible(index) { + let child = layout + .slot(index)? + .ok_or_else(|| vortex_err!("Missing chunk {index}"))?; + collect_layout_starts(&child, offset, None, starts)?; + } + } + } else if layout.is::() || layout.is::() || layout.is::() { + let slot = usize::from(layout.is::()); + let child = layout + .slot(slot)? + .ok_or_else(|| vortex_err!("Missing layout child {slot}"))?; + collect_layout_starts(&child, offset, indices, starts)?; + } else if layout.is::() { + for child in layout.children()? { + collect_layout_starts(&child, offset, indices, starts)?; + } + } + Ok(()) +} + +/// The row positions where the chunks of `plans` start, over their shared row domain, sorted and +/// without duplicates. +/// +/// This follows plans whose children cover the same rows, or known parts of them, and stops at any +/// other plan: a take's values or a list's elements cover other rows. +pub(super) fn chunk_starts<'a>( + plans: impl IntoIterator, +) -> VortexResult> { + let mut starts = Vec::new(); + for plan in plans { + collect_starts(plan, 0, &mut starts)?; + } + starts.sort_unstable(); + starts.dedup(); + Ok(starts) +} + +fn collect_starts(plan: &PlanRef, offset: u64, starts: &mut Vec) -> VortexResult<()> { + if let Some(concat) = plan.as_opt::() { + for (index, &child_offset) in concat.row_offsets().iter().enumerate() { + starts.push(offset + child_offset); + collect_starts( + plan.child_ref_required(index)?, + offset + child_offset, + starts, + )?; + } + } else if plan.is::() { + collect_starts(plan.child_ref_required(0)?, offset, starts)?; + } else if plan.is::() || plan.is::() || plan.is::() { + for child in plan.children().iter_refs() { + collect_starts(child?, offset, starts)?; + } + } + Ok(()) +} + +/// The filter splits of `rows`, as every cut including `rows.start` and `rows.end`. +/// +/// Rows are cut at `filter_starts`, the chunk starts of the columns the filter reads. A split +/// longer than `max_rows` is cut further at `finer_starts`, the chunk starts of every column the +/// scan reads, so a scan with coarse filter chunks still gives every thread several splits. No +/// split is cut inside a chunk, which every split reading it would decode again. A cut closer than +/// a quarter of `max_rows` to its neighbour is dropped, so misaligned chunks of different columns +/// do not leave slivers. +pub(super) fn filter_split_boundaries( + filter_starts: &[u64], + finer_starts: &[u64], + rows: Range, + max_rows: u64, +) -> Vec { + if rows.is_empty() { + return Vec::new(); + } + let min_rows = max_rows / 4; + let cuts = cut_at(filter_starts, rows.clone(), min_rows); + let mut boundaries = vec![rows.start]; + for (start, end) in cuts.into_iter().tuple_windows() { + if end - start > max_rows { + let finer = cut_at(finer_starts, start..end, min_rows); + boundaries.extend(&finer[1..]); + } else { + boundaries.push(end); + } + } + boundaries +} + +/// Coarsens sparse identity scans only when their selected rows would otherwise occupy many +/// small tasks. Keep two active splits per worker, with a minimum budget of eight, so small +/// scans do not become one large nested assembly on machines with few workers. +pub(super) fn sparse_filter_split_boundaries( + filter_starts: &[u64], + finer_starts: &[u64], + row_count: u64, + indices: &[u64], + threads: usize, +) -> Vec { + let threads = threads.max(1); + let boundaries = filter_split_boundaries( + filter_starts, + finer_starts, + 0..row_count, + max_split_rows(row_count, threads), + ); + let active = indices + .iter() + .map(|row| boundaries.partition_point(|boundary| boundary <= row)) + .dedup() + .count(); + if active <= threads.saturating_mul(2).max(8) { + return boundaries; + } + // The minimum split length is a quarter of max_rows. Aim for roughly one task per worker; + // retaining chunk boundaries avoids introducing additional partial-chunk decodes. + filter_split_boundaries( + filter_starts, + finer_starts, + 0..row_count, + row_count.div_ceil(threads as u64).saturating_mul(4), + ) +} + +/// `rows` cut at every one of `starts` inside it that is at least `min_rows` from the cut before +/// and from `rows.end`, as every cut including `rows.start` and `rows.end`. +fn cut_at(starts: &[u64], rows: Range, min_rows: u64) -> Vec { + let mut cuts = vec![rows.start]; + for &start in starts { + let previous = cuts[cuts.len() - 1]; + if rows.start < start && start < rows.end && start - previous >= min_rows { + cuts.push(start); + } + } + // The last split may be a sliver; fold it into the one before. + if cuts.len() > 1 && rows.end - cuts[cuts.len() - 1] < min_rows { + cuts.pop(); + } + cuts.push(rows.end); + cuts +} + +/// The projection splits of `rows`: cut at every one of `starts` inside it. +pub(crate) fn projection_splits(starts: &[u64], rows: Range) -> Vec> { + let first = starts.partition_point(|&start| start <= rows.start); + let last = starts.partition_point(|&start| start < rows.end); + iter::once(rows.start) + .chain(starts[first..last].iter().copied()) + .chain(iter::once(rows.end)) + .tuple_windows() + .map(|(start, end)| start..end) + .filter(|range| !range.is_empty()) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::few_active(8, 2, 32)] + #[case::two_per_worker(8, 16, 32)] + #[case::many_active(8, 32, 8)] + #[case::few_single_worker(1, 8, 32)] + #[case::many_single_worker(1, 32, 1)] + fn sparse_splits_preserve_parallelism( + #[case] threads: usize, + #[case] active: u64, + #[case] expected: usize, + ) { + let starts = (0..32).map(|i| i * 32_768).collect::>(); + let indices = (0..active).map(|i| i * 32_768 + 1).collect::>(); + let boundaries = + sparse_filter_split_boundaries(&starts, &starts, 1_048_576, &indices, threads); + assert_eq!(boundaries.len() - 1, expected); + } + + #[rstest] + fn sparse_splits_do_not_merge_an_already_parallel_short_tail(#[values(1, 8)] threads: usize) { + let starts = (0..8).map(|i| i * 131_072).collect::>(); + let indices = starts.iter().map(|i| i + 1).collect::>(); + let mut expected = starts.clone(); + expected.push(1_000_000); + assert_eq!( + sparse_filter_split_boundaries(&starts, &starts, 1_000_000, &indices, threads), + expected, + ); + } + + #[rstest] + #[case::whole_when_no_starts(&[], &[], 0..150_000, 65_536, vec![0, 150_000])] + #[case::at_starts(&[40_000, 80_000], &[], 0..120_000, 65_536, vec![0, 40_000, 80_000, 120_000])] + #[case::drops_slivers(&[40_000, 41_000, 80_000], &[], 0..120_000, 65_536, vec![0, 40_000, 80_000, 120_000])] + #[case::folds_short_tail(&[40_000, 115_000], &[], 0..120_000, 131_072, vec![0, 40_000, 120_000])] + #[case::long_split_cut_at_finer_starts(&[40_000], &[40_000, 60_000, 80_000, 100_000], 0..120_000, 65_536, vec![0, 40_000, 60_000, 80_000, 100_000, 120_000])] + #[case::short_split_ignores_finer_starts(&[60_000], &[30_000, 60_000, 90_000], 0..120_000, 65_536, vec![0, 60_000, 120_000])] + #[case::never_inside_a_chunk(&[], &[], 0..1_000_000, 65_536, vec![0, 1_000_000])] + #[case::only_inside_rows(&[10_000, 50_000, 200_000], &[], 20_000..100_000, 65_536, vec![20_000, 50_000, 100_000])] + #[case::empty(&[10], &[], 5..5, 65_536, vec![])] + fn filter_splits( + #[case] filter_starts: &[u64], + #[case] finer_starts: &[u64], + #[case] rows: Range, + #[case] max_rows: u64, + #[case] expected: Vec, + ) { + assert_eq!( + filter_split_boundaries(filter_starts, finer_starts, rows, max_rows), + expected + ); + } + + #[rstest] + #[case::small_scan_floors(150_000, 14, 8_192)] + #[case::medium_scan(1_500_000, 14, 26_785)] + #[case::large_scan_caps(6_000_000, 14, 65_536)] + #[case::no_threads(100_000, 0, 25_000)] + fn split_rows(#[case] rows: u64, #[case] threads: usize, #[case] expected: u64) { + assert_eq!(max_split_rows(rows, threads), expected); + } + + #[rstest] + #[case::whole(&[], 10..20, vec![10..20])] + #[case::cut(&[0, 12, 15, 30], 10..20, vec![10..12, 12..15, 15..20])] + #[case::start_on_boundary(&[10, 15], 10..20, vec![10..15, 15..20])] + fn projection( + #[case] starts: &[u64], + #[case] rows: Range, + #[case] expected: Vec>, + ) { + assert_eq!(projection_splits(starts, rows), expected); + } +} diff --git a/vortex-layout/src/scan/v2/stream.rs b/vortex-layout/src/scan/v2/stream.rs new file mode 100644 index 00000000000..5a339eea934 --- /dev/null +++ b/vortex-layout/src/scan/v2/stream.rs @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; +use std::task::ready; + +use futures::Stream; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::future::BoxFuture; +use futures::future::Either; +use futures::stream::BoxStream; +use vortex_error::VortexError; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_io::runtime::Handle; +use vortex_io::runtime::Task; +use vortex_io::session::RuntimeSessionExt; +use vortex_utils::parallelism::get_available_parallelism; + +use crate::scan::scan_builder; +use crate::scan::v2::ScanBuilder; +use crate::scan::v2::ScanFile; +use crate::scan::v2::pruning::file_pruning_enabled; + +/// Returns a stream that prepares the scan over `file` on first poll and runs its split tasks. +/// +/// The replacement for the default +/// [`ScanBuilder::into_stream`](scan_builder::ScanBuilder::into_stream). Equivalent to copying +/// `builder` into a [`ScanBuilder`] and streaming that. +pub fn into_stream( + builder: scan_builder::ScanBuilder, + file: ScanFile, +) -> VortexResult> + Send + 'static + use> { + ScanBuilder::from_default(builder, file).into_stream() +} + +type Tasks = Vec>>>; + +enum State { + Builder(Option>>), + Preparing { + ordered: bool, + concurrency: usize, + handle: Handle, + task: Task>>, + }, + Stream(BoxStream<'static, VortexResult>), + Error(Option), +} + +/// Prepares a scan on first poll, then runs its split tasks and yields their batches. +pub(super) struct LazyScanStream { + state: State, +} + +impl LazyScanStream { + pub(super) fn new(builder: ScanBuilder) -> Self { + Self { + state: State::Builder(Some(Box::new(builder))), + } + } +} + +impl Unpin for LazyScanStream {} + +impl Stream for LazyScanStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match &mut self.state { + State::Builder(builder) => { + let builder = builder.take().vortex_expect("polled after completion"); + let ordered = builder.ordered(); + let num_workers = get_available_parallelism().unwrap_or(1); + let concurrency = builder.concurrency() * num_workers; + let handle = builder.session().handle(); + let task = if file_pruning_enabled() { + let prepared = handle.spawn_cpu(move || builder.prepare()); + handle.spawn( + async move { prepared.await?.execute_batches_pruned(None).await }, + ) + } else { + handle.spawn_cpu(move || { + builder + .prepare() + .and_then(|scan| scan.execute_batches(None)) + }) + }; + self.state = State::Preparing { + ordered, + concurrency, + handle, + task, + }; + } + State::Preparing { + ordered, + concurrency, + handle, + task, + } => match ready!(Pin::new(task).poll(cx)) { + Ok(tasks) => { + let concurrency = *concurrency; + let handle = handle.clone(); + // A single split already runs on the polling worker. Spawning it adds + // a scheduler handoff without exposing any more parallelism. + let spawn = tasks.len() > 1 && concurrency > 1; + let stream = futures::stream::iter(tasks).map(move |task| { + if spawn { + Either::Right(handle.spawn(task)) + } else { + Either::Left(task) + } + }); + let stream = if *ordered { + stream.buffered(concurrency).boxed() + } else { + stream.buffer_unordered(concurrency).boxed() + }; + // A task returns one batch per projection split of its filter split. + self.state = State::Stream( + stream + .map_ok(|batches| { + futures::stream::iter(batches.into_iter().map(Ok)) + }) + .try_flatten() + .boxed(), + ); + } + Err(err) => self.state = State::Error(Some(err)), + }, + State::Stream(stream) => return stream.as_mut().poll_next(cx), + State::Error(err) => return Poll::Ready(err.take().map(Err)), + } + } + } +} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs new file mode 100644 index 00000000000..f9792259dc5 --- /dev/null +++ b/vortex-layout/src/scan/v2/tests.rs @@ -0,0 +1,1401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[path = "legacy_tests.rs"] +mod legacy; +#[path = "list_tests.rs"] +mod lists; + +use std::collections::BTreeSet; +use std::num::NonZeroUsize; +use std::ops::Range; +use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::thread; + +use futures::FutureExt; +use futures::TryStreamExt; +use futures::stream; +use parking_lot::Mutex; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::dict::dict_encode; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; +use vortex_array::expr::and; +use vortex_array::expr::byte_length; +use vortex_array::expr::cast; +use vortex_array::expr::dynamic; +use vortex_array::expr::eq; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::like; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::not_eq; +use vortex_array::expr::or; +use vortex_array::expr::root; +use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_session::VortexSession; + +use crate::LayoutChildren; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::chunked::Chunked; +use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; +use crate::layouts::dict::Dict; +use crate::layouts::dict::writer::DictLayoutOptions; +use crate::layouts::dict::writer::DictStrategy; +use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::row_idx::row_idx; +use crate::layouts::struct_::StructLayout; +use crate::layouts::zoned::writer::ZonedLayoutOptions; +use crate::layouts::zoned::writer::ZonedStrategy; +use crate::plan::physical::PlanRef; +use crate::scan::planning::SegmentLocation; +use crate::scan::scan_builder::ScanBuilder; +use crate::scan::v2; +use crate::scan::v2::FilePlans; +use crate::scan::v2::ScanFile; +use crate::scan::v2::file::shared_file; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt as _; +use crate::test::new_session; + +const CHUNK_ROWS: i32 = 1000; +const DTYPE: DType = DType::Primitive(PType::I32, NonNullable); + +#[derive(Clone)] +struct CountingChildren { + inner: Arc, + accessed: Arc, + row_counts: Arc, + hints: Arc, +} + +impl LayoutChildren for CountingChildren { + fn to_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn child(&self, index: usize, dtype: &DType) -> VortexResult { + self.accessed.fetch_add(1, Ordering::Relaxed); + self.inner.child(index, dtype) + } + + fn child_row_count(&self, index: usize) -> u64 { + self.row_counts.fetch_add(1, Ordering::Relaxed); + self.inner.child_row_count(index) + } + + fn nchildren(&self) -> usize { + self.inner.nchildren() + } + + fn child_is_indivisible(&self, index: usize) -> bool { + self.hints.fetch_add(1, Ordering::Relaxed); + self.inner.child_is_indivisible(index) + } +} + +#[rstest] +#[case::empty(&[], 0)] +#[case::sparse(&[1, 2001], 2)] +#[case::same_chunk(&[1001, 1002], 1)] +#[case::last_row(&[3999], 1)] +#[tokio::test] +async fn identity_preparation_keeps_flat_chunks_lazy( + #[case] indices: &'static [u64], + #[case] selected_chunks: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let mut expected_segments = BTreeSet::new(); + for &index in indices { + let child = layout + .slot(usize::try_from(index / CHUNK_ROWS as u64)?)? + .ok_or_else(|| vortex_err!("Missing selected chunk"))?; + expected_segments.extend(segment_ids(&child)?); + } + let registered = Arc::default(); + let segments: Arc = Arc::new(RecordingSegments { + inner: segments, + registered: Arc::clone(®istered), + reads: Arc::default(), + }); + let accessed = Arc::new(AtomicUsize::new(0)); + let row_counts = Arc::new(AtomicUsize::new(0)); + let hints = Arc::new(AtomicUsize::new(0)); + let layout = ChunkedLayout::new( + layout.row_count(), + DTYPE, + Arc::new(CountingChildren { + inner: Arc::clone(layout.as_::().children()), + accessed: Arc::clone(&accessed), + row_counts: Arc::clone(&row_counts), + hints: Arc::clone(&hints), + }), + ) + .into_layout(); + let file = scan_file(&segments, &layout)?; + accessed.store(0, Ordering::Relaxed); + row_counts.store(0, Ordering::Relaxed); + hints.store(0, Ordering::Relaxed); + let case = Case { + rows: Some(indices), + ..Default::default() + }; + let scan = v2::prepare(builder(&session, &segments, &layout, &case)?, file)?; + assert_eq!(accessed.load(Ordering::Relaxed), 0); + assert_eq!(row_counts.load(Ordering::Relaxed), 0); + assert_eq!(hints.load(Ordering::Relaxed), selected_chunks); + + let actual = await_tasks(DTYPE, scan.execute(None)?).await?; + assert_eq!(*registered.lock(), expected_segments); + let expected = indices + .iter() + .copied() + .map(i32::try_from) + .collect::, _>>()?; + assert_arrays_eq!( + actual, + Buffer::from(expected).into_array(), + &mut session.create_execution_ctx() + ); + Ok(()) +} + +/// Four chunks of consecutive integers, `0..4000`. +async fn write_layout( + session: &VortexSession, +) -> VortexResult<(Arc, LayoutRef)> { + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunks = (0..4) + .map(|chunk| { + let values = Buffer::from_iter(chunk * CHUNK_ROWS..(chunk + 1) * CHUNK_ROWS); + Ok((sequence_id.advance(), values.into_array())) + }) + .collect::>(); + let layout = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(DTYPE, stream::iter(chunks)).sendable(), + eof, + session, + ) + .await?; + Ok((segments, layout)) +} + +/// A [`ScanFile`] over test segments. They have no byte offsets, and the scan fetches them by id, +/// so each location is a placeholder whose offset is its id, with the alignment the test segments +/// already have. +fn scan_file(segments: &Arc, layout: &LayoutRef) -> VortexResult { + let mut count = 0; + for layout in layout.depth_first_traversal() { + for id in layout?.segment_ids() { + count = count.max(*id as usize + 1); + } + } + let locations = (0..count as u64) + .map(|offset| SegmentLocation { + offset, + length: 0, + alignment: Alignment::none(), + }) + .collect(); + Ok(ScanFile { + layout: Arc::clone(layout), + locations, + segments: Arc::clone(segments), + io: None, + plans: None, + }) +} + +#[derive(Clone, Default)] +struct Case { + filter: bool, + row_range: Option>, + limit: Option, + rows: Option<&'static [u64]>, + row_idx: bool, + /// Adds a second conjunct to the filter, `$ < below`. + below: Option, +} + +fn builder( + session: &VortexSession, + segments: &Arc, + layout: &LayoutRef, + case: &Case, +) -> VortexResult> { + let reader = layout.new_reader( + "".into(), + Arc::clone(segments), + session, + &Default::default(), + )?; + let mut builder = ScanBuilder::new(session.clone(), reader); + if case.filter { + let filter = gt(root(), lit(1500_i32)); + let filter = match case.below { + Some(below) => and(filter, lt(root(), lit(below))), + None => filter, + }; + builder = builder.with_filter(filter.bind(&DTYPE)?); + } + if let Some(row_range) = case.row_range.clone() { + builder = builder.with_row_range(row_range); + } + if let Some(limit) = case.limit { + builder = builder.with_limit(limit); + } + if let Some(rows) = case.rows { + builder = builder.with_row_indices(StrictSortedBuffer::try_new(Buffer::copy_from(rows))?); + } + if case.row_idx { + builder = builder + .with_projection(row_idx().bind(&DTYPE)?) + .with_row_offset(10_000); + } + Ok(builder) +} + +async fn await_tasks( + dtype: DType, + tasks: Vec>>>, +) -> VortexResult { + let mut chunks = Vec::new(); + for task in tasks { + if let Some(chunk) = task.await? { + chunks.push(chunk); + } + } + Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) +} + +fn case(filter: bool, row_range: Option>, limit: Option) -> Case { + Case { + filter, + row_range, + limit, + ..Case::default() + } +} + +fn rows(filter: bool, rows: &'static [u64]) -> Case { + Case { + filter, + rows: Some(rows), + ..Case::default() + } +} + +fn conjuncts(below: i32) -> Case { + Case { + filter: true, + below: Some(below), + ..Case::default() + } +} + +fn row_idx_case(filter: bool) -> Case { + Case { + filter, + row_idx: true, + ..Case::default() + } +} + +#[rstest] +#[case::everything(case(false, None, None))] +#[case::filter(case(true, None, None))] +#[case::row_range(case(false, Some(500..2500), None))] +#[case::filter_and_row_range(case(true, Some(500..2500), None))] +#[case::limit(case(false, None, Some(1200)))] +#[case::limit_and_row_range(case(false, Some(900..3100), Some(1500)))] +#[case::selection(rows(false, &[0, 5, 999, 1000, 2500, 3999]))] +#[case::selection_and_filter(rows(true, &[0, 1499, 1501, 2500, 3999]))] +#[case::row_idx(row_idx_case(false))] +#[case::row_idx_and_filter(row_idx_case(true))] +#[case::two_conjuncts(conjuncts(3000))] +#[case::conjuncts_matching_nothing(conjuncts(1000))] +#[tokio::test(flavor = "multi_thread")] +async fn stream_matches_default(#[case] case: Case) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let dtype = builder(&session, &segments, &layout, &case)?.dtype()?; + + let expected = builder(&session, &segments, &layout, &case)? + .into_stream()? + .try_collect::>() + .await?; + let actual = v2::into_stream( + builder(&session, &segments, &layout, &case)?, + scan_file(&segments, &layout)?, + )? + .try_collect::>() + .await?; + + assert_arrays_eq!( + ChunkedArray::try_new(actual, dtype.clone())?, + ChunkedArray::try_new(expected, dtype)?, + &mut session.create_execution_ctx() + ); + Ok(()) +} + +#[rstest] +#[case::single_split(250..750)] +#[case::empty_range(250..250)] +#[tokio::test(flavor = "multi_thread")] +async fn stream_without_parallel_splits( + #[case] row_range: Range, + #[values(true, false)] ordered: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let case = case(false, Some(row_range), None); + let dtype = builder(&session, &segments, &layout, &case)?.dtype()?; + let expected = builder(&session, &segments, &layout, &case)? + .with_ordered(ordered) + .into_stream()? + .try_collect::>() + .await?; + let actual = v2::into_stream( + builder(&session, &segments, &layout, &case)?.with_ordered(ordered), + scan_file(&segments, &layout)?, + )? + .try_collect::>() + .await?; + assert_arrays_eq!( + ChunkedArray::try_new(actual, dtype.clone())?, + ChunkedArray::try_new(expected, dtype)?, + &mut session.create_execution_ctx() + ); + Ok(()) +} + +/// The executor's own builder, configured directly rather than copied from a default builder, +/// returns what the default scan returns. +#[rstest] +#[case::everything(None, None)] +#[case::filter_and_row_range(Some(gt(root(), lit(1500_i32))), Some(500..2500))] +#[tokio::test(flavor = "multi_thread")] +async fn own_builder_matches_default( + #[case] filter: Option, + #[case] row_range: Option>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let reader = || { + layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + ) + }; + let filter = filter.map(|filter| filter.bind(&DTYPE)).transpose()?; + + let mut default = ScanBuilder::new(session.clone(), reader()?).with_some_filter(filter.clone()); + let mut own = v2::ScanBuilder::new(session.clone(), reader()?, scan_file(&segments, &layout)?) + .with_some_filter(filter); + if let Some(row_range) = row_range { + default = default.with_row_range(row_range.clone()); + own = own.with_row_range(row_range); + } + assert_eq!(own.dtype()?, default.dtype()?); + + let expected = default.into_stream()?.try_collect::>().await?; + let actual = own.into_stream()?.try_collect::>().await?; + assert_arrays_eq!( + ChunkedArray::try_new(actual, DTYPE)?, + ChunkedArray::try_new(expected, DTYPE)?, + &mut session.create_execution_ctx() + ); + Ok(()) +} + +#[rstest] +#[case::whole_scan(case(true, None, None), None)] +#[case::execute_range(case(true, None, None), Some(700..3300))] +#[case::both_ranges(case(false, Some(500..2500), None), Some(2000..4000))] +#[case::empty_range(case(true, None, None), Some(1200..1200))] +#[tokio::test(flavor = "multi_thread")] +async fn execute_matches_default( + #[case] case: Case, + #[case] execute_range: Option>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + + let default = builder(&session, &segments, &layout, &case)?.prepare()?; + let replacement = v2::prepare( + builder(&session, &segments, &layout, &case)?, + scan_file(&segments, &layout)?, + )?; + assert_eq!(replacement.dtype(), default.dtype()); + let dtype = default.dtype().clone(); + + let expected = await_tasks(dtype.clone(), default.execute(execute_range.clone())?).await?; + let actual = await_tasks(dtype, replacement.execute(execute_range)?).await?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn filter_with_limit_is_rejected() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let case = case(true, None, Some(10)); + let file = scan_file(&segments, &layout)?; + assert!(v2::prepare(builder(&session, &segments, &layout, &case)?, file).is_err()); + Ok(()) +} + +/// Records every segment whose read is awaited. +/// +/// A scan also registers segments it may read, so the source can coalesce them. Registration +/// alone reads nothing, so only awaited reads are recorded. +struct RecordingSegments { + inner: Arc, + registered: Arc>>, + reads: Arc>>, +} + +impl SegmentSource for RecordingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.registered.lock().insert(*id); + let read = self.inner.request(id); + let reads = Arc::clone(&self.reads); + async move { + reads.lock().insert(*id); + read.await + } + .boxed() + } +} + +struct CountedSegments { + inner: Arc, + reads: Arc>>, +} + +impl SegmentSource for CountedSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + let read = self.inner.request(id); + let reads = Arc::clone(&self.reads); + async move { + reads.lock().push(id); + read.await + } + .boxed() + } +} + +#[rstest] +#[case::whole(None, None, &[0, 1000, 2000, 3999], 5)] +#[case::range(Some(500..3500), None, &[1000, 2000], 3)] +#[case::limit(None, Some(2), &[0, 1000], 3)] +#[tokio::test] +async fn sparse_projection_reads_each_column_chunk_once( + #[case] row_range: Option>, + #[case] limit: Option, + #[case] expected_indices: &[usize], + #[case] expected_reads: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let written = Arc::new(TestSegments::default()); + let values = PrimitiveArray::from_option_iter((0..4000_i32).map(|i| (i % 7 != 0).then_some(i))) + .into_array(); + let expected = + StructArray::from_fields(&[("a", values.clone()), ("b", values.clone())])?.into_array(); + let mut fields = Vec::new(); + for chunk_rows in [1000, 4000] { + let (mut sequence, eof) = SequenceId::root().split(); + let chunks = (0..4000) + .step_by(chunk_rows) + .map(|start| Ok((sequence.advance(), values.slice(start..start + chunk_rows)?))) + .collect::>(); + fields.push( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&written), + SequentialStreamAdapter::new(values.dtype().clone(), stream::iter(chunks)) + .sendable(), + eof, + &session, + ) + .await?, + ); + } + let layout = StructLayout::new(4000, expected.dtype().clone(), fields).into_layout(); + let reads = Arc::default(); + let segments: Arc = Arc::new(CountedSegments { + inner: written, + reads: Arc::clone(&reads), + }); + let case = Case { + rows: Some(&[0, 1000, 2000, 3999]), + row_range, + limit, + ..Default::default() + }; + let scan = v2::prepare( + builder(&session, &segments, &layout, &case)?, + scan_file(&segments, &layout)?, + )?; + let expected = expected.filter(Mask::from_indices(4000, expected_indices.iter().copied()))?; + for _ in 0..2 { + reads.lock().clear(); + let actual = await_tasks(expected.dtype().clone(), scan.execute(None)?).await?; + assert_eq!(reads.lock().len(), expected_reads); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} + +#[rstest] +#[case::projection(false)] +#[case::filter_and_projection(true)] +#[tokio::test] +async fn ordinary_segment_reads_match_v1_on_each_execution( + #[case] filter: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let reads = Arc::default(); + let segments: Arc = Arc::new(CountedSegments { + inner: segments, + reads: Arc::clone(&reads), + }); + let case = case(filter, Some(1000..2000), None); + let default = builder(&session, &segments, &layout, &case)?.prepare()?; + let replacement = v2::prepare( + builder(&session, &segments, &layout, &case)?, + scan_file(&segments, &layout)?, + )?; + for _ in 0..2 { + reads.lock().clear(); + let expected = await_tasks(DTYPE, default.execute(None)?).await?; + let mut expected_reads = std::mem::take(&mut *reads.lock()); + expected_reads.sort_unstable(); + assert_eq!(expected_reads.len(), if filter { 2 } else { 1 }); + + let actual = await_tasks(DTYPE, replacement.execute(None)?).await?; + let mut actual_reads = std::mem::take(&mut *reads.lock()); + actual_reads.sort_unstable(); + assert_eq!(actual_reads, expected_reads); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} + +/// Building tasks announces all their ranges without polling a single read. +#[tokio::test] +async fn tasks_announce_before_they_run() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let registered = Arc::new(Mutex::new(BTreeSet::new())); + let reads = Arc::new(Mutex::new(BTreeSet::new())); + let segments: Arc = Arc::new(RecordingSegments { + inner: segments, + registered: Arc::clone(®istered), + reads: Arc::clone(&reads), + }); + let scan = v2::prepare( + builder(&session, &segments, &layout, &Case::default())?, + scan_file(&segments, &layout)?, + )?; + let tasks = scan.execute(None)?; + assert_eq!(*registered.lock(), segment_ids(&layout)?); + assert!(reads.lock().is_empty()); + let actual = await_tasks(DTYPE, tasks).await?; + assert_arrays_eq!( + actual, + Buffer::from_iter(0..4 * CHUNK_ROWS).into_array(), + &mut session.create_execution_ctx() + ); + Ok(()) +} + +/// The ids of the segments below `layout`. +fn segment_ids(layout: &LayoutRef) -> VortexResult> { + let mut ids = BTreeSet::new(); + for layout in layout.depth_first_traversal() { + ids.extend(layout?.segment_ids().into_iter().map(|id| *id)); + } + Ok(ids) +} + +/// Four 1000-row chunks of `0..4000`, one zone each. +async fn write_zoned_layout( + session: &VortexSession, +) -> VortexResult<(Arc, LayoutRef)> { + let written = Arc::new(TestSegments::default()); + let segments: Arc = Arc::clone(&written) as _; + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunks = (0..4) + .map(|chunk| { + let values = Buffer::from_iter(chunk * CHUNK_ROWS..(chunk + 1) * CHUNK_ROWS); + Ok((sequence_id.advance(), values.into_array())) + }) + .collect::>(); + let layout = ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(CHUNK_ROWS as usize).unwrap_or(NonZeroUsize::MIN), + ..Default::default() + }, + ) + .write_stream( + ArrayContext::empty().into(), + written, + SequentialStreamAdapter::new(DTYPE, stream::iter(chunks)).sendable(), + eof, + session, + ) + .await?; + Ok((segments, layout)) +} + +/// Pruning drops the zones whose statistics prove the filter false: of four 1000-row zones over +/// `0..4000`, the first cannot hold a value above 1500, so V2 reads the zone table and the other +/// three chunks, never the first, and still returns what the default path returns. +#[tokio::test(flavor = "multi_thread")] +async fn zone_pruning_reads_only_zones_that_can_match() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_zoned_layout(&session).await?; + + let data = layout + .slot(0)? + .ok_or_else(|| vortex_error::vortex_err!("no data"))?; + let zones = layout + .slot(1)? + .ok_or_else(|| vortex_error::vortex_err!("no zones"))?; + let data_chunks = data.children()?; + let mut expected_reads = segment_ids(&zones)?; + for chunk in &data_chunks[1..] { + expected_reads.extend(segment_ids(chunk)?); + } + + let case = case(true, None, None); + let dtype = builder(&session, &segments, &layout, &case)?.dtype()?; + let expected = builder(&session, &segments, &layout, &case)? + .into_stream()? + .try_collect::>() + .await?; + + let recording = Arc::new(RecordingSegments { + registered: Default::default(), + inner: Arc::clone(&segments), + reads: Arc::default(), + }); + let mut file = scan_file(&segments, &layout)?; + file.segments = Arc::clone(&recording) as _; + let actual = v2::into_stream(builder(&session, &segments, &layout, &case)?, file)? + .try_collect::>() + .await?; + + assert_arrays_eq!( + ChunkedArray::try_new(actual, dtype.clone())?, + ChunkedArray::try_new(expected, dtype)?, + &mut session.create_execution_ctx() + ); + assert_eq!(*recording.reads.lock(), expected_reads); + Ok(()) +} + +const WORDS: [&str; 4] = ["", "apple", "banana", "cherry"]; + +/// Four chunks of dictionary-encoded strings cycling through [`WORDS`]. +async fn write_dict_layout( + session: &VortexSession, +) -> VortexResult<(Arc, LayoutRef)> { + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunks = (0..4) + .map(|chunk| { + let words = (0..CHUNK_ROWS).map(|row| WORDS[((row * 7 + chunk) % 4) as usize]); + let values = VarBinArray::from_iter_nonnull(words, DType::Utf8(NonNullable)); + Ok((sequence_id.advance(), values.into_array())) + }) + .collect::>(); + let layout = DictStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + DictLayoutOptions::default(), + // The fixture must stay dictionary encoded regardless of compressor heuristics. + Arc::new(|chunk: &ArrayRef, ctx: &mut ExecutionCtx| { + Ok(dict_encode(chunk, ctx)?.into_array()) + }), + ) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(DType::Utf8(NonNullable), stream::iter(chunks)).sendable(), + eof, + session, + ) + .await?; + Ok((segments, layout)) +} + +/// Expressions over dictionary values run on shared canonical values when the scan also reads +/// the values whole, and on the encoded values otherwise, including the negative-cost part of a +/// projection split off from the rest; each returns what the default path returns. +#[rstest] +#[case::values_read_whole(root(), Some(eq(root(), lit("apple"))))] +#[case::byte_length_split( + cast(byte_length(root()), DType::Primitive(PType::I64, NonNullable)), + Some(not_eq(root(), lit(""))) +)] +#[case::byte_length_whole(byte_length(root()), Some(like(root(), lit("%an%"))))] +#[case::no_filter(byte_length(root()), None)] +#[tokio::test(flavor = "multi_thread")] +async fn dictionary_expressions_match_default( + #[case] projection: Expression, + #[case] filter: Option, + #[values(false, true)] with_selection: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_dict_layout(&session).await?; + assert!( + layout + .depth_first_traversal() + .any(|layout| layout.is_ok_and(|layout| layout.is::())), + "the strings must be dictionary encoded" + ); + let dtype = DType::Utf8(NonNullable); + let builder = || -> VortexResult> { + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let builder = + ScanBuilder::new(session.clone(), reader).with_projection(projection.bind(&dtype)?); + let builder = if with_selection { + builder.with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![ + 0_u64, 103, 999, 2001, 2500, 3002, + ]))?) + } else { + builder + }; + Ok(match &filter { + Some(filter) => builder.with_filter(filter.bind(&dtype)?), + None => builder, + }) + }; + let result_dtype = builder()?.dtype()?; + + let expected = builder()?.into_stream()?.try_collect::>().await?; + let actual = v2::into_stream(builder()?, scan_file(&segments, &layout)?)? + .try_collect::>() + .await?; + assert_arrays_eq!( + ChunkedArray::try_new(actual, result_dtype.clone())?, + ChunkedArray::try_new(expected, result_dtype)?, + &mut session.create_execution_ctx() + ); + Ok(()) +} + +/// Scans over one layout reader share its lowered file, and with it what executions learn, such +/// as dictionary values; scans over another reader of the same layout do not. +#[tokio::test(flavor = "multi_thread")] +async fn scans_over_one_reader_share_the_file() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let reader = || { + layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + ) + }; + let first = reader()?; + let shared = shared_file(&first, scan_file(&segments, &layout)?)?; + assert!(Arc::ptr_eq( + &shared, + &shared_file(&first, scan_file(&segments, &layout)?)? + )); + assert!(!Arc::ptr_eq( + &shared, + &shared_file(&reader()?, scan_file(&segments, &layout)?)? + )); + Ok(()) +} + +#[tokio::test] +async fn registry_does_not_retain_inactive_scan_caches() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let shared = shared_file(&reader, scan_file(&segments, &layout)?)?; + let weak = Arc::downgrade(&shared); + drop(shared); + assert!(weak.upgrade().is_none()); + assert_eq!(Arc::strong_count(&reader), 1); + Ok(()) +} + +/// A scan keeps no dictionary values once it has finished, even over an explicitly cached reader. +#[rstest] +#[tokio::test] +async fn reader_cache_retains_no_dictionary_values( + #[values(false, true)] cached: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_dict_layout(&session).await?; + let mut values = BTreeSet::new(); + for node in layout.depth_first_traversal() { + let node = node?; + if node.is::() { + values.extend(segment_ids( + &node.slot(0)?.ok_or_else(|| vortex_err!("missing values"))?, + )?); + } + } + assert!(!values.is_empty()); + let recording = Arc::new(RecordingSegments { + inner: segments, + reads: Default::default(), + registered: Default::default(), + }); + let source: Arc = Arc::::clone(&recording); + let reader = layout.new_reader( + "".into(), + Arc::clone(&source), + &session, + &Default::default(), + )?; + let mut file = scan_file(&source, &layout)?; + file.plans = cached.then(|| Arc::new(FilePlans::default())); + let mut first_reads = BTreeSet::new(); + for iteration in 0..2 { + recording.reads.lock().clear(); + let arrays = v2::into_stream( + ScanBuilder::new(session.clone(), Arc::clone(&reader)), + file.clone(), + )? + .try_collect::>() + .await?; + let expected = VarBinArray::from_iter_nonnull( + (0..4).flat_map(|chunk| { + (0..CHUNK_ROWS).map(move |row| WORDS[((row * 7 + chunk) % 4) as usize]) + }), + DType::Utf8(NonNullable), + ); + assert_arrays_eq!( + ChunkedArray::try_new(arrays, DType::Utf8(NonNullable))?, + expected, + &mut session.create_execution_ctx() + ); + let reads = recording.reads.lock().clone(); + if iteration == 0 { + assert!(values.is_subset(&reads)); + first_reads = reads; + } else { + // Values live no longer than the split that read them, cached reader or not. + assert_eq!(reads, first_reads); + } + } + Ok(()) +} + +#[tokio::test] +async fn reader_plan_cache_has_no_ownership_cycle() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let mut file = scan_file(&segments, &layout)?; + let plans = Arc::new(FilePlans::default()); + let weak_plans = Arc::downgrade(&plans); + file.plans = Some(plans); + let shared = shared_file(&reader, file.clone())?; + let root = shared.root.clone(); + let weak_shared = Arc::downgrade(&shared); + drop(shared); + assert!(weak_shared.upgrade().is_none()); + let shared = shared_file(&reader, file.clone())?; + assert!(PlanRef::ptr_eq(&root, &shared.root)); + drop(shared); + drop(file); + assert!(weak_plans.upgrade().is_none()); + assert_eq!(Arc::strong_count(&reader), 1); + Ok(()) +} + +#[tokio::test] +async fn cached_dictionary_accepts_different_queries() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_dict_layout(&session).await?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let mut file = scan_file(&segments, &layout)?; + file.plans = Some(Arc::default()); + let dtype = DType::Utf8(NonNullable); + for (projection, filter) in [ + (root(), eq(root(), lit("apple"))), + (byte_length(root()), gt(byte_length(root()), lit(3u64))), + (root(), eq(root(), lit("banana"))), + ] { + let builder = || -> VortexResult<_> { + Ok(ScanBuilder::new(session.clone(), Arc::clone(&reader)) + .with_projection(projection.bind(&dtype)?) + .with_filter(filter.bind(&dtype)?)) + }; + let result_dtype = builder()?.dtype()?; + let expected = builder()?.into_stream()?.try_collect::>().await?; + let actual = v2::into_stream(builder()?, file.clone())? + .try_collect::>() + .await?; + assert_arrays_eq!( + ChunkedArray::try_new(actual, result_dtype.clone())?, + ChunkedArray::try_new(expected, result_dtype)?, + &mut session.create_execution_ctx() + ); + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn concurrent_scans_over_one_reader_share_initialization() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_layout(&session).await?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let file = scan_file(&segments, &layout)?; + let barrier = Barrier::new(8); + let shared = thread::scope(|scope| { + let threads = (0..8) + .map(|_| { + let file = file.clone(); + let barrier = &barrier; + let reader = &reader; + scope.spawn(move || { + barrier.wait(); + shared_file(reader, file) + }) + }) + .collect::>(); + threads + .into_iter() + .map(|thread| { + thread + .join() + .map_err(|_| vortex_err!("file preparation thread panicked"))? + }) + .collect::>>() + })?; + for other in &shared[1..] { + assert!(Arc::ptr_eq(&shared[0], other)); + } + Ok(()) +} + +/// Zone pruning follows a dynamic comparison the engine tightens during the scan: with no bound +/// every zone is read, and once the bound is above 2500 a later execution skips the two zones +/// below it, without reading the zone table again. +#[tokio::test(flavor = "multi_thread")] +async fn zone_pruning_follows_dynamic_comparisons() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_zoned_layout(&session).await?; + let data = layout + .slot(0)? + .ok_or_else(|| vortex_error::vortex_err!("no data"))?; + let data_chunks = data.children()?; + + let bound = Arc::new(Mutex::new(None::)); + let filter = { + let bound = Arc::clone(&bound); + dynamic( + CompareOperator::Gt, + move || bound.lock().map(Into::into), + DTYPE, + true, + root(), + ) + }; + let recording = Arc::new(RecordingSegments { + registered: Default::default(), + inner: Arc::clone(&segments), + reads: Arc::default(), + }); + let mut file = scan_file(&segments, &layout)?; + file.segments = Arc::clone(&recording) as _; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let scan = v2::prepare( + ScanBuilder::new(session.clone(), reader).with_filter(filter.bind(&DTYPE)?), + file, + )?; + + let everything = await_tasks(DTYPE, scan.execute(None)?).await?; + assert_eq!(everything.len(), 4000); + for chunk in &data_chunks { + assert!(segment_ids(chunk)?.is_subset(&recording.reads.lock())); + } + + *bound.lock() = Some(2500); + recording.reads.lock().clear(); + let above = await_tasks(DTYPE, scan.execute(None)?).await?; + assert_arrays_eq!( + above, + Buffer::from_iter(2501..4000).into_array(), + &mut session.create_execution_ctx() + ); + let mut expected_reads = segment_ids(&data_chunks[2])?; + expected_reads.extend(segment_ids(&data_chunks[3])?); + assert_eq!(*recording.reads.lock(), expected_reads); + Ok(()) +} + +/// Four 1000-row chunks of a struct `{a, b}` with `a` over `0..4000` and `b = 2a`. +async fn write_struct_layout( + session: &VortexSession, + field_count: usize, +) -> VortexResult<(Arc, LayoutRef)> { + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunks = (0..4) + .map(|chunk| { + let a = Buffer::from_iter(chunk * CHUNK_ROWS..(chunk + 1) * CHUNK_ROWS); + let b = + Buffer::from_iter((chunk * CHUNK_ROWS..(chunk + 1) * CHUNK_ROWS).map(|a| 2 * a)); + let mut fields = vec![ + ("a".to_owned(), a.into_array()), + ("b".to_owned(), b.into_array()), + ]; + for field in 2..field_count { + let increment = i32::try_from(field)?; + fields.push(( + format!("f{field}"), + Buffer::from_iter( + (chunk * CHUNK_ROWS..(chunk + 1) * CHUNK_ROWS).map(|a| a + increment), + ) + .into_array(), + )); + } + let array = StructArray::from_fields(&fields)?; + Ok((sequence_id.advance(), array.into_array())) + }) + .collect::>>()?; + let dtype = chunks[0].1.dtype().clone(); + let layout = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(chunks.into_iter().map(Ok))) + .sendable(), + eof, + session, + ) + .await?; + Ok((segments, layout)) +} + +/// Narrow and wide projections retain one batch per selected cut. Both the joined result and +/// the stream preserve the reference executor's values while unused IO interests are withdrawn. +#[rstest::rstest] +#[case(2)] +#[case(8)] +#[tokio::test(flavor = "multi_thread")] +async fn struct_projection_splits_match_default(#[case] field_count: usize) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_struct_layout(&session, field_count).await?; + let builder = || -> VortexResult> { + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let filter = gt(get_item("a", root()), lit(1500_i32)).bind(layout.dtype())?; + Ok(ScanBuilder::new(session.clone(), reader).with_filter(filter)) + }; + let dtype = builder()?.dtype()?; + let expected = await_tasks(dtype.clone(), builder()?.prepare()?.execute(None)?).await?; + + let scan = v2::prepare(builder()?, scan_file(&segments, &layout)?)?; + let tasks = scan.execute(None)?; + assert_eq!(tasks.len(), 1, "4000 rows make one filter split"); + let joined = await_tasks(dtype.clone(), tasks).await?; + assert_arrays_eq!(joined, expected, &mut session.create_execution_ctx()); + + let streamed = v2::into_stream(builder()?, scan_file(&segments, &layout)?)? + .try_collect::>() + .await?; + assert_eq!( + streamed.len(), + 3, + "the three chunks with a > 1500 are projection splits" + ); + assert_arrays_eq!( + ChunkedArray::try_new(streamed, dtype)?, + expected, + &mut session.create_execution_ctx() + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn wide_projection_with_mismatched_chunks_and_mask_holes_matches_default() -> VortexResult<()> +{ + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let mut children = Vec::new(); + for field in 0..8_i32 { + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunk_rows = 700 + field * 91; + let chunks = (0..4000) + .step_by(usize::try_from(chunk_rows)?) + .map(|start| { + let rows = start..(start + chunk_rows).min(4000); + let values = Buffer::from_iter(rows.map(|row| row + field * 10_000)); + Ok((sequence_id.advance(), values.into_array())) + }) + .collect::>(); + children.push( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(DTYPE, stream::iter(chunks)).sendable(), + eof, + &session, + ) + .await?, + ); + } + let fields = StructFields::from_iter((0..8).map(|field| { + ( + if field == 0 { + "a".to_owned() + } else { + format!("f{field}") + }, + DTYPE, + ) + })); + let dtype = DType::Struct(fields, NonNullable); + let layout = StructLayout::new(4000, dtype.clone(), children).into_layout(); + let segments = segments as Arc; + let a = || get_item("a", root()); + let filter = or( + lt(a(), lit(10_i32)), + or( + and(gt(a(), lit(1499_i32)), lt(a(), lit(1510_i32))), + gt(a(), lit(3500_i32)), + ), + ) + .bind(&dtype)?; + let builder = || -> VortexResult> { + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + Ok(ScanBuilder::new(session.clone(), reader).with_filter(filter.clone())) + }; + let expected = await_tasks(dtype.clone(), builder()?.prepare()?.execute(None)?).await?; + let actual = v2::into_stream(builder()?, scan_file(&segments, &layout)?)? + .try_collect::>() + .await?; + assert_arrays_eq!( + ChunkedArray::try_new(actual, dtype)?, + expected, + &mut session.create_execution_ctx() + ); + Ok(()) +} + +#[rstest] +#[case(1500, Some(1000..4000))] +#[case(4000, None)] +#[case(-1, Some(0..4000))] +#[tokio::test(flavor = "multi_thread")] +async fn file_pruning_precedes_data_split_creation( + #[case] threshold: i32, + #[case] expected_scope: Option>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_zoned_layout(&session).await?; + let recording = Arc::new(RecordingSegments { + registered: Default::default(), + inner: Arc::clone(&segments), + reads: Arc::default(), + }); + let mut file = scan_file(&segments, &layout)?; + file.segments = Arc::clone(&recording) as _; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let scan = v2::prepare( + ScanBuilder::new(session.clone(), reader) + .with_filter(gt(root(), lit(threshold)).bind(&DTYPE)?), + file, + )?; + let splits = scan.pruned_split_plans(None).await?; + assert_eq!( + splits + .iter() + .map(|split| split.scope.rows.clone()) + .collect::>(), + expected_scope.into_iter().collect::>() + ); + let zones = layout + .slot(1)? + .ok_or_else(|| vortex_error::vortex_err!("no zones"))?; + let zone_segments = segment_ids(&zones)?; + assert_eq!(*recording.reads.lock(), zone_segments); + assert!(recording.registered.lock().is_subset(&zone_segments)); + let mut arrays = Vec::new(); + for split in splits { + arrays.extend(split.run(Arc::clone(scan.io())).await?); + } + assert_arrays_eq!( + ChunkedArray::try_new(arrays, DTYPE)?, + Buffer::from_iter((0..4000).filter(|value| *value > threshold)).into_array(), + &mut session.create_execution_ctx() + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn file_pruning_rechecks_dynamic_bounds_after_splitting() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_zoned_layout(&session).await?; + let bound = Arc::new(Mutex::new(None::)); + let filter = { + let bound = Arc::clone(&bound); + dynamic( + CompareOperator::Gt, + move || bound.lock().map(Into::into), + DTYPE, + true, + root(), + ) + }; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let scan = v2::prepare( + ScanBuilder::new(session.clone(), reader).with_filter(filter.bind(&DTYPE)?), + scan_file(&segments, &layout)?, + )?; + let splits = scan.pruned_split_plans(None).await?; + assert_eq!(splits[0].scope.rows, 0..4000); + *bound.lock() = Some(2500); + let mut arrays = Vec::new(); + for split in splits { + arrays.extend(split.run(Arc::clone(scan.io())).await?); + } + assert_arrays_eq!( + ChunkedArray::try_new(arrays, DTYPE)?, + Buffer::from_iter(2501..4000).into_array(), + &mut session.create_execution_ctx() + ); + let splits = scan.pruned_split_plans(None).await?; + assert_eq!(splits[0].scope.rows, 2000..4000); + *bound.lock() = None; + let splits = scan.pruned_split_plans(None).await?; + assert_eq!(splits[0].scope.rows, 0..4000); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn file_pruning_keeps_disjoint_zones_in_one_filter_task() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write_zoned_layout(&session).await?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segments), + &session, + &Default::default(), + )?; + let filter = or(lt(root(), lit(500i32)), gt(root(), lit(3500i32))).bind(&DTYPE)?; + let scan = v2::prepare( + ScanBuilder::new(session.clone(), reader).with_filter(filter), + scan_file(&segments, &layout)?, + )?; + let splits = scan.pruned_split_plans(None).await?; + assert_eq!(splits.len(), 1); + assert_eq!(splits[0].scope.rows, 0..4000); + let mut arrays = Vec::new(); + for split in splits { + arrays.extend(split.run(Arc::clone(scan.io())).await?); + } + assert_arrays_eq!( + ChunkedArray::try_new(arrays, DTYPE)?, + Buffer::from_iter((0..500).chain(3501..4000)).into_array(), + &mut session.create_execution_ctx() + ); + Ok(()) +}