diff --git a/Cargo.lock b/Cargo.lock index 5224e8d45f3..342c9883b29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12379,7 +12379,9 @@ version = "0.1.0" dependencies = [ "async-trait", "futures", + "parking_lot", "roaring", + "rstest", "tracing", "vortex-array", "vortex-buffer", @@ -12387,6 +12389,7 @@ dependencies = [ "vortex-io", "vortex-mask", "vortex-session", + "vortex-utils", ] [[package]] diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index 26a1bee31fa..fbdeb027554 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -19,6 +19,7 @@ use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; +use vortex::layout::scan::v2::ScanBuilder; use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; @@ -81,8 +82,9 @@ impl Compressor for VortexCompressor { /// Scan every record batch out of an in-memory Vortex file. async fn vortex_decompress_read(data: Bytes) -> Result<()> { - let mut scan = SESSION.open_options().open_buffer(data)?.scan()?; - let source_dtype = scan.dtype()?; + let file = SESSION.open_options().open_buffer(data)?; + let mut scan = ScanBuilder::new(SESSION.clone(), file.scan_file()); + let source_dtype = scan.dtype().clone(); let root_columns = source_dtype .as_struct_fields_opt() .map_or(0, |fields| fields.nfields()); @@ -94,7 +96,7 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> { .optimize_recursive()?; scan = scan.with_projection(projection); } - let schema = Arc::new(SESSION.arrow().to_arrow_schema(&scan.dtype()?)?); + let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?); let stream = scan.into_record_batch_stream(schema)?; pin_mut!(stream); diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 72a5e8787e0..9444d8cb1a5 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -21,6 +21,8 @@ use vortex_layout::LayoutReader; use vortex_layout::scan::layout::LayoutReaderDataSource; use vortex_layout::scan::scan_builder::ScanBuilder; use vortex_layout::scan::split_by::SplitBy; +use vortex_layout::scan::v2::ScanFile; +use vortex_layout::scan::v2::SegmentLocation; use vortex_layout::segments::SegmentSource; use vortex_scan::DataSourceRef; use vortex_session::VortexSession; @@ -221,6 +223,28 @@ impl VortexFile { )) } + /// The parts of this file the planner-driven scan executor reads through: its root layout, + /// where each of its segments lives, and its segment source. + /// + /// Pass this to [`vortex_layout::scan::v2::ScanBuilder::new`] to build a V2 scan explicitly. + pub fn scan_file(&self) -> ScanFile { + let locations = self + .footer + .segment_specs_with_metadata() + .iter() + .map(|spec| SegmentLocation { + offset: spec.offset, + length: spec.length, + alignment: spec.alignment, + }) + .collect(); + ScanFile { + layout: Arc::clone(self.footer.layout()), + locations, + segments: Arc::clone(&self.segment_source), + } + } + /// Returns `true` if file-level statistics prove the expression cannot /// match any rows in this file. /// diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 1b286be4671..22768d7e4f1 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -92,6 +92,10 @@ harness = false name = "aggregate_partial_merge" harness = false +[[bench]] +name = "pipeline_scan" +harness = false + [lints] workspace = true diff --git a/vortex-layout/benches/pipeline_scan.rs b/vortex-layout/benches/pipeline_scan.rs new file mode 100644 index 00000000000..db4beeaa043 --- /dev/null +++ b/vortex-layout/benches/pipeline_scan.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Isolates scan preparation, compilation, and segment delivery from Arrow conversion and IO. + +#![expect(clippy::unwrap_used)] + +use std::hint::black_box; + +use divan::Bencher; +use mimalloc::MiMalloc; +use vortex_array::ArrayContext; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::serde::SerializeOptions; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBufferMut; +use vortex_layout::plan::physical::PackPlan; +use vortex_layout::plan::physical::PlanRef; +use vortex_layout::plan::physical::SegmentScanPlan; +use vortex_layout::plan::physical::pipeline::Scan; +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_session::VortexSession; +use vortex_session::registry::ReadContext; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +const ROWS: u64 = 1024; + +fn main() { + divan::main(); +} + +struct Fixture { + session: VortexSession, + plan: PlanRef, + segment: BufferHandle, +} + +impl Fixture { + fn new(columns: u32) -> Self { + let session = vortex_array::array_session().with::(); + let array = PrimitiveArray::from_iter(0_i32..1024).into_array(); + 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()); + } + let fields = StructFields::from_iter( + (0..columns).map(|index| (format!("c{index}"), array.dtype().clone())), + ); + let children = (0..columns) + .map(|index| { + SegmentScanPlan::new( + array.dtype().clone(), + ROWS, + SegmentId::from(index), + ReadContext::new(context.to_ids()), + None, + ) + .into_plan() + }) + .collect(); + Self { + session, + plan: PackPlan::try_new(fields, Nullability::NonNullable, ROWS, children, None) + .unwrap() + .into_plan(), + segment: BufferHandle::new_host(bytes.freeze()), + } + } + + fn scan(&self) -> Scan { + Scan::try_new( + self.session.clone(), + self.plan.clone(), + vec![Split::all(0..ROWS)], + ) + .unwrap() + } +} + +#[divan::bench(args = [1, 100, 1000])] +fn prepare(bencher: Bencher, columns: u32) { + let fixture = Fixture::new(columns); + bencher.bench_local(|| black_box(fixture.scan())); +} + +#[divan::bench(args = [1, 100, 1000])] +fn prepare_and_compile(bencher: Bencher, columns: u32) { + let fixture = Fixture::new(columns); + bencher.bench_local(|| { + let mut scan = fixture.scan(); + assert!(matches!(scan.step().unwrap(), Turn::Read(_))); + black_box(scan); + }); +} + +#[divan::bench(args = [1, 100, 1000])] +fn scan(bencher: Bencher, columns: u32) { + let fixture = Fixture::new(columns); + bencher.bench_local(|| { + let mut scan = fixture.scan(); + loop { + match scan.step().unwrap() { + Turn::Read(read) => scan.deliver(read.id, fixture.segment.clone()).unwrap(), + Turn::Output(_, array) => { + black_box(array); + } + Turn::Waiting => unreachable!("every read is delivered immediately"), + Turn::Done => break, + } + } + }); +} diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 992830a4eaf..74f5296cd19 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -192,7 +192,7 @@ impl Compiler<'_> { } return port; } - let ranges = std::mem::take(shares.reaches.get_mut(*segment as usize)?); + let ranges = shares.shared.remove(&segment)?; let mut claimants: SmallVec<[usize; 4]> = SmallVec::new(); for rows in &ranges { claimants.extend( @@ -240,10 +240,11 @@ pub(crate) struct Shares { ordered: bool, /// Splits that have finished, and so will claim nothing. finished: Vec, - /// By segment id, which a file numbers densely: for each segment more than one reader will - /// read and whose decoding pipeline is not built yet, the range of plan rows each of its - /// readers reads it for. Every split overlapping a range reads it once for that range. - reaches: Vec; 1]>>, + /// Segment readers collected before the scan starts. A split may read only a few of a + /// file's segments, so recording a reader must not allocate by the file-wide segment id. + reaches: Vec<(SegmentId, Range)>, + /// The row ranges of segments read more than once, until their decoding pipeline is built. + shared: FxHashMap; 1]>>, /// By split: the unclaimed ports of each shared segment the split reads. ports: Vec>>, } @@ -257,6 +258,7 @@ impl Shares { splits, ordered, reaches: Vec::new(), + shared: FxHashMap::default(), } } @@ -290,7 +292,7 @@ impl Shares { pub(crate) fn add(&mut self, plan: &PlanRef, rows: Range) -> VortexResult<()> { let reaches = &mut self.reaches; plan.reach(rows, &Reach::Offset(0), &mut |segment, rows| { - record(reaches, segment, rows) + reaches.push((segment, rows)) }) } @@ -298,23 +300,25 @@ impl Shares { /// pipeline that needs them. pub(crate) fn retain_shared(&mut self) { let mut reaches = std::mem::take(&mut self.reaches); - for ranges in &mut reaches { - let shared = match ranges.as_slice() { - [] => false, - // One reader per split overlapping the range: shared when the range crosses a - // split boundary. - [rows] => self + reaches.sort_unstable_by_key(|(segment, _)| *segment); + for readers in reaches.chunk_by(|(left, _), (right, _)| left == right) { + let (segment, rows) = &readers[0]; + // One reader per split overlapping the range: shared when the range crosses a + // split boundary, or when multiple plan nodes read the segment. + if readers.len() == 1 + && self .overlapping(rows) .filter(|&split| self.overlaps(split, rows)) .nth(1) - .is_some(), - _ => true, - }; - if !shared { - *ranges = SmallVec::new(); + .is_none() + { + continue; } + self.shared.insert( + *segment, + readers.iter().map(|(_, rows)| rows.clone()).collect(), + ); } - self.reaches = reaches; } /// Every port made for a reader that has not claimed it yet. @@ -371,15 +375,6 @@ impl Reach { } } -/// Adds a reader of `segment` reading it for `rows`. -fn record(reaches: &mut Vec; 1]>>, segment: SegmentId, rows: Range) { - let index = *segment as usize; - if index >= reaches.len() { - reaches.resize_with(index + 1, SmallVec::new); - } - reaches[index].push(rows); -} - /// The chunks of `concat` overlapping `rows`: each chunk's index, where it starts, and the rows /// of `rows` it holds, in the concatenation's domain. Found by binary search, so a split pays /// for the chunks it overlaps, not for the chunks of the file. diff --git a/vortex-layout/src/plan/physical/pipeline/port.rs b/vortex-layout/src/plan/physical/pipeline/port.rs index a5c4c38a6e4..e089a9f0f76 100644 --- a/vortex-layout/src/plan/physical/pipeline/port.rs +++ b/vortex-layout/src/plan/physical/pipeline/port.rs @@ -34,9 +34,10 @@ pub(crate) enum SourceEnd { /// A queue of batches with one sink and one source. /// -/// Its buffer is allocated once, for the capacity. Only an unbounded port, a share's waiting for a -/// source compiled later, grows past the default. +/// The first batch is inline: a segment scan commonly emits only one batch. Additional batches +/// use a deque, whose buffer is retained when the port is reused. pub(crate) struct Queue { + front: Option, batches: VecDeque, /// Batches the queue may hold before its sink is blocked. capacity: usize, @@ -57,7 +58,8 @@ pub(crate) struct Queue { impl Queue { fn new(capacity: usize, sink: Option, source: SourceEnd) -> Self { Self { - batches: VecDeque::with_capacity(capacity.clamp(1, DEFAULT_CAPACITY)), + front: None, + batches: VecDeque::new(), capacity: capacity.max(1), closed: false, changed: false, @@ -69,7 +71,7 @@ impl Queue { } pub(crate) fn has_room(&self) -> bool { - self.batches.len() < self.capacity + self.len() < self.capacity } pub(crate) fn closed(&self) -> bool { @@ -83,7 +85,7 @@ impl Queue { } pub(crate) fn len(&self) -> usize { - self.batches.len() + usize::from(self.front.is_some()) + self.batches.len() } #[cfg(debug_assertions)] @@ -96,7 +98,7 @@ impl Queue { } pub(crate) fn pop(&mut self) -> Option { - self.batches.pop_front() + std::mem::replace(&mut self.front, self.batches.pop_front()) } } @@ -108,12 +110,12 @@ pub struct Inlet<'a> { impl Inlet<'_> { /// The oldest batch, in place, so the source can slice it down without removing it. pub fn peek_mut(&mut self) -> Option<&mut ArrayRef> { - self.queue.batches.front_mut() + self.queue.front.as_mut() } /// The oldest batch, removed. Removing a batch is what frees its slot. pub fn take(&mut self) -> Option { - self.queue.batches.pop_front() + self.queue.pop() } /// Whether the sink has finished. Closed and empty is the end of the inlet. @@ -123,23 +125,28 @@ impl Inlet<'_> { /// Whether no batch is queued. pub fn is_empty(&self) -> bool { - self.queue.batches.is_empty() + self.queue.front.is_none() } /// Batches queued. pub fn len(&self) -> usize { - self.queue.batches.len() + self.queue.len() } /// Whether the sink can push no more until the source takes a batch. An unbounded inlet /// counts as full at the default capacity, so a source waiting for it to fill still runs. pub fn full(&self) -> bool { - self.queue.batches.len() >= self.queue.capacity.min(DEFAULT_CAPACITY) + self.queue.len() >= self.queue.capacity.min(DEFAULT_CAPACITY) } /// Rows queued, over every batch. pub fn rows(&self) -> usize { - self.queue.batches.iter().map(|batch| batch.len()).sum() + self.queue + .front + .iter() + .chain(self.queue.batches.iter()) + .map(|batch| batch.len()) + .sum() } } @@ -214,10 +221,9 @@ impl Arena { ) -> PortId { match self.free.pop() { Some(id) => { - // A freed slot keeps its buffer, so reusing it allocates only for a larger port. + // A freed slot keeps any overflow buffer for reuse. let queue = &mut self.queues[id]; let capacity = capacity.max(1); - queue.batches.reserve_exact(capacity.min(DEFAULT_CAPACITY)); queue.capacity = capacity; queue.closed = false; queue.changed = false; @@ -257,7 +263,11 @@ impl Arena { "push into port {id} with no consumer" ); if !batch.is_empty() { - queue.batches.push_back(batch); + if queue.front.is_none() { + queue.front = Some(batch); + } else { + queue.batches.push_back(batch); + } queue.changed = true; } } @@ -278,6 +288,7 @@ impl Arena { let queue = &mut self.queues[id]; queue.source_gone = true; queue.source = SourceEnd::Unclaimed; + queue.front = None; queue.batches.clear(); self.maybe_free(id); } @@ -292,6 +303,7 @@ impl Arena { ); queue.closed = true; queue.source_gone = true; + queue.front = None; queue.batches.clear(); self.maybe_free(id); } @@ -300,6 +312,7 @@ impl Arena { let queue = &mut self.queues[id]; if queue.closed && queue.source_gone && !queue.free { queue.free = true; + queue.front = None; queue.batches.clear(); self.free.push(id); } @@ -317,3 +330,43 @@ impl Arena { self.queues.len() - self.free.len() } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_buffer::buffer; + + use super::Arena; + use super::SourceEnd; + + #[rstest] + #[case::single(1)] + #[case::bounded(4)] + #[case::unbounded(usize::MAX)] + fn queued_batches_keep_order_and_reused_ports_are_empty(#[case] capacity: usize) { + let mut arena = Arena::default(); + let port = arena.create(capacity, None, SourceEnd::Scan(0)); + let arrays: Vec<_> = (0..capacity.min(4)) + .map(|i| buffer![i as u64].into_array()) + .collect(); + for array in &arrays { + assert!(arena.has_room(port)); + arena.push(port, array.clone()); + } + assert_eq!(arena.inlet(port).len(), arrays.len()); + assert_eq!(arena.inlet(port).rows(), arrays.len()); + arena.close(port); + for expected in arrays { + let actual = arena.inlet(port).take().expect("queued batch"); + assert!(ArrayRef::ptr_eq(&actual, &expected)); + } + assert!(arena.inlet(port).is_empty()); + arena.drop_source(port); + let reused = arena.create(1, None, SourceEnd::Scan(0)); + assert_eq!(reused, port); + assert!(arena.inlet(reused).is_empty()); + assert!(arena.has_room(reused)); + } +} diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index fc444340e83..9d650ebbd26 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -626,7 +626,7 @@ struct Active { /// reads of later splits are in flight while earlier ones compute. pub struct Scan { core: Core, - root: Root, + root: Option, splits: Vec, next: usize, active: Vec>, @@ -685,7 +685,7 @@ impl Scan { core.dirty_flags = vec![false; 1]; Ok(Self { core, - root, + root: Some(root), splits, next: 0, active: vec![None], @@ -766,12 +766,21 @@ impl Scan { let index = self.next; self.next += 1; let split = self.splits[index].clone(); - let output = match &self.root { + let output = match self + .root + .as_ref() + .ok_or_else(|| vortex_err!("Scan has no plan for its next split"))? + { Root::Plan(plan) => { self.core .compile_stage(plan, split.rows.clone(), &split.mask, (slot, index))? } }; + if self.next == self.splits.len() { + // Compiled operators own what they need. The last split no longer needs the + // lowering tree or the source layouts retained by its lazy children. + self.root = None; + } if output.is_none() { self.core.finish_split(index); } else { diff --git a/vortex-layout/src/plan/physical/pipeline/values.rs b/vortex-layout/src/plan/physical/pipeline/values.rs index b4dd7623336..18c2116b637 100644 --- a/vortex-layout/src/plan/physical/pipeline/values.rs +++ b/vortex-layout/src/plan/physical/pipeline/values.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Whole values, such as dictionary values, kept between the scans of one split. +//! Whole values, such as dictionary values, kept between the scans that share them. use std::sync::Arc; @@ -15,9 +15,9 @@ use crate::plan::physical::PlanRef; /// Values a plan produces over its whole domain whatever rows it is compiled for, kept for the /// later scans that share this cache. /// -/// A dictionary's values are read whole by every reader of its codes, so the filter and the -/// projection of one split decode them once. The cache lives as long as the split that owns it -/// and holds only the values that split read: nothing is kept for the rest of a scan. +/// A dictionary's values are read whole by every reader of its codes, so readers sharing this +/// cache can reuse them across filters, projections, and splits. The caller controls its lifetime; +/// only whole values that a reader requested are retained. #[derive(Clone, Default)] pub struct ValueCache(Arc>>); diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index 825897e33ee..25f12380493 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -235,10 +235,10 @@ impl PlanVTable for Pack { return Ok(None); } let mut chains = Vec::with_capacity(count); - for child in plan.children().iter() { + for child in plan.children().iter_refs() { chains.push( compiler - .compile(&child?, rows.clone(), mask)? + .compile(child?, rows.clone(), mask)? .ok_or_else(|| vortex_err!("A Pack field produced no rows"))?, ); } @@ -258,7 +258,7 @@ impl PlanVTable for Pack { at: &Reach, visit: &mut dyn FnMut(SegmentId, Range), ) -> VortexResult<()> { - for child in plan.children().iter() { + for child in plan.children().iter_refs() { child?.reach(rows.clone(), at, visit)?; } Ok(()) diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 224035a45fe..b128b92b452 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -64,7 +64,7 @@ impl ScanBuilder { } } -fn to_record_batch( +pub(super) fn to_record_batch( chunk: ArrayRef, data_type: &Field, ctx: &mut ExecutionCtx, diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs new file mode 100644 index 00000000000..fc6053d55ca --- /dev/null +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Adapts pipeline execution to the planner I/O and compute protocol. + +use std::ops::Range; +use std::sync::Arc; + +use rustc_hash::FxHashMap; +use vortex_array::ArrayRef; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_io::request; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoResult; +use vortex_scan::planning::planner::State; + +use crate::plan::physical::pipeline::ReadId; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::Turn; +use crate::scan::execution::SegmentLocation; +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. + pub(super) rows: Range, + /// The selected rows the piece holds. + pub(super) array: ArrayRef, +} + +/// What one [`ScanAdaptor::compute`] produced. +pub(super) enum AdaptorStep { + /// Control returns to the owner with CPU work still ready. + Yield, + /// Requests to publish, each returned once. + NeedsIO(IoBatch), + /// Rows the plan produced. + Piece(Piece), +} + +/// Adapts a pipeline [`Scan`] to the planner I/O and compute protocol. +/// +/// The scan names segments by id; the protocol names byte ranges. Each scan read becomes a +/// `Fetch` of the segment's range at the segment's alignment, and each delivery is handed back to +/// the scan under its own id. +pub(super) struct ScanAdaptor { + scan: Scan, + /// A turn already taken while batching reads, retained for the next compute. + pending: Option, + done: bool, + /// Orders the pieces handed out. + next_piece: u64, + locations: Arc<[SegmentLocation]>, + /// Published and undelivered requests, with the scan read and segment each one answers. + outstanding: FxHashMap, + /// A delivery the scan could not take, reported by the next compute. + failed: Option, + /// The protocol id the next published request gets. + next_id: u32, +} + +impl ScanAdaptor { + /// Wraps a scan, numbering requests uniquely within this adaptor's lifetime. + pub(super) fn new(scan: Scan, locations: Arc<[SegmentLocation]>) -> Self { + Self { + scan, + pending: None, + done: false, + next_piece: 0, + locations, + outstanding: FxHashMap::default(), + failed: None, + next_id: 0, + } + } + + /// The owner's protocol state while the scan is running. + pub(super) fn state(&self) -> State { + match &self.pending { + _ if self.failed.is_some() => State::NeedsCompute, + _ if self.done => State::Done, + Some(Turn::Waiting) => State::Waiting, + _ => State::NeedsCompute, + } + } + + pub(super) fn compute(&mut self) -> VortexResult { + if let Some(error) = self.failed.take() { + return Err(error); + } + Ok(match self.step()? { + Turn::Waiting | Turn::Done => AdaptorStep::Yield, + Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Read(read) => { + // Publish every read the scan has ready together, so the owner can coalesce them. + let mut batch = vec![self.request(read.id, read.segment_id)?]; + loop { + match self.step()? { + Turn::Read(read) => batch.push(self.request(read.id, read.segment_id)?), + turn => { + if !matches!(turn, Turn::Waiting) { + self.pending = Some(turn); + self.done = false; + } + break; + } + } + } + AdaptorStep::NeedsIO(batch) + } + }) + } + + fn step(&mut self) -> VortexResult { + let turn = match self.pending.take() { + Some(Turn::Waiting) | None => self.scan.step()?, + Some(turn) => turn, + }; + match &turn { + Turn::Done => self.done = true, + Turn::Waiting => self.pending = Some(Turn::Waiting), + Turn::Read(_) | Turn::Output(..) => {} + } + Ok(turn) + } + + fn piece(&mut self, array: ArrayRef) -> Piece { + let start = self.next_piece; + self.next_piece += 1; + Piece { + rows: start..start + 1, + array, + } + } + + fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { + let location = self.location(segment_id)?; + let id = self.next_id; + self.next_id = self + .next_id + .checked_add(1) + .ok_or_else(|| vortex_err!("ScanAdaptor ran out of request ids"))?; + self.outstanding.insert(id, (read, segment_id)); + Ok(request::IoRequest { + intent: IoIntent::Fetch, + request: request::IoRequestId(id), + target: location.target(), + }) + } + + /// Hands a delivered read to the scan. The driver only delivers requests this adaptor + /// published, so anything else is a driver bug. Bytes the scan cannot take fail the owner's + /// next compute. + pub(super) fn set_io_result(&mut self, id: request::IoRequestId, result: IoResult) { + let Some((read, segment_id)) = self.outstanding.remove(&id.0) else { + vortex_panic!("ScanAdaptor: delivery of {id:?}, which is not outstanding"); + }; + let IoResult::Bytes(bytes) = result else { + vortex_panic!("ScanAdaptor: segment {segment_id} answered with a size"); + }; + if let Err(error) = self.deliver(read, segment_id, bytes) + && self.failed.is_none() + { + self.failed = Some(error.with_context(format!("delivering segment {segment_id}"))); + } + } + + fn deliver( + &mut self, + read: ReadId, + segment_id: SegmentId, + bytes: BufferHandle, + ) -> VortexResult<()> { + // A source that honours the target's alignment makes this a check, not a copy. + let alignment = self.location(segment_id)?.alignment; + let bytes = BufferHandle::new_host(bytes.try_into_host_sync()?.aligned(alignment)); + self.scan.deliver(read, bytes)?; + if matches!(self.pending, Some(Turn::Waiting)) { + self.pending = None; + } + Ok(()) + } + + fn location(&self, id: SegmentId) -> VortexResult { + self.locations + .get(*id as usize) + .copied() + .ok_or_else(|| vortex_err!("segment {id} has no known location")) + } +} diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs new file mode 100644 index 00000000000..9d8b8897ef3 --- /dev/null +++ b/vortex-layout/src/scan/execution/mod.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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. + +mod adaptor; +mod projection; + +use std::sync::Arc; + +use vortex_buffer::Alignment; +use vortex_io::request::IoTarget; +use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::WorkScope; +use vortex_session::VortexSession; + +use self::projection::ProjectionPlanner; +use crate::plan::physical::PlanRef; +use crate::plan::physical::pipeline::ValueCache; + +/// Where a segment's bytes live in its source. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SegmentLocation { + /// Byte offset of the segment from the start of the source. + pub offset: u64, + /// Length of the segment in bytes. + pub length: u32, + /// Alignment the segment's bytes must have once delivered. + pub alignment: Alignment, +} + +impl SegmentLocation { + /// The byte range a read of the segment asks for. + fn target(&self) -> IoTarget { + IoTarget::Range { + offset: self.offset, + len: self.length as usize, + alignment: self.alignment, + } + } +} + +/// What every split of a scan shares: the projection plan and where its segments live. +#[derive(Clone)] +pub(super) struct ScanPlans { + /// The session used for decoding and expression evaluation. + pub(super) session: VortexSession, + /// The location of every segment, indexed by segment id. + pub(super) locations: Arc<[SegmentLocation]>, + /// The optimized projection plan, over the scan's root row domain. + pub(super) projection: PlanRef, + /// The global row index of the root row domain's first row. + pub(super) row_offset: u64, + /// 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, +} + +/// The rows of a split selected for projection. +struct SelectedRows { + /// The split the rows belong to. + 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 })) +} diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs new file mode 100644 index 00000000000..b6180725c27 --- /dev/null +++ b/vortex-layout/src/scan/execution/projection.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_scan::planning::morsel::Morsel; +use vortex_scan::planning::morsel::MorselOutput; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; + +use crate::plan::physical::pipeline::Scan; +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)), + )) + } +} + +/// 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 { + plans: Option, + dtype: DType, + selected: SelectedRows, + adaptor: Option, + pieces: Vec, + done: bool, +} + +impl ProjectionMorsel { + /// Creates a morsel projecting `selected`. + fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + Self { + dtype: plans.projection.dtype().clone(), + plans: Some(plans), + selected, + adaptor: None, + pieces: Vec::new(), + done: false, + } + } + + fn finish(&mut self) -> VortexResult { + self.done = true; + self.pieces.sort_by_key(|piece| piece.rows.start); + let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) + .into_iter() + .map(|piece| piece.array) + .filter(|array| !array.is_empty()) + .collect(); + let array = match arrays.len() { + 0 => return Ok(MorselOutput::Done), + 1 => arrays.remove(0), + _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), + }; + Ok(MorselOutput::Batch(array)) + } +} + +impl IoConsumer for ProjectionMorsel { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some(adaptor) = self.adaptor.as_mut() { + adaptor.set_io_result(request, result); + } + } +} + +impl Morsel for ProjectionMorsel { + fn state(&self) -> State { + match &self.adaptor { + _ if self.done => State::Done, + None => State::NeedsCompute, + Some(adaptor) => match adaptor.state() { + State::Done => State::NeedsCompute, + state => state, + }, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + vortex_bail!("ProjectionMorsel: compute called after Done"); + } + 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); + self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); + return Ok(MorselOutput::Continue); + }; + if adaptor.state() == State::Done { + return self.finish(); + } + Ok(match adaptor.compute()? { + AdaptorStep::Yield => MorselOutput::Continue, + AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + self.pieces.push(piece); + MorselOutput::Continue + } + }) + } +} diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index e08e4cb3f31..2d3d3b70031 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub mod arrow; +mod execution; mod filter; pub mod layout; mod limit; @@ -13,6 +14,7 @@ mod splits; mod tasks; #[cfg(test)] mod test; +pub mod v2; /// A heuristic for an ideal split size. /// diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs new file mode 100644 index 00000000000..ce565846d7e --- /dev/null +++ b/vortex-layout/src/scan/v2/io.rs @@ -0,0 +1,389 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::VecDeque; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use futures::FutureExt; +use futures::StreamExt; +use futures::TryFutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use futures::stream::FuturesUnordered; +use parking_lot::Mutex; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoResult; +use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::scan::execution::SegmentLocation; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// Byte ranges of a file's segments, keyed by `(offset, length)`, back to their segment ids. +pub(super) type SegmentRanges = Arc>; + +pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { + let ranges = locations.iter().enumerate().map(|(id, location)| { + let id = SegmentId::from(u32::try_from(id).unwrap_or(u32::MAX)); + ((location.offset, location.length as usize), id) + }); + Arc::new(ranges.collect()) +} + +/// Serves splits from any [`SegmentSource`], by segment id. +pub(super) struct SegmentScanIo { + segments: Arc, + ranges: SegmentRanges, +} + +impl SegmentScanIo { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self { segments, ranges } + } +} + +impl IoService for SegmentScanIo { + fn session(&self) -> Arc { + Arc::new(SegmentIoSource::new( + Arc::clone(&self.segments), + Arc::clone(&self.ranges), + )) + } +} + +/// Serves a split's range requests from a segment source, without blocking. +/// +/// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once, +/// [`poll`](IoSource::poll) takes a finished one if there is any, and the split's future awaits +/// the next with [`poll_completion`](IoSource::poll_completion) when the driver has nothing else +/// to run. [`wait`](IoSource::wait) is never called. +/// +/// An announcement requests the segment and holds the request, never polled, until a prefetch or +/// fetch consumes it, so a source that coalesces requests can fold it into a nearby read. +pub(super) struct SegmentIoSource { + segments: Arc, + ranges: SegmentRanges, + reads: Mutex>>, + ready: Mutex>, + /// Reads started by prefetches, consumed by the next fetch of their segment. + prefetched: Mutex>, + /// Announced segments, requested and never polled. + announced: Mutex>>, +} + +/// A read started ahead of the fetch that needs it. +type Prefetched = Shared>>>; + +impl SegmentIoSource { + pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self { + segments, + ranges, + reads: Mutex::new(FuturesUnordered::new()), + ready: Mutex::default(), + prefetched: Mutex::default(), + announced: Mutex::default(), + } + } +} + +impl IoSource for SegmentIoSource { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + let reads = self.reads.lock(); + let mut ready = self.ready.lock(); + let mut prefetched = self.prefetched.lock(); + for request in batch { + let IoTarget::Range { offset, len, .. } = request.target else { + vortex_bail!( + "SegmentIoSource only serves byte ranges, not {:?}", + request.target + ); + }; + let segment = *self + .ranges + .get(&(offset, len)) + .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; + if request.intent == IoIntent::Announce { + self.announced + .lock() + .entry(segment) + .or_default() + .push_back(self.segments.request(segment)); + continue; + } + if request.intent == IoIntent::Prefetch { + prefetched.entry(segment).or_insert_with(|| { + let read = self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared(); + // Polling registers the read as wanted, so the source starts it now. + drop( + read.clone() + .poll_unpin(&mut Context::from_waker(Waker::noop())), + ); + read + }); + continue; + } + let mut bytes = match prefetched.remove(&segment) { + Some(read) => read + .map_err(move |err| { + vortex_err!("prefetched read of segment {segment} failed: {err}") + }) + .boxed(), + None => self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)), + }; + let id = request.request; + // In-memory and cached segments need no allocated asynchronous task. A pending + // read is polled again by FuturesUnordered before the driver can go to sleep. + if let Poll::Ready(result) = bytes.poll_unpin(&mut Context::from_waker(Waker::noop())) { + ready.push_back(Completion { + owner, + request: id, + result: result.map(IoResult::Bytes), + }); + continue; + } + reads.push( + async move { + Completion { + owner, + request: id, + result: bytes.await.map(IoResult::Bytes), + } + } + .boxed(), + ); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + if let Some(completion) = self.ready.lock().pop_front() { + return Ok(Some(completion)); + } + // The split's future polls again with its own waker before it waits, so a read that + // finishes after this check is not missed. + let mut cx = Context::from_waker(Waker::noop()); + match self.reads.lock().poll_next_unpin(&mut cx) { + Poll::Ready(completion) => Ok(completion), + Poll::Pending => Ok(None), + } + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + if let Some(completion) = self.ready.lock().pop_front() { + return Poll::Ready(Ok(completion)); + } + match self.reads.lock().poll_next_unpin(cx) { + Poll::Ready(Some(completion)) => Poll::Ready(Ok(completion)), + Poll::Ready(None) => Poll::Ready(Err(vortex_err!( + "The split's driver is waiting with no reads in flight" + ))), + Poll::Pending => Poll::Pending, + } + } + + fn wait(&self) -> VortexResult { + vortex_bail!("SegmentIoSource is awaited through poll_completion, not waited on") + } + + fn release(&self, _owner: IoOwnerId) {} + + fn clear(&self) { + self.reads.lock().clear(); + self.ready.lock().clear(); + self.prefetched.lock().clear(); + self.announced.lock().clear(); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use futures::FutureExt; + use futures::channel::oneshot; + use futures::future; + use parking_lot::Mutex; + use vortex_array::buffer::BufferHandle; + use vortex_buffer::Alignment; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_io::request::IoIntent; + use vortex_io::request::IoOwnerId; + use vortex_io::request::IoRequest; + use vortex_io::request::IoRequestId; + use vortex_io::request::IoSource; + use vortex_io::request::IoTarget; + + use super::SegmentIoSource; + use super::segment_ranges; + use crate::scan::execution::SegmentLocation; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; + use crate::segments::SegmentSource; + + /// Serves four bytes per segment and counts the reads it is asked for. + #[derive(Default)] + struct CountingSegments(AtomicUsize); + + impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.0.fetch_add(1, Ordering::Relaxed); + let bytes = ByteBuffer::from(vec![u8::try_from(*id).unwrap_or(u8::MAX); 4]); + future::ready(Ok(BufferHandle::new_host(bytes))).boxed() + } + } + + fn request(intent: IoIntent, id: u32) -> IoRequest { + IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::range(0, 4), + } + } + + struct MixedSegments(Mutex>>); + + impl SegmentSource for MixedSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + if *id == 0 { + return future::ready(Ok(BufferHandle::new_host(ByteBuffer::from(vec![0; 4])))) + .boxed(); + } + let receiver = self.0.lock().take().expect("one pending read"); + async move { + receiver + .await + .map(BufferHandle::new_host) + .map_err(|err| vortex_err!("{err}")) + } + .boxed() + } + } + + #[tokio::test] + async fn pending_read_wakes_after_an_immediate_completion() -> VortexResult<()> { + let (sender, receiver) = oneshot::channel(); + let segments = Arc::new(MixedSegments(Mutex::new(Some(receiver)))); + let locations = [0, 4].map(|offset| SegmentLocation { + offset, + length: 4, + alignment: Alignment::none(), + }); + let io = SegmentIoSource::new(segments, segment_ranges(&locations)); + let mut pending = request(IoIntent::Fetch, 1); + pending.target = IoTarget::range(4, 4); + io.submit(IoOwnerId(0), vec![pending, request(IoIntent::Fetch, 0)])?; + let immediate = io.poll()?.expect("the in-memory read is already ready"); + assert_eq!(immediate.request, IoRequestId(0)); + assert!(immediate.result.is_ok()); + assert!(io.poll()?.is_none()); + + let delivery = tokio::spawn(async move { + tokio::task::yield_now().await; + sender + .send(ByteBuffer::from(vec![1; 4])) + .expect("read is live"); + }); + let completed = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completed.request, IoRequestId(1)); + assert!(completed.result.is_ok()); + delivery.await.expect("delivery task succeeds"); + assert!(io.poll()?.is_none()); + Ok(()) + } + + /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, + /// and only the fetch is delivered. + #[tokio::test] + async fn fetch_reuses_a_prefetched_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Prefetch, 0)])?; + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 1)])?; + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(1)); + assert!(completion.result.is_ok()); + assert!(io.poll()?.is_none()); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2)])?; + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(2)); + assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[tokio::test] + async fn announcements_keep_independent_consumers() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let io = SegmentIoSource::new(Arc::clone(&segments) as _, segment_ranges(&[location])); + io.submit( + IoOwnerId(0), + vec![ + request(IoIntent::Announce, 0), + request(IoIntent::Announce, 1), + ], + )?; + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + for id in 2..5 { + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); + assert_eq!( + segments.0.load(Ordering::Relaxed), + if id < 4 { 2 } else { 3 } + ); + } + Ok(()) + } +} diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs new file mode 100644 index 00000000000..46cbde987fc --- /dev/null +++ b/vortex-layout/src/scan/v2/mod.rs @@ -0,0 +1,387 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The planner-driven scan executor. +//! +//! A V2 scan lowers the file's layout once into a physical plan and runs every split of it +//! through the planning protocol: the [`Run`](vortex_scan::planning::driver::Run) driver advances +//! the split's planners and morsels, and every segment read is a protocol request served from the +//! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. +//! +//! 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. + +mod io; +mod split; +#[cfg(test)] +mod tests; + +use std::ops::Range; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::ArrowError; +use arrow_schema::Field; +use arrow_schema::SchemaRef; +use futures::Stream; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use futures::stream::BoxStream; +use vortex_array::ArrayRef; +use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_io::request::IoService; +use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_session::VortexSession; +use vortex_utils::parallelism::get_available_parallelism; + +use crate::LayoutRef; +use crate::layouts::chunked::Chunked; +use crate::layouts::dict::Dict; +use crate::layouts::flat::Flat; +use crate::layouts::struct_::Struct; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; +use crate::plan::physical::lower; +use crate::plan::physical::optimize; +use crate::plan::physical::plan_row_idx_expression; +use crate::scan::arrow::to_record_batch; +use crate::scan::execution::ScanPlans; +pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::plan_projection; +use crate::scan::v2::io::SegmentScanIo; +use crate::scan::v2::io::segment_ranges; +use crate::scan::v2::split::SplitPlan; +use crate::segments::SegmentSource; + +const ROW_BLOCK_SIZE: usize = 8192; +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 { + if row_count <= MAX_ROWS_PER_CHUNK as u64 { + return MAX_ROWS_PER_CHUNK; + } + let bytes: u64 = locations + .iter() + .map(|location| u64::from(location.length)) + .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 +} + +/// Fixed boundaries in the file's row domain, clipped to the requested range. +fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { + if rows.is_empty() { + return Vec::new(); + } + let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; + (first..rows.end) + .step_by(rows_per_chunk) + .map(|start| { + start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) + }) + .collect() +} + +/// The parts of a file the planner-driven executor reads through. +#[derive(Clone)] +pub struct ScanFile { + /// The root layout to scan. + pub layout: LayoutRef, + /// The location of every segment, indexed by segment id. + pub locations: Arc<[SegmentLocation]>, + /// Fetches segment bytes by id. + pub segments: Arc, +} + +/// 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. +pub struct ScanBuilder { + session: VortexSession, + file: ScanFile, + projection: BoundExpression, + ordered: bool, + row_range: Option>, + selection: Selection, + concurrency: usize, + map_fn: Arc VortexResult + Send + Sync>, + row_offset: u64, +} + +impl ScanBuilder { + /// Creates a V2 scan over the file's layout and segment source. + pub fn new(session: VortexSession, file: ScanFile) -> Self { + Self { + projection: BoundExpression::new_root(file.layout.dtype().clone()), + session, + file, + ordered: true, + row_range: None, + selection: Selection::default(), + concurrency: 4, + map_fn: Arc::new(Ok), + row_offset: 0, + } + } + + /// Converts each output batch to Arrow on the split's task. + pub fn into_record_batch_stream( + self, + schema: SchemaRef, + ) -> VortexResult> + Send + 'static> { + let struct_field = Field::new_struct("", schema.fields().clone(), false); + let session = self.session.clone(); + Ok(self + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + to_record_batch(chunk, &struct_field, &mut ctx) + }) + .into_stream()? + .map_err(|error| ArrowError::ExternalError(Box::new(error)))) + } +} + +impl ScanBuilder { + /// The output dtype after projection, before any output mapping. + pub fn dtype(&self) -> &DType { + self.projection.dtype() + } + + /// Sets a projection bound against the file's dtype. + pub fn with_projection(mut self, projection: BoundExpression) -> Self { + self.projection = projection; + self + } + + /// Configures whether batches are yielded in file order. + pub fn with_ordered(mut self, ordered: bool) -> Self { + self.ordered = ordered; + self + } + + /// Restricts scanning to a contiguous file-local row range. + pub fn with_row_range(mut self, row_range: Range) -> Self { + self.row_range = Some(row_range); + self + } + + /// Applies a row selection within the scan's row range. + pub fn with_selection(mut self, selection: Selection) -> Self { + self.selection = selection; + self + } + + /// Selects rows by strictly sorted file-local indices. + pub fn with_row_indices(self, row_indices: StrictSortedBuffer) -> Self { + self.with_selection(Selection::IncludeByIndex(row_indices)) + } + + /// Sets the root row offset used by row-index expressions. + pub fn with_row_offset(mut self, row_offset: u64) -> Self { + self.row_offset = row_offset; + self + } + + /// Sets the number of active splits per worker thread. Must be greater than zero. + pub fn with_concurrency(mut self, concurrency: usize) -> Self { + assert!(concurrency > 0); + self.concurrency = concurrency; + self + } + + /// Maps each output batch on its split's task. + fn map( + self, + map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, + ) -> ScanBuilder { + let old_map_fn = self.map_fn; + ScanBuilder { + session: self.session, + file: self.file, + projection: self.projection, + ordered: self.ordered, + row_range: self.row_range, + selection: self.selection, + concurrency: self.concurrency, + row_offset: self.row_offset, + map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)), + } + } + + /// Runs the scan on the session's runtime, preparing it on the first poll. + /// Preparation errors are returned as stream items. + pub fn into_stream(self) -> VortexResult>> { + let handle = self.session.handle(); + Ok(stream::once(async move { + handle + .spawn_cpu(move || PreparedScan::try_new(self)?.execute_stream()) + .await + }) + .try_flatten() + .boxed()) + } +} + +/// Whether every required layout lowers to plans with a pipeline implementation. +/// +/// List layouts lower to a list assembly, which cannot compile before its offsets are read. +fn supports_layout(layout: &LayoutRef, identity: bool) -> VortexResult { + if !(layout.is::() + || layout.is::() + || layout.is::() + || layout.is::() + || layout.is::() + || layout.is::()) + { + return Ok(false); + } + // An identity projection cannot use zone statistics. Leave their layouts uninitialized. + if identity && (layout.is::() || layout.is::()) { + return match layout.slot(0)? { + Some(data) => supports_layout(&data, identity), + None => Ok(false), + }; + } + for child in layout.children()? { + if !supports_layout(&child, identity)? { + return Ok(false); + } + } + Ok(true) +} + +/// A scan planned over its file, ready to run its splits. +struct PreparedScan { + session: VortexSession, + plans: ScanPlans, + /// Serves the splits' reads. + io: Arc, + /// The scan's splits, over the root row domain, in row order. + splits: Vec>, + selection: Selection, + map_fn: Arc VortexResult + Send + Sync>, + ordered: bool, + concurrency: usize, +} + +impl PreparedScan { + /// Lowers the file's layout, plans the projection, and partitions the selected range into + /// 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 + // statistics that a projection-only scan never reads. + let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); + let projection = if identity { + root + } else { + optimize(plan_row_idx_expression(builder.projection, root)?)? + }; + let chunk_rows = if identity { + rows_per_chunk(row_count, &file.locations) + } else { + MAX_ROWS_PER_CHUNK + }; + let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); + let io = Arc::new(SegmentScanIo::new( + file.segments, + segment_ranges(&file.locations), + )); + Ok(Self { + session: builder.session.clone(), + plans: ScanPlans { + session: builder.session, + locations: file.locations, + projection, + row_offset: builder.row_offset, + values: Default::default(), + }, + io, + splits, + selection: builder.selection, + map_fn: builder.map_fn, + ordered: builder.ordered, + concurrency: builder.concurrency, + }) + } + + /// The splits with selected rows, each ready to be admitted to a planning run. + fn split_plans(&self) -> Vec { + self.splits + .iter() + .filter_map(|rows| { + let row_mask = self.selection.row_mask(rows); + if row_mask.mask().all_false() { + return None; + } + let scope = WorkScope { + file_ordinal: 0, + rows: row_mask.row_range(), + }; + Some(SplitPlan { + root: plan_projection( + self.plans.clone(), + scope.clone(), + row_mask.mask().clone(), + ), + scope, + }) + }) + .collect() + } + + /// Runs every split as a task on the session's runtime and yields their mapped batches. + /// + /// Like the default executor, up to `concurrency` splits per thread run at once, and their + /// batches are yielded in row order when the scan is ordered. + fn execute_stream(self) -> VortexResult>> { + let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); + let handle = self.session.handle(); + let split_plans = self.split_plans(); + let io = self.io; + let map_fn = self.map_fn; + // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. + let tasks = stream::iter(split_plans).map(move |split| { + let io = Arc::clone(&io); + let map_fn = Arc::clone(&map_fn); + handle.spawn(async move { + split + .run(io) + .await? + .into_iter() + .map(|array| map_fn(array)) + .collect::>>() + }) + }); + let batches = if self.ordered { + tasks.buffered(concurrency).boxed() + } else { + tasks.buffer_unordered(concurrency).boxed() + }; + Ok(batches + .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) + .try_flatten() + .boxed()) + } +} diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs new file mode 100644 index 00000000000..b9371735d39 --- /dev/null +++ b/vortex-layout/src/scan/v2/split.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::future; +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_io::request::IoService; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::WorkScope; + +/// One split of a scan, ready to be admitted to a planning run. +/// +/// The split's planner projects the rows of `scope`, and its morsels emit one batch per +/// projection split with selected rows, in whatever order their reads arrive. +pub(super) struct SplitPlan { + /// The rows of the scan's root row domain the split covers. + pub(super) scope: WorkScope, + /// The split's first planner. + pub(super) root: Box, +} + +impl SplitPlan { + /// Runs the split on a planning run of its own, reading through a session of `io`, and + /// returns its batches in row order. + /// + /// The run advances inside this future, on whichever thread polls it, and the future awaits + /// its reads between steps. Nothing is handed to another thread. + pub(super) async fn run(self, io: Arc) -> VortexResult> { + let mut run = Run::new(); + run.admit(self.root, self.scope, io.session()); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + // Morsels finish in whatever order their reads arrive. + batches.sort_by_key(|batch| batch.scope.rows.start); + Ok(batches.into_iter().map(|batch| batch.array).collect()) + } +} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs new file mode 100644 index 00000000000..a910172f87b --- /dev/null +++ b/vortex-layout/src/scan/v2/tests.rs @@ -0,0 +1,483 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::num::NonZeroUsize; +use std::ops::Range; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use arrow_array::cast::AsArray; +use arrow_array::types::Int32Type; +use futures::TryStreamExt; +use futures::stream; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::expr::pack; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_arrow::ArrowSessionExt; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_io::session::RuntimeSessionExt; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex_session::VortexSession; + +use super::MAX_ROWS_PER_CHUNK; +use super::PreparedScan; +use super::ScanBuilder; +use super::ScanFile; +use super::row_chunks; +use super::rows_per_chunk; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; +use crate::layouts::dict::DictLayout; +use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::row_idx::row_idx; +use crate::layouts::table::TableStrategy; +use crate::layouts::zoned::writer::ZonedLayoutOptions; +use crate::layouts::zoned::writer::ZonedStrategy; +use crate::scan::execution::SegmentLocation; +use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt as _; +use crate::test::new_session; + +const CHUNK_ROWS: usize = 1000; +const CHUNKS: usize = 4; + +/// 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 { + let rows = start..start + CHUNK_ROWS; + let ints = Buffer::from_iter( + rows.clone() + .map(|row| i32::try_from(row).unwrap_or(i32::MAX)), + ); + let strings = VarBinArray::from_iter( + rows.map(|row| (row % 3 != 0).then(|| format!("row {row}"))), + DType::Utf8(true.into()), + ); + let elements = Buffer::from_iter((0..2 * CHUNK_ROWS).map(|i| i64::try_from(i).unwrap_or(0))); + let offsets = Buffer::from_iter((0..=CHUNK_ROWS).map(|i| u32::try_from(2 * i).unwrap_or(0))); + let lists = ListArray::try_new( + elements.into_array(), + offsets.into_array(), + Validity::NonNullable, + )?; + Ok(StructArray::from_fields(&[ + ("ints", ints.into_array()), + ("strings", strings.into_array()), + ("lists", lists.into_array()), + ])? + .into_array()) +} + +/// Writes `CHUNKS` table chunks with `strategy`. +async fn write( + session: &VortexSession, + strategy: &dyn LayoutStrategy, +) -> VortexResult<(Arc, LayoutRef)> { + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + let chunks = (0..CHUNKS) + .map(|chunk| Ok((sequence_id.advance(), table_chunk(chunk * CHUNK_ROWS)?))) + .collect::>(); + let dtype = table_chunk(0)?.dtype().clone(); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(chunks)).sendable(), + eof, + session, + ) + .await?; + Ok((segments, layout)) +} + +fn chunked_flat() -> Arc { + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())) +} + +/// Every column is shredded into its own chunked leaves. +fn table() -> Arc { + Arc::new(TableStrategy::new(chunked_flat(), chunked_flat())) +} + +/// Every column is shredded into chunked leaves with zone statistics. +fn zoned_table() -> Arc { + let zoned = ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(CHUNK_ROWS).unwrap_or(NonZeroUsize::MIN), + ..Default::default() + }, + ); + Arc::new(TableStrategy::new(chunked_flat(), Arc::new(zoned))) +} + +/// A [`ScanFile`] over test segments. They have no byte offsets, and the scan fetches them by id, +/// so each location is a placeholder whose offset is its id. +fn scan_file(segments: &Arc, layout: &LayoutRef) -> VortexResult { + let mut count = 0; + for layout in layout.depth_first_traversal() { + for id in layout?.segment_ids() { + count = count.max(*id as usize + 1); + } + } + let locations = (0..count as u64) + .map(|offset| SegmentLocation { + offset, + length: 0, + alignment: Alignment::none(), + }) + .collect(); + Ok(ScanFile { + layout: Arc::clone(layout), + locations, + segments: Arc::clone(segments), + }) +} + +fn builder( + session: &VortexSession, + segments: &Arc, + layout: &LayoutRef, +) -> VortexResult> { + Ok(ScanBuilder::new( + session.clone(), + scan_file(segments, layout)?, + )) +} + +async fn collect( + dtype: DType, + stream: impl futures::Stream>, +) -> VortexResult { + let chunks: Vec = stream.try_collect().await?; + Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) +} + +/// Runs `builder` on both executors and checks the planner-driven one returns the same rows. +async fn assert_same_as_v1( + session: &VortexSession, + make: impl Fn() -> VortexResult>, +) -> VortexResult<()> { + let builder = make()?; + let dtype = builder.dtype().clone(); + let reader = builder.file.layout.new_reader( + "".into(), + Arc::clone(&builder.file.segments), + session, + &Default::default(), + )?; + let mut v1 = V1ScanBuilder::new(session.clone(), reader) + .with_projection(builder.projection.clone()) + .with_selection(builder.selection.clone()) + .with_row_offset(builder.row_offset) + .with_ordered(builder.ordered) + .with_concurrency(builder.concurrency); + if let Some(rows) = builder.row_range.clone() { + v1 = v1.with_row_range(rows); + } + let expected = collect(dtype.clone(), v1.into_stream()?).await?; + let actual = collect(dtype, builder.into_stream()?).await?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) +} + +#[derive(Clone, Copy, Debug)] +enum Strategy { + Flat, + Table, + ZonedTable, +} + +impl Strategy { + fn get(self) -> Arc { + match self { + Strategy::Flat => chunked_flat(), + Strategy::Table => table(), + Strategy::ZonedTable => zoned_table(), + } + } +} + +#[rstest] +#[case::whole(None, None)] +#[case::row_range(Some(250..3_100), None)] +#[case::empty_row_range(Some(10..10), None)] +#[case::row_indices(None, Some(&[0, 1, 999, 1000, 2500, 3999][..]))] +#[case::row_range_and_indices(Some(900..2600), Some(&[0, 999, 1000, 2500, 3999][..]))] +#[tokio::test] +async fn matches_v1( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, + #[case] row_range: Option>, + #[case] rows: Option<&'static [u64]>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + assert_same_as_v1(&session, || { + let mut builder = builder(&session, &segments, &layout)?; + if let Some(row_range) = row_range.clone() { + builder = builder.with_row_range(row_range); + } + if let Some(rows) = rows { + builder = + builder.with_row_indices(StrictSortedBuffer::try_new(Buffer::copy_from(rows))?); + } + Ok(builder) + }) + .await +} + +#[rstest] +#[tokio::test] +async fn projection_matches_v1( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + let dtype = layout.dtype().clone(); + assert_same_as_v1(&session, || { + let projection = pack( + [ + ("strings", get_item("strings", root())), + ("idx", row_idx()), + ("big", gt(get_item("ints", root()), lit(1500i32))), + ], + false.into(), + ); + Ok(builder(&session, &segments, &layout)? + .with_projection(projection.bind(&dtype)?) + .with_row_offset(100)) + }) + .await +} + +#[tokio::test] +async fn arrow_stream_projects_across_chunks() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let projection = + pack([("value", get_item("ints", root()))], false.into()).bind(layout.dtype())?; + let scan = builder(&session, &segments, &layout)? + .with_projection(projection) + .with_row_range(990..1010); + let schema = Arc::new(session.arrow().to_arrow_schema(scan.dtype())?); + let batches: Vec<_> = scan.into_record_batch_stream(schema)?.try_collect().await?; + let values: Vec<_> = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, (990_i32..1010).collect::>()); + Ok(()) +} + +#[rstest] +#[case::whole(0..150_000, vec![0..65_536, 65_536..131_072, 131_072..150_000])] +#[case::unaligned(10..140_000, vec![10..65_536, 65_536..131_072, 131_072..140_000])] +#[case::boundary(65_536..131_072, vec![65_536..131_072])] +#[case::empty(5..5, vec![])] +#[case::near_max(u64::MAX - 3..u64::MAX, vec![u64::MAX - 3..u64::MAX])] +fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) { + assert_eq!(row_chunks(rows, 65_536), expected); +} + +#[rstest] +#[case::empty(0, 0, 65_536)] +#[case::single_chunk(1000, 100_000_000, 65_536)] +#[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) { + let locations = [SegmentLocation { + offset: 0, + length: bytes, + alignment: Alignment::none(), + }]; + assert_eq!(rows_per_chunk(rows, &locations), expected); +} + +#[test] +fn smaller_chunks_keep_file_alignment() { + assert_eq!( + row_chunks(7500..33_000, 16_384), + [7500..16_384, 16_384..32_768, 32_768..33_000] + ); +} + +struct CountingSegments { + inner: Arc, + reads: AtomicUsize, +} + +impl SegmentSource for CountingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.reads.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } +} + +#[tokio::test] +async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { + let session = new_session().with_tokio(); + const STORED_CHUNK_ROWS: usize = 70_000; + const STORED_CHUNKS: usize = 3; + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let arrays = (0..STORED_CHUNKS) + .map(|index| { + let start = i32::try_from(index * STORED_CHUNK_ROWS)?; + let end = start + i32::try_from(STORED_CHUNK_ROWS)?; + Ok(( + sequence.advance(), + PrimitiveArray::from_iter(start..end).into_array(), + )) + }) + .collect::>>()?; + let dtype = DType::Primitive(vortex_array::dtype::PType::I32, false.into()); + let ints = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))) + .sendable(), + eof, + &session, + ) + .await?; + // Codes index the same values across fixed work chunks and stored chunk boundaries. + let layout = DictLayout::new(Arc::clone(&ints), ints).into_layout(); + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments = Arc::clone(&counted) as Arc; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; + let splits = scan.split_plans(); + assert_eq!( + splits.len(), + (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) + ); + for (index, split) in splits.into_iter().enumerate() { + let rows = split.scope.rows.clone(); + let before = counted.reads.load(Ordering::Relaxed); + let arrays = split.run(Arc::clone(&scan.io)).await?; + let reads = counted.reads.load(Ordering::Relaxed) - before; + let code_chunks = + (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; + assert_eq!( + reads, + if index == 0 { + STORED_CHUNKS + } else { + usize::try_from(code_chunks)? + } + ); + let expected = + PrimitiveArray::from_iter(i32::try_from(rows.start)?..i32::try_from(rows.end)?); + assert_arrays_eq!( + ChunkedArray::try_new(arrays, layout.dtype().clone())?.into_array(), + expected, + &mut session.create_execution_ctx() + ); + } + 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<()> { + 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")); + Ok(()) +} + +#[tokio::test] +async fn unordered_returns_every_row() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; + let chunks: Vec = scan.execute_stream()?.try_collect().await?; + let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); + assert_eq!(rows, CHUNKS * CHUNK_ROWS); + Ok(()) +} + +#[tokio::test] +async fn primitive_column() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let dtype = layout.dtype().clone(); + let scan = PreparedScan::try_new( + builder(&session, &segments, &layout)? + .with_projection(get_item("ints", root()).bind(&dtype)?) + .with_row_range(990..1010), + )?; + let actual = collect( + DType::Primitive(vortex_array::dtype::PType::I32, false.into()), + scan.execute_stream()?, + ) + .await?; + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter(990i32..1010), + &mut session.create_execution_ctx() + ); + Ok(()) +} + +/// Stored chunk boundaries do not subdivide a fixed work chunk. +#[rstest] +#[tokio::test] +async fn work_chunks_ignore_layout_boundaries( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; + assert_eq!(scan.splits, vec![0..(CHUNKS * CHUNK_ROWS) as u64]); + let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; + assert_eq!(arrays.len(), 1); + assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); + Ok(()) +} diff --git a/vortex-scan/Cargo.toml b/vortex-scan/Cargo.toml index d706ec20e22..d12c207b049 100644 --- a/vortex-scan/Cargo.toml +++ b/vortex-scan/Cargo.toml @@ -21,11 +21,17 @@ vortex-error = { workspace = true } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } +vortex-utils = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } roaring = { workspace = true } tracing = { workspace = true } +[dev-dependencies] +parking_lot = { workspace = true } +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } + [lints] workspace = true diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs new file mode 100644 index 00000000000..8cb78099272 --- /dev/null +++ b/vortex-scan/src/planning/driver/mod.rs @@ -0,0 +1,551 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A single-worker driver for the protocol. +//! +//! A [`Run`] drives any number of root planners, each with its own IO session, and hands their +//! 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 +//! 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 +//! a completion it took from a source to [`Run::complete`]. [`Driver::run`] is the blocking +//! convenience for one root. + +use std::collections::BTreeMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::Entry; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::planning::morsel::Morsel; +use crate::planning::morsel::MorselOutput; +use crate::planning::planner::Planner; +use crate::planning::planner::PlannerOutput; +use crate::planning::planner::State; +use crate::planning::planner::WorkScope; + +/// Identifies one root admitted to a [`Run`]. Unique within the run, in admission order. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct RootId(pub u64); + +/// One array produced by a morsel, with the scope the morsel was created for. +pub struct Batch { + /// The root the morsel descends from. + pub root: RootId, + /// Scope of the morsel that produced the array. + pub scope: WorkScope, + /// The non-empty array. + pub array: ArrayRef, +} + +/// What a call to [`Run::advance`] left behind. +pub enum Progress { + /// A morsel's batch. One morsel's batches arrive in the order it emitted them. + Batch(Batch), + /// Every item descended from the root has retired, and its IO session was cleared. + RootDone(RootId), + /// Every live item waits for IO. Await [`Run::poll_completion`], or pass the next completion + /// to [`Run::complete`]. + Waiting, + /// Nothing is live: every admitted root has finished. + Idle, +} + +enum Item { + Planner(Box), + Morsel(Box), +} + +/// What one item's `compute()` produced. +enum Output { + Planner(PlannerOutput), + Morsel(MorselOutput), +} + +/// One live planner or morsel. Boxed in the run queue and the parked map, so scheduling moves a +/// pointer. +struct Work { + id: IoOwnerId, + root: RootId, + scope: WorkScope, + item: Item, + /// Stable child-ordinal path, beginning with the root's admission ordinal. + path: Vec, + /// 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. + outstanding: HashMap, +} + +impl Work { + fn priority(&self) -> Vec { + let mut priority = self.path.clone(); + priority.push(self.next_child); + priority + } + + // 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 { + let path = self.priority(); + self.next_child += 1; + path + } + + fn state(&self) -> State { + match &self.item { + Item::Planner(planner) => planner.state(), + Item::Morsel(morsel) => morsel.state(), + } + } + + fn consumer(&mut self) -> &mut dyn IoConsumer { + match &mut self.item { + Item::Planner(planner) => &mut **planner, + Item::Morsel(morsel) => &mut **morsel, + } + } + + /// Checks `result` against the outstanding request it answers and hands it to the item. + fn deliver(&mut self, request: IoRequestId, result: IoResult) -> VortexResult<()> { + let Some(target) = self.outstanding.remove(&request) else { + vortex_bail!( + "completion for {request:?} of {:?}, which is not outstanding", + self.id + ); + }; + if !result.matches(&target) { + vortex_bail!("IO source answered {target:?} with {}", result.kind()); + } + self.consumer().set_io_result(request, result); + Ok(()) + } +} + +/// Keeps scheduling order separate from lookup for deliveries to runnable work. +#[derive(Default)] +struct ReadyQueue { + priorities: BTreeMap, IoOwnerId>, + work: HashMap>, +} + +impl ReadyQueue { + fn insert(&mut self, work: Box) { + self.priorities.insert(work.priority(), work.id); + self.work.insert(work.id, work); + } + + fn pop(&mut self) -> Option> { + let (_, owner) = self.priorities.pop_first()?; + self.work.remove(&owner) + } + + fn is_empty(&self) -> bool { + self.priorities.is_empty() + } + + fn cancel(&mut self, root: RootId) { + self.work.retain(|_, work| work.root != root); + self.priorities + .retain(|_, owner| self.work.contains_key(owner)); + } +} + +/// Poll IO even when one high-priority branch keeps producing CPU work. +const IO_POLL_INTERVAL: usize = 64; + +/// One admitted root: the session its descendants read through, and how many of them are live. +struct Root { + io: Arc, + live: usize, +} + +/// A set of root planners advanced by its caller on one thread at a time. +/// +/// Dropping a run clears the IO session of every root still live, even when it ends in an error. +#[derive(Default)] +pub struct Run { + roots: HashMap, + next_root: u64, + next_owner: u64, + queue: ReadyQueue, + parked: HashMap>, + /// Batches and finished roots not yet returned by [`advance`](Self::advance). + events: VecDeque, + /// Visits left before the sources are polled for completions again. + visits_until_poll: usize, + step_limit: Option, + steps: usize, +} + +impl Drop for Run { + fn drop(&mut self) { + for root in self.roots.values() { + root.io.clear(); + } + } +} + +impl Run { + /// Creates a run with nothing admitted. + pub fn new() -> Self { + Self::default() + } + + /// Fails with an error naming the limit if more than `steps` visits are needed. Tests use it + /// to catch loops. + pub fn with_step_limit(mut self, steps: usize) -> Self { + self.step_limit = Some(steps); + self + } + + /// Admits `root` and everything it spawns, reading through `io`. + /// + /// The session belongs to the root: the run submits the requests of every descendant to it, + /// and clears it once they have all retired or the root is cancelled. + pub fn admit( + &mut self, + root: Box, + scope: WorkScope, + io: Arc, + ) -> RootId { + 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]); + id + } + + /// Drops every live item descended from `root` and clears its session. Batches the root + /// already produced are still returned, followed by no [`Progress::RootDone`]. + pub fn cancel(&mut self, root: RootId) { + let Some(cancelled) = self.roots.remove(&root) else { + return; + }; + self.queue.cancel(root); + self.parked.retain(|_, work| work.root != root); + cancelled.io.clear(); + } + + /// The number of admitted roots that have not finished. + pub fn live_roots(&self) -> usize { + self.roots.len() + } + + /// Runs ready work, taking completions the sources already have, until there is something to + /// hand back or everything left waits for IO. Never blocks. + pub fn advance(&mut self) -> VortexResult { + loop { + if let Some(event) = self.events.pop_front() { + return Ok(event); + } + if self.visits_until_poll == 0 || self.queue.is_empty() { + self.take_completions()?; + self.visits_until_poll = IO_POLL_INTERVAL; + if self.queue.is_empty() { + return Ok(if self.parked.is_empty() { + Progress::Idle + } else { + Progress::Waiting + }); + } + } + self.visits_until_poll -= 1; + let Some(work) = self.queue.pop() else { + continue; + }; + self.steps += 1; + if let Some(limit) = self.step_limit + && self.steps > limit + { + vortex_bail!("driver exceeded its step limit of {limit}"); + } + self.visit(work)?; + } + } + + /// Delivers a completion taken from the session of one of the run's roots. + pub fn complete(&mut self, completion: Completion) -> VortexResult<()> { + let Completion { + owner, + request, + result, + } = completion; + let result = result?; + if let Entry::Occupied(mut parked) = self.parked.entry(owner) { + let work = parked.get_mut(); + work.deliver(request, result)?; + // An item that still waits for other requests stays parked. + if work.state() != State::Waiting { + let work = parked.remove(); + self.queue.insert(work); + } else if work.outstanding.is_empty() { + vortex_bail!( + "work for scope {:?} waits with no outstanding fetch", + work.scope + ); + } + return Ok(()); + } + // An item that published fetches and had CPU work left is still in the run queue. + let Some(work) = self.queue.work.get_mut(&owner) else { + vortex_bail!("completion for unknown work {owner:?}"); + }; + work.deliver(request, result) + } + + /// Waits for a completion from any root's session, registering `cx` to be woken by the next + /// one, and delivers every completion that is ready. + /// + /// Call it after [`advance`](Self::advance) returned [`Progress::Waiting`]. + pub fn poll_completion(&mut self, cx: &mut Context<'_>) -> Poll> { + let sources: Vec<_> = self + .roots + .values() + .map(|root| Arc::clone(&root.io)) + .collect(); + if sources.is_empty() { + return Poll::Ready(Err(vortex_err!( + "the run is waiting for IO with no live root" + ))); + } + let mut delivered = false; + for io in sources { + match io.poll_completion(cx) { + Poll::Ready(Ok(completion)) => { + if let Err(error) = self.complete(completion) { + return Poll::Ready(Err(error)); + } + delivered = true; + } + Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), + Poll::Pending => {} + } + } + if delivered { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + + /// Delivers the completions every root's session has ready, without blocking. + fn take_completions(&mut self) -> VortexResult<()> { + 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 { + return Ok(()); + }; + while let Some(completion) = io.poll()? { + self.complete(completion)?; + } + return Ok(()); + } + let sources: Vec<_> = self + .roots + .values() + .map(|root| Arc::clone(&root.io)) + .collect(); + for io in sources { + while let Some(completion) = io.poll()? { + self.complete(completion)?; + } + } + Ok(()) + } + + fn spawn(&mut self, root: RootId, scope: WorkScope, item: Item, path: Vec) { + let id = IoOwnerId(self.next_owner); + self.next_owner += 1; + if let Some(root) = self.roots.get_mut(&root) { + root.live += 1; + } + let work = Box::new(Work { + id, + root, + scope, + item, + path, + next_child: 0, + outstanding: HashMap::default(), + }); + self.queue.insert(work); + } + + /// Retires a finished item, and its root once nothing descended from it is live. + fn retire(&mut self, work: &Work) { + let Entry::Occupied(mut root) = self.roots.entry(work.root) else { + return; + }; + root.get().io.release(work.id); + root.get_mut().live -= 1; + if root.get().live == 0 { + root.remove().io.clear(); + self.events.push_back(Progress::RootDone(work.root)); + } + } + + /// Parks an item that waits for IO, which is a protocol error with no fetch outstanding. + fn park(&mut self, work: Box) -> VortexResult<()> { + if work.outstanding.is_empty() { + vortex_bail!( + "work for scope {:?} waits with no outstanding fetch", + work.scope + ); + } + self.parked.insert(work.id, work); + Ok(()) + } + + /// Requeues, parks, or retires an item according to the state it reports. + fn settle(&mut self, work: Box) -> VortexResult<()> { + match work.state() { + State::Done => self.retire(&work), + State::NeedsCompute => self.queue.insert(work), + State::Waiting => self.park(work)?, + } + Ok(()) + } + + /// Visits one item: computes once if it is ready, and settles it otherwise. + fn visit(&mut self, mut work: Box) -> VortexResult<()> { + let output = match &mut work.item { + Item::Planner(planner) => match planner.state() { + State::NeedsCompute => Output::Planner(planner.compute()?), + _ => return self.settle(work), + }, + Item::Morsel(morsel) => match morsel.state() { + State::NeedsCompute => Output::Morsel(morsel.compute()?), + _ => return self.settle(work), + }, + }; + match output { + Output::Planner(PlannerOutput::Done) | Output::Morsel(MorselOutput::Done) => { + self.retire(&work); + } + Output::Planner(PlannerOutput::Continue) | Output::Morsel(MorselOutput::Continue) => { + self.queue.insert(work); + } + Output::Planner(PlannerOutput::NeedsIO(batch)) + | Output::Morsel(MorselOutput::NeedsIO(batch)) => { + self.submit(&mut work, batch)?; + self.settle(work)?; + } + Output::Planner(PlannerOutput::Planner(scope, child)) => { + let path = work.child_path(); + self.spawn(work.root, scope, Item::Planner(child), path); + self.queue.insert(work); + } + Output::Planner(PlannerOutput::Morsel(scope, morsel)) => { + let path = work.child_path(); + self.spawn(work.root, scope, Item::Morsel(morsel), path); + self.queue.insert(work); + } + Output::Morsel(MorselOutput::Batch(array)) => { + if array.is_empty() { + vortex_bail!("morsel for scope {:?} returned an empty batch", work.scope); + } + self.events.push_back(Progress::Batch(Batch { + root: work.root, + scope: work.scope.clone(), + array, + })); + self.queue.insert(work); + } + } + Ok(()) + } + + /// Registers the requests an item published with its root's session. + fn submit(&mut self, work: &mut Work, batch: IoBatch) -> VortexResult<()> { + if batch.is_empty() { + vortex_bail!("compute() published an empty batch"); + } + for request in &batch { + if request.intent == IoIntent::Fetch + && work + .outstanding + .insert(request.request, request.target) + .is_some() + { + vortex_bail!( + "fetch {:?} of {:?} was published while it is outstanding", + request.request, + work.id + ); + } + } + let Some(root) = self.roots.get(&work.root) else { + vortex_bail!("work {:?} outlived its root {:?}", work.id, work.root); + }; + root.io.submit(work.id, batch) + } +} + +/// Drives one root to completion on the calling thread, blocking while it waits for IO. +pub struct Driver { + io: Arc, + step_limit: Option, +} + +impl Driver { + /// Creates a driver that submits every request to `io`. + pub fn new(io: Arc) -> Self { + Self { + io, + step_limit: None, + } + } + + /// Fails with an error naming the limit if more than `steps` visits are needed. Tests use it + /// to catch loops. + pub fn with_step_limit(self, steps: usize) -> Self { + Self { + step_limit: Some(steps), + ..self + } + } + + /// Runs `root` and everything it spawns, returning batches in emission order. + pub fn run(&self, root: Box) -> VortexResult> { + let mut run = Run::new(); + run.step_limit = self.step_limit; + let scope = WorkScope { + file_ordinal: 0, + rows: 0..0, + }; + run.admit(root, scope, Arc::clone(&self.io)); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::RootDone(_) | Progress::Idle => return Ok(batches), + Progress::Waiting => run.complete(self.io.wait()?)?, + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs new file mode 100644 index 00000000000..389ad24372d --- /dev/null +++ b/vortex-scan/src/planning/driver/tests.rs @@ -0,0 +1,813 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Waker; + +use rstest::rstest; +use vortex_array::assert_arrays_eq; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_io::request::IoRequest; + +use super::*; +use crate::planning::tests::scripted::Log; +use crate::planning::tests::scripted::MorselStep; +use crate::planning::tests::scripted::PlannerStep; +use crate::planning::tests::scripted::RecordingIoSource; +use crate::planning::tests::scripted::ScriptedPlanner; +use crate::planning::tests::scripted::array_of; +use crate::planning::tests::scripted::ctx; +use crate::planning::tests::scripted::scope; + +fn request(id: u32, offset: u64) -> IoRequest { + IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(id), + target: IoTarget::range(offset, 4), + } +} + +fn driver(source: &Arc) -> Driver { + Driver::new(Arc::clone(source) as Arc).with_step_limit(1_000) +} + +fn source() -> Arc { + let source = RecordingIoSource::default(); + source.canned_bytes(0, buffer![0u8, 1, 2, 3].into_byte_buffer()); + source.canned_bytes(4, buffer![4u8, 5, 6, 7].into_byte_buffer()); + Arc::new(source) +} + +fn lifo_source() -> Arc { + let source = RecordingIoSource::lifo(); + source.canned_bytes(0, buffer![0u8, 1, 2, 3].into_byte_buffer()); + source.canned_bytes(4, buffer![4u8, 5, 6, 7].into_byte_buffer()); + Arc::new(source) +} + +fn one_batch_morsel(scope: WorkScope, rows: u64) -> PlannerStep { + PlannerStep::Morsel( + scope, + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(rows))), + MorselStep::Compute(MorselOutput::Done), + ], + ) +} + +#[test] +fn root_done_produces_nothing() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed("root", vec![PlannerStep::Done], &log); + let batches = driver(&source()).run(root)?; + assert!(batches.is_empty()); + assert_eq!(log.events(), vec!["compute root"]); + Ok(()) +} + +#[test] +fn one_morsel_one_batch() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![one_batch_morsel(scope(10..20), 3), PlannerStep::Done], + &log, + ); + let batches = driver(&source()).run(root)?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].scope, scope(10..20)); + assert_arrays_eq!(batches[0].array, array_of(3), &mut ctx()); + Ok(()) +} + +#[test] +fn two_morsels_emit_in_order() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + one_batch_morsel(scope(0..1), 1), + one_batch_morsel(scope(1..3), 2), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source()).run(root)?; + let scopes: Vec<_> = batches.iter().map(|batch| batch.scope.clone()).collect(); + assert_eq!(scopes, vec![scope(0..1), scope(1..3)]); + assert_arrays_eq!(batches[0].array, array_of(1), &mut ctx()); + assert_arrays_eq!(batches[1].array, array_of(2), &mut ctx()); + Ok(()) +} + +#[test] +fn dfs_finishes_child_subtree_before_parent_continues() -> VortexResult<()> { + let log = Log::default(); + let child = PlannerStep::Planner( + scope(0..5), + vec![one_batch_morsel(scope(0..5), 1), PlannerStep::Done], + ); + let root = ScriptedPlanner::boxed( + "parent", + vec![ + child, + PlannerStep::Continue, + one_batch_morsel(scope(5..10), 2), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source()).run(root)?; + let scopes: Vec<_> = batches.iter().map(|batch| batch.scope.clone()).collect(); + assert_eq!(scopes, vec![scope(0..5), scope(5..10)]); + assert_eq!( + log.events(), + vec![ + "compute parent", + "compute child", + "compute morsel", + "compute morsel", + "compute child", + "compute parent", + "compute parent", + "compute morsel", + "compute morsel", + "compute parent", + ] + ); + Ok(()) +} + +#[test] +fn needs_io_delivers_every_request_in_order() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Io(vec![request(0, 4), request(1, 0)]), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source).run(root)?; + assert!(batches.is_empty()); + assert_eq!( + source.performed(), + vec![IoTarget::range(4, 4), IoTarget::range(0, 4)] + ); + assert_eq!( + log.events(), + vec![ + "publish root", + "deliver root 0 bytes", + "deliver root 1 bytes", + "compute root", + ] + ); + Ok(()) +} + +/// A fetch published while CPU work is left is performed once, and delivered when the planner +/// gets round to waiting for it. +#[test] +fn fetch_published_ahead_is_delivered_when_awaited() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::NeedsIO(vec![request(0, 0)]), + PlannerStep::Continue, + PlannerStep::Await(vec![IoRequestId(0)]), + PlannerStep::Done, + ], + &log, + ); + driver(&source).run(root)?; + assert_eq!(source.performed(), vec![IoTarget::range(0, 4)]); + assert_eq!( + log.events(), + vec![ + "compute root", + "compute root", + "deliver root 0 bytes", + "compute root", + ] + ); + Ok(()) +} + +#[test] +fn continue_requeues() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Continue, + PlannerStep::Continue, + PlannerStep::Done, + ], + &log, + ); + driver(&source()).run(root)?; + assert_eq!( + log.events() + .iter() + .filter(|event| event.as_str() == "compute root") + .count(), + 3 + ); + Ok(()) +} + +#[test] +fn io_failure_propagates_after_two_performs() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + source.fail(IoTarget::range(4, 4), "disk on fire"); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Io(vec![request(0, 0), request(1, 4)]), + PlannerStep::Done, + ], + &log, + ); + let err = driver(&source).run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref().is_some_and(|m| m.contains("disk on fire")), + "{err:?}" + ); + assert_eq!(source.performed().len(), 2); + Ok(()) +} + +#[test] +fn completions_out_of_order_within_one_batch() -> VortexResult<()> { + let log = Log::default(); + let source = lifo_source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Io(vec![request(0, 0), request(1, 4)]), + PlannerStep::Done, + ], + &log, + ); + driver(&source).run(root)?; + assert_eq!( + source.performed(), + vec![IoTarget::range(0, 4), IoTarget::range(4, 4)], + "submitted in request order" + ); + assert_eq!( + log.events(), + vec![ + "publish root", + "deliver root 1 bytes", + "deliver root 0 bytes", + "compute root", + ] + ); + Ok(()) +} + +#[test] +fn completions_out_of_order_across_parked_items() -> VortexResult<()> { + let log = Log::default(); + let source = lifo_source(); + let child = |id: u32, offset: u64| { + PlannerStep::Planner( + scope(0..1), + vec![ + PlannerStep::Io(vec![request(id, offset)]), + PlannerStep::Done, + ], + ) + }; + let root = ScriptedPlanner::boxed( + "root", + vec![child(5, 0), child(6, 4), PlannerStep::Done], + &log, + ); + driver(&source).run(root)?; + let deliveries: Vec<_> = log + .events() + .into_iter() + .filter(|event| event.starts_with("deliver")) + .collect(); + assert_eq!( + deliveries, + vec!["deliver child 6 bytes", "deliver child 5 bytes"], + "both children were parked before either completed" + ); + Ok(()) +} + +#[test] +fn lying_source_fails_before_delivery() -> VortexResult<()> { + let log = Log::default(); + let source = source(); + source.answer_with_size(IoTarget::range(0, 4), 4); + let root = ScriptedPlanner::boxed( + "root", + vec![PlannerStep::Io(vec![request(0, 0)]), PlannerStep::Done], + &log, + ); + let err = driver(&source).run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref().is_some_and(|m| m.contains("answered")), + "{err:?}" + ); + assert_eq!(log.events(), vec!["publish root"]); + Ok(()) +} + +#[test] +fn morsel_needs_io_midway() -> VortexResult<()> { + let log = Log::default(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Morsel( + scope(0..4), + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(1))), + MorselStep::Io(vec![request(0, 0)]), + MorselStep::Compute(MorselOutput::Batch(array_of(2))), + MorselStep::Compute(MorselOutput::Done), + ], + ), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source()).run(root)?; + assert_eq!(batches.len(), 2); + assert_arrays_eq!(batches[0].array, array_of(1), &mut ctx()); + assert_arrays_eq!(batches[1].array, array_of(2), &mut ctx()); + Ok(()) +} + +#[rstest] +#[case::empty_batch( + vec![PlannerStep::Morsel(scope(0..1), vec![MorselStep::Compute(MorselOutput::Batch(array_of(0)))])], + "empty batch" +)] +#[case::empty_io(vec![PlannerStep::Io(vec![])], "empty batch")] +#[case::republished_fetch( + vec![ + PlannerStep::NeedsIO(vec![request(0, 0)]), + PlannerStep::NeedsIO(vec![request(0, 0)]), + ], + "published while it is outstanding" +)] +#[case::waiting_for_nothing(vec![PlannerStep::Wait], "waits with no outstanding fetch")] +#[case::still_waiting_after_the_last_delivery( + vec![PlannerStep::Io(vec![request(0, 0)]), PlannerStep::Wait], + "waits with no outstanding fetch" +)] +fn protocol_errors(#[case] mut steps: Vec, #[case] message: &str) -> VortexResult<()> { + let log = Log::default(); + steps.push(PlannerStep::Done); + let root = ScriptedPlanner::boxed("root", steps, &log); + let err = driver(&source()).run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref().is_some_and(|m| m.contains(message)), + "{err:?}" + ); + Ok(()) +} + +#[test] +fn step_limit_is_enforced() -> VortexResult<()> { + let log = Log::default(); + let mut steps: Vec<_> = (0..100).map(|_| PlannerStep::Continue).collect(); + steps.push(PlannerStep::Done); + let root = ScriptedPlanner::boxed("root", steps, &log); + let driver = Driver::new(source() as Arc).with_step_limit(10); + let err = driver.run(root).err().map(|e| e.to_string()); + assert!( + err.as_deref() + .is_some_and(|m| m.contains("step limit of 10")), + "{err:?}" + ); + Ok(()) +} + +#[rstest] +#[case(IoIntent::Announce)] +#[case(IoIntent::Prefetch)] +fn optional_publication_does_not_park_or_receive_bytes( + #[case] intent: IoIntent, +) -> VortexResult<()> { + let log = Log::default(); + let source = source(); + let mut optional = request(0, 0); + optional.intent = intent; + let root = ScriptedPlanner::boxed( + "warm", + vec![ + PlannerStep::NeedsIO(vec![optional]), + one_batch_morsel(scope(0..1), 1), + PlannerStep::Done, + ], + &log, + ); + let batches = driver(&source).run(root)?; + assert_eq!(batches.len(), 1); + let submitted = source.submissions(); + assert_eq!(submitted.len(), 1); + assert_eq!(submitted[0].1.intent, intent); + assert!(source.performed().is_empty()); + assert!( + !log.events() + .iter() + .any(|event| event.starts_with("deliver")) + ); + Ok(()) +} + +#[test] +fn optional_requests_cannot_be_wait_dependencies() -> VortexResult<()> { + let source = source(); + let mut optional = request(0, 0); + optional.intent = IoIntent::Announce; + let root = ScriptedPlanner::boxed( + "root", + vec![PlannerStep::Io(vec![optional])], + &Log::default(), + ); + let error = driver(&source) + .run(root) + .err() + .map(|error| error.to_string()); + assert!(error.is_some_and(|error| error.contains("waits with no outstanding fetch"))); + assert!(source.performed().is_empty()); + Ok(()) +} + +#[test] +fn optional_registration_can_be_promoted_to_fetch() -> VortexResult<()> { + let source = source(); + let mut optional = request(0, 0); + optional.intent = IoIntent::Announce; + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::NeedsIO(vec![optional]), + PlannerStep::Io(vec![request(0, 0)]), + PlannerStep::Done, + ], + &Log::default(), + ); + driver(&source).run(root)?; + assert_eq!(source.submissions().len(), 2); + assert_eq!(source.performed().len(), 1); + Ok(()) +} + +const _: () = { + const fn assert_send() {} + assert_send::(); +}; + +fn ready_root(name: &'static str, log: &Log, rows: Range) -> Box { + ScriptedPlanner::boxed( + name, + vec![one_batch_morsel(scope(rows), 1), PlannerStep::Done], + log, + ) +} + +fn reading_root(name: &'static str, log: &Log, rows: Range) -> Box { + ScriptedPlanner::boxed( + name, + vec![ + PlannerStep::Io(vec![request(0, 0)]), + one_batch_morsel(scope(rows), 2), + PlannerStep::Done, + ], + log, + ) +} + +/// Drains `run` until it waits or idles, describing what it handed back. +fn drain(run: &mut Run) -> VortexResult> { + let mut events = Vec::new(); + loop { + events.push(match run.advance()? { + Progress::Batch(batch) => format!("batch {:?} {:?}", batch.root, batch.scope.rows), + Progress::RootDone(root) => format!("done {root:?}"), + Progress::Waiting => { + events.push("waiting".to_string()); + return Ok(events); + } + Progress::Idle => { + events.push("idle".to_string()); + return Ok(events); + } + }); + } +} + +/// Roots run side by side: a root with ready work finishes while another waits for its read, and +/// each root reads through, and finally clears, its own session. +#[test] +fn roots_finish_independently() -> VortexResult<()> { + let log = Log::default(); + let (reading_io, ready_io) = (source(), source()); + let mut run = Run::new().with_step_limit(1_000); + let reading = run.admit( + reading_root("reading", &log, 0..2), + scope(0..2), + Arc::clone(&reading_io) as Arc, + ); + let ready = run.admit( + ready_root("ready", &log, 2..3), + scope(2..3), + Arc::clone(&ready_io) as Arc, + ); + assert_eq!((reading, ready), (RootId(0), RootId(1))); + assert_eq!(run.live_roots(), 2); + + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(1) 2..3", "done RootId(1)", "waiting"] + ); + assert_eq!(run.live_roots(), 1); + assert_eq!((reading_io.clears(), ready_io.clears()), (0, 1)); + assert!(ready_io.submissions().is_empty()); + + run.complete(reading_io.wait()?)?; + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(0) 0..2", "done RootId(0)", "idle"] + ); + assert_eq!(run.live_roots(), 0); + assert_eq!(reading_io.clears(), 1); + Ok(()) +} + +#[test] +fn poll_completion_delivers_from_every_waiting_root() -> VortexResult<()> { + let log = Log::default(); + let sources = [source(), source()]; + let mut run = Run::new().with_step_limit(1_000); + for (index, io) in sources.iter().enumerate() { + let rows = index as u64..index as u64 + 1; + run.admit( + reading_root("root", &log, rows.clone()), + scope(rows), + Arc::clone(io) as Arc, + ); + } + assert_eq!(drain(&mut run)?, vec!["waiting"]); + let mut cx = Context::from_waker(Waker::noop()); + assert!(matches!(run.poll_completion(&mut cx), Poll::Ready(Ok(())))); + let mut events = drain(&mut run)?; + events.sort(); + assert_eq!( + events, + vec![ + "batch RootId(0) 0..1", + "batch RootId(1) 1..2", + "done RootId(0)", + "done RootId(1)", + "idle", + ] + ); + Ok(()) +} + +#[test] +fn cancel_drops_a_roots_work_and_clears_its_session() -> VortexResult<()> { + let log = Log::default(); + let (cancelled_io, kept_io) = (source(), source()); + let mut run = Run::new().with_step_limit(1_000); + let cancelled = run.admit( + reading_root("cancelled", &log, 0..2), + scope(0..2), + Arc::clone(&cancelled_io) as Arc, + ); + run.admit( + reading_root("kept", &log, 2..4), + scope(2..4), + Arc::clone(&kept_io) as Arc, + ); + assert_eq!(drain(&mut run)?, vec!["waiting"]); + + run.cancel(cancelled); + assert_eq!(run.live_roots(), 1); + assert_eq!(cancelled_io.clears(), 1); + run.complete(kept_io.wait()?)?; + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(1) 2..4", "done RootId(1)", "idle"] + ); + assert!( + !log.events() + .contains(&"deliver cancelled 0 bytes".to_string()) + ); + Ok(()) +} + +#[test] +fn dropping_a_run_clears_every_live_session() -> VortexResult<()> { + let log = Log::default(); + let io = source(); + let mut run = Run::new(); + run.admit( + reading_root("root", &log, 0..2), + scope(0..2), + Arc::clone(&io) as Arc, + ); + assert_eq!(drain(&mut run)?, vec!["waiting"]); + drop(run); + assert_eq!(io.clears(), 1); + Ok(()) +} + +/// Waiting for M fetches costs the driver one delivery each, not a visit that re-lists the rest. +#[test] +fn a_waiting_item_is_not_revisited_per_delivery() -> VortexResult<()> { + const FETCHES: u32 = 64; + let log = Log::default(); + let io = RecordingIoSource::default(); + let batch: Vec = (0..FETCHES) + .map(|id| { + io.canned_bytes(u64::from(id) * 4, buffer![0u8; 4].into_byte_buffer()); + request(id, u64::from(id) * 4) + }) + .collect(); + let io = Arc::new(io); + let root = ScriptedPlanner::boxed( + "root", + vec![PlannerStep::Io(batch), PlannerStep::Done], + &log, + ); + // Publish and the final compute: the deliveries in between cost no visit. + Driver::new(Arc::clone(&io) as Arc) + .with_step_limit(2) + .run(root)?; + assert_eq!(io.performed().len(), FETCHES as usize); + Ok(()) +} + +#[test] +fn waking_subtree_precedes_an_already_discovered_sibling() -> VortexResult<()> { + let log = Log::default(); + let io = source(); + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Planner( + scope(0..1), + vec![ + PlannerStep::Io(vec![request(0, 0)]), + one_batch_morsel(scope(0..1), 1), + PlannerStep::Done, + ], + ), + PlannerStep::Morsel( + scope(1..2), + vec![ + MorselStep::Compute(MorselOutput::Batch(array_of(1))), + MorselStep::Compute(MorselOutput::Batch(array_of(1))), + MorselStep::Compute(MorselOutput::Done), + ], + ), + PlannerStep::Done, + ], + &log, + ); + let mut run = Run::new().with_step_limit(1_000); + run.admit(root, scope(0..2), Arc::clone(&io) as Arc); + let Progress::Batch(batch) = run.advance()? else { + panic!("the later sibling should run while the first child waits"); + }; + assert_eq!(batch.scope, scope(1..2)); + + run.complete(io.wait()?)?; + // The waking child creates a grandchild after the sibling already exists. Its subtree + // must still win, including over the sibling's continuation after its first batch. + let Progress::Batch(batch) = run.advance()? else { + panic!("the waking subtree should produce a batch first"); + }; + assert_eq!(batch.scope, scope(0..1)); + assert_eq!( + drain(&mut run)?, + vec!["batch RootId(0) 1..2", "done RootId(0)", "idle"] + ); + Ok(()) +} + +#[test] +fn runnable_roots_follow_admission_order() -> VortexResult<()> { + let log = Log::default(); + let mut run = Run::new().with_step_limit(1_000); + for rows in [0..1, 1..2] { + let root = ScriptedPlanner::boxed( + "root", + vec![ + PlannerStep::Continue, + one_batch_morsel(scope(rows.clone()), 1), + PlannerStep::Done, + ], + &log, + ); + run.admit(root, scope(rows), source()); + } + assert_eq!( + drain(&mut run)?, + vec![ + "batch RootId(0) 0..1", + "done RootId(0)", + "batch RootId(1) 1..2", + "done RootId(1)", + "idle", + ] + ); + Ok(()) +} + +/// Makes the recording source's completions available to nonblocking polling. +struct PollingSource { + inner: Arc, + outstanding: AtomicUsize, +} + +impl IoSource for PollingSource { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + self.outstanding.fetch_add( + batch + .iter() + .filter(|request| request.intent == IoIntent::Fetch) + .count(), + Ordering::Relaxed, + ); + self.inner.submit(owner, batch) + } + + fn poll(&self) -> VortexResult> { + if self.outstanding.load(Ordering::Relaxed) == 0 { + return Ok(None); + } + self.wait().map(Some) + } + + fn poll_completion(&self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.wait()) + } + + fn wait(&self) -> VortexResult { + let completion = self.inner.wait()?; + self.outstanding.fetch_sub(1, Ordering::Relaxed); + Ok(completion) + } + + fn release(&self, owner: IoOwnerId) { + self.inner.release(owner); + } + + fn clear(&self) { + self.outstanding.store(0, Ordering::Relaxed); + self.inner.clear(); + } +} + +#[test] +fn polls_io_while_a_later_branch_keeps_computing() -> VortexResult<()> { + let log = Log::default(); + let io = Arc::new(PollingSource { + inner: source(), + outstanding: AtomicUsize::new(0), + }); + let mut run = Run::new().with_step_limit(1_000); + run.admit(reading_root("earlier", &log, 0..1), scope(0..1), io); + let mut script: Vec<_> = (0..IO_POLL_INTERVAL * 2) + .map(|_| PlannerStep::Continue) + .collect(); + script.push(one_batch_morsel(scope(1..2), 1)); + script.push(PlannerStep::Done); + run.admit( + ScriptedPlanner::boxed("later", script, &log), + scope(1..2), + source(), + ); + let Progress::Batch(batch) = run.advance()? else { + panic!("ready IO should wake the earlier root before the later root finishes"); + }; + assert_eq!(batch.root, RootId(0)); + let later_visits = log + .events() + .iter() + .filter(|event| event.as_str() == "compute later") + .count(); + assert!(later_visits > 0 && later_visits <= IO_POLL_INTERVAL); + Ok(()) +} diff --git a/vortex-scan/src/planning/mod.rs b/vortex-scan/src/planning/mod.rs index 5e293dae237..3095712f161 100644 --- a/vortex-scan/src/planning/mod.rs +++ b/vortex-scan/src/planning/mod.rs @@ -6,6 +6,10 @@ //! Planners and morsels are `Send` and move between threads with the run that owns them; //! stage-specific phases remain private to those implementations. +pub mod driver; pub mod morsel; pub mod next; pub mod planner; + +#[cfg(test)] +mod tests; diff --git a/vortex-scan/src/planning/tests/mod.rs b/vortex-scan/src/planning/tests/mod.rs new file mode 100644 index 00000000000..59229044cbc --- /dev/null +++ b/vortex-scan/src/planning/tests/mod.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scripted work and IO sources for the planning driver tests. + +pub mod scripted; diff --git a/vortex-scan/src/planning/tests/scripted.rs b/vortex-scan/src/planning/tests/scripted.rs new file mode 100644 index 00000000000..4035bb8ac61 --- /dev/null +++ b/vortex-scan/src/planning/tests/scripted.rs @@ -0,0 +1,418 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scripted planners, morsels, and a recording IO source for driver tests. +//! +//! A scripted object follows its script exactly and panics when the driver calls it against the +//! protocol, so a driver bug becomes a test failure that names the script position. + +use std::ops::Range; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::planning::morsel::Morsel; +use crate::planning::morsel::MorselOutput; +use crate::planning::planner::Planner; +use crate::planning::planner::PlannerOutput; +use crate::planning::planner::State; +use crate::planning::planner::WorkScope; + +/// Shared record of the calls the driver made, as readable one-line events. +#[derive(Clone, Default)] +pub struct Log(Arc>>); + +impl Log { + fn push(&self, event: String) { + self.0.lock().push(event); + } + + pub fn events(&self) -> Vec { + self.0.lock().clone() + } +} + +/// One scripted planner step. Every step but `Await` and `Wait` is the output of one `compute()`. +pub enum PlannerStep { + Done, + Continue, + /// Publishes the batch and goes straight on to the next step. + NeedsIO(IoBatch), + Planner(WorkScope, Vec), + Morsel(WorkScope, Vec), + /// Publishes the batch, then waits until every request in it is delivered. + Io(IoBatch), + /// Waits for requests an earlier `NeedsIO` step published. + Await(Vec), + /// Reports `Waiting` without having published anything. + Wait, +} + +/// One scripted morsel step. +pub enum MorselStep { + Compute(MorselOutput), + /// Publishes the batch, then waits until every request in it is delivered. + Io(IoBatch), +} + +/// The ids a waiting step still expects. +struct Pending(Vec); + +impl Pending { + fn of(batch: &IoBatch) -> Self { + Self(batch.iter().map(|request| request.request).collect()) + } + + /// Removes a delivered id and returns whether the step has everything it waited for. + fn deliver(&mut self, name: &str, position: usize, request: IoRequestId) -> bool { + let index = self + .0 + .iter() + .position(|pending| *pending == request) + .unwrap_or_else(|| { + panic!("{name}: delivery of unknown request {request:?} at step {position}") + }); + self.0.remove(index); + self.0.is_empty() + } +} + +pub struct ScriptedPlanner { + name: &'static str, + script: Vec, + position: usize, + /// What the step at `position` waits for, once it has published or received something. + pending: Option, + log: Log, +} + +impl ScriptedPlanner { + pub fn new(name: &'static str, script: Vec, log: &Log) -> Self { + Self { + name, + script, + position: 0, + pending: None, + log: log.clone(), + } + } + + /// A fresh scripted planner, boxed as the driver takes it. + pub fn boxed(name: &'static str, script: Vec, log: &Log) -> Box { + Box::new(Self::new(name, script, log)) + } +} + +impl IoConsumer for ScriptedPlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.log.push(format!( + "deliver {} {} {}", + self.name, + request.0, + result.kind() + )); + let position = self.position; + if let Some(PlannerStep::Await(ids)) = self.script.get(position) + && self.pending.is_none() + { + self.pending = Some(Pending(ids.clone())); + } + let Some(pending) = self.pending.as_mut() else { + panic!( + "{}: unexpected delivery at script step {position}", + self.name + ); + }; + if pending.deliver(self.name, position, request) { + self.pending = None; + self.position += 1; + } + } +} + +impl Planner for ScriptedPlanner { + fn state(&self) -> State { + match self.script.get(self.position) { + None => State::Done, + Some(PlannerStep::Io(_)) if self.pending.is_none() => State::NeedsCompute, + Some(PlannerStep::Io(_) | PlannerStep::Await(_) | PlannerStep::Wait) => State::Waiting, + Some(_) => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + let position = self.position; + if let Some(PlannerStep::Io(batch)) = self.script.get(position) { + assert!( + self.pending.is_none(), + "{}: compute called while script step {position} waits for IO", + self.name + ); + self.log.push(format!("publish {}", self.name)); + // A batch with nothing to wait for is left for the driver to reject. + self.pending = Some(Pending::of(batch)); + return Ok(PlannerOutput::NeedsIO(batch.clone())); + } + self.log.push(format!("compute {}", self.name)); + self.position += 1; + let Some(slot) = self.script.get_mut(position) else { + panic!("{}: compute past the end of the script", self.name); + }; + let step = std::mem::replace(slot, PlannerStep::Done); + Ok(match step { + PlannerStep::Done => PlannerOutput::Done, + PlannerStep::Continue => PlannerOutput::Continue, + PlannerStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + PlannerStep::Planner(scope, script) => { + PlannerOutput::Planner(scope, Self::boxed("child", script, &self.log)) + } + PlannerStep::Morsel(scope, script) => PlannerOutput::Morsel( + scope, + Box::new(ScriptedMorsel::new("morsel", script, &self.log)), + ), + PlannerStep::Io(_) | PlannerStep::Await(_) | PlannerStep::Wait => panic!( + "{}: compute called while script step {position} waits for IO", + self.name + ), + }) + } +} + +pub struct ScriptedMorsel { + name: &'static str, + script: Vec, + position: usize, + pending: Option, + log: Log, +} + +impl ScriptedMorsel { + pub fn new(name: &'static str, script: Vec, log: &Log) -> Self { + Self { + name, + script, + position: 0, + pending: None, + log: log.clone(), + } + } +} + +impl IoConsumer for ScriptedMorsel { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.log.push(format!( + "deliver {} {} {}", + self.name, + request.0, + result.kind() + )); + let position = self.position; + let Some(pending) = self.pending.as_mut() else { + panic!( + "{}: unexpected delivery at script step {position}", + self.name + ); + }; + if pending.deliver(self.name, position, request) { + self.pending = None; + self.position += 1; + } + } +} + +impl Morsel for ScriptedMorsel { + fn state(&self) -> State { + match self.script.get(self.position) { + None => State::Done, + Some(MorselStep::Io(_)) if self.pending.is_some() => State::Waiting, + Some(_) => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + let position = self.position; + match self.script.get_mut(position) { + Some(MorselStep::Io(batch)) => { + assert!( + self.pending.is_none(), + "{}: compute called while script step {position} waits for IO", + self.name + ); + let batch = batch.clone(); + self.log.push(format!("publish {}", self.name)); + self.pending = Some(Pending::of(&batch)); + Ok(MorselOutput::NeedsIO(batch)) + } + Some(MorselStep::Compute(output)) => { + self.log.push(format!("compute {}", self.name)); + self.position += 1; + Ok(std::mem::replace(output, MorselOutput::Done)) + } + None => panic!("{}: compute past the end of the script", self.name), + } + } +} + +enum Canned { + Bytes(ByteBuffer), + Size(u64), + Fail(String), +} + +/// Which submitted request `wait` completes first. +#[derive(Clone, Copy, Default)] +pub enum Order { + /// Oldest submission first. + #[default] + Fifo, + /// Newest submission first, so a batch and concurrently parked items complete in reverse. + Lifo, +} + +/// An [`IoSource`] answering from canned buffers. Every submission is resolved at once and +/// held until `wait`, which hands completions back in the configured [`Order`]. It records the +/// targets in submission order, and how often it was cleared. +#[derive(Default)] +pub struct RecordingIoSource { + canned: Mutex>, + performed: Mutex>, + ready: Mutex>, + order: Order, + submissions: Mutex>, + clears: Mutex, +} + +impl RecordingIoSource { + pub fn lifo() -> Self { + Self { + order: Order::Lifo, + ..Default::default() + } + } + + /// Answers a range at `offset` with `bytes`. + pub fn canned_bytes(&self, offset: u64, bytes: ByteBuffer) { + self.canned + .lock() + .insert(IoTarget::range(offset, bytes.len()), Canned::Bytes(bytes)); + } + + /// Answers `target` with a size regardless of what it asked for, to test the driver's check. + pub fn answer_with_size(&self, target: IoTarget, size: u64) { + self.canned.lock().insert(target, Canned::Size(size)); + } + + pub fn fail(&self, target: IoTarget, message: &str) { + self.canned + .lock() + .insert(target, Canned::Fail(message.to_string())); + } + + pub fn submissions(&self) -> Vec<(IoOwnerId, IoRequest)> { + self.submissions.lock().clone() + } + + pub fn performed(&self) -> Vec { + self.performed.lock().clone() + } + + /// How many times the source was cleared. + pub fn clears(&self) -> usize { + *self.clears.lock() + } +} + +impl IoSource for RecordingIoSource { + fn submit(&self, work: IoOwnerId, batch: Vec) -> VortexResult<()> { + for request in batch { + self.submissions.lock().push((work, request)); + if request.intent != IoIntent::Fetch { + continue; + } + self.performed.lock().push(request.target); + let result = match self.canned.lock().get(&request.target) { + Some(Canned::Bytes(bytes)) => { + Ok(IoResult::Bytes(BufferHandle::new_host(bytes.clone()))) + } + Some(Canned::Size(size)) => Ok(IoResult::Size(*size)), + Some(Canned::Fail(message)) => Err(vortex_err!("{message}")), + None => Err(vortex_err!("no canned answer for {:?}", request.target)), + }; + self.ready.lock().push(Completion { + owner: work, + request: request.request, + result, + }); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + Ok(None) + } + + fn poll_completion(&self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.wait()) + } + + fn release(&self, work: IoOwnerId) { + self.ready + .lock() + .retain(|completion| completion.owner != work); + } + + fn clear(&self) { + *self.clears.lock() += 1; + self.ready.lock().clear(); + } + + fn wait(&self) -> VortexResult { + let mut ready = self.ready.lock(); + let completion = match self.order { + Order::Fifo if !ready.is_empty() => Some(ready.remove(0)), + Order::Fifo => None, + Order::Lifo => ready.pop(), + }; + completion.ok_or_else(|| vortex_err!("RecordingIoSource: wait with nothing outstanding")) + } +} + +/// A small non-nullable `u32` array of `n` rows. +pub fn array_of(n: u64) -> ArrayRef { + PrimitiveArray::from_iter(0..u32::try_from(n).expect("small test array")).into_array() +} + +/// An execution context over a bare array session, for `assert_arrays_eq!`. +pub fn ctx() -> ExecutionCtx { + array_session().create_execution_ctx() +} + +pub fn scope(rows: Range) -> WorkScope { + WorkScope { + file_ordinal: 0, + rows, + } +}