diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 357e688ed53..3d72ad91264 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -229,7 +229,11 @@ fn measurement_name( format: Format, remote: bool, ) -> String { - let fmt = format.ext(); + let fmt = if format == Format::VortexCompact { + format.name() + } else { + format.ext() + }; let suffix = source_suffix(remote); match pattern { Some(p) => format!("random-access/{}/{}/{}-{}", dataset, p.name(), fmt, suffix), diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 805529851a7..38d6c8bef72 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -184,6 +184,7 @@ pub const ALLOWED_FORMATS: &[Format] = &[ Format::ArrowIpc, Format::Parquet, Format::OnDiskVortex, + Format::VortexCompact, Format::Lance, ]; diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 93a66ccd49e..b45d550743b 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -29,10 +29,12 @@ use tokio::fs::File; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; +use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::layout::scan::v2::ScanBuilder; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex::utils::aliases::hash_map::HashMap; @@ -128,11 +130,14 @@ impl RandomAccessor for ArrowIpcRandomAccessor { /// Random accessor for Vortex format files. /// -/// The file handle is opened at construction time and reused across `take()` calls. +/// The file handle is opened at construction time and reused across `take()` calls. Set +/// `VORTEX_BENCH_SCAN=v1` to compare with V1 using the same open options and materialization. +/// V2 requires files without list layouts. pub struct VortexRandomAccessor { name: String, format: Format, file: VortexFile, + v2: bool, } impl VortexRandomAccessor { @@ -150,6 +155,7 @@ impl VortexRandomAccessor { Ok(Self { name: name.into(), format, + v2: !std::env::var("VORTEX_BENCH_SCAN").is_ok_and(|scan| scan == "v1"), file, }) } @@ -169,6 +175,7 @@ impl VortexRandomAccessor { Ok(Self { name: name.into(), format, + v2: !std::env::var("VORTEX_BENCH_SCAN").is_ok_and(|scan| scan == "v1"), file, }) } @@ -185,16 +192,25 @@ impl RandomAccessor for VortexRandomAccessor { } async fn take(&self, indices: &[u64]) -> anyhow::Result { - let indices_buf: Buffer = Buffer::from(indices.to_vec()); - let array = self - .file - .scan()? - .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?) - .into_array_stream()? - .read_all() - .await?; + let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?; + let array = if self.v2 { + let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file()) + .with_row_indices(indices) + .into_stream()?; + ArrayStreamAdapter::new(self.file.dtype().clone(), stream) + .read_all() + .await? + } else { + self.file + .scan()? + .with_row_indices(indices) + .into_array_stream()? + .read_all() + .await? + }; - // We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es. + // Keep the historical benchmark contract: canonicalize the outer array, leaving nested + // values lazy. Both scan implementations must perform the same materialization. let mut ctx = SESSION.create_execution_ctx(); let canonical = array.execute::(&mut ctx)?.into_array(); Ok(RandomAccessorRet::ArrayRef(canonical)) diff --git a/vortex-layout/src/plan/physical/mod.rs b/vortex-layout/src/plan/physical/mod.rs index e1f0451cd0e..bedd3a7cad8 100644 --- a/vortex-layout/src/plan/physical/mod.rs +++ b/vortex-layout/src/plan/physical/mod.rs @@ -19,6 +19,7 @@ mod optimize; pub mod optimizer; pub mod pipeline; mod plans; +pub(crate) mod selection; mod typed; mod vtable; diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 8fc9c67b900..5a8daa39665 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -43,9 +43,11 @@ use crate::layouts::chunked::ChunkedLayout; use crate::layouts::dict::DictLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::struct_::StructLayout; +use crate::plan::physical::ConcatPlan; use crate::plan::physical::EvalPlan; use crate::plan::physical::Filter; use crate::plan::physical::PackPlan; +use crate::plan::physical::PlanChildren; use crate::plan::physical::PlanRef; use crate::plan::physical::SegmentScan; use crate::plan::physical::Take; @@ -60,6 +62,7 @@ use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; use crate::plan::physical::operators::PackSource; use crate::plan::physical::optimize; +use crate::plan::physical::selection::selected; use crate::segments::SegmentId; use crate::test::SESSION; @@ -834,3 +837,88 @@ fn pack_readiness_tracks_changes_and_consumption( } Ok(()) } + +#[rstest] +#[case::sparse(2..19, vec![0, 5, 14, 16])] +#[case::contiguous(2..19, (5..10).collect())] +#[case::empty(2..19, vec![])] +fn bound_selection_matches_projection( + #[case] rows: Range, + #[case] indices: Vec, +) -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + let mask = Mask::from_indices((rows.end - rows.start) as usize, indices); + let plan = selected(plan, rows.clone(), mask.clone()); + let result = run( + &store, + &plan, + rows.clone(), + Mask::new_true(mask.len()), + delivery(Delivery::Lifo), + )?; + assert_view(&expected, &rows, &mask, result.arrays) +} + +#[test] +fn bound_selection_never_initializes_skipped_chunks() -> VortexResult<()> { + let mut store = Store::default(); + let array = PrimitiveArray::from_iter(0..20i32).into_array(); + let first = lower(&store.flat(&array.slice(0..5)?)?)?; + let last = lower(&store.flat(&array.slice(15..20)?)?)?; + let children = PlanChildren::lazy(4, move |index| match index { + 0 => Ok(first.clone()), + 3 => Ok(last.clone()), + _ => Err(vortex_err!("unselected chunk {index} was initialized")), + }); + // SAFETY: All four logical chunks have five rows of the same primitive dtype. The two + // unselected initializers deliberately fail to detect discovery as well as compilation. + let plan = unsafe { + ConcatPlan::from_children_unchecked( + array.dtype().clone(), + 20, + vec![0, 5, 10, 15].into(), + children, + ) + } + .into_plan(); + let mask = Mask::from_indices(20, [2, 17]); + let plan = selected(plan, 0..20, mask.clone()); + let result = run( + &store, + &plan, + 0..20, + Mask::new_true(20), + delivery(Delivery::Fifo), + )?; + assert_eq!(reads(&result.events), 2); + assert_view(&array, &(0..20), &mask, result.arrays) +} + +#[test] +fn bound_selection_keeps_whole_dictionary_values() -> VortexResult<()> { + let mut store = Store::default(); + let values = PrimitiveArray::from_iter(100..120i32).into_array(); + let codes = PrimitiveArray::from_iter((0..20u32).rev()).into_array(); + let layout = DictLayout::new( + store.chunked(&values, &[5, 5, 5, 5])?, + store.chunked(&codes, &[5, 5, 5, 5])?, + ) + .into_layout(); + let mask = Mask::from_indices(20, [1, 18]); + let plan = selected(lower(&layout)?, 0..20, mask); + let result = run( + &store, + &plan, + 0..20, + Mask::new_true(20), + delivery(Delivery::Lifo), + )?; + assert_eq!(reads(&result.events), 6, "all values and two code chunks"); + assert_arrays_eq!( + join(values.dtype(), result.arrays)?, + PrimitiveArray::from_iter([118, 101]).into_array(), + &mut SESSION.create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-layout/src/plan/physical/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs index b2e1ffea0f7..c4c6b41dc11 100644 --- a/vortex-layout/src/plan/physical/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -26,6 +26,7 @@ use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; use crate::plan::physical::pipeline::Reach; use crate::plan::physical::pipeline::overlapping; +use crate::plan::physical::selection::selected_slice; use crate::segments::SegmentId; /// Concatenates its children row-wise. @@ -157,20 +158,26 @@ impl PlanVTable for Concat { mask: &Mask, compiler: &mut Compiler<'_>, ) -> VortexResult> { - let mut chains = Vec::with_capacity(overlapping(plan, &rows).size_hint().1.unwrap_or(0)); + let capacity = if mask.all_true() { + overlapping(plan, &rows).size_hint().1.unwrap_or(0) + } else { + 0 + }; + let mut chains = Vec::with_capacity(capacity); for (index, start, local) in overlapping(plan, &rows) { - let local_mask = mask.slice( + let Some((span, local_mask)) = selected_slice( + mask, usize::try_from(local.start - rows.start)? ..usize::try_from(local.end - rows.start)?, - ); - // A chunk the mask selects nothing of is never built, so never read. - if local_mask.all_false() { + ) else { continue; - } + }; let chunk = plan.child_required(index)?; - if let Some(chain) = - compiler.compile(&chunk, local.start - start..local.end - start, &local_mask)? - { + if let Some(chain) = compiler.compile( + &chunk, + rows.start + span.start as u64 - start..rows.start + span.end as u64 - start, + &local_mask, + )? { chains.push(chain); } } diff --git a/vortex-layout/src/plan/physical/plans/filter.rs b/vortex-layout/src/plan/physical/plans/filter.rs index 4850674c67b..6e8f032845a 100644 --- a/vortex-layout/src/plan/physical/plans/filter.rs +++ b/vortex-layout/src/plan/physical/plans/filter.rs @@ -103,13 +103,13 @@ impl PlanVTable for Filter { mask: &Mask, compiler: &mut Compiler<'_>, ) -> VortexResult> { + if mask.all_false() { + return Ok(None); + } let child = plan.child_plan()?; if let Some(scan) = child.as_opt::() { return compiler.scan(scan, rows, Some(mask.clone())); } - if mask.all_false() { - return Ok(None); - } let len = usize::try_from(rows.end - rows.start)?; let predicate = plan.dtype().is_boolean(); let chain = compiler.compile(&child, rows, &Mask::new_true(len))?; diff --git a/vortex-layout/src/plan/physical/selection.rs b/vortex-layout/src/plan/physical/selection.rs new file mode 100644 index 00000000000..177f800afcd --- /dev/null +++ b/vortex-layout/src/plan/physical/selection.rs @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selection-specific planning without changing the pipeline or planner contracts. + +use std::ops::Range; + +use vortex_array::EmptyMetadata; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use super::Concat; +use super::Eval; +use super::LegacyStats; +use super::Pack; +use super::Plan; +use super::PlanId; +use super::PlanParts; +use super::PlanRef; +use super::PlanVTable; +use super::Take; +use super::Zoned; +use super::pipeline::Chain; +use super::pipeline::Compiler; +use super::pipeline::Reach; +use super::pipeline::overlapping; +use crate::segments::SegmentId; + +#[derive(Clone, Debug)] +struct Selected; + +#[derive(Clone, Debug)] +struct SelectedData { + rows: Range, + mask: Mask, +} + +/// Binds a projection's selection before segment discovery can initialize unselected children. +/// Whole-value dependencies retain their original plans and cache identities. +pub(crate) fn selected(plan: PlanRef, rows: Range, mask: Mask) -> PlanRef { + if mask.all_true() { + return plan; + } + PlanParts { + vtable: Selected, + dtype: plan.dtype().clone(), + row_count: plan.row_count(), + children: vec![plan].into(), + data: SelectedData { rows, mask }, + } + .into_typed() + .into_plan() +} + +impl PlanVTable for Selected { + type PlanData = SelectedData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.selected"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + None + } + + fn compile( + plan: &Plan, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + let selected = plan.data().mask(&rows)?; + let mask = mask & &selected; + if mask.all_false() { + return Ok(None); + } + compiler.compile(&plan.child_required(0)?, rows, &mask) + } + + fn reach( + plan: &Plan, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + let mask = plan.data().mask(&rows)?; + reach_selected(&plan.child_required(0)?, rows, &mask, at, visit) + } +} + +impl SelectedData { + fn mask(&self, rows: &Range) -> VortexResult { + vortex_ensure!( + rows.start >= self.rows.start && rows.end <= self.rows.end, + "Selected plan requested outside its bound row range" + ); + Ok(self.mask.slice( + usize::try_from(rows.start - self.rows.start)? + ..usize::try_from(rows.end - self.rows.start)?, + )) + } +} + +/// Trims a child range to selected rows, without allocating a mask for an empty or contiguous +/// selection. The returned range remains relative to the input mask. +pub(crate) fn selected_slice(mask: &Mask, rows: Range) -> Option<(Range, Mask)> { + if rows.is_empty() { + return None; + } + match mask { + Mask::AllFalse(_) => None, + Mask::AllTrue(_) => { + let len = rows.end - rows.start; + Some((rows, Mask::new_true(len))) + } + Mask::Values(values) => { + let bits = values.bit_buffer().slice(rows.clone()); + let first = bits.set_indices().next()?; + let end = bits.last_set_index()? + 1; + if rows.start + first == 0 && rows.start + end == mask.len() { + // Keep the derived indices shared across fields when no rows were trimmed. + return Some((0..mask.len(), mask.clone())); + } + let len = end - first; + let selected = if bits.count_range(first, end) == len { + Mask::new_true(len) + } else { + Mask::from_buffer(bits.slice(first..end)) + }; + Some((rows.start + first..rows.start + end, selected)) + } + } +} + +fn reach_selected( + plan: &PlanRef, + rows: Range, + mask: &Mask, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), +) -> VortexResult<()> { + if mask.all_false() || rows.is_empty() { + return Ok(()); + } + if mask.all_true() { + return plan.reach(rows, at, visit); + } + if let Some(concat) = plan.as_opt::() { + for (index, start, local) in overlapping(concat, &rows) { + let Some((span, mask)) = selected_slice( + mask, + usize::try_from(local.start - rows.start)? + ..usize::try_from(local.end - rows.start)?, + ) else { + continue; + }; + reach_selected( + &concat.child_required(index)?, + rows.start + span.start as u64 - start..rows.start + span.end as u64 - start, + &mask, + &at.shift(start), + visit, + )?; + } + return Ok(()); + } + if plan.is::() || plan.is::() { + for child in plan.children().iter() { + reach_selected(&child?, rows.clone(), mask, at, visit)?; + } + return Ok(()); + } + if let Some(take) = plan.as_opt::() { + let values = take.values()?; + values.reach(0..values.row_count(), &at.fixed(&rows), visit)?; + return reach_selected(&take.codes()?, rows, mask, at, visit); + } + if let Some(zoned) = plan.as_opt::() + && !zoned.is_pruning() + { + return match zoned.data_plan()? { + Some(data) => reach_selected(&data, rows, mask, at, visit), + None => Ok(()), + }; + } + if let Some(stats) = plan.as_opt::() { + return reach_selected(&stats.data_plan()?, rows, mask, at, visit); + } + // Filters over non-segment children consume every input row. Share, pruning, and unknown + // operators may also read whole dependencies under a sparse mask. + plan.reach(rows, at, visit) +} diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 9d8b8897ef3..02e02ccc3b1 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -10,17 +10,17 @@ mod adaptor; mod projection; +mod selection; use std::sync::Arc; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; use vortex_mask::Mask; -use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; use vortex_session::VortexSession; -use self::projection::ProjectionPlanner; +pub(super) use self::selection::FilePlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; @@ -69,8 +69,3 @@ struct SelectedRows { /// Which rows of `scope.rows` are selected. mask: Mask, } - -/// The work for one split: project the rows of `scope` selected by `mask`. -pub(super) fn plan_projection(plans: ScanPlans, scope: WorkScope, mask: Mask) -> Box { - Box::new(ProjectionPlanner::new(plans, SelectedRows { scope, mask })) -} diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index b6180725c27..4941cacb0a7 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -10,6 +10,7 @@ use vortex_error::vortex_bail; use vortex_io::request::IoConsumer; use vortex_io::request::IoRequestId; use vortex_io::request::IoResult; +use vortex_mask::Mask; use vortex_scan::planning::morsel::Morsel; use vortex_scan::planning::morsel::MorselOutput; use vortex_scan::planning::planner::Planner; @@ -18,6 +19,7 @@ use vortex_scan::planning::planner::State; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; +use crate::plan::physical::selection::selected; use crate::scan::execution::ScanPlans; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; @@ -137,10 +139,14 @@ impl Morsel for ProjectionMorsel { .vortex_expect("an uncompiled projection still has its plans"); let scan = Scan::try_new( plans.session, - plans.projection, + selected( + plans.projection, + self.selected.scope.rows.clone(), + self.selected.mask.clone(), + ), vec![Split { rows: self.selected.scope.rows.clone(), - mask: self.selected.mask.clone(), + mask: Mask::new_true(self.selected.mask.len()), }], )? .with_row_offset(plans.row_offset) diff --git a/vortex-layout/src/scan/execution/selection/mod.rs b/vortex-layout/src/scan/execution/selection/mod.rs new file mode 100644 index 00000000000..44c752cc2c6 --- /dev/null +++ b/vortex-layout/src/scan/execution/selection/mod.rs @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! One file planner walks the selection and hands children to independently scheduled tasks. + +use std::ops::Range; + +use futures::SinkExt; +use futures::StreamExt; +use futures::channel::mpsc; +use futures::stream::BoxStream; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; +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_scan::selection::Selection; + +use super::ScanPlans; +use super::SelectedRows; +use super::projection::ProjectionPlanner; + +/// Advances the file's selection one work chunk at a time, emitting only selected projections. +pub(crate) struct FilePlanner { + plans: ScanPlans, + chunks: SelectionChunks, +} + +impl FilePlanner { + pub(crate) fn new( + plans: ScanPlans, + rows: Range, + chunk_rows: usize, + selection: Selection, + ) -> Self { + Self { + plans, + chunks: SelectionChunks { + rows, + chunk_rows: chunk_rows as u64, + selection, + }, + } + } + + /// Runs the file planner on its own task. The receiver schedules each emitted child, while + /// bounded handoff prevents the producer from allocating the rest of the scan in advance. + pub(crate) fn into_stream( + mut self, + lookahead: usize, + ) -> BoxStream<'static, VortexResult<(WorkScope, Box)>> { + let (mut send, mut recv) = mpsc::channel(lookahead); + let planning = self.plans.session.handle().spawn(async move { + while self.state() == State::NeedsCompute { + let output = match self.compute() { + Ok(PlannerOutput::Planner(scope, child)) => Ok((scope, child)), + Ok(PlannerOutput::Done) => break, + Ok(PlannerOutput::Continue) => continue, + Ok(output) => Err(vortex_err!("Unexpected file planner output: {output:?}")), + Err(error) => Err(error), + }; + let failed = output.is_err(); + if send.send(output).await.is_err() || failed { + break; + } + } + }); + async_stream::stream! { + while let Some(child) = recv.next().await { + yield child; + } + // Keep the task owned by the stream, propagate panics, and cancel it on early drop. + planning.await; + } + .boxed() + } +} + +impl IoConsumer for FilePlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for FilePlanner { + fn state(&self) -> State { + if self.chunks.rows.is_empty() { + State::Done + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + loop { + let Some(rows) = self.chunks.next() else { + return Ok(PlannerOutput::Done); + }; + if let Some((scope, child)) = plan_selection( + &self.plans, + WorkScope { + file_ordinal: 0, + rows, + }, + &self.chunks.selection, + ) { + return Ok(PlannerOutput::Planner(scope, child)); + } + } + } +} + +/// A lazy cursor over fixed row boundaries. Indexed selections seek directly to occupied chunks. +struct SelectionChunks { + rows: Range, + chunk_rows: u64, + selection: Selection, +} + +impl Iterator for SelectionChunks { + type Item = Range; + + fn next(&mut self) -> Option { + if self.rows.is_empty() { + return None; + } + let first = match &self.selection { + Selection::IncludeByIndex(indices) => indices + .get(indices.partition_point(|&row| row < self.rows.start)) + .copied(), + Selection::IncludeRoaring(indices) => { + let mut iter = indices.iter(); + iter.advance_to(self.rows.start); + iter.next() + } + _ => Some(self.rows.start), + }; + let Some(first) = first.filter(|&row| row < self.rows.end) else { + self.rows.start = self.rows.end; + return None; + }; + let aligned_start = first / self.chunk_rows * self.chunk_rows; + let start = aligned_start.max(self.rows.start); + let end = aligned_start + .saturating_add(self.chunk_rows) + .min(self.rows.end); + self.rows.start = end; + Some(start..end) + } +} + +/// Plans only the selected span, retaining one projection planner and IO session per work chunk. +fn plan_selection( + plans: &ScanPlans, + scope: WorkScope, + selection: &Selection, +) -> Option<(WorkScope, Box)> { + let selected = select_rows(scope, selection)?; + Some(( + selected.scope.clone(), + Box::new(ProjectionPlanner::new(plans.clone(), selected)), + )) +} + +fn select_rows(mut scope: WorkScope, selection: &Selection) -> Option { + let count = match selection { + Selection::IncludeByIndex(indices) => { + let first = indices.partition_point(|&row| row < scope.rows.start); + let end = indices.partition_point(|&row| row < scope.rows.end); + if first == end { + return None; + } + scope.rows = indices[first]..indices[end - 1] + 1; + (end - first) as u64 + } + _ => selection.row_count_in_range(&scope.rows), + }; + if count == 0 { + return None; + } + // Contiguous selections need neither a bitmap nor a filtering operator. + let mask = if count == scope.rows.end - scope.rows.start { + Mask::new_true(usize::try_from(count).vortex_expect("a work chunk fits usize")) + } else { + selection.row_mask(&scope.rows).mask().clone() + }; + let first = mask.first()?; + let end = mask.last()? + 1; + scope.rows = scope.rows.start + first as u64..scope.rows.start + end as u64; + Some(SelectedRows { + scope, + mask: mask.slice(first..end), + }) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/scan/execution/selection/tests.rs b/vortex-layout/src/scan/execution/selection/tests.rs new file mode 100644 index 00000000000..17e1b1744c7 --- /dev/null +++ b/vortex-layout/src/scan/execution/selection/tests.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; + +use super::SelectionChunks; +use super::select_rows; + +#[rstest] +#[case::point(vec![7], 7..8, true)] +#[case::contiguous(vec![7, 8, 9], 7..10, true)] +#[case::sparse(vec![7, 9], 7..10, false)] +#[case::clipped(vec![0, 7, 9, 100], 7..10, false)] +fn indexed_selection_is_bounded_before_mask_allocation( + #[case] indices: Vec, + #[case] expected: Range, + #[case] all_true: bool, +) -> VortexResult<()> { + let selection = Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(indices))?); + let selected = select_rows( + WorkScope { + file_ordinal: 3, + rows: 5..20, + }, + &selection, + ) + .expect("selected rows"); + assert_eq!(selected.scope.file_ordinal, 3); + assert_eq!(selected.scope.rows, expected); + assert_eq!(selected.mask.all_true(), all_true); + Ok(()) +} + +#[test] +fn excluded_rows_can_prove_the_whole_chunk_empty() -> VortexResult<()> { + let selection = + Selection::ExcludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![5, 6, 7]))?); + assert!( + select_rows( + WorkScope { + file_ordinal: 0, + rows: 5..8 + }, + &selection + ) + .is_none() + ); + Ok(()) +} + +#[rstest] +#[case::whole(0..150_000, vec![0..65_536, 65_536..131_072, 131_072..150_000])] +#[case::unaligned(10..140_000, vec![10..65_536, 65_536..131_072, 131_072..140_000])] +#[case::boundary(65_536..131_072, vec![65_536..131_072])] +#[case::empty(5..5, vec![])] +#[case::near_max(u64::MAX - 3..u64::MAX, vec![u64::MAX - 3..u64::MAX])] +fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) { + let actual: Vec<_> = SelectionChunks { + rows, + chunk_rows: 65_536, + selection: Selection::All, + } + .collect(); + assert_eq!(actual, expected); +} + +#[test] +fn first_chunk_does_not_enumerate_the_scan() { + let mut chunks = SelectionChunks { + rows: 0..u64::MAX, + chunk_rows: 8192, + selection: Selection::All, + }; + assert_eq!(chunks.next(), Some(0..8192)); + assert_eq!(chunks.next(), Some(8192..16_384)); +} + +#[rstest] +fn sparse_indices_jump_over_empty_chunks(#[values(false, true)] roaring: bool) -> VortexResult<()> { + let indices = [7, 8, u64::MAX - 2, u64::MAX - 1]; + let selection = if roaring { + Selection::IncludeRoaring(indices.into_iter().collect()) + } else { + Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from_iter(indices))?) + }; + let actual: Vec<_> = SelectionChunks { + rows: 0..u64::MAX, + chunk_rows: 8192, + selection, + } + .collect(); + assert_eq!(actual, [0..8192, (u64::MAX / 8192 * 8192)..u64::MAX]); + Ok(()) +} + +#[rstest] +#[case::dense((0..200).collect(), 0..200)] +#[case::clustered(vec![0, 1, 2, 3, 7, 8, 15, 16, 17, 190, 199], 0..200)] +#[case::clipped(vec![0, 1, 4, 7, 8, 16, 17, 190, 199], 3..193)] +#[case::outside(vec![0, 1, 2, 199], 3..193)] +#[case::empty(vec![], 0..200)] +fn chunk_ownership_and_selection( + #[case] indices: Vec, + #[case] rows: Range, + #[values(false, true)] roaring: bool, + #[values(false, true)] exclude: bool, +) -> VortexResult<()> { + let selection = match (roaring, exclude) { + (false, false) => { + Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(indices.clone()))?) + } + (false, true) => { + Selection::ExcludeByIndex(StrictSortedBuffer::try_new(Buffer::from(indices.clone()))?) + } + (true, false) => Selection::IncludeRoaring(indices.iter().copied().collect()), + (true, true) => Selection::ExcludeRoaring(indices.iter().copied().collect()), + }; + let expected: Vec<_> = rows + .clone() + .filter(|row| indices.contains(row) != exclude) + .collect(); + let mut actual = Vec::new(); + let mut owners = Vec::new(); + let mut chunks = SelectionChunks { + rows, + chunk_rows: 8, + selection, + }; + while let Some(chunk) = chunks.next() { + owners.push(chunk.start / 8); + if let Some(selected) = select_rows( + WorkScope { + file_ordinal: 0, + rows: chunk, + }, + &chunks.selection, + ) { + actual.extend(selected.scope.rows.clone().filter(|row| { + selected + .mask + .value(usize::try_from(row - selected.scope.rows.start).unwrap()) + })); + } + } + assert!(owners.windows(2).all(|pair| pair[0] < pair[1])); + assert_eq!(actual, expected); + Ok(()) +} + +#[test] +fn smaller_chunks_keep_file_alignment() { + let actual: Vec<_> = SelectionChunks { + rows: 7500..33_000, + chunk_rows: 16_384, + selection: Selection::All, + } + .collect(); + assert_eq!(actual, [7500..16_384, 16_384..32_768, 32_768..33_000]); +} diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 46cbde987fc..c6f78a4805a 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -13,6 +13,8 @@ //! error; filters and limits are not part of this builder's API. Work is partitioned into fixed //! chunks of up to 65,536 rows, independently of the stored layout's chunk boundaries. //! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. +//! A file planner emits selected chunk planners incrementally to tasks on the runtime, without +//! an upfront walk or allocation of all selected chunk planners. mod io; mod split; @@ -36,29 +38,21 @@ use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; -use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; -use crate::layouts::chunked::Chunked; -use crate::layouts::dict::Dict; -use crate::layouts::flat::Flat; -use crate::layouts::struct_::Struct; -use crate::layouts::zoned::LegacyStats; -use crate::layouts::zoned::Zoned; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; +use crate::scan::execution::FilePlanner; use crate::scan::execution::ScanPlans; pub use crate::scan::execution::SegmentLocation; -use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; @@ -82,20 +76,6 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } -/// Fixed boundaries in the file's row domain, clipped to the requested range. -fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { - if rows.is_empty() { - return Vec::new(); - } - let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; - (first..rows.end) - .step_by(rows_per_chunk) - .map(|start| { - start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) - }) - .collect() -} - /// The parts of a file the planner-driven executor reads through. #[derive(Clone)] pub struct ScanFile { @@ -197,7 +177,7 @@ impl ScanBuilder { self } - /// Sets the number of active splits per worker thread. Must be greater than zero. + /// Sets the number of active tasks per worker thread. Must be greater than zero. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); self.concurrency = concurrency; @@ -237,42 +217,14 @@ impl ScanBuilder { } } -/// Whether every required layout lowers to plans with a pipeline implementation. -/// -/// List layouts lower to a list assembly, which cannot compile before its offsets are read. -fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { - if !(layout.is::() - || layout.is::() - || layout.is::() - || layout.is::() - || layout.is::() - || layout.is::()) - { - return Ok(false); - } - // An identity projection cannot use zone statistics. Leave their layouts uninitialized. - if identity && (layout.is::() || layout.is::()) { - return match layout.slot(0)? { - Some(data) => supports_layout(&data, identity), - None => Ok(false), - }; - } - for child in layout.children()? { - if !supports_layout(&child, identity)? { - return Ok(false); - } - } - Ok(true) -} - /// A scan planned over its file, ready to run its splits. struct PreparedScan { session: VortexSession, plans: ScanPlans, /// Serves the splits' reads. io: Arc, - /// The scan's splits, over the root row domain, in row order. - splits: Vec>, + rows: Range, + chunk_rows: usize, selection: Selection, map_fn: Arc VortexResult + Send + Sync>, ordered: bool, @@ -280,14 +232,9 @@ struct PreparedScan { } impl PreparedScan { - /// Lowers the file's layout, plans the projection, and partitions the selected range into - /// fixed-size row chunks. + /// Lowers the file's layout and plans the projection without enumerating work chunks. fn try_new(builder: ScanBuilder) -> VortexResult { let file = builder.file; - vortex_ensure!( - supports_layout(&file.layout, builder.projection.is_root())?, - "V2 scans do not support this file's layout" - ); let root = lower(&file.layout)?; let row_count = root.row_count(); // An identity projection needs no rewrites. Keep the lowered children lazy, including @@ -303,7 +250,6 @@ impl PreparedScan { } else { MAX_ROWS_PER_CHUNK }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); let io = Arc::new(SegmentScanIo::new( file.segments, segment_ranges(&file.locations), @@ -318,7 +264,8 @@ impl PreparedScan { values: Default::default(), }, io, - splits, + rows: builder.row_range.unwrap_or(0..row_count), + chunk_rows, selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, @@ -326,58 +273,31 @@ impl PreparedScan { }) } - /// The splits with selected rows, each ready to be admitted to a planning run. - fn split_plans(&self) -> Vec { - self.splits - .iter() - .filter_map(|rows| { - let row_mask = self.selection.row_mask(rows); - if row_mask.mask().all_false() { - return None; - } - let scope = WorkScope { - file_ordinal: 0, - rows: row_mask.row_range(), - }; - Some(SplitPlan { - root: plan_projection( - self.plans.clone(), - scope.clone(), - row_mask.mask().clone(), - ), - scope, - }) - }) - .collect() - } - - /// Runs every split as a task on the session's runtime and yields their mapped batches. - /// - /// Like the default executor, up to `concurrency` splits per thread run at once, and their - /// batches are yielded in row order when the scan is ordered. + /// Runs one file planner and schedules its children independently as they are emitted. fn execute_stream(self) -> VortexResult>> { let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans(); + let planner = FilePlanner::new(self.plans, self.rows, self.chunk_rows, self.selection); let io = self.io; let map_fn = self.map_fn; - // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. - let tasks = stream::iter(split_plans).map(move |split| { - let io = Arc::clone(&io); - let map_fn = Arc::clone(&map_fn); - handle.spawn(async move { - split - .run(io) - .await? - .into_iter() - .map(|array| map_fn(array)) - .collect::>>() - }) - }); + let tasks = planner + .into_stream(concurrency) + .map_ok(move |(scope, root)| { + let io = Arc::clone(&io); + let map_fn = Arc::clone(&map_fn); + handle.spawn(async move { + SplitPlan { scope, root } + .run(io) + .await? + .into_iter() + .map(|array| map_fn(array)) + .collect::>>() + }) + }); let batches = if self.ordered { - tasks.buffered(concurrency).boxed() + tasks.try_buffered(concurrency).boxed() } else { - tasks.buffer_unordered(concurrency).boxed() + tasks.try_buffer_unordered(concurrency).boxed() }; Ok(batches .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index a910172f87b..c4fd1aaf0fa 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -6,11 +6,14 @@ use std::ops::Range; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::time::Duration; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; use futures::TryStreamExt; use futures::stream; +use parking_lot::Condvar; +use parking_lot::Mutex; use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::ArrayRef; @@ -33,7 +36,11 @@ use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; @@ -41,7 +48,6 @@ use super::MAX_ROWS_PER_CHUNK; use super::PreparedScan; use super::ScanBuilder; use super::ScanFile; -use super::row_chunks; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; @@ -52,8 +58,10 @@ use crate::layouts::row_idx::row_idx; use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; +use crate::scan::execution::FilePlanner; use crate::scan::execution::SegmentLocation; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; +use crate::scan::v2::split::SplitPlan; use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; @@ -229,8 +237,10 @@ impl Strategy { #[case::whole(None, None)] #[case::row_range(Some(250..3_100), None)] #[case::empty_row_range(Some(10..10), None)] +#[case::empty_indices(None, Some(&[][..]))] #[case::row_indices(None, Some(&[0, 1, 999, 1000, 2500, 3999][..]))] #[case::row_range_and_indices(Some(900..2600), Some(&[0, 999, 1000, 2500, 3999][..]))] +#[case::indices_outside_range(Some(900..2600), Some(&[0, 3999][..]))] #[tokio::test] async fn matches_v1( #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, @@ -303,16 +313,6 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { Ok(()) } -#[rstest] -#[case::whole(0..150_000, vec![0..65_536, 65_536..131_072, 131_072..150_000])] -#[case::unaligned(10..140_000, vec![10..65_536, 65_536..131_072, 131_072..140_000])] -#[case::boundary(65_536..131_072, vec![65_536..131_072])] -#[case::empty(5..5, vec![])] -#[case::near_max(u64::MAX - 3..u64::MAX, vec![u64::MAX - 3..u64::MAX])] -fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) { - assert_eq!(row_chunks(rows, 65_536), expected); -} - #[rstest] #[case::empty(0, 0, 65_536)] #[case::single_chunk(1000, 100_000_000, 65_536)] @@ -329,19 +329,62 @@ fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: us assert_eq!(rows_per_chunk(rows, &locations), expected); } -#[test] -fn smaller_chunks_keep_file_alignment() { - assert_eq!( - row_chunks(7500..33_000, 16_384), - [7500..16_384, 16_384..32_768, 32_768..33_000] - ); -} - struct CountingSegments { inner: Arc, reads: AtomicUsize, } +struct ConcurrentSegments { + inner: Arc, + started: Mutex, + ready: Condvar, +} + +impl SegmentSource for ConcurrentSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + let mut started = self.started.lock(); + *started += 1; + self.ready.notify_all(); + let timeout = + self.ready + .wait_while_for(&mut started, |started| *started < 2, Duration::from_secs(5)); + drop(started); + assert!( + !timeout.timed_out(), + "a second child must run while the first child is blocked" + ); + self.inner.request(id) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn file_planner_schedules_children_on_independent_workers() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let segments = Arc::new(ConcurrentSegments { + inner: segments, + started: Mutex::new(0), + ready: Condvar::new(), + }) as Arc; + let mut scan = PreparedScan::try_new( + builder(&session, &segments, &layout)? + .with_projection(get_item("ints", root()).bind(layout.dtype())?) + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 3007]))?), + )?; + scan.chunk_rows = CHUNK_ROWS; + let actual = collect( + scan.plans.projection.dtype().clone(), + scan.execute_stream()?, + ) + .await?; + assert_arrays_eq!( + actual, + Buffer::from(vec![7_i32, 3007]).into_array(), + &mut session.create_execution_ctx() + ); + Ok(()) +} + impl SegmentSource for CountingSegments { fn request(&self, id: SegmentId) -> SegmentFuture { self.reads.fetch_add(1, Ordering::Relaxed); @@ -385,12 +428,18 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans(); - assert_eq!( - splits.len(), - (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) + let mut planner = FilePlanner::new( + scan.plans.clone(), + scan.rows.clone(), + scan.chunk_rows, + scan.selection.clone(), ); - for (index, split) in splits.into_iter().enumerate() { + let splits = (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK); + for index in 0..splits { + let PlannerOutput::Planner(scope, root) = planner.compute()? else { + vortex_bail!("all rows must have a projection planner"); + }; + let split = SplitPlan { scope, root }; let rows = split.scope.rows.clone(); let before = counted.reads.load(Ordering::Relaxed); let arrays = split.run(Arc::clone(&scan.io)).await?; @@ -413,6 +462,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { &mut session.create_execution_ctx() ); } + assert_eq!(planner.state(), State::Done); Ok(()) } @@ -428,7 +478,7 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { .try_next() .await .expect_err("V2 must reject an unsupported layout"); - assert!(error.to_string().contains("V2 scans do not support")); + assert!(error.to_string().contains("vortex.plan.list_pack")); Ok(()) } @@ -475,9 +525,50 @@ async fn work_chunks_ignore_layout_boundaries( let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - assert_eq!(scan.splits, vec![0..(CHUNKS * CHUNK_ROWS) as u64]); + assert_eq!(scan.rows, 0..(CHUNKS * CHUNK_ROWS) as u64); let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) } + +#[tokio::test] +async fn selected_split_scope_is_bounded_before_projection() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let scan = PreparedScan::try_new( + builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?), + )?; + let mut planner = FilePlanner::new(scan.plans, scan.rows, scan.chunk_rows, scan.selection); + let PlannerOutput::Planner(scope, _) = planner.compute()? else { + vortex_bail!("selected chunk must have a projection planner"); + }; + assert_eq!(scope.rows, 7..2002); + assert_eq!(planner.state(), State::Done); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn dropping_file_planner_stream_releases_pending_work( + #[values(false, true)] consume_first: bool, + #[values(1, 4)] lookahead: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; + let locations = Arc::downgrade(&scan.plans.locations); + let mut children = + FilePlanner::new(scan.plans, scan.rows, 1, scan.selection).into_stream(lookahead); + if consume_first { + assert!(children.try_next().await?.is_some()); + } + drop(children); + tokio::task::yield_now().await; + assert!( + locations.upgrade().is_none(), + "dropping the stream must cancel the producer and release queued planners" + ); + Ok(()) +}