diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 93a66ccd49e..a59a88d34ff 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,7 +130,8 @@ 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. Rows are +/// read with the V2 scan, which requires files without list layouts. pub struct VortexRandomAccessor { name: String, format: Format, @@ -142,11 +145,7 @@ impl VortexRandomAccessor { name: impl Into, format: Format, ) -> anyhow::Result { - let file = SESSION - .open_options() - .with_layout_reader_cache() - .open_path(path.as_ref()) - .await?; + let file = SESSION.open_options().open_path(path.as_ref()).await?; Ok(Self { name: name.into(), format, @@ -163,7 +162,6 @@ impl VortexRandomAccessor { ) -> anyhow::Result { let file = SESSION .open_options() - .with_layout_reader_cache() .open_object_store(remote.store(), remote.key(path)?) .await?; Ok(Self { @@ -185,16 +183,16 @@ 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()? + let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?; + let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file()) + .with_row_indices(indices) + .into_stream()?; + let array = ArrayStreamAdapter::new(self.file.dtype().clone(), 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. 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..f5fdafd3345 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::plan_selection; 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.rs b/vortex-layout/src/scan/execution/selection.rs new file mode 100644 index 00000000000..b699eaaf98f --- /dev/null +++ b/vortex-layout/src/scan/execution/selection.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Resolves known selections before allocating projection planners or their masks. + +use vortex_error::VortexExpect; +use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; + +use super::ScanPlans; +use super::SelectedRows; +use super::projection::ProjectionPlanner; + +/// Plans only the selected span, retaining one projection planner and IO session per work chunk. +pub(crate) 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 { + use std::ops::Range; + + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_error::VortexResult; + use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; + + use super::*; + + #[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(()) + } +} diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 46cbde987fc..739cb7ec06e 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -36,7 +36,6 @@ 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; @@ -46,19 +45,13 @@ 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::ScanPlans; pub use crate::scan::execution::SegmentLocation; -use crate::scan::execution::plan_projection; +use crate::scan::execution::plan_selection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::scan::v2::split::SplitPlan; @@ -237,34 +230,6 @@ 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, @@ -284,10 +249,6 @@ impl PreparedScan { /// fixed-size row 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 @@ -331,22 +292,15 @@ impl PreparedScan { 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, - }) + let (scope, root) = plan_selection( + &self.plans, + WorkScope { + file_ordinal: 0, + rows: rows.clone(), + }, + &self.selection, + )?; + Some(SplitPlan { scope, root }) }) .collect() } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index a910172f87b..c1a939083e5 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -428,7 +428,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(()) } @@ -481,3 +481,17 @@ async fn work_chunks_ignore_layout_boundaries( 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 splits = scan.split_plans(); + assert_eq!(splits.len(), 1); + assert_eq!(splits[0].scope.rows, 7..2002); + Ok(()) +}