diff --git a/Cargo.lock b/Cargo.lock index b0f98336b0b..03d34e78ed3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10556,6 +10556,26 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-scan-v2" +version = "0.1.0" +dependencies = [ + "futures", + "itertools 0.14.0", + "tracing", + "tracing-subscriber", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-file", + "vortex-io", + "vortex-layout", + "vortex-mask", + "vortex-scan", + "vortex-session", + "vortex-utils", +] + [[package]] name = "vortex-sequence" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..cefacf71138 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ members = [ "vortex-btrblocks", "vortex-layout", "vortex-scan", + "vortex-scan-v2", "vortex-file", "vortex-ipc", "vortex", @@ -325,6 +326,7 @@ vortex-proto = { version = "0.1.0", path = "./vortex-proto", default-features = vortex-row = { version = "0.1.0", path = "./vortex-row", default-features = false } vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-features = false } vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = false } +vortex-scan-v2 = { version = "0.1.0", path = "./vortex-scan-v2", default-features = false } vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-features = false } vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false } vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false } diff --git a/docs/developer-guide/internals/scan-planning.md b/docs/developer-guide/internals/scan-planning.md index e03131924ac..0ae2a605efe 100644 --- a/docs/developer-guide/internals/scan-planning.md +++ b/docs/developer-guide/internals/scan-planning.md @@ -18,7 +18,7 @@ on demand. | `Take` | index `values` by `codes` | | `ListPack` | assemble a list from elements and offsets, plus optional validity | | `Eval` | apply an expression to its child | -| `RowIdx` | offset row numbers into the file's row domain | +| `RowIdx` | generate row numbers for the current execution row domain | Naming operators for what they compute is what lets one rule cover every case. `Concat` of `Concat` flattens on shape alone, and `Take` over `SegmentScan` is the dictionary pushdown, @@ -55,9 +55,14 @@ ID, dtype, row count, and lazy children. Only the unsized tail containing the vt already serialize their metadata; the ones holding a read context or a bound expression return `None` until those codecs exist. +## Execution + +Each operator executes over a row range and selection mask. `SegmentScan` reads its segment, +structural operators combine their children, and `Eval` applies the remaining derived work. +`vortex-scan-v2` copies the existing scan orchestration around this API, so the original +`LayoutReader` scanner is untouched while the plan-native path is developed. + ## Future work -Plans currently stop at construction and optimization. Still to come: a plan registry and foreign -operator placeholder so third-party operators survive a round trip, a serialization envelope, and -an execution stage that walks an optimized plan, reads the referenced segments, and returns the -query result. +Still to come: a plan registry and foreign operator placeholder so third-party operators survive +a round trip, and a serialization envelope. diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index e7c83ec2950..d4ce3912a40 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -297,7 +297,7 @@ fn row_idx_dtype() -> DType { } // Returns a SequenceArray representing the row indices for the given row range, -fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { +pub(crate) fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { Sequence::try_new( PValue::U64(row_offset + row_range.start), PValue::U64(1), diff --git a/vortex-layout/src/plan/execution.rs b/vortex-layout/src/plan/execution.rs new file mode 100644 index 00000000000..e9e4a5e06d8 --- /dev/null +++ b/vortex-layout/src/plan/execution.rs @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::future::BoxFuture; +use vortex_array::ArrayRef; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::segments::SegmentSource; + +/// Future resolving to the array produced by a physical plan. +pub type PlanArrayFuture = BoxFuture<'static, VortexResult>; + +/// Runtime dependencies shared by every node in a plan execution. +#[derive(Clone)] +pub struct PlanExecutionContext { + segment_source: Arc, + session: VortexSession, + row_offset: u64, +} + +impl PlanExecutionContext { + /// Creates an execution context over a segment source and Vortex session. + pub fn new(segment_source: Arc, session: VortexSession) -> Self { + Self { + segment_source, + session, + row_offset: 0, + } + } + + /// Sets the global row index of the first row in the root plan's row domain. + pub fn with_row_offset(mut self, row_offset: u64) -> Self { + self.row_offset = row_offset; + self + } + + /// Returns the global row index of the first row in the current plan's row domain. + pub fn row_offset(&self) -> u64 { + self.row_offset + } + + /// Derives the execution context for a child whose row domain starts within this one. + pub(crate) fn child_row_domain(&self, relative_row_offset: u64) -> VortexResult { + let row_offset = self + .row_offset + .checked_add(relative_row_offset) + .ok_or_else(|| vortex_error::vortex_err!("Plan row-domain offset overflow"))?; + Ok(Self { + segment_source: Arc::clone(&self.segment_source), + session: self.session.clone(), + row_offset, + }) + } + + /// Returns the segment source used to satisfy leaf reads. + pub fn segment_source(&self) -> &Arc { + &self.segment_source + } + + /// Returns the Vortex session used for array decoding and expression execution. + pub fn session(&self) -> &VortexSession { + &self.session + } +} diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index 7aa8c41ea55..f7b01fe5b54 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -126,6 +126,7 @@ fn lower_dict(layout: &DictLayout) -> VortexResult { TakePlan::from_children_unchecked( layout.dtype().clone(), layout.row_count(), + layout.has_all_values_referenced(), lazy_children(layout.to_layout(), vec![1, 0]), ) }) diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 39fffd924d5..c8ccea89df8 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -9,6 +9,7 @@ mod children; mod display; +mod execution; mod lower; mod optimize; pub mod optimizer; @@ -22,6 +23,8 @@ pub use display::PlanSummaryExtractor; pub use display::PlanTreeContext; pub use display::PlanTreeDisplay; pub use display::PlanTreeExtractor; +pub use execution::PlanArrayFuture; +pub use execution::PlanExecutionContext; pub use lower::lower; pub use optimize::optimize; pub use plans::Concat; @@ -43,6 +46,7 @@ pub use plans::SegmentScan; pub use plans::SegmentScanData; pub use plans::SegmentScanPlan; pub use plans::Take; +pub use plans::TakeData; pub use plans::TakePlan; pub use plans::Zoned; pub use plans::ZonedData; diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 2e814156309..b9029f5a518 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -2,18 +2,32 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::future; +use std::ops::Range; use std::sync::Arc; +use futures::FutureExt; +use futures::TryStreamExt; +use futures::stream::FuturesOrdered; +use vortex_array::Canonical; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::plan::Eval; use crate::plan::EvalPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -139,6 +153,62 @@ impl PlanVTable for Concat { Ok(()) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "Concat row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + vortex_ensure!( + mask.len() == usize::try_from(row_range.end - row_range.start)?, + "Concat mask length mismatch" + ); + if row_range.is_empty() { + let empty = Canonical::empty(plan.dtype()).into_array(); + return Ok(future::ready(Ok(empty)).boxed()); + } + + let mut chunk_futures = Vec::new(); + for (chunk, &chunk_offset) in plan.children().iter().zip(plan.row_offsets()) { + let chunk = chunk?; + let chunk_end = chunk_offset + .checked_add(chunk.row_count()) + .ok_or_else(|| vortex_err!("Chunk row offset overflow"))?; + let start = row_range.start.max(chunk_offset); + let end = row_range.end.min(chunk_end); + if start < end { + let child_range = start - chunk_offset..end - chunk_offset; + let mask_range = usize::try_from(start - row_range.start)? + ..usize::try_from(end - row_range.start)?; + let child_ctx = ctx.child_row_domain(chunk_offset)?; + chunk_futures.push(chunk.execute( + &child_ctx, + &child_range, + mask.slice(mask_range), + )?); + } + } + + Ok(async move { + let chunks: Vec<_> = FuturesOrdered::from_iter(chunk_futures) + .try_collect() + .await?; + vortex_ensure!(!chunks.is_empty(), "Non-empty row range selected no chunks"); + if chunks.len() == 1 { + return Ok(chunks.into_iter().next().vortex_expect("one chunk")); + } + let dtype = chunks[0].dtype().clone(); + Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) + } + .boxed()) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { Cow::Owned(format!("chunks[{index}]")) } diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index f1edb967e41..505614f4a2d 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -3,15 +3,20 @@ use std::borrow::Cow; use std::fmt; +use std::ops::Range; +use futures::FutureExt; use vortex_array::EmptyMetadata; +use vortex_array::MaskFuture; use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -106,6 +111,17 @@ impl PlanVTable for Eval { Ok(()) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let child = plan.child_plan()?.execute(ctx, row_range, mask)?; + let expression = plan.expression().clone(); + Ok(async move { child.await?.apply_bound(&expression) }.boxed()) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { if index == 0 { Cow::Borrowed("child") diff --git a/vortex-layout/src/plan/plans/list_pack.rs b/vortex-layout/src/plan/plans/list_pack.rs index d9a23de818b..1c2992e3d74 100644 --- a/vortex-layout/src/plan/plans/list_pack.rs +++ b/vortex-layout/src/plan/plans/list_pack.rs @@ -2,18 +2,35 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use std::sync::Arc; +use futures::FutureExt; +use futures::try_join; +use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListArray; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -116,6 +133,75 @@ impl PlanVTable for ListPack { validate_children(plan.dtype(), plan.row_count(), children) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "ListPack row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + let row_count = usize::try_from(row_range.end - row_range.start)?; + vortex_ensure!(mask.len() == row_count, "ListPack mask length mismatch"); + + let offsets_range = row_range.start + ..row_range + .end + .checked_add(1) + .ok_or_else(|| vortex_err!("List offsets range overflow"))?; + let offsets = plan.offsets()?.execute( + ctx, + &offsets_range, + MaskFuture::new_true(row_count.saturating_add(1)), + )?; + let validity = plan + .validity()? + .map(|validity| validity.execute(ctx, row_range, MaskFuture::new_true(row_count))) + .transpose()?; + let elements = plan.elements()?; + let execution = ctx.clone(); + let dtype = plan.dtype().clone(); + let nullability = dtype.nullability(); + + Ok(async move { + let (offsets, mask) = try_join!(offsets, mask)?; + if mask.all_false() { + return Ok(Canonical::empty(&dtype).into_array()); + } + + let elements_range = elements_range_from_offsets(&offsets, execution.session())?; + let elements_count = usize::try_from(elements_range.end - elements_range.start)?; + let elements = elements + .execute( + &execution, + &elements_range, + MaskFuture::new_true(elements_count), + )? + .await?; + let validity = match validity { + Some(validity) => Some(validity.await?), + None => None, + }; + let offsets = rebase_offsets(offsets, elements_range.start)?; + // SAFETY: lowering from a list layout guarantees compatible elements and monotonically + // increasing offsets. Rebasing preserves the represented list lengths. + let list = unsafe { + ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) + } + .into_array(); + if mask.all_true() { + Ok(list) + } else { + list.filter(mask) + } + } + .boxed()) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { match index { ELEMENTS => Cow::Borrowed("elements"), @@ -188,3 +274,42 @@ fn validate_children(dtype: &DType, row_count: u64, children: &PlanChildren) -> } Ok(()) } + +fn elements_range_from_offsets( + offsets: &ArrayRef, + session: &vortex_session::VortexSession, +) -> VortexResult> { + if offsets.is_empty() { + return Ok(0..0); + } + let mut ctx = session.create_execution_ctx(); + let start = offsets + .execute_scalar(0, &mut ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value must fit in u64"); + let end = offsets + .execute_scalar(offsets.len() - 1, &mut ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value must fit in u64"); + Ok(start..end) +} + +fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { + if first == 0 { + return Ok(offsets); + } + let constant = ConstantArray::new(first, offsets.len()) + .into_array() + .cast(offsets.dtype().clone())?; + offsets.binary(constant, Operator::Sub) +} + +fn create_validity(validity: Option, nullability: Nullability) -> Validity { + match validity { + Some(validity) => Validity::Array(validity), + None if nullability.is_nullable() => Validity::AllValid, + None => Validity::NonNullable, + } +} diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index f24381c8951..7a50c4fe827 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -35,6 +35,7 @@ pub use segment_scan::SegmentScanData; pub use segment_scan::SegmentScanPlan; pub(crate) use take::ExpressionTakeRule; pub use take::Take; +pub use take::TakeData; pub use take::TakePlan; pub(crate) use zoned::ExpressionZonedRule; pub use zoned::Zoned; diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 6497cc63d45..dc341b0b57a 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -2,8 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; +use futures::FutureExt; +use futures::try_join; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; @@ -22,6 +28,7 @@ use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::pack::Pack as PackFn; use vortex_array::scalar_fn::fns::pack::PackOptions; use vortex_array::scalar_fn::fns::select::Select; +use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -32,7 +39,9 @@ use vortex_session::registry::CachedId; use crate::plan::Eval; use crate::plan::EvalPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -193,6 +202,51 @@ impl PlanVTable for Pack { Ok(()) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "Pack row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + vortex_ensure!( + mask.len() == usize::try_from(row_range.end - row_range.start)?, + "Pack mask length mismatch" + ); + let names = plan.fields().names().clone(); + let field_count = plan.nfields(); + let mut field_futures = Vec::with_capacity(field_count); + for index in 0..field_count { + let child = field_plan(plan, index)?; + field_futures.push(child.execute(ctx, row_range, mask.clone())?); + } + let validity = plan + .validity()? + .map(|validity| validity.execute(ctx, row_range, mask.clone())) + .transpose()?; + let output_mask = mask; + + Ok(async move { + let fields = futures::future::try_join_all(field_futures); + let validity = async move { + match validity { + Some(validity) => validity.await.map(Some), + None => Ok(None), + } + }; + let (fields, validity) = try_join!(fields, validity)?; + let len = output_mask.await?.true_count(); + let validity = validity.map_or(Validity::NonNullable, Validity::Array); + Ok(StructArray::try_new(names, fields, len, validity)?.into_array()) + } + .boxed()) + } + fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { assert!( index < plan.children().len(), diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index 893980825d3..336eddf949a 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -3,8 +3,12 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::ops::Range; +use futures::FutureExt; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; @@ -22,10 +26,13 @@ use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::layouts::row_idx::RowIdx as RowIdxFn; +use crate::layouts::row_idx::idx_array; use crate::plan::EvalPlan; use crate::plan::PackPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -86,6 +93,40 @@ impl PlanVTable for RowIdx { ) -> VortexResult<()> { check_child_count("RowIdx", children, 0) } + + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "RowIdx row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + vortex_ensure!( + mask.len() == usize::try_from(row_range.end - row_range.start)?, + "RowIdx mask length mismatch" + ); + let row_offset = ctx.row_offset(); + vortex_ensure!( + row_offset.checked_add(row_range.start).is_some() + && (row_range.is_empty() || row_offset.checked_add(row_range.end - 1).is_some()), + "RowIdx offset overflows u64" + ); + let array = idx_array(row_offset, row_range).into_array(); + Ok(async move { + let mask = mask.await?; + if mask.all_true() { + Ok(array) + } else { + array.filter(mask) + } + } + .boxed()) + } } /// Plans an expression over a data source and its global row-index domain. diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/plans/segment_scan.rs index d2b20df89ad..bbb2fdecdbd 100644 --- a/vortex-layout/src/plan/plans/segment_scan.rs +++ b/vortex-layout/src/plan/plans/segment_scan.rs @@ -1,15 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; + +use futures::FutureExt; use vortex_array::EmptyMetadata; +use vortex_array::MaskFuture; use vortex_array::dtype::DType; +use vortex_array::serde::SerializedArray; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanVTable; @@ -92,4 +100,49 @@ impl PlanVTable for SegmentScan { check_child_count("SegmentScan", children, 0)?; Ok(()) } + + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= plan.row_count(), + "SegmentScan row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + let row_count = usize::try_from(plan.row_count())?; + let row_range = usize::try_from(row_range.start)?..usize::try_from(row_range.end)?; + vortex_ensure!( + mask.len() == row_range.len(), + "SegmentScan mask length mismatch" + ); + + let segment = ctx.segment_source().request(plan.segment_id()); + let array_ctx = plan.array_ctx().clone(); + let array_tree = plan.array_tree().cloned(); + let dtype = plan.dtype().clone(); + let session = ctx.session().clone(); + + Ok(async move { + let segment = segment.await?; + let serialized = if let Some(array_tree) = array_tree { + SerializedArray::from_flatbuffer_and_segment(array_tree, segment)? + } else { + SerializedArray::try_from(segment)? + }; + let mut array = serialized.decode(&dtype, row_count, &array_ctx, &session)?; + if row_range.start > 0 || row_range.end < array.len() { + array = array.slice(row_range)?; + } + let mask = mask.await?; + if !mask.all_true() { + array = array.filter(mask)?; + } + Ok(array) + } + .boxed()) + } } diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/plans/take.rs index af9eca55da1..892ea63fe7e 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/plans/take.rs @@ -2,18 +2,27 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; +use futures::FutureExt; +use futures::try_join; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::DictArray; use vortex_array::dtype::DType; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::label_bound_tree; +use vortex_array::optimizer::ArrayOptimizer; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use crate::plan::Eval; use crate::plan::EvalPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -28,6 +37,14 @@ const VALUES: usize = 1; #[derive(Clone, Debug)] pub struct Take; +/// Whether every dictionary value is referenced by at least one code. +/// +/// Lowering carries this over from the dictionary layout, since the operator does not hold one. +#[derive(Clone, Debug)] +pub struct TakeData { + all_values_referenced: bool, +} + /// A plan that indexes one child by another. pub type TakePlan = Plan; @@ -41,6 +58,7 @@ impl TakePlan { pub(crate) unsafe fn from_children_unchecked( dtype: DType, row_count: u64, + all_values_referenced: bool, children: PlanChildren, ) -> Self { PlanParts { @@ -48,7 +66,9 @@ impl TakePlan { dtype, row_count, children, - data: (), + data: TakeData { + all_values_referenced, + }, } .into_typed() } @@ -57,12 +77,33 @@ impl TakePlan { /// /// The row domain is that of `codes`, and the output dtype is that of `values`. pub fn new(codes: PlanRef, values: PlanRef) -> Self { + Self::new_with_all_values_referenced(codes, values, false) + } + + /// Creates a take that records whether every value is referenced by some code. + pub fn new_with_all_values_referenced( + codes: PlanRef, + values: PlanRef, + all_values_referenced: bool, + ) -> Self { let dtype = values .dtype() .union_nullability(codes.dtype().nullability()); let row_count = codes.row_count(); // SAFETY: Parent metadata is derived from the ordered children immediately above. - unsafe { Self::from_children_unchecked(dtype, row_count, vec![codes, values].into()) } + unsafe { + Self::from_children_unchecked( + dtype, + row_count, + all_values_referenced, + vec![codes, values].into(), + ) + } + } + + /// Returns whether every value is referenced by at least one code. + pub fn all_values_referenced(&self) -> bool { + self.data().all_values_referenced } /// Returns the plan producing indices. @@ -77,7 +118,7 @@ impl TakePlan { } impl PlanVTable for Take { - type PlanData = (); + type PlanData = TakeData; type Metadata = EmptyMetadata; fn id(&self) -> PlanId { @@ -110,6 +151,37 @@ impl PlanVTable for Take { Ok(()) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let codes_plan = plan.codes()?; + let values_plan = plan.values()?; + let codes = codes_plan.execute(ctx, row_range, mask)?; + let values_len = usize::try_from(values_plan.row_count())?; + let values = values_plan.execute( + ctx, + &(0..values_plan.row_count()), + MaskFuture::new_true(values_len), + )?; + let all_values_referenced = plan.all_values_referenced(); + + Ok(async move { + let (codes, values) = try_join!(codes, values)?; + // SAFETY: lowering from a dict layout guarantees integer codes and matching dtypes. + let dictionary = unsafe { + DictArray::new_unchecked(codes, values) + .set_all_values_referenced(all_values_referenced) + } + .into_array() + .optimize()?; + Ok(dictionary) + } + .boxed()) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { match index { CODES => Cow::Borrowed("codes"), diff --git a/vortex-layout/src/plan/plans/zoned.rs b/vortex-layout/src/plan/plans/zoned.rs index f458e093ce5..08eb119386a 100644 --- a/vortex-layout/src/plan/plans/zoned.rs +++ b/vortex-layout/src/plan/plans/zoned.rs @@ -3,8 +3,10 @@ use std::borrow::Cow; use std::fmt; +use std::ops::Range; use vortex_array::EmptyMetadata; +use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::traversal::NodeExt; @@ -16,7 +18,9 @@ use vortex_session::registry::CachedId; use crate::plan::Eval; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -161,6 +165,17 @@ impl PlanVTable for Zoned { Ok(()) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + plan.data_plan()? + .ok_or_else(|| vortex_error::vortex_err!("Zoned pruning execution is not available"))? + .execute(ctx, row_range, mask) + } + fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { if plan.is_pruning() { return if index == 0 { diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index b25d8fd3ca3..eb5153e3959 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -5,10 +5,17 @@ use std::fmt; use std::num::NonZeroUsize; use std::sync::Arc; +use vortex_array::ArrayContext; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::max::Max; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -25,6 +32,8 @@ use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_io::runtime::single::block_on; +use vortex_io::session::RuntimeSessionExt; use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; @@ -32,10 +41,12 @@ use super::*; use crate::LayoutBuildContext; use crate::LayoutEncoding; use crate::LayoutRef; +use crate::LayoutStrategy; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::dict::DictLayout; use crate::layouts::flat::FlatLayout; +use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::foreign::new_foreign_layout; use crate::layouts::list::ListLayout; use crate::layouts::row_idx::row_idx; @@ -43,6 +54,9 @@ use crate::layouts::struct_::StructLayout; use crate::layouts::zoned::LegacyStatsLayoutEncoding; use crate::layouts::zoned::ZonedLayout; use crate::segments::SegmentId; +use crate::segments::TestSegments; +use crate::sequence::SequenceId; +use crate::sequence::SequentialArrayStreamExt; fn primitive(ptype: PType, nullability: Nullability) -> DType { DType::Primitive(ptype, nullability) @@ -525,6 +539,51 @@ fn row_idx_only_expression_uses_row_idx_source() -> VortexResult<()> { Ok(()) } +#[test] +fn row_idx_source_adds_the_execution_range_start() -> VortexResult<()> { + block_on(|handle| async move { + let session = crate::test::new_session().with_handle(handle); + let execution = + PlanExecutionContext::new(Arc::new(TestSegments::default()), session.clone()) + .with_row_offset(100); + let plan = RowIdxPlan::new(6).into_plan(); + + let actual = plan + .execute(&execution, &(2..5), MaskFuture::new_true(3))? + .await?; + let expected = PrimitiveArray::from_iter([102_u64, 103, 104]).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +} + +#[test] +fn row_idx_source_uses_concat_child_row_domains() -> VortexResult<()> { + block_on(|handle| async move { + let session = crate::test::new_session().with_handle(handle); + let execution = + PlanExecutionContext::new(Arc::new(TestSegments::default()), session.clone()) + .with_row_offset(100); + let plan = ConcatPlan::try_new( + row_idx_dtype(), + vec![ + RowIdxPlan::new(2).into_plan(), + RowIdxPlan::new(3).into_plan(), + ], + )? + .into_plan(); + + let actual = plan + .execute(&execution, &(1..4), MaskFuture::new_true(3))? + .await?; + let expected = PrimitiveArray::from_iter([101_u64, 102, 103]).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +} + #[test] fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { let value_dtype = primitive(PType::I32, Nullability::NonNullable); @@ -1269,3 +1328,68 @@ fn legacy_stats_layout_uses_zoned_plan() -> VortexResult<()> { "); Ok(()) } + +#[test] +fn multi_field_struct_expression_does_not_read_unused_fields() -> VortexResult<()> { + block_on(|handle| async move { + let session = crate::test::new_session().with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let strategy = FlatLayoutStrategy::default(); + + let (a_sequence, a_eof) = SequenceId::root().split(); + let a = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + PrimitiveArray::from_iter([1_i32, 6, 8]) + .into_array() + .to_array_stream() + .sequenced(a_sequence), + a_eof, + &session, + ) + .await?; + let (b_sequence, b_eof) = SequenceId::root().split(); + let b = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + PrimitiveArray::from_iter([10_i32, 8, 9]) + .into_array() + .to_array_stream() + .sequenced(b_sequence), + b_eof, + &session, + ) + .await?; + + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([ + ("a", value_dtype.clone()), + ("b", value_dtype.clone()), + ("c", value_dtype.clone()), + ]), + Nullability::NonNullable, + ), + vec![a, b, flat(3, value_dtype, 2)], + ) + .into_layout(); + let expression = and( + gt(get_item("a", root()), lit(5_i32)), + gt(get_item("b", root()), lit(7_i32)), + ); + let optimized = optimize(make_eval(expression, make_plan(layout)?)?.into_plan())?; + let execution = PlanExecutionContext::new(segments, session.clone()); + + let actual = optimized + .execute(&execution, &(0..3), MaskFuture::new_true(3))? + .await?; + let expected = BoolArray::from_iter([false, true, true]).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +} diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/typed.rs index bf544801263..0c74aaf951c 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/typed.rs @@ -9,15 +9,19 @@ use std::fmt::Display; use std::fmt::Formatter; use std::marker::PhantomData; use std::ops::Deref; +use std::ops::Range; use std::sync::Arc; +use vortex_array::MaskFuture; use vortex_array::SerializeMetadata; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanVTable; use crate::plan::display::PlanTreeDisplay; @@ -105,6 +109,16 @@ impl PlanRef { self.dyn_plan().dyn_metadata(self) } + /// Executes this plan over `row_range`, returning the values selected by `mask`. + pub fn execute( + &self, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + self.dyn_plan().dyn_execute(self, ctx, row_range, mask) + } + /// Returns whether this plan uses vtable `V`. pub fn is(&self) -> bool { self.dyn_plan().as_any().is::>() @@ -337,6 +351,15 @@ pub trait DynPlan: 'static + Send + Sync + Debug { /// Serializes operator-specific metadata, or `None` when the operator is not serializable. fn dyn_metadata(&self, plan: &PlanRef) -> Option>; + + /// Executes this plan over `row_range`, returning the values selected by `mask`. + fn dyn_execute( + &self, + plan: &PlanRef, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult; } impl DynPlan for PlanData { @@ -368,4 +391,14 @@ impl DynPlan for PlanData { fn dyn_metadata(&self, plan: &PlanRef) -> Option> { V::metadata(plan.as_::()).map(SerializeMetadata::serialize) } + + fn dyn_execute( + &self, + plan: &PlanRef, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + V::execute(plan.as_::(), ctx, row_range, mask) + } } diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/vtable.rs index caf004383b5..dd9f3169e36 100644 --- a/vortex-layout/src/plan/vtable.rs +++ b/vortex-layout/src/plan/vtable.rs @@ -4,13 +4,18 @@ use std::borrow::Cow; use std::fmt; use std::fmt::Debug; +use std::ops::Range; use vortex_array::DeserializeMetadata; +use vortex_array::MaskFuture; use vortex_array::SerializeMetadata; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_session::registry::Id; +use crate::plan::PlanArrayFuture; use crate::plan::PlanChildren; +use crate::plan::PlanExecutionContext; use crate::plan::typed::Plan; /// A unique identifier for a plan operator. @@ -63,6 +68,23 @@ pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { Ok(()) } + /// Executes this operator over `row_range`, returning the values selected by `mask`. + /// + /// The row range is expressed in this plan's row domain. The returned array has one row for + /// every true value in `mask`. + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + drop((ctx, row_range, mask)); + vortex_bail!( + "Plan execution is not implemented for '{}'", + plan.vtable().id() + ) + } + /// Returns the display name of the child at `index`. fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { let _ = plan; diff --git a/vortex-scan-v2/Cargo.toml b/vortex-scan-v2/Cargo.toml new file mode 100644 index 00000000000..8f4e97cad0e --- /dev/null +++ b/vortex-scan-v2/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "vortex-scan-v2" +authors.workspace = true +description = "Plan-native scanning for Vortex layouts" +edition = { workspace = true } +homepage = { workspace = true } +categories = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[dependencies] +futures = { workspace = true, features = ["alloc", "async-await"] } +itertools = { workspace = true } +tracing = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-io = { workspace = true } +vortex-layout = { workspace = true } +vortex-mask = { workspace = true } +vortex-scan = { workspace = true } +vortex-session = { workspace = true } +vortex-utils = { workspace = true } + +[dev-dependencies] +tracing-subscriber = { workspace = true, features = ["env-filter"] } +vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-file = { workspace = true, features = ["tokio"] } +vortex-io = { workspace = true, features = ["tokio"] } +vortex-layout = { workspace = true, features = ["_test-harness"] } + +[lints] +workspace = true diff --git a/vortex-scan-v2/examples/tpch_scan.rs b/vortex-scan-v2/examples/tpch_scan.rs new file mode 100644 index 00000000000..67b00e1102d --- /dev/null +++ b/vortex-scan-v2/examples/tpch_scan.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::env; +use std::path::PathBuf; + +use tracing_subscriber::EnvFilter; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::assert_arrays_eq; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::expr::select; +use vortex_array::stream::ArrayStreamExt; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_io::runtime::single::block_on; +use vortex_io::session::RuntimeSession; +use vortex_io::session::RuntimeSessionExt; +use vortex_layout::session::LayoutSession; +use vortex_scan_v2::ScanBuilder; + +fn main() -> VortexResult<()> { + let filter = EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("vortex_scan_v2=debug")); + tracing_subscriber::fmt() + .with_env_filter(filter) + .with_target(true) + .without_time() + .init(); + + let path = env::args_os().nth(1).map_or_else( + || PathBuf::from("vortex-bench/data/tpch/0.01/vortex-file-compressed/lineitem.vortex"), + PathBuf::from, + ); + + block_on(|handle| async move { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + vortex_file::register_default_encodings(&session); + + let file = session.open_options().open_path(&path).await?; + println!( + "opened {}: rows={}, dtype={}", + path.display(), + file.row_count(), + file.dtype() + ); + + let filter = gt(get_item("l_linenumber", root()), lit(5_i32)); + let projection = select(["l_orderkey", "l_linenumber"], root()); + let result = ScanBuilder::try_new( + file.footer().layout(), + file.segment_source(), + session.clone(), + )? + .with_filter(filter.clone()) + .with_projection(projection.clone()) + .into_array_stream()? + .read_all() + .await?; + + println!( + "scan result: rows={}, dtype={}", + result.len(), + result.dtype() + ); + let expected = file + .scan()? + .with_filter(filter.bind(file.dtype())?) + .with_projection(projection.bind(file.dtype())?) + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(result, expected, &mut session.create_execution_ctx()); + println!("validated every result value against the LayoutReader scan"); + Ok(()) + }) +} diff --git a/vortex-scan-v2/src/lib.rs b/vortex-scan-v2/src/lib.rs new file mode 100644 index 00000000000..52249d77cee --- /dev/null +++ b/vortex-scan-v2/src/lib.rs @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Plan-native scanning for Vortex layouts. +//! +//! This crate intentionally owns a separate copy of the scan orchestration. It executes +//! [`vortex_layout::plan::Plan`] trees and never constructs a +//! [`vortex_layout::LayoutReader`]. +//! +//! Set `RUST_LOG=vortex_scan_v2=debug` to log source and optimized plan trees and selected scan +//! splits. Use `trace` to also log execution of each split. + +mod repeated_scan; +mod scan_builder; +mod splits; +mod tasks; + +#[cfg(test)] +mod tests; + +pub use repeated_scan::RepeatedScan; +pub use scan_builder::ScanBuilder; +pub use splits::SplitBy; diff --git a/vortex-scan-v2/src/repeated_scan.rs b/vortex-scan-v2/src/repeated_scan.rs new file mode 100644 index 00000000000..8e51097c78e --- /dev/null +++ b/vortex-scan-v2/src/repeated_scan.rs @@ -0,0 +1,210 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cmp; +use std::iter; +use std::ops::Range; +use std::sync::Arc; + +use futures::Stream; +use futures::future::BoxFuture; +use itertools::Either; +use itertools::Itertools; +use vortex_array::ArrayRef; +use vortex_array::dtype::DType; +use vortex_array::iter::ArrayIterator; +use vortex_array::iter::ArrayIteratorAdapter; +use vortex_array::stream::ArrayStream; +use vortex_array::stream::ArrayStreamAdapter; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_io::runtime::BlockingRuntime; +use vortex_io::session::RuntimeSessionExt; +use vortex_layout::plan::PlanExecutionContext; +use vortex_layout::plan::PlanRef; +use vortex_scan::selection::Selection; +use vortex_utils::parallelism::get_available_parallelism; + +use crate::splits::Splits; +use crate::tasks::TaskContext; +use crate::tasks::split_exec; + +/// A prepared plan-native scan that can be executed repeatedly over narrower row ranges. +pub struct RepeatedScan { + execution: PlanExecutionContext, + projection: PlanRef, + filter: Option, + ordered: bool, + row_range: Option>, + selection: Selection, + splits: Splits, + concurrency: usize, + map_fn: Arc VortexResult + Send + Sync>, + limit: Option, + dtype: DType, +} + +impl RepeatedScan { + /// Returns the dtype produced by this scan. + pub fn dtype(&self) -> &DType { + &self.dtype + } + + /// Executes the scan as a blocking array iterator. + pub fn execute_array_iter( + &self, + row_range: Option>, + runtime: &B, + ) -> VortexResult { + let dtype = self.dtype.clone(); + let stream = self.execute_stream(row_range)?; + Ok(ArrayIteratorAdapter::new( + dtype, + runtime.block_on_stream(stream), + )) + } + + /// Executes the scan as an asynchronous array stream. + pub fn execute_array_stream( + &self, + row_range: Option>, + ) -> VortexResult { + let dtype = self.dtype.clone(); + let stream = self.execute_stream(row_range)?; + Ok(ArrayStreamAdapter::new(dtype, stream)) + } +} + +impl RepeatedScan { + #[expect(clippy::too_many_arguments, reason = "scan construction state")] + pub(crate) fn new( + execution: PlanExecutionContext, + projection: PlanRef, + filter: Option, + ordered: bool, + row_range: Option>, + selection: Selection, + splits: Splits, + concurrency: usize, + map_fn: Arc VortexResult + Send + Sync>, + limit: Option, + ) -> Self { + let dtype = projection.dtype().clone(); + Self { + execution, + projection, + filter, + ordered, + row_range, + selection, + splits, + concurrency, + map_fn, + limit, + dtype, + } + } + + /// Constructs one execution future per selected row split. + pub fn execute( + &self, + row_range: Option>, + ) -> VortexResult>>>> { + let selection_range = match &self.selection { + Selection::IncludeByIndex(indices) if !indices.is_empty() => { + Some(indices[0]..indices[indices.len() - 1] + 1) + } + Selection::IncludeRoaring(indices) if !indices.is_empty() => Some( + indices.min().vortex_expect("non-empty selection") + ..indices.max().vortex_expect("non-empty selection") + 1, + ), + _ => None, + }; + let row_range = intersect_ranges(self.row_range.as_ref(), row_range); + let row_range = intersect_ranges(row_range.as_ref(), selection_range); + + let ranges = match &self.splits { + Splits::Natural(boundaries) => { + let boundaries = match row_range { + None => Either::Left(boundaries.iter().copied()), + Some(range) => { + if range.is_empty() { + return Ok(Vec::new()); + } + let start = boundaries.partition_point(|&point| point < range.start); + let end = boundaries.partition_point(|&point| point < range.end); + Either::Right( + iter::once(range.start) + .chain(boundaries[start..end].iter().copied()) + .chain(iter::once(range.end)), + ) + } + }; + Either::Left(boundaries.tuple_windows().map(|(start, end)| start..end)) + } + Splits::Ranges(ranges) => Either::Right(match row_range { + None => Either::Left(ranges.iter().cloned()), + Some(range) => { + if range.is_empty() { + return Ok(Vec::new()); + } + Either::Right(ranges.iter().filter_map(move |candidate| { + let start = cmp::max(candidate.start, range.start); + let end = cmp::min(candidate.end, range.end); + (start < end).then_some(start..end) + })) + } + }), + }; + + let ctx = Arc::new(TaskContext { + execution: self.execution.clone(), + filter: self.filter.clone(), + projection: self.projection.clone(), + mapper: Arc::clone(&self.map_fn), + }); + let mut limit = self.limit; + let mut tasks = Vec::new(); + for range in ranges { + let row_mask = self.selection.row_mask(&range); + if row_mask.mask().all_false() { + continue; + } + tasks.push(split_exec(Arc::clone(&ctx), row_mask, limit.as_mut())?); + if limit.is_some_and(|limit| limit == 0) { + break; + } + } + Ok(tasks) + } + + /// Executes all selected row splits with the configured ordering and concurrency. + pub fn execute_stream( + &self, + row_range: Option>, + ) -> VortexResult> + Send + 'static + use> { + use futures::StreamExt; + + let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); + let handle = self.execution.session().handle(); + let stream = + futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task)); + let stream = if self.ordered { + stream.buffered(concurrency).boxed() + } else { + stream.buffer_unordered(concurrency).boxed() + }; + Ok(stream.filter_map(|chunk| async move { chunk.transpose() })) + } +} + +fn intersect_ranges(left: Option<&Range>, right: Option>) -> Option> { + match (left, right) { + (None, None) => None, + (None, Some(right)) => Some(right), + (Some(left), None) => Some(left.clone()), + (Some(left), Some(right)) => { + Some(cmp::max(left.start, right.start)..cmp::min(left.end, right.end)) + } + } +} diff --git a/vortex-scan-v2/src/scan_builder.rs b/vortex-scan-v2/src/scan_builder.rs new file mode 100644 index 00000000000..99c7630146e --- /dev/null +++ b/vortex-scan-v2/src/scan_builder.rs @@ -0,0 +1,449 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::ready; + +use futures::Stream; +use futures::StreamExt; +use futures::future::BoxFuture; +use futures::stream::BoxStream; +use vortex_array::ArrayRef; +use vortex_array::dtype::DType; +use vortex_array::expr::Expression; +use vortex_array::expr::root; +use vortex_array::iter::ArrayIterator; +use vortex_array::iter::ArrayIteratorAdapter; +use vortex_array::stream::ArrayStream; +use vortex_array::stream::ArrayStreamAdapter; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_io::runtime::BlockingRuntime; +use vortex_io::runtime::Handle; +use vortex_io::runtime::Task; +use vortex_io::session::RuntimeSessionExt; +use vortex_layout::LayoutRef; +use vortex_layout::plan::PlanExecutionContext; +use vortex_layout::plan::PlanRef; +use vortex_layout::plan::lower; +use vortex_layout::plan::optimize; +use vortex_layout::plan::plan_row_idx_expression; +use vortex_layout::segments::SegmentSource; +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::RepeatedScan; +use crate::splits::SplitBy; +use crate::splits::Splits; +use crate::splits::attempt_split_ranges; + +/// Builds a plan-native scan without constructing a layout reader. +pub struct ScanBuilder { + execution: PlanExecutionContext, + base_plan: PlanRef, + projection: Expression, + filter: Option, + ordered: bool, + row_range: Option>, + selection: Selection, + split_by: SplitBy, + concurrency: usize, + map_fn: Arc VortexResult + Send + Sync>, + limit: Option, + row_offset: u64, +} + +impl ScanBuilder { + /// Creates a plan-native scan directly from a stored layout. + pub fn try_new( + layout: &LayoutRef, + segment_source: Arc, + session: VortexSession, + ) -> VortexResult { + tracing::debug!( + target: "vortex_scan_v2::planner", + layout = %layout.display_tree(), + "building a plan-native scan from a layout" + ); + let plan = lower(layout)?; + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %plan.display_tree(), + "constructed the source physical plan" + ); + Ok(Self::from_plan( + plan, + PlanExecutionContext::new(segment_source, session), + )) + } + + /// Creates a scan from an already constructed physical plan. + pub fn from_plan(base_plan: PlanRef, execution: PlanExecutionContext) -> Self { + Self { + execution, + base_plan, + projection: root(), + filter: None, + ordered: true, + row_range: None, + selection: Selection::default(), + split_by: SplitBy::default(), + concurrency: 4, + map_fn: Arc::new(Ok), + limit: None, + row_offset: 0, + } + } + + /// Returns an asynchronous stream of Vortex arrays. + pub fn into_array_stream(self) -> VortexResult { + let dtype = self.dtype()?; + Ok(ArrayStreamAdapter::new(dtype, self.into_stream()?)) + } + + /// Returns a blocking iterator of Vortex arrays. + pub fn into_array_iter( + self, + runtime: &B, + ) -> VortexResult { + let stream = self.into_array_stream()?; + let dtype = stream.dtype().clone(); + Ok(ArrayIteratorAdapter::new( + dtype, + runtime.block_on_stream(stream), + )) + } +} + +impl ScanBuilder { + /// Sets the filter expression. + pub fn with_filter(mut self, filter: Expression) -> Self { + self.filter = Some(filter); + self + } + + /// Sets or clears the filter expression. + pub fn with_some_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Sets the projection expression. + pub fn with_projection(mut self, projection: Expression) -> Self { + self.projection = projection; + self + } + + /// Returns whether output splits retain row order. + pub fn ordered(&self) -> bool { + self.ordered + } + + /// Configures whether output splits retain row order. + pub fn with_ordered(mut self, ordered: bool) -> Self { + self.ordered = ordered; + self + } + + /// Restricts the scan to a contiguous row range. + pub fn with_row_range(mut self, row_range: Range) -> Self { + self.row_range = Some(row_range); + self + } + + /// Applies an additional row selection. + pub fn with_selection(mut self, selection: Selection) -> Self { + self.selection = selection; + self + } + + /// Selects strictly sorted absolute row indices relative to the scan input. + pub fn with_row_indices(mut self, row_indices: StrictSortedBuffer) -> Self { + self.selection = Selection::IncludeByIndex(row_indices); + self + } + + /// Sets the global offset used by row-index expressions. + pub fn with_row_offset(mut self, row_offset: u64) -> Self { + self.row_offset = row_offset; + self + } + + /// Configures how scan work is split into tasks. + pub fn with_split_by(mut self, split_by: SplitBy) -> Self { + self.split_by = split_by; + self + } + + /// Returns the per-worker split concurrency. + pub fn concurrency(&self) -> usize { + self.concurrency + } + + /// Sets the per-worker split concurrency. + pub fn with_concurrency(mut self, concurrency: usize) -> Self { + assert!(concurrency > 0, "scan concurrency must be non-zero"); + self.concurrency = concurrency; + self + } + + /// Sets a maximum number of output rows. + pub fn with_limit(mut self, limit: u64) -> Self { + self.limit = Some(limit); + self + } + + /// Sets or clears the maximum number of output rows. + pub fn with_some_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + /// Returns the dtype produced by the projection expression. + pub fn dtype(&self) -> VortexResult { + self.projection.return_dtype(self.base_plan.dtype()) + } + + /// Returns the session used by plan execution. + pub fn session(&self) -> &VortexSession { + self.execution.session() + } + + /// Maps every output array into another result type. + pub fn map( + self, + map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, + ) -> ScanBuilder { + let old_map_fn = self.map_fn; + ScanBuilder { + execution: self.execution, + base_plan: self.base_plan, + projection: self.projection, + filter: self.filter, + ordered: self.ordered, + row_range: self.row_range, + selection: self.selection, + split_by: self.split_by, + concurrency: self.concurrency, + map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)), + limit: self.limit, + row_offset: self.row_offset, + } + } + + /// Constructs and optimizes the projection and filter plans. + pub fn prepare(self) -> VortexResult> { + if self.filter.is_some() && self.limit.is_some() { + vortex_bail!("Vortex doesn't support scans with both a filter and a limit") + } + + let source = self.base_plan.clone(); + tracing::debug!( + target: "vortex_scan_v2::planner", + row_offset = self.row_offset, + plan = %source.display_tree(), + "planning expressions over the source and its row-index domain" + ); + let projection = optimize_projection_plan(self.projection, &source)?; + let filter = optimize_filter_plan(self.filter, &source)?; + + let splits = + if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) { + Splits::Ranges(ranges) + } else { + let row_range = self + .row_range + .clone() + .unwrap_or_else(|| 0..self.base_plan.row_count()); + let mut plans = vec![&projection]; + plans.extend(filter.as_ref()); + Splits::Natural(self.split_by.splits(&plans, &row_range)?) + }; + match &splits { + Splits::Natural(boundaries) => tracing::debug!( + target: "vortex_scan_v2::planner", + split_count = boundaries.len().saturating_sub(1), + ?boundaries, + "selected natural plan scan splits" + ), + Splits::Ranges(ranges) => tracing::debug!( + target: "vortex_scan_v2::planner", + split_count = ranges.len(), + ?ranges, + "selected sparse plan scan ranges" + ), + } + + Ok(RepeatedScan::new( + self.execution.with_row_offset(self.row_offset), + projection, + filter, + self.ordered, + self.row_range, + self.selection, + splits, + self.concurrency, + self.map_fn, + self.limit, + )) + } + + /// Builds one future per scan split. + pub fn build(self) -> VortexResult>>>> { + if self.limit.is_some_and(|limit| limit == 0) { + return Ok(Vec::new()); + } + self.prepare()?.execute(None) + } + + /// Returns an asynchronous stream that schedules scan splits on the session runtime. + pub fn into_stream( + self, + ) -> VortexResult> + Send + 'static + use> { + Ok(LazyScanStream::new(self)) + } + + /// Returns a blocking iterator over mapped scan outputs. + pub fn into_iter( + self, + runtime: &B, + ) -> VortexResult> + 'static> { + Ok(runtime.block_on_stream(self.into_stream()?)) + } +} + +fn optimize_projection_plan(expression: Expression, source: &PlanRef) -> VortexResult { + tracing::debug!( + target: "vortex_scan_v2::planner", + %expression, + "optimizing the projection expression" + ); + let expression = expression + .optimize_recursive(source.dtype())? + .bind(source.dtype())?; + let projection = optimize(plan_row_idx_expression(expression, source.clone())?)?; + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %projection.display_tree(), + "optimized the projection physical plan" + ); + Ok(projection) +} + +fn optimize_filter_plan( + filter: Option, + source: &PlanRef, +) -> VortexResult> { + let filter = filter + .map(|expression| -> VortexResult { + tracing::debug!( + target: "vortex_scan_v2::planner", + %expression, + "optimizing the filter expression" + ); + let expression = expression + .optimize_recursive(source.dtype())? + .bind(source.dtype())?; + let filter = optimize(plan_row_idx_expression(expression, source.clone())?)?; + vortex_ensure!( + filter.dtype().is_boolean(), + "Filter plan must produce booleans" + ); + Ok(filter) + }) + .transpose()?; + if let Some(filter) = &filter { + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %filter.display_tree(), + "optimized the filter physical plan" + ); + } + Ok(filter) +} + +enum LazyScanState { + Builder(Option>>), + Preparing(PreparingScan), + Stream(BoxStream<'static, VortexResult>), + Error(Option), +} + +type PreparedScanTasks = Vec>>>; + +struct PreparingScan { + ordered: bool, + concurrency: usize, + handle: Handle, + task: Task>>, +} + +struct LazyScanStream { + state: LazyScanState, +} + +impl LazyScanStream { + fn new(builder: ScanBuilder) -> Self { + Self { + state: LazyScanState::Builder(Some(Box::new(builder))), + } + } +} + +impl Unpin for LazyScanStream {} + +impl Stream for LazyScanStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match &mut self.state { + LazyScanState::Builder(builder) => { + let builder = builder.take().vortex_expect("polled after completion"); + let ordered = builder.ordered; + let concurrency = + builder.concurrency * get_available_parallelism().unwrap_or(1); + let handle = builder.execution.session().handle(); + let task = handle + .spawn_cpu(move || builder.prepare().and_then(|scan| scan.execute(None))); + self.state = LazyScanState::Preparing(PreparingScan { + ordered, + concurrency, + handle, + task, + }); + } + LazyScanState::Preparing(preparing) => { + match ready!(Pin::new(&mut preparing.task).poll(cx)) { + Ok(tasks) => { + let handle = preparing.handle.clone(); + let stream = + futures::stream::iter(tasks).map(move |task| handle.spawn(task)); + let stream = if preparing.ordered { + stream.buffered(preparing.concurrency).boxed() + } else { + stream.buffer_unordered(preparing.concurrency).boxed() + }; + self.state = LazyScanState::Stream( + stream + .filter_map(|chunk| async move { chunk.transpose() }) + .boxed(), + ); + } + Err(error) => self.state = LazyScanState::Error(Some(error)), + } + } + LazyScanState::Stream(stream) => return stream.as_mut().poll_next(cx), + LazyScanState::Error(error) => return Poll::Ready(error.take().map(Err)), + } + } + } +} diff --git a/vortex-scan-v2/src/splits.rs b/vortex-scan-v2/src/splits.rs new file mode 100644 index 00000000000..b0f6111e399 --- /dev/null +++ b/vortex-scan-v2/src/splits.rs @@ -0,0 +1,224 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter::once; +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_layout::plan::Concat; +use vortex_layout::plan::Eval; +use vortex_layout::plan::Pack; +use vortex_layout::plan::PlanRef; +use vortex_layout::plan::Take; +use vortex_layout::plan::Zoned; +use vortex_scan::selection::Selection; + +const IDEAL_SPLIT_SIZE: u64 = 100_000; +const MAX_RANGE_SIZE: u64 = IDEAL_SPLIT_SIZE / 25; +const MIN_GAP_BETWEEN_RANGES: u64 = IDEAL_SPLIT_SIZE / 2; + +/// Defines how a plan scan is divided into independently executable row ranges. +#[derive(Default, Copy, Clone, Debug)] +pub enum SplitBy { + /// Uses boundaries exposed by the optimized physical plan. + #[default] + Layout, + /// Splits every `n` rows. + RowCount(usize), +} + +impl SplitBy { + pub(crate) fn splits( + &self, + plans: &[&PlanRef], + row_range: &Range, + ) -> VortexResult> { + let mut boundaries = match *self { + Self::Layout => { + let mut boundaries = vec![row_range.start]; + for plan in plans { + collect_plan_splits(plan, 0, row_range, &mut boundaries)?; + } + boundaries + } + Self::RowCount(row_count) => { + vortex_ensure!(row_count > 0, "Row-count split size must be non-zero"); + row_range + .clone() + .step_by(row_count) + .chain(once(row_range.end)) + .collect() + } + }; + boundaries.sort_unstable(); + boundaries.dedup(); + Ok(subdivide_large_spans(boundaries, IDEAL_SPLIT_SIZE)) + } +} + +fn collect_plan_splits( + plan: &PlanRef, + row_offset: u64, + row_range: &Range, + boundaries: &mut Vec, +) -> VortexResult<()> { + if plan.is::() || plan.is::() { + if let Some(child) = plan.child(0)? { + collect_plan_splits(&child, row_offset, row_range, boundaries)?; + } + return Ok(()); + } + + if plan.is::() { + if let Some(codes) = plan.child(0)? { + collect_plan_splits(&codes, row_offset, row_range, boundaries)?; + } + return Ok(()); + } + + // A childless Pack preserves row count even though it exposes no leaf boundaries. + if plan.is::() && plan.child_count() == 0 { + boundaries.push(row_offset + row_range.end); + return Ok(()); + } + + if plan.is::() { + for index in 0..plan.child_count() { + if let Some(child) = plan.child(index)? + && child.row_count() == plan.row_count() + { + collect_plan_splits(&child, row_offset, row_range, boundaries)?; + } + } + return Ok(()); + } + + if plan.is::() { + let mut chunk_offset = 0_u64; + for index in 0..plan.child_count() { + let Some(chunk) = plan.child(index)? else { + continue; + }; + let chunk_end = chunk_offset + .checked_add(chunk.row_count()) + .ok_or_else(|| vortex_error::vortex_err!("Chunk row offset overflow"))?; + let start = row_range.start.max(chunk_offset); + let end = row_range.end.min(chunk_end); + if start < end { + let child_range = start - chunk_offset..end - chunk_offset; + collect_plan_splits(&chunk, row_offset + chunk_offset, &child_range, boundaries)?; + boundaries.push(row_offset + end); + } + chunk_offset = chunk_end; + } + return Ok(()); + } + + boundaries.push(row_offset + row_range.end); + Ok(()) +} + +fn subdivide_large_spans(boundaries: Vec, max_span: u64) -> Vec { + if boundaries.len() < 2 + || boundaries + .windows(2) + .all(|window| window[1] - window[0] <= max_span) + { + return boundaries; + } + + let mut output = Vec::with_capacity(boundaries.len() * 2); + for window in boundaries.windows(2) { + let start = window[0]; + let end = window[1]; + output.push(start); + let span = end - start; + if span > max_span { + let split_count = span.div_ceil(max_span); + let split_size = span.div_ceil(split_count); + let mut point = start + split_size; + while point < end { + output.push(point); + point = point.saturating_add(split_size); + } + } + } + if let Some(&last) = boundaries.last() { + output.push(last); + } + output +} + +pub(crate) enum Splits { + Natural(Vec), + Ranges(Vec>), +} + +pub(crate) fn attempt_split_ranges( + selection: &Selection, + row_range: Option<&Range>, +) -> Option>> { + let Selection::IncludeByIndex(buffer) = selection else { + return None; + }; + if row_range.is_some() { + return None; + } + let indices = buffer.as_slice(); + if indices.is_empty() { + return Some(Vec::new()); + } + + let mut ranges = Vec::with_capacity((indices.len() as u64 / MAX_RANGE_SIZE) as usize); + let mut current_start = indices[0]; + let mut current_end = indices[0] + 1; + for &index in &indices[1..] { + let new_range_size = (index + 1) - current_start; + let gap = (index + 1) - current_end; + if new_range_size >= MAX_RANGE_SIZE { + if gap < MIN_GAP_BETWEEN_RANGES { + return None; + } + ranges.push(current_start..current_end); + current_start = index; + } + current_end = index + 1; + } + ranges.push(current_start..current_end); + Some(ranges) +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::Nullability; + use vortex_array::dtype::StructFields; + use vortex_error::VortexResult; + use vortex_layout::plan::PackPlan; + use vortex_layout::plan::RowIdxPlan; + + use super::SplitBy; + + #[test] + fn childless_pack_preserves_its_row_count() -> VortexResult<()> { + let plan = PackPlan::try_new( + StructFields::empty(), + Nullability::NonNullable, + 6, + Vec::new(), + None, + )? + .into_plan(); + + assert_eq!(SplitBy::Layout.splits(&[&plan], &(0..6))?, vec![0, 6]); + Ok(()) + } + + #[test] + fn row_idx_source_preserves_its_row_count() -> VortexResult<()> { + let plan = RowIdxPlan::new(6).into_plan(); + + assert_eq!(SplitBy::Layout.splits(&[&plan], &(0..6))?, vec![0, 6]); + Ok(()) + } +} diff --git a/vortex-scan-v2/src/tasks.rs b/vortex-scan-v2/src/tasks.rs new file mode 100644 index 00000000000..e3fecf928a8 --- /dev/null +++ b/vortex-scan-v2/src/tasks.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex_array::ArrayRef; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_error::VortexResult; +use vortex_layout::plan::PlanExecutionContext; +use vortex_layout::plan::PlanRef; +use vortex_mask::Mask; +use vortex_scan::row_mask::RowMask; + +pub(crate) type TaskFuture = BoxFuture<'static, VortexResult>; + +pub(crate) fn split_exec( + ctx: Arc>, + read_mask: RowMask, + limit: Option<&mut u64>, +) -> VortexResult>> { + let row_range = read_mask.row_range(); + let row_mask = read_mask.mask().clone(); + tracing::trace!( + target: "vortex_scan_v2::execution", + ?row_range, + selected_rows = row_mask.true_count(), + has_filter = ctx.filter.is_some(), + "executing a plan scan split" + ); + + let filter_mask = match &ctx.filter { + None => { + let row_mask = match limit { + Some(limit) if *limit == 0 => Mask::new_false(row_mask.len()), + Some(limit) => { + let true_count = row_mask.true_count(); + let mask_limit = usize::try_from(*limit) + .map(|limit| limit.min(true_count)) + .unwrap_or(true_count); + let row_mask = row_mask.limit(mask_limit); + *limit -= mask_limit as u64; + row_mask + } + None => row_mask, + }; + MaskFuture::ready(row_mask) + } + Some(filter) => { + let predicate = filter.execute( + &ctx.execution, + &row_range, + MaskFuture::ready(row_mask.clone()), + )?; + let session = ctx.execution.session().clone(); + MaskFuture::new(row_mask.len(), async move { + let predicate = predicate.await?; + let mut execution = session.create_execution_ctx(); + let predicate = predicate.null_as_false().execute(&mut execution)?; + Ok(row_mask.intersect_by_rank(&predicate)) + }) + } + }; + + let projection = ctx + .projection + .execute(&ctx.execution, &row_range, filter_mask.clone())?; + let mapper = Arc::clone(&ctx.mapper); + Ok(async move { + if filter_mask.await?.all_false() { + tracing::trace!( + target: "vortex_scan_v2::execution", + ?row_range, + "plan scan split produced no matching rows" + ); + return Ok(None); + } + let array = projection.await?; + tracing::trace!( + target: "vortex_scan_v2::execution", + ?row_range, + output_rows = array.len(), + dtype = %array.dtype(), + "completed a plan scan split" + ); + mapper(array).map(Some) + } + .boxed()) +} + +pub(crate) struct TaskContext { + pub(crate) execution: PlanExecutionContext, + pub(crate) filter: Option, + pub(crate) projection: PlanRef, + pub(crate) mapper: Arc VortexResult + Send + Sync>, +} diff --git a/vortex-scan-v2/src/tests.rs b/vortex-scan-v2/src/tests.rs new file mode 100644 index 00000000000..155675ce216 --- /dev/null +++ b/vortex-scan-v2/src/tests.rs @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_array::ArrayContext; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::assert_arrays_eq; +use vortex_array::expr::and; +use vortex_array::expr::checked_add; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::stream::ArrayStreamExt; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_io::runtime::single::block_on; +use vortex_io::session::RuntimeSession; +use vortex_io::session::RuntimeSessionExt; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::row_idx::row_idx; +use vortex_layout::layouts::table::TableStrategy; +use vortex_layout::segments::TestSegments; +use vortex_layout::sequence::SequenceId; +use vortex_layout::sequence::SequentialArrayStreamExt; +use vortex_layout::session::LayoutSession; +use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; + +use crate::ScanBuilder; +use crate::SplitBy; + +#[test] +fn scans_layout_through_optimized_plans() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = PrimitiveArray::from_iter(0_i32..10).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + + let actual = ScanBuilder::try_new(&layout, segments, session.clone())? + .with_filter(gt(root(), lit(4_i32))) + .with_projection(checked_add(root(), lit(1_i32))) + .with_split_by(SplitBy::RowCount(3)) + .into_array_stream()? + .read_all() + .await?; + let expected = PrimitiveArray::from_iter(6_i32..11).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +} + +#[test] +fn scans_row_idx_and_struct_expression_partitions() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = StructArray::from_fields( + [ + ("a", buffer![1_i32, 6, 7, 8, 9, 2].into_array()), + ("b", buffer![10_i32, 20, 3, 40, 5, 60].into_array()), + ] + .as_slice(), + )? + .into_array(); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let strategy = TableStrategy::new(Arc::clone(&flat), flat); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + + let filter = and( + gt(row_idx(), lit(102_u64)), + and( + gt(get_item("a", root()), lit(5_i32)), + gt(get_item("b", root()), lit(10_i32)), + ), + ); + let actual = ScanBuilder::try_new(&layout, segments, session.clone())? + .with_row_offset(100) + .with_filter(filter) + .with_projection(row_idx()) + .with_split_by(SplitBy::RowCount(2)) + .into_array_stream()? + .read_all() + .await?; + let expected = PrimitiveArray::from_iter([103_u64]).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +} + +#[test] +fn scans_selected_rows_from_a_list_plan() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = ListArray::try_new( + buffer![1_i32, 2, 3, 4, 5, 6].into_array(), + buffer![0_u32, 2, 2, 5, 6].into_array(), + Validity::NonNullable, + )? + .into_array(); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let strategy = TableStrategy::new(Arc::clone(&flat), flat).with_list_layout(); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + + let actual = ScanBuilder::try_new(&layout, segments, session.clone())? + .with_row_indices(StrictSortedBuffer::try_new(buffer![1_u64, 3])?) + .into_array_stream()? + .read_all() + .await?; + let expected = ListArray::try_new( + buffer![6_i32].into_array(), + buffer![0_u32, 0, 1].into_array(), + Validity::NonNullable, + )? + .into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) + }) +}