From 2c281ff3e5983dc39557fb178a7e748544619b76 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 7 Aug 2026 12:49:49 +0000 Subject: [PATCH 1/2] Add plan-native scan execution Give `PlanVTable` an `execute` hook taking a row range and a selection mask, and implement it for every operator: `SegmentScan` reads and decodes its segment, the structural operators combine their children, and `Eval` applies the expression to its child's output. Add `vortex-scan-v2`, which copies the existing scan orchestration around this API so the `LayoutReader` scanner is untouched while the plan-native path is developed. `Take` now records whether every dictionary value is referenced by some code. That fact previously came from the dict layout; since operators no longer hold a layout, lowering carries it into the operator's data. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_012obBhJ8oPZoBbKyeS79yMv Signed-off-by: Joe Isaacs --- Cargo.lock | 20 + Cargo.toml | 2 + .../internals/scan-planning.md | 13 +- vortex-layout/src/layouts/row_idx/mod.rs | 2 +- vortex-layout/src/plan/execution.rs | 41 ++ vortex-layout/src/plan/lower.rs | 6 +- vortex-layout/src/plan/mod.rs | 4 + vortex-layout/src/plan/plans/concat.rs | 64 +++ vortex-layout/src/plan/plans/eval.rs | 16 + vortex-layout/src/plan/plans/list_pack.rs | 126 +++++ vortex-layout/src/plan/plans/mod.rs | 1 + vortex-layout/src/plan/plans/pack.rs | 54 +++ vortex-layout/src/plan/plans/row_idx.rs | 13 + .../src/plan/plans/row_idx_partition.rs | 35 ++ .../src/plan/plans/row_idx_values.rs | 43 ++ vortex-layout/src/plan/plans/segment_scan.rs | 53 +++ vortex-layout/src/plan/plans/take.rs | 74 ++- vortex-layout/src/plan/plans/zoned.rs | 13 + vortex-layout/src/plan/tests.rs | 79 ++++ vortex-layout/src/plan/typed.rs | 31 ++ vortex-layout/src/plan/vtable.rs | 22 + vortex-scan-v2/Cargo.toml | 38 ++ vortex-scan-v2/examples/tpch_scan.rs | 83 ++++ vortex-scan-v2/src/lib.rs | 23 + vortex-scan-v2/src/repeated_scan.rs | 210 +++++++++ vortex-scan-v2/src/scan_builder.rs | 442 ++++++++++++++++++ vortex-scan-v2/src/splits.rs | 186 ++++++++ vortex-scan-v2/src/tasks.rs | 98 ++++ vortex-scan-v2/src/tests.rs | 167 +++++++ 29 files changed, 1949 insertions(+), 10 deletions(-) create mode 100644 vortex-layout/src/plan/execution.rs create mode 100644 vortex-scan-v2/Cargo.toml create mode 100644 vortex-scan-v2/examples/tpch_scan.rs create mode 100644 vortex-scan-v2/src/lib.rs create mode 100644 vortex-scan-v2/src/repeated_scan.rs create mode 100644 vortex-scan-v2/src/scan_builder.rs create mode 100644 vortex-scan-v2/src/splits.rs create mode 100644 vortex-scan-v2/src/tasks.rs create mode 100644 vortex-scan-v2/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index f54ffce4671..3e11af3875c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10559,6 +10559,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 00a45b03618..e73b79e53b3 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", @@ -324,6 +325,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 b08e2e6d9f6..fc346556c40 100644 --- a/docs/developer-guide/internals/scan-planning.md +++ b/docs/developer-guide/internals/scan-planning.md @@ -57,12 +57,17 @@ holding the common fields — dtype, row count, and children. Operator-specific 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. Lowering does not yet take a projection or row range, so it lowers the whole layout tree. Once it does, an unsupported layout in a column the query never reads will no longer fail the scan. 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..5d9a3c5829f --- /dev/null +++ b/vortex-layout/src/plan/execution.rs @@ -0,0 +1,41 @@ +// 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, +} + +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, + } + } + + /// 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 bb31b61febe..9f1f8748f51 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -123,7 +123,11 @@ fn lower_dict(layout: &DictLayout) -> VortexResult { .slot(0)? .ok_or_else(|| vortex_err!("Dictionary values child is absent"))?, )?; - Ok(TakePlan::new(codes, values)) + Ok(TakePlan::new_with_all_values_referenced( + codes, + values, + layout.has_all_values_referenced(), + )) } fn lower_list(layout: &ListLayout) -> VortexResult { diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 060b73fe74a..d6360587dc6 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -8,6 +8,7 @@ //! can reason about a plan's shape alone. [`lower`] is the one-way bridge from a stored layout. mod display; +mod execution; mod lower; mod optimize; pub mod optimizer; @@ -21,6 +22,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; @@ -49,6 +52,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::ZonedPlan; diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index fcba700a1e9..bc211f02d6a 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -2,20 +2,34 @@ // 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_array::expr::ExactBoundExpr; use vortex_array::expr::label_bound_tree; +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::layouts::row_idx::RowIdx as RowIdxFn; use crate::plan::Eval; use crate::plan::EvalPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -95,6 +109,56 @@ impl PlanVTable for Concat { ConcatPlan::try_new(plan.dtype().clone(), 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(), + "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_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)?; + chunk_futures.push(chunk.execute(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 c1e9ba11de6..8bd011035c8 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -3,8 +3,11 @@ 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::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::expr::BoundExpression; @@ -17,6 +20,8 @@ use vortex_error::VortexResult; use vortex_session::registry::CachedId; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -88,6 +93,17 @@ impl PlanVTable for Eval { )) } + 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 a5775e35cdf..10c45641e4c 100644 --- a/vortex-layout/src/plan/plans/list_pack.rs +++ b/vortex-layout/src/plan/plans/list_pack.rs @@ -2,16 +2,34 @@ // 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::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -116,6 +134,75 @@ impl PlanVTable for ListPack { ) } + 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().clone(); + 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"), @@ -125,3 +212,42 @@ impl PlanVTable for ListPack { } } } + +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 8c8dfd6feb0..ffaf3f75f82 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -43,6 +43,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 use zoned::Zoned; pub use zoned::ZonedPlan; diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index e3b39635745..709688f6c9c 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,6 +39,8 @@ use vortex_session::registry::CachedId; use crate::plan::Eval; use crate::plan::EvalPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -140,6 +149,51 @@ impl PlanVTable for Pack { ) } + 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> { if let Some(name) = plan.data().fields.field_name(index) { return Cow::Borrowed(name.as_ref()); diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index c445331f672..2dc3f56f42d 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -4,7 +4,9 @@ use std::borrow::Cow; use std::fmt::Display; use std::fmt::Formatter; +use std::ops::Range; +use vortex_array::MaskFuture; use vortex_array::ProstMetadata; use vortex_array::dtype::FieldName; use vortex_array::expr::BoundExpression; @@ -20,6 +22,8 @@ use crate::layouts::row_idx::RowIdx as RowIdxFn; use crate::plan::Eval; use crate::plan::EvalPlan; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -90,6 +94,15 @@ impl PlanVTable for RowIdx { Ok(RowIdxPlan::new(plan.data().row_offset, children.remove(0))) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + plan.child_plan().execute(ctx, row_range, mask) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { if index == 0 { Cow::Borrowed("child") diff --git a/vortex-layout/src/plan/plans/row_idx_partition.rs b/vortex-layout/src/plan/plans/row_idx_partition.rs index 1f14d5c24d5..33f4589449d 100644 --- a/vortex-layout/src/plan/plans/row_idx_partition.rs +++ b/vortex-layout/src/plan/plans/row_idx_partition.rs @@ -2,16 +2,26 @@ // 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::Nullability; use vortex_array::dtype::StructFields; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; 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::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -91,6 +101,31 @@ impl PlanVTable for RowIdxPartition { RowIdxPartitionPlan::try_new(row_idx, child) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let row_idx = plan.row_idx_plan().execute(ctx, row_range, mask.clone())?; + let child = plan.child_plan().execute(ctx, row_range, mask)?; + let names = plan + .dtype() + .as_struct_fields_opt() + .vortex_expect("RowIdxPartition dtype must be a struct") + .names() + .clone(); + Ok(async move { + let (row_idx, child) = try_join!(row_idx, child)?; + let len = child.len(); + Ok( + StructArray::try_new(names, vec![row_idx, child], len, Validity::NonNullable)? + .into_array(), + ) + } + .boxed()) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { match index { ROW_IDX => Cow::Borrowed(ROW_IDX_PARTITION_NAME), diff --git a/vortex-layout/src/plan/plans/row_idx_values.rs b/vortex-layout/src/plan/plans/row_idx_values.rs index ac2ca64a746..ff05e7f618e 100644 --- a/vortex-layout/src/plan/plans/row_idx_values.rs +++ b/vortex-layout/src/plan/plans/row_idx_values.rs @@ -1,14 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; + +use futures::FutureExt; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; use vortex_array::ProstMetadata; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::registry::CachedId; +use crate::layouts::row_idx::idx_array; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -71,6 +80,40 @@ impl PlanVTable for RowIdxValues { check_child_count("RowIdxValues", &children, 0)?; Ok(plan.clone()) } + + 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(), + "RowIdxValues row range {:?} is outside 0..{}", + row_range, + plan.row_count() + ); + vortex_ensure!( + mask.len() == usize::try_from(row_range.end - row_range.start)?, + "RowIdxValues mask length mismatch" + ); + let row_offset = plan.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()), + "RowIdxValues 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()) + } } /// Serialized metadata for a [`RowIdxValues`] plan. diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/plans/segment_scan.rs index 3b10df8fa66..52ecf02d40f 100644 --- a/vortex-layout/src/plan/plans/segment_scan.rs +++ b/vortex-layout/src/plan/plans/segment_scan.rs @@ -1,14 +1,22 @@ // 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::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -84,6 +92,51 @@ impl PlanVTable for SegmentScan { None } + 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()) + } + fn with_children(plan: &Plan, children: Vec) -> VortexResult> { check_child_count("SegmentScan", &children, 0)?; Ok(plan.clone()) diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/plans/take.rs index b4848d827a0..dcffa085be7 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/plans/take.rs @@ -2,16 +2,25 @@ // 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::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::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -26,6 +35,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; @@ -34,16 +51,32 @@ 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 { PlanParts { vtable: Take, dtype: values.dtype().clone(), row_count: codes.row_count(), children: vec![codes, values], - data: (), + data: TakeData { + all_values_referenced, + }, } .into_typed() } + /// 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. pub fn codes(&self) -> &PlanRef { &self.children()[CODES] @@ -56,7 +89,7 @@ impl TakePlan { } impl PlanVTable for Take { - type PlanData = (); + type PlanData = TakeData; type Metadata = EmptyMetadata; fn id(&self) -> PlanId { @@ -68,11 +101,44 @@ impl PlanVTable for Take { Some(EmptyMetadata) } - fn with_children(_plan: &Plan, mut children: Vec) -> VortexResult> { + fn with_children(plan: &Plan, mut children: Vec) -> VortexResult> { check_child_count("Take", &children, 2)?; let values = children.remove(VALUES); let codes = children.remove(CODES); - Ok(TakePlan::new(codes, values)) + Ok(TakePlan::new_with_all_values_referenced( + codes, + values, + plan.all_values_referenced(), + )) + } + + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let codes = plan.codes().execute(ctx, row_range, mask)?; + let values_len = usize::try_from(plan.values().row_count())?; + let values = plan.values().execute( + ctx, + &(0..plan.values().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> { diff --git a/vortex-layout/src/plan/plans/zoned.rs b/vortex-layout/src/plan/plans/zoned.rs index 15fef3a4c8e..7823f5b4a02 100644 --- a/vortex-layout/src/plan/plans/zoned.rs +++ b/vortex-layout/src/plan/plans/zoned.rs @@ -2,13 +2,17 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use vortex_array::EmptyMetadata; +use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use crate::plan::Plan; +use crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; @@ -74,6 +78,15 @@ impl PlanVTable for Zoned { Ok(ZonedPlan::new(data, zones)) } + fn execute( + plan: &Plan, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + plan.data_plan().execute(ctx, row_range, mask) + } + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { match index { DATA => Cow::Borrowed("data"), diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 6b3de94ea68..6b04026f256 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -5,7 +5,14 @@ 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::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; @@ -20,6 +27,8 @@ use vortex_array::expr::lit; 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; @@ -27,10 +36,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; @@ -38,6 +49,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) @@ -919,3 +933,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 539be3adddd..4bf47268824 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/typed.rs @@ -8,13 +8,17 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; 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 crate::plan::PlanArrayFuture; +use crate::plan::PlanExecutionContext; use crate::plan::PlanId; use crate::plan::PlanVTable; use crate::plan::display::PlanTreeDisplay; @@ -224,6 +228,14 @@ pub trait DynPlan: 'static + Send + Sync { /// Serializes operator-specific metadata, or `None` when the operator is not serializable. fn dyn_metadata(&self) -> Option>; + + /// Executes this plan over `row_range`, returning the values selected by `mask`. + fn dyn_execute( + &self, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult; } impl DynPlan for Plan { @@ -262,6 +274,15 @@ impl DynPlan for Plan { fn dyn_metadata(&self) -> Option> { V::metadata(self).map(SerializeMetadata::serialize) } + + fn dyn_execute( + &self, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + V::execute(self, ctx, row_range, mask) + } } impl dyn DynPlan + '_ { @@ -300,6 +321,16 @@ impl dyn DynPlan + '_ { self.dyn_metadata() } + /// 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_execute(ctx, row_range, mask) + } + /// Returns the number of children of this plan. pub fn child_count(&self) -> usize { self.children().len() diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/vtable.rs index 3ab6a838b47..3a958df8646 100644 --- a/vortex-layout/src/plan/vtable.rs +++ b/vortex-layout/src/plan/vtable.rs @@ -4,12 +4,17 @@ 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::PlanExecutionContext; use crate::plan::PlanRef; use crate::plan::typed::Plan; @@ -54,6 +59,23 @@ pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { /// an error. This is the hook every generic rewrite goes through. fn with_children(plan: &Plan, children: Vec) -> VortexResult>; + /// 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..96e4b52c071 --- /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) + .with_projection(projection) + .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..be893560e53 --- /dev/null +++ b/vortex-scan-v2/src/scan_builder.rs @@ -0,0 +1,442 @@ +// 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::EvalPlan; +use vortex_layout::plan::PlanExecutionContext; +use vortex_layout::plan::PlanRef; +use vortex_layout::plan::RowIdxPlan; +use vortex_layout::plan::lower; +use vortex_layout::plan::optimize; +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 = RowIdxPlan::new(self.row_offset, self.base_plan.clone()).into_plan(); + tracing::debug!( + target: "vortex_scan_v2::planner", + row_offset = self.row_offset, + plan = %source.display_tree(), + "planning expressions over the row-index-aware source" + ); + tracing::debug!( + target: "vortex_scan_v2::planner", + expression = %self.projection, + "optimizing the projection expression" + ); + let projection = self + .projection + .optimize_recursive(source.dtype())? + .bind(source.dtype())?; + let projection: PlanRef = EvalPlan::new(projection, source.clone()).into_plan(); + let projection = optimize(projection)?; + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %projection.display_tree(), + "optimized the projection physical plan" + ); + + let filter = self + .filter + .map(|filter| -> VortexResult { + tracing::debug!( + target: "vortex_scan_v2::planner", + expression = %filter, + "optimizing the filter expression" + ); + let filter = filter + .optimize_recursive(source.dtype())? + .bind(source.dtype())?; + let filter: PlanRef = EvalPlan::new(filter, source.clone()).into_plan(); + let filter = optimize(filter)?; + 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" + ); + } + + 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, + 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()?)) + } +} + +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..b55ab31d127 --- /dev/null +++ b/vortex-scan-v2/src/splits.rs @@ -0,0 +1,186 @@ +// 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::RowIdx; +use vortex_layout::plan::RowIdxPartition; +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::() || 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(()); + } + + if plan.is::() || 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) +} 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(()) + }) +} From 9260967d130848bfa754e17b275cc8c8d529fd39 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 7 Aug 2026 16:26:21 +0100 Subject: [PATCH 2/2] Bind comparison scan expressions in example Signed-off-by: Joe Isaacs --- vortex-scan-v2/examples/tpch_scan.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-scan-v2/examples/tpch_scan.rs b/vortex-scan-v2/examples/tpch_scan.rs index 96e4b52c071..67b00e1102d 100644 --- a/vortex-scan-v2/examples/tpch_scan.rs +++ b/vortex-scan-v2/examples/tpch_scan.rs @@ -71,8 +71,8 @@ fn main() -> VortexResult<()> { ); let expected = file .scan()? - .with_filter(filter) - .with_projection(projection) + .with_filter(filter.bind(file.dtype())?) + .with_projection(projection.bind(file.dtype())?) .into_array_stream()? .read_all() .await?;