diff --git a/Cargo.lock b/Cargo.lock index 22aace294d8..6f3d8741fbf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11323,6 +11323,7 @@ dependencies = [ "arrow-schema 59.3.0", "codspeed-divan-compat", "fastlanes 0.7.2", + "futures", "mimalloc", "parquet 59.3.0", "rand 0.10.3", @@ -12396,6 +12397,7 @@ dependencies = [ "parking_lot", "roaring", "rstest", + "smallvec", "tracing", "vortex-array", "vortex-buffer", diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 73f52f5b219..f845bb8ee11 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -28,6 +28,8 @@ work in each timed iteration. CI drives the full matrix via cargo run -p random-access-bench --profile release_debug --features lance ``` +Vortex measurements always use the planner-driven V2 scan builder. + ## Running against S3 The same benchmark can read its data from an object store instead of local disk. The remote diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index 93a66ccd49e..2291905825d 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; @@ -142,11 +144,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 +161,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 { @@ -186,11 +183,10 @@ 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()? + let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file()) .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?) - .into_array_stream()? + .into_stream()?; + let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream) .read_all() .await?; diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs index db4beeaa043..de952cf6b91 100644 --- a/vortex-layout/benches/pipeline_scan.rs +++ b/vortex-layout/benches/pipeline_scan.rs @@ -10,6 +10,7 @@ use std::hint::black_box; use divan::Bencher; use mimalloc::MiMalloc; use vortex_array::ArrayContext; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -17,7 +18,11 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::serde::SerializeOptions; use vortex_buffer::Alignment; +use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; +use vortex_layout::plan::physical::ConcatPlan; +use vortex_layout::plan::physical::FilterPlan; +use vortex_layout::plan::physical::ListPackPlan; use vortex_layout::plan::physical::PackPlan; use vortex_layout::plan::physical::PlanRef; use vortex_layout::plan::physical::SegmentScanPlan; @@ -26,6 +31,7 @@ use vortex_layout::plan::physical::pipeline::Split; use vortex_layout::plan::physical::pipeline::Turn; use vortex_layout::segments::SegmentId; use vortex_layout::session::LayoutSession; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; @@ -124,9 +130,140 @@ fn scan(bencher: Bencher, columns: u32) { Turn::Output(_, array) => { black_box(array); } + Turn::Child(_) => unreachable!("benchmark has no asynchronous child work"), Turn::Waiting => unreachable!("every read is delivered immediately"), Turn::Done => break, } } }); } + +struct ListFixture { + session: VortexSession, + plan: PlanRef, + segments: Vec, + mask: Mask, +} + +impl ListFixture { + fn new(chunks: u32, uniform: bool) -> Self { + const LISTS: usize = 65_536; + const ELEMENTS_PER_LIST: usize = 8; + let session = vortex_array::array_session().with::(); + let chunk_len = LISTS * ELEMENTS_PER_LIST / chunks as usize; + let elements = Buffer::from_iter(0..chunk_len as u64).into_array(); + let offsets = Buffer::from_iter((0..=LISTS as u64).map(|row| row * 8)).into_array(); + let serialize = |array: &ArrayRef| { + let context = ArrayContext::empty(); + let buffers = array + .serialize( + &context, + &session, + &SerializeOptions { + offset: 0, + include_padding: true, + }, + ) + .unwrap(); + let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64)); + for buffer in buffers { + bytes.extend_from_slice(buffer.as_ref()); + } + (BufferHandle::new_host(bytes.freeze()), context) + }; + let (element_bytes, element_context) = serialize(&elements); + let (offset_bytes, offset_context) = serialize(&offsets); + let children = (0..chunks) + .map(|id| { + FilterPlan::new( + SegmentScanPlan::new( + elements.dtype().clone(), + chunk_len as u64, + SegmentId::from(id), + ReadContext::new(element_context.to_ids()), + None, + ) + .into_plan(), + ) + .into_plan() + }) + .collect(); + let elements = ConcatPlan::try_new(elements.dtype().clone(), children).unwrap(); + let offsets = FilterPlan::new( + SegmentScanPlan::new( + offsets.dtype().clone(), + LISTS as u64 + 1, + SegmentId::from(chunks), + ReadContext::new(offset_context.to_ids()), + None, + ) + .into_plan(), + ); + let plan = ListPackPlan::try_new( + Nullability::NonNullable, + LISTS as u64, + elements.into_plan(), + offsets.into_plan(), + None, + ) + .unwrap() + .into_plan(); + let mut segments = vec![element_bytes; chunks as usize]; + segments.push(offset_bytes); + let indices = (0..100).map(|index| { + if uniform { + index * (LISTS / 100) + } else { + index + } + }); + Self { + session, + plan, + segments, + mask: Mask::from_indices(LISTS, indices), + } + } + + fn scan(&self) { + let split = Split { + rows: 0..self.mask.len() as u64, + mask: self.mask.clone(), + }; + let scan = Scan::try_new(self.session.clone(), self.plan.clone(), vec![split]).unwrap(); + let arrays = self.execute(scan); + assert_eq!(arrays.iter().map(|array| array.len()).sum::(), 100); + black_box(arrays); + } + + /// A minimal in-memory planner driving each selected child graph through the same protocol. + fn execute(&self, mut scan: Scan) -> Vec { + let mut output = Vec::new(); + loop { + match scan.step().unwrap() { + Turn::Read(read) => { + scan.deliver(read.id, self.segments[*read.segment_id as usize].clone()) + .unwrap(); + } + Turn::Output(split, array) => output.push((split, array)), + Turn::Child(child) => { + let (id, child_scan) = child.into_scan(self.session.clone()).unwrap(); + let arrays = self.execute(child_scan); + scan.complete_child(id, Some(arrays)).unwrap(); + } + Turn::Waiting => unreachable!("every read is delivered immediately"), + Turn::Done => break, + } + } + output.sort_by_key(|(split, _)| *split); + output.into_iter().map(|(_, array)| array).collect() + } +} + +/// Fixed selected row count, varying physical chunk count. Compare in separate processes with +/// `VORTEX_LIST_PIPELINE` unset and set to `1`; no disk IO or Arrow conversion is timed. +#[divan::bench(args = [8, 128, 1024], consts = [false, true])] +fn list_scan(bencher: Bencher, chunks: u32) { + let fixture = ListFixture::new(chunks, UNIFORM); + bencher.bench_local(|| fixture.scan()); +} diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 459429046c2..55c268b295e 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -125,6 +125,11 @@ impl Layout { OwnedLayoutChildren::layout_children(children), ) } + + /// Validated chunk boundaries, including the terminal row count. + pub(crate) fn chunk_offsets(&self) -> &[u64] { + &self.chunk_offsets + } } fn chunk_offsets(children: &dyn LayoutChildren) -> VortexResult> { diff --git a/vortex-layout/src/plan/physical/lower.rs b/vortex-layout/src/plan/physical/lower.rs index 1ea284cc1ca..9f965206544 100644 --- a/vortex-layout/src/plan/physical/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -83,21 +83,13 @@ fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan { } fn lower_chunked(layout: &ChunkedLayout) -> VortexResult { - let mut row_offsets = Vec::with_capacity(layout.nchildren()); - let mut row_count = 0u64; - for index in 0..layout.nchildren() { - row_offsets.push(row_count); - row_count = row_count - .checked_add(layout.child_row_count(index)) - .ok_or_else(|| vortex_err!("Chunked row count overflow"))?; - } // SAFETY: Chunked layout construction validates that every child has the parent dtype and - // that their row counts sum to the parent; offsets were computed with checked addition above. + // that these checked cumulative offsets end at the parent's row count. Ok(unsafe { ConcatPlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), - row_offsets.into(), + layout.chunk_offsets()[..layout.nchildren()].into(), lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()), ) }) diff --git a/vortex-layout/src/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index 96a8d0cfbd7..de1b1e35560 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -15,6 +15,7 @@ use vortex_session::VortexSession; use super::pipeline::Arena; use super::pipeline::Inlet; use super::pipeline::PortId; +use crate::plan::physical::pipeline::ChildProjection; use crate::segments::SegmentId; /// The output of one non-blocking compute call. @@ -43,6 +44,8 @@ pub enum SourceReady { pub enum SourceBlocked { /// Wait for bytes from newly queued or outstanding reads. Io, + /// Wait for a child projection, possibly alongside pending reads. + Child, /// The inlet needs more batches or its sink to close it. It must be open and have room /// for its sink to push; debug builds check that the sink can still wake this source. Inlet(usize), @@ -157,9 +160,28 @@ pub struct Cx<'a> { pub(super) exec: &'a mut ExecutionCtx, pub(super) touched: &'a mut InletSet, pub(super) changed: &'a mut InletSet, + pub(super) child_requests: &'a mut Vec, + pub(super) child_inbox: &'a mut Option>, +} + +#[derive(Default)] +pub(super) struct ChildInbox { + pub(super) pending: usize, + pub(super) completed: VecDeque<(u64, Option>)>, } impl Cx<'_> { + /// Describes child work for the enclosing planner. Return [`SourceBlocked::Child`] while awaiting its completion. + /// Identifiers need only be unique among this source's outstanding child projections. + pub(crate) fn request_child(&mut self, child: ChildProjection) { + self.child_requests.push(child); + } + + /// Takes a completed child projection, oldest completion first. + pub(crate) fn take_child(&mut self) -> Option<(u64, Option>)> { + self.child_inbox.as_mut()?.completed.pop_front() + } + /// Takes an inlet index whose sink appended data or closed since the last readiness /// check. Each index is reported once, and all inlets are reported on the first check. /// Remaining notifications are cleared after [`Source::ready`] returns. A source must diff --git a/vortex-layout/src/plan/physical/operators/list.rs b/vortex-layout/src/plan/physical/operators/list.rs new file mode 100644 index 00000000000..7b94351b2fa --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/list.rs @@ -0,0 +1,400 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Resolves selected list offsets before scanning their element ranges. + +use std::collections::VecDeque; +use std::ops::Range; + +use rustc_hash::FxHashMap; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_utils::parallelism::get_available_parallelism; + +use super::join; +use super::join_inlet; +use crate::plan::physical::ListPackPlan; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; +use crate::plan::physical::pipeline::ChildProjection; +use crate::plan::physical::pipeline::ReadId; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::SelectedPlan; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::Turn; +use crate::plan::physical::pipeline::ValueCache; +use crate::segments::SegmentId; + +mod work; + +/// A list's offsets and validity arrive before its selected element ranges can be planned. +pub(crate) struct ListSource { + plan: ListPackPlan, + mask: Mask, + values: ValueCache, + plan_children: bool, + planner_owned: bool, + awaiting_child: bool, + serial_splits: Option<(Vec, bool)>, + initialized: bool, + offsets: Option, + validity: Option, + elements: Option, + reads: FxHashMap>, + pending_children: usize, + chunks: Vec<(usize, ArrayRef)>, + output: Option, + done: bool, +} + +impl ListSource { + pub(crate) fn new( + plan: ListPackPlan, + mask: Mask, + values: ValueCache, + plan_children: bool, + planner_owned: bool, + ) -> Self { + Self { + plan, + mask, + values, + plan_children, + planner_owned, + awaiting_child: false, + serial_splits: None, + initialized: false, + offsets: None, + validity: None, + elements: None, + reads: FxHashMap::default(), + pending_children: 0, + chunks: Vec::new(), + output: None, + done: false, + } + } + + fn initialize(&mut self, cx: &mut Cx<'_>) -> VortexResult<()> { + let offsets = join_inlet(cx, 0, self.plan.offsets()?.dtype())? + .cast(DType::Primitive(PType::U64, false.into()))? + .execute::(cx.exec())?; + let validity = if self.plan.dtype().is_nullable() { + Validity::Array(join_inlet(cx, 1, &DType::Bool(false.into()))?) + } else { + Validity::NonNullable + }; + let valid = validity.execute_mask(self.mask.true_count(), cx.exec())?; + let elements = self.plan.elements()?; + let offsets = offsets.as_slice::(); + let (mut selected_offsets, mut ranges) = match self.mask.indices() { + AllOr::All => { + selected_ranges(0..self.mask.len(), offsets, &valid, elements.row_count())? + } + AllOr::Some(indices) => selected_ranges( + indices.iter().copied(), + offsets, + &valid, + elements.row_count(), + )?, + AllOr::None => (Buffer::from(vec![0_u64]), Vec::new()), + }; + // An all-selected scope can retain payload under null rows. For dense, bounded spans, + // this avoids filtering every element just to remove values hidden by list validity. + if self.mask.all_true() + && !valid.all_true() + && let Some(offsets) = work::preserve_null_span( + elements.dtype(), + offsets, + selected_offsets.last().copied().unwrap_or(0), + &mut ranges, + ) + { + selected_offsets = offsets; + } + let offsets = selected_offsets; + let last_offset = offsets.last().copied().unwrap_or(0); + let offsets = offsets.into_array(); + // These offsets start at zero and are checked sums of nonnegative lengths. Retain + // those facts so ListArray validation need not compute aggregates over them again. + let stats = offsets.statistics(); + stats.set(Stat::IsSorted, Precision::Exact(true.into())); + stats.set(Stat::Min, Precision::Exact(0_u64.into())); + stats.set(Stat::Max, Precision::Exact(last_offset.into())); + self.offsets = Some(offsets); + self.validity = Some(validity); + let overlap = work::can_overlap( + &elements, + &ranges, + self.mask.true_count(), + get_available_parallelism().unwrap_or(1), + )?; + let splits = work::element_splits(elements.dtype(), ranges)?; + // The baseline only hands off substantial fixed-width work. The selected-view path + // lets the planner own every child graph, including tiny and nested selections. + let width = match elements.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => 0, + }; + let bytes = splits.iter().fold(0_u64, |bytes, split| { + bytes.saturating_add((split.mask.true_count() as u64).saturating_mul(width)) + }); + if self.planner_owned || (self.plan_children && bytes >= 512 * 1024) { + if !self.planner_owned { + self.serial_splits = Some((splits.clone(), overlap)); + } + self.awaiting_child = true; + cx.request_child(ChildProjection { + id: 0, + view: SelectedPlan::new(elements, splits, if overlap { 2 } else { 1 }), + planner_owned: self.planner_owned, + values: self.values.clone(), + }); + self.initialized = true; + return Ok(()); + } + self.start_elements(cx, splits, overlap)?; + self.initialized = true; + Ok(()) + } + + fn start_elements( + &mut self, + cx: &Cx<'_>, + splits: Vec, + overlap: bool, + ) -> VortexResult<()> { + let scan = Scan::try_new(cx.session().clone(), self.plan.elements()?, splits)? + .with_value_cache(self.values.clone()) + .with_child_planning(self.plan_children); + self.elements = Some(if overlap { + scan.with_max_active(2) + } else { + scan + }); + Ok(()) + } + + fn finish(&mut self) -> VortexResult<()> { + // Independent element ranges may finish out of order. Preserve batch order within + // each range while putting the ranges back in list order. + self.chunks.sort_by_key(|(split, _)| *split); + let elements = join( + self.plan.elements()?.dtype(), + std::mem::take(&mut self.chunks) + .into_iter() + .map(|(_, array)| array) + .collect(), + )?; + let offsets = self + .offsets + .take() + .ok_or_else(|| vortex_err!("Missing list offsets"))?; + let validity = self + .validity + .take() + .ok_or_else(|| vortex_err!("Missing list validity"))?; + self.output = Some(ListArray::try_new(elements, offsets, validity)?.into_array()); + self.elements = None; + Ok(()) + } +} + +impl Source for ListSource { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if self.done || self.output.is_some() { + return Ok(SourceReady::Ready); + } + if !self.initialized { + for index in 0..1 + usize::from(self.plan.dtype().is_nullable()) { + if !cx.inlet(index).closed() { + return Ok(SourceReady::Blocked(SourceBlocked::Inlet(index))); + } + } + self.initialize(cx)?; + } + if self.awaiting_child { + let Some((_, arrays)) = cx.take_child() else { + return Ok(SourceReady::Blocked(SourceBlocked::Child)); + }; + self.awaiting_child = false; + if let Some(arrays) = arrays { + self.chunks.extend(arrays.into_iter().enumerate()); + self.finish()?; + return Ok(SourceReady::Ready); + } + vortex_ensure!(!self.planner_owned, "Planner declined a selected list view"); + let (splits, overlap) = self + .serial_splits + .take() + .ok_or_else(|| vortex_err!("Missing deferred list ranges"))?; + self.start_elements(cx, splits, overlap)?; + } + let scan = self + .elements + .as_mut() + .ok_or_else(|| vortex_err!("Missing list element scan"))?; + while let Some((segment, bytes)) = cx.take_segment() { + let request = self + .reads + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .ok_or_else(|| vortex_err!("Unexpected list element segment {segment}"))?; + scan.deliver(request, bytes)?; + } + while let Some((id, arrays)) = cx.take_child() { + scan.complete_child(id, arrays)?; + self.pending_children -= 1; + } + loop { + match scan.step()? { + Turn::Read(read) => { + self.reads + .entry(read.segment_id) + .or_default() + .push_back(read.id); + cx.request_read(read.segment_id); + } + Turn::Output(split, array) => self.chunks.push((split, array)), + Turn::Child(child) => { + self.pending_children += 1; + cx.request_child(child); + } + Turn::Waiting => { + return Ok(SourceReady::Blocked(if self.pending_children > 0 { + SourceBlocked::Child + } else { + SourceBlocked::Io + })); + } + Turn::Done => { + self.finish()?; + return Ok(SourceReady::Ready); + } + } + } + } +} + +impl Operator for ListSource { + fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { + if self.done { + return Ok(Step::End); + } + self.done = true; + self.output + .take() + .map(Step::Final) + .ok_or_else(|| vortex_err!("List computed before its elements were ready")) + } +} + +/// Compacts selected offsets and merges adjacent element ranges. No mask spans the gaps between +/// distant lists, so planning memory scales with selected rows rather than their element span. +fn selected_ranges( + selected: impl Iterator, + offsets: &[u64], + valid: &Mask, + element_count: u64, +) -> VortexResult<(Buffer, Vec>)> { + let mut output_offsets = Vec::with_capacity(valid.len() + 1); + output_offsets.push(0_u64); + let mut ranges: Vec> = Vec::new(); + let mut previous = None; + let mut cursor = 0; + let mut total = 0_u64; + let mut last_end = 0; + for (index, is_valid) in selected.zip(valid.iter()) { + // Consecutive selected lists share one boundary in the offsets selection. + if previous == Some(index) { + cursor -= 1; + } + let pair = offsets + .get(cursor..cursor + 2) + .ok_or_else(|| vortex_err!("Missing selected list boundaries"))?; + let start = pair[0]; + let end = pair[1]; + vortex_ensure!( + last_end <= start && start <= end && end <= element_count, + "Invalid list element range {start}..{end}" + ); + last_end = end; + if is_valid { + total = total + .checked_add(end - start) + .ok_or_else(|| vortex_err!("Selected list elements overflow"))?; + if start < end { + match ranges.last_mut() { + Some(last) if last.end == start => last.end = end, + _ => ranges.push(start..end), + } + } + } + output_offsets.push(total); + previous = Some(index + 1); + cursor += 2; + } + vortex_ensure!( + cursor == offsets.len(), + "Unexpected selected list boundaries" + ); + Ok((Buffer::from(output_offsets), ranges)) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::adjacent(vec![0, 1], vec![0, 2, 5], vec![true, true], vec![0, 2, 5], vec![0..5])] + #[case::gaps(vec![0, 3], vec![2, 4, 100, 103], vec![true, true], vec![0, 2, 5], vec![2..4, 100..103])] + #[case::null(vec![0, 1, 4], vec![0, 3, 8, 100, 105], vec![true, false, true], vec![0, 3, 3, 8], vec![0..3, 100..105])] + #[case::empty_lists(vec![0, 1], vec![2, 2, 2], vec![true, true], vec![0, 0, 0], vec![])] + fn compacts_selected_elements( + #[case] selected: Vec, + #[case] boundaries: Vec, + #[case] valid: Vec, + #[case] offsets: Vec, + #[case] ranges: Vec>, + ) -> VortexResult<()> { + let (actual_offsets, actual_ranges) = selected_ranges( + selected.into_iter(), + &boundaries, + &Mask::from_iter(valid), + 1000, + )?; + assert_eq!(actual_offsets.as_slice(), offsets); + assert_eq!(actual_ranges, ranges); + Ok(()) + } + + #[rstest] + #[case::missing(vec![0, 1])] + #[case::negative_length(vec![0, 10, 9])] + #[case::out_of_bounds(vec![0, 9, 11])] + fn rejects_invalid_boundaries(#[case] boundaries: Vec) { + assert!(selected_ranges(0..2, &boundaries, &Mask::new_true(2), 10).is_err()); + } +} diff --git a/vortex-layout/src/plan/physical/operators/list/work.rs b/vortex-layout/src/plan/physical/operators/list/work.rs new file mode 100644 index 00000000000..e93595683d6 --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/list/work.rs @@ -0,0 +1,411 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bounded grouping and admission of list element work. + +use std::ops::Range; + +use vortex_array::dtype::DType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::plan::physical::Concat; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::Zoned; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::overlapping; + +const MAX_ELEMENT_BYTES: u64 = 8 * 1024 * 1024; +const MAX_MASK_FRAGMENT_BYTES: u64 = 32 * 1024; +/// Retain a dense, fixed-width element span across null lists. The caller has validated every +/// offset and selected every row in this scope; only null rows may occupy gaps in `ranges`. +pub(super) fn preserve_null_span( + dtype: &DType, + offsets: &[u64], + selected: u64, + ranges: &mut Vec>, +) -> Option> { + if ranges.len() < 2 { + return None; + } + let start = ranges.first()?.start; + let end = ranges.last()?.end; + let span = end - start; + let width = match dtype { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return None, + }; + let validity = if dtype.is_nullable() { + span.div_ceil(8) + } else { + 0 + }; + if selected < span.div_ceil(2) + || span.saturating_mul(width).saturating_add(validity) > MAX_ELEMENT_BYTES + { + return None; + } + // Trim payload under leading/trailing nulls while retaining internal gaps. + let offsets = offsets + .iter() + .map(|offset| (*offset).clamp(start, end) - start) + .collect(); + ranges.clear(); + ranges.push(start..end); + Some(offsets) +} + +/// Combine nearby ranges into a mask, keeping holes out of the output without a scan per hole. +/// Only dense fixed-width spans of at most 8 MiB are grouped. Large contiguous fragments keep +/// their slice path: applying a mask costs more than the scan setup it saves for those ranges. +pub(super) fn element_splits(dtype: &DType, ranges: Vec>) -> VortexResult> { + if ranges.len() < 2 { + return Ok(ranges.into_iter().map(Split::all).collect()); + } + let width = match dtype { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(ranges.into_iter().map(Split::all).collect()), + }; + let mut splits = Vec::new(); + let mut first = 0; + while first < ranges.len() { + let start = ranges[first].start; + let mut selected = ranges[first].end - start; + let mut end = first + 1; + while let Some(next) = ranges.get(end) { + let span = next.end - start; + let validity = if dtype.is_nullable() { + span.div_ceil(8) + } else { + 0 + }; + let bytes = span.saturating_mul(width).saturating_add(validity); + let total = selected + (next.end - next.start); + if bytes > MAX_ELEMENT_BYTES || total < span.div_ceil(2) { + break; + } + selected = total; + end += 1; + } + if end == first + 1 + || selected + .saturating_mul(width) + .div_ceil((end - first) as u64) + > MAX_MASK_FRAGMENT_BYTES + { + splits.extend(ranges[first..end].iter().cloned().map(Split::all)); + } else { + let rows = start..ranges[end - 1].end; + let slices = ranges[first..end] + .iter() + .map(|range| { + Ok(( + usize::try_from(range.start - start)?, + usize::try_from(range.end - start)?, + )) + }) + .collect::>>()?; + let mask = Mask::from_slices(usize::try_from(rows.end - start)?, slices); + splits.push(Split { rows, mask }); + } + first = end; + } + Ok(splits) +} + +/// Overlap grouped rows only when all touched chunks fit the decoded-element estimate. +/// This limits additional range concurrency, not total heap use or coalesced IO buffers. +pub(super) fn can_overlap( + elements: &PlanRef, + ranges: &[Range], + selected: usize, + parallelism: usize, +) -> VortexResult { + if parallelism <= 1 || ranges.len() <= 1 || selected / ranges.len() < 2 { + return Ok(false); + } + let mut bytes = 0_u64; + for range in ranges { + let Some(estimate) = range_bytes(elements, range)? else { + return Ok(false); + }; + // Charging a shared chunk again is conservative and avoids a per-scan set. + bytes = bytes.saturating_add(estimate); + if bytes > MAX_ELEMENT_BYTES { + return Ok(false); + } + } + Ok(true) +} + +fn range_bytes(plan: &PlanRef, rows: &Range) -> VortexResult> { + if let Some(filter) = plan.as_opt::() { + return range_bytes(&filter.child_plan()?, rows); + } + if let Some(zoned) = plan.as_opt::() { + return match zoned.data_plan()? { + Some(data) => range_bytes(&data, rows), + None => Ok(None), + }; + } + if let Some(stats) = plan.as_opt::() { + return range_bytes(&stats.data_plan()?, rows); + } + if let Some(concat) = plan.as_opt::() { + let mut bytes = 0_u64; + for (index, start, local) in overlapping(concat, rows) { + let Some(estimate) = range_bytes( + &concat.child_required(index)?, + &(local.start - start..local.end - start), + )? + else { + return Ok(None); + }; + bytes = bytes.saturating_add(estimate); + if bytes > MAX_ELEMENT_BYTES { + break; + } + } + return Ok(Some(bytes)); + } + if !plan.is::() { + return Ok(None); + } + let width = match plan.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(None), + }; + let validity = if plan.dtype().is_nullable() { + plan.row_count().div_ceil(8) + } else { + 0 + }; + Ok(Some( + plan.row_count() + .saturating_mul(width) + .saturating_add(validity), + )) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::PType; + use vortex_mask::AllOr; + use vortex_session::registry::ReadContext; + + use super::*; + use crate::plan::physical::ConcatPlan; + use crate::plan::physical::LegacyStatsPlan; + use crate::plan::physical::SegmentScanPlan; + use crate::segments::SegmentId; + + fn segment(rows: u64, dtype: DType) -> PlanRef { + SegmentScanPlan::new(dtype, rows, SegmentId::from(0), ReadContext::new([]), None) + .into_plan() + } + + #[test] + fn retained_span_trims_null_edges_and_rebases_internal_offsets() { + let mut ranges = vec![10..18, 22..30]; + let offsets = preserve_null_span( + &DType::Primitive(PType::I64, true.into()), + &[0, 10, 18, 22, 30, 50], + 16, + &mut ranges, + ) + .unwrap(); + assert_eq!(offsets.as_slice(), &[0, 0, 8, 12, 20, 20]); + assert_eq!(ranges, vec![10..30]); + } + + #[rstest] + #[case::at_limit(MAX_ELEMENT_BYTES / 8, false, true)] + #[case::over_limit(MAX_ELEMENT_BYTES / 8 + 2, false, false)] + #[case::validity_counts(MAX_ELEMENT_BYTES / 8, true, false)] + fn retained_span_includes_validity_in_byte_estimate( + #[case] rows: u64, + #[case] nullable: bool, + #[case] retained: bool, + ) { + let original = vec![0..rows / 2, rows / 2 + 2..rows]; + let mut ranges = original.clone(); + let actual = preserve_null_span( + &DType::Primitive(PType::I64, nullable.into()), + &[0, rows / 2, rows / 2 + 2, rows], + rows - 2, + &mut ranges, + ); + assert_eq!(actual.is_some(), retained); + if !retained { + assert_eq!(ranges, original); + } + } + + #[rstest] + #[case::sparse(DType::Primitive(PType::I64, false.into()), vec![0..4, 100..104])] + #[case::variable_width(DType::Utf8(true.into()), vec![0..4, 6..10])] + #[case::one_range(DType::Bool(false.into()), vec![0..4])] + #[case::empty(DType::Bool(false.into()), vec![])] + fn unsuitable_null_spans_stay_compact(#[case] dtype: DType, #[case] ranges: Vec>) { + let mut actual = ranges.clone(); + let selected = ranges.iter().map(|range| range.end - range.start).sum(); + assert!(preserve_null_span(&dtype, &[], selected, &mut actual).is_none()); + assert_eq!(actual, ranges); + } + + #[rstest] + fn dense_ranges_use_relative_masks(#[values(0, u64::MAX - 20)] base: u64) -> VortexResult<()> { + let splits = element_splits( + &DType::Primitive(PType::I64, true.into()), + vec![base + 1..base + 5, base + 7..base + 9, base + 12..base + 13], + )?; + assert_eq!(splits.len(), 1); + assert_eq!(splits[0].rows, base + 1..base + 13); + assert_eq!( + splits[0].mask.slices(), + AllOr::Some(&[(0, 4), (6, 8), (11, 12)][..]) + ); + Ok(()) + } + + #[test] + fn sparse_gaps_and_unknown_widths_stay_separate() -> VortexResult<()> { + let ranges = vec![0..2, 100..104, 105..109]; + let splits = element_splits(&DType::Bool(false.into()), ranges.clone())?; + assert_eq!(splits.len(), 2); + assert_eq!(splits[0].rows, 0..2); + assert_eq!(splits[1].rows, 100..109); + assert_eq!(splits[1].mask.true_count(), 8); + let unknown = element_splits(&DType::Utf8(false.into()), ranges.clone())?; + assert_eq!( + unknown + .iter() + .map(|split| split.rows.clone()) + .collect::>(), + ranges + ); + Ok(()) + } + + #[rstest] + #[case::at_limit(MAX_ELEMENT_BYTES / 8, false, 1)] + #[case::over_limit(MAX_ELEMENT_BYTES / 8 + 2, false, 2)] + #[case::validity_counts(MAX_ELEMENT_BYTES / 8, true, 2)] + fn grouping_bounds_the_mask_span( + #[case] rows: u64, + #[case] nullable: bool, + #[case] count: usize, + ) -> VortexResult<()> { + let ranges = (0..rows) + .step_by(1024) + .map(|start| start + u64::from(start != 0)..(start + 1024).min(rows)) + .collect::>(); + let selected = rows - (ranges.len() as u64 - 1); + let splits = element_splits(&DType::Primitive(PType::I64, nullable.into()), ranges)?; + assert_eq!(splits.len(), count); + assert_eq!( + splits + .iter() + .map(|split| split.mask.true_count()) + .sum::(), + usize::try_from(selected)? + ); + Ok(()) + } + + #[test] + fn large_fragments_keep_contiguous_slices() -> VortexResult<()> { + let ranges = vec![0..8192, 8193..16384]; + let splits = element_splits(&DType::Primitive(PType::I64, false.into()), ranges.clone())?; + assert_eq!( + splits + .iter() + .map(|split| split.rows.clone()) + .collect::>(), + ranges + ); + assert!( + splits + .iter() + .all(|split| matches!(split.mask, Mask::AllTrue(_))) + ); + Ok(()) + } + + #[test] + #[allow(clippy::single_range_in_vec_init)] + fn contiguous_ranges_need_no_mask_allocation() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, false.into()); + assert!(element_splits(&dtype, vec![])?.is_empty()); + let splits = element_splits(&dtype, vec![0..MAX_ELEMENT_BYTES])?; + assert!(matches!(splits[0].mask, Mask::AllTrue(_))); + Ok(()) + } + + #[rstest] + #[case::grouped(4, 30, true)] + #[case::isolated(2, 30, false)] + #[case::single_worker(4, 1, false)] + fn overlap_requires_grouped_parallel_work( + #[case] selected: usize, + #[case] workers: usize, + #[case] expected: bool, + ) -> VortexResult<()> { + let elements = segment(1024, DType::Primitive(PType::I64, false.into())); + assert_eq!( + can_overlap(&elements, &[0..2, 10..12], selected, workers)?, + expected + ); + Ok(()) + } + + #[rstest] + #[case::within_limit(MAX_ELEMENT_BYTES / 16, false, true)] + #[case::large_chunk(MAX_ELEMENT_BYTES / 8, false, false)] + #[case::validity_counts(MAX_ELEMENT_BYTES / 16, true, false)] + fn charges_whole_chunks_and_validity( + #[case] rows: u64, + #[case] nullable: bool, + #[case] expected: bool, + ) -> VortexResult<()> { + let elements = segment(rows, DType::Primitive(PType::I64, nullable.into())); + assert_eq!(can_overlap(&elements, &[0..2, 10..12], 4, 30)?, expected); + Ok(()) + } + + #[test] + fn skips_untouched_chunks_and_zone_data() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, false.into()); + let data = ConcatPlan::try_new( + dtype.clone(), + vec![ + segment(1024, dtype.clone()), + segment(10_000_000, dtype.clone()), + ], + )? + .into_plan(); + let elements = + LegacyStatsPlan::new(data, segment(10_000_000, dtype), 1024, [].into()).into_plan(); + assert!(can_overlap(&elements, &[0..2, 10..12], 4, 30)?); + assert!(!can_overlap(&elements, &[0..2, 1024..1026], 4, 30)?); + Ok(()) + } + + #[test] + fn unknown_width_and_overflow_stay_sequential() -> VortexResult<()> { + let unknown = segment(10, DType::Utf8(false.into())); + let overflow = segment(u64::MAX, DType::Primitive(PType::I64, true.into())); + for plan in [unknown, overflow] { + assert!(!can_overlap(&plan, &[0..2, 5..7], 4, 30)?); + } + Ok(()) + } +} diff --git a/vortex-layout/src/plan/physical/operators/mod.rs b/vortex-layout/src/plan/physical/operators/mod.rs index b0f5937d71d..c4545ec5e19 100644 --- a/vortex-layout/src/plan/physical/operators/mod.rs +++ b/vortex-layout/src/plan/physical/operators/mod.rs @@ -9,6 +9,7 @@ mod concat; mod eval; mod filter; +mod list; mod once; mod pack; mod scan; @@ -18,6 +19,7 @@ mod whole; pub(crate) use concat::*; pub(crate) use eval::*; pub(crate) use filter::*; +pub(crate) use list::*; pub(crate) use once::*; pub(crate) use pack::*; pub(crate) use scan::*; diff --git a/vortex-layout/src/plan/physical/operators/pack.rs b/vortex-layout/src/plan/physical/operators/pack.rs index 4fefa368f05..7a4e83efc40 100644 --- a/vortex-layout/src/plan/physical/operators/pack.rs +++ b/vortex-layout/src/plan/physical/operators/pack.rs @@ -117,7 +117,7 @@ impl Operator for PackSource { } else { Validity::NonNullable }; - let array = assemble(&self.fields, arrays, validity).into_array(); + let array = assemble(&self.fields, arrays, len, validity).into_array(); Ok(Step::Final(array)) } } @@ -138,8 +138,9 @@ impl Operator for WrapStage { let Input::Batch(batch) = input else { return Ok(Step::End); }; + let len = batch.len(); Ok(Step::Final( - assemble(&self.fields, vec![batch], Validity::NonNullable).into_array(), + assemble(&self.fields, vec![batch], len, Validity::NonNullable).into_array(), )) } } @@ -149,8 +150,12 @@ impl Operator for WrapStage { /// The type is not derived or validated per batch: a `Pack` plan checks when it is built that /// each field's plan produces the field's dtype over the struct's rows, and a pack source takes /// the same rows from every field. -fn assemble(fields: &StructFields, arrays: Vec, validity: Validity) -> StructArray { - let len = arrays.first().map_or(0, |array| array.len()); +fn assemble( + fields: &StructFields, + arrays: Vec, + len: usize, + validity: Validity, +) -> StructArray { debug_assert!( arrays.iter().all(|array| array.len() == len) && fields diff --git a/vortex-layout/src/plan/physical/pipeline/children.rs b/vortex-layout/src/plan/physical/pipeline/children.rs new file mode 100644 index 00000000000..a2685079ed8 --- /dev/null +++ b/vortex-layout/src/plan/physical/pipeline/children.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selected plan views handed to the enclosing planner for execution. + +use std::fmt; + +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::Scan; +use super::Split; +use super::ValueCache; +use crate::plan::physical::PlanRef; + +/// A shared physical plan and its selected row domains. Constructing the view does not lower +/// children or compile a graph; the planner compiles only children touched by these splits. +#[derive(Debug)] +pub(crate) struct SelectedPlan { + pub(crate) plan: PlanRef, + pub(crate) splits: Vec, + max_active: usize, +} + +impl SelectedPlan { + pub(crate) fn new(plan: PlanRef, splits: Vec, max_active: usize) -> Self { + Self { + plan, + splits, + max_active, + } + } + + pub(crate) fn into_scan(self, session: VortexSession, values: ValueCache) -> VortexResult { + Ok(Scan::try_new(session, self.plan, self.splits)? + .with_value_cache(values) + .with_max_active(self.max_active)) + } +} + +/// A source's selected child projection, handed to its enclosing planner for execution. +pub struct ChildProjection { + pub(crate) id: u64, + pub(crate) view: SelectedPlan, + /// The planner must execute this view even when parallel subdivision is unsuitable. + pub(crate) planner_owned: bool, + pub(crate) values: ValueCache, +} + +impl ChildProjection { + /// Creates the selected child scan at the planner boundary, returning its completion id. + /// Drive this scan and pass its output to [`Scan::complete_child`] on the requesting scan. + pub fn into_scan(self, session: VortexSession) -> VortexResult<(u64, Scan)> { + Ok(( + self.id, + self.view + .into_scan(session, self.values)? + .with_child_planning(true), + )) + } +} + +impl fmt::Debug for ChildProjection { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ChildProjection") + .field("id", &self.id) + .field("view", &self.view) + .field("planner_owned", &self.planner_owned) + .finish_non_exhaustive() + } +} diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 74f5296cd19..cf2dcb10184 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -8,6 +8,7 @@ use std::ops::Range; use rustc_hash::FxHashMap; use smallvec::SmallVec; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -143,6 +144,10 @@ impl Compiler<'_> { &self.core.values } + pub(crate) fn plan_children(&self) -> bool { + self.core.plan_children + } + /// The rows `rows` of `scan`'s segment, filtered by `filter`, read by the split compiling. /// A segment several readers read is decoded once and shared. pub(crate) fn scan( @@ -296,6 +301,30 @@ impl Shares { }) } + /// Counts a single split's readers without lowering chunks its mask excludes. + pub(crate) fn add_selected( + &mut self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + ) -> VortexResult<()> { + if mask.all_false() { + return Ok(()); + } + if mask.all_true() { + return self.add(plan, rows); + } + let at = Reach::Selected { + offset: 0, + rows: rows.clone(), + mask: mask.clone(), + }; + let reaches = &mut self.reaches; + plan.reach(rows, &at, &mut |segment, rows| { + reaches.push((segment, rows)) + }) + } + /// Keeps only the segments more than one reader reads. The rest are read by the one /// pipeline that needs them. pub(crate) fn retain_shared(&mut self) { @@ -350,13 +379,24 @@ pub enum Reach { Offset(u64), /// Every row is read for these scanned rows, as a dictionary's values are for its codes. Fixed(Range), + /// A linear mapping restricted to selected rows in the scan's root domain. + Selected { + /// Root row index of this plan's row zero. + offset: u64, + /// The root row range the mask covers. + rows: Range, + /// Selected positions relative to `rows.start`. + mask: Mask, + }, } impl Reach { /// The scanned rows `rows` are read for. pub fn root(&self, rows: &Range) -> Range { match self { - Reach::Offset(offset) => rows.start + offset..rows.end + offset, + Reach::Offset(offset) | Reach::Selected { offset, .. } => { + rows.start + offset..rows.end + offset + } Reach::Fixed(root) => root.clone(), } } @@ -366,6 +406,11 @@ impl Reach { match self { Reach::Offset(offset) => Reach::Offset(offset + by), Reach::Fixed(root) => Reach::Fixed(root.clone()), + Reach::Selected { offset, rows, mask } => Reach::Selected { + offset: offset + by, + rows: rows.clone(), + mask: mask.clone(), + }, } } @@ -373,6 +418,30 @@ impl Reach { pub fn fixed(&self, rows: &Range) -> Self { Reach::Fixed(self.root(rows)) } + + /// Whether the selection needs any row of this child, before lowering it. + pub(crate) fn selects(&self, local: &Range) -> bool { + let Reach::Selected { offset, rows, mask } = self else { + return true; + }; + let start = (local.start + offset).max(rows.start); + let end = (local.end + offset).min(rows.end); + if start >= end { + return false; + } + let start = usize::try_from(start - rows.start).vortex_expect("offset fits selection mask"); + let end = usize::try_from(end - rows.start).vortex_expect("offset fits selection mask"); + match mask { + Mask::AllTrue(_) => true, + Mask::AllFalse(_) => false, + Mask::Values(values) => match values.cached_indices() { + Some(indices) => indices + .get(indices.partition_point(|&index| index < start)) + .is_some_and(|&index| index < end), + None => values.bit_buffer().count_range(start, end) > 0, + }, + } + } } /// The chunks of `concat` overlapping `rows`: each chunk's index, where it starts, and the rows diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs index b6037c9e6ea..c6deed171e8 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -38,12 +38,15 @@ //! read, and a claimed port is dropped once its source has drained it, so a decoded segment //! lives only as long as a split that may still read it. +mod children; mod compile; mod port; mod scan; pub mod synthetic; mod values; +pub use self::children::ChildProjection; +pub(crate) use self::children::SelectedPlan; pub use self::compile::Chain; pub use self::compile::Compiler; pub use self::compile::Reach; diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 9d650ebbd26..9ed12d08a5e 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -31,6 +31,7 @@ use super::port::Slab; use super::port::SourceEnd; use super::values::ValueCache; use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::ChildInbox; use crate::plan::physical::operator::Cx; use crate::plan::physical::operator::InletSet; use crate::plan::physical::operator::Input; @@ -39,6 +40,7 @@ use crate::plan::physical::operator::Source; use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; +use crate::plan::physical::pipeline::ChildProjection; use crate::segments::SegmentId; #[cfg(debug_assertions)] @@ -82,9 +84,11 @@ impl Split { pub enum Turn { /// Read this segment and [`deliver`](Scan::deliver) its bytes. Read(ReadRequest), + /// Plan and execute this child projection, then return its arrays to the scan. + Child(ChildProjection), /// The next rows of split `split`, in row order within the split. Output(usize, ArrayRef), - /// Nothing can run until a read is delivered. + /// Nothing can run until a read or child projection completes. Waiting, /// Every split has produced all its rows. Done, @@ -111,6 +115,7 @@ pub(crate) struct Pipeline { /// Reads it asked for that have not arrived, so retiring it cancels them only when any /// are left. pending: usize, + child_inbox: Option>, /// Once its source has finished, the next stage to finish: stages are finished one at a /// time, each when the outlets have room for what it flushes. finishing: Option, @@ -141,6 +146,7 @@ fn drive( session: &VortexSession, exec: &mut ExecutionCtx, reads: &mut SmallVec<[SegmentId; 1]>, + child_requests: &mut Vec, ) -> VortexResult { let Pipeline { source, @@ -153,6 +159,7 @@ fn drive( source_more, touched, changed, + child_inbox, .. } = pipe; let mut cx = Cx { @@ -164,6 +171,8 @@ fn drive( exec, touched, changed, + child_requests, + child_inbox, }; loop { // Closing needs no room, so a pipeline with nothing left to flush closes at once. @@ -191,9 +200,17 @@ fn drive( let ready = source.ready(&mut cx)?; cx.changed.clear(); vortex_ensure!( - cx.reads.is_empty() || ready == SourceReady::Blocked(SourceBlocked::Io), + cx.reads.is_empty() + || matches!( + ready, + SourceReady::Blocked(SourceBlocked::Io | SourceBlocked::Child) + ), "A source that queues reads must block on IO" ); + vortex_ensure!( + cx.child_requests.is_empty() || ready == SourceReady::Blocked(SourceBlocked::Child), + "A source that queues asynchronous work must block on its completion" + ); match ready { SourceReady::Ready => {} SourceReady::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), @@ -281,6 +298,9 @@ pub(crate) struct Core { reads: FxHashMap, SegmentId)>, next_read: u64, new_reads: VecDeque, + child_requests: FxHashMap, + next_child: u64, + new_children: VecDeque, /// Split slots whose stage output has something new, with a flag per slot so a slot is /// queued once. dirty: Vec, @@ -291,6 +311,12 @@ pub(crate) struct Core { pub(crate) split: usize, /// Whole values kept for later scans of the same split. pub(crate) values: ValueCache, + pub(crate) plan_children: bool, +} + +struct PendingChild { + owner: Option, + local_id: u64, } impl Core { @@ -306,11 +332,15 @@ impl Core { reads: FxHashMap::default(), next_read: 0, new_reads: VecDeque::new(), + child_requests: FxHashMap::default(), + next_child: 0, + new_children: VecDeque::new(), dirty: Vec::new(), dirty_flags: Vec::new(), shares: Shares::new(splits), split: usize::MAX, values: ValueCache::default(), + plan_children: false, } } @@ -331,6 +361,7 @@ impl Core { outlets, bytes: VecDeque::new(), pending: 0, + child_inbox: None, finishing: None, touched, changed, @@ -405,16 +436,21 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let mut reads = SmallVec::new(); + let mut child_requests = Vec::new(); let progress = drive( pipe, &mut self.arena, &self.session, &mut self.exec, &mut reads, + &mut child_requests, )?; for segment_id in reads { self.ask(id, segment_id); } + for child in child_requests { + self.ask_child(id, child)?; + } #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { @@ -422,6 +458,16 @@ impl Core { Progress::Blocked(SourceBlocked::Io) if !self.pipeline(id).bytes.is_empty() => { self.make_runnable(id); } + Progress::Blocked(SourceBlocked::Child) + if !self.pipeline(id).bytes.is_empty() + || self + .pipeline(id) + .child_inbox + .as_ref() + .is_some_and(|inbox| !inbox.completed.is_empty()) => + { + self.make_runnable(id); + } // Nothing changes a port while its source's pipeline runs, so the condition the run reported // holds until a neighbour wakes it. Progress::Blocked(blocked) => self.pipeline_mut(id).state = State::Blocked(*blocked), @@ -441,6 +487,52 @@ impl Core { self.pipeline_mut(id).pending += 1; } + fn ask_child(&mut self, owner: PipelineId, child: ChildProjection) -> VortexResult<()> { + vortex_ensure!( + !self + .child_requests + .values() + .any(|pending| pending.owner == Some(owner) && pending.local_id == child.id), + "A source published child {} twice", + child.id + ); + let id = self.next_child; + self.next_child = id + .checked_add(1) + .ok_or_else(|| vortex_err!("Scan child identifiers exhausted"))?; + self.child_requests.insert( + id, + PendingChild { + owner: Some(owner), + local_id: child.id, + }, + ); + self.pipeline_mut(owner) + .child_inbox + .get_or_insert_with(Default::default) + .pending += 1; + self.new_children.push_back(ChildProjection { id, ..child }); + Ok(()) + } + + fn complete_child(&mut self, id: u64, arrays: Option>) -> VortexResult<()> { + let child = self + .child_requests + .remove(&id) + .ok_or_else(|| vortex_err!("Unknown scan child {id}"))?; + if let Some(owner) = child.owner { + let inbox = self + .pipeline_mut(owner) + .child_inbox + .as_mut() + .ok_or_else(|| vortex_err!("Child completed for a source with no child inbox"))?; + inbox.pending -= 1; + inbox.completed.push_back((child.local_id, arrays)); + self.make_runnable(owner); + } + Ok(()) + } + /// Wakes the sources of what pipeline `id` pushed and the sinks of what it took. fn wake_neighbours(&mut self, id: PipelineId) { for index in 0..self.pipeline(id).outlets.len() { @@ -529,6 +621,13 @@ impl Core { if pipe.pending > 0 { self.cancel_reads(id); } + if pipe.child_inbox.is_some_and(|inbox| inbox.pending > 0) { + for child in self.child_requests.values_mut() { + if child.owner == Some(id) { + child.owner = None; + } + } + } self.free_pipelines.push(id); } } @@ -671,15 +770,40 @@ impl Scan { 0, splits.iter().map(|split| split.rows.clone()).collect(), ); - // Only the rows some split reads can share a segment between readers. - let start = splits - .iter() - .map(|split| split.rows.start) - .min() - .unwrap_or(0); - let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0); match &root { - Root::Plan(plan) => core.shares.add(plan, start..end)?, + Root::Plan(plan) => { + if let [split] = splits.as_slice() { + core.shares + .add_selected(plan, split.rows.clone(), &split.mask)?; + } else if splits + .windows(2) + .all(|pair| pair[0].rows.end <= pair[1].rows.start) + { + // List selections can leave large gaps between element ranges. Keep those + // chunks lazy, merging adjacent ranges so dense scans still visit once. + let mut rows = 0..0; + for split in &splits { + if split.rows.start > rows.end { + if !rows.is_empty() { + core.shares.add(plan, rows.clone())?; + } + rows.start = split.rows.start; + } + rows.end = split.rows.end; + } + if !rows.is_empty() { + core.shares.add(plan, rows)?; + } + } else { + let start = splits + .iter() + .map(|split| split.rows.start) + .min() + .unwrap_or(0); + let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0); + core.shares.add(plan, start..end)?; + } + } } core.shares.retain_shared(); core.dirty_flags = vec![false; 1]; @@ -714,6 +838,11 @@ impl Scan { self } + pub(crate) fn with_child_planning(mut self, enabled: bool) -> Self { + self.core.plan_children = enabled; + self + } + /// Runs pipelines until the owner has something to do. pub fn step(&mut self) -> VortexResult { loop { @@ -723,6 +852,18 @@ impl Scan { if let Some(read) = self.core.new_reads.pop_front() { return Ok(Turn::Read(read)); } + while let Some(child) = self.core.new_children.pop_front() { + if self + .core + .child_requests + .get(&child.id) + .is_some_and(|pending| pending.owner.is_none()) + { + self.core.child_requests.remove(&child.id); + continue; + } + return Ok(Turn::Child(child)); + } if let Some(slot) = self.core.dirty.pop() { self.core.dirty_flags[slot] = false; self.advance(slot)?; @@ -741,7 +882,7 @@ impl Scan { if self.live == 0 { return Ok(Turn::Done); } - if !self.core.reads.is_empty() { + if !self.core.reads.is_empty() || !self.core.child_requests.is_empty() { return Ok(Turn::Waiting); } vortex_bail!( @@ -756,6 +897,12 @@ impl Scan { self.core.deliver(read, bytes) } + /// Delivers a child projection to its source, with arrays in selected-split order. + /// `None` declines optional parallel work; a planner-owned selected view requires `Some`. + pub fn complete_child(&mut self, id: u64, arrays: Option>) -> VortexResult<()> { + self.core.complete_child(id, arrays) + } + /// Starts the next split in a free slot. fn activate(&mut self) -> VortexResult<()> { let slot = self diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index df404bd34ea..d1749972d65 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -88,6 +88,12 @@ impl Core { pipe.pending > 0 || !pipe.bytes.is_empty(), "pipeline {id} waits for a read with none in flight or delivered" ), + Progress::Blocked(SourceBlocked::Child) => assert!( + pipe.child_inbox + .as_ref() + .is_some_and(|inbox| inbox.pending > 0 || !inbox.completed.is_empty()), + "pipeline {id} waits for a child with none in flight or completed" + ), Progress::Done => { for (index, &port) in pipe.inlets.iter().enumerate() { assert!( diff --git a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index 19bcb5504db..e8bd51e1ac9 100644 --- a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -64,6 +64,7 @@ fn run( events.push(None); arrays.push(array); } + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { if inflight.is_empty() { return Err(vortex_err!("scan waits with no reads in flight")); @@ -193,6 +194,7 @@ fn splits_produce_their_own_rows() -> VortexResult<()> { match scan.step()? { Turn::Read(read) => inflight.push(read), Turn::Output(split, array) => by_split[split].push(array), + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { let read: ReadRequest = inflight.remove(0); scan.deliver(read.id, empty_segment())?; @@ -258,6 +260,7 @@ fn a_source_has_several_reads_in_flight(#[values(1, 3)] deliveries: usize) -> Vo match scan.step()? { Turn::Read(read) => inflight.push(read), Turn::Output(_, array) => arrays.push(array), + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { if first_wait { assert_eq!( @@ -417,6 +420,7 @@ fn finishing_cancels_reads_in_flight() -> VortexResult<()> { match scan.step()? { Turn::Read(read) => inflight.push(read), Turn::Output(split, array) => by_split[split].push(array), + Turn::Child(_) => unreachable!("child planning is disabled"), Turn::Waiting => { let read = inflight.remove(0); scan.deliver(read.id, empty_segment())?; diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 8fc9c67b900..c5460cef9c5 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -6,6 +6,7 @@ use std::collections::VecDeque; use std::ops::Range; +use std::sync::Arc; use rstest::rstest; use vortex_array::ArrayContext; @@ -14,9 +15,12 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::StructFields; @@ -25,6 +29,7 @@ use vortex_array::expr::gt; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::serde::SerializeOptions; +use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; @@ -32,9 +37,12 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::Mask; +use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; +use vortex_utils::parallelism::get_available_parallelism; use super::port::SourceEnd; +use super::synthetic::SyntheticPlan; use super::synthetic::row_dtype; use super::*; use crate::LayoutRef; @@ -42,9 +50,13 @@ use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::dict::DictLayout; use crate::layouts::flat::FlatLayout; +use crate::layouts::foreign::new_foreign_layout; use crate::layouts::struct_::StructLayout; +use crate::plan::physical::ConcatPlan; use crate::plan::physical::EvalPlan; use crate::plan::physical::Filter; +use crate::plan::physical::ListPack; +use crate::plan::physical::ListPackPlan; use crate::plan::physical::PackPlan; use crate::plan::physical::PlanRef; use crate::plan::physical::SegmentScan; @@ -198,8 +210,16 @@ struct Run { /// Drives a scan the way an owner does, completing reads only while the scan waits. fn drive( store: &Store, - mut scan: Scan, + scan: Scan, mut pick: impl FnMut(&[ReadRequest]) -> usize, +) -> VortexResult { + drive_inner(store, scan, &mut pick) +} + +fn drive_inner( + store: &Store, + mut scan: Scan, + pick: &mut dyn FnMut(&[ReadRequest]) -> usize, ) -> VortexResult { let mut inflight: Vec = Vec::new(); let mut batch: Vec = Vec::new(); @@ -224,6 +244,19 @@ fn drive( splits[split].push(array.clone()); arrays.push(array); } + Turn::Child(child) => { + let (id, child_scan) = child.into_scan(SESSION.clone())?; + let child = drive_inner(store, child_scan, pick)?; + assert_eq!(child.leftover, (0, 0)); + events.extend( + child + .events + .into_iter() + .filter(|event| !matches!(event, Event::Piece(_))), + ); + let arrays = child.splits.into_iter().flatten().collect(); + scan.complete_child(id, Some(arrays))?; + } Turn::Waiting => { if !batch.is_empty() { events.push(Event::Io(std::mem::take(&mut batch))); @@ -288,6 +321,171 @@ fn delivery(order: Delivery) -> impl FnMut(&[ReadRequest]) -> usize { } } +#[rstest] +fn grouped_list_ranges_preserve_order_and_shared_reads( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] shared: bool, + #[values(false, true)] planned: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = Buffer::from_iter(0_i64..20).into_array(); + let offsets = Buffer::from_iter((0_u64..=10).map(|row| row * 2)).into_array(); + let input = + ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?.into_array(); + let layout = if shared { + store.flat(&elements)? + } else { + store.chunked(&elements, &[5, 5, 5, 5])? + }; + let plan = ListPackPlan::try_new( + false.into(), + 10, + lower(&layout)?, + lower(&store.flat(&offsets)?)?, + None, + )? + .into_plan(); + let plan = if planned { selected_list(plan) } else { plan }; + let rows = [0_u64, 1, 7, 8]; + let mask = Mask::from_iter((0..10).map(|row| rows.contains(&row))); + let result = run(&store, &plan, 0..10, mask, delivery(order))?; + let actual = ChunkedArray::try_new(result.arrays, input.dtype().clone())?.into_array(); + let expected = input.take(Buffer::from_iter(rows).into_array())?; + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); + if get_available_parallelism().unwrap_or(1) > 1 { + if shared { + let reads = result + .events + .iter() + .filter(|event| **event == Event::Delivered(0)); + assert_eq!(reads.count(), 1); + } else { + assert!( + result.events.iter().any(|event| { + matches!(event, Event::Io(ids) if ids.contains(&0) && ids.contains(&2)) + }), + "{:?}", + result.events + ); + } + } + Ok(()) +} + +fn selected_list(plan: PlanRef) -> PlanRef { + SyntheticPlan::new( + plan.dtype().clone(), + plan.row_count(), + Vec::new(), + move |rows, mask, compiler| { + plan.as_::() + .compile_selected(rows, mask, compiler, true) + }, + ) + .into_plan() +} + +#[rstest] +fn selected_list_view_skips_unselected_children( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] empty: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = Buffer::from_iter(10_i64..12).into_array(); + let selected = lower(&store.flat(&elements)?)?; + let unused = SyntheticPlan::new(elements.dtype().clone(), 2, Vec::new(), |_, _, _| { + Err(vortex_err!("Unselected child must not be compiled")) + }) + .into_plan(); + let chunks = ConcatPlan::try_new( + elements.dtype().clone(), + vec![unused.clone(), selected, unused], + )?; + let offsets = Buffer::from(vec![0_u64, 2, 4, 6]).into_array(); + let plan = selected_list( + ListPackPlan::try_new( + false.into(), + 3, + chunks.into_plan(), + lower(&store.flat(&offsets)?)?, + None, + )? + .into_plan(), + ); + let result = run( + &store, + &plan, + 0..3, + Mask::from_iter([false, !empty, false]), + delivery(order), + )?; + let expected = ListArray::try_new( + elements, + Buffer::from(vec![0_u64, 2]).into_array(), + Validity::NonNullable, + )? + .into_array(); + let expected = if empty { + expected.slice(0..0)? + } else { + expected + }; + let actual = ChunkedArray::try_new(result.arrays, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); + if !empty { + assert_eq!(segments_read(&result.events), vec![0, 1]); + } + Ok(()) +} + +#[rstest] +fn dense_nullable_list_ranges_share_one_scan( + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, + #[values(false, true)] partial: bool, + #[values(false, true)] planned: bool, +) -> VortexResult<()> { + let mut store = Store::default(); + let elements = Buffer::from_iter(0_i64..20).into_array(); + let offsets = Buffer::from_iter((0_u64..=10).map(|row| row * 2)).into_array(); + let valid = BoolArray::from_iter((0..10).map(|row| row != 2 && row != 7)).into_array(); + let input = ListArray::try_new( + elements.clone(), + offsets.clone(), + Validity::Array(valid.clone()), + )? + .into_array(); + let plan = ListPackPlan::try_new( + true.into(), + 10, + lower(&store.chunked(&elements, &[5, 5, 5, 5])?)?, + lower(&store.flat(&offsets)?)?, + Some(lower(&store.flat(&valid)?)?), + )? + .into_plan(); + let plan = if planned { selected_list(plan) } else { plan }; + let mask = Mask::from_iter((0..10).map(|row| !partial || row != 4)); + let result = run(&store, &plan, 0..10, mask.clone(), delivery(order))?; + if !partial { + assert_eq!(result.arrays[0].as_::().elements().len(), 20); + } + let actual = ChunkedArray::try_new(result.arrays, input.dtype().clone())?.into_array(); + assert_arrays_eq!( + actual, + input.filter(mask)?, + &mut SESSION.create_execution_ctx() + ); + assert!( + result + .events + .iter() + .any(|event| { matches!(event, Event::Io(ids) if (0..4).all(|id| ids.contains(&id))) }), + "{:?}", + result.events + ); + assert_eq!(segments_read(&result.events), vec![0, 1, 2, 3, 4, 5]); + Ok(()) +} + /// Delivers segments in the given order. fn scripted(order: &[u32]) -> impl FnMut(&[ReadRequest]) -> usize + '_ { let mut next = order.iter(); @@ -432,6 +630,135 @@ fn splits_of_one_scan_read_each_segment_once( Ok(()) } +#[rstest] +#[case::gaps(vec![0..2, 8..10, 18..20])] +#[case::overlapping(vec![0..8, 4..12, 10..20])] +#[case::unordered(vec![15..20, 0..4, 7..12])] +#[case::empty(vec![0..0, 4..8, 20..20])] +fn split_ranges_preserve_shared_reads( + #[case] ranges: Vec>, + #[values(1, 3)] active: usize, + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + let splits: Vec<_> = ranges.into_iter().map(Split::all).collect(); + let scan = Scan::try_new(SESSION.clone(), plan, splits.clone())?.with_max_active(active); + let run = drive(&store, scan, delivery(order))?; + assert_eq!(run.leftover, (0, 0)); + let reads = segments_read(&run.events); + let mut unique = reads.clone(); + unique.dedup(); + assert_eq!(reads, unique, "each shared segment is read once"); + for (index, split) in splits.iter().enumerate() { + assert_view( + &expected, + &split.rows, + &split.mask, + run.splits.get(index).cloned().unwrap_or_default(), + )?; + } + Ok(()) +} + +#[rstest] +fn chunk_in_gap_between_splits_stays_lazy(#[values(1, 2)] active: usize) -> VortexResult<()> { + static UNKNOWN_ID: CachedId = CachedId::new("test.unread-gap"); + let mut store = Store::default(); + let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let known = store.flat(&values)?; + let unknown = new_foreign_layout( + *UNKNOWN_ID, + values.dtype().clone(), + ROWS, + Vec::new(), + Vec::new(), + Vec::new(), + ); + let layout = ChunkedLayout::new( + 3 * ROWS, + values.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![Arc::clone(&known), unknown, known]), + ) + .into_layout(); + let scan = Scan::try_new( + SESSION.clone(), + lower(&layout)?, + vec![Split::all(0..2), Split::all(2 * ROWS..2 * ROWS + 2)], + )? + .with_max_active(active); + let run = drive(&store, scan, delivery(Delivery::Lifo))?; + assert_eq!(segments_read(&run.events), vec![0]); + assert_eq!(run.leftover, (0, 0)); + assert_eq!(run.splits.len(), 2); + for pieces in run.splits { + assert_view(&values, &(0..2), &Mask::new_true(2), pieces)?; + } + Ok(()) +} + +#[rstest] +fn chunk_excluded_by_mask_stays_lazy( + #[values(false, true)] nested: bool, + #[values(false, true)] empty: bool, + #[values(false, true)] cached_indices: bool, +) -> VortexResult<()> { + static UNKNOWN_ID: CachedId = CachedId::new("test.masked-gap"); + let mut store = Store::default(); + let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let known = store.flat(&values)?; + let unknown = new_foreign_layout( + *UNKNOWN_ID, + values.dtype().clone(), + ROWS, + Vec::new(), + Vec::new(), + Vec::new(), + ); + let mut layout = ChunkedLayout::new( + 3 * ROWS, + values.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![Arc::clone(&known), unknown, Arc::clone(&known)]), + ) + .into_layout(); + let offset = if nested { ROWS } else { 0 }; + if nested { + layout = ChunkedLayout::new( + 4 * ROWS, + values.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![known, layout]), + ) + .into_layout(); + } + let rows = offset + 5..offset + 3 * ROWS - 5; + let len = usize::try_from(rows.end - rows.start)?; + let mut mask = if empty { + Mask::new_false(len) + } else { + Mask::from_indices(len, [0, len - 1]) + }; + if !cached_indices { + mask = Mask::from_buffer(mask.to_bit_buffer()); + } + let scan = Scan::try_new(SESSION.clone(), lower(&layout)?, vec![Split { rows, mask }])?; + let run = drive(&store, scan, delivery(Delivery::Lifo))?; + assert_eq!(run.leftover, (0, 0)); + if empty { + assert!(segments_read(&run.events).is_empty()); + assert!(run.splits.iter().all(Vec::is_empty)); + } else { + assert_eq!(segments_read(&run.events), vec![0]); + let expected = PrimitiveArray::from_iter([5, ROWS as i32 - 6]).into_array(); + assert_view( + &expected, + &(0..2), + &Mask::new_true(2), + run.splits[0].clone(), + )?; + } + Ok(()) +} + #[test] fn every_read_is_issued_before_any_is_delivered() -> VortexResult<()> { let mut store = Store::default(); @@ -762,7 +1089,11 @@ fn pack_readiness_tracks_changes_and_consumption( let mut exec = SESSION.create_execution_ctx(); let mut touched = InletSet::new(2); let mut changed = InletSet::all(2); + let mut child_requests = Vec::new(); + let mut child_inbox = None; let mut cx = Cx { + child_requests: &mut child_requests, + child_inbox: &mut child_inbox, arena: &mut arena, inlets: &inlets, bytes: &mut bytes, @@ -834,3 +1165,29 @@ fn pack_readiness_tracks_changes_and_consumption( } Ok(()) } + +#[rstest] +fn nullable_struct_without_fields_keeps_its_rows( + #[values(Sel::All, Sel::EveryOther)] selection: Sel, +) -> VortexResult<()> { + let mut store = Store::default(); + let validity = BoolArray::from_iter((0..ROWS).map(|row| row % 3 != 0)).into_array(); + let expected = StructArray::try_new_with_dtype( + Vec::new(), + StructFields::empty(), + ROWS as usize, + Validity::Array(validity.clone()), + )? + .into_array(); + let layout = StructLayout::new(ROWS, expected.dtype().clone(), vec![store.flat(&validity)?]) + .into_layout(); + let mask = selection.mask(ROWS as usize); + let run = run( + &store, + &lower(&layout)?, + 0..ROWS, + mask.clone(), + delivery(Delivery::Fifo), + )?; + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} diff --git a/vortex-layout/src/plan/physical/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs index b2e1ffea0f7..377e4bf270c 100644 --- a/vortex-layout/src/plan/physical/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -189,6 +189,9 @@ impl PlanVTable for Concat { visit: &mut dyn FnMut(SegmentId, Range), ) -> VortexResult<()> { for (index, start, local) in overlapping(plan, &rows) { + if !at.selects(&local) { + continue; + } plan.child_required(index)?.reach( local.start - start..local.end - start, &at.shift(start), diff --git a/vortex-layout/src/plan/physical/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs index bb075c7f4e2..4be082730ae 100644 --- a/vortex-layout/src/plan/physical/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use std::sync::Arc; use vortex_array::EmptyMetadata; @@ -10,6 +11,8 @@ use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::plan::physical::Plan; @@ -18,6 +21,12 @@ use crate::plan::physical::PlanId; use crate::plan::physical::PlanParts; use crate::plan::physical::PlanRef; use crate::plan::physical::PlanVTable; +use crate::plan::physical::operators::ListSource; +use crate::plan::physical::operators::UNBOUNDED; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; +use crate::segments::SegmentId; const ELEMENTS: usize = 0; const OFFSETS: usize = 1; @@ -78,6 +87,39 @@ impl ListPackPlan { Ok(unsafe { Self::from_children_unchecked(dtype, row_count, children) }) } + pub(crate) fn compile_selected( + &self, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + planner_owned: bool, + ) -> VortexResult> { + let plan = self; + if mask.all_false() { + return Ok(None); + } + let boundaries = offset_mask(mask); + let offsets = compiler + .compile(&plan.offsets()?, rows.start..rows.end + 1, &boundaries)? + .ok_or_else(|| vortex_err!("Selected list offsets produced no rows"))?; + let mut chains = vec![offsets]; + if let Some(validity) = plan.validity()? { + chains.push( + compiler + .compile(&validity, rows, mask)? + .ok_or_else(|| vortex_err!("Selected list validity produced no rows"))?, + ); + } + let source = ListSource::new( + plan.clone(), + mask.clone(), + compiler.values().clone(), + compiler.plan_children(), + planner_owned, + ); + Ok(Some(compiler.join_with(chains, |_| UNBOUNDED, source))) + } + /// Returns the plan producing list elements. pub fn elements(&self) -> VortexResult { self.child_required(ELEMENTS) @@ -124,6 +166,57 @@ impl PlanVTable for ListPack { _ => Cow::Owned(format!("child[{index}]")), } } + + fn compile( + plan: &Plan, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + plan.compile_selected( + rows, + mask, + compiler, + std::env::var_os("VORTEX_LIST_PIPELINE").is_some(), + ) + } + + fn reach( + plan: &Plan, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + // Elements have a different row domain, known only after the selected offsets arrive. + // Their scan records its own segment readers when those ranges have been resolved. + let root_rows = at.root(&rows); + plan.offsets()? + .reach(rows.start..rows.end + 1, &at.fixed(&rows), &mut |id, _| { + visit(id, root_rows.clone()) + })?; + if let Some(validity) = plan.validity()? { + validity.reach(rows, at, visit)?; + } + Ok(()) + } +} + +/// A selected list needs both its starting and ending offsets. +fn offset_mask(mask: &Mask) -> Mask { + match mask.indices() { + AllOr::None => Mask::new_false(mask.len() + 1), + AllOr::All => Mask::new_true(mask.len() + 1), + AllOr::Some(indices) => { + let mut boundaries = Vec::with_capacity(indices.len() * 2); + for &index in indices { + if boundaries.last() != Some(&index) { + boundaries.push(index); + } + boundaries.push(index + 1); + } + Mask::from_indices(mask.len() + 1, boundaries) + } + } } fn validate_children(dtype: &DType, row_count: u64, children: &PlanChildren) -> VortexResult<()> { diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index da6fa56ccb9..2b74e23d17a 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -4,7 +4,10 @@ use std::fmt; use std::num::NonZeroUsize; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; @@ -43,6 +46,7 @@ use vortex_session::registry::ReadContext; use super::*; use crate::LayoutBuildContext; +use crate::LayoutChildren; use crate::LayoutEncoding; use crate::LayoutRef; use crate::OwnedLayoutChildren; @@ -150,6 +154,60 @@ fn chunked_plan_exposes_chunks() -> VortexResult<()> { Ok(()) } +#[derive(Clone)] +struct CountedChunkRows { + rows: Arc<[u64]>, + calls: Arc, +} + +impl LayoutChildren for CountedChunkRows { + fn to_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn child(&self, _idx: usize, _dtype: &DType) -> VortexResult { + Err(vortex_err!("Lowering must leave chunk children lazy")) + } + + fn child_row_count(&self, idx: usize) -> u64 { + self.calls.fetch_add(1, Ordering::Relaxed); + self.rows[idx] + } + + fn nchildren(&self) -> usize { + self.rows.len() + } +} + +#[rstest] +#[case::normal(vec![3, 2], vec![0, 3])] +#[case::empty_chunks(vec![0, 2, 0, 3, 0], vec![0, 0, 2, 2, 5])] +#[case::empty(vec![], vec![])] +#[case::zero_rows(vec![0], vec![0])] +fn lowering_reuses_validated_chunk_boundaries( + #[case] rows: Vec, + #[case] offsets: Vec, +) -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let row_count = rows.iter().sum(); + let layout = ChunkedLayout::new( + row_count, + DType::Null, + Arc::new(CountedChunkRows { + rows: rows.into(), + calls: Arc::clone(&calls), + }), + ) + .into_layout(); + calls.store(0, Ordering::Relaxed); + let plan = lower(&layout)?; + assert_eq!(plan.as_::().row_offsets(), offsets); + assert_eq!(plan.row_count(), row_count); + assert_eq!(plan.child_count(), offsets.len()); + assert_eq!(calls.load(Ordering::Relaxed), 0); + Ok(()) +} + #[test] fn chunked_plan_lowers_each_chunk_on_access() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index fc6053d55ca..dce9beb9743 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -19,6 +19,7 @@ use vortex_io::request::IoIntent; use vortex_io::request::IoResult; use vortex_scan::planning::planner::State; +use crate::plan::physical::pipeline::ChildProjection; use crate::plan::physical::pipeline::ReadId; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Turn; @@ -28,7 +29,9 @@ use crate::segments::SegmentId; /// Rows a running plan produced, in row order. #[derive(Clone)] pub(super) struct Piece { - /// Orders the piece among the outputs of one scan. + /// The selected split that produced this piece. + pub(super) split: usize, + /// Orders the piece within its split. pub(super) rows: Range, /// The selected rows the piece holds. pub(super) array: ArrayRef, @@ -40,6 +43,8 @@ pub(super) enum AdaptorStep { Yield, /// Requests to publish, each returned once. NeedsIO(IoBatch), + /// A child projection for the enclosing planner to schedule. + Child(ChildProjection), /// Rows the plan produced. Piece(Piece), } @@ -96,7 +101,8 @@ impl ScanAdaptor { } Ok(match self.step()? { Turn::Waiting | Turn::Done => AdaptorStep::Yield, - Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Output(split, array) => AdaptorStep::Piece(self.piece(split, array)), + Turn::Child(child) => AdaptorStep::Child(child), 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)?]; @@ -125,20 +131,33 @@ impl ScanAdaptor { match &turn { Turn::Done => self.done = true, Turn::Waiting => self.pending = Some(Turn::Waiting), - Turn::Read(_) | Turn::Output(..) => {} + Turn::Read(_) | Turn::Child(_) | Turn::Output(..) => {} } Ok(turn) } - fn piece(&mut self, array: ArrayRef) -> Piece { + fn piece(&mut self, split: usize, array: ArrayRef) -> Piece { let start = self.next_piece; self.next_piece += 1; Piece { + split, rows: start..start + 1, array, } } + pub(super) fn complete_child( + &mut self, + id: u64, + arrays: Option>, + ) -> VortexResult<()> { + self.scan.complete_child(id, arrays)?; + if matches!(self.pending, Some(Turn::Waiting)) { + self.pending = None; + } + Ok(()) + } + fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { let location = self.location(segment_id)?; let id = self.next_id; diff --git a/vortex-layout/src/scan/execution/children.rs b/vortex-layout/src/scan/execution/children.rs new file mode 100644 index 00000000000..47255bc85a0 --- /dev/null +++ b/vortex-layout/src/scan/execution/children.rs @@ -0,0 +1,273 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Query-wide admission and execution of large child projections. +//! +//! Permits cover the estimated decoded bytes of whole touched chunks until their task completes. +//! Retained output, encoded buffers and allocator overhead are outside this estimate; this is not +//! a limit on the query's total heap. Cancellation drops child handles and their IO sessions. + +use std::future::poll_fn; +use std::sync::Arc; +use std::sync::OnceLock; +#[cfg(test)] +use std::sync::atomic::AtomicUsize; +#[cfg(test)] +use std::sync::atomic::Ordering; + +use async_lock::Mutex; +use async_lock::Semaphore; +use async_lock::SemaphoreGuardArc; +use futures::FutureExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::future::BoxFuture; +use futures::stream; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_io::request::IoService; +use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::planning::selection::SelectedRows; +use vortex_session::VortexSession; +use vortex_utils::parallelism::get_available_parallelism; + +use super::ScanPlans; +use super::SegmentLocation; +use super::projection::ProjectionPlanner; +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::SelectedPlan; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::ValueCache; + +const BUDGET_BYTES: u64 = 32 * 1024 * 1024; +const PERMIT_BYTES: u64 = 256 * 1024; + +pub(super) mod splits; + +#[cfg(test)] +mod tests; + +#[derive(Debug)] +pub(super) struct ChildWork { + pub(super) split: Split, + /// Estimated bytes of whole touched chunks, including validity. + pub(super) bytes: u64, +} + +struct ByteBudget { + permits: Arc, + acquisition: Mutex<()>, +} + +impl Default for ByteBudget { + fn default() -> Self { + Self { + permits: Arc::new(Semaphore::new((BUDGET_BYTES / PERMIT_BYTES) as usize)), + acquisition: Mutex::new(()), + } + } +} + +impl ByteBudget { + async fn acquire(&self, bytes: u64) -> VortexResult> { + vortex_ensure!( + bytes > 0 && bytes <= BUDGET_BYTES, + "Child work exceeds its admission budget" + ); + // Acquire a whole weight before admitting another claimant, so partial acquisitions + // cannot fill the budget and deadlock one another. Completed tasks release independently. + let _acquisition = self.acquisition.lock().await; + let count = usize::try_from(bytes.div_ceil(PERMIT_BYTES))?; + let mut permits = Vec::with_capacity(count); + for _ in 0..count { + permits.push(self.permits.acquire_arc().await); + } + Ok(permits) + } +} + +pub(in crate::scan) struct Children { + session: VortexSession, + locations: Arc<[SegmentLocation]>, + io: Arc, + workers: usize, + budget: OnceLock, + #[cfg(test)] + admitted: AtomicUsize, +} + +pub(in crate::scan) fn child_executor( + session: VortexSession, + locations: Arc<[SegmentLocation]>, + io: Arc, +) -> Option> { + let workers = get_available_parallelism().unwrap_or(1); + (workers > 1 || std::env::var_os("VORTEX_LIST_PIPELINE").is_some()).then(|| { + Arc::new(Children { + session, + locations, + io, + workers, + budget: OnceLock::new(), + #[cfg(test)] + admitted: Default::default(), + }) + }) +} + +impl Children { + #[cfg(test)] + pub(in crate::scan) fn admitted(&self) -> usize { + self.admitted.load(Ordering::Relaxed) + } + + pub(super) fn can_parallelize(&self) -> bool { + self.workers > 1 + } + + /// Execute a selected graph without holding the fixed-width parallel admission budget. + /// Nested lists may ask their enclosing planner for further selected views. + pub(super) fn execute_view( + self: Arc, + view: SelectedPlan, + values: ValueCache, + ) -> BoxFuture<'static, VortexResult>> { + async move { + let rows = view.splits.first().map_or(0, |split| split.rows.start) + ..view.splits.last().map_or(0, |split| split.rows.end); + let scope = WorkScope { + file_ordinal: 0, + rows, + }; + let plans = ScanPlans { + session: self.session.clone(), + locations: Arc::clone(&self.locations), + projection: view.plan.clone(), + row_offset: 0, + values, + children: Some(Arc::clone(&self)), + }; + let root = ProjectionPlanner::from_view(plans, view, scope.clone()); + let mut run = Run::new(); + run.admit(Box::new(root), scope, self.io.session()); + let mut arrays = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => arrays.push(batch.array), + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + Ok(arrays) + } + .boxed() + } + + pub(super) fn execute( + self: Arc, + plan: PlanRef, + work: Vec, + values: ValueCache, + ) -> BoxFuture<'static, VortexResult>> { + async move { + let workers = self.workers; + stream::iter(work) + .map(|work| { + let children = Arc::clone(&self); + let plan = plan.clone(); + let values = values.clone(); + async move { + let permits = children + .budget + .get_or_init(ByteBudget::default) + .acquire(work.bytes) + .await?; + #[cfg(test)] + children.admitted.fetch_add(1, Ordering::Relaxed); + let handle = children.session.handle(); + handle + .spawn(async move { + let _permits = permits; + let scope = WorkScope { + file_ordinal: 0, + rows: work.split.rows, + }; + let plans = ScanPlans { + session: children.session.clone(), + locations: Arc::clone(&children.locations), + projection: plan, + row_offset: 0, + values, + // Eligible child plans contain only fixed-width segment leaves. + // They cannot recursively wait for this same byte budget. + children: None, + }; + let root = ProjectionPlanner::new( + plans, + SelectedRows { + scope: scope.clone(), + mask: work.split.mask, + }, + ); + let mut run = Run::new(); + run.admit(Box::new(root), scope, children.io.session()); + let mut result = Vec::new(); + let mut exec = children.session.create_execution_ctx(); + loop { + match run.advance()? { + Progress::Batch(batch) => { + materialize(batch.array, &mut exec, &mut result)?; + } + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => { + poll_fn(|cx| run.poll_completion(cx)).await? + } + } + } + Ok::<_, vortex_error::VortexError>(result) + }) + .await + } + }) + .buffered(workers) + .try_collect::>>() + .await + .map(|arrays| arrays.into_iter().flatten().collect()) + } + .boxed() + } +} + +/// Decode on the child task, preserving chunk boundaries so the parent need only concatenate +/// once. Canonicalizing the whole job would copy its chunks into an intermediate buffer first. +fn materialize( + array: ArrayRef, + exec: &mut ExecutionCtx, + output: &mut Vec, +) -> VortexResult<()> { + if let Some(chunks) = array.as_opt::() { + for chunk in chunks.iter_chunks() { + materialize(chunk.clone(), exec, output)?; + } + return Ok(()); + } + output.push(match array.dtype() { + DType::Primitive(..) => array.execute::(exec)?.into_array(), + DType::Bool(..) => array.execute::(exec)?.into_array(), + _ => vortex_bail!("Child projection is not fixed-width"), + }); + Ok(()) +} diff --git a/vortex-layout/src/scan/execution/children/splits.rs b/vortex-layout/src/scan/execution/children/splits.rs new file mode 100644 index 00000000000..89bf01eff97 --- /dev/null +++ b/vortex-layout/src/scan/execution/children/splits.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Split selected child projections at stored segment boundaries before admitting them. + +use std::ops::Range; + +use rustc_hash::FxHashSet; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; + +use super::ChildWork; +use crate::plan::physical::Concat; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::Zoned; +use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::overlapping; +use crate::segments::SegmentId; + +const MAX_ELEMENT_BYTES: u64 = 8 * 1024 * 1024; +const CHILD_TARGET_BYTES: u64 = 256 * 1024; + +#[cfg(test)] +mod tests; + +/// Divide substantial fixed-width selections at stored segment boundaries. Each segment is +/// owned by at most one task, so independent execution cannot duplicate its IO and decoding. +pub(in crate::scan::execution) fn parallel_splits( + plan: &PlanRef, + splits: &[Split], +) -> VortexResult>> { + let width = match plan.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(None), + }; + let selected = splits.iter().fold(0_u64, |total, split| { + total.saturating_add((split.mask.true_count() as u64).saturating_mul(width)) + }); + if selected < CHILD_TARGET_BYTES * 2 { + return Ok(None); + } + let mut seen = FxHashSet::default(); + let mut work = Vec::new(); + for split in splits { + let mut leaves = Vec::new(); + if !selected_leaves(plan, split, 0, &mut seen, &mut leaves)? { + return Ok(None); + } + let mut first = 0; + while first < leaves.len() { + let start = leaves[first].0.start; + let mut end = leaves[first].0.end; + let mut bytes = leaves[first].1; + first += 1; + while bytes < CHILD_TARGET_BYTES + && first < leaves.len() + && bytes.saturating_add(leaves[first].1) <= MAX_ELEMENT_BYTES + { + end = leaves[first].0.end; + bytes += leaves[first].1; + first += 1; + } + work.push(ChildWork { + split: Split { + rows: start..end, + mask: split.mask.slice( + usize::try_from(start - split.rows.start)? + ..usize::try_from(end - split.rows.start)?, + ), + }, + bytes, + }); + } + } + Ok((work.len() > 1).then_some(work)) +} + +fn selected_leaves( + plan: &PlanRef, + split: &Split, + offset: u64, + seen: &mut FxHashSet, + leaves: &mut Vec<(Range, u64)>, +) -> VortexResult { + if split.mask.all_false() { + return Ok(true); + } + if let Some(filter) = plan.as_opt::() { + return selected_leaves(&filter.child_plan()?, split, offset, seen, leaves); + } + if let Some(zoned) = plan.as_opt::() { + return match zoned.data_plan()? { + Some(data) => selected_leaves(&data, split, offset, seen, leaves), + None => Ok(false), + }; + } + if let Some(stats) = plan.as_opt::() { + return selected_leaves(&stats.data_plan()?, split, offset, seen, leaves); + } + if let Some(concat) = plan.as_opt::() { + for (index, start, rows) in overlapping(concat, &split.rows) { + let mask = split.mask.slice( + usize::try_from(rows.start - split.rows.start)? + ..usize::try_from(rows.end - split.rows.start)?, + ); + // Do not lower an untouched lazy child, including unknown layouts in holes. + if mask.all_false() { + continue; + } + let child = Split { + rows: rows.start - start..rows.end - start, + mask, + }; + if !selected_leaves( + &concat.child_required(index)?, + &child, + offset + start, + seen, + leaves, + )? { + return Ok(false); + } + } + return Ok(true); + } + let Some(segment) = plan.as_opt::() else { + return Ok(false); + }; + let width = match plan.dtype() { + DType::Primitive(ptype, _) => ptype.byte_width() as u64, + DType::Bool(_) => 1, + _ => return Ok(false), + }; + let validity = if plan.dtype().is_nullable() { + plan.row_count().div_ceil(8) + } else { + 0 + }; + let bytes = plan + .row_count() + .saturating_mul(width) + .saturating_add(validity); + if bytes == 0 || bytes > MAX_ELEMENT_BYTES || !seen.insert(segment.segment_id()) { + return Ok(false); + } + leaves.push((offset + split.rows.start..offset + split.rows.end, bytes)); + Ok(true) +} diff --git a/vortex-layout/src/scan/execution/children/splits/tests.rs b/vortex-layout/src/scan/execution/children/splits/tests.rs new file mode 100644 index 00000000000..d7ba1c8f7e8 --- /dev/null +++ b/vortex-layout/src/scan/execution/children/splits/tests.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::dtype::PType; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::SegmentScanPlan; + +fn chunks(count: u32, rows: u64, dtype: DType, duplicate: bool) -> VortexResult { + ConcatPlan::try_new( + dtype.clone(), + (0..count) + .map(|id| { + SegmentScanPlan::new( + dtype.clone(), + rows, + SegmentId::from(if duplicate { 0 } else { id }), + ReadContext::new([]), + None, + ) + .into_plan() + }) + .collect(), + ) + .map(|plan| plan.into_plan()) +} + +#[rstest] +fn contiguous_list_row_splits_at_stored_boundaries( + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let rows = 8192; + let plan = chunks( + 32, + rows, + DType::Primitive(PType::I64, nullable.into()), + false, + )?; + let jobs = parallel_splits(&plan, &[Split::all(100..rows * 32 - 100)])? + .expect("large row is parallel"); + assert_eq!(jobs.len(), 8); + assert_eq!(jobs[0].split.rows, 100..rows * 4); + assert_eq!(jobs[7].split.rows, rows * 28..rows * 32 - 100); + assert!(jobs.iter().all(|job| job.split.mask.all_true())); + assert_eq!( + jobs.iter().map(|job| job.bytes).sum::(), + rows * 32 * 8 + if nullable { rows * 4 } else { 0 } + ); + Ok(()) +} + +#[test] +fn masks_preserve_holes_and_do_not_charge_unselected_segments() -> VortexResult<()> { + let rows = 32768; + let plan = chunks(5, rows, DType::Primitive(PType::I64, false.into()), false)?; + let mask = Mask::from_slices( + usize::try_from(rows * 5)?, + vec![(0, 32768), (65536, 98304), (131072, 163840)], + ); + let jobs = parallel_splits( + &plan, + &[Split { + rows: 0..rows * 5, + mask, + }], + )? + .expect("three selected chunks"); + assert_eq!( + jobs.iter() + .map(|job| job.split.rows.clone()) + .collect::>(), + vec![0..rows, rows * 2..rows * 3, rows * 4..rows * 5] + ); + assert_eq!(jobs.iter().map(|job| job.bytes).sum::(), rows * 3 * 8); + Ok(()) +} + +#[test] +fn partial_masks_survive_grouping_and_boundary_clipping() -> VortexResult<()> { + let plan = chunks(32, 8192, DType::Primitive(PType::I64, false.into()), false)?; + let mask = Mask::from_iter((0..200_000).map(|i| i % 7 != 0)); + let split = Split { + rows: 999..200_999, + mask: mask.clone(), + }; + let jobs = parallel_splits(&plan, &[split])?.expect("large selected payload"); + let actual: Vec<_> = jobs + .iter() + .flat_map(|job| { + job.split + .mask + .iter() + .enumerate() + .filter_map(|(i, selected)| selected.then_some(job.split.rows.start + i as u64)) + }) + .collect(); + let expected: Vec<_> = mask + .iter() + .enumerate() + .filter_map(|(i, selected)| selected.then_some(999 + i as u64)) + .collect(); + assert_eq!(actual, expected); + Ok(()) +} + +#[test] +fn repeated_segments_and_repeated_ranges_stay_serial() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, false.into()); + let repeated = chunks(2, 65536, dtype.clone(), true)?; + assert!(parallel_splits(&repeated, &[Split::all(0..131072)])?.is_none()); + let distinct = chunks(2, 65536, dtype, false)?; + assert!( + parallel_splits( + &distinct, + &[Split::all(0..32000), Split::all(32001..131072)] + )? + .is_none() + ); + Ok(()) +} + +#[rstest] +#[case::single_segment(1, 131072, DType::Primitive(PType::I64, false.into()))] +#[case::oversized_segment(2, MAX_ELEMENT_BYTES / 8 + 1, DType::Primitive(PType::I64, false.into()))] +#[case::tiny_selection(8, 1024, DType::Primitive(PType::I64, false.into()))] +#[case::variable_width(8, 131072, DType::Utf8(false.into()))] +fn unsuitable_child_work_stays_serial( + #[case] count: u32, + #[case] rows: u64, + #[case] dtype: DType, +) -> VortexResult<()> { + let plan = chunks(count, rows, dtype, false)?; + assert!(parallel_splits(&plan, &[Split::all(0..rows * u64::from(count))])?.is_none()); + Ok(()) +} diff --git a/vortex-layout/src/scan/execution/children/tests.rs b/vortex-layout/src/scan/execution/children/tests.rs new file mode 100644 index 00000000000..c6c166730a6 --- /dev/null +++ b/vortex-layout/src/scan/execution/children/tests.rs @@ -0,0 +1,413 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::VecDeque; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; +use std::time::Duration; +use std::time::Instant; + +use futures::executor::block_on; +use futures::task::AtomicWaker; +use parking_lot::Mutex as SyncMutex; +use vortex_array::ArrayContext; +use vortex_array::arrays::ChunkedArray; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::PType; +use vortex_array::serde::SerializeOptions; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +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::IoSource; +use vortex_io::request::IoTarget; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::pipeline::Split; +use crate::segments::SegmentId; +use crate::test::new_session; + +#[test] +fn byte_budget_prevents_over_admission_and_cancellation_releases_partial_acquisitions() +-> VortexResult<()> { + let budget = ByteBudget::default(); + let held = block_on(budget.acquire(BUDGET_BYTES * 3 / 4))?; + let mut waiting = Box::pin(budget.acquire(BUDGET_BYTES / 2)); + let mut cx = Context::from_waker(Waker::noop()); + assert!(waiting.poll_unpin(&mut cx).is_pending()); + // The waiting acquisition has taken the remaining quarter. Cancelling it must return + // those permits and the acquisition lock, even while the first job is still running. + drop(waiting); + let mut quarter = Box::pin(budget.acquire(BUDGET_BYTES / 4)); + let Poll::Ready(permits) = quarter.poll_unpin(&mut cx) else { + panic!("cancelled admission leaked permits") + }; + let quarter = permits?; + let mut extra = Box::pin(budget.acquire(1)); + assert!(extra.poll_unpin(&mut cx).is_pending()); + drop(quarter); + let Poll::Ready(permits) = extra.poll_unpin(&mut cx) else { + panic!("released permits did not wake admission") + }; + drop(permits?); + drop(held); + let mut whole = Box::pin(budget.acquire(BUDGET_BYTES)); + assert!(matches!(whole.poll_unpin(&mut cx), Poll::Ready(Ok(_)))); + Ok(()) +} + +#[test] +fn admission_rejects_zero_and_oversized_work() { + let budget = ByteBudget::default(); + assert!(block_on(budget.acquire(0)).is_err()); + assert!(block_on(budget.acquire(BUDGET_BYTES + 1)).is_err()); +} + +#[derive(Default)] +struct Gates { + sources: SyncMutex>>, + bytes: Vec, +} + +impl IoService for Gates { + fn session(&self) -> Arc { + let gate = Arc::new(Gate { + bytes: self.bytes.clone(), + ready: AtomicBool::new(false), + fail: AtomicBool::new(false), + requests: SyncMutex::new(VecDeque::new()), + wake: AtomicWaker::new(), + cleared: AtomicUsize::new(0), + }); + self.sources.lock().push(Arc::clone(&gate)); + gate + } +} + +struct Gate { + bytes: Vec, + ready: AtomicBool, + fail: AtomicBool, + requests: SyncMutex>, + wake: AtomicWaker, + cleared: AtomicUsize, +} + +impl Gate { + fn open(&self) { + self.ready.store(true, Ordering::Release); + self.wake.wake(); + } +} + +impl IoSource for Gate { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + for request in batch + .into_iter() + .filter(|request| request.intent == IoIntent::Fetch) + { + let IoTarget::Range { offset, .. } = request.target else { + return Err(vortex_err!("unexpected size request")); + }; + self.requests.lock().push_back(Completion { + owner, + request: request.request, + result: Ok(IoResult::Bytes( + self.bytes[usize::try_from(offset)?].clone(), + )), + }); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + if self.fail.load(Ordering::Acquire) { + return Err(vortex_err!("injected child IO failure")); + } + Ok(if self.ready.load(Ordering::Acquire) { + self.requests.lock().pop_front() + } else { + None + }) + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + self.wake.register(cx.waker()); + match self.poll() { + Ok(Some(completion)) => Poll::Ready(Ok(completion)), + Ok(None) => Poll::Pending, + Err(error) => Poll::Ready(Err(error)), + } + } + + fn wait(&self) -> VortexResult { + Err(vortex_err!("child executor must not block on IO")) + } + fn release(&self, owner: IoOwnerId) { + self.requests + .lock() + .retain(|request| request.owner != owner); + } + fn clear(&self) { + self.requests.lock().clear(); + self.cleared.fetch_add(1, Ordering::Relaxed); + } +} + +type Fixture = (Arc, Arc, PlanRef, Vec); + +#[test] +fn materialization_preserves_chunk_buffers_instead_of_concatenating_twice() -> VortexResult<()> { + let session = new_session(); + let chunks = [ + PrimitiveArray::from_option_iter([Some(1_i64), None]), + PrimitiveArray::from_option_iter([None, Some(4_i64)]), + ]; + let array = ChunkedArray::try_new( + chunks.iter().cloned().map(IntoArray::into_array), + chunks[0].dtype().clone(), + )? + .into_array(); + let mut output = Vec::new(); + materialize(array, &mut session.create_execution_ctx(), &mut output)?; + assert_eq!(output.len(), chunks.len()); + for (actual, expected) in output.into_iter().zip(chunks) { + let actual = actual.execute::(&mut session.create_execution_ctx())?; + assert_eq!( + actual.as_slice::().as_ptr(), + expected.as_slice::().as_ptr() + ); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} + +fn fixture() -> VortexResult { + let session = new_session().with_tokio(); + let mut service = Gates::default(); + let mut plans = Vec::new(); + for id in 0_u32..3 { + let array = Buffer::from_iter(u64::from(id) * 8..u64::from(id + 1) * 8).into_array(); + let ctx = ArrayContext::empty(); + let buffers = array.serialize( + &ctx, + &session, + &SerializeOptions { + offset: 0, + include_padding: true, + }, + )?; + let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64)); + for buffer in buffers { + bytes.extend_from_slice(buffer.as_ref()); + } + service.bytes.push(BufferHandle::new_host(bytes.freeze())); + plans.push( + SegmentScanPlan::new( + array.dtype().clone(), + 8, + SegmentId::from(id), + ReadContext::new(ctx.to_ids()), + None, + ) + .into_plan(), + ); + } + let service = Arc::new(service); + let children = Arc::new(Children { + session, + io: Arc::::clone(&service), + workers: 4, + budget: OnceLock::new(), + admitted: Default::default(), + locations: (0..3) + .map(|offset| SegmentLocation { + offset, + length: 0, + alignment: Alignment::new(64), + }) + .collect(), + }); + let plan = ConcatPlan::try_new(DType::Primitive(PType::U64, false.into()), plans)?.into_plan(); + // Inflate the estimate to exercise admission without a 48 MiB test fixture. + let work = (0..3) + .map(|id| ChildWork { + split: Split::all(id * 8..(id + 1) * 8), + bytes: BUDGET_BYTES / 2, + }) + .collect(); + Ok((children, service, plan, work)) +} + +async fn until(mut condition: impl FnMut() -> bool) { + let start = Instant::now(); + while !condition() { + assert!( + start.elapsed() < Duration::from_secs(5), + "child execution did not make progress" + ); + tokio::task::yield_now().await; + } +} + +#[tokio::test] +async fn independent_children_respect_the_shared_budget_and_return_in_order() -> VortexResult<()> { + let (children, service, plan, work) = fixture()?; + let task = tokio::spawn(Arc::clone(&children).execute(plan, work, ValueCache::default())); + until(|| service.sources.lock().len() == 2).await; + let first = Arc::clone(&service.sources.lock()[0]); + let second = Arc::clone(&service.sources.lock()[1]); + assert!(!task.is_finished()); + assert_eq!(first.requests.lock().len(), 1); + assert_eq!(second.requests.lock().len(), 1); + // The later job releases capacity while the first job is still held. The third job + // must start without losing output order or exceeding the two-job byte budget. + second.open(); + until(|| service.sources.lock().len() == 3).await; + assert!(!first.ready.load(Ordering::Relaxed)); + service.sources.lock()[2].open(); + first.open(); + let arrays = task.await.expect("child task panicked")?; + assert_eq!(arrays.len(), 3); + for (index, array) in arrays.into_iter().enumerate() { + let start = u64::try_from(index)? * 8; + assert_arrays_eq!( + array, + Buffer::from_iter(start..start + 8).into_array(), + &mut children.session.create_execution_ctx() + ); + } + assert!( + service + .sources + .lock() + .iter() + .all(|source| source.cleared.load(Ordering::Relaxed) == 1) + ); + Ok(()) +} + +#[tokio::test] +async fn selected_view_preserves_split_order_and_skips_unselected_segments() -> VortexResult<()> { + let (children, service, plan, _) = fixture()?; + let view = SelectedPlan::new(plan, vec![Split::all(1..4), Split::all(17..20)], 2); + let task = tokio::spawn(Arc::clone(&children).execute_view(view, ValueCache::default())); + until(|| { + service + .sources + .lock() + .first() + .is_some_and(|source| source.requests.lock().len() == 2) + }) + .await; + let source = Arc::clone(&service.sources.lock()[0]); + source.requests.lock().make_contiguous().reverse(); + source.open(); + let arrays = task.await.expect("selected view task panicked")?; + let actual = + ChunkedArray::try_new(arrays, DType::Primitive(PType::U64, false.into()))?.into_array(); + let expected = Buffer::from_iter((1_u64..4).chain(17..20)).into_array(); + assert_arrays_eq!(actual, expected, &mut children.session.create_execution_ctx()); + assert_eq!(children.admitted(), 0); + assert!(children.budget.get().is_none()); + assert_eq!(service.sources.lock().len(), 1); + assert_eq!(source.cleared.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[rstest::rstest] +#[tokio::test] +async fn cancelled_or_failed_selected_view_releases_its_io( + #[values(false, true)] failure: bool, +) -> VortexResult<()> { + let (children, service, plan, _) = fixture()?; + let view = SelectedPlan::new(plan, vec![Split::all(1..4)], 1); + let task = tokio::spawn(Arc::clone(&children).execute_view(view, ValueCache::default())); + until(|| !service.sources.lock().is_empty()).await; + let source = Arc::clone(&service.sources.lock()[0]); + if failure { + source.fail.store(true, Ordering::Release); + source.open(); + let error = task + .await + .expect("selected view task panicked") + .expect_err("injected IO failure must propagate"); + assert!(error.to_string().contains("injected child IO failure")); + } else { + task.abort(); + assert!(task.await.expect_err("task must be cancelled").is_cancelled()); + } + until(|| source.cleared.load(Ordering::Relaxed) == 1).await; + assert!(source.requests.lock().is_empty()); + assert!(children.budget.get().is_none()); + Ok(()) +} + +#[rstest::rstest] +#[tokio::test] +async fn cancelled_or_failed_parent_releases_child_io_and_budget( + #[values(false, true)] failure: bool, +) -> VortexResult<()> { + let (children, service, plan, work) = fixture()?; + let task = tokio::spawn(Arc::clone(&children).execute(plan, work, ValueCache::default())); + until(|| service.sources.lock().len() == 2).await; + if failure { + let first = Arc::clone(&service.sources.lock()[0]); + first.fail.store(true, Ordering::Release); + first.open(); + assert!( + task.await + .expect("task panicked") + .unwrap_err() + .to_string() + .contains("injected child IO failure") + ); + } else { + task.abort(); + assert!( + task.await + .expect_err("task should be cancelled") + .is_cancelled() + ); + } + until(|| { + service + .sources + .lock() + .iter() + .all(|source| source.cleared.load(Ordering::Relaxed) == 1) + }) + .await; + assert!( + service + .sources + .lock() + .iter() + .all(|source| source.requests.lock().is_empty()) + ); + let mut all = Box::pin( + children + .budget + .get() + .expect("budget initialized") + .acquire(BUDGET_BYTES), + ); + assert!(matches!( + all.poll_unpin(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok(_)) + )); + Ok(()) +} diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 9d8b8897ef3..3f919d3a5b5 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,23 +3,29 @@ //! Scan execution over physical plans through the planner I/O and compute protocol. //! -//! A `ProjectionPlanner` emits one `ProjectionMorsel` for each fixed-size work chunk. The morsel -//! runs the physical projection plan over the chunk's selected rows through -//! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published -//! through a `ScanAdaptor`, so the planning driver can advance other work while reads wait. +//! A `ProjectionPlanner` drives the physical projection over selected rows and schedules +//! independently admitted child projections. Every segment read is published through a +//! `ScanAdaptor`, so the driver advances other work while reads wait. Only the assembled result +//! becomes a morsel, which publishes the selected rows in order. mod adaptor; +mod children; mod projection; use std::sync::Arc; +pub(super) use children::child_executor; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; -use vortex_mask::Mask; +use vortex_scan::planning::next::next_fn; use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; +use vortex_scan::planning::selection::SelectedRows; +use vortex_scan::planning::selection::SelectionPlanner; +use vortex_scan::selection::Selection; use vortex_session::VortexSession; +use self::children::Children; use self::projection::ProjectionPlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; @@ -60,17 +66,18 @@ pub(super) struct ScanPlans { /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans /// always read the same complete row domain, independently of a split's selection. pub(super) values: ValueCache, + pub(super) children: Option>, } -/// The rows of a split selected for projection. -struct SelectedRows { - /// The split the rows belong to. +/// Selects the rows of a fixed work chunk before handing them to its projection planner. +pub(super) fn plan_projection( + plans: ScanPlans, scope: WorkScope, - /// 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 })) + selection: Arc, +) -> Box { + Box::new(SelectionPlanner::new( + scope, + selection, + next_fn(move |selected| Ok(ProjectionPlanner::new(plans.clone(), selected))), + )) } diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index b6180725c27..f07054487b5 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,12 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + +use futures::FutureExt; +use futures::channel::oneshot; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_error::VortexExpect; 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; @@ -15,98 +21,82 @@ 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 crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::SelectedPlan; use crate::plan::physical::pipeline::Split; use crate::scan::execution::ScanPlans; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; - -/// Hands one fixed-size work chunk to a projection morsel. -pub(super) struct ProjectionPlanner { - work: Option<(ScanPlans, SelectedRows)>, -} - -impl ProjectionPlanner { - /// Creates a planner for `selected`. - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { - Self { - work: Some((plans, selected)), - } - } -} - -impl IoConsumer for ProjectionPlanner { - fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} -} - -impl Planner for ProjectionPlanner { - fn state(&self) -> State { - if self.work.is_some() { - State::NeedsCompute - } else { - State::Done - } - } - - fn compute(&mut self) -> VortexResult { - let Some((plans, selected)) = self.work.take() else { - return Ok(PlannerOutput::Done); - }; - if selected.mask.all_false() { - return Ok(PlannerOutput::Done); - } - Ok(PlannerOutput::Morsel( - selected.scope.clone(), - Box::new(ProjectionMorsel::new(plans, selected)), - )) - } -} +use crate::scan::execution::children::Children; +use crate::scan::execution::children::splits::parallel_splits; /// 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. -struct ProjectionMorsel { +pub(super) struct ProjectionPlanner { plans: Option, dtype: DType, - selected: SelectedRows, + scope: WorkScope, + view: Option, adaptor: Option, pieces: Vec, done: bool, + children: Option>, + pending: Option<(u64, oneshot::Receiver>)>, } -impl ProjectionMorsel { - /// Creates a morsel projecting `selected`. - fn new(plans: ScanPlans, selected: SelectedRows) -> Self { +impl ProjectionPlanner { + /// Creates a planner projecting `selected`. + pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + let view = SelectedPlan::new( + plans.projection.clone(), + vec![Split { + rows: selected.scope.rows.clone(), + mask: selected.mask, + }], + 1, + ); + Self::from_view(plans, view, selected.scope) + } + + pub(super) fn from_view(plans: ScanPlans, view: SelectedPlan, scope: WorkScope) -> Self { Self { dtype: plans.projection.dtype().clone(), + children: plans.children.clone(), + pending: None, plans: Some(plans), - selected, + scope, + view: Some(view), adaptor: None, pieces: Vec::new(), done: false, } } - fn finish(&mut self) -> VortexResult { + fn finish(&mut self) -> VortexResult { self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); + self.pieces + .sort_by_key(|piece| (piece.split, 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), + 0 => return Ok(PlannerOutput::Done), 1 => arrays.remove(0), _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), }; - Ok(MorselOutput::Batch(array)) + Ok(PlannerOutput::Morsel( + self.scope.clone(), + Box::new(ProjectedBatch(Some(array))), + )) } } -impl IoConsumer for ProjectionMorsel { +impl IoConsumer for ProjectionPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { if let Some(adaptor) = self.adaptor.as_mut() { adaptor.set_io_result(request, result); @@ -114,10 +104,11 @@ impl IoConsumer for ProjectionMorsel { } } -impl Morsel for ProjectionMorsel { +impl Planner for ProjectionPlanner { fn state(&self) -> State { match &self.adaptor { _ if self.done => State::Done, + _ if self.pending.is_some() => State::NeedsCompute, None => State::NeedsCompute, Some(adaptor) => match adaptor.state() { State::Done => State::NeedsCompute, @@ -126,38 +117,101 @@ impl Morsel for ProjectionMorsel { } } - fn compute(&mut self) -> VortexResult { + fn compute(&mut self) -> VortexResult { if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); + vortex_bail!("ProjectionPlanner: compute called after Done"); + } + if let Some((id, mut result)) = self.pending.take() { + let arrays = result + .try_recv() + .map_err(|_| vortex_err!("Child projection was cancelled"))? + .ok_or_else(|| vortex_err!("Child projection resumed before completion"))?; + self.adaptor + .as_mut() + .ok_or_else(|| vortex_err!("Child completed before projection started"))? + .complete_child(id, Some(arrays))?; } let Some(adaptor) = self.adaptor.as_mut() else { let plans = self .plans .take() .vortex_expect("an uncompiled projection still has its plans"); - let scan = Scan::try_new( - plans.session, - plans.projection, - vec![Split { - rows: self.selected.scope.rows.clone(), - mask: self.selected.mask.clone(), - }], - )? - .with_row_offset(plans.row_offset) - .with_value_cache(plans.values); + let view = self + .view + .take() + .vortex_expect("an uncompiled projection has a view"); + let scan = view + .into_scan(plans.session, plans.values)? + .with_row_offset(plans.row_offset) + .with_child_planning(self.children.is_some()); self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); - return Ok(MorselOutput::Continue); + return Ok(PlannerOutput::Continue); }; if adaptor.state() == State::Done { return self.finish(); } Ok(match adaptor.compute()? { - AdaptorStep::Yield => MorselOutput::Continue, - AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Child(child) => { + let children = Arc::clone( + self.children + .as_ref() + .ok_or_else(|| vortex_err!("Child projection has no scheduler"))?, + ); + let work = if children.can_parallelize() { + parallel_splits(&child.view.plan, &child.view.splits)? + } else { + None + }; + if work.is_none() && !child.planner_owned { + adaptor.complete_child(child.id, None)?; + return Ok(PlannerOutput::Continue); + } + let (send, recv) = oneshot::channel(); + self.pending = Some((child.id, recv)); + PlannerOutput::Await( + async move { + let arrays = match work { + Some(work) => { + children.execute(child.view.plan, work, child.values).await? + } + None => children.execute_view(child.view, child.values).await?, + }; + send.send(arrays) + .map_err(|_| vortex_err!("Projection planner was cancelled")) + } + .boxed(), + ) + } AdaptorStep::Piece(piece) => { self.pieces.push(piece); - MorselOutput::Continue + PlannerOutput::Continue } }) } } + +/// Planning resolves the projection; this morsel only publishes the assembled batch. +struct ProjectedBatch(Option); + +impl IoConsumer for ProjectedBatch { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Morsel for ProjectedBatch { + fn state(&self) -> State { + if self.0.is_some() { + State::NeedsCompute + } else { + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + Ok(match self.0.take() { + Some(array) => MorselOutput::Batch(array), + None => MorselOutput::Done, + }) + } +} diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index ce565846d7e..2c4ee9c121f 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -34,25 +34,55 @@ 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()) +/// File segments are normally ordered, so range lookup can reuse the location table directly. +/// Custom segment sources can supply arbitrary order and retain indexed lookup. +pub(super) enum SegmentRanges { + Sorted(Arc<[SegmentLocation]>), + Indexed(HashMap<(u64, usize), SegmentId>), +} + +pub(super) fn segment_ranges(locations: Arc<[SegmentLocation]>) -> Arc { + if locations.is_sorted_by_key(|location| (location.offset, location.length)) { + return Arc::new(SegmentRanges::Sorted(locations)); + } + Arc::new(SegmentRanges::Indexed( + 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) + }) + .collect(), + )) +} + +impl SegmentRanges { + fn get(&self, offset: u64, len: usize) -> Option { + match self { + Self::Sorted(locations) => { + let end = locations.partition_point(|location| { + (location.offset, location.length as usize) <= (offset, len) + }); + let index = end.checked_sub(1)?; + let location = locations[index]; + // Match the indexed path's last-entry rule for duplicate ranges. + (location.offset == offset && location.length as usize == len) + .then(|| SegmentId::from(u32::try_from(index).unwrap_or(u32::MAX))) + } + Self::Indexed(ranges) => ranges.get(&(offset, len)).copied(), + } + } } /// Serves splits from any [`SegmentSource`], by segment id. pub(super) struct SegmentScanIo { segments: Arc, - ranges: SegmentRanges, + ranges: Arc, } impl SegmentScanIo { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + pub(super) fn new(segments: Arc, ranges: Arc) -> Self { Self { segments, ranges } } } @@ -77,7 +107,7 @@ impl IoService for SegmentScanIo { /// fetch consumes it, so a source that coalesces requests can fold it into a nearby read. pub(super) struct SegmentIoSource { segments: Arc, - ranges: SegmentRanges, + ranges: Arc, reads: Mutex>>, ready: Mutex>, /// Reads started by prefetches, consumed by the next fetch of their segment. @@ -90,7 +120,7 @@ pub(super) struct SegmentIoSource { type Prefetched = Shared>>>; impl SegmentIoSource { - pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + pub(super) fn new(segments: Arc, ranges: Arc) -> Self { Self { segments, ranges, @@ -114,9 +144,9 @@ impl IoSource for SegmentIoSource { request.target ); }; - let segment = *self + let segment = self .ranges - .get(&(offset, len)) + .get(offset, len) .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; if request.intent == IoIntent::Announce { self.announced @@ -234,6 +264,7 @@ mod tests { use futures::channel::oneshot; use futures::future; use parking_lot::Mutex; + use rstest::rstest; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -247,12 +278,47 @@ mod tests { use vortex_io::request::IoTarget; use super::SegmentIoSource; + use super::SegmentRanges; use super::segment_ranges; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; + #[rstest] + #[case::sorted(vec![(0, 0), (0, 4), (8, 4), (u64::MAX, 0)], true, vec![Some(0), Some(1), Some(2), Some(3)])] + #[case::unordered(vec![(8, 4), (0, 4), (0, 0), (u64::MAX, 0)], false, vec![Some(2), Some(1), Some(0), Some(3)])] + #[case::duplicates(vec![(0, 0), (0, 0), (0, 4), (8, 4), (8, 4)], true, vec![Some(1), Some(2), Some(4), None])] + #[case::empty(vec![], true, vec![None, None, None, None])] + fn range_lookup_preserves_segment_ids( + #[case] entries: Vec<(u64, u32)>, + #[case] sorted: bool, + #[case] expected: Vec>, + ) { + let locations: Arc<[_]> = entries + .into_iter() + .map(|(offset, length)| SegmentLocation { + offset, + length, + alignment: Alignment::none(), + }) + .collect(); + let ranges = segment_ranges(Arc::clone(&locations)); + assert_eq!(matches!(ranges.as_ref(), SegmentRanges::Sorted(_)), sorted); + if let SegmentRanges::Sorted(reused) = ranges.as_ref() { + assert!(Arc::ptr_eq(reused, &locations)); + } + for ((offset, len), expected) in [(0, 0), (0, 4), (8, 4), (u64::MAX, 0)] + .into_iter() + .zip(expected) + { + assert_eq!(ranges.get(offset, len), expected.map(SegmentId::from)); + } + for (offset, len) in [(0, 1), (1, 4), (8, 0), (8, 5)] { + assert_eq!(ranges.get(offset, len), None); + } + } + /// Serves four bytes per segment and counts the reads it is asked for. #[derive(Default)] struct CountingSegments(AtomicUsize); @@ -301,7 +367,7 @@ mod tests { length: 4, alignment: Alignment::none(), }); - let io = SegmentIoSource::new(segments, segment_ranges(&locations)); + let io = SegmentIoSource::new(segments, segment_ranges(Arc::from(locations))); let mut pending = request(IoIntent::Fetch, 1); pending.target = IoTarget::range(4, 4); io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; @@ -334,7 +400,10 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let io = SegmentIoSource::new( + Arc::clone(&segments) as _, + segment_ranges(Arc::from([location])), + ); io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; assert_eq!(segments.0.load(Ordering::Relaxed), 1); @@ -362,7 +431,10 @@ mod tests { length: 4, alignment: Alignment::none(), }; - let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + let io = SegmentIoSource::new( + Arc::clone(&segments) as _, + segment_ranges(Arc::from([location])), + ); io.submit( IoOwnerId(0), vec![ diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 46cbde987fc..b488b3bca74 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -11,8 +11,10 @@ //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] //! directly and supports projection, row ranges, and selections. Unsupported layouts return an //! 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. +//! chunks independently of the stored layout's chunk boundaries. Dense scans use up to 65,536 +//! rows per chunk, in 8,192-row multiples, targeting 8 MiB of serialized data. Sparse identity +//! projections can use larger aligned chunks within that estimate, targeting four tasks per +//! worker to reduce repeated segment reads. mod io; mod split; @@ -35,8 +37,8 @@ use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; +use vortex_error::VortexExpect; 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,18 +48,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::child_executor; use crate::scan::execution::plan_projection; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; @@ -69,7 +66,11 @@ const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; /// Limit estimated bytes per task without changing the writer-aligned row granularity. -fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { +fn rows_per_chunk( + row_count: u64, + locations: &[SegmentLocation], + sparse_target: Option, +) -> usize { if row_count <= MAX_ROWS_PER_CHUNK as u64 { return MAX_ROWS_PER_CHUNK; } @@ -79,7 +80,15 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { .sum(); let bytes_per_row = bytes.div_ceil(row_count).max(1); let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE + let max_blocks = sparse_target.map_or(8, |rows| { + // Grow sparse tasks only within the byte target, keeping power-of-two boundaries + // aligned with the original work chunks and enough tasks to occupy the workers. + let blocks = (rows / ROW_BLOCK_SIZE as u64).min(blocks).max(8); + 1_u64 << blocks.ilog2() + }); + usize::try_from(blocks.clamp(1, max_blocks)) + .vortex_expect("the byte target bounds the block count to at most 1024") + * ROW_BLOCK_SIZE } /// Fixed boundaries in the file's row domain, clipped to the requested range. @@ -109,8 +118,10 @@ pub struct ScanFile { /// Builds a projection scan on the planner-driven V2 executor. /// -/// This builder reads the file's layout and segments directly. Unsupported layouts return an -/// error when the stream is polled. Filters and row limits are not supported. +/// This builder reads the file's layout and segments directly. Layouts are checked as their +/// plans are lowered, so unsupported layouts outside the requested range can remain unread. +/// Unsupported plans return an error when the stream is polled. Filters and row limits are +/// not supported. pub struct ScanBuilder { session: VortexSession, file: ScanFile, @@ -237,34 +248,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 +267,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 @@ -299,15 +278,29 @@ impl PreparedScan { optimize(plan_row_idx_expression(builder.projection, root)?)? }; let chunk_rows = if identity { - rows_per_chunk(row_count, &file.locations) + let rows = builder.row_range.clone().unwrap_or(0..row_count); + let selected = builder.selection.row_count_in_range(&rows); + let span = rows.end.saturating_sub(rows.start); + let sparse_target = (selected <= span / 16).then(|| { + // Tiny selections cannot occupy the full worker budget. Larger aligned tasks + // let nearby selected rows share their stored segments. + let tasks = (4 * get_available_parallelism().unwrap_or(1)) as u64; + span / tasks.min(selected).max(1) + }); + rows_per_chunk(row_count, &file.locations, sparse_target) } 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( + let io: Arc = Arc::new(SegmentScanIo::new( file.segments, - segment_ranges(&file.locations), + segment_ranges(Arc::clone(&file.locations)), )); + let children = child_executor( + builder.session.clone(), + Arc::clone(&file.locations), + Arc::clone(&io), + ); Ok(Self { session: builder.session.clone(), plans: ScanPlans { @@ -316,6 +309,7 @@ impl PreparedScan { projection, row_offset: builder.row_offset, values: Default::default(), + children, }, io, splits, @@ -328,22 +322,22 @@ impl PreparedScan { /// The splits with selected rows, each ready to be admitted to a planning run. fn split_plans(&self) -> Vec { + let selection = Arc::new(self.selection.clone()); self.splits .iter() .filter_map(|rows| { - let row_mask = self.selection.row_mask(rows); - if row_mask.mask().all_false() { + if self.selection.row_count_in_range(rows) == 0 { return None; } let scope = WorkScope { file_ordinal: 0, - rows: row_mask.row_range(), + rows: rows.clone(), }; Some(SplitPlan { root: plan_projection( self.plans.clone(), scope.clone(), - row_mask.mask().clone(), + Arc::clone(&selection), ), scope, }) diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index a910172f87b..43bb6392757 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -16,6 +16,7 @@ use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; @@ -36,6 +37,7 @@ use vortex_error::VortexResult; use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use super::MAX_ROWS_PER_CHUNK; use super::PreparedScan; @@ -45,9 +47,13 @@ use super::row_chunks; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::children::OwnedLayoutChildren; +use crate::layouts::chunked::ChunkedLayout; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::dict::DictLayout; use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::foreign::new_foreign_layout; +use crate::layouts::list::writer::ListLayoutStrategy; use crate::layouts::row_idx::row_idx; use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; @@ -66,6 +72,9 @@ use crate::test::new_session; const CHUNK_ROWS: usize = 1000; const CHUNKS: usize = 4; +mod child_work; +static UNSUPPORTED_ID: CachedId = CachedId::new("test.unsupported"); + /// A table chunk of `CHUNK_ROWS` rows starting at row `start`: an integer, a nullable string and a /// list column. fn table_chunk(start: usize) -> VortexResult { @@ -314,19 +323,29 @@ fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) } #[rstest] -#[case::empty(0, 0, 65_536)] -#[case::single_chunk(1000, 100_000_000, 65_536)] -#[case::no_bytes(100_000, 0, 65_536)] -#[case::small_rows(1_000_000, 64_000_000, 65_536)] -#[case::large_rows(1_000_000, 384_000_000, 16_384)] -#[case::minimum(1_000_000, 2_000_000_000, 8192)] -fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: usize) { +#[case::empty(0, 0, None, 65_536)] +#[case::single_chunk(1000, 100_000_000, None, 65_536)] +#[case::no_bytes(100_000, 0, None, 65_536)] +#[case::small_rows(1_000_000, 64_000_000, None, 65_536)] +#[case::large_rows(1_000_000, 384_000_000, None, 16_384)] +#[case::minimum(1_000_000, 2_000_000_000, None, 8192)] +#[case::sparse_small_rows(3_339_715, 49_808_116, Some(834_928), 524_288)] +#[case::sparse_medium_rows(1_000_000, 35_640_884, Some(250_000), 131_072)] +#[case::sparse_more_workers(3_339_715, 49_808_116, Some(27_830), 65_536)] +#[case::sparse_lists(1_000_000, 86_983_648, Some(250_000), 65_536)] +#[case::sparse_wide_rows(1_000_000, 3_354_532_688, Some(250_000), 8192)] +fn chunk_byte_budget( + #[case] rows: u64, + #[case] bytes: u32, + #[case] sparse_target: Option, + #[case] expected: usize, +) { let locations = [SegmentLocation { offset: 0, length: bytes, alignment: Alignment::none(), }]; - assert_eq!(rows_per_chunk(rows, &locations), expected); + assert_eq!(rows_per_chunk(rows, &locations, sparse_target), expected); } #[test] @@ -349,6 +368,46 @@ impl SegmentSource for CountingSegments { } } +#[tokio::test] +async fn tiny_selection_shares_a_stored_segment() -> VortexResult<()> { + let session = new_session().with_tokio(); + let input = PrimitiveArray::from_iter(0_i32..1_048_576).into_array(); + let dtype = input.dtype().clone(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new( + dtype.clone(), + stream::iter([Ok((sequence.advance(), input))]), + ) + .sendable(), + eof, + &session, + ) + .await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments: Arc = Arc::::clone(&counted); + // These rows share one stored segment but fall across the old worker-sized task boundary, + // even on one core. Selection-aware sizing lets one projection share its read. + let scan = builder(&session, &segments, &layout)?.with_row_indices( + StrictSortedBuffer::try_new(Buffer::from(vec![10, 350_000]))?, + ); + let actual = collect(dtype, scan.into_stream()?).await?; + assert_arrays_eq!( + actual, + Buffer::from(vec![10_i32, 350_000]).into_array(), + &mut session.create_execution_ctx() + ); + assert_eq!(counted.reads.load(Ordering::Relaxed), 1); + Ok(()) +} + #[tokio::test] async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -416,19 +475,94 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } -/// V2 rejects list layouts: their element ranges are only known once their offsets are read. #[tokio::test] -async fn list_layouts_return_an_error() -> VortexResult<()> { +async fn list_layouts_match_v1() -> VortexResult<()> { let session = new_session().with_tokio(); let strategy = TableStrategy::new(chunked_flat(), chunked_flat()).with_list_layout(); let (segments, layout) = write(&session, &strategy).await?; - let builder = builder(&session, &segments, &layout)?; - let error = builder - .into_stream()? - .try_next() - .await - .expect_err("V2 must reject an unsupported layout"); - assert!(error.to_string().contains("V2 scans do not support")); + assert_same_as_v1(&session, || builder(&session, &segments, &layout)).await +} + +#[rstest] +#[case::empty(vec![])] +#[case::null_only(vec![3, 6, 9])] +#[case::contiguous(vec![1, 2])] +#[case::contiguous_with_empty(vec![7, 8])] +#[case::mixed(vec![0, 1, 2, 3, 8, 31, 63])] +#[case::all((0..64).collect())] +#[tokio::test] +async fn selected_list_values_match_input( + #[case] rows: Vec, + #[values(false, true)] nested: bool, + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let mut offsets = vec![0_u32]; + let mut total = 0; + for row in 0..64 { + total += row % 4; + offsets.push(total); + } + let inner_strategy = ListLayoutStrategy::default() + .with_elements(chunked_flat()) + .with_offsets(chunked_flat()) + .with_validity(chunked_flat()); + let (elements, element_strategy) = if nested { + let elements = ListArray::try_new( + Buffer::from_iter(0..2 * total).into_array(), + Buffer::from_iter((0..=total).map(|row| 2 * row)).into_array(), + Validity::Array(BoolArray::from_iter((0..total).map(|row| row % 3 != 0)).into_array()), + )? + .into_array(); + ( + elements, + Arc::new(inner_strategy) as Arc, + ) + } else { + (Buffer::from_iter(0..total).into_array(), chunked_flat()) + }; + let validity = if nullable { + Validity::Array(BoolArray::from_iter((0..64).map(|row| row % 3 != 0)).into_array()) + } else { + Validity::NonNullable + }; + let input = + ListArray::try_new(elements, Buffer::from(offsets).into_array(), validity)?.into_array(); + let dtype = input.dtype().clone(); + let expected = input.take(Buffer::from(rows.clone()).into_array())?; + let strategy = ListLayoutStrategy::default() + .with_elements(element_strategy) + .with_offsets(chunked_flat()) + .with_validity(chunked_flat()); + let segments = Arc::new(TestSegments::default()); + let (mut id, eof) = SequenceId::root().split(); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype.clone(), stream::iter([Ok((id.advance(), input))])) + .sendable(), + eof, + &session, + ) + .await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments: Arc = Arc::::clone(&counted); + let empty = rows.is_empty(); + let null_only = nullable && !empty && rows.iter().all(|row| row % 3 == 0); + let scan = builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(rows))?); + let actual = collect(dtype, scan.into_stream()?).await?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + if empty { + assert_eq!(counted.reads.load(Ordering::Relaxed), 0); + } else if null_only { + // Only offsets and validity are needed, even when null lists store nonempty elements. + assert_eq!(counted.reads.load(Ordering::Relaxed), 2); + } Ok(()) } @@ -481,3 +615,72 @@ async fn work_chunks_ignore_layout_boundaries( assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); Ok(()) } + +#[tokio::test] +async fn unknown_layout_returns_an_error() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, known) = write(&session, chunked_flat().as_ref()).await?; + let layout = new_foreign_layout( + *UNSUPPORTED_ID, + known.dtype().clone(), + known.row_count(), + Vec::new(), + Vec::new(), + Vec::new(), + ); + let error = builder(&session, &segments, &layout)? + .into_stream()? + .try_next() + .await + .expect_err("V2 must reject an unsupported layout"); + assert!( + error + .to_string() + .contains("No physical plan implementation for layout 'test.unsupported'") + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn unsupported_chunk_is_checked_when_reached( + #[values(false, true)] selected: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, known) = write(&session, chunked_flat().as_ref()).await?; + let rows = known.row_count(); + let unknown = new_foreign_layout( + *UNSUPPORTED_ID, + known.dtype().clone(), + rows, + Vec::new(), + Vec::new(), + Vec::new(), + ); + let layout = ChunkedLayout::new( + 2 * rows, + known.dtype().clone(), + OwnedLayoutChildren::layout_children(vec![Arc::clone(&known), unknown]), + ) + .into_layout(); + let range = if selected { rows..2 * rows } else { 0..rows }; + let result = collect( + known.dtype().clone(), + builder(&session, &segments, &layout)? + .with_row_range(range) + .into_stream()?, + ) + .await; + if selected { + let error = result.expect_err("a reached unsupported chunk must fail"); + assert!(error.to_string().contains("test.unsupported")); + } else { + let expected = collect( + known.dtype().clone(), + builder(&session, &segments, &known)?.into_stream()?, + ) + .await?; + assert_arrays_eq!(result?, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} diff --git a/vortex-layout/src/scan/v2/tests/child_work.rs b/vortex-layout/src/scan/v2/tests/child_work.rs new file mode 100644 index 00000000000..d26a2d97c8c --- /dev/null +++ b/vortex-layout/src/scan/v2/tests/child_work.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_utils::parallelism::get_available_parallelism; + +use super::*; +use crate::layouts::list::ListLayout; +async fn store_array( + session: &VortexSession, + segments: &Arc, + array: ArrayRef, + chunk_rows: usize, +) -> VortexResult { + let dtype = array.dtype().clone(); + let (mut id, eof) = SequenceId::root().split(); + let chunks = (0..array.len()) + .step_by(chunk_rows) + .map(|start| { + Ok(( + id.advance(), + array.slice(start..(start + chunk_rows).min(array.len()))?, + )) + }) + .collect::>(); + chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(segments), + SequentialStreamAdapter::new(dtype, stream::iter(chunks)).sendable(), + eof, + session, + ) + .await +} + +#[rstest] +#[case::empty(vec![], false)] +#[case::empty_rows(vec![0, 2], false)] +#[case::null_row(vec![3], false)] +#[case::large_row(vec![1], true)] +#[case::mixed(vec![0, 1, 3], true)] +#[case::all(vec![0, 1, 2, 3], true)] +#[tokio::test] +async fn large_list_children_match_input_and_v1( + #[case] selected: Vec, + #[case] parallel: bool, + #[values(false, true)] nested: bool, + #[values(8192, 262144)] chunk_rows: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let primitive = PrimitiveArray::from_option_iter( + (0_i64..262144).map(|value| (value % 7 != 0).then_some(value)), + ) + .into_array(); + let mut elements = primitive.clone(); + let mut layout = store_array(&session, &segments, primitive, chunk_rows).await?; + if nested { + let offsets = Buffer::from_iter((0_u64..=2048).map(|row| row * 128)).into_array(); + let valid = BoolArray::from_iter((0..2048).map(|row| row % 11 != 0)).into_array(); + elements = ListArray::try_new(elements, offsets.clone(), Validity::Array(valid.clone()))? + .into_array(); + layout = ListLayout::new( + elements.dtype().clone(), + layout, + store_array(&session, &segments, offsets, 2049).await?, + Some(store_array(&session, &segments, valid, 2048).await?), + ) + .into_layout(); + } + let half = u64::try_from(elements.len())? / 2; + let offsets = Buffer::from(vec![0_u64, 0, half, half, 2 * half]).into_array(); + let valid = BoolArray::from_iter([true, true, true, false]).into_array(); + let input = + ListArray::try_new(elements, offsets.clone(), Validity::Array(valid.clone()))?.into_array(); + let layout = ListLayout::new( + input.dtype().clone(), + layout, + store_array(&session, &segments, offsets, 5).await?, + Some(store_array(&session, &segments, valid, 4).await?), + ) + .into_layout(); + let segments: Arc = segments; + let make = || { + Ok(builder(&session, &segments, &layout)? + .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(selected.clone()))?)) + }; + let scan = PreparedScan::try_new(make()?)?; + let children = scan.plans.children.clone(); + let actual = collect(input.dtype().clone(), scan.execute_stream()?).await?; + assert_arrays_eq!( + actual, + input.take(Buffer::from(selected.clone()).into_array())?, + &mut session.create_execution_ctx() + ); + assert_eq!( + children.is_some_and(|children| children.admitted() > 1), + parallel && chunk_rows == 8192 && get_available_parallelism().unwrap_or(1) > 1 + ); + assert_same_as_v1(&session, make).await +} diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d12c207b049..63352a21566 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -26,6 +26,7 @@ vortex-utils = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } roaring = { workspace = true } +smallvec = { workspace = true } tracing = { workspace = true } [dev-dependencies] diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs index 8cb78099272..024748774fc 100644 --- a/vortex-scan/src/planning/driver/mod.rs +++ b/vortex-scan/src/planning/driver/mod.rs @@ -7,12 +7,12 @@ //! batches to its caller as they are produced. Runnable items are prioritised in DFS preorder, //! with roots ordered by admission and children by emission. Each visit takes at most one CPU //! step. A parent's continuation follows its emitted child's entire subtree. Waiting work is -//! skipped and regains its priority when IO makes it runnable. Published IO is registered +//! skipped and regains its priority when IO or an asynchronous task makes it runnable. Published IO is registered //! immediately and state is inspected again before parking. Output follows emission order, //! not global row order. An earlier runnable branch can starve later branches. //! //! A run never blocks: [`Run::advance`] returns [`Progress::Waiting`] when every live item waits -//! for IO, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing +//! for IO or tasks, and the caller decides how to wait, by awaiting [`Run::poll_completion`] or by handing //! a completion it took from a source to [`Run::complete`]. [`Driver::run`] is the blocking //! convenience for one root. @@ -21,7 +21,12 @@ 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::future::BoxFuture; +use smallvec::SmallVec; +use smallvec::smallvec; use vortex_array::ArrayRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -83,6 +88,9 @@ enum Output { Morsel(MorselOutput), } +// Root, selection, projection, and continuation fit inline. Deeper compositions can spill. +type Priority = SmallVec<[u64; 4]>; + /// One live planner or morsel. Boxed in the run queue and the parked map, so scheduling moves a /// pointer. struct Work { @@ -91,7 +99,7 @@ struct Work { scope: WorkScope, item: Item, /// Stable child-ordinal path, beginning with the root's admission ordinal. - path: Vec, + path: Priority, /// The parent continuation sorts after all previously emitted child subtrees. next_child: u64, /// Published `Fetch` requests not yet delivered, with the target each must be answered with. @@ -99,7 +107,7 @@ struct Work { } impl Work { - fn priority(&self) -> Vec { + fn priority(&self) -> Priority { let mut priority = self.path.clone(); priority.push(self.next_child); priority @@ -107,7 +115,7 @@ impl Work { // Child k takes the continuation's old key as its path. Its subtree then sorts before // the parent's new continuation at k + 1, even if descendants are discovered after IO. - fn child_path(&mut self) -> Vec { + fn child_path(&mut self) -> Priority { let path = self.priority(); self.next_child += 1; path @@ -146,7 +154,7 @@ impl Work { /// Keeps scheduling order separate from lookup for deliveries to runnable work. #[derive(Default)] struct ReadyQueue { - priorities: BTreeMap, IoOwnerId>, + priorities: BTreeMap, work: HashMap>, } @@ -191,6 +199,8 @@ pub struct Run { next_owner: u64, queue: ReadyQueue, parked: HashMap>, + /// Allocated only for owners that publish asynchronous work. + tasks: HashMap>>, /// Batches and finished roots not yet returned by [`advance`](Self::advance). events: VecDeque, /// Visits left before the sources are polled for completions again. @@ -233,7 +243,7 @@ impl Run { let id = RootId(self.next_root); self.next_root += 1; self.roots.insert(id, Root { io, live: 0 }); - self.spawn(id, scope, Item::Planner(root), vec![id.0]); + self.spawn(id, scope, Item::Planner(root), smallvec![id.0]); id } @@ -245,6 +255,9 @@ impl Run { }; self.queue.cancel(root); self.parked.retain(|_, work| work.root != root); + self.tasks.retain(|owner, _| { + self.queue.work.contains_key(owner) || self.parked.contains_key(owner) + }); cancelled.io.clear(); } @@ -297,10 +310,10 @@ impl Run { let work = parked.get_mut(); work.deliver(request, result)?; // An item that still waits for other requests stays parked. - if work.state() != State::Waiting { + if work.state() != State::Waiting && !self.tasks.contains_key(&owner) { let work = parked.remove(); self.queue.insert(work); - } else if work.outstanding.is_empty() { + } else if work.outstanding.is_empty() && !self.tasks.contains_key(&owner) { vortex_bail!( "work for scope {:?} waits with no outstanding fetch", work.scope @@ -332,6 +345,23 @@ impl Run { } let mut delivered = false; for io in sources { + // A root may now wait solely for child work. IO sources need not accept a poll + // with no fetches outstanding. + if !self.tasks.is_empty() + && !self + .parked + .values() + .chain(self.queue.work.values()) + .any(|work| { + !work.outstanding.is_empty() + && self + .roots + .get(&work.root) + .is_some_and(|root| Arc::ptr_eq(&root.io, &io)) + }) + { + continue; + } match io.poll_completion(cx) { Poll::Ready(Ok(completion)) => { if let Err(error) = self.complete(completion) { @@ -343,6 +373,10 @@ impl Run { Poll::Pending => {} } } + match self.take_tasks(cx) { + Ok(completed) => delivered |= completed, + Err(error) => return Poll::Ready(Err(error)), + } if delivered { Poll::Ready(Ok(())) } else { @@ -352,6 +386,7 @@ impl Run { /// Delivers the completions every root's session has ready, without blocking. fn take_completions(&mut self) -> VortexResult<()> { + self.take_tasks(&mut Context::from_waker(Waker::noop()))?; if self.roots.len() == 1 { // The common case, without collecting the sessions. let Some(io) = self.roots.values().next().map(|root| Arc::clone(&root.io)) else { @@ -375,7 +410,35 @@ impl Run { Ok(()) } - fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { + fn take_tasks(&mut self, cx: &mut Context<'_>) -> VortexResult { + if self.tasks.is_empty() { + return Ok(false); + } + let mut completed = Vec::new(); + for (&owner, future) in &mut self.tasks { + if let Poll::Ready(result) = future.poll_unpin(cx) { + result?; + completed.push(owner); + } + } + let delivered = !completed.is_empty(); + for owner in completed { + self.tasks.remove(&owner); + let work = self + .parked + .remove(&owner) + .ok_or_else(|| vortex_err!("future completed for unknown planner {owner:?}"))?; + self.settle(work)?; + } + for work in self.parked.values() { + if work.outstanding.is_empty() && !self.tasks.contains_key(&work.id) { + vortex_bail!("work waits after its last task and fetch completed"); + } + } + Ok(delivered) + } + + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Priority) { let id = IoOwnerId(self.next_owner); self.next_owner += 1; if let Some(root) = self.roots.get_mut(&root) { @@ -395,6 +458,7 @@ impl Run { /// Retires a finished item, and its root once nothing descended from it is live. fn retire(&mut self, work: &Work) { + self.tasks.remove(&work.id); let Entry::Occupied(mut root) = self.roots.entry(work.root) else { return; }; @@ -406,9 +470,9 @@ impl Run { } } - /// Parks an item that waits for IO, which is a protocol error with no fetch outstanding. + /// Parks an item that waits for IO or tasks; at least one must be outstanding. fn park(&mut self, work: Box) -> VortexResult<()> { - if work.outstanding.is_empty() { + if work.outstanding.is_empty() && !self.tasks.contains_key(&work.id) { vortex_bail!( "work for scope {:?} waits with no outstanding fetch", work.scope @@ -452,6 +516,12 @@ impl Run { self.submit(&mut work, batch)?; self.settle(work)?; } + Output::Planner(PlannerOutput::Await(future)) => { + self.tasks.insert(work.id, future); + // Await suspends this planner even if its next compute is already prepared. + // Outstanding IO still belongs to it and can be delivered while it is suspended. + self.parked.insert(work.id, work); + } Output::Planner(PlannerOutput::Planner(scope, child)) => { let path = work.child_path(); self.spawn(work.root, scope, Item::Planner(child), path); @@ -541,7 +611,10 @@ impl Driver { match run.advance()? { Progress::Batch(batch) => batches.push(batch), Progress::RootDone(_) | Progress::Idle => return Ok(batches), - Progress::Waiting => run.complete(self.io.wait()?)?, + Progress::Waiting if run.tasks.is_empty() => run.complete(self.io.wait()?)?, + Progress::Waiting => { + futures::executor::block_on(std::future::poll_fn(|cx| run.poll_completion(cx)))? + } } } } diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs index 389ad24372d..d27fa0a5be7 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -22,6 +22,8 @@ use crate::planning::tests::scripted::array_of; use crate::planning::tests::scripted::ctx; use crate::planning::tests::scripted::scope; +mod tasks; + fn request(id: u32, offset: u64) -> IoRequest { IoRequest { intent: IoIntent::Fetch, @@ -141,6 +143,34 @@ fn dfs_finishes_child_subtree_before_parent_continues() -> VortexResult<()> { Ok(()) } +#[rstest] +fn nested_planner_chains_preserve_sibling_order( + #[values(1, 3, 8, 16)] depth: u64, +) -> VortexResult<()> { + let log = Log::default(); + let mut steps = vec![one_batch_morsel(scope(0..1), 1), PlannerStep::Done]; + for row in 1..=depth { + steps = vec![ + PlannerStep::Planner(scope(0..row), steps), + one_batch_morsel(scope(row..row + 1), 1), + PlannerStep::Done, + ]; + } + let root = ScriptedPlanner::boxed("root", steps, &log); + let batches = driver(&source()).run(root)?; + let scopes: Vec<_> = batches.iter().map(|batch| batch.scope.clone()).collect(); + assert_eq!( + scopes, + (0..=depth) + .map(|row| scope(row..row + 1)) + .collect::>() + ); + for batch in batches { + assert_arrays_eq!(batch.array, array_of(1), &mut ctx()); + } + Ok(()) +} + #[test] fn needs_io_delivers_every_request_in_order() -> VortexResult<()> { let log = Log::default(); diff --git a/vortex-scan/src/planning/driver/tests/tasks.rs b/vortex-scan/src/planning/driver/tests/tasks.rs new file mode 100644 index 00000000000..140a06d1a3d --- /dev/null +++ b/vortex-scan/src/planning/driver/tests/tasks.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use futures::channel::oneshot; +use futures::task::ArcWake; +use futures::task::waker; +use vortex_error::vortex_err; + +use super::*; + +struct Awaiting { + tasks: VecDeque>>, + read: Option, + waiting_read: bool, + emitted: bool, +} + +impl IoConsumer for Awaiting { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) { + self.waiting_read = false; + } +} + +impl Planner for Awaiting { + fn state(&self) -> State { + if self.emitted { + State::Done + } else if self.tasks.is_empty() && self.read.is_none() && self.waiting_read { + State::Waiting + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + if let Some(read) = self.read.take() { + return Ok(PlannerOutput::NeedsIO(vec![read])); + } + if let Some(task) = self.tasks.pop_front() { + return Ok(PlannerOutput::Await(task)); + } + self.emitted = true; + Ok(PlannerOutput::Planner( + scope(0..3), + ScriptedPlanner::boxed( + "after task", + vec![one_batch_morsel(scope(0..3), 3), PlannerStep::Done], + &Log::default(), + ), + )) + } +} + +fn awaiting(tasks: Vec>>, read: bool) -> Box { + Box::new(Awaiting { + tasks: tasks.into(), + read: read.then(|| request(0, 0)), + waiting_read: read, + emitted: false, + }) +} + +#[derive(Default)] +struct WakeCount(AtomicUsize); + +impl ArcWake for WakeCount { + fn wake_by_ref(arc_self: &Arc) { + arc_self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[rstest] +fn completion_wakes_a_planner_and_preserves_its_continuation( + #[values(false, true)] read: bool, +) -> VortexResult<()> { + let (send, recv) = oneshot::channel::<()>(); + let task = async move { recv.await.map_err(|_| vortex_err!("cancelled")) }.boxed(); + let root = awaiting(vec![task], read); + let mut run = Run::new(); + run.admit(root, scope(0..3), source()); + assert!(matches!(run.advance()?, Progress::Waiting)); + let wakes = Arc::new(WakeCount::default()); + let waker = waker(Arc::clone(&wakes)); + if read { + let Poll::Ready(result) = run.poll_completion(&mut Context::from_waker(&waker)) else { + panic!("the canned read should complete") + }; + result?; + assert!(matches!(run.advance()?, Progress::Waiting)); + } + assert!( + run.poll_completion(&mut Context::from_waker(&waker)) + .is_pending() + ); + send.send(()).map_err(|_| vortex_err!("receiver dropped"))?; + assert!(wakes.0.load(Ordering::Relaxed) > 0); + let Progress::Batch(batch) = run.advance()? else { + panic!("missing continuation output") + }; + assert_arrays_eq!(batch.array, array_of(3), &mut ctx()); + assert!(matches!(run.advance()?, Progress::RootDone(_))); + Ok(()) +} + +#[test] +fn completed_future_keeps_waiting_for_outstanding_io() -> VortexResult<()> { + let root = awaiting(vec![async { Ok(()) }.boxed()], true); + let mut run = Run::new(); + run.admit(root, scope(0..3), source()); + assert!(matches!(run.advance()?, Progress::Waiting)); + assert!(run.tasks.is_empty()); + let Poll::Ready(result) = run.poll_completion(&mut Context::from_waker(Waker::noop())) else { + panic!("the outstanding read should complete") + }; + result?; + assert!(matches!(run.advance()?, Progress::Batch(_))); + Ok(()) +} + +#[test] +fn repeated_awaits_preserve_the_planner_continuation() -> VortexResult<()> { + let tasks = (0..3).map(|_| async { Ok(()) }.boxed()).collect(); + let root = awaiting(tasks, false); + assert_eq!(driver(&source()).run(root)?.len(), 1); + Ok(()) +} + +struct DropCount(Arc); + +impl Drop for DropCount { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[rstest] +fn cancellation_drops_pending_child_work( + #[values(false, true)] cancel_root: bool, +) -> VortexResult<()> { + let drops = Arc::new(AtomicUsize::new(0)); + let guard = DropCount(Arc::clone(&drops)); + let root = awaiting( + vec![ + async move { + let _guard = guard; + futures::future::pending().await + } + .boxed(), + ], + false, + ); + let mut run = Run::new(); + let root = run.admit(root, scope(0..3), source()); + assert!(matches!(run.advance()?, Progress::Waiting)); + assert_eq!(drops.load(Ordering::Relaxed), 0); + if cancel_root { + run.cancel(root); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + drop(run); + assert_eq!(drops.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn task_failure_fails_the_run() { + let root = awaiting( + vec![async { Err(vortex_err!("child task failed")) }.boxed()], + false, + ); + let mut run = Run::new(); + run.admit(root, scope(0..3), source()); + let error = run.advance().err().expect("child error must fail the run"); + assert!(error.to_string().contains("child task failed")); +} diff --git a/vortex-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 3095712f161..2c69e734294 100644 --- a/vortex-scan/src/planning/mod.rs +++ b/vortex-scan/src/planning/mod.rs @@ -10,6 +10,7 @@ pub mod driver; pub mod morsel; pub mod next; pub mod planner; +pub mod selection; #[cfg(test)] mod tests; diff --git a/vortex-scan/src/planning/planner.rs b/vortex-scan/src/planning/planner.rs index 91d875ff317..cf0b71ad569 100644 --- a/vortex-scan/src/planning/planner.rs +++ b/vortex-scan/src/planning/planner.rs @@ -6,6 +6,7 @@ use std::fmt; use std::ops::Range; +use futures::future::BoxFuture; use vortex_error::VortexResult; use vortex_io::request::IoBatch; use vortex_io::request::IoConsumer; @@ -49,6 +50,10 @@ pub enum PlannerOutput { Continue, /// Requests the planner publishes: fetches it needs, and optional hints. NeedsIO(IoBatch), + /// Suspends this planner until the future completes. The future owns its child tasks; + /// dropping the planner or cancelling its root drops the future. The planner retains + /// results through its own channel and resumes without a separate completion callback. + Await(BoxFuture<'static, VortexResult<()>>), /// A child planner, with the scope it is responsible for. Planner(WorkScope, Box), /// A morsel authorised to produce arrays for the scope. @@ -74,6 +79,7 @@ impl fmt::Debug for PlannerOutput { Self::Done => f.write_str("Done"), Self::Continue => f.write_str("Continue"), Self::NeedsIO(batch) => f.debug_tuple("NeedsIO").field(batch).finish(), + Self::Await(_) => f.write_str("Await"), Self::Planner(scope, _) => f.debug_tuple("Planner").field(scope).finish(), Self::Morsel(scope, _) => f.debug_tuple("Morsel").field(scope).finish(), } diff --git a/vortex-scan/src/planning/selection.rs b/vortex-scan/src/planning/selection.rs new file mode 100644 index 00000000000..75b389690bb --- /dev/null +++ b/vortex-scan/src/planning/selection.rs @@ -0,0 +1,210 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A selection stage that hands its surviving rows to another planner. + +use std::sync::Arc; + +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_mask::Mask; + +use crate::planning::next::Next; +use crate::planning::planner::Planner; +use crate::planning::planner::PlannerOutput; +use crate::planning::planner::State; +use crate::planning::planner::WorkScope; +use crate::selection::Selection; + +/// Selected rows passed between planning stages. +/// +/// The mask is relative to `scope.rows.start`, and its length equals the scope's row count. +/// A stage can narrow the scope while rebasing the mask to preserve absolute row coordinates. +pub struct SelectedRows { + /// The input row domain covering the selected rows, including any unselected gaps. + pub scope: WorkScope, + /// The rows within that domain that the successor should process. + pub mask: Mask, +} + +/// Applies a scan selection to a scope, then constructs the next planner with the surviving rows. +/// +/// The successor may project, filter, or further partition the rows. This stage neither reads +/// data nor decides how those rows are executed. An empty selection creates no successor. +/// Compact selections exclude unselected leading and trailing rows from the successor's scope. +pub struct SelectionPlanner { + scope: Option, + selection: Arc, + next: Option>, +} + +impl SelectionPlanner { + /// Creates an unevaluated selection stage. The mask and successor are built in `compute`. + pub fn new(scope: WorkScope, selection: Arc, next: Next) -> Self { + Self { + scope: Some(scope), + selection, + next: Some(next), + } + } +} + +impl IoConsumer for SelectionPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for SelectionPlanner { + fn state(&self) -> State { + if self.scope.is_some() { + State::NeedsCompute + } else { + State::Done + } + } + + fn compute(&mut self) -> VortexResult { + let mut scope = self + .scope + .take() + .ok_or_else(|| vortex_err!("SelectionPlanner computed after completion"))?; + let selected = self.selection.row_mask(&scope.rows); + let Some(start) = selected.mask().first() else { + return Ok(PlannerOutput::Done); + }; + let end = selected.mask().last().vortex_expect("nonempty selection") + 1; + let mut mask = selected.mask().clone(); + // Narrow clusters to avoid planning unrelated chunks. Preserve alignment for sparse + // gathers: slicing their encoded inputs at arbitrary boundaries can cost more than it saves. + if end - start <= mask.true_count().saturating_mul(16) { + mask = mask.slice(start..end); + scope.rows = scope.rows.start + start as u64..scope.rows.start + end as u64; + } + let next = self + .next + .take() + .ok_or_else(|| vortex_err!("SelectionPlanner has no successor"))?; + let child = next(SelectedRows { + scope: scope.clone(), + mask, + })?; + Ok(PlannerOutput::Planner(scope, child)) + } +} + +#[cfg(test)] +mod tests { + use parking_lot::Mutex; + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_error::vortex_bail; + + use super::*; + use crate::planning::next::next_fn; + use crate::strict_sorted_buffer::StrictSortedBuffer; + + struct Finished; + + impl IoConsumer for Finished { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} + } + + impl Planner for Finished { + fn state(&self) -> State { + State::Done + } + + fn compute(&mut self) -> VortexResult { + Ok(PlannerOutput::Done) + } + } + + #[rstest] + #[case::all(Selection::All, vec![10, 11, 12, 13, 14])] + #[case::include(Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![1, 10, 13, 20]))?), vec![10, 13])] + #[case::exclude(Selection::ExcludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![1, 10, 13, 20]))?), vec![11, 12, 14])] + #[case::roaring(Selection::IncludeRoaring([1, 10, 13, 20].into_iter().collect()), vec![10, 13])] + #[case::exclude_roaring(Selection::ExcludeRoaring([1, 10, 13, 20].into_iter().collect()), vec![11, 12, 14])] + #[case::interior(Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![11, 12, 13]))?), vec![11, 12, 13])] + #[case::singleton(Selection::IncludeRoaring([12].into_iter().collect()), vec![12])] + #[case::excluded_edges(Selection::ExcludeRoaring([10, 14].into_iter().collect()), vec![11, 12, 13])] + #[case::empty(Selection::IncludeRoaring(Default::default()), vec![])] + fn hands_selected_rows_to_successor( + #[case] selection: Selection, + #[case] expected: Vec, + ) -> VortexResult<()> { + let scope = WorkScope { + file_ordinal: 7, + rows: 10..15, + }; + let received = Arc::new(Mutex::new(None)); + let output = Arc::clone(&received); + let next = next_fn(move |selected: SelectedRows| { + *output.lock() = Some(selected); + Ok(Finished) + }); + let mut planner = SelectionPlanner::new(scope.clone(), Arc::new(selection), next); + assert_eq!(planner.state(), State::NeedsCompute); + assert!(received.lock().is_none()); + let result = planner.compute()?; + assert_eq!(planner.state(), State::Done); + let selected = received.lock().take(); + if expected.is_empty() { + assert!(matches!(result, PlannerOutput::Done)); + assert!(selected.is_none()); + } else { + let PlannerOutput::Planner(child_scope, _) = result else { + vortex_bail!("selection should emit its successor"); + }; + assert_eq!(child_scope.file_ordinal, scope.file_ordinal); + assert_eq!( + child_scope.rows, + expected[0]..expected[expected.len() - 1] + 1 + ); + let selected = selected.ok_or_else(|| vortex_err!("missing successor input"))?; + assert_eq!(selected.scope, child_scope); + assert_eq!( + selected.mask.len() as u64, + child_scope.rows.end - child_scope.rows.start + ); + let rows = selected + .mask + .iter() + .enumerate() + .filter(|(_, selected)| *selected) + .map(|(index, _)| selected.scope.rows.start + index as u64) + .collect::>(); + assert_eq!(rows, expected); + } + assert!(planner.compute().is_err()); + Ok(()) + } + + #[test] + fn scattered_selection_keeps_its_aligned_scope() -> VortexResult<()> { + let scope = WorkScope { + file_ordinal: 0, + rows: 0..65_536, + }; + let expected = scope.clone(); + let next = next_fn(move |selected: SelectedRows| { + assert_eq!(selected.scope, expected); + assert_eq!(selected.mask.len(), 65_536); + assert_eq!( + selected.mask.indices(), + vortex_mask::AllOr::Some(&[100, 50_000][..]) + ); + Ok(Finished) + }); + let selection = + Selection::IncludeByIndex(StrictSortedBuffer::try_new(Buffer::from(vec![ + 100, 50_000, + ]))?); + let mut planner = SelectionPlanner::new(scope, Arc::new(selection), next); + assert!(matches!(planner.compute()?, PlannerOutput::Planner(..))); + Ok(()) + } +} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index fb961a63dab..9b9138fc46d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -56,6 +56,7 @@ arrow-array = { workspace = true } arrow-schema = { workspace = true } divan = { workspace = true } fastlanes = { workspace = true } +futures = { workspace = true } mimalloc = { workspace = true } parquet = { workspace = true } rand = { workspace = true }