From cc4c7705dd8cc69ac2ea6b789b5af797fb99dc9e Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 8 Oct 2026 14:50:23 +0000 Subject: [PATCH 01/20] feat(layout): add a Filter plan over dense segment scans Add `Filter`, a plan operator with no predicate that keeps only the rows of the selection it is executed with. Its child produces every row of the same row domain, and values outside the selection are unspecified. Flat layouts now lower to a `Filter` over their `SegmentScan`, so a bare scan can stay dense and return every row of its range while the filter above it keeps the selected ones. Plan display snapshots are updated for the extra level. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/lower.rs | 6 +- vortex-layout/src/plan/mod.rs | 3 + vortex-layout/src/plan/plans/filter.rs | 89 ++++++++++++ vortex-layout/src/plan/plans/mod.rs | 4 + vortex-layout/src/plan/tests.rs | 190 ++++++++++++++++--------- 5 files changed, 223 insertions(+), 69 deletions(-) create mode 100644 vortex-layout/src/plan/plans/filter.rs diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index e8911734549..d04e3e128d4 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -25,6 +25,7 @@ use crate::layouts::list::VALIDITY_CHILD_INDEX; use crate::layouts::struct_::Struct; use crate::layouts::struct_::StructLayout; use crate::plan::ConcatPlan; +use crate::plan::FilterPlan; use crate::plan::ListPackPlan; use crate::plan::PackPlan; use crate::plan::PlanChildren; @@ -36,9 +37,12 @@ use crate::plan::TakePlan; /// /// The root operator is built immediately. Its child container owns a hidden clone of the source /// layout and lowers each child independently on first access. +/// +/// A flat layout lowers to a [`Filter`](crate::plan::Filter) over its segment scan, so the scan +/// returns only the rows it is executed with. pub fn lower(layout: &LayoutRef) -> VortexResult { if let Some(layout) = layout.as_opt::() { - return Ok(lower_flat(layout).into_plan()); + return Ok(FilterPlan::new(lower_flat(layout).into_plan()).into_plan()); } if let Some(layout) = layout.as_opt::() { return Ok(lower_chunked(layout)?.into_plan()); diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index b4782481b49..7e45533ed84 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -31,6 +31,9 @@ pub use plans::ConcatPlan; pub use plans::Eval; pub use plans::EvalData; pub use plans::EvalPlan; +pub use plans::Filter; +pub use plans::FilterData; +pub use plans::FilterPlan; pub use plans::ListPack; pub use plans::ListPackData; pub use plans::ListPackPlan; diff --git a/vortex-layout/src/plan/plans/filter.rs b/vortex-layout/src/plan/plans/filter.rs new file mode 100644 index 00000000000..da5907f3e80 --- /dev/null +++ b/vortex-layout/src/plan/plans/filter.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; + +use vortex_array::EmptyMetadata; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::plan::check_child_count; + +/// Keeps only the selected rows of its child. +/// +/// A filter holds no predicate. The rows it keeps are the selection it is executed with; its child +/// produces every row of the same row domain, and values outside the selection are unspecified. +/// The filter's dtype and row domain are its child's. +#[derive(Clone, Debug)] +pub struct Filter; + +/// Operator-specific data for a [`Filter`] plan. +#[derive(Clone, Debug)] +pub struct FilterData; + +/// A plan that keeps only the selected rows of its child. +pub type FilterPlan = Plan; + +impl FilterPlan { + /// Creates a filter over `child`. + pub fn new(child: PlanRef) -> Self { + PlanParts { + vtable: Filter, + dtype: child.dtype().clone(), + row_count: child.row_count(), + children: vec![child].into(), + data: FilterData, + } + .into_typed() + } + + /// Returns the plan whose rows are filtered. + pub fn child_plan(&self) -> VortexResult { + self.child_required(0) + } +} + +impl PlanVTable for Filter { + type PlanData = FilterData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.filter"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + Some(EmptyMetadata) + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + check_child_count("Filter", children, 1)?; + let child = children + .get(0)? + .ok_or_else(|| vortex_err!("Filter child is absent"))?; + if child.dtype() != plan.dtype() || child.row_count() != plan.row_count() { + vortex_bail!("Filter child does not match the filter's dtype and row count"); + } + Ok(()) + } + + fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + if index == 0 { + Cow::Borrowed("child") + } else { + Cow::Owned(format!("child[{index}]")) + } + } +} diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index 7b6a1893d0b..b20731b8d34 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -3,6 +3,7 @@ mod concat; pub(crate) mod eval; +mod filter; mod list_pack; mod pack; mod row_idx; @@ -17,6 +18,9 @@ pub use eval::Eval; pub use eval::EvalData; pub(crate) use eval::EvalIdentityRule; pub use eval::EvalPlan; +pub use filter::Filter; +pub use filter::FilterData; +pub use filter::FilterPlan; pub use list_pack::ListPack; pub use list_pack::ListPackData; pub use list_pack::ListPackPlan; diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 783eb455327..fccc06ebb2f 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -94,12 +94,14 @@ fn unsupported_layout_has_no_plan() -> VortexResult<()> { } #[test] -fn flat_plan_has_no_children() -> VortexResult<()> { +fn flat_plan_is_a_filtered_segment_scan_without_children() -> VortexResult<()> { let plan = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; - assert!(plan.is::()); - assert_eq!(plan.child_count(), 0); - assert!(plan.child(0)?.is_none()); + assert!(plan.is::()); + let scan = child_of(&plan, 0)?; + assert!(scan.is::()); + assert_eq!(scan.child_count(), 0); + assert!(scan.child(0)?.is_none()); Ok(()) } @@ -323,7 +325,7 @@ fn optimize_drops_identity_expressions() -> VortexResult<()> { let expression = root().bind(child.dtype())?; let plan: PlanRef = EvalPlan::try_new(expression, child)?.into_plan(); - assert!(optimize(plan)?.is::()); + assert!(optimize(plan)?.is::()); Ok(()) } @@ -346,8 +348,8 @@ fn optimize_rewrites_nested_children() -> VortexResult<()> { let optimized = optimize(wrapped)?; assert!(optimized.is::()); - assert!(child_of(&optimized, 0)?.is::()); - assert!(child_of(&optimized, 1)?.is::()); + assert!(child_of(&optimized, 0)?.is::()); + assert!(child_of(&optimized, 1)?.is::()); Ok(()) } @@ -371,11 +373,13 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { let plan: PlanRef = plan.into_plan(); assert_eq!(plan.to_string(), "vortex.plan.eval(i32, rows=3) expr=$.a"); - insta::assert_snapshot!(plan.display_tree(), @r" + insta::assert_snapshot!(plan.display_tree(), @" root: vortex.plan.eval(i32, rows=3) expr=$.a child: vortex.plan.pack({a=i32, b=i32}, rows=3) - a: vortex.plan.segment_scan(i32, rows=3) - b: vortex.plan.segment_scan(i32, rows=3) + a: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) + b: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); struct DepthExtractor; @@ -391,11 +395,13 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { } } - insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @r" + insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @" root: depth=0 child: depth=1 a: depth=2 + child: depth=3 b: depth=2 + child: depth=3 "); let nullable_fields = StructFields::from_iter([ @@ -413,11 +419,14 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { ) .into_layout(); let nullable = make_plan(nullable_layout)?; - insta::assert_snapshot!(nullable.tree_display_builder(), @r" + insta::assert_snapshot!(nullable.tree_display_builder(), @" root: a: + child: b: + child: validity: + child: "); Ok(()) } @@ -433,10 +442,12 @@ fn chunked_plan_display_names_chunks() -> VortexResult<()> { .into_layout(); let plan = make_plan(layout)?; - insta::assert_snapshot!(plan.display_tree(), @r" + insta::assert_snapshot!(plan.display_tree(), @" root: vortex.plan.concat(i32, rows=3) - chunks[0]: vortex.plan.segment_scan(i32, rows=2) - chunks[1]: vortex.plan.segment_scan(i32, rows=1) + chunks[0]: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) + chunks[1]: vortex.plan.filter(i32, rows=1) + child: vortex.plan.segment_scan(i32, rows=1) "); Ok(()) } @@ -450,10 +461,12 @@ fn dict_plan_display_names_logical_children() -> VortexResult<()> { .into_layout(); let plan = make_plan(layout)?; - insta::assert_snapshot!(plan.display_tree(), @r" + insta::assert_snapshot!(plan.display_tree(), @" root: vortex.plan.take(i32, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.segment_scan(i32, rows=2) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -471,10 +484,12 @@ fn list_plan_display_handles_optional_validity() -> VortexResult<()> { .into_layout(); let non_nullable = make_plan(non_nullable_layout)?; - insta::assert_snapshot!(non_nullable.display_tree(), @r" + insta::assert_snapshot!(non_nullable.display_tree(), @" root: vortex.plan.list_pack(list(i32), rows=2) - elements: vortex.plan.segment_scan(i32, rows=4) - offsets: vortex.plan.segment_scan(u32, rows=3) + elements: vortex.plan.filter(i32, rows=4) + child: vortex.plan.segment_scan(i32, rows=4) + offsets: vortex.plan.filter(u32, rows=3) + child: vortex.plan.segment_scan(u32, rows=3) "); let nullable_layout = ListLayout::new( @@ -486,11 +501,14 @@ fn list_plan_display_handles_optional_validity() -> VortexResult<()> { .into_layout(); let nullable = make_plan(nullable_layout)?; - insta::assert_snapshot!(nullable.display_tree(), @r" + insta::assert_snapshot!(nullable.display_tree(), @" root: vortex.plan.list_pack(list(i32)?, rows=2) - elements: vortex.plan.segment_scan(i32, rows=4) - offsets: vortex.plan.segment_scan(u32, rows=3) - validity: vortex.plan.segment_scan(bool, rows=2) + elements: vortex.plan.filter(i32, rows=4) + child: vortex.plan.segment_scan(i32, rows=4) + offsets: vortex.plan.filter(u32, rows=3) + child: vortex.plan.segment_scan(u32, rows=3) + validity: vortex.plan.filter(bool, rows=2) + child: vortex.plan.segment_scan(bool, rows=2) "); Ok(()) } @@ -546,11 +564,14 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { child: vortex.plan.eval({child_0=bool, child_1=bool}, rows=3) expr=pack(child_0: $.a, child_1: $.b) child: vortex.plan.pack({a=bool, b=bool}, rows=3) a: vortex.plan.take(bool, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) b: vortex.plan.eval(bool, rows=3) expr=($ > 7i32) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -600,16 +621,18 @@ fn row_idx_and_data_expression_pushes_data_into_chunks() -> VortexResult<()> { let expression = and(gt(row_idx(), lit(10_u64)), gt(root(), lit(5_i32))); let plan = optimize(make_row_idx_plan(expression, child)?)?; - insta::assert_snapshot!(plan.display_tree(), @r" + insta::assert_snapshot!(plan.display_tree(), @" root: vortex.plan.eval(bool, rows=3) expr=($.row_idx and $.child) child: vortex.plan.pack({row_idx=bool, child=bool}, rows=3) row_idx: vortex.plan.eval(bool, rows=3) expr=($ > 10u64) child: vortex.plan.row_idx(u64, rows=3) child: vortex.plan.concat(bool, rows=3) chunks[0]: vortex.plan.eval(bool, rows=1) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=1) + child: vortex.plan.filter(i32, rows=1) + child: vortex.plan.segment_scan(i32, rows=1) chunks[1]: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -635,17 +658,22 @@ fn expression_pushes_through_struct_field_and_dictionary_values() -> VortexResul root: vortex.plan.eval(bool, rows=3) expr=($.a > 5i32) child: vortex.plan.pack({a=i32, b=i32}, rows=3) a: vortex.plan.take(i32, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.segment_scan(i32, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); let optimized = optimize(plan)?; insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.take(bool, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -680,21 +708,28 @@ fn expression_pushes_through_struct_field_with_heterogeneous_chunks() -> VortexR child: vortex.plan.pack({a=i32, b=i32}, rows=5) a: vortex.plan.concat(i32, rows=5) chunks[0]: vortex.plan.take(i32, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.segment_scan(i32, rows=2) - chunks[1]: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.segment_scan(i32, rows=5) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) + chunks[1]: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.filter(i32, rows=5) + child: vortex.plan.segment_scan(i32, rows=5) "); let optimized = optimize(plan)?; insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.concat(bool, rows=5) chunks[0]: vortex.plan.take(bool, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) chunks[1]: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -728,9 +763,10 @@ fn expression_pushes_through_nested_struct_fields_in_one_pass() -> VortexResult< let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @r" + insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -756,17 +792,19 @@ fn expression_pushes_through_single_field_nested_structs() -> VortexResult<()> { let expression = gt(get_item("b", get_item("a", root())), lit(5_i32)); let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); - insta::assert_snapshot!(plan.display_tree(), @r" + insta::assert_snapshot!(plan.display_tree(), @" root: vortex.plan.eval(bool, rows=3) expr=($.a.b > 5i32) child: vortex.plan.pack({a={b=i32}}, rows=3) a: vortex.plan.pack({b=i32}, rows=3) - b: vortex.plan.segment_scan(i32, rows=3) + b: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @r" + insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -800,9 +838,10 @@ fn expression_pushes_through_three_single_field_structs() -> VortexResult<()> { let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @r" + insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -847,15 +886,18 @@ fn compound_expression_pushes_through_nested_struct_and_dictionary() -> VortexRe let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @r" + insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.eval(bool, rows=3) expr=($.b and $.c) child: vortex.plan.pack({b=bool, c=bool}, rows=3) b: vortex.plan.take(bool, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) c: vortex.plan.eval(bool, rows=3) expr=($ > 7i32) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -894,11 +936,13 @@ fn expression_pushes_through_dictionary_of_struct_values() -> VortexResult<()> { let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @r" + insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.take(bool, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -938,10 +982,14 @@ fn multi_field_struct_expression_pushes_into_each_field() -> VortexResult<()> { root: vortex.plan.eval(bool, rows=3) expr=(($.a > 5i32) and ($.b > 7i32)) child: vortex.plan.pack({a=i32, b=i32, c=i32}, rows=3) a: vortex.plan.take(i32, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.segment_scan(i32, rows=3) - c: vortex.plan.segment_scan(i32, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) + c: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); let optimized = optimize(plan)?; @@ -949,11 +997,14 @@ fn multi_field_struct_expression_pushes_into_each_field() -> VortexResult<()> { root: vortex.plan.eval(bool, rows=3) expr=($.a and $.b) child: vortex.plan.pack({a=bool, b=bool}, rows=3) a: vortex.plan.take(bool, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) b: vortex.plan.eval(bool, rows=3) expr=($ > 7i32) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); let reoptimized = optimize(optimized.clone())?; assert!(PlanRef::ptr_eq(&optimized, &reoptimized)); @@ -1022,9 +1073,12 @@ fn multi_field_struct_expression_keeps_cross_field_refinement() -> VortexResult< root: vortex.plan.eval(bool, rows=3) expr=(($.a + $.b) > 10i32) child: vortex.plan.pack({a=i32, b=i32}, rows=3) a: vortex.plan.take(i32, rows=3) - codes: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.segment_scan(i32, rows=3) + codes: vortex.plan.filter(u8, rows=3) + child: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.filter(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); assert_eq!( optimized.display_tree().to_string(), From 8e49f35f75ae42a1aa2bb0d929e4d7bfe3170dbf Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 8 Oct 2026 20:13:45 +0000 Subject: [PATCH 02/20] feat(layout): add the ExecNode trait and push-based exec graph Add `plan::exec`, a push-based executor interface for physical plans. A node is built by `PlanVTable::exec` for a range of its plan's rows and a selection over them. It has input ports, one output, and two entry points: `start`, called once to spawn children and request segments, and `compute`, called whenever something changed for the node and its readiness rule holds. A port is an ordered queue of arrays fed by one child plus a closed flag, owned by the graph and read by the node through `StepCx::input`. Every node emits its selected rows in row order as non-empty arrays and returns `Done` to close its own output; there are no close events, no empty arrays, and no row ranges on the data. `Ready` is the readiness rule: `Any` for nodes that stream input through, `Closed(ports)` for nodes that need one input whole before using the others, as a dictionary needs its values before its codes, and `AllClosed` for nodes that need everything. Chains of single-port `Any` nodes are pipelines; the other rules are the barriers between them. `ExecGraph` is the flat node arena. It wires ports, routes segment deliveries to the leaves that requested them, runs ready nodes last-in first-out so data drains toward the root, and hands its owner the root's arrays and batches of `IoRequest`s. IO may complete in any order; a node whose children are backed by independent reads holds the early ones in their ports and emits the ordered prefix. `plan::exec::synthetic` has plans whose nodes are built by hand, so the graph can be exercised without layouts. The scheduling tests drive sources that cut their rows into pieces, request in-memory segments, and yield, under a probe node that records when its readiness rule let it run. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- Cargo.lock | 1 + vortex-layout/Cargo.toml | 1 + vortex-layout/src/plan/exec/mod.rs | 672 ++++++++++++++++++ .../src/plan/exec/scheduling_tests.rs | 214 ++++++ vortex-layout/src/plan/exec/synthetic.rs | 295 ++++++++ vortex-layout/src/plan/mod.rs | 1 + vortex-layout/src/plan/typed.rs | 45 ++ vortex-layout/src/plan/vtable.rs | 21 + 8 files changed, 1250 insertions(+) create mode 100644 vortex-layout/src/plan/exec/mod.rs create mode 100644 vortex-layout/src/plan/exec/scheduling_tests.rs create mode 100644 vortex-layout/src/plan/exec/synthetic.rs diff --git a/Cargo.lock b/Cargo.lock index 58506db4952..ed7cfc974ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12129,6 +12129,7 @@ dependencies = [ "rstest", "rustc-hash", "sketches-ddsketch", + "smallvec", "temp-env", "termtree", "tokio", diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index eab8869bf8b..9d235191556 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -35,6 +35,7 @@ pin-project-lite = { workspace = true } prost = { workspace = true } rustc-hash = { workspace = true } sketches-ddsketch = { workspace = true } +smallvec = { workspace = true } termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } diff --git a/vortex-layout/src/plan/exec/mod.rs b/vortex-layout/src/plan/exec/mod.rs new file mode 100644 index 00000000000..5cb10f6e569 --- /dev/null +++ b/vortex-layout/src/plan/exec/mod.rs @@ -0,0 +1,672 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Push-based execution of physical plans. +//! +//! A [`PlanRef`] is turned into an [`ExecGraph`] for one row range and selection. The graph is a +//! flat arena of [`ExecNode`]s wired child-to-parent through [`Input`] ports. A node's output is +//! an ordered sequence of arrays: each node emits its rows in row order, and its parent reads +//! them from the port the child feeds. The root's arrays are handed to the owner one per +//! [`ExecGraph::compute`], in row order. +//! +//! IO may complete in any order. A node whose children are backed by independent reads, such as +//! a concatenation of chunks, holds the early ones in their ports and emits only the ordered +//! prefix, so nothing above it sees the disorder. Nodes never talk to the IO service: they +//! publish requests through [`StepCx::request`], the graph returns them from +//! [`ExecGraph::compute`], and the owner delivers results through [`ExecGraph::set_io_result`]. +//! +//! # Interfaces +//! +//! The module has one interface for each side of a graph: +//! +//! - A plan operator runs as an [`ExecNode`], built by +//! [`PlanVTable::exec`](crate::plan::PlanVTable::exec). A node sees the graph only through its +//! [`StepCx`]: it reads its ports, and spawns children, requests segments, and emits arrays +//! through it. [`ExecNode`] says what a node must do, and [`Ready`] says when it runs. +//! - An owner drives an [`ExecGraph`]: it calls [`ExecGraph::compute`], answers the +//! [`IoRequest`]s that returns, and collects the root's arrays. +//! +//! How a graph stores, wires, and schedules its nodes belongs to neither interface, and nor do +//! the nodes of this crate's own operators. +//! +//! # Pipelines +//! +//! A chain of nodes that are [`Ready::Any`] over one port is a pipeline: an array emitted at +//! its bottom can be pushed through every node in the chain in turn. Nodes with several ports +//! merge pipelines, and nodes that wait for a port to close, [`Ready::Closed`] and +//! [`Ready::AllClosed`], are barriers between them, as a join build is in a query engine. + +pub mod synthetic; + +use std::collections::VecDeque; +use std::mem; +use std::ops::Range; +use std::sync::Arc; + +use parking_lot::Mutex; +use rustc_hash::FxHashMap; +use smallvec::SmallVec; +use vortex_array::ArrayRef; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::plan::Concat; +use crate::plan::PlanRef; +use crate::segments::SegmentId; + +/// Identifies one IO request within an [`ExecGraph`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct IoRequestId(pub(crate) u64); + +/// A read published by a leaf node. +#[derive(Clone, Debug)] +pub struct IoRequest { + /// Identifies the request when its result is delivered. + pub id: IoRequestId, + /// The segment to read. + pub segment_id: SegmentId, +} + +/// Requests discovered together. +pub type IoBatch = Vec; + +/// An input of a node. +/// +/// A node names the port when it spawns the child that feeds it, and reads the child's output +/// from it through [`StepCx::input`]. +pub type Port = usize; + +/// One input port of a node: the arrays its child has produced and not yet been taken, in row +/// order, and whether the child has closed. +/// +/// A port that was never spawned reads as closed with nothing available, so a node can treat an +/// optional child it did not spawn like one that finished at once. +#[derive(Default)] +pub struct Input { + spawned: bool, + closed: bool, + pieces: VecDeque, + available: usize, +} + +impl Input { + /// Rows the child has produced that the node has not taken yet. + pub fn available(&self) -> usize { + self.available + } + + /// Whether the child will produce nothing more. True for a port that was never spawned. + pub fn closed(&self) -> bool { + !self.spawned || self.closed + } + + /// Whether the child has closed and every row it produced has been taken. + pub fn finished(&self) -> bool { + self.closed() && self.pieces.is_empty() + } + + /// Takes every available array, in row order. + pub fn take_all(&mut self) -> Vec { + self.available = 0; + self.pieces.drain(..).collect() + } + + /// Takes the next `n` rows, in row order, as the arrays that hold them. The array that + /// spans the boundary is sliced, and its remainder stays available. + /// + /// `n` must not exceed [`available`](Self::available). + pub fn take(&mut self, n: usize) -> VortexResult> { + if n > self.available { + vortex_bail!( + "Taking {n} rows from a port with {} available", + self.available + ); + } + let mut taken = Vec::new(); + let mut remaining = n; + while remaining > 0 { + let Some(front) = self.pieces.pop_front() else { + vortex_bail!("Port ran out of arrays with {remaining} rows still to take"); + }; + if front.len() <= remaining { + remaining -= front.len(); + taken.push(front); + } else { + taken.push(front.slice(0..remaining)?); + self.pieces.push_front(front.slice(remaining..front.len())?); + remaining = 0; + } + } + self.available -= n; + Ok(taken) + } + + fn push(&mut self, array: ArrayRef) { + self.available += array.len(); + self.pieces.push_back(array); + } +} + +/// When a node's [`compute`](ExecNode::compute) runs. +/// +/// A node is considered whenever something changes for it: a child produced rows or closed, or +/// a requested segment was delivered. The rule then decides whether the node runs. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Ready { + /// On every change. For nodes that stream their input through. + Any, + /// Only once each of these ports has closed. For nodes that need one input whole before + /// they can use the others, as a dictionary needs its values before its codes. + Closed(&'static [Port]), + /// Only once every spawned port has closed. For nodes that need all their input at once. + AllClosed, +} + +/// One running plan operator inside an [`ExecGraph`]. +/// +/// A node is built by [`PlanVTable::exec`](crate::plan::PlanVTable::exec) for a range of its +/// plan's rows and a selection over them, and produces the selected rows as arrays, in row +/// order. +/// +/// # Lifecycle +/// +/// The graph calls [`start`](Self::start) once, right after the node is built. It then calls +/// [`compute`](Self::compute) whenever something changed for the node and the node's +/// [`ready`](Self::ready) rule holds, and again after `compute` returned [`NodeState::Yield`]. +/// A node is not called again once it has returned [`NodeState::Done`] or an error. +/// +/// Everything a node does goes through its [`StepCx`] and takes effect once the call returns, +/// in this order: its children are spawned and started, its requests are published, its arrays +/// are pushed to its parent's port, and, if it is done, its port closes. +/// +/// # Output +/// +/// A node emits the selected rows of the range it was built for, in row order, as any number of +/// non-empty arrays. It returns [`NodeState::Done`] once it has emitted them all. Its children +/// are not cancelled, so it returns `Done` only once every child it spawned has closed, and it +/// returns [`NodeState::Wait`] only while a child is open or a request is outstanding, since +/// nothing else runs it again. +/// +/// The graph checks none of this. +/// +/// # Errors +/// +/// An error from `start` or `compute` fails the whole graph. +pub trait ExecNode: Send { + /// When [`compute`](Self::compute) runs. May change as the node progresses. + fn ready(&self) -> Ready { + Ready::Any + } + + /// Spawns children, publishes requests, and emits whatever needs no input. + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult; + + /// Takes what its ports hold, does work, and reports what happens to the node next. + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult; +} + +/// What happens to a node after a [`start`](ExecNode::start) or [`compute`](ExecNode::compute). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NodeState { + /// Run again once something changes: a child produces rows or closes, or a delivery lands. + Wait, + /// Run again without waiting for a change, after other work: CPU work is left. + Yield, + /// Finished: the node's port closes and it never runs again. + Done, +} + +/// What a graph gives every node it builds. +#[derive(Clone)] +pub struct ExecContext { + session: VortexSession, + row_offset: u64, + decoded: DecodeCache, +} + +impl ExecContext { + pub(crate) fn new(session: VortexSession, row_offset: u64, decoded: DecodeCache) -> Self { + Self { + session, + row_offset, + decoded, + } + } + + /// The session used for decoding and expression evaluation. + pub fn session(&self) -> &VortexSession { + &self.session + } + + /// The global row index of the graph's first plan row. + pub fn row_offset(&self) -> u64 { + self.row_offset + } + + /// Decoded segments shared by the graphs reading the same rows. + pub fn decoded(&self) -> &DecodeCache { + &self.decoded + } +} + +/// Decoded segments shared by the graphs that read one row range, so a segment several of them +/// read is fetched and decoded once. +/// +/// A segment decodes the same way wherever it appears, so entries are keyed by segment id alone. +#[derive(Clone, Default)] +pub struct DecodeCache { + decoded: Arc>>, +} + +impl DecodeCache { + /// The whole decoded array of `id`, if a graph sharing this cache decoded it. + pub fn get(&self, id: SegmentId) -> Option { + self.decoded.lock().get(&id).cloned() + } + + /// Shares the whole decoded array of `id` with the graphs sharing this cache. + pub fn insert(&self, id: SegmentId, array: ArrayRef) { + self.decoded.lock().insert(id, array); + } +} + +/// A child a node spawned: its port, plan, rows, and selection. +type Spawn = (Port, PlanRef, Range, Mask); + +/// A node's view of its graph during one [`start`](ExecNode::start) or +/// [`compute`](ExecNode::compute): its ports, a delivered segment, and what the node does about +/// them. +/// +/// Nothing recorded here happens until the call returns. [`ExecNode`] gives the order. +pub struct StepCx<'a> { + inputs: &'a mut SmallVec<[Input; 2]>, + delivered: Option, + next_io_id: &'a mut u64, + spawned: SmallVec<[Spawn; 1]>, + requests: SmallVec<[IoRequest; 1]>, + emitted: SmallVec<[ArrayRef; 1]>, +} + +impl<'a> StepCx<'a> { + fn new( + inputs: &'a mut SmallVec<[Input; 2]>, + delivered: Option, + next_io_id: &'a mut u64, + ) -> Self { + Self { + inputs, + delivered, + next_io_id, + spawned: SmallVec::new(), + requests: SmallVec::new(), + emitted: SmallVec::new(), + } + } + + /// The port fed by the child spawned on `port`. + pub fn input(&mut self, port: Port) -> &mut Input { + if port >= self.inputs.len() { + self.inputs.resize_with(port + 1, Input::default); + } + &mut self.inputs[port] + } + + /// Every port, in port order, including ports never spawned. + pub fn inputs(&mut self) -> &mut [Input] { + self.inputs + } + + /// Whether every spawned port has closed and been drained. + pub fn all_finished(&self) -> bool { + self.inputs.iter().all(Input::finished) + } + + /// The bytes of the segment this node requested, if they arrived since it last ran. + pub fn take_delivery(&mut self) -> Option { + self.delivered.take() + } + + /// Spawns `plan` as a child feeding `port`, over `rows` of the child's row domain restricted + /// to `mask`, which is as long as `rows`. + pub fn spawn(&mut self, port: Port, plan: PlanRef, rows: Range, mask: Mask) { + self.spawned.push((port, plan, rows, mask)); + } + + /// Publishes a read of `segment_id`. Its bytes arrive through + /// [`take_delivery`](Self::take_delivery). A node has at most one read outstanding. + pub fn request(&mut self, segment_id: SegmentId) -> IoRequestId { + let id = IoRequestId(*self.next_io_id); + *self.next_io_id += 1; + self.requests.push(IoRequest { id, segment_id }); + id + } + + /// Pushes the node's next rows to its parent, or from the root to the graph's owner. Empty + /// arrays are dropped. + pub fn emit(&mut self, array: ArrayRef) { + if !array.is_empty() { + self.emitted.push(array); + } + } +} + +/// What an [`ExecGraph`] needs next. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExecState { + /// A node is ready, or an array or requests are waiting to be returned. + NeedsCompute, + /// Every remaining node waits on a delivered request. + Waiting, + /// The root has closed and every array has been returned. + Done, +} + +/// The result of one [`ExecGraph::compute`] call. +#[derive(Debug)] +pub enum ExecOutput { + /// Control returns to the owner. Check [`ExecGraph::state`] for what is needed next. + Yield, + /// Requests published since the last flush, each returned exactly once. + NeedsIO(IoBatch), + /// The root's next rows. + Piece(ArrayRef), +} + +type NodeId = usize; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Status { + /// Not in the ready queue. + Idle, + /// In the ready queue. + Queued, + /// Closed; never run again. + Done, +} + +struct Node { + exec: Box, + parent: Option<(NodeId, Port)>, + inputs: SmallVec<[Input; 2]>, + delivered: Option, + status: Status, +} + +/// A running plan: a flat arena of nodes wired child-to-parent. +pub struct ExecGraph { + ctx: ExecContext, + nodes: Vec, + /// Ready nodes, run last-in first-out so a woken parent runs before its child's next call. + ready: Vec, + io_routes: SmallVec<[Option; 2]>, + new_io: IoBatch, + outputs: VecDeque, + root_closed: bool, + next_io_id: u64, +} + +impl ExecGraph { + /// Builds the graph for `plan` over `rows`, restricted to `mask`, and starts each node. + /// `row_offset` is the global row index of the plan's first row, used by row-index plans. + /// Segments found in `decoded` are neither read nor decoded again, and segments the graph + /// decodes are added to it. + /// + /// Construction does no IO. Every leaf's first request is returned by the first + /// [`compute`](Self::compute). + pub fn try_new( + session: VortexSession, + plan: &PlanRef, + rows: Range, + mask: Mask, + row_offset: u64, + decoded: DecodeCache, + ) -> VortexResult { + let mut graph = Self { + ctx: ExecContext::new(session, row_offset, decoded), + nodes: Vec::new(), + ready: Vec::new(), + io_routes: SmallVec::new(), + new_io: Vec::new(), + outputs: VecDeque::new(), + root_closed: false, + next_io_id: 0, + }; + graph.spawn(None, plan, rows, mask)?; + Ok(graph) + } + + /// Reports what the graph needs next. Cheap and side-effect free. + pub fn state(&self) -> ExecState { + if !self.outputs.is_empty() || !self.ready.is_empty() || !self.new_io.is_empty() { + ExecState::NeedsCompute + } else if self.root_closed { + ExecState::Done + } else { + ExecState::Waiting + } + } + + /// Ids of requests returned by [`compute`](Self::compute) whose results are not yet delivered. + pub fn outstanding(&self) -> impl Iterator + '_ { + self.io_routes + .iter() + .enumerate() + .filter_map(|(index, node)| { + let id = IoRequestId(index as u64); + (node.is_some() && !self.new_io.iter().any(|request| request.id == id)) + .then_some(id) + }) + } + + /// Runs ready nodes until an array reaches the root, requests should be flushed, a node + /// yields, or nothing is ready. + /// + /// # Errors + /// + /// Returns the error of a node that failed. The graph cannot be run further. + pub fn compute(&mut self) -> VortexResult { + loop { + if let Some(array) = self.outputs.pop_front() { + return Ok(ExecOutput::Piece(array)); + } + let Some(node) = self.ready.pop() else { + return Ok(self.flush_or_yield()); + }; + let yielded = self.run(node, false)?; + + // Get reads in flight before spending more CPU on other nodes. + if !self.new_io.is_empty() && self.outputs.is_empty() { + return Ok(ExecOutput::NeedsIO(mem::take(&mut self.new_io))); + } + if yielded && self.outputs.is_empty() { + return Ok(ExecOutput::Yield); + } + } + } + + /// Delivers the result of a request returned by [`compute`](Self::compute). + pub fn set_io_result(&mut self, id: IoRequestId, result: BufferHandle) -> VortexResult<()> { + let node = self + .io_routes + .get_mut(usize::try_from(id.0)?) + .and_then(Option::take) + .ok_or_else(|| vortex_err!("Unknown IO request {id:?}"))?; + self.nodes[node].delivered = Some(result); + self.enqueue(node); + Ok(()) + } + + fn flush_or_yield(&mut self) -> ExecOutput { + if self.new_io.is_empty() { + ExecOutput::Yield + } else { + ExecOutput::NeedsIO(mem::take(&mut self.new_io)) + } + } + + /// Builds the node for `plan` under `parent` and starts it. + /// + /// A concatenation whose requested rows fall inside one chunk is skipped: the chunk is + /// wired to the parent directly, since its rows are the concatenation's rows in order. + fn spawn( + &mut self, + parent: Option<(NodeId, Port)>, + plan: &PlanRef, + rows: Range, + mask: Mask, + ) -> VortexResult<()> { + if let Some(concat) = plan.as_opt::() + && rows.start < rows.end + { + let offsets = concat.row_offsets(); + let index = offsets + .partition_point(|&offset| offset <= rows.start) + .saturating_sub(1); + let end = offsets + .get(index + 1) + .copied() + .unwrap_or_else(|| concat.row_count()); + if let Some(&start) = offsets.get(index) + && rows.end <= end + { + return self.spawn( + parent, + &concat.child_required(index)?, + rows.start - start..rows.end - start, + mask, + ); + } + } + if let Some((parent, port)) = parent { + let inputs = &mut self.nodes[parent].inputs; + if port >= inputs.len() { + inputs.resize_with(port + 1, Input::default); + } + inputs[port].spawned = true; + } + let id = self.nodes.len(); + self.nodes.push(Node { + exec: plan.exec(rows, mask, &self.ctx)?, + parent, + inputs: SmallVec::new(), + delivered: None, + status: Status::Queued, + }); + self.run(id, true)?; + Ok(()) + } + + /// Runs one node, its `start` or its `compute`, and applies the effects. Returns whether + /// control goes back to the owner: when the node yields, and after a node handles a + /// delivery, since decoding and evaluating newly read data is the costly work between reads + /// that other owners' work should interleave with. + fn run(&mut self, id: NodeId, start: bool) -> VortexResult { + let mut next_io_id = self.next_io_id; + let node = &mut self.nodes[id]; + let delivered = node.delivered.take(); + let had_delivery = delivered.is_some(); + let mut cx = StepCx::new(&mut node.inputs, delivered, &mut next_io_id); + let state = if start { + node.exec.start(&mut cx)? + } else { + node.exec.compute(&mut cx)? + }; + let effects = Effects::from(cx); + self.next_io_id = next_io_id; + + // Settle the node before waking its parent, so a re-queued node runs after the parent. + self.nodes[id].status = match state { + NodeState::Done => Status::Done, + NodeState::Wait | NodeState::Yield => Status::Idle, + }; + if state == NodeState::Yield { + self.enqueue(id); + } + self.apply(id, effects, state == NodeState::Done)?; + Ok(state == NodeState::Yield || had_delivery) + } + + fn apply(&mut self, id: NodeId, effects: Effects, done: bool) -> VortexResult<()> { + for (port, plan, rows, mask) in effects.spawned { + self.spawn(Some((id, port)), &plan, rows, mask)?; + } + for request in effects.requests { + let index = usize::try_from(request.id.0)?; + if index >= self.io_routes.len() { + self.io_routes.resize(index + 1, None); + } + self.io_routes[index] = Some(id); + self.new_io.push(request); + } + for array in effects.emitted { + self.deliver(id, Some(array)); + } + if done { + self.deliver(id, None); + } + Ok(()) + } + + /// Passes an array, or the close when `array` is `None`, from `from` to its parent's port, + /// and wakes the parent if its readiness rule holds. + fn deliver(&mut self, from: NodeId, array: Option) { + match self.nodes[from].parent { + None => match array { + Some(array) => self.outputs.push_back(array), + None => self.root_closed = true, + }, + Some((parent, port)) => { + let input = &mut self.nodes[parent].inputs[port]; + match array { + Some(array) => input.push(array), + None => input.closed = true, + } + if self.is_ready(parent) { + self.enqueue(parent); + } + } + } + } + + fn is_ready(&self, id: NodeId) -> bool { + let node = &self.nodes[id]; + match node.exec.ready() { + Ready::Any => true, + Ready::Closed(ports) => ports + .iter() + .all(|&port| node.inputs.get(port).is_none_or(Input::closed)), + Ready::AllClosed => node.inputs.iter().all(Input::closed), + } + } + + fn enqueue(&mut self, id: NodeId) { + if self.nodes[id].status == Status::Idle { + self.nodes[id].status = Status::Queued; + self.ready.push(id); + } + } +} + +/// The owned effects of a [`StepCx`], detached from its borrows. +struct Effects { + spawned: SmallVec<[Spawn; 1]>, + requests: SmallVec<[IoRequest; 1]>, + emitted: SmallVec<[ArrayRef; 1]>, +} + +impl From> for Effects { + fn from(cx: StepCx<'_>) -> Self { + Self { + spawned: cx.spawned, + requests: cx.requests, + emitted: cx.emitted, + } + } +} + +#[cfg(test)] +mod scheduling_tests; diff --git a/vortex-layout/src/plan/exec/scheduling_tests.rs b/vortex-layout/src/plan/exec/scheduling_tests.rs new file mode 100644 index 00000000000..6fb2ee452f7 --- /dev/null +++ b/vortex-layout/src/plan/exec/scheduling_tests.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests of the graph's scheduling, driven through hand-built nodes over in-memory segments. +//! +//! Every source emits its own row indices, so what reaches the root can be checked against a +//! sequence. + +#![allow(clippy::cast_possible_truncation)] + +use std::ops::Range; +use std::sync::Arc; + +use parking_lot::Mutex; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use super::synthetic::Probe; +use super::synthetic::RowSource; +use super::synthetic::empty_segment; +use super::*; +use crate::test::SESSION; + +#[derive(Debug, PartialEq, Eq)] +enum Event { + /// A batch returned by `compute`, as segment ids. + Io(Vec), + /// A result delivered for this segment. + Delivered(u32), + /// An array of this many rows reached the root. + Piece(usize), + /// A node yielded to the owner. + Yield, +} + +struct Run { + arrays: Vec, + events: Vec, +} + +/// Drives a graph the way an owner does, answering reads with empty segments in the order +/// `pick` chooses, only while the graph waits. +fn run( + plan: &PlanRef, + rows: Range, + mask: Mask, + mut pick: impl FnMut(&[IoRequest]) -> usize, +) -> VortexResult { + let mut graph = + ExecGraph::try_new(SESSION.clone(), plan, rows, mask, 0, DecodeCache::default())?; + let mut inflight: Vec = Vec::new(); + let mut arrays = Vec::new(); + let mut events = Vec::new(); + loop { + match graph.state() { + ExecState::Done => break, + ExecState::NeedsCompute => match graph.compute()? { + ExecOutput::Piece(array) => { + assert!(!array.is_empty(), "the root never emits an empty array"); + events.push(Event::Piece(array.len())); + arrays.push(array); + } + ExecOutput::NeedsIO(batch) => { + events.push(Event::Io(batch.iter().map(|r| *r.segment_id).collect())); + inflight.extend(batch); + } + ExecOutput::Yield => events.push(Event::Yield), + }, + ExecState::Waiting => { + if inflight.is_empty() { + return Err(vortex_err!("graph waits with no reads in flight")); + } + let request = inflight.remove(pick(&inflight)); + events.push(Event::Delivered(*request.segment_id)); + graph.set_io_result(request.id, empty_segment())?; + } + } + } + assert!(inflight.is_empty(), "graph finished with reads in flight"); + Ok(Run { arrays, events }) +} + +fn fifo(_: &[IoRequest]) -> usize { + 0 +} + +/// Delivers segments in the given order. +fn scripted(order: &[u32]) -> impl FnMut(&[IoRequest]) -> usize + '_ { + let mut next = order.iter(); + move |inflight| { + let segment = *next.next().expect("script covers every read"); + inflight + .iter() + .position(|r| *r.segment_id == segment) + .expect("scripted segment is in flight") + } +} + +fn pieces(events: &[Event]) -> Vec { + events + .iter() + .filter_map(|event| match event { + Event::Piece(len) => Some(*len), + _ => None, + }) + .collect() +} + +fn indices(values: impl IntoIterator) -> ArrayRef { + Buffer::from_iter(values).into_array() +} + +/// Joins the arrays a run produced, in order, into one array. +fn join(dtype: &DType, arrays: Vec) -> VortexResult { + Ok(ChunkedArray::try_new(arrays, dtype.clone())?.into_array()) +} + +fn assert_rows(expected: ArrayRef, arrays: Vec) -> VortexResult<()> { + let actual = join(expected.dtype(), arrays)?; + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +fn segment(id: u32) -> Option { + Some(SegmentId::from(id)) +} + +#[rstest] +#[case::one_piece(vec![], false)] +#[case::many_pieces(vec![3, 1, 4, 1, 5], false)] +#[case::yielding(vec![3, 1, 4, 1, 5], true)] +#[case::oversized_pieces(vec![100], true)] +fn source_pieces_come_out_in_order( + #[case] cut: Vec, + #[case] yielding: bool, +) -> VortexResult<()> { + let plan = RowSource::plan(20, cut, None, yielding, false); + let mask = Mask::from_iter((0..10).map(|i| i != 4)); + let run = run(&plan, 5..15, mask, fifo)?; + assert_rows(indices((5..15).filter(|i| *i != 9)), run.arrays) +} + +/// A node runs only when its readiness rule holds, whatever order its inputs finish in. +#[rstest] +#[case::any(Ready::Any)] +#[case::closed_second(Ready::Closed(&[1]))] +#[case::all_closed(Ready::AllClosed)] +fn probe_runs_only_when_ready(#[case] ready: Ready) -> VortexResult<()> { + let runs = Arc::new(Mutex::new(Vec::new())); + let children = vec![ + RowSource::plan(10, vec![2, 2, 2], segment(0), false, false), + RowSource::plan(10, vec![5], segment(1), false, false), + ]; + let plan = Probe::plan(children, ready, Arc::clone(&runs)); + // The first child lands first, so under `Any` the node runs while the second is open. + let run = run(&plan, 0..10, Mask::new_true(10), scripted(&[0, 1]))?; + assert_rows(indices(0..10), run.arrays)?; + + let runs = runs.lock(); + assert!(!runs.is_empty()); + match ready { + Ready::Any => assert!(runs[0] == [true, false], "ran as soon as port 0 closed"), + Ready::Closed(_) => assert!(runs.iter().all(|ports| ports[1])), + Ready::AllClosed => { + assert_eq!(runs.len(), 1); + assert_eq!(runs[0], [true, true]); + } + } + Ok(()) +} + +/// A source that yields between pieces is run once per piece, and its pieces reach the root in +/// order, one per run. +#[test] +fn yielding_source_is_run_once_per_piece() -> VortexResult<()> { + let plan = RowSource::plan(12, vec![4, 4, 4], None, true, false); + let run = run(&plan, 0..12, Mask::new_true(12), fifo)?; + assert_eq!(pieces(&run.events), [4, 4, 4]); + assert_rows(indices(0..12), run.arrays) +} + +#[test] +fn input_take_slices_at_the_boundary() -> VortexResult<()> { + let mut input = Input { + spawned: true, + ..Default::default() + }; + input.push(PrimitiveArray::from_iter(0..4i32).into_array()); + input.push(PrimitiveArray::from_iter(4..10i32).into_array()); + assert_eq!(input.available(), 10); + + let first = input.take(6)?; + assert_eq!(first.iter().map(|a| a.len()).collect::>(), [4, 2]); + assert_eq!(input.available(), 4); + assert!(!input.finished()); + + input.closed = true; + let rest = input.take_all(); + assert_eq!(rest.iter().map(|a| a.len()).collect::>(), [4]); + assert!(input.finished()); + assert!(input.take(1).is_err()); + Ok(()) +} diff --git a/vortex-layout/src/plan/exec/synthetic.rs b/vortex-layout/src/plan/exec/synthetic.rs new file mode 100644 index 00000000000..1d6b74e44ed --- /dev/null +++ b/vortex-layout/src/plan/exec/synthetic.rs @@ -0,0 +1,295 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Plans whose exec nodes are built by hand, for exercising the graph without layouts. +//! +//! A [`SyntheticPlan`] carries a closure that builds its node. The graph treats it like any +//! other operator, so a test or benchmark can wire hand-written sources under the real +//! operators (Pack, Concat, Filter, Eval) and drive the whole thing through [`ExecGraph`], +//! reading segments from memory. +//! +//! [`ExecGraph`]: crate::plan::exec::ExecGraph + +use std::fmt; +use std::ops::Range; +use std::sync::Arc; + +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; +use crate::plan::exec::NodeState; +use crate::plan::exec::Ready; +use crate::plan::exec::StepCx; +use crate::segments::SegmentId; + +/// Builds the node of a [`SyntheticPlan`] for one execution. +pub type NodeBuilder = + dyn Fn(Range, Mask, &ExecContext) -> VortexResult> + Send + Sync; + +/// A plan operator whose node is built by a closure. +#[derive(Clone, Debug)] +pub struct Synthetic; + +/// The node builder of a [`Synthetic`] plan. +#[derive(Clone)] +pub struct SyntheticData { + builder: Arc, +} + +impl fmt::Debug for SyntheticData { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("SyntheticData") + } +} + +/// A plan whose node is built by a closure. +pub type SyntheticPlan = Plan; + +impl SyntheticPlan { + /// A plan of `row_count` rows of `dtype`, with `children`, whose node `builder` makes. + pub fn new( + dtype: DType, + row_count: u64, + children: Vec, + builder: impl Fn(Range, Mask, &ExecContext) -> VortexResult> + + Send + + Sync + + 'static, + ) -> Self { + PlanParts { + vtable: Synthetic, + dtype, + row_count, + children: children.into(), + data: SyntheticData { + builder: Arc::new(builder), + }, + } + .into_typed() + } +} + +impl PlanVTable for Synthetic { + type PlanData = SyntheticData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.synthetic"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + None + } + + fn with_children( + _plan: &Plan, + _children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + Ok(()) + } + + fn exec( + plan: &Plan, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult> { + (plan.data().builder)(rows, mask, ctx) + } +} + +/// The dtype every synthetic source produces: its rows' indices, as `u64`. +pub fn row_dtype() -> DType { + DType::Primitive(PType::U64, Nullability::NonNullable) +} + +/// A source that emits its rows' indices, in order, cut into pieces of the given lengths, +/// optionally after a round trip through the IO service and optionally yielding between +/// pieces. +/// +/// A selected source emits only the rows its mask selects, as every operator does. A dense +/// source ignores its mask and emits every row of its range, as a bare segment scan does, so +/// a filter plan can sit over it. The indices are the plan's own rows, so +/// a parent that reorders or filters can be checked against a sequence. Piece lengths count +/// emitted rows; the last piece takes whatever remains. +pub struct RowSource { + /// Remaining indices to emit, in order. + remaining: Vec, + /// Lengths of the pieces still to emit; the last piece takes whatever remains. + pieces: Vec, + /// A segment to request before emitting anything, standing in for a read. + segment: Option, + /// Whether to return `Yield` between pieces instead of emitting them all at once. + yielding: bool, + requested: bool, +} + +impl RowSource { + /// A source plan of `row_count` rows that emits pieces of the given lengths. + pub fn plan( + row_count: u64, + pieces: Vec, + segment: Option, + yielding: bool, + dense: bool, + ) -> PlanRef { + SyntheticPlan::new(row_dtype(), row_count, Vec::new(), move |rows, mask, _| { + let remaining = if dense || mask.all_true() { + (rows.start..rows.end).collect() + } else { + (rows.start..rows.end) + .zip(mask.iter()) + .filter_map(|(row, selected)| selected.then_some(row)) + .collect() + }; + Ok(Box::new(RowSource { + remaining, + pieces: pieces.iter().rev().copied().collect(), + segment, + yielding, + requested: false, + })) + }) + .into_plan() + } + + /// Emits the next piece, or everything left when no lengths remain. + fn emit_next(&mut self, cx: &mut StepCx<'_>) { + let len = self + .pieces + .pop() + .unwrap_or(self.remaining.len()) + .min(self.remaining.len()); + let piece: Vec = self.remaining.drain(..len).collect(); + cx.emit(Buffer::from(piece).into_array()); + } + + fn emit(&mut self, cx: &mut StepCx<'_>) -> NodeState { + if self.yielding { + self.emit_next(cx); + if self.remaining.is_empty() { + return NodeState::Done; + } + return NodeState::Yield; + } + while !self.remaining.is_empty() { + self.emit_next(cx); + } + NodeState::Done + } +} + +impl ExecNode for RowSource { + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + if self.remaining.is_empty() { + return Ok(NodeState::Done); + } + if let Some(segment) = self.segment { + self.requested = true; + cx.request(segment); + return Ok(NodeState::Wait); + } + Ok(self.emit(cx)) + } + + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + if self.requested { + self.requested = false; + cx.take_delivery() + .ok_or_else(|| vortex_err!("RowSource ran without its delivery"))?; + } + Ok(self.emit(cx)) + } +} + +/// A node with a chosen readiness rule that records the state of its ports every time it runs, +/// then passes its first port through. +/// +/// It spawns one child per plan child, each over the node's own rows and selection. +pub struct Probe { + ready: Ready, + children: Vec, + rows: Range, + mask: Mask, + /// For each compute: whether each port was closed when the node ran. + pub runs: Arc>>>, +} + +impl Probe { + /// A plan over `children` whose node runs under `ready` and logs into `runs`. + pub fn plan( + children: Vec, + ready: Ready, + runs: Arc>>>, + ) -> PlanRef { + let row_count = children.first().map_or(0, |child| child.row_count()); + SyntheticPlan::new( + row_dtype(), + row_count, + children.clone(), + move |rows, mask, _| { + Ok(Box::new(Probe { + ready, + children: children.clone(), + rows, + mask, + runs: Arc::clone(&runs), + })) + }, + ) + .into_plan() + } +} + +impl ExecNode for Probe { + fn ready(&self) -> Ready { + self.ready + } + + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + for (port, child) in self.children.iter().enumerate() { + cx.spawn(port, child.clone(), self.rows.clone(), self.mask.clone()); + } + Ok(NodeState::Wait) + } + + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + let closed = cx.inputs().iter().map(|input| input.closed()).collect(); + self.runs.lock().push(closed); + for array in cx.input(0).take_all() { + cx.emit(array); + } + for port in 1..self.children.len() { + cx.input(port).take_all(); + } + if cx.all_finished() { + return Ok(NodeState::Done); + } + Ok(NodeState::Wait) + } +} + +/// An in-memory segment store that answers any segment id with an empty buffer, for sources +/// whose read is only a scheduling event. +pub fn empty_segment() -> BufferHandle { + BufferHandle::new_host(Buffer::::empty().into_byte_buffer()) +} diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 802e210b07e..b4782481b49 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -9,6 +9,7 @@ mod children; mod display; +pub mod exec; mod lower; mod optimize; pub mod optimizer; diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/typed.rs index bf544801263..625ec127165 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/typed.rs @@ -9,18 +9,23 @@ 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::SerializeMetadata; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::Mask; use crate::plan::PlanChildren; use crate::plan::PlanId; use crate::plan::PlanVTable; use crate::plan::display::PlanTreeDisplay; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; /// The combined allocation behind [`PlanRef`]. /// @@ -126,6 +131,27 @@ impl PlanRef { Some(unsafe { &*(std::ptr::from_ref(self).cast::>()) }) } + /// Builds the exec node that runs this plan over `rows` of its row domain, restricted to + /// `mask`, with the graph's `ctx`. + pub fn exec( + &self, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult> { + vortex_ensure!( + rows.start <= rows.end && rows.end <= self.row_count(), + "Exec rows {rows:?} exceed plan row count {}", + self.row_count() + ); + vortex_ensure!( + mask.len() as u64 == rows.end - rows.start, + "Exec mask length {} does not match rows {rows:?}", + mask.len() + ); + self.dyn_plan().dyn_exec(self, rows, mask, ctx) + } + /// Displays this plan and its descendants with the default plan extractors. pub fn display_tree(&self) -> PlanTreeDisplay<'_> { PlanTreeDisplay::default_display(self) @@ -337,6 +363,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>; + + /// Builds the exec node for this operator. + fn dyn_exec( + &self, + plan: &PlanRef, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult>; } impl DynPlan for PlanData { @@ -368,4 +403,14 @@ impl DynPlan for PlanData { fn dyn_metadata(&self, plan: &PlanRef) -> Option> { V::metadata(plan.as_::()).map(SerializeMetadata::serialize) } + + fn dyn_exec( + &self, + plan: &PlanRef, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult> { + V::exec(plan.as_::(), rows, mask, ctx) + } } diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/vtable.rs index caf004383b5..416fbf6ce50 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::SerializeMetadata; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; use vortex_session::registry::Id; use crate::plan::PlanChildren; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; use crate::plan::typed::Plan; /// A unique identifier for a plan operator. @@ -68,4 +73,20 @@ pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { let _ = plan; Cow::Owned(format!("child[{index}]")) } + + /// Builds the push-based exec node that runs this plan over `rows` of its row domain, + /// restricted to `mask`, with the graph's `ctx`. `rows` lies within the plan's row domain + /// and `mask` is as long as `rows`. + /// + /// Construction does no IO. Children are spawned and requests published when the graph + /// starts the node. [`ExecNode`] says what the node must then do. + fn exec( + plan: &Plan, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult> { + drop((rows, mask, ctx)); + vortex_bail!("Plan {} has no exec implementation", plan.vtable().id()) + } } From c31fbcf6c1703f84290fe56362f6e31a7805db08 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 8 Oct 2026 20:14:36 +0000 Subject: [PATCH 03/20] feat(layout): add SegmentScan, Filter, Concat, and Pack exec nodes Implement `PlanVTable::exec` for the operators a projection of a struct of chunked columns needs: - `SegmentScanNode` requests its segment at start, decodes the delivered bytes, and emits the node's rows as one dense array. A `Filter` over a scan fuses into the same node with a filter mask, so the array holds exactly the selected rows. Decoded segments are shared through the graph's `DecodeCache`. - `FilterNode` keeps the selected rows of any other child's dense arrays as they stream through, tracking its position with a cursor. - `ConcatNode` spawns only the chunks overlapping its rows, one port each in chunk order, and emits the current chunk's rows as they arrive. Chunks whose reads land early wait in their ports. - `PackNode` zips its fields as their rows arrive: whenever every port has rows, it takes as many as the shortest holds from each and emits one struct. Fields chunked alike stream one struct per chunk. The layout tests drive lowered plans over in-memory segments with scripted IO ordering and check what reaches the root. The scheduling tests add hand-built sources under the real operators, including random trees over random views with reads completing in random order, checked against a model of the same tree. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/exec/concat.rs | 97 ++++ vortex-layout/src/plan/exec/filter.rs | 98 ++++ vortex-layout/src/plan/exec/mod.rs | 13 + vortex-layout/src/plan/exec/pack.rs | 111 ++++ .../src/plan/exec/scheduling_tests.rs | 324 ++++++++++- vortex-layout/src/plan/exec/segment_scan.rs | 118 ++++ vortex-layout/src/plan/exec/selection.rs | 55 ++ vortex-layout/src/plan/exec/synthetic.rs | 2 +- vortex-layout/src/plan/exec/tests.rs | 548 ++++++++++++++++++ vortex-layout/src/plan/plans/concat.rs | 18 + vortex-layout/src/plan/plans/filter.rs | 33 ++ vortex-layout/src/plan/plans/pack.rs | 18 + vortex-layout/src/plan/plans/segment_scan.rs | 22 + 13 files changed, 1454 insertions(+), 3 deletions(-) create mode 100644 vortex-layout/src/plan/exec/concat.rs create mode 100644 vortex-layout/src/plan/exec/filter.rs create mode 100644 vortex-layout/src/plan/exec/pack.rs create mode 100644 vortex-layout/src/plan/exec/segment_scan.rs create mode 100644 vortex-layout/src/plan/exec/selection.rs create mode 100644 vortex-layout/src/plan/exec/tests.rs diff --git a/vortex-layout/src/plan/exec/concat.rs b/vortex-layout/src/plan/exec/concat.rs new file mode 100644 index 00000000000..e6dbd4ec00d --- /dev/null +++ b/vortex-layout/src/plan/exec/concat.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_error::VortexResult; + +use crate::plan::ConcatPlan; +use crate::plan::exec::ExecNode; +use crate::plan::exec::NodeState; +use crate::plan::exec::StepCx; +use crate::plan::exec::selection::Selection; + +/// Emits its chunks' rows in chunk order. +/// +/// Only chunks overlapping the selection's rows are spawned, one port each in chunk order, and +/// a chunk whose slice of the mask is all false is skipped. Chunks are read independently and +/// may finish in any order; the node emits whatever the current chunk has produced, moves on +/// once that chunk has finished, and leaves what later chunks produced early waiting in their +/// ports. Nothing above it sees the disorder. +pub(crate) struct ConcatNode { + plan: ConcatPlan, + selection: Selection, + /// The port whose chunk is being emitted. Ports are numbered in chunk order. + current: usize, + ports: usize, +} + +impl ConcatNode { + pub(crate) fn new(plan: ConcatPlan, selection: Selection) -> Self { + Self { + plan, + selection, + current: 0, + ports: 0, + } + } + + fn chunk_rows(&self, index: usize) -> Range { + let offsets = self.plan.row_offsets(); + let end = offsets + .get(index + 1) + .copied() + .unwrap_or_else(|| self.plan.row_count()); + offsets[index]..end + } +} + +impl ExecNode for ConcatNode { + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + let rows = self.selection.rows().clone(); + let offsets = self.plan.row_offsets(); + // Every split visits a small part of a file; skip the chunks before and after it. + let first = offsets + .partition_point(|&offset| offset <= rows.start) + .saturating_sub(1); + let end = offsets.partition_point(|&offset| offset < rows.end); + for index in first..end { + let chunk = self.chunk_rows(index); + let local = rows.start.max(chunk.start)..rows.end.min(chunk.end); + if local.start >= local.end { + continue; + } + let mask = self.selection.slice(&local)?; + if mask.all_false() { + continue; + } + cx.spawn( + self.ports, + self.plan.child_required(index)?, + local.start - chunk.start..local.end - chunk.start, + mask, + ); + self.ports += 1; + } + if self.ports == 0 { + return Ok(NodeState::Done); + } + Ok(NodeState::Wait) + } + + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + while self.current < self.ports { + let input = cx.input(self.current); + let arrays = input.take_all(); + let finished = input.finished(); + for array in arrays { + cx.emit(array); + } + if !finished { + return Ok(NodeState::Wait); + } + self.current += 1; + } + Ok(NodeState::Done) + } +} diff --git a/vortex-layout/src/plan/exec/filter.rs b/vortex-layout/src/plan/exec/filter.rs new file mode 100644 index 00000000000..e2294c074e1 --- /dev/null +++ b/vortex-layout/src/plan/exec/filter.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::plan::FilterPlan; +use crate::plan::exec::ExecNode; +use crate::plan::exec::NodeState; +use crate::plan::exec::StepCx; +use crate::plan::exec::selection::Selection; + +const CHILD: usize = 0; + +/// The selected fraction of an array at or above which a predicate runs over the whole array. +const EXPR_EVAL_THRESHOLD: f64 = 0.2; + +/// Keeps the selected rows of the dense arrays its child produces. +/// +/// The child runs over the same rows with the selection as its care hint, and returns every row +/// of its range in order. A cursor tracks how far along the range the arrays taken so far reach, +/// so each one is filtered by its own slice of the mask as it arrives. +pub(crate) struct FilterNode { + plan: FilterPlan, + selection: Selection, + session: VortexSession, + /// Rows of the child's range consumed so far. + cursor: usize, +} + +impl FilterNode { + pub(crate) fn new(plan: FilterPlan, selection: Selection, session: VortexSession) -> Self { + Self { + plan, + selection, + session, + cursor: 0, + } + } +} + +/// Keeps the rows of `array` that `mask` selects. `predicate` says the array is a predicate's +/// result. +pub(crate) fn keep_selected( + array: ArrayRef, + mask: Mask, + predicate: bool, + session: &VortexSession, +) -> VortexResult { + if mask.all_true() { + return Ok(array); + } + if predicate && mask.density() >= EXPR_EVAL_THRESHOLD { + // A predicate over a mostly selected array runs over every row and its result is + // filtered, as the default scan's flat reader does. Filtering lazily would push the + // filter back through the predicate onto the encoded input, which for some encodings + // costs more than the predicate. + let mut ctx = session.create_execution_ctx(); + return array + .execute::(&mut ctx)? + .into_array() + .filter(mask); + } + array.filter(mask) +} + +impl ExecNode for FilterNode { + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + cx.spawn( + CHILD, + self.plan.child_plan()?, + self.selection.rows().clone(), + self.selection.mask().clone(), + ); + Ok(NodeState::Wait) + } + + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + let predicate = self.plan.dtype().is_boolean(); + for array in cx.input(CHILD).take_all() { + let mask = self + .selection + .mask() + .slice(self.cursor..self.cursor + array.len()); + self.cursor += array.len(); + cx.emit(keep_selected(array, mask, predicate, &self.session)?); + } + if cx.input(CHILD).finished() { + return Ok(NodeState::Done); + } + Ok(NodeState::Wait) + } +} diff --git a/vortex-layout/src/plan/exec/mod.rs b/vortex-layout/src/plan/exec/mod.rs index 5cb10f6e569..127f5898b8d 100644 --- a/vortex-layout/src/plan/exec/mod.rs +++ b/vortex-layout/src/plan/exec/mod.rs @@ -36,6 +36,11 @@ //! merge pipelines, and nodes that wait for a port to close, [`Ready::Closed`] and //! [`Ready::AllClosed`], are barriers between them, as a join build is in a query engine. +mod concat; +mod filter; +mod pack; +mod segment_scan; +mod selection; pub mod synthetic; use std::collections::VecDeque; @@ -668,5 +673,13 @@ impl From> for Effects { } } +pub(crate) use concat::ConcatNode; +pub(crate) use filter::FilterNode; +pub(crate) use pack::PackNode; +pub(crate) use segment_scan::SegmentScanNode; +pub(crate) use selection::Selection; + #[cfg(test)] mod scheduling_tests; +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/plan/exec/pack.rs b/vortex-layout/src/plan/exec/pack.rs new file mode 100644 index 00000000000..cedec0eb215 --- /dev/null +++ b/vortex-layout/src/plan/exec/pack.rs @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::plan::PackPlan; +use crate::plan::exec::ExecNode; +use crate::plan::exec::NodeState; +use crate::plan::exec::StepCx; +use crate::plan::exec::selection::Selection; +use crate::plan::exec::selection::join; + +/// Zips its fields into struct arrays as their rows arrive. +/// +/// Every field runs over the same rows and selection, so the fields' outputs have the same +/// length and line up row for row. Whenever every port has rows available, the node takes as +/// many rows as the shortest port holds from each of them and emits one struct. Fields chunked +/// alike therefore stream one struct per chunk; fields chunked differently are sliced at the +/// boundaries they share. +pub(crate) struct PackNode { + plan: PackPlan, + selection: Selection, + ports: usize, +} + +impl PackNode { + pub(crate) fn new(plan: PackPlan, selection: Selection) -> Self { + let ports = plan.children().len(); + Self { + plan, + selection, + ports, + } + } + + /// Builds a struct of `len` rows from `arrays`, one per port, the last being the validity + /// when the struct is nullable. + fn assemble(&self, mut arrays: Vec, len: usize) -> VortexResult { + let validity = if self.plan.dtype().is_nullable() { + Validity::Array( + arrays + .pop() + .ok_or_else(|| vortex_err!("Nullable Pack is missing its validity port"))?, + ) + } else { + Validity::NonNullable + }; + StructArray::try_new_with_dtype(arrays, self.plan.fields().clone(), len, validity) + } + + fn field_dtype(&self, port: usize) -> DType { + self.plan + .fields() + .field_by_index(port) + .unwrap_or(DType::Bool(Nullability::NonNullable)) + } +} + +impl ExecNode for PackNode { + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + if self.ports == 0 { + // A struct with no fields still has rows, so no child can carry them. + let len = self.selection.mask().true_count(); + cx.emit(self.assemble(Vec::new(), len)?.into_array()); + return Ok(NodeState::Done); + } + if self.selection.mask().all_false() { + return Ok(NodeState::Done); + } + for port in 0..self.ports { + cx.spawn( + port, + self.plan.child_required(port)?, + self.selection.rows().clone(), + self.selection.mask().clone(), + ); + } + Ok(NodeState::Wait) + } + + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + loop { + let len = cx + .inputs() + .iter() + .map(|input| input.available()) + .min() + .unwrap_or(0); + if len == 0 { + break; + } + let mut arrays = Vec::with_capacity(self.ports); + for port in 0..self.ports { + let dtype = self.field_dtype(port); + arrays.push(join(&dtype, cx.input(port).take(len)?)?); + } + cx.emit(self.assemble(arrays, len)?.into_array()); + } + if cx.all_finished() { + return Ok(NodeState::Done); + } + Ok(NodeState::Wait) + } +} diff --git a/vortex-layout/src/plan/exec/scheduling_tests.rs b/vortex-layout/src/plan/exec/scheduling_tests.rs index 6fb2ee452f7..afe0fe27cc2 100644 --- a/vortex-layout/src/plan/exec/scheduling_tests.rs +++ b/vortex-layout/src/plan/exec/scheduling_tests.rs @@ -3,8 +3,8 @@ //! Tests of the graph's scheduling, driven through hand-built nodes over in-memory segments. //! -//! Every source emits its own row indices, so what reaches the root can be checked against a -//! sequence. +//! Every source emits its own row indices, so whatever the tree does to them can be checked +//! against a model of the same tree built from plain arrays. #![allow(clippy::cast_possible_truncation)] @@ -12,14 +12,21 @@ use std::ops::Range; use std::sync::Arc; use parking_lot::Mutex; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -28,7 +35,11 @@ use vortex_mask::Mask; use super::synthetic::Probe; use super::synthetic::RowSource; use super::synthetic::empty_segment; +use super::synthetic::row_dtype; use super::*; +use crate::plan::ConcatPlan; +use crate::plan::FilterPlan; +use crate::plan::PackPlan; use crate::test::SESSION; #[derive(Debug, PartialEq, Eq)] @@ -94,6 +105,10 @@ fn fifo(_: &[IoRequest]) -> usize { 0 } +fn lifo(inflight: &[IoRequest]) -> usize { + inflight.len() - 1 +} + /// Delivers segments in the given order. fn scripted(order: &[u32]) -> impl FnMut(&[IoRequest]) -> usize + '_ { let mut next = order.iter(); @@ -136,6 +151,39 @@ fn segment(id: u32) -> Option { Some(SegmentId::from(id)) } +fn fields(dtypes: impl IntoIterator) -> StructFields { + StructFields::from_iter( + dtypes + .into_iter() + .enumerate() + .map(|(i, dtype)| (format!("f{i}"), dtype)), + ) +} + +fn struct_of(arrays: Vec, len: usize) -> VortexResult { + let names = fields(arrays.iter().map(|array| array.dtype().clone())) + .names() + .clone(); + Ok(StructArray::try_new(names, arrays, len, Validity::NonNullable)?.into_array()) +} + +fn pack(children: Vec) -> VortexResult { + let row_count = children[0].row_count(); + Ok(PackPlan::try_new( + fields(children.iter().map(|child| child.dtype().clone())), + Nullability::NonNullable, + row_count, + children, + None, + )? + .into_plan()) +} + +fn concat(children: Vec) -> VortexResult { + let dtype = children[0].dtype().clone(); + Ok(ConcatPlan::try_new(dtype, children)?.into_plan()) +} + #[rstest] #[case::one_piece(vec![], false)] #[case::many_pieces(vec![3, 1, 4, 1, 5], false)] @@ -151,6 +199,62 @@ fn source_pieces_come_out_in_order( assert_rows(indices((5..15).filter(|i| *i != 9)), run.arrays) } +/// Chunks whose reads land in reverse order come out in chunk order, and nothing comes out +/// before the first chunk has landed. +#[rstest] +#[case::fifo(&[0, 1, 2, 3])] +#[case::lifo(&[3, 2, 1, 0])] +#[case::first_last(&[1, 2, 3, 0])] +fn concat_reorders_out_of_order_reads(#[case] order: &[u32]) -> VortexResult<()> { + let chunks = (0..4) + .map(|i| RowSource::plan(10, vec![4, 6], segment(i), false, false)) + .collect(); + let plan = concat(chunks)?; + let run = run(&plan, 0..40, Mask::new_true(40), scripted(order))?; + + let first_piece = run + .events + .iter() + .position(|e| matches!(e, Event::Piece(_))) + .expect("something came out"); + let first_chunk_landed = run + .events + .iter() + .position(|e| *e == Event::Delivered(0)) + .expect("chunk 0 landed"); + assert!(first_piece > first_chunk_landed); + // Each chunk's own rows are local to the chunk. + assert_rows(indices((0..4).flat_map(|_| 0..10)), run.arrays) +} + +/// Fields cut at different places zip into structs at the boundaries they share, and the +/// structs come out in row order under any delivery order. +#[rstest] +fn pack_zips_misaligned_fields(#[values(true, false)] reverse: bool) -> VortexResult<()> { + let a = concat(vec![ + RowSource::plan(7, vec![], segment(0), false, false), + RowSource::plan(13, vec![5], segment(1), false, false), + ])?; + let b = RowSource::plan(20, vec![10], segment(2), false, false); + let c = concat(vec![ + RowSource::plan(3, vec![], segment(3), false, false), + RowSource::plan(17, vec![], segment(4), false, false), + ])?; + let plan = pack(vec![a, b, c])?; + let pick: fn(&[IoRequest]) -> usize = if reverse { lifo } else { fifo }; + let run = run(&plan, 0..20, Mask::new_true(20), pick)?; + + let expected = struct_of( + vec![ + indices((0..7).chain(0..13)), + indices(0..20), + indices((0..3).chain(0..17)), + ], + 20, + )?; + assert_rows(expected, run.arrays) +} + /// A node runs only when its readiness rule holds, whatever order its inputs finish in. #[rstest] #[case::any(Ready::Any)] @@ -180,6 +284,18 @@ fn probe_runs_only_when_ready(#[case] ready: Ready) -> VortexResult<()> { Ok(()) } +/// A filter over a dense source keeps the selected rows as the pieces stream through, however +/// the source cuts them. +#[rstest] +#[case::one_piece(vec![])] +#[case::odd_cuts(vec![1, 3, 2, 7])] +fn filter_streams_dense_pieces(#[case] cut: Vec) -> VortexResult<()> { + let plan = FilterPlan::new(RowSource::plan(20, cut, None, true, true)).into_plan(); + let mask = Mask::from_iter((0..14).map(|i| i % 3 == 0)); + let run = run(&plan, 3..17, mask.clone(), fifo)?; + assert_rows(indices(3..17).filter(mask)?, run.arrays) +} + /// A source that yields between pieces is run once per piece, and its pieces reach the root in /// order, one per run. #[test] @@ -190,6 +306,210 @@ fn yielding_source_is_run_once_per_piece() -> VortexResult<()> { assert_rows(indices(0..12), run.arrays) } +/// A tree of operators over sources, mirrored as plain arrays. +#[derive(Debug)] +enum Tree { + /// A source of this many rows, cut into these pieces, read from this segment, yielding. + Source { + rows: u64, + cut: Vec, + segment: Option, + yielding: bool, + }, + Concat(Vec), + Pack(Vec), + /// A filter over a dense source. + Filter(Box), +} + +impl Tree { + fn rows(&self) -> u64 { + match self { + Tree::Source { rows, .. } => *rows, + Tree::Concat(children) => children.iter().map(Tree::rows).sum(), + Tree::Pack(children) => children[0].rows(), + Tree::Filter(child) => child.rows(), + } + } + + fn dtype(&self) -> DType { + match self { + Tree::Pack(children) => DType::Struct( + fields(children.iter().map(Tree::dtype)), + Nullability::NonNullable, + ), + _ => row_dtype(), + } + } + + fn plan(&self, dense: bool) -> VortexResult { + Ok(match self { + Tree::Source { + rows, + cut, + segment, + yielding, + } => RowSource::plan( + *rows, + cut.clone(), + segment.map(SegmentId::from), + *yielding, + dense, + ), + Tree::Concat(children) => concat( + children + .iter() + .map(|child| child.plan(false)) + .collect::>()?, + )?, + Tree::Pack(children) => pack( + children + .iter() + .map(|child| child.plan(false)) + .collect::>()?, + )?, + Tree::Filter(child) => FilterPlan::new(child.plan(true)?).into_plan(), + }) + } + + /// The selected rows of `rows`, as the tree's plan should produce them. + fn expected(&self, rows: Range, mask: &Mask) -> VortexResult { + match self { + Tree::Source { .. } | Tree::Filter(_) => indices(rows).filter(mask.clone()), + Tree::Concat(children) => { + let mut parts = Vec::new(); + let mut offset = 0; + for child in children { + let chunk = offset..offset + child.rows(); + offset = chunk.end; + let local = rows.start.max(chunk.start)..rows.end.min(chunk.end); + if local.start >= local.end { + continue; + } + let mask = mask.slice( + (local.start - rows.start) as usize..(local.end - rows.start) as usize, + ); + parts.push( + child + .expected(local.start - chunk.start..local.end - chunk.start, &mask)?, + ); + } + join(&children[0].dtype(), parts) + } + Tree::Pack(children) => { + let arrays = children + .iter() + .map(|child| child.expected(rows.clone(), mask)) + .collect::>>()?; + struct_of(arrays, mask.true_count()) + } + } + } +} + +/// A random tree. Sources get distinct segments so delivery order can be scripted. +fn random_tree(rng: &mut StdRng, depth: usize, rows: Option, next_segment: &mut u32) -> Tree { + let kind = if depth == 0 { + Kind::Source + } else { + [Kind::Source, Kind::Concat, Kind::Pack, Kind::Filter][rng.random_range(0..4)] + }; + random_tree_of(rng, kind, depth, rows, next_segment) +} + +#[derive(Clone, Copy)] +enum Kind { + Source, + Concat, + Pack, + Filter, +} + +fn random_tree_of( + rng: &mut StdRng, + kind: Kind, + depth: usize, + rows: Option, + next_segment: &mut u32, +) -> Tree { + let rows = rows.unwrap_or_else(|| rng.random_range(1..24)); + match kind { + Kind::Concat => { + // Chunks summing to `rows`. Chunks of one column are leaves: a source, or a filter + // over one, so every chunk has the column's dtype. + let mut children = Vec::new(); + let mut left = rows; + while left > 0 { + let chunk = rng.random_range(1..=left); + let leaf = if rng.random_bool(0.3) { + Kind::Filter + } else { + Kind::Source + }; + children.push(random_tree_of(rng, leaf, 0, Some(chunk), next_segment)); + left -= chunk; + } + Tree::Concat(children) + } + Kind::Pack => Tree::Pack( + (0..rng.random_range(1..4)) + .map(|_| random_tree(rng, depth - 1, Some(rows), next_segment)) + .collect(), + ), + Kind::Filter => Tree::Filter(Box::new(random_tree_of( + rng, + Kind::Source, + 0, + Some(rows), + next_segment, + ))), + Kind::Source => { + let segment = rng.random_bool(0.7).then(|| { + *next_segment += 1; + *next_segment - 1 + }); + let cut = (0..rng.random_range(0..4)) + .map(|_| rng.random_range(1..8)) + .collect(); + Tree::Source { + rows, + cut, + segment, + yielding: rng.random_bool(0.3), + } + } + } +} + +/// Random trees over random views, with reads completing in random order, come out as the +/// model says. +#[test] +fn random_trees_match_their_model() -> VortexResult<()> { + let mut rng = StdRng::seed_from_u64(7); + for _ in 0..300 { + let mut next_segment = 0; + let tree = random_tree(&mut rng, 3, None, &mut next_segment); + let total = tree.rows(); + let start = rng.random_range(0..=total); + let end = rng.random_range(start..=total); + let density = rng.random_range(0.0..=1.0); + let mask = Mask::from_iter((start..end).map(|_| rng.random_bool(density))); + + let plan = tree.plan(false)?; + let mut order = StdRng::seed_from_u64(rng.random()); + let run = run(&plan, start..end, mask.clone(), |inflight| { + order.random_range(0..inflight.len()) + }) + .map_err(|e| vortex_err!("{tree:?} over {start}..{end}: {e}"))?; + + let expected = tree.expected(start..end, &mask)?; + assert_eq!(expected.dtype(), &tree.dtype()); + assert_rows(expected, run.arrays) + .map_err(|e| vortex_err!("{tree:?} over {start}..{end}: {e}"))?; + } + Ok(()) +} + #[test] fn input_take_slices_at_the_boundary() -> VortexResult<()> { let mut input = Input { diff --git a/vortex-layout/src/plan/exec/segment_scan.rs b/vortex-layout/src/plan/exec/segment_scan.rs new file mode 100644 index 00000000000..17c1ff0508d --- /dev/null +++ b/vortex-layout/src/plan/exec/segment_scan.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::buffer::BufferHandle; +use vortex_array::serde::SerializedArray; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::plan::SegmentScanPlan; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; +use crate::plan::exec::NodeState; +use crate::plan::exec::StepCx; +use crate::plan::exec::selection::Selection; + +/// Reads one segment, then emits its rows as a single array. +/// +/// Two masks cover the node's rows: +/// +/// - The selection is always given and says which rows the parent cares about. It is a hint: rows +/// outside it may be skipped or decoded, and the node may ignore it entirely. +/// - The optional filter says which rows to return. With a filter the array holds exactly the +/// filter's rows, in order; without one it holds every row, dense, and values at rows outside +/// the selection are unspecified. A filter only selects rows the selection cares about. +/// +/// `start` publishes the read, unless the filter selects nothing or the segment is already +/// decoded, and the compute that receives the bytes decodes them, slices them to the node's +/// rows, and applies the filter. +pub(crate) struct SegmentScanNode { + plan: SegmentScanPlan, + selection: Selection, + filter: Option, + ctx: ExecContext, +} + +impl SegmentScanNode { + pub(crate) fn try_new( + plan: SegmentScanPlan, + selection: Selection, + filter: Option, + ctx: ExecContext, + ) -> VortexResult { + if let Some(filter) = &filter { + vortex_ensure!( + filter.len() == selection.mask().len(), + "SegmentScan filter of length {} does not cover rows {:?}", + filter.len(), + selection.rows() + ); + } + Ok(Self { + plan, + selection, + filter, + ctx, + }) + } + + /// Decodes the whole segment. + fn decode(&self, segment: BufferHandle) -> VortexResult { + let serialized = match self.plan.array_tree() { + Some(tree) => SerializedArray::from_flatbuffer_and_segment(tree.clone(), segment)?, + None => SerializedArray::try_from(segment)?, + }; + let row_count = usize::try_from(self.plan.row_count())?; + serialized.decode( + self.plan.dtype(), + row_count, + self.plan.array_ctx(), + self.ctx.session(), + ) + } + + /// Slices the whole decoded segment to the node's rows, filtered when the node has a filter. + fn select(&self, array: ArrayRef) -> VortexResult { + let rows = self.selection.rows(); + let mut array = if rows.start == 0 && rows.end == self.plan.row_count() { + array + } else { + array.slice(usize::try_from(rows.start)?..usize::try_from(rows.end)?)? + }; + if let Some(filter) = &self.filter + && !filter.all_true() + { + array = array.filter(filter.clone())?; + } + Ok(array) + } +} + +impl ExecNode for SegmentScanNode { + fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + if self.filter.as_ref().is_some_and(Mask::all_false) { + return Ok(NodeState::Done); + } + if let Some(array) = self.ctx.decoded().get(self.plan.segment_id()) { + cx.emit(self.select(array)?); + return Ok(NodeState::Done); + } + cx.request(self.plan.segment_id()); + Ok(NodeState::Wait) + } + + fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { + let segment = cx + .take_delivery() + .ok_or_else(|| vortex_err!("SegmentScan ran without its delivery"))?; + let array = self.decode(segment)?; + self.ctx + .decoded() + .insert(self.plan.segment_id(), array.clone()); + cx.emit(self.select(array)?); + Ok(NodeState::Done) + } +} diff --git a/vortex-layout/src/plan/exec/selection.rs b/vortex-layout/src/plan/exec/selection.rs new file mode 100644 index 00000000000..ce6f1c34033 --- /dev/null +++ b/vortex-layout/src/plan/exec/selection.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +/// A node's selection over `rows` of its plan domain, indexed by plan row. +#[derive(Clone, Debug)] +pub(crate) struct Selection { + rows: Range, + mask: Mask, +} + +impl Selection { + pub(crate) fn try_new(rows: Range, mask: Mask) -> VortexResult { + vortex_ensure!( + rows.start <= rows.end && mask.len() as u64 == rows.end - rows.start, + "Mask of length {} does not cover rows {rows:?}", + mask.len() + ); + Ok(Self { rows, mask }) + } + + pub(crate) fn rows(&self) -> &Range { + &self.rows + } + + pub(crate) fn mask(&self) -> &Mask { + &self.mask + } + + /// The mask over `rows`, which must lie within this selection. + pub(crate) fn slice(&self, rows: &Range) -> VortexResult { + let start = usize::try_from(rows.start - self.rows.start)?; + let end = usize::try_from(rows.end - self.rows.start)?; + Ok(self.mask.slice(start..end)) + } +} + +/// Joins arrays covering consecutive rows into one. An empty list joins to an empty array. +pub(crate) fn join(dtype: &DType, mut arrays: Vec) -> VortexResult { + match arrays.len() { + 0 => Ok(Canonical::empty(dtype).into_array()), + 1 => Ok(arrays.remove(0)), + _ => Ok(ChunkedArray::try_new(arrays, dtype.clone())?.into_array()), + } +} diff --git a/vortex-layout/src/plan/exec/synthetic.rs b/vortex-layout/src/plan/exec/synthetic.rs index 1d6b74e44ed..a32b4350149 100644 --- a/vortex-layout/src/plan/exec/synthetic.rs +++ b/vortex-layout/src/plan/exec/synthetic.rs @@ -128,7 +128,7 @@ pub fn row_dtype() -> DType { /// /// A selected source emits only the rows its mask selects, as every operator does. A dense /// source ignores its mask and emits every row of its range, as a bare segment scan does, so -/// a filter plan can sit over it. The indices are the plan's own rows, so +/// a [`Filter`](crate::plan::Filter) can sit over it. The indices are the plan's own rows, so /// a parent that reorders or filters can be checked against a sequence. Piece lengths count /// emitted rows; the last piece takes whatever remains. pub struct RowSource { diff --git a/vortex-layout/src/plan/exec/tests.rs b/vortex-layout/src/plan/exec/tests.rs new file mode 100644 index 00000000000..390a6e9734c --- /dev/null +++ b/vortex-layout/src/plan/exec/tests.rs @@ -0,0 +1,548 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +// Fixtures use a 20-row domain and name columns after their struct fields. +#![allow(clippy::cast_possible_truncation, clippy::many_single_char_names)] + +use std::ops::Range; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::serde::SerializeOptions; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::LayoutRef; +use crate::OwnedLayoutChildren; +use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::flat::FlatLayout; +use crate::layouts::struct_::StructLayout; +use crate::plan::Filter; +use crate::plan::SegmentScan; +use crate::plan::exec::selection::join; +use crate::plan::lower; +use crate::test::SESSION; + +const ROWS: u64 = 20; + +/// In-memory segments, standing in for the IO service. +#[derive(Default)] +struct Store { + segments: Vec, +} + +impl Store { + fn flat(&mut self, array: &ArrayRef) -> VortexResult { + let ctx = ArrayContext::empty(); + let buffers = array.serialize( + &ctx, + &SESSION, + &SerializeOptions { + offset: 0, + include_padding: true, + }, + )?; + let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64)); + for buffer in buffers { + bytes.extend_from_slice(buffer.as_ref()); + } + let segment_id = SegmentId::from(u32::try_from(self.segments.len())?); + self.segments.push(BufferHandle::new_host(bytes.freeze())); + Ok(FlatLayout::new( + array.len() as u64, + array.dtype().clone(), + segment_id, + ReadContext::new(ctx.to_ids()), + ) + .into_layout()) + } + + fn chunked(&mut self, array: &ArrayRef, sizes: &[usize]) -> VortexResult { + let mut chunks = Vec::with_capacity(sizes.len()); + let mut start = 0; + for size in sizes { + chunks.push(self.flat(&array.slice(start..start + size)?)?); + start += size; + } + assert_eq!(start, array.len(), "chunk sizes must cover the array"); + Ok(ChunkedLayout::new( + array.len() as u64, + array.dtype().clone(), + OwnedLayoutChildren::layout_children(chunks), + ) + .into_layout()) + } + + fn read(&self, segment_id: SegmentId) -> BufferHandle { + self.segments[*segment_id as usize].clone() + } +} + +/// A struct whose columns are chunked at different row boundaries. +/// +/// ```text +/// a i32 chunks [0,7) [7,12) [12,20) +/// b i64 chunks [0,10) [10,20) +/// c utf8 flat [0,20) +/// d i32? chunks [0,1) [1,3) [3,6) [6,20) +/// e struct x: chunks [0,5) [5,20), y: flat +/// ``` +fn fixture(store: &mut Store) -> VortexResult<(PlanRef, ArrayRef)> { + let a = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let b = PrimitiveArray::from_iter((0..ROWS as i64).map(|v| v * 10)).into_array(); + let c = VarBinViewArray::from_iter_str((0..ROWS).map(|v| format!("row-{v}"))).into_array(); + let d = PrimitiveArray::from_option_iter((0..ROWS as i32).map(|v| (v % 3 != 0).then_some(v))) + .into_array(); + let x = PrimitiveArray::from_iter((0..ROWS).map(|v| v as u8)).into_array(); + let y = BoolArray::from_iter((0..ROWS).map(|v| v % 2 == 0)).into_array(); + let e = StructArray::from_fields(&[("x", x.clone()), ("y", y.clone())])?.into_array(); + let expected = StructArray::from_fields(&[ + ("a", a.clone()), + ("b", b.clone()), + ("c", c.clone()), + ("d", d.clone()), + ("e", e.clone()), + ])? + .into_array(); + + let e_layout = StructLayout::new( + ROWS, + e.dtype().clone(), + vec![store.chunked(&x, &[5, 15])?, store.flat(&y)?], + ) + .into_layout(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![ + store.chunked(&a, &[7, 5, 8])?, + store.chunked(&b, &[10, 10])?, + store.flat(&c)?, + store.chunked(&d, &[1, 2, 3, 14])?, + e_layout, + ], + ) + .into_layout(); + Ok((lower(&layout)?, expected)) +} + +/// Which outstanding request the scripted IO service completes next. +#[derive(Clone, Copy, Debug)] +enum Delivery { + /// Oldest request first. + Fifo, + /// Newest request first. + Lifo, +} + +#[derive(Debug, PartialEq, Eq)] +enum Event { + /// A batch returned by `compute`, as segment ids. + Io(Vec), + /// A result delivered for this segment. + Delivered(u32), + /// An array of this many rows reached the root. + Piece(usize), +} + +struct Run { + arrays: Vec, + events: Vec, +} + +/// Drives a graph the way an owner does, completing reads only while the graph waits. +fn run( + store: &Store, + plan: &PlanRef, + rows: Range, + mask: Mask, + pick: impl FnMut(&[IoRequest]) -> usize, +) -> VortexResult { + run_with(store, plan, rows, mask, pick, DecodeCache::default()) +} + +/// Like [`run`], sharing `decoded` with other graphs. +fn run_with( + store: &Store, + plan: &PlanRef, + rows: Range, + mask: Mask, + mut pick: impl FnMut(&[IoRequest]) -> usize, + decoded: DecodeCache, +) -> VortexResult { + let mut graph = ExecGraph::try_new(SESSION.clone(), plan, rows, mask, 0, decoded)?; + let mut inflight: Vec = Vec::new(); + let mut arrays = Vec::new(); + let mut events = Vec::new(); + loop { + match graph.state() { + ExecState::Done => break, + ExecState::NeedsCompute => match graph.compute()? { + ExecOutput::Piece(array) => { + assert!(!array.is_empty(), "the root never emits an empty array"); + events.push(Event::Piece(array.len())); + arrays.push(array); + } + ExecOutput::NeedsIO(batch) => { + events.push(Event::Io(batch.iter().map(|r| *r.segment_id).collect())); + inflight.extend(batch); + } + ExecOutput::Yield => {} + }, + ExecState::Waiting => { + if inflight.is_empty() { + return Err(vortex_err!("graph waits with no reads in flight")); + } + let request = inflight.remove(pick(&inflight)); + events.push(Event::Delivered(*request.segment_id)); + graph.set_io_result(request.id, store.read(request.segment_id))?; + } + } + } + assert!(inflight.is_empty(), "graph finished with reads in flight"); + Ok(Run { arrays, events }) +} + +fn delivery(order: Delivery) -> impl FnMut(&[IoRequest]) -> usize { + move |inflight| match order { + Delivery::Fifo => 0, + Delivery::Lifo => inflight.len() - 1, + } +} + +/// Delivers segments in the given order. +fn scripted(order: &[u32]) -> impl FnMut(&[IoRequest]) -> usize + '_ { + let mut next = order.iter(); + move |inflight| { + let segment = *next.next().expect("script covers every read"); + inflight + .iter() + .position(|r| *r.segment_id == segment) + .expect("scripted segment is in flight") + } +} + +fn reads(events: &[Event]) -> usize { + events + .iter() + .map(|event| match event { + Event::Io(batch) => batch.len(), + _ => 0, + }) + .sum() +} + +fn pieces(events: &[Event]) -> Vec { + events + .iter() + .filter_map(|event| match event { + Event::Piece(len) => Some(*len), + _ => None, + }) + .collect() +} + +/// A view's selection over its rows. +#[derive(Clone, Copy, Debug)] +enum Sel { + All, + None, + EveryOther, + /// Rows relative to the view start. + Rows(&'static [usize]), +} + +impl Sel { + fn mask(self, len: usize) -> Mask { + match self { + Sel::All => Mask::new_true(len), + Sel::None => Mask::new_false(len), + Sel::EveryOther => Mask::from_iter((0..len).map(|i| i % 2 == 0)), + Sel::Rows(rows) => Mask::from_indices(len, rows.iter().copied()), + } + } +} + +/// Checks that the arrays, in the order they came out, are the selected rows of the view. +fn assert_view( + expected: &ArrayRef, + rows: &Range, + mask: &Mask, + arrays: Vec, +) -> VortexResult<()> { + let actual = join(expected.dtype(), arrays)?; + let expected = expected + .slice(rows.start as usize..rows.end as usize)? + .filter(mask.clone())?; + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::full(0..20, Sel::All)] +#[case::crosses_every_boundary(3..17, Sel::All)] +#[case::every_other(5..15, Sel::EveryOther)] +#[case::sparse(0..20, Sel::Rows(&[0, 6, 7, 19]))] +#[case::one_row(8..9, Sel::All)] +#[case::inside_one_chunk_per_column(13..16, Sel::All)] +#[case::nothing_selected(2..18, Sel::None)] +#[case::empty_range(5..5, Sel::All)] +fn views_of_one_plan( + #[case] rows: Range, + #[case] sel: Sel, + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + let mask = sel.mask((rows.end - rows.start) as usize); + + let run = run(&store, &plan, rows.clone(), mask.clone(), delivery(order))?; + assert_view(&expected, &rows, &mask, run.arrays) +} + +#[test] +fn many_views_share_one_plan() -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + + // Split the row domain into views that cut through every column at different places. + for split in [[0, 4, 11, 20], [0, 9, 13, 20], [0, 1, 2, 20]] { + for window in split.windows(2) { + let rows = window[0]..window[1]; + let mask = Sel::EveryOther.mask((rows.end - rows.start) as usize); + let run = run( + &store, + &plan, + rows.clone(), + mask.clone(), + delivery(Delivery::Lifo), + )?; + assert_view(&expected, &rows, &mask, run.arrays)?; + } + } + Ok(()) +} + +#[test] +fn first_compute_publishes_every_read_in_one_batch() -> VortexResult<()> { + let mut store = Store::default(); + let (plan, _) = fixture(&mut store)?; + + let run = run( + &store, + &plan, + 0..ROWS, + Mask::new_true(ROWS as usize), + delivery(Delivery::Fifo), + )?; + let Some(Event::Io(first)) = run.events.first() else { + return Err(vortex_err!("first event must be a read batch")); + }; + assert_eq!(first.len(), store.segments.len()); + assert_eq!(reads(&run.events), store.segments.len()); + Ok(()) +} + +#[test] +fn unselected_chunks_are_never_read() -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + // Rows 7..12 only: a reads one chunk, b both, c one, d one, e.x one, e.y one. + let mask = Mask::from_indices(ROWS as usize, 7..12); + + let run = run( + &store, + &plan, + 0..ROWS, + mask.clone(), + delivery(Delivery::Fifo), + )?; + assert_eq!(reads(&run.events), 7); + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} + +#[test] +fn nothing_selected_issues_no_io_and_emits_nothing() -> VortexResult<()> { + let mut store = Store::default(); + let (plan, _) = fixture(&mut store)?; + let mask = Mask::new_false(ROWS as usize); + + let run = run(&store, &plan, 0..ROWS, mask, delivery(Delivery::Fifo))?; + assert_eq!(reads(&run.events), 0); + assert!(run.arrays.is_empty()); + Ok(()) +} + +/// `{a, b}` with `a` chunked `[0,7) [7,12) [12,20)` and `b` flat. +fn two_columns(store: &mut Store) -> VortexResult<(PlanRef, ArrayRef)> { + let a = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let b = PrimitiveArray::from_iter((0..ROWS as i64).map(|v| v * 10)).into_array(); + let expected = StructArray::from_fields(&[("a", a.clone()), ("b", b.clone())])?.into_array(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![store.chunked(&a, &[7, 5, 8])?, store.flat(&b)?], + ) + .into_layout(); + Ok((lower(&layout)?, expected)) +} + +/// Pack emits a struct as soon as every field has rows, however the reads complete, and the +/// structs come out in row order. +/// +/// Segments: a0=0, a1=1, a2=2, b=3. +#[rstest] +// `b` arrives last: nothing can be emitted before it, then everything at once. +#[case::b_last(&[1, 0, 2, 3], &[20])] +// `b` and `a0` first: rows 0..7 go out; `a2` arrives before `a1` and waits in its port. +#[case::a_chunk_last(&[3, 0, 2, 1], &[7, 13])] +// Chunks in order: one struct per chunk of `a`. +#[case::in_order(&[3, 0, 1, 2], &[7, 5, 8])] +fn pack_streams_in_row_order( + #[case] order: &[u32], + #[case] expected_pieces: &[usize], +) -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = two_columns(&mut store)?; + let mask = Mask::new_true(ROWS as usize); + + let run = run(&store, &plan, 0..ROWS, mask.clone(), scripted(order))?; + assert_eq!(pieces(&run.events), expected_pieces); + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} + +/// Fields chunked alike stream one struct per chunk. +#[test] +fn aligned_fields_stream_one_struct_per_chunk() -> VortexResult<()> { + let mut store = Store::default(); + let a = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let b = PrimitiveArray::from_iter((0..ROWS as i64).map(|v| v * 10)).into_array(); + let expected = StructArray::from_fields(&[("a", a.clone()), ("b", b.clone())])?.into_array(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![ + store.chunked(&a, &[7, 5, 8])?, + store.chunked(&b, &[7, 5, 8])?, + ], + ) + .into_layout(); + let plan = lower(&layout)?; + let mask = Mask::new_true(ROWS as usize); + + let run = run( + &store, + &plan, + 0..ROWS, + mask.clone(), + delivery(Delivery::Fifo), + )?; + assert_eq!(pieces(&run.events), [7, 5, 8]); + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} + +#[test] +fn state_is_side_effect_free() -> VortexResult<()> { + let mut store = Store::default(); + let (plan, _) = two_columns(&mut store)?; + let graph = ExecGraph::try_new( + SESSION.clone(), + &plan, + 0..ROWS, + Mask::new_true(ROWS as usize), + 0, + DecodeCache::default(), + )?; + for _ in 0..3 { + assert_eq!(graph.state(), ExecState::NeedsCompute); + } + Ok(()) +} + +/// A bare segment scan returns every row of its range whatever it is told to care about, and a +/// filter over the same scan returns only the selected rows. +#[rstest] +#[case::every_other(Sel::EveryOther)] +#[case::sparse(Sel::Rows(&[1, 4]))] +#[case::nothing(Sel::None)] +fn bare_scan_is_dense_and_filter_keeps_the_selection(#[case] sel: Sel) -> VortexResult<()> { + let mut store = Store::default(); + let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let filtered = lower(&store.flat(&values)?)?; + assert!(filtered.is::()); + let scan = filtered.child_required(0)?; + assert!(scan.is::()); + + let rows = 3..9; + let mask = sel.mask(6); + let mut ctx = SESSION.create_execution_ctx(); + + let dense = run( + &store, + &scan, + rows.clone(), + mask.clone(), + delivery(Delivery::Fifo), + )?; + assert_eq!(dense.arrays.len(), 1); + assert_arrays_eq!(dense.arrays[0], values.slice(3..9)?, &mut ctx); + + let kept = run( + &store, + &filtered, + rows.clone(), + mask.clone(), + delivery(Delivery::Fifo), + )?; + assert_view(&values, &rows, &mask, kept.arrays)?; + Ok(()) +} + +/// A graph sharing a decode cache with one that already ran over the same plan reads nothing and +/// returns the same rows, even under a different selection. +#[test] +fn shared_decode_cache_skips_reads() -> VortexResult<()> { + let mut store = Store::default(); + let (plan, expected) = fixture(&mut store)?; + let decoded = DecodeCache::default(); + + let rows = 0..ROWS; + let all = Mask::new_true(ROWS as usize); + let first = run_with( + &store, + &plan, + rows.clone(), + all, + delivery(Delivery::Fifo), + decoded.clone(), + )?; + assert!(reads(&first.events) > 0); + + let mask = Sel::EveryOther.mask(ROWS as usize); + let second = run_with( + &store, + &plan, + rows.clone(), + mask.clone(), + delivery(Delivery::Fifo), + decoded, + )?; + assert_eq!(reads(&second.events), 0); + assert_view(&expected, &rows, &mask, second.arrays)?; + Ok(()) +} diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 2e814156309..f8b704c2131 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -2,12 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use std::sync::Arc; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::plan::Eval; @@ -18,6 +20,10 @@ use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::exec::ConcatNode; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; +use crate::plan::exec::Selection; use crate::plan::optimizer::PlanParentReduceRule; /// Concatenates its children row-wise. @@ -142,6 +148,18 @@ impl PlanVTable for Concat { fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { Cow::Owned(format!("chunks[{index}]")) } + + fn exec( + plan: &Plan, + rows: Range, + mask: Mask, + _ctx: &ExecContext, + ) -> VortexResult> { + Ok(Box::new(ConcatNode::new( + plan.clone(), + Selection::try_new(rows, mask)?, + ))) + } } /// Pushes an expression into every chunk of a [`Concat`]. diff --git a/vortex-layout/src/plan/plans/filter.rs b/vortex-layout/src/plan/plans/filter.rs index da5907f3e80..e7424b236ba 100644 --- a/vortex-layout/src/plan/plans/filter.rs +++ b/vortex-layout/src/plan/plans/filter.rs @@ -2,11 +2,13 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use vortex_array::EmptyMetadata; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::plan::Plan; @@ -15,7 +17,13 @@ use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::SegmentScan; use crate::plan::check_child_count; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; +use crate::plan::exec::FilterNode; +use crate::plan::exec::SegmentScanNode; +use crate::plan::exec::Selection; /// Keeps only the selected rows of its child. /// @@ -86,4 +94,29 @@ impl PlanVTable for Filter { Cow::Owned(format!("child[{index}]")) } } + + /// Runs fused with a segment-scan child, as one node that keeps the selected rows itself; + /// over any other child, as a filter node that filters the child's whole pieces. + fn exec( + plan: &Plan, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult> { + let child = plan.child_plan()?; + let Some(scan) = child.as_opt::() else { + return Ok(Box::new(FilterNode::new( + plan.clone(), + Selection::try_new(rows, mask)?, + ctx.session().clone(), + ))); + }; + let filter = Some(mask.clone()); + Ok(Box::new(SegmentScanNode::try_new( + scan.clone(), + Selection::try_new(rows, mask)?, + filter, + ctx.clone(), + )?)) + } } diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 6497cc63d45..81311678376 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; @@ -27,6 +28,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::plan::Eval; @@ -37,6 +39,10 @@ use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; +use crate::plan::exec::PackNode; +use crate::plan::exec::Selection; use crate::plan::optimizer::PlanParentReduceRule; /// Assembles a struct from one child per field, plus an optional trailing validity child. @@ -204,6 +210,18 @@ impl PlanVTable for Pack { } Cow::Borrowed("validity") } + + fn exec( + plan: &Plan, + rows: Range, + mask: Mask, + _ctx: &ExecContext, + ) -> VortexResult> { + Ok(Box::new(PackNode::new( + plan.clone(), + Selection::try_new(rows, mask)?, + ))) + } } fn validate_field_child( diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/plans/segment_scan.rs index d2b20df89ad..521d44b589f 100644 --- a/vortex-layout/src/plan/plans/segment_scan.rs +++ b/vortex-layout/src/plan/plans/segment_scan.rs @@ -1,10 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; + use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; @@ -14,6 +17,10 @@ use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanVTable; use crate::plan::check_child_count; +use crate::plan::exec::ExecContext; +use crate::plan::exec::ExecNode; +use crate::plan::exec::SegmentScanNode; +use crate::plan::exec::Selection; use crate::segments::SegmentId; /// Reads one serialized array segment. @@ -92,4 +99,19 @@ impl PlanVTable for SegmentScan { check_child_count("SegmentScan", children, 0)?; Ok(()) } + + fn exec( + plan: &Plan, + rows: Range, + mask: Mask, + ctx: &ExecContext, + ) -> VortexResult> { + // A bare scan returns every row; a Filter over it runs the same node with a filter. + Ok(Box::new(SegmentScanNode::try_new( + plan.clone(), + Selection::try_new(rows, mask)?, + None, + ctx.clone(), + )?)) + } } From 6abc26e01496434011d79e2c71dc9918312d079a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 18:50:01 +0000 Subject: [PATCH 04/20] feat(layout): replace the exec graph with push-based pipelines Replaces the `ExecNode` graph with `plan::pipeline`, the executor a plan now runs on: - `Operator` and `Source`: one `compute` method returning a `Step` (`More`, `Last`, `Consumed`, `Blocked`, `Finished`). Only a source blocks, on a read, an empty inlet, or a full outlet. - Ports with one writer and one reader, in an arena the scan owns; the reading source declares each inlet's capacity, and `Inlet` exposes `peek_mut`, `take`, and `closed`. - A driver that pushes each batch through a pipeline's stages into its outlets, and a scheduler that wakes a blocked pipeline only when its condition clears, checking only the ports a run touched. - `Scan`, which compiles a plan per split, runs several splits at once, and returns reads to its owner as `Turn::Read`. - `PlanVTable::compile` and `PlanVTable::reach` replace `PlanVTable::exec`: the hooks a plan implements to compile into a chain and to report the segments it reads. Synthetic plans exercise the runtime without layouts. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/exec/mod.rs | 672 ------------------ .../src/plan/exec/scheduling_tests.rs | 214 ------ vortex-layout/src/plan/exec/synthetic.rs | 295 -------- vortex-layout/src/plan/mod.rs | 2 +- vortex-layout/src/plan/pipeline/compile.rs | 141 ++++ vortex-layout/src/plan/pipeline/mod.rs | 178 +++++ vortex-layout/src/plan/pipeline/port.rs | 197 +++++ vortex-layout/src/plan/pipeline/scan.rs | 671 +++++++++++++++++ .../src/plan/pipeline/scheduling_tests.rs | 198 ++++++ vortex-layout/src/plan/pipeline/synthetic.rs | 290 ++++++++ vortex-layout/src/plan/typed.rs | 77 +- vortex-layout/src/plan/vtable.rs | 45 +- 12 files changed, 1764 insertions(+), 1216 deletions(-) delete mode 100644 vortex-layout/src/plan/exec/mod.rs delete mode 100644 vortex-layout/src/plan/exec/scheduling_tests.rs delete mode 100644 vortex-layout/src/plan/exec/synthetic.rs create mode 100644 vortex-layout/src/plan/pipeline/compile.rs create mode 100644 vortex-layout/src/plan/pipeline/mod.rs create mode 100644 vortex-layout/src/plan/pipeline/port.rs create mode 100644 vortex-layout/src/plan/pipeline/scan.rs create mode 100644 vortex-layout/src/plan/pipeline/scheduling_tests.rs create mode 100644 vortex-layout/src/plan/pipeline/synthetic.rs diff --git a/vortex-layout/src/plan/exec/mod.rs b/vortex-layout/src/plan/exec/mod.rs deleted file mode 100644 index 5cb10f6e569..00000000000 --- a/vortex-layout/src/plan/exec/mod.rs +++ /dev/null @@ -1,672 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Push-based execution of physical plans. -//! -//! A [`PlanRef`] is turned into an [`ExecGraph`] for one row range and selection. The graph is a -//! flat arena of [`ExecNode`]s wired child-to-parent through [`Input`] ports. A node's output is -//! an ordered sequence of arrays: each node emits its rows in row order, and its parent reads -//! them from the port the child feeds. The root's arrays are handed to the owner one per -//! [`ExecGraph::compute`], in row order. -//! -//! IO may complete in any order. A node whose children are backed by independent reads, such as -//! a concatenation of chunks, holds the early ones in their ports and emits only the ordered -//! prefix, so nothing above it sees the disorder. Nodes never talk to the IO service: they -//! publish requests through [`StepCx::request`], the graph returns them from -//! [`ExecGraph::compute`], and the owner delivers results through [`ExecGraph::set_io_result`]. -//! -//! # Interfaces -//! -//! The module has one interface for each side of a graph: -//! -//! - A plan operator runs as an [`ExecNode`], built by -//! [`PlanVTable::exec`](crate::plan::PlanVTable::exec). A node sees the graph only through its -//! [`StepCx`]: it reads its ports, and spawns children, requests segments, and emits arrays -//! through it. [`ExecNode`] says what a node must do, and [`Ready`] says when it runs. -//! - An owner drives an [`ExecGraph`]: it calls [`ExecGraph::compute`], answers the -//! [`IoRequest`]s that returns, and collects the root's arrays. -//! -//! How a graph stores, wires, and schedules its nodes belongs to neither interface, and nor do -//! the nodes of this crate's own operators. -//! -//! # Pipelines -//! -//! A chain of nodes that are [`Ready::Any`] over one port is a pipeline: an array emitted at -//! its bottom can be pushed through every node in the chain in turn. Nodes with several ports -//! merge pipelines, and nodes that wait for a port to close, [`Ready::Closed`] and -//! [`Ready::AllClosed`], are barriers between them, as a join build is in a query engine. - -pub mod synthetic; - -use std::collections::VecDeque; -use std::mem; -use std::ops::Range; -use std::sync::Arc; - -use parking_lot::Mutex; -use rustc_hash::FxHashMap; -use smallvec::SmallVec; -use vortex_array::ArrayRef; -use vortex_array::buffer::BufferHandle; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_err; -use vortex_mask::Mask; -use vortex_session::VortexSession; - -use crate::plan::Concat; -use crate::plan::PlanRef; -use crate::segments::SegmentId; - -/// Identifies one IO request within an [`ExecGraph`]. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct IoRequestId(pub(crate) u64); - -/// A read published by a leaf node. -#[derive(Clone, Debug)] -pub struct IoRequest { - /// Identifies the request when its result is delivered. - pub id: IoRequestId, - /// The segment to read. - pub segment_id: SegmentId, -} - -/// Requests discovered together. -pub type IoBatch = Vec; - -/// An input of a node. -/// -/// A node names the port when it spawns the child that feeds it, and reads the child's output -/// from it through [`StepCx::input`]. -pub type Port = usize; - -/// One input port of a node: the arrays its child has produced and not yet been taken, in row -/// order, and whether the child has closed. -/// -/// A port that was never spawned reads as closed with nothing available, so a node can treat an -/// optional child it did not spawn like one that finished at once. -#[derive(Default)] -pub struct Input { - spawned: bool, - closed: bool, - pieces: VecDeque, - available: usize, -} - -impl Input { - /// Rows the child has produced that the node has not taken yet. - pub fn available(&self) -> usize { - self.available - } - - /// Whether the child will produce nothing more. True for a port that was never spawned. - pub fn closed(&self) -> bool { - !self.spawned || self.closed - } - - /// Whether the child has closed and every row it produced has been taken. - pub fn finished(&self) -> bool { - self.closed() && self.pieces.is_empty() - } - - /// Takes every available array, in row order. - pub fn take_all(&mut self) -> Vec { - self.available = 0; - self.pieces.drain(..).collect() - } - - /// Takes the next `n` rows, in row order, as the arrays that hold them. The array that - /// spans the boundary is sliced, and its remainder stays available. - /// - /// `n` must not exceed [`available`](Self::available). - pub fn take(&mut self, n: usize) -> VortexResult> { - if n > self.available { - vortex_bail!( - "Taking {n} rows from a port with {} available", - self.available - ); - } - let mut taken = Vec::new(); - let mut remaining = n; - while remaining > 0 { - let Some(front) = self.pieces.pop_front() else { - vortex_bail!("Port ran out of arrays with {remaining} rows still to take"); - }; - if front.len() <= remaining { - remaining -= front.len(); - taken.push(front); - } else { - taken.push(front.slice(0..remaining)?); - self.pieces.push_front(front.slice(remaining..front.len())?); - remaining = 0; - } - } - self.available -= n; - Ok(taken) - } - - fn push(&mut self, array: ArrayRef) { - self.available += array.len(); - self.pieces.push_back(array); - } -} - -/// When a node's [`compute`](ExecNode::compute) runs. -/// -/// A node is considered whenever something changes for it: a child produced rows or closed, or -/// a requested segment was delivered. The rule then decides whether the node runs. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Ready { - /// On every change. For nodes that stream their input through. - Any, - /// Only once each of these ports has closed. For nodes that need one input whole before - /// they can use the others, as a dictionary needs its values before its codes. - Closed(&'static [Port]), - /// Only once every spawned port has closed. For nodes that need all their input at once. - AllClosed, -} - -/// One running plan operator inside an [`ExecGraph`]. -/// -/// A node is built by [`PlanVTable::exec`](crate::plan::PlanVTable::exec) for a range of its -/// plan's rows and a selection over them, and produces the selected rows as arrays, in row -/// order. -/// -/// # Lifecycle -/// -/// The graph calls [`start`](Self::start) once, right after the node is built. It then calls -/// [`compute`](Self::compute) whenever something changed for the node and the node's -/// [`ready`](Self::ready) rule holds, and again after `compute` returned [`NodeState::Yield`]. -/// A node is not called again once it has returned [`NodeState::Done`] or an error. -/// -/// Everything a node does goes through its [`StepCx`] and takes effect once the call returns, -/// in this order: its children are spawned and started, its requests are published, its arrays -/// are pushed to its parent's port, and, if it is done, its port closes. -/// -/// # Output -/// -/// A node emits the selected rows of the range it was built for, in row order, as any number of -/// non-empty arrays. It returns [`NodeState::Done`] once it has emitted them all. Its children -/// are not cancelled, so it returns `Done` only once every child it spawned has closed, and it -/// returns [`NodeState::Wait`] only while a child is open or a request is outstanding, since -/// nothing else runs it again. -/// -/// The graph checks none of this. -/// -/// # Errors -/// -/// An error from `start` or `compute` fails the whole graph. -pub trait ExecNode: Send { - /// When [`compute`](Self::compute) runs. May change as the node progresses. - fn ready(&self) -> Ready { - Ready::Any - } - - /// Spawns children, publishes requests, and emits whatever needs no input. - fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult; - - /// Takes what its ports hold, does work, and reports what happens to the node next. - fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult; -} - -/// What happens to a node after a [`start`](ExecNode::start) or [`compute`](ExecNode::compute). -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum NodeState { - /// Run again once something changes: a child produces rows or closes, or a delivery lands. - Wait, - /// Run again without waiting for a change, after other work: CPU work is left. - Yield, - /// Finished: the node's port closes and it never runs again. - Done, -} - -/// What a graph gives every node it builds. -#[derive(Clone)] -pub struct ExecContext { - session: VortexSession, - row_offset: u64, - decoded: DecodeCache, -} - -impl ExecContext { - pub(crate) fn new(session: VortexSession, row_offset: u64, decoded: DecodeCache) -> Self { - Self { - session, - row_offset, - decoded, - } - } - - /// The session used for decoding and expression evaluation. - pub fn session(&self) -> &VortexSession { - &self.session - } - - /// The global row index of the graph's first plan row. - pub fn row_offset(&self) -> u64 { - self.row_offset - } - - /// Decoded segments shared by the graphs reading the same rows. - pub fn decoded(&self) -> &DecodeCache { - &self.decoded - } -} - -/// Decoded segments shared by the graphs that read one row range, so a segment several of them -/// read is fetched and decoded once. -/// -/// A segment decodes the same way wherever it appears, so entries are keyed by segment id alone. -#[derive(Clone, Default)] -pub struct DecodeCache { - decoded: Arc>>, -} - -impl DecodeCache { - /// The whole decoded array of `id`, if a graph sharing this cache decoded it. - pub fn get(&self, id: SegmentId) -> Option { - self.decoded.lock().get(&id).cloned() - } - - /// Shares the whole decoded array of `id` with the graphs sharing this cache. - pub fn insert(&self, id: SegmentId, array: ArrayRef) { - self.decoded.lock().insert(id, array); - } -} - -/// A child a node spawned: its port, plan, rows, and selection. -type Spawn = (Port, PlanRef, Range, Mask); - -/// A node's view of its graph during one [`start`](ExecNode::start) or -/// [`compute`](ExecNode::compute): its ports, a delivered segment, and what the node does about -/// them. -/// -/// Nothing recorded here happens until the call returns. [`ExecNode`] gives the order. -pub struct StepCx<'a> { - inputs: &'a mut SmallVec<[Input; 2]>, - delivered: Option, - next_io_id: &'a mut u64, - spawned: SmallVec<[Spawn; 1]>, - requests: SmallVec<[IoRequest; 1]>, - emitted: SmallVec<[ArrayRef; 1]>, -} - -impl<'a> StepCx<'a> { - fn new( - inputs: &'a mut SmallVec<[Input; 2]>, - delivered: Option, - next_io_id: &'a mut u64, - ) -> Self { - Self { - inputs, - delivered, - next_io_id, - spawned: SmallVec::new(), - requests: SmallVec::new(), - emitted: SmallVec::new(), - } - } - - /// The port fed by the child spawned on `port`. - pub fn input(&mut self, port: Port) -> &mut Input { - if port >= self.inputs.len() { - self.inputs.resize_with(port + 1, Input::default); - } - &mut self.inputs[port] - } - - /// Every port, in port order, including ports never spawned. - pub fn inputs(&mut self) -> &mut [Input] { - self.inputs - } - - /// Whether every spawned port has closed and been drained. - pub fn all_finished(&self) -> bool { - self.inputs.iter().all(Input::finished) - } - - /// The bytes of the segment this node requested, if they arrived since it last ran. - pub fn take_delivery(&mut self) -> Option { - self.delivered.take() - } - - /// Spawns `plan` as a child feeding `port`, over `rows` of the child's row domain restricted - /// to `mask`, which is as long as `rows`. - pub fn spawn(&mut self, port: Port, plan: PlanRef, rows: Range, mask: Mask) { - self.spawned.push((port, plan, rows, mask)); - } - - /// Publishes a read of `segment_id`. Its bytes arrive through - /// [`take_delivery`](Self::take_delivery). A node has at most one read outstanding. - pub fn request(&mut self, segment_id: SegmentId) -> IoRequestId { - let id = IoRequestId(*self.next_io_id); - *self.next_io_id += 1; - self.requests.push(IoRequest { id, segment_id }); - id - } - - /// Pushes the node's next rows to its parent, or from the root to the graph's owner. Empty - /// arrays are dropped. - pub fn emit(&mut self, array: ArrayRef) { - if !array.is_empty() { - self.emitted.push(array); - } - } -} - -/// What an [`ExecGraph`] needs next. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum ExecState { - /// A node is ready, or an array or requests are waiting to be returned. - NeedsCompute, - /// Every remaining node waits on a delivered request. - Waiting, - /// The root has closed and every array has been returned. - Done, -} - -/// The result of one [`ExecGraph::compute`] call. -#[derive(Debug)] -pub enum ExecOutput { - /// Control returns to the owner. Check [`ExecGraph::state`] for what is needed next. - Yield, - /// Requests published since the last flush, each returned exactly once. - NeedsIO(IoBatch), - /// The root's next rows. - Piece(ArrayRef), -} - -type NodeId = usize; - -#[derive(Clone, Copy, PartialEq, Eq)] -enum Status { - /// Not in the ready queue. - Idle, - /// In the ready queue. - Queued, - /// Closed; never run again. - Done, -} - -struct Node { - exec: Box, - parent: Option<(NodeId, Port)>, - inputs: SmallVec<[Input; 2]>, - delivered: Option, - status: Status, -} - -/// A running plan: a flat arena of nodes wired child-to-parent. -pub struct ExecGraph { - ctx: ExecContext, - nodes: Vec, - /// Ready nodes, run last-in first-out so a woken parent runs before its child's next call. - ready: Vec, - io_routes: SmallVec<[Option; 2]>, - new_io: IoBatch, - outputs: VecDeque, - root_closed: bool, - next_io_id: u64, -} - -impl ExecGraph { - /// Builds the graph for `plan` over `rows`, restricted to `mask`, and starts each node. - /// `row_offset` is the global row index of the plan's first row, used by row-index plans. - /// Segments found in `decoded` are neither read nor decoded again, and segments the graph - /// decodes are added to it. - /// - /// Construction does no IO. Every leaf's first request is returned by the first - /// [`compute`](Self::compute). - pub fn try_new( - session: VortexSession, - plan: &PlanRef, - rows: Range, - mask: Mask, - row_offset: u64, - decoded: DecodeCache, - ) -> VortexResult { - let mut graph = Self { - ctx: ExecContext::new(session, row_offset, decoded), - nodes: Vec::new(), - ready: Vec::new(), - io_routes: SmallVec::new(), - new_io: Vec::new(), - outputs: VecDeque::new(), - root_closed: false, - next_io_id: 0, - }; - graph.spawn(None, plan, rows, mask)?; - Ok(graph) - } - - /// Reports what the graph needs next. Cheap and side-effect free. - pub fn state(&self) -> ExecState { - if !self.outputs.is_empty() || !self.ready.is_empty() || !self.new_io.is_empty() { - ExecState::NeedsCompute - } else if self.root_closed { - ExecState::Done - } else { - ExecState::Waiting - } - } - - /// Ids of requests returned by [`compute`](Self::compute) whose results are not yet delivered. - pub fn outstanding(&self) -> impl Iterator + '_ { - self.io_routes - .iter() - .enumerate() - .filter_map(|(index, node)| { - let id = IoRequestId(index as u64); - (node.is_some() && !self.new_io.iter().any(|request| request.id == id)) - .then_some(id) - }) - } - - /// Runs ready nodes until an array reaches the root, requests should be flushed, a node - /// yields, or nothing is ready. - /// - /// # Errors - /// - /// Returns the error of a node that failed. The graph cannot be run further. - pub fn compute(&mut self) -> VortexResult { - loop { - if let Some(array) = self.outputs.pop_front() { - return Ok(ExecOutput::Piece(array)); - } - let Some(node) = self.ready.pop() else { - return Ok(self.flush_or_yield()); - }; - let yielded = self.run(node, false)?; - - // Get reads in flight before spending more CPU on other nodes. - if !self.new_io.is_empty() && self.outputs.is_empty() { - return Ok(ExecOutput::NeedsIO(mem::take(&mut self.new_io))); - } - if yielded && self.outputs.is_empty() { - return Ok(ExecOutput::Yield); - } - } - } - - /// Delivers the result of a request returned by [`compute`](Self::compute). - pub fn set_io_result(&mut self, id: IoRequestId, result: BufferHandle) -> VortexResult<()> { - let node = self - .io_routes - .get_mut(usize::try_from(id.0)?) - .and_then(Option::take) - .ok_or_else(|| vortex_err!("Unknown IO request {id:?}"))?; - self.nodes[node].delivered = Some(result); - self.enqueue(node); - Ok(()) - } - - fn flush_or_yield(&mut self) -> ExecOutput { - if self.new_io.is_empty() { - ExecOutput::Yield - } else { - ExecOutput::NeedsIO(mem::take(&mut self.new_io)) - } - } - - /// Builds the node for `plan` under `parent` and starts it. - /// - /// A concatenation whose requested rows fall inside one chunk is skipped: the chunk is - /// wired to the parent directly, since its rows are the concatenation's rows in order. - fn spawn( - &mut self, - parent: Option<(NodeId, Port)>, - plan: &PlanRef, - rows: Range, - mask: Mask, - ) -> VortexResult<()> { - if let Some(concat) = plan.as_opt::() - && rows.start < rows.end - { - let offsets = concat.row_offsets(); - let index = offsets - .partition_point(|&offset| offset <= rows.start) - .saturating_sub(1); - let end = offsets - .get(index + 1) - .copied() - .unwrap_or_else(|| concat.row_count()); - if let Some(&start) = offsets.get(index) - && rows.end <= end - { - return self.spawn( - parent, - &concat.child_required(index)?, - rows.start - start..rows.end - start, - mask, - ); - } - } - if let Some((parent, port)) = parent { - let inputs = &mut self.nodes[parent].inputs; - if port >= inputs.len() { - inputs.resize_with(port + 1, Input::default); - } - inputs[port].spawned = true; - } - let id = self.nodes.len(); - self.nodes.push(Node { - exec: plan.exec(rows, mask, &self.ctx)?, - parent, - inputs: SmallVec::new(), - delivered: None, - status: Status::Queued, - }); - self.run(id, true)?; - Ok(()) - } - - /// Runs one node, its `start` or its `compute`, and applies the effects. Returns whether - /// control goes back to the owner: when the node yields, and after a node handles a - /// delivery, since decoding and evaluating newly read data is the costly work between reads - /// that other owners' work should interleave with. - fn run(&mut self, id: NodeId, start: bool) -> VortexResult { - let mut next_io_id = self.next_io_id; - let node = &mut self.nodes[id]; - let delivered = node.delivered.take(); - let had_delivery = delivered.is_some(); - let mut cx = StepCx::new(&mut node.inputs, delivered, &mut next_io_id); - let state = if start { - node.exec.start(&mut cx)? - } else { - node.exec.compute(&mut cx)? - }; - let effects = Effects::from(cx); - self.next_io_id = next_io_id; - - // Settle the node before waking its parent, so a re-queued node runs after the parent. - self.nodes[id].status = match state { - NodeState::Done => Status::Done, - NodeState::Wait | NodeState::Yield => Status::Idle, - }; - if state == NodeState::Yield { - self.enqueue(id); - } - self.apply(id, effects, state == NodeState::Done)?; - Ok(state == NodeState::Yield || had_delivery) - } - - fn apply(&mut self, id: NodeId, effects: Effects, done: bool) -> VortexResult<()> { - for (port, plan, rows, mask) in effects.spawned { - self.spawn(Some((id, port)), &plan, rows, mask)?; - } - for request in effects.requests { - let index = usize::try_from(request.id.0)?; - if index >= self.io_routes.len() { - self.io_routes.resize(index + 1, None); - } - self.io_routes[index] = Some(id); - self.new_io.push(request); - } - for array in effects.emitted { - self.deliver(id, Some(array)); - } - if done { - self.deliver(id, None); - } - Ok(()) - } - - /// Passes an array, or the close when `array` is `None`, from `from` to its parent's port, - /// and wakes the parent if its readiness rule holds. - fn deliver(&mut self, from: NodeId, array: Option) { - match self.nodes[from].parent { - None => match array { - Some(array) => self.outputs.push_back(array), - None => self.root_closed = true, - }, - Some((parent, port)) => { - let input = &mut self.nodes[parent].inputs[port]; - match array { - Some(array) => input.push(array), - None => input.closed = true, - } - if self.is_ready(parent) { - self.enqueue(parent); - } - } - } - } - - fn is_ready(&self, id: NodeId) -> bool { - let node = &self.nodes[id]; - match node.exec.ready() { - Ready::Any => true, - Ready::Closed(ports) => ports - .iter() - .all(|&port| node.inputs.get(port).is_none_or(Input::closed)), - Ready::AllClosed => node.inputs.iter().all(Input::closed), - } - } - - fn enqueue(&mut self, id: NodeId) { - if self.nodes[id].status == Status::Idle { - self.nodes[id].status = Status::Queued; - self.ready.push(id); - } - } -} - -/// The owned effects of a [`StepCx`], detached from its borrows. -struct Effects { - spawned: SmallVec<[Spawn; 1]>, - requests: SmallVec<[IoRequest; 1]>, - emitted: SmallVec<[ArrayRef; 1]>, -} - -impl From> for Effects { - fn from(cx: StepCx<'_>) -> Self { - Self { - spawned: cx.spawned, - requests: cx.requests, - emitted: cx.emitted, - } - } -} - -#[cfg(test)] -mod scheduling_tests; diff --git a/vortex-layout/src/plan/exec/scheduling_tests.rs b/vortex-layout/src/plan/exec/scheduling_tests.rs deleted file mode 100644 index 6fb2ee452f7..00000000000 --- a/vortex-layout/src/plan/exec/scheduling_tests.rs +++ /dev/null @@ -1,214 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Tests of the graph's scheduling, driven through hand-built nodes over in-memory segments. -//! -//! Every source emits its own row indices, so what reaches the root can be checked against a -//! sequence. - -#![allow(clippy::cast_possible_truncation)] - -use std::ops::Range; -use std::sync::Arc; - -use parking_lot::Mutex; -use rstest::rstest; -use vortex_array::ArrayRef; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::ChunkedArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::assert_arrays_eq; -use vortex_array::dtype::DType; -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::synthetic::Probe; -use super::synthetic::RowSource; -use super::synthetic::empty_segment; -use super::*; -use crate::test::SESSION; - -#[derive(Debug, PartialEq, Eq)] -enum Event { - /// A batch returned by `compute`, as segment ids. - Io(Vec), - /// A result delivered for this segment. - Delivered(u32), - /// An array of this many rows reached the root. - Piece(usize), - /// A node yielded to the owner. - Yield, -} - -struct Run { - arrays: Vec, - events: Vec, -} - -/// Drives a graph the way an owner does, answering reads with empty segments in the order -/// `pick` chooses, only while the graph waits. -fn run( - plan: &PlanRef, - rows: Range, - mask: Mask, - mut pick: impl FnMut(&[IoRequest]) -> usize, -) -> VortexResult { - let mut graph = - ExecGraph::try_new(SESSION.clone(), plan, rows, mask, 0, DecodeCache::default())?; - let mut inflight: Vec = Vec::new(); - let mut arrays = Vec::new(); - let mut events = Vec::new(); - loop { - match graph.state() { - ExecState::Done => break, - ExecState::NeedsCompute => match graph.compute()? { - ExecOutput::Piece(array) => { - assert!(!array.is_empty(), "the root never emits an empty array"); - events.push(Event::Piece(array.len())); - arrays.push(array); - } - ExecOutput::NeedsIO(batch) => { - events.push(Event::Io(batch.iter().map(|r| *r.segment_id).collect())); - inflight.extend(batch); - } - ExecOutput::Yield => events.push(Event::Yield), - }, - ExecState::Waiting => { - if inflight.is_empty() { - return Err(vortex_err!("graph waits with no reads in flight")); - } - let request = inflight.remove(pick(&inflight)); - events.push(Event::Delivered(*request.segment_id)); - graph.set_io_result(request.id, empty_segment())?; - } - } - } - assert!(inflight.is_empty(), "graph finished with reads in flight"); - Ok(Run { arrays, events }) -} - -fn fifo(_: &[IoRequest]) -> usize { - 0 -} - -/// Delivers segments in the given order. -fn scripted(order: &[u32]) -> impl FnMut(&[IoRequest]) -> usize + '_ { - let mut next = order.iter(); - move |inflight| { - let segment = *next.next().expect("script covers every read"); - inflight - .iter() - .position(|r| *r.segment_id == segment) - .expect("scripted segment is in flight") - } -} - -fn pieces(events: &[Event]) -> Vec { - events - .iter() - .filter_map(|event| match event { - Event::Piece(len) => Some(*len), - _ => None, - }) - .collect() -} - -fn indices(values: impl IntoIterator) -> ArrayRef { - Buffer::from_iter(values).into_array() -} - -/// Joins the arrays a run produced, in order, into one array. -fn join(dtype: &DType, arrays: Vec) -> VortexResult { - Ok(ChunkedArray::try_new(arrays, dtype.clone())?.into_array()) -} - -fn assert_rows(expected: ArrayRef, arrays: Vec) -> VortexResult<()> { - let actual = join(expected.dtype(), arrays)?; - let mut ctx = SESSION.create_execution_ctx(); - assert_arrays_eq!(actual, expected, &mut ctx); - Ok(()) -} - -fn segment(id: u32) -> Option { - Some(SegmentId::from(id)) -} - -#[rstest] -#[case::one_piece(vec![], false)] -#[case::many_pieces(vec![3, 1, 4, 1, 5], false)] -#[case::yielding(vec![3, 1, 4, 1, 5], true)] -#[case::oversized_pieces(vec![100], true)] -fn source_pieces_come_out_in_order( - #[case] cut: Vec, - #[case] yielding: bool, -) -> VortexResult<()> { - let plan = RowSource::plan(20, cut, None, yielding, false); - let mask = Mask::from_iter((0..10).map(|i| i != 4)); - let run = run(&plan, 5..15, mask, fifo)?; - assert_rows(indices((5..15).filter(|i| *i != 9)), run.arrays) -} - -/// A node runs only when its readiness rule holds, whatever order its inputs finish in. -#[rstest] -#[case::any(Ready::Any)] -#[case::closed_second(Ready::Closed(&[1]))] -#[case::all_closed(Ready::AllClosed)] -fn probe_runs_only_when_ready(#[case] ready: Ready) -> VortexResult<()> { - let runs = Arc::new(Mutex::new(Vec::new())); - let children = vec![ - RowSource::plan(10, vec![2, 2, 2], segment(0), false, false), - RowSource::plan(10, vec![5], segment(1), false, false), - ]; - let plan = Probe::plan(children, ready, Arc::clone(&runs)); - // The first child lands first, so under `Any` the node runs while the second is open. - let run = run(&plan, 0..10, Mask::new_true(10), scripted(&[0, 1]))?; - assert_rows(indices(0..10), run.arrays)?; - - let runs = runs.lock(); - assert!(!runs.is_empty()); - match ready { - Ready::Any => assert!(runs[0] == [true, false], "ran as soon as port 0 closed"), - Ready::Closed(_) => assert!(runs.iter().all(|ports| ports[1])), - Ready::AllClosed => { - assert_eq!(runs.len(), 1); - assert_eq!(runs[0], [true, true]); - } - } - Ok(()) -} - -/// A source that yields between pieces is run once per piece, and its pieces reach the root in -/// order, one per run. -#[test] -fn yielding_source_is_run_once_per_piece() -> VortexResult<()> { - let plan = RowSource::plan(12, vec![4, 4, 4], None, true, false); - let run = run(&plan, 0..12, Mask::new_true(12), fifo)?; - assert_eq!(pieces(&run.events), [4, 4, 4]); - assert_rows(indices(0..12), run.arrays) -} - -#[test] -fn input_take_slices_at_the_boundary() -> VortexResult<()> { - let mut input = Input { - spawned: true, - ..Default::default() - }; - input.push(PrimitiveArray::from_iter(0..4i32).into_array()); - input.push(PrimitiveArray::from_iter(4..10i32).into_array()); - assert_eq!(input.available(), 10); - - let first = input.take(6)?; - assert_eq!(first.iter().map(|a| a.len()).collect::>(), [4, 2]); - assert_eq!(input.available(), 4); - assert!(!input.finished()); - - input.closed = true; - let rest = input.take_all(); - assert_eq!(rest.iter().map(|a| a.len()).collect::>(), [4]); - assert!(input.finished()); - assert!(input.take(1).is_err()); - Ok(()) -} diff --git a/vortex-layout/src/plan/exec/synthetic.rs b/vortex-layout/src/plan/exec/synthetic.rs deleted file mode 100644 index 1d6b74e44ed..00000000000 --- a/vortex-layout/src/plan/exec/synthetic.rs +++ /dev/null @@ -1,295 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Plans whose exec nodes are built by hand, for exercising the graph without layouts. -//! -//! A [`SyntheticPlan`] carries a closure that builds its node. The graph treats it like any -//! other operator, so a test or benchmark can wire hand-written sources under the real -//! operators (Pack, Concat, Filter, Eval) and drive the whole thing through [`ExecGraph`], -//! reading segments from memory. -//! -//! [`ExecGraph`]: crate::plan::exec::ExecGraph - -use std::fmt; -use std::ops::Range; -use std::sync::Arc; - -use vortex_array::EmptyMetadata; -use vortex_array::IntoArray; -use vortex_array::buffer::BufferHandle; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; -use vortex_session::registry::CachedId; - -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::exec::ExecContext; -use crate::plan::exec::ExecNode; -use crate::plan::exec::NodeState; -use crate::plan::exec::Ready; -use crate::plan::exec::StepCx; -use crate::segments::SegmentId; - -/// Builds the node of a [`SyntheticPlan`] for one execution. -pub type NodeBuilder = - dyn Fn(Range, Mask, &ExecContext) -> VortexResult> + Send + Sync; - -/// A plan operator whose node is built by a closure. -#[derive(Clone, Debug)] -pub struct Synthetic; - -/// The node builder of a [`Synthetic`] plan. -#[derive(Clone)] -pub struct SyntheticData { - builder: Arc, -} - -impl fmt::Debug for SyntheticData { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("SyntheticData") - } -} - -/// A plan whose node is built by a closure. -pub type SyntheticPlan = Plan; - -impl SyntheticPlan { - /// A plan of `row_count` rows of `dtype`, with `children`, whose node `builder` makes. - pub fn new( - dtype: DType, - row_count: u64, - children: Vec, - builder: impl Fn(Range, Mask, &ExecContext) -> VortexResult> - + Send - + Sync - + 'static, - ) -> Self { - PlanParts { - vtable: Synthetic, - dtype, - row_count, - children: children.into(), - data: SyntheticData { - builder: Arc::new(builder), - }, - } - .into_typed() - } -} - -impl PlanVTable for Synthetic { - type PlanData = SyntheticData; - type Metadata = EmptyMetadata; - - fn id(&self) -> PlanId { - static ID: CachedId = CachedId::new("vortex.plan.synthetic"); - *ID - } - - fn metadata(_plan: &Plan) -> Option { - None - } - - fn with_children( - _plan: &Plan, - _children: &PlanChildren, - _data: &mut Self::PlanData, - ) -> VortexResult<()> { - Ok(()) - } - - fn exec( - plan: &Plan, - rows: Range, - mask: Mask, - ctx: &ExecContext, - ) -> VortexResult> { - (plan.data().builder)(rows, mask, ctx) - } -} - -/// The dtype every synthetic source produces: its rows' indices, as `u64`. -pub fn row_dtype() -> DType { - DType::Primitive(PType::U64, Nullability::NonNullable) -} - -/// A source that emits its rows' indices, in order, cut into pieces of the given lengths, -/// optionally after a round trip through the IO service and optionally yielding between -/// pieces. -/// -/// A selected source emits only the rows its mask selects, as every operator does. A dense -/// source ignores its mask and emits every row of its range, as a bare segment scan does, so -/// a filter plan can sit over it. The indices are the plan's own rows, so -/// a parent that reorders or filters can be checked against a sequence. Piece lengths count -/// emitted rows; the last piece takes whatever remains. -pub struct RowSource { - /// Remaining indices to emit, in order. - remaining: Vec, - /// Lengths of the pieces still to emit; the last piece takes whatever remains. - pieces: Vec, - /// A segment to request before emitting anything, standing in for a read. - segment: Option, - /// Whether to return `Yield` between pieces instead of emitting them all at once. - yielding: bool, - requested: bool, -} - -impl RowSource { - /// A source plan of `row_count` rows that emits pieces of the given lengths. - pub fn plan( - row_count: u64, - pieces: Vec, - segment: Option, - yielding: bool, - dense: bool, - ) -> PlanRef { - SyntheticPlan::new(row_dtype(), row_count, Vec::new(), move |rows, mask, _| { - let remaining = if dense || mask.all_true() { - (rows.start..rows.end).collect() - } else { - (rows.start..rows.end) - .zip(mask.iter()) - .filter_map(|(row, selected)| selected.then_some(row)) - .collect() - }; - Ok(Box::new(RowSource { - remaining, - pieces: pieces.iter().rev().copied().collect(), - segment, - yielding, - requested: false, - })) - }) - .into_plan() - } - - /// Emits the next piece, or everything left when no lengths remain. - fn emit_next(&mut self, cx: &mut StepCx<'_>) { - let len = self - .pieces - .pop() - .unwrap_or(self.remaining.len()) - .min(self.remaining.len()); - let piece: Vec = self.remaining.drain(..len).collect(); - cx.emit(Buffer::from(piece).into_array()); - } - - fn emit(&mut self, cx: &mut StepCx<'_>) -> NodeState { - if self.yielding { - self.emit_next(cx); - if self.remaining.is_empty() { - return NodeState::Done; - } - return NodeState::Yield; - } - while !self.remaining.is_empty() { - self.emit_next(cx); - } - NodeState::Done - } -} - -impl ExecNode for RowSource { - fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { - if self.remaining.is_empty() { - return Ok(NodeState::Done); - } - if let Some(segment) = self.segment { - self.requested = true; - cx.request(segment); - return Ok(NodeState::Wait); - } - Ok(self.emit(cx)) - } - - fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { - if self.requested { - self.requested = false; - cx.take_delivery() - .ok_or_else(|| vortex_err!("RowSource ran without its delivery"))?; - } - Ok(self.emit(cx)) - } -} - -/// A node with a chosen readiness rule that records the state of its ports every time it runs, -/// then passes its first port through. -/// -/// It spawns one child per plan child, each over the node's own rows and selection. -pub struct Probe { - ready: Ready, - children: Vec, - rows: Range, - mask: Mask, - /// For each compute: whether each port was closed when the node ran. - pub runs: Arc>>>, -} - -impl Probe { - /// A plan over `children` whose node runs under `ready` and logs into `runs`. - pub fn plan( - children: Vec, - ready: Ready, - runs: Arc>>>, - ) -> PlanRef { - let row_count = children.first().map_or(0, |child| child.row_count()); - SyntheticPlan::new( - row_dtype(), - row_count, - children.clone(), - move |rows, mask, _| { - Ok(Box::new(Probe { - ready, - children: children.clone(), - rows, - mask, - runs: Arc::clone(&runs), - })) - }, - ) - .into_plan() - } -} - -impl ExecNode for Probe { - fn ready(&self) -> Ready { - self.ready - } - - fn start(&mut self, cx: &mut StepCx<'_>) -> VortexResult { - for (port, child) in self.children.iter().enumerate() { - cx.spawn(port, child.clone(), self.rows.clone(), self.mask.clone()); - } - Ok(NodeState::Wait) - } - - fn compute(&mut self, cx: &mut StepCx<'_>) -> VortexResult { - let closed = cx.inputs().iter().map(|input| input.closed()).collect(); - self.runs.lock().push(closed); - for array in cx.input(0).take_all() { - cx.emit(array); - } - for port in 1..self.children.len() { - cx.input(port).take_all(); - } - if cx.all_finished() { - return Ok(NodeState::Done); - } - Ok(NodeState::Wait) - } -} - -/// An in-memory segment store that answers any segment id with an empty buffer, for sources -/// whose read is only a scheduling event. -pub fn empty_segment() -> BufferHandle { - BufferHandle::new_host(Buffer::::empty().into_byte_buffer()) -} diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index b4782481b49..92ece0c38dd 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -9,10 +9,10 @@ mod children; mod display; -pub mod exec; mod lower; mod optimize; pub mod optimizer; +pub mod pipeline; mod plans; mod typed; mod vtable; diff --git a/vortex-layout/src/plan/pipeline/compile.rs b/vortex-layout/src/plan/pipeline/compile.rs new file mode 100644 index 00000000000..0f5735b86aa --- /dev/null +++ b/vortex-layout/src/plan/pipeline/compile.rs @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compiling plans into pipelines, and counting the readers of each segment so segments read +//! more than once are decoded once. + +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::Operator; +use super::Source; +use super::port::PortId; +use super::port::Reader; +use super::scan::Core; +use crate::plan::PlanRef; + +/// A pipeline under construction: a source and the stages after it, not yet given outlets, so a +/// parent plan can add stages before it becomes a pipeline. +pub struct Chain { + pub(crate) source: Box, + pub(crate) stages: Vec>, + pub(crate) inlets: Vec, +} + +impl Chain { + /// A chain of `source` alone. + pub fn new(source: impl Source + 'static) -> Self { + Self { + source: Box::new(source), + stages: Vec::new(), + inlets: Vec::new(), + } + } + + /// This chain with `stage` after its last stage. + pub fn with(mut self, stage: impl Operator + 'static) -> Self { + self.stages.push(Box::new(stage)); + self + } +} + +impl Core { + /// Compiles `plan` over `rows`, restricted to `mask`. See [`Compiler::compile`]. + pub(crate) fn compile( + &mut self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + ) -> VortexResult> { + Compiler { core: self }.compile(plan, rows, mask) + } +} + +/// What a plan compiles its children through. See +/// [`PlanVTable::compile`](crate::plan::PlanVTable::compile). +pub struct Compiler<'a> { + pub(crate) core: &'a mut Core, +} + +impl Compiler<'_> { + /// Compiles `plan` over `rows` of its domain, restricted to `mask`, into a chain producing + /// the rows `mask` selects, in order, or `None` when that is no row. + pub fn compile( + &mut self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + ) -> VortexResult> { + if rows.start >= rows.end { + return Ok(None); + } + plan.compile(rows, mask, self) + } + + /// A chain reading `chains` through `source`: each chain becomes a pipeline writing a new + /// port, inlet `i` of `source` reading chain `i`, with the capacity `source` asks for. + pub fn join(&mut self, chains: Vec, source: impl Source + 'static) -> Chain { + let inlets = chains + .into_iter() + .enumerate() + .map(|(index, chain)| { + let port = self + .core + .arena + .create(source.capacity(index), None, Reader::Unclaimed); + self.core.add_pipeline(chain, smallvec::smallvec![port]); + port + }) + .collect(); + Chain { + source: Box::new(source), + stages: Vec::new(), + inlets, + } + } + + /// The global row index of the scanned plan's row zero. + pub fn row_offset(&self) -> u64 { + self.core.row_offset + } + + /// The session used for decoding and evaluation. + pub fn session(&self) -> &VortexSession { + &self.core.session + } +} + +/// How a plan's rows map to the rows of the plan a scan's splits range over. +#[derive(Clone, Debug)] +pub enum Reach { + /// Row `r` is scanned row `r + offset`. + Offset(u64), + /// Every row is read for these scanned rows, as a dictionary's values are for its codes. + Fixed(Range), +} + +impl Reach { + /// The scanned rows `rows` are read for. + pub fn root(&self, rows: &Range) -> Range { + match self { + Reach::Offset(offset) => rows.start + offset..rows.end + offset, + Reach::Fixed(root) => root.clone(), + } + } + + /// The mapping of a child whose row zero is this plan's row `by`. + pub fn shift(&self, by: u64) -> Self { + match self { + Reach::Offset(offset) => Reach::Offset(offset + by), + Reach::Fixed(root) => Reach::Fixed(root.clone()), + } + } + + /// The mapping of a child read whole for `rows` of this plan, as dictionary values are. + pub fn fixed(&self, rows: &Range) -> Self { + Reach::Fixed(self.root(rows)) + } +} diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs new file mode 100644 index 00000000000..47161d379ad --- /dev/null +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Push-based execution of physical plans as pipelines. +//! +//! A [`Scan`] runs a [`PlanRef`](crate::plan::PlanRef) over a list of splits, row ranges of the +//! plan's domain. Each split is compiled into pipelines: a pipeline is one [`Source`], zero or +//! more [`Operator`] stages, and the ports its last stage writes. The driver asks the source for +//! a batch and pushes it through the stages into the ports. Every port has exactly one writer and +//! one reader, and the reader's source declares how many batches it may hold. +//! +//! Only a source blocks, and only on one of three things: a segment read it requested, an empty +//! inlet, or a full outlet, which the driver detects. The scheduler re-runs a blocked pipeline +//! when that one condition clears. +//! +//! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and +//! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. + +mod compile; +mod port; +mod scan; +pub mod synthetic; + +use smallvec::SmallVec; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +pub use self::compile::Chain; +pub use self::compile::Compiler; +pub use self::compile::Reach; +use self::port::Arena; +pub use self::port::Inlet; +use self::port::PortId; +pub use self::scan::ReadId; +pub use self::scan::ReadRequest; +pub use self::scan::Scan; +pub use self::scan::Split; +pub use self::scan::Turn; +use crate::plan::PlanRef; +use crate::segments::SegmentId; + +/// Batches an inlet holds before its writer is blocked, unless its reader says otherwise. +pub const DEFAULT_CAPACITY: usize = 2; + +/// What a pipeline stage is handed on each call. +#[derive(Debug)] +pub enum Input { + /// One batch from the stage below. + Chunk(ArrayRef), + /// The stage below has finished. Flush whatever is held. + End, + /// Nothing new. Sent to a source on every call, and to a stage that returned + /// [`Step::More`]. + None, +} + +/// What a stage says back. +#[derive(Debug)] +pub enum Step { + /// A batch, and more is ready now without further input. The driver pushes the batch on + /// and calls again with [`Input::None`]. + More(ArrayRef), + /// A batch, and the input is spent. + Last(ArrayRef), + /// The input was absorbed and nothing is ready. + Consumed, + /// Nothing can happen until the condition clears. Sources only. + Blocked(Blocked), + /// No more output will ever come. The driver sends [`Input::End`] to the next stage. + Finished, +} + +/// Why a pipeline cannot run. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Blocked { + /// A read was requested and has not been delivered. + Io, + /// The inlet is empty and its writer has not closed it. + Inlet(usize), + /// An outlet is full. Produced by the driver, never by a stage. + Outlet, +} + +/// A stage of a pipeline: takes one input at a time and pushes what it produces on. +/// +/// A stage never blocks. A stage returning [`Step::More`] must eventually return something +/// else when called again with [`Input::None`]. A stage that has returned [`Step::Finished`] is +/// not called again. A stage receiving [`Input::End`] flushes, returning [`Step::More`] while it +/// has more to flush. +pub trait Operator: Send { + /// Takes `input` and says what happens next. + fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult; +} + +/// The start of a pipeline. Reads its inlets and requests segments rather than taking input +/// from a stage below. Always called with [`Input::None`]. +pub trait Source: Operator { + /// The inlets this source reads, numbered from zero. + fn inlet_count(&self) -> usize { + 0 + } + + /// How many batches inlet `inlet` may hold before its writer is blocked. Asked once, when + /// the port is made. An empty port always has room. + fn capacity(&self, inlet: usize) -> usize { + let _ = inlet; + DEFAULT_CAPACITY + } + + /// The segment this source wants read next, asked before every + /// [`compute`](Operator::compute). A source that returns one is not computed until its + /// bytes have arrived through [`Cx::take_bytes`]. + fn request(&mut self) -> Option { + None + } +} + +/// A plan a source asks to have compiled into a new inlet. +pub(crate) struct Spawn { + plan: PlanRef, + rows: std::ops::Range, + mask: Mask, +} + +/// What a stage may touch during one call: its pipeline's inlets, the bytes it asked for, and +/// the session. +pub struct Cx<'a> { + arena: &'a mut Arena, + inlets: &'a [PortId], + bytes: &'a mut Option, + session: &'a VortexSession, + exec: &'a mut ExecutionCtx, + spawns: &'a mut Vec, + /// The inlets read during this run, so only their writers are checked for room after it. + touched: &'a mut SmallVec<[usize; 4]>, +} + +impl Cx<'_> { + /// The reading side of inlet `index` of this pipeline's source. + pub fn inlet(&mut self, index: usize) -> Inlet<'_> { + if self.touched.last() != Some(&index) { + self.touched.push(index); + } + self.arena.inlet(self.inlets[index]) + } + + /// The bytes of the segment the source requested, once they have arrived. + pub fn take_bytes(&mut self) -> Option { + self.bytes.take() + } + + /// The session used for decoding and evaluation. + pub fn session(&self) -> &VortexSession { + self.session + } + + /// The execution context for executing arrays. + pub fn exec(&mut self) -> &mut ExecutionCtx { + self.exec + } + + /// Asks for `plan` over `rows`, restricted to `mask`, to be compiled into a new inlet of this + /// source once the call returns, and returns the inlet's index. + /// + /// For a source that learns what it must read only from what it has read, as a list learns + /// its elements' range from its offsets. + pub fn spawn(&mut self, plan: PlanRef, rows: std::ops::Range, mask: Mask) -> usize { + self.spawns.push(Spawn { plan, rows, mask }); + self.inlets.len() + self.spawns.len() - 1 + } +} + +#[cfg(test)] +mod scheduling_tests; diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/pipeline/port.rs new file mode 100644 index 00000000000..5172e719dc2 --- /dev/null +++ b/vortex-layout/src/plan/pipeline/port.rs @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Ports: single-writer, single-reader queues of batches, kept in an arena the scan owns so a +//! port can outlive the stage or split that wrote it. + +use std::collections::VecDeque; + +use vortex_array::ArrayRef; + +/// Identifies a port in the scan's arena. +pub(crate) type PortId = usize; + +/// Identifies a pipeline in the scan's table. +pub(crate) type PipelineId = usize; + +/// Who reads a port. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Reader { + /// Nobody: a share's port waiting for the reader it was made for, or a port whose reader + /// has gone. + Unclaimed, + /// A pipeline's source. + Pipeline(PipelineId), + /// The scan, for the output of a split's stage. + Split(usize), +} + +/// A queue of batches with one writer and one reader. +pub(crate) struct Queue { + batches: VecDeque, + /// Batches the queue may hold before its writer is blocked. + capacity: usize, + /// The writer will push nothing more. + closed: bool, + pub(crate) writer: Option, + pub(crate) reader: Reader, + /// The reader will take nothing more, so pushes are dropped. + reader_gone: bool, + /// The slot is free for reuse. + free: bool, +} + +impl Queue { + fn new(capacity: usize, writer: Option, reader: Reader) -> Self { + Self { + batches: VecDeque::new(), + capacity: capacity.max(1), + closed: false, + writer, + reader, + reader_gone: false, + free: false, + } + } + + pub(crate) fn has_room(&self) -> bool { + self.reader_gone || self.batches.len() < self.capacity + } + + pub(crate) fn closed(&self) -> bool { + self.closed + } + + /// Whether the reader would find something new: a batch, or the close. + pub(crate) fn readable(&self) -> bool { + !self.batches.is_empty() || self.closed + } + + pub(crate) fn pop(&mut self) -> Option { + self.batches.pop_front() + } +} + +/// The reading side of a port, borrowed for one call. +pub struct Inlet<'a> { + queue: &'a mut Queue, +} + +impl Inlet<'_> { + /// The oldest batch, in place, so the reader can slice it down without removing it. + pub fn peek_mut(&mut self) -> Option<&mut ArrayRef> { + self.queue.batches.front_mut() + } + + /// The oldest batch, removed. Removing a batch is what frees its slot. + pub fn take(&mut self) -> Option { + self.queue.batches.pop_front() + } + + /// Whether the writer has finished. Closed and empty is the end of the inlet. + pub fn closed(&self) -> bool { + self.queue.closed + } + + /// Whether no batch is queued. + pub fn is_empty(&self) -> bool { + self.queue.batches.is_empty() + } + + /// Batches queued. + pub fn len(&self) -> usize { + self.queue.batches.len() + } + + /// Whether the inlet has ended: closed with nothing queued. + pub fn finished(&self) -> bool { + self.queue.closed && self.queue.batches.is_empty() + } +} + +/// Every port of a scan. Slots are reused once a port's writer has closed it and its reader has +/// gone, so no pipeline still refers to it. +#[derive(Default)] +pub(crate) struct Arena { + queues: Vec, + free: Vec, +} + +impl Arena { + pub(crate) fn create( + &mut self, + capacity: usize, + writer: Option, + reader: Reader, + ) -> PortId { + let queue = Queue::new(capacity, writer, reader); + match self.free.pop() { + Some(id) => { + self.queues[id] = queue; + id + } + None => { + self.queues.push(queue); + self.queues.len() - 1 + } + } + } + + pub(crate) fn get(&self, id: PortId) -> &Queue { + &self.queues[id] + } + + pub(crate) fn get_mut(&mut self, id: PortId) -> &mut Queue { + &mut self.queues[id] + } + + pub(crate) fn inlet(&mut self, id: PortId) -> Inlet<'_> { + Inlet { + queue: &mut self.queues[id], + } + } + + pub(crate) fn has_room(&self, id: PortId) -> bool { + self.queues[id].has_room() + } + + /// Appends a batch, unless the reader has gone. Empty batches are dropped. + pub(crate) fn push(&mut self, id: PortId, batch: ArrayRef) { + let queue = &mut self.queues[id]; + if !queue.reader_gone && !batch.is_empty() { + queue.batches.push_back(batch); + } + } + + /// The writer will push nothing more. + pub(crate) fn close(&mut self, id: PortId) { + let queue = &mut self.queues[id]; + queue.closed = true; + queue.writer = None; + self.maybe_free(id); + } + + /// The reader will take nothing more: what is queued is dropped, and so is what is pushed. + pub(crate) fn drop_reader(&mut self, id: PortId) { + let queue = &mut self.queues[id]; + queue.reader_gone = true; + queue.reader = Reader::Unclaimed; + queue.batches.clear(); + self.maybe_free(id); + } + + fn maybe_free(&mut self, id: PortId) { + let queue = &mut self.queues[id]; + if queue.closed && queue.reader_gone && !queue.free { + queue.free = true; + queue.batches = VecDeque::new(); + self.free.push(id); + } + } + + /// Ports in use. + #[cfg(test)] + pub(crate) fn live(&self) -> usize { + self.queues.len() - self.free.len() + } +} diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs new file mode 100644 index 00000000000..bb6adc45ec5 --- /dev/null +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -0,0 +1,671 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The scan: its pipelines, the driver that runs one, the scheduler that picks which, and the +//! splits whose stages it compiles. + +use std::collections::VecDeque; +use std::fmt; +use std::ops::Range; + +use rustc_hash::FxHashMap; +use smallvec::SmallVec; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::VortexSessionExecute; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::Blocked; +use super::Cx; +use super::Input; +use super::Operator; +use super::Source; +use super::Spawn; +use super::Step; +use super::compile::Chain; +use super::port::Arena; +use super::port::PipelineId; +use super::port::PortId; +use super::port::Reader; +use crate::plan::PlanRef; +use crate::segments::SegmentId; + +/// Identifies a read a scan asked its owner for. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ReadId(u64); + +/// A segment the scan needs read. +#[derive(Clone, Debug)] +pub struct ReadRequest { + /// Identifies the read when its bytes are delivered. + pub id: ReadId, + /// The segment to read. + pub segment_id: SegmentId, +} + +/// A row range of the plan's domain, and the rows of it to produce. +#[derive(Clone, Debug)] +pub struct Split { + /// The rows. + pub rows: Range, + /// Which of them to produce; as long as `rows`. + pub mask: Mask, +} + +impl Split { + /// Every row of `rows`. + pub fn all(rows: Range) -> Self { + let len = usize::try_from(rows.end - rows.start).unwrap_or(usize::MAX); + Self { + mask: Mask::new_true(len), + rows, + } + } +} + +/// What the owner of a [`Scan`] does next. +#[derive(Debug)] +pub enum Turn { + /// Read this segment and [`deliver`](Scan::deliver) its bytes. + Read(ReadRequest), + /// The next rows of split `split`, in row order within the split. + Output(usize, ArrayRef), + /// Nothing can run until a read is delivered. + Waiting, + /// Every split has produced all its rows. + Done, +} + +/// Where a pipeline is in its life. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum State { + Runnable, + Blocked(Blocked), +} + +/// A source, its stages, and the ports they read and write. +pub(crate) struct Pipeline { + source: Box, + stages: Vec>, + inlets: Vec, + outlets: SmallVec<[PortId; 1]>, + bytes: Option, + spawns: Vec, + /// The inlets the last run read. + touched: SmallVec<[usize; 4]>, + state: State, + queued: bool, +} + +/// What one run of a pipeline came to. +enum Progress { + /// The source wants this segment read. + Read(SegmentId), + /// The pipeline did work and may do more. + Ran, + /// The pipeline cannot run until the condition clears. + Blocked(Blocked), + /// The pipeline has finished and closed its outlets. + Done, +} + +/// Runs `pipe` until it blocks, finishes, or requests a read. +fn drive( + pipe: &mut Pipeline, + arena: &mut Arena, + session: &VortexSession, + exec: &mut ExecutionCtx, +) -> VortexResult { + let Pipeline { + source, + stages, + inlets, + outlets, + bytes, + spawns, + touched, + .. + } = pipe; + let mut cx = Cx { + arena, + inlets, + bytes, + session, + exec, + spawns, + touched, + }; + loop { + if !outlets.iter().all(|&outlet| cx.arena.has_room(outlet)) { + return Ok(Progress::Blocked(Blocked::Outlet)); + } + if let Some(segment) = source.request() { + return Ok(Progress::Read(segment)); + } + match source.compute(Input::None, &mut cx)? { + // A source's input is its inlets and reads, which nothing changes while it runs, so + // it is asked again at once: it blocks or finishes without a trip through the + // scheduler, and the outlets' room still bounds what it produces. + Step::More(batch) | Step::Last(batch) => push(stages, batch, outlets, &mut cx)?, + Step::Consumed => return Ok(Progress::Ran), + Step::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), + Step::Finished => { + push_end(stages, outlets, &mut cx)?; + return Ok(Progress::Done); + } + } + } +} + +/// Pushes `batch` through `stages` into `outlets`. +fn push( + stages: &mut [Box], + batch: ArrayRef, + outlets: &[PortId], + cx: &mut Cx<'_>, +) -> VortexResult<()> { + let Some((stage, rest)) = stages.split_first_mut() else { + push_out(batch, outlets, cx.arena); + return Ok(()); + }; + let mut input = Input::Chunk(batch); + loop { + match stage.compute(input, cx)? { + Step::More(batch) => { + push(rest, batch, outlets, cx)?; + input = Input::None; + } + Step::Last(batch) => return push(rest, batch, outlets, cx), + Step::Consumed => return Ok(()), + Step::Finished => return push_end(rest, outlets, cx), + Step::Blocked(_) => vortex_bail!("A pipeline stage blocked; only sources block"), + } + } +} + +/// Sends the end through `stages`, then closes `outlets`. +fn push_end( + stages: &mut [Box], + outlets: &[PortId], + cx: &mut Cx<'_>, +) -> VortexResult<()> { + let Some((stage, rest)) = stages.split_first_mut() else { + for &outlet in outlets { + cx.arena.close(outlet); + } + return Ok(()); + }; + loop { + match stage.compute(Input::End, cx)? { + Step::More(batch) => push(rest, batch, outlets, cx)?, + Step::Last(batch) => { + push(rest, batch, outlets, cx)?; + return push_end(rest, outlets, cx); + } + Step::Consumed | Step::Finished => return push_end(rest, outlets, cx), + Step::Blocked(_) => vortex_bail!("A pipeline stage blocked; only sources block"), + } + } +} + +/// Writes `batch` into every outlet: one for a plain pipeline, one per reader for a share. +fn push_out(batch: ArrayRef, outlets: &[PortId], arena: &mut Arena) { + if let Some((&last, rest)) = outlets.split_last() { + for &outlet in rest { + arena.push(outlet, batch.clone()); + } + arena.push(last, batch); + } +} + +/// Everything a scan's pipelines share: the ports, the pipeline table, the scheduler, the reads +/// in flight, and the shares. +pub(crate) struct Core { + pub(crate) session: VortexSession, + pub(crate) exec: ExecutionCtx, + /// The global row index of plan row zero. + pub(crate) row_offset: u64, + pub(crate) arena: Arena, + pipelines: Vec>, + free_pipelines: Vec, + /// Runnable pipelines, last in first out, so data a pipeline just produced is consumed + /// before more is produced. + ready: Vec, + reads: FxHashMap, + next_read: u64, + new_reads: VecDeque, + /// Split slots whose stage output has something new, with a flag per slot so a slot is + /// queued once. + dirty: Vec, + dirty_flags: Vec, +} + +impl Core { + fn new(session: VortexSession, row_offset: u64) -> Self { + Self { + exec: session.create_execution_ctx(), + session, + row_offset, + arena: Arena::default(), + pipelines: Vec::new(), + free_pipelines: Vec::new(), + ready: Vec::new(), + reads: FxHashMap::default(), + next_read: 0, + new_reads: VecDeque::new(), + dirty: Vec::new(), + dirty_flags: Vec::new(), + } + } + + /// Makes `chain` a pipeline writing `outlets`, and queues it to run. + pub(crate) fn add_pipeline( + &mut self, + chain: Chain, + outlets: SmallVec<[PortId; 1]>, + ) -> PipelineId { + let pipeline = Pipeline { + source: chain.source, + stages: chain.stages, + inlets: chain.inlets, + outlets, + bytes: None, + spawns: Vec::new(), + touched: SmallVec::new(), + state: State::Runnable, + queued: true, + }; + let id = match self.free_pipelines.pop() { + Some(id) => id, + None => { + self.pipelines.push(None); + self.pipelines.len() - 1 + } + }; + for &inlet in &pipeline.inlets { + self.arena.get_mut(inlet).reader = Reader::Pipeline(id); + } + for &outlet in &pipeline.outlets { + self.arena.get_mut(outlet).writer = Some(id); + } + self.pipelines[id] = Some(pipeline); + self.ready.push(id); + id + } + + /// Compiles `plan` over `rows`, restricted to `mask`, into pipelines whose output the scan + /// reads for split `split` in slot `slot`, and returns the output port. `None` when nothing + /// is selected. + pub(crate) fn compile_stage( + &mut self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + (slot, split): (usize, usize), + ) -> VortexResult> { + let _ = split; + let Some(chain) = self.compile(plan, rows, mask)? else { + return Ok(None); + }; + let port = self.arena.create(usize::MAX, None, Reader::Split(slot)); + self.add_pipeline(chain, smallvec::smallvec![port]); + Ok(Some(port)) + } + + fn pipeline(&self, id: PipelineId) -> &Pipeline { + self.pipelines[id] + .as_ref() + .unwrap_or_else(|| unreachable!("pipeline {id} is not live")) + } + + fn pipeline_mut(&mut self, id: PipelineId) -> &mut Pipeline { + self.pipelines[id] + .as_mut() + .unwrap_or_else(|| unreachable!("pipeline {id} is not live")) + } + + fn make_runnable(&mut self, id: PipelineId) { + let pipe = self.pipeline_mut(id); + pipe.state = State::Runnable; + if !pipe.queued { + pipe.queued = true; + self.ready.push(id); + } + } + + /// Runs pipeline `id` once and settles what it did. + fn run(&mut self, id: PipelineId) -> VortexResult<()> { + let pipe = self.pipelines[id] + .as_mut() + .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; + pipe.queued = false; + let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; + if !pipe.spawns.is_empty() { + for spawn in std::mem::take(&mut pipe.spawns) { + self.spawn_inlet(id, spawn)?; + } + } + match progress { + Progress::Read(segment_id) => { + let read = ReadId(self.next_read); + self.next_read += 1; + self.reads.insert(read, id); + self.new_reads.push_back(ReadRequest { + id: read, + segment_id, + }); + self.pipeline_mut(id).state = State::Blocked(Blocked::Io); + } + Progress::Ran => self.make_runnable(id), + // Nothing changes a port while its reader computes, so the condition the source + // reported holds until a neighbour wakes it. + Progress::Blocked(blocked) => self.pipeline_mut(id).state = State::Blocked(blocked), + Progress::Done => {} + } + self.wake_neighbours(id); + if matches!(progress, Progress::Done) { + self.retire(id); + } + Ok(()) + } + + /// Compiles what a source asked for into a new inlet of it. + fn spawn_inlet(&mut self, id: PipelineId, spawn: Spawn) -> VortexResult<()> { + let index = self.pipeline(id).inlets.len(); + let capacity = self.pipeline(id).source.capacity(index); + let port = self.arena.create(capacity, None, Reader::Pipeline(id)); + self.pipeline_mut(id).inlets.push(port); + match self.compile(&spawn.plan, spawn.rows, &spawn.mask)? { + Some(chain) => { + self.add_pipeline(chain, smallvec::smallvec![port]); + } + None => self.arena.close(port), + } + Ok(()) + } + + /// Wakes the readers of what pipeline `id` wrote and the writers of what it took. + fn wake_neighbours(&mut self, id: PipelineId) { + let mut wake: SmallVec<[PipelineId; 4]> = SmallVec::new(); + let mut dirty: SmallVec<[usize; 1]> = SmallVec::new(); + let pipe = self.pipeline(id); + for &outlet in &pipe.outlets { + let queue = self.arena.get(outlet); + match queue.reader { + Reader::Pipeline(reader) => { + if queue.readable() + && matches!( + self.pipeline(reader).state, + State::Blocked(Blocked::Inlet(_)) + ) + { + wake.push(reader); + } + } + Reader::Split(slot) => { + if queue.readable() { + dirty.push(slot); + } + } + Reader::Unclaimed => {} + } + } + // Only an inlet the run read can have gained room, so a source with many inlets, as a + // concatenation of many chunks, pays for the inlets it read, not for all of them. + for &index in &pipe.touched { + if let Some(writer) = self.arena.get(pipe.inlets[index]).writer + && writer != id + && self.writer_has_room(writer) + { + wake.push(writer); + } + } + for id in wake { + self.make_runnable(id); + } + self.pipeline_mut(id).touched.clear(); + for slot in dirty { + if slot >= self.dirty_flags.len() { + self.dirty_flags.resize(slot + 1, false); + } + if !self.dirty_flags[slot] { + self.dirty_flags[slot] = true; + self.dirty.push(slot); + } + } + } + + /// Whether `writer` is blocked on its outlets and every one of them has room now. + fn writer_has_room(&self, writer: PipelineId) -> bool { + let pipe = self.pipeline(writer); + pipe.state == State::Blocked(Blocked::Outlet) + && pipe + .outlets + .iter() + .all(|&outlet| self.arena.has_room(outlet)) + } + + /// Drops a finished pipeline. Its outlets are closed; its inlets lose their reader. + fn retire(&mut self, id: PipelineId) { + if let Some(pipe) = self.pipelines[id].take() { + for inlet in pipe.inlets { + self.arena.drop_reader(inlet); + } + self.free_pipelines.push(id); + } + } + + /// The scan has read a stage's output: wakes its writer if the output was full. + fn read_output(&mut self, port: PortId) { + if let Some(writer) = self.arena.get(port).writer + && self.writer_has_room(writer) + { + self.make_runnable(writer); + } + } + + fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { + let id = self + .reads + .remove(&read) + .ok_or_else(|| vortex_err!("Unknown read {read:?}"))?; + self.pipeline_mut(id).bytes = Some(bytes); + self.make_runnable(id); + Ok(()) + } + + fn describe_blocked(&self) -> String { + self.pipelines + .iter() + .enumerate() + .filter_map(|(id, pipe)| pipe.as_ref().map(|pipe| format!("{id}: {:?}", pipe.state))) + .collect::>() + .join(", ") + } +} + +/// What a scan runs over each split. +#[derive(Clone)] +enum Root { + /// The plan, under the split's selection. + Plan(PlanRef), +} + +/// A split being run. +struct Active { + index: usize, + /// The output port of the stage running, read by the scan. + output: Option, +} + +/// Runs a plan over a list of splits. +/// +/// A [`QueryPlan`] at the root is run as stages: zone pruning, then each conjunct under the rows +/// the earlier ones kept, then the projection under the rows they all kept. Any other plan runs +/// as one stage under each split's selection. +/// +/// Up to [`with_max_active`](Self::with_max_active) splits are compiled at once, in order, so +/// reads of later splits are in flight while earlier ones compute. +pub struct Scan { + core: Core, + root: Root, + splits: Vec, + next: usize, + active: Vec>, + live: usize, + outputs: VecDeque<(usize, ArrayRef)>, +} + +impl fmt::Debug for Scan { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Scan") + .field("splits", &self.splits.len()) + .field("next", &self.next) + .field("live", &self.live) + .finish() + } +} + +impl Scan { + /// A scan of `plan` over `splits`, with one split active at a time. + /// + /// Counts, from the plan, how often each split will read each segment, so segments read + /// more than once are decoded once. Builds no pipeline and issues no read. + pub fn try_new( + session: VortexSession, + plan: PlanRef, + splits: Vec, + ) -> VortexResult { + for split in &splits { + vortex_ensure!( + split.rows.start <= split.rows.end + && split.rows.end <= plan.row_count() + && split.mask.len() as u64 == split.rows.end - split.rows.start, + "Split {:?} with a mask of {} rows does not fit a plan of {} rows", + split.rows, + split.mask.len(), + plan.row_count() + ); + } + let root = Root::Plan(plan); + let core = Core::new(session, 0); + Ok(Self { + core, + root, + splits, + next: 0, + active: vec![None], + live: 0, + outputs: VecDeque::new(), + }) + } + + /// Runs up to `max_active` splits at once. + pub fn with_max_active(mut self, max_active: usize) -> Self { + self.active.resize_with(max_active.max(1), || None); + self + } + + /// Sets the global row index of plan row zero, which row-index plans add to their rows. + pub fn with_row_offset(mut self, row_offset: u64) -> Self { + self.core.row_offset = row_offset; + self + } + + /// Runs pipelines until the owner has something to do. + pub fn step(&mut self) -> VortexResult { + loop { + if let Some((split, array)) = self.outputs.pop_front() { + return Ok(Turn::Output(split, array)); + } + if let Some(read) = self.core.new_reads.pop_front() { + return Ok(Turn::Read(read)); + } + if let Some(slot) = self.core.dirty.pop() { + self.core.dirty_flags[slot] = false; + self.advance(slot)?; + continue; + } + if self.live < self.active.len() && self.next < self.splits.len() { + self.activate()?; + continue; + } + if let Some(id) = self.core.ready.pop() { + self.core.run(id)?; + continue; + } + if !self.core.reads.is_empty() { + return Ok(Turn::Waiting); + } + if self.live == 0 { + return Ok(Turn::Done); + } + vortex_bail!( + "Scan is stuck with nothing runnable and no read in flight: {}", + self.core.describe_blocked() + ); + } + } + + /// Delivers the bytes of a read [`step`](Self::step) returned. + pub fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { + self.core.deliver(read, bytes) + } + + /// Starts the next split in a free slot. + fn activate(&mut self) -> VortexResult<()> { + let slot = self + .active + .iter() + .position(Option::is_none) + .ok_or_else(|| vortex_err!("No free split slot"))?; + let index = self.next; + self.next += 1; + let split = self.splits[index].clone(); + let output = match &self.root { + Root::Plan(plan) => { + self.core + .compile_stage(plan, split.rows.clone(), &split.mask, (slot, index))? + } + }; + if output.is_some() { + self.active[slot] = Some(Active { index, output }); + self.live += 1; + } + Ok(()) + } + + /// Takes what the stage running in `slot` produced, and moves the split on when the stage + /// has finished. + fn advance(&mut self, slot: usize) -> VortexResult<()> { + let Some(active) = self.active[slot].as_mut() else { + return Ok(()); + }; + while let Some(port) = active.output { + while let Some(batch) = self.core.arena.get_mut(port).pop() { + self.outputs.push_back((active.index, batch)); + } + self.core.read_output(port); + if !self.core.arena.get(port).closed() { + return Ok(()); + } + self.core.arena.drop_reader(port); + active.output = None; + } + self.active[slot] = None; + self.live -= 1; + Ok(()) + } + + /// Ports in use. + #[cfg(test)] + pub(crate) fn live_ports(&self) -> usize { + self.core.arena.live() + } +} diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs new file mode 100644 index 00000000000..02a1d54aef7 --- /dev/null +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The runtime over hand-built sources: order, blocking, backpressure, and wakeups. + +use std::ops::Range; +use std::sync::Arc; + +use parking_lot::Mutex; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::assert_arrays_eq; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use super::synthetic::Probe; +use super::synthetic::ProbeRun; +use super::synthetic::RowSource; +use super::synthetic::empty_segment; +use super::synthetic::row_dtype; +use super::*; +use crate::plan::PlanRef; +use crate::segments::SegmentId; +use crate::test::SESSION; + +/// The arrays a scan produced, the segments it asked for in order, and when each array came +/// out relative to the deliveries. +struct Run { + arrays: Vec, + /// `Some(segment)` for a delivery, `None` for an array. + events: Vec>, +} + +/// Drives one split of `plan`, delivering reads in the order `pick` chooses, only while the +/// scan waits. +fn run( + plan: &PlanRef, + rows: Range, + mask: Mask, + mut pick: impl FnMut(&[ReadRequest]) -> usize, +) -> VortexResult { + let mut scan = Scan::try_new(SESSION.clone(), plan.clone(), vec![Split { rows, mask }])?; + let mut inflight = Vec::new(); + let mut arrays = Vec::new(); + let mut events = Vec::new(); + loop { + match scan.step()? { + Turn::Read(read) => inflight.push(read), + Turn::Output(_, array) => { + events.push(None); + arrays.push(array); + } + Turn::Waiting => { + if inflight.is_empty() { + return Err(vortex_err!("scan waits with no reads in flight")); + } + let read = inflight.remove(pick(&inflight)); + events.push(Some(*read.segment_id)); + scan.deliver(read.id, empty_segment())?; + } + Turn::Done => break, + } + } + assert_eq!(scan.live_ports(), 0, "a finished scan holds no port"); + Ok(Run { arrays, events }) +} + +fn fifo(_: &[ReadRequest]) -> usize { + 0 +} + +/// Delivers segments in the given order. +fn scripted(order: &[u32]) -> impl FnMut(&[ReadRequest]) -> usize + '_ { + let mut next = order.iter(); + move |inflight| { + let segment = *next.next().expect("script covers every read"); + inflight + .iter() + .position(|read| *read.segment_id == segment) + .expect("scripted segment is in flight") + } +} + +fn indices(values: impl IntoIterator) -> ArrayRef { + Buffer::from_iter(values).into_array() +} + +fn assert_rows(expected: ArrayRef, arrays: Vec) -> VortexResult<()> { + let actual = ChunkedArray::try_new(arrays, row_dtype())?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +fn segment(id: u32) -> Option { + Some(SegmentId::from(id)) +} + +/// A source's pieces come out in order, one array each, after its read if it has one. +#[rstest] +fn source_pieces_come_out_in_order(#[values(None, Some(7))] read: Option) -> VortexResult<()> { + let plan = RowSource::plan(10, vec![3, 3, 4], read.and_then(segment), false); + let run = run(&plan, 0..10, Mask::new_true(10), fifo)?; + assert_eq!( + run.arrays.iter().map(|a| a.len()).collect::>(), + [3, 3, 4] + ); + if let Some(read) = read { + assert_eq!(run.events.first(), Some(&Some(read))); + } + assert_rows(indices(0..10), run.arrays) +} + +/// A selected source emits only the selected rows; a dense one emits every row. +#[rstest] +fn masks_are_applied_by_the_source(#[values(false, true)] dense: bool) -> VortexResult<()> { + let plan = RowSource::plan(8, Vec::new(), None, dense); + let mask = Mask::from_iter((0..8).map(|row| row % 2 == 0)); + let run = run(&plan, 0..8, mask, fifo)?; + let expected = if dense { + indices(0..8) + } else { + indices((0..8).step_by(2)) + }; + assert_rows(expected, run.arrays) +} + +/// A reader runs only when an inlet has changed: every run of the probe after its first finds +/// a batch it has not seen, or an inlet newly closed. +#[rstest] +fn a_reader_runs_only_when_an_inlet_changed( + #[values(&[1, 2], &[2, 1])] order: &[u32], +) -> VortexResult<()> { + let runs: Arc>> = Arc::default(); + let plan = Probe::plan( + vec![ + RowSource::plan(6, vec![2, 2, 2], segment(1), false), + RowSource::plan(6, vec![6], segment(2), false), + ], + 8, + Arc::clone(&runs), + ); + let run = run(&plan, 0..6, Mask::new_true(6), scripted(order))?; + assert_rows(indices(0..6), run.arrays)?; + let runs = runs.lock(); + for pair in runs.windows(2) { + assert_ne!(pair[0], pair[1], "a run saw nothing new: {runs:?}"); + } + Ok(()) +} + +/// A writer is blocked once its reader's inlet holds as many batches as the reader allows, so +/// a fast source never runs far ahead of a slow reader. +#[test] +fn a_full_inlet_blocks_its_writer() -> VortexResult<()> { + let runs: Arc>> = Arc::default(); + let plan = Probe::plan( + vec![RowSource::plan(20, vec![1; 20], None, false)], + 2, + Arc::clone(&runs), + ); + let run = run(&plan, 0..20, Mask::new_true(20), fifo)?; + assert_rows(indices(0..20), run.arrays)?; + // An empty inlet always has room, and the batch in flight may land in a full one. + let deepest = runs.lock().iter().map(|run| run[0].0).max().unwrap_or(0); + assert!(deepest <= 3, "inlet held {deepest} batches"); + Ok(()) +} + +/// Several splits of one scan each produce their own rows. +#[test] +fn splits_produce_their_own_rows() -> VortexResult<()> { + let plan = RowSource::plan(12, vec![2; 6], segment(3), false); + let splits = vec![Split::all(0..4), Split::all(4..9), Split::all(9..12)]; + let mut scan = Scan::try_new(SESSION.clone(), plan, splits)?.with_max_active(2); + let mut by_split: Vec> = vec![Vec::new(); 3]; + let mut inflight = Vec::new(); + loop { + match scan.step()? { + Turn::Read(read) => inflight.push(read), + Turn::Output(split, array) => by_split[split].push(array), + Turn::Waiting => { + let read: ReadRequest = inflight.remove(0); + scan.deliver(read.id, empty_segment())?; + } + Turn::Done => break, + } + } + for (split, rows) in [(0, 0..4), (1, 4..9), (2, 9..12)] { + assert_rows(indices(rows), std::mem::take(&mut by_split[split]))?; + } + Ok(()) +} diff --git a/vortex-layout/src/plan/pipeline/synthetic.rs b/vortex-layout/src/plan/pipeline/synthetic.rs new file mode 100644 index 00000000000..c8f3be38f62 --- /dev/null +++ b/vortex-layout/src/plan/pipeline/synthetic.rs @@ -0,0 +1,290 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Plans whose chains are built by hand, for exercising pipelines without layouts. +//! +//! A [`SyntheticPlan`] carries a closure that compiles it. A scan treats it like any other +//! plan, so a test or benchmark can wire hand-written sources under the real operators and run +//! the whole thing through a [`Scan`](super::Scan), answering reads from memory. + +use std::fmt; +use std::ops::Range; +use std::sync::Arc; + +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use super::Blocked; +use super::Chain; +use super::Compiler; +use super::Cx; +use super::Input; +use super::Operator; +use super::Source; +use super::Step; +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::segments::SegmentId; + +/// Compiles a [`SyntheticPlan`] over some rows, restricted to a mask. +pub type ChainBuilder = + dyn Fn(Range, &Mask, &mut Compiler<'_>) -> VortexResult> + Send + Sync; + +/// A plan operator whose chain is built by a closure. +#[derive(Clone, Debug)] +pub struct Synthetic; + +/// The chain builder of a [`Synthetic`] plan. +#[derive(Clone)] +pub struct SyntheticData { + builder: Arc, +} + +impl fmt::Debug for SyntheticData { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("SyntheticData") + } +} + +/// A plan whose chain is built by a closure. +pub type SyntheticPlan = Plan; + +impl SyntheticPlan { + /// A plan of `row_count` rows of `dtype`, with `children`, compiled by `builder`. + pub fn new( + dtype: DType, + row_count: u64, + children: Vec, + builder: impl Fn(Range, &Mask, &mut Compiler<'_>) -> VortexResult> + + Send + + Sync + + 'static, + ) -> Self { + PlanParts { + vtable: Synthetic, + dtype, + row_count, + children: children.into(), + data: SyntheticData { + builder: Arc::new(builder), + }, + } + .into_typed() + } +} + +impl PlanVTable for Synthetic { + type PlanData = SyntheticData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.synthetic"); + *ID + } + + fn metadata(_plan: &Plan) -> Option { + None + } + + fn with_children( + _plan: &Plan, + _children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + Ok(()) + } + + fn compile( + plan: &Plan, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + (plan.data().builder)(rows, mask, compiler) + } +} + +/// The dtype every synthetic source produces: its rows' indices, as `u64`. +pub fn row_dtype() -> DType { + DType::Primitive(PType::U64, Nullability::NonNullable) +} + +/// A source that emits its rows' indices, in order, cut into pieces of the given lengths, one +/// piece per call, optionally after a round trip through the IO service. +/// +/// A selected source emits only the rows its mask selects, as every plan does. A dense source +/// ignores its mask and emits every row of its range, as a bare segment scan does. Piece +/// lengths count emitted rows; the last piece takes whatever remains. +pub struct RowSource { + /// Remaining indices to emit, in order. + remaining: Vec, + /// Lengths of the pieces still to emit, last first. + pieces: Vec, + /// A segment to read before emitting anything, standing in for a real read. + segment: Option, + waiting: bool, +} + +impl RowSource { + /// A plan of `row_count` rows compiled to a [`RowSource`] emitting pieces of the given + /// lengths. + pub fn plan( + row_count: u64, + pieces: Vec, + segment: Option, + dense: bool, + ) -> PlanRef { + SyntheticPlan::new(row_dtype(), row_count, Vec::new(), move |rows, mask, _| { + let remaining: Vec = if dense || mask.all_true() { + (rows.start..rows.end).collect() + } else { + (rows.start..rows.end) + .zip(mask.iter()) + .filter_map(|(row, selected)| selected.then_some(row)) + .collect() + }; + if remaining.is_empty() { + return Ok(None); + } + Ok(Some(Chain::new(RowSource { + remaining, + pieces: pieces.iter().rev().copied().collect(), + segment, + waiting: false, + }))) + }) + .into_plan() + } +} + +impl Operator for RowSource { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + if self.waiting { + self.waiting = false; + cx.take_bytes() + .ok_or_else(|| vortex_err!("RowSource ran without its bytes"))?; + } + if self.remaining.is_empty() { + return Ok(Step::Finished); + } + let len = self + .pieces + .pop() + .unwrap_or(self.remaining.len()) + .min(self.remaining.len()); + let piece: Vec = self.remaining.drain(..len).collect(); + Ok(Step::More(Buffer::from(piece).into_array())) + } +} + +impl Source for RowSource { + fn request(&mut self) -> Option { + let segment = self.segment.take()?; + self.waiting = true; + Some(segment) + } +} + +/// What a [`Probe`] saw of its inlets on one call: for each, the batches queued and whether +/// its writer had closed it. +pub type ProbeRun = Vec<(usize, bool)>; + +/// A source over one inlet per child that records what it sees each time it runs, then passes +/// its first inlet through, one batch per call, and drops what the others hold. +pub struct Probe { + inlets: usize, + capacity: usize, + runs: Arc>>, +} + +impl Probe { + /// A plan over `children` compiled to a probe whose inlets hold `capacity` batches and + /// which logs into `runs`. + pub fn plan( + children: Vec, + capacity: usize, + runs: Arc>>, + ) -> PlanRef { + let row_count = children.first().map_or(0, |child| child.row_count()); + SyntheticPlan::new( + row_dtype(), + row_count, + children.clone(), + move |rows, mask, compiler| { + let mut chains = Vec::with_capacity(children.len()); + for child in &children { + match compiler.compile(child, rows.clone(), mask)? { + Some(chain) => chains.push(chain), + None => return Ok(None), + } + } + let probe = Probe { + inlets: chains.len(), + capacity, + runs: Arc::clone(&runs), + }; + Ok(Some(compiler.join(chains, probe))) + }, + ) + .into_plan() + } +} + +impl Operator for Probe { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + let run: ProbeRun = (0..self.inlets) + .map(|index| { + let inlet = cx.inlet(index); + (inlet.len(), inlet.closed()) + }) + .collect(); + self.runs.lock().push(run); + for index in 1..self.inlets { + while cx.inlet(index).take().is_some() {} + } + let mut first = cx.inlet(0); + let batch: Option = first.take(); + match batch { + Some(batch) => Ok(Step::Last(batch)), + None if first.closed() => { + for index in 1..self.inlets { + if !cx.inlet(index).closed() { + return Ok(Step::Blocked(Blocked::Inlet(index))); + } + } + Ok(Step::Finished) + } + None => Ok(Step::Blocked(Blocked::Inlet(0))), + } + } +} + +impl Source for Probe { + fn inlet_count(&self) -> usize { + self.inlets + } + + fn capacity(&self, _inlet: usize) -> usize { + self.capacity + } +} + +/// An empty segment, for sources whose read is only a scheduling event. +pub fn empty_segment() -> BufferHandle { + BufferHandle::new_host(Buffer::::empty().into_byte_buffer()) +} diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/typed.rs index 625ec127165..bd00ebffc9c 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/typed.rs @@ -24,8 +24,10 @@ use crate::plan::PlanChildren; use crate::plan::PlanId; use crate::plan::PlanVTable; use crate::plan::display::PlanTreeDisplay; -use crate::plan::exec::ExecContext; -use crate::plan::exec::ExecNode; +use crate::plan::pipeline::Chain; +use crate::plan::pipeline::Compiler; +use crate::plan::pipeline::Reach; +use crate::segments::SegmentId; /// The combined allocation behind [`PlanRef`]. /// @@ -131,25 +133,39 @@ impl PlanRef { Some(unsafe { &*(std::ptr::from_ref(self).cast::>()) }) } - /// Builds the exec node that runs this plan over `rows` of its row domain, restricted to - /// `mask`, with the graph's `ctx`. - pub fn exec( + /// Compiles this plan over `rows` of its row domain, restricted to `mask`. See + /// [`PlanVTable::compile`]. + pub fn compile( &self, rows: Range, - mask: Mask, - ctx: &ExecContext, - ) -> VortexResult> { + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { vortex_ensure!( rows.start <= rows.end && rows.end <= self.row_count(), - "Exec rows {rows:?} exceed plan row count {}", + "Compile rows {rows:?} exceed plan row count {}", self.row_count() ); vortex_ensure!( mask.len() as u64 == rows.end - rows.start, - "Exec mask length {} does not match rows {rows:?}", + "Compile mask length {} does not match rows {rows:?}", mask.len() ); - self.dyn_plan().dyn_exec(self, rows, mask, ctx) + self.dyn_plan().dyn_compile(self, rows, mask, compiler) + } + + /// Visits the segments compiling this plan over `rows` could read. See + /// [`PlanVTable::reach`]. + pub fn reach( + &self, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + if rows.start >= rows.end { + return Ok(()); + } + self.dyn_plan().dyn_reach(self, rows, at, visit) } /// Displays this plan and its descendants with the default plan extractors. @@ -364,14 +380,23 @@ 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>; - /// Builds the exec node for this operator. - fn dyn_exec( + /// Compiles this operator. See [`PlanVTable::compile`]. + fn dyn_compile( + &self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult>; + + /// Visits this operator's segments. See [`PlanVTable::reach`]. + fn dyn_reach( &self, plan: &PlanRef, rows: Range, - mask: Mask, - ctx: &ExecContext, - ) -> VortexResult>; + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()>; } impl DynPlan for PlanData { @@ -404,13 +429,23 @@ impl DynPlan for PlanData { V::metadata(plan.as_::()).map(SerializeMetadata::serialize) } - fn dyn_exec( + fn dyn_compile( + &self, + plan: &PlanRef, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + V::compile(plan.as_::(), rows, mask, compiler) + } + + fn dyn_reach( &self, plan: &PlanRef, rows: Range, - mask: Mask, - ctx: &ExecContext, - ) -> VortexResult> { - V::exec(plan.as_::(), rows, mask, ctx) + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + V::reach(plan.as_::(), rows, at, visit) } } diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/vtable.rs index 416fbf6ce50..74cf44bfcd4 100644 --- a/vortex-layout/src/plan/vtable.rs +++ b/vortex-layout/src/plan/vtable.rs @@ -14,9 +14,11 @@ use vortex_mask::Mask; use vortex_session::registry::Id; use crate::plan::PlanChildren; -use crate::plan::exec::ExecContext; -use crate::plan::exec::ExecNode; +use crate::plan::pipeline::Chain; +use crate::plan::pipeline::Compiler; +use crate::plan::pipeline::Reach; use crate::plan::typed::Plan; +use crate::segments::SegmentId; /// A unique identifier for a plan operator. pub type PlanId = Id; @@ -74,19 +76,36 @@ pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { Cow::Owned(format!("child[{index}]")) } - /// Builds the push-based exec node that runs this plan over `rows` of its row domain, - /// restricted to `mask`, with the graph's `ctx`. `rows` lies within the plan's row domain - /// and `mask` is as long as `rows`. + /// Compiles this plan over `rows` of its row domain, restricted to `mask`, into a chain: + /// a source and the stages after it, producing the rows `mask` selects in row order, or + /// `None` when that is no row. `rows` lies within the plan's row domain and `mask` is as long + /// as `rows`. /// - /// Construction does no IO. Children are spawned and requests published when the graph - /// starts the node. [`ExecNode`] says what the node must then do. - fn exec( + /// A plan compiles its children through `compiler` and either adds stages to a child's + /// chain or joins several children's chains with a source of its own. Compiling does no IO. + fn compile( plan: &Plan, rows: Range, - mask: Mask, - ctx: &ExecContext, - ) -> VortexResult> { - drop((rows, mask, ctx)); - vortex_bail!("Plan {} has no exec implementation", plan.vtable().id()) + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + let _ = (rows, mask, compiler); + vortex_bail!("Plan {} has no pipeline implementation", plan.vtable().id()) + } + + /// Visits every segment compiling this plan over `rows` could read, under any mask, with + /// the rows of the scanned plan it is read for, so a segment more than one reader reads is + /// decoded once. `at` maps this plan's rows to the scanned plan's. + /// + /// The default visits nothing, so the segments under a plan that does not override it are + /// read by each reader alone. + fn reach( + plan: &Plan, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + let _ = (plan, rows, at, visit); + Ok(()) } } From d092cfd1f885e68efd0a673e934a107ad4de7ed2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 19:16:15 +0000 Subject: [PATCH 05/20] refactor(layout): keep Filter out of lower, and let a segment scan keep its own rows A flat layout lowers to a bare segment scan again, as on develop. Every plan now produces only the rows its split selects, a segment scan included, so the scan's source applies the mask; a Filter over a scan compiles to the same source. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/lower.rs | 6 +- vortex-layout/src/plan/pipeline/tests.rs | 45 ++--- vortex-layout/src/plan/plans/segment_scan.rs | 5 +- vortex-layout/src/plan/tests.rs | 190 +++++++------------ 4 files changed, 89 insertions(+), 157 deletions(-) diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index d04e3e128d4..e8911734549 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -25,7 +25,6 @@ use crate::layouts::list::VALIDITY_CHILD_INDEX; use crate::layouts::struct_::Struct; use crate::layouts::struct_::StructLayout; use crate::plan::ConcatPlan; -use crate::plan::FilterPlan; use crate::plan::ListPackPlan; use crate::plan::PackPlan; use crate::plan::PlanChildren; @@ -37,12 +36,9 @@ use crate::plan::TakePlan; /// /// The root operator is built immediately. Its child container owns a hidden clone of the source /// layout and lowers each child independently on first access. -/// -/// A flat layout lowers to a [`Filter`](crate::plan::Filter) over its segment scan, so the scan -/// returns only the rows it is executed with. pub fn lower(layout: &LayoutRef) -> VortexResult { if let Some(layout) = layout.as_opt::() { - return Ok(FilterPlan::new(lower_flat(layout).into_plan()).into_plan()); + return Ok(lower_flat(layout).into_plan()); } if let Some(layout) = layout.as_opt::() { return Ok(lower_chunked(layout)?.into_plan()); diff --git a/vortex-layout/src/plan/pipeline/tests.rs b/vortex-layout/src/plan/pipeline/tests.rs index 3aa732b7e45..41d44dcc3a1 100644 --- a/vortex-layout/src/plan/pipeline/tests.rs +++ b/vortex-layout/src/plan/pipeline/tests.rs @@ -33,7 +33,6 @@ use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::struct_::StructLayout; -use crate::plan::Filter; use crate::plan::PlanRef; use crate::plan::SegmentScan; use crate::plan::lower; @@ -521,42 +520,34 @@ fn aligned_fields_stream_one_struct_per_chunk() -> VortexResult<()> { assert_view(&expected, &(0..ROWS), &mask, run.arrays) } -/// A bare segment scan returns every row of its range whatever it is told to care about, and a -/// filter over the same scan returns only the selected rows. +/// A segment scan keeps the rows its split selects, and so does a filter over the same scan, +/// which runs as the same source. #[rstest] #[case::every_other(Sel::EveryOther)] #[case::sparse(Sel::Rows(&[1, 4]))] #[case::nothing(Sel::None)] -fn bare_scan_is_dense_and_filter_keeps_the_selection(#[case] sel: Sel) -> VortexResult<()> { +fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> VortexResult<()> { let mut store = Store::default(); let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); - let filtered = lower(&store.flat(&values)?)?; - assert!(filtered.is::()); - let scan = filtered.child_required(0)?; + let scan = lower(&store.flat(&values)?)?; assert!(scan.is::()); + let filtered = crate::plan::FilterPlan::new(scan.clone()).into_plan(); let rows = 3..9; let mask = sel.mask(6); - let mut ctx = SESSION.create_execution_ctx(); - - let dense = run( - &store, - &scan, - rows.clone(), - mask.clone(), - delivery(Delivery::Fifo), - )?; - assert_eq!(dense.arrays.len(), 1); - assert_arrays_eq!(dense.arrays[0], values.slice(3..9)?, &mut ctx); - - let kept = run( - &store, - &filtered, - rows.clone(), - mask.clone(), - delivery(Delivery::Fifo), - )?; - assert_view(&values, &rows, &mask, kept.arrays)?; + for plan in [&scan, &filtered] { + let kept = run( + &store, + plan, + rows.clone(), + mask.clone(), + delivery(Delivery::Fifo), + )?; + if mask.all_false() { + assert_eq!(reads(&kept.events), 0, "nothing selected reads nothing"); + } + assert_view(&values, &rows, &mask, kept.arrays)?; + } Ok(()) } diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/plans/segment_scan.rs index 5f455fae561..9ace7e43ff9 100644 --- a/vortex-layout/src/plan/plans/segment_scan.rs +++ b/vortex-layout/src/plan/plans/segment_scan.rs @@ -105,9 +105,8 @@ impl PlanVTable for SegmentScan { mask: &Mask, compiler: &mut Compiler<'_>, ) -> VortexResult> { - let _ = mask; - // A bare scan produces every row of its range; a filter over it keeps the selected ones. - compiler.scan(plan, rows, None) + // A scan keeps the selected rows itself, as every plan produces only those. + compiler.scan(plan, rows, Some(mask.clone())) } fn reach( diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index fccc06ebb2f..783eb455327 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -94,14 +94,12 @@ fn unsupported_layout_has_no_plan() -> VortexResult<()> { } #[test] -fn flat_plan_is_a_filtered_segment_scan_without_children() -> VortexResult<()> { +fn flat_plan_has_no_children() -> VortexResult<()> { let plan = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; - assert!(plan.is::()); - let scan = child_of(&plan, 0)?; - assert!(scan.is::()); - assert_eq!(scan.child_count(), 0); - assert!(scan.child(0)?.is_none()); + assert!(plan.is::()); + assert_eq!(plan.child_count(), 0); + assert!(plan.child(0)?.is_none()); Ok(()) } @@ -325,7 +323,7 @@ fn optimize_drops_identity_expressions() -> VortexResult<()> { let expression = root().bind(child.dtype())?; let plan: PlanRef = EvalPlan::try_new(expression, child)?.into_plan(); - assert!(optimize(plan)?.is::()); + assert!(optimize(plan)?.is::()); Ok(()) } @@ -348,8 +346,8 @@ fn optimize_rewrites_nested_children() -> VortexResult<()> { let optimized = optimize(wrapped)?; assert!(optimized.is::()); - assert!(child_of(&optimized, 0)?.is::()); - assert!(child_of(&optimized, 1)?.is::()); + assert!(child_of(&optimized, 0)?.is::()); + assert!(child_of(&optimized, 1)?.is::()); Ok(()) } @@ -373,13 +371,11 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { let plan: PlanRef = plan.into_plan(); assert_eq!(plan.to_string(), "vortex.plan.eval(i32, rows=3) expr=$.a"); - insta::assert_snapshot!(plan.display_tree(), @" + insta::assert_snapshot!(plan.display_tree(), @r" root: vortex.plan.eval(i32, rows=3) expr=$.a child: vortex.plan.pack({a=i32, b=i32}, rows=3) - a: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) - b: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + a: vortex.plan.segment_scan(i32, rows=3) + b: vortex.plan.segment_scan(i32, rows=3) "); struct DepthExtractor; @@ -395,13 +391,11 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { } } - insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @" + insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @r" root: depth=0 child: depth=1 a: depth=2 - child: depth=3 b: depth=2 - child: depth=3 "); let nullable_fields = StructFields::from_iter([ @@ -419,14 +413,11 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { ) .into_layout(); let nullable = make_plan(nullable_layout)?; - insta::assert_snapshot!(nullable.tree_display_builder(), @" + insta::assert_snapshot!(nullable.tree_display_builder(), @r" root: a: - child: b: - child: validity: - child: "); Ok(()) } @@ -442,12 +433,10 @@ fn chunked_plan_display_names_chunks() -> VortexResult<()> { .into_layout(); let plan = make_plan(layout)?; - insta::assert_snapshot!(plan.display_tree(), @" + insta::assert_snapshot!(plan.display_tree(), @r" root: vortex.plan.concat(i32, rows=3) - chunks[0]: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - chunks[1]: vortex.plan.filter(i32, rows=1) - child: vortex.plan.segment_scan(i32, rows=1) + chunks[0]: vortex.plan.segment_scan(i32, rows=2) + chunks[1]: vortex.plan.segment_scan(i32, rows=1) "); Ok(()) } @@ -461,12 +450,10 @@ fn dict_plan_display_names_logical_children() -> VortexResult<()> { .into_layout(); let plan = make_plan(layout)?; - insta::assert_snapshot!(plan.display_tree(), @" + insta::assert_snapshot!(plan.display_tree(), @r" root: vortex.plan.take(i32, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + codes: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -484,12 +471,10 @@ fn list_plan_display_handles_optional_validity() -> VortexResult<()> { .into_layout(); let non_nullable = make_plan(non_nullable_layout)?; - insta::assert_snapshot!(non_nullable.display_tree(), @" + insta::assert_snapshot!(non_nullable.display_tree(), @r" root: vortex.plan.list_pack(list(i32), rows=2) - elements: vortex.plan.filter(i32, rows=4) - child: vortex.plan.segment_scan(i32, rows=4) - offsets: vortex.plan.filter(u32, rows=3) - child: vortex.plan.segment_scan(u32, rows=3) + elements: vortex.plan.segment_scan(i32, rows=4) + offsets: vortex.plan.segment_scan(u32, rows=3) "); let nullable_layout = ListLayout::new( @@ -501,14 +486,11 @@ fn list_plan_display_handles_optional_validity() -> VortexResult<()> { .into_layout(); let nullable = make_plan(nullable_layout)?; - insta::assert_snapshot!(nullable.display_tree(), @" + insta::assert_snapshot!(nullable.display_tree(), @r" root: vortex.plan.list_pack(list(i32)?, rows=2) - elements: vortex.plan.filter(i32, rows=4) - child: vortex.plan.segment_scan(i32, rows=4) - offsets: vortex.plan.filter(u32, rows=3) - child: vortex.plan.segment_scan(u32, rows=3) - validity: vortex.plan.filter(bool, rows=2) - child: vortex.plan.segment_scan(bool, rows=2) + elements: vortex.plan.segment_scan(i32, rows=4) + offsets: vortex.plan.segment_scan(u32, rows=3) + validity: vortex.plan.segment_scan(bool, rows=2) "); Ok(()) } @@ -564,14 +546,11 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { child: vortex.plan.eval({child_0=bool, child_1=bool}, rows=3) expr=pack(child_0: $.a, child_1: $.b) child: vortex.plan.pack({a=bool, b=bool}, rows=3) a: vortex.plan.take(bool, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) b: vortex.plan.eval(bool, rows=3) expr=($ > 7i32) - child: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -621,18 +600,16 @@ fn row_idx_and_data_expression_pushes_data_into_chunks() -> VortexResult<()> { let expression = and(gt(row_idx(), lit(10_u64)), gt(root(), lit(5_i32))); let plan = optimize(make_row_idx_plan(expression, child)?)?; - insta::assert_snapshot!(plan.display_tree(), @" + insta::assert_snapshot!(plan.display_tree(), @r" root: vortex.plan.eval(bool, rows=3) expr=($.row_idx and $.child) child: vortex.plan.pack({row_idx=bool, child=bool}, rows=3) row_idx: vortex.plan.eval(bool, rows=3) expr=($ > 10u64) child: vortex.plan.row_idx(u64, rows=3) child: vortex.plan.concat(bool, rows=3) chunks[0]: vortex.plan.eval(bool, rows=1) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=1) - child: vortex.plan.segment_scan(i32, rows=1) + child: vortex.plan.segment_scan(i32, rows=1) chunks[1]: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -658,22 +635,17 @@ fn expression_pushes_through_struct_field_and_dictionary_values() -> VortexResul root: vortex.plan.eval(bool, rows=3) expr=($.a > 5i32) child: vortex.plan.pack({a=i32, b=i32}, rows=3) a: vortex.plan.take(i32, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.segment_scan(i32, rows=3) "); let optimized = optimize(plan)?; insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.take(bool, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -708,28 +680,21 @@ fn expression_pushes_through_struct_field_with_heterogeneous_chunks() -> VortexR child: vortex.plan.pack({a=i32, b=i32}, rows=5) a: vortex.plan.concat(i32, rows=5) chunks[0]: vortex.plan.take(i32, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - chunks[1]: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.filter(i32, rows=5) - child: vortex.plan.segment_scan(i32, rows=5) + codes: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.segment_scan(i32, rows=2) + chunks[1]: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.segment_scan(i32, rows=5) "); let optimized = optimize(plan)?; insta::assert_snapshot!(optimized.display_tree(), @" root: vortex.plan.concat(bool, rows=5) chunks[0]: vortex.plan.take(bool, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - chunks[1]: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) child: vortex.plan.segment_scan(i32, rows=2) + chunks[1]: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -763,10 +728,9 @@ fn expression_pushes_through_nested_struct_fields_in_one_pass() -> VortexResult< let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @" + insta::assert_snapshot!(optimized.display_tree(), @r" root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -792,19 +756,17 @@ fn expression_pushes_through_single_field_nested_structs() -> VortexResult<()> { let expression = gt(get_item("b", get_item("a", root())), lit(5_i32)); let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); - insta::assert_snapshot!(plan.display_tree(), @" + insta::assert_snapshot!(plan.display_tree(), @r" root: vortex.plan.eval(bool, rows=3) expr=($.a.b > 5i32) child: vortex.plan.pack({a={b=i32}}, rows=3) a: vortex.plan.pack({b=i32}, rows=3) - b: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + b: vortex.plan.segment_scan(i32, rows=3) "); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @" + insta::assert_snapshot!(optimized.display_tree(), @r" root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -838,10 +800,9 @@ fn expression_pushes_through_three_single_field_structs() -> VortexResult<()> { let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @" + insta::assert_snapshot!(optimized.display_tree(), @r" root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -886,18 +847,15 @@ fn compound_expression_pushes_through_nested_struct_and_dictionary() -> VortexRe let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @" + insta::assert_snapshot!(optimized.display_tree(), @r" root: vortex.plan.eval(bool, rows=3) expr=($.b and $.c) child: vortex.plan.pack({b=bool, c=bool}, rows=3) b: vortex.plan.take(bool, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) c: vortex.plan.eval(bool, rows=3) expr=($ > 7i32) - child: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); Ok(()) } @@ -936,13 +894,11 @@ fn expression_pushes_through_dictionary_of_struct_values() -> VortexResult<()> { let plan = make_eval(expression, make_plan(layout)?)?.into_plan(); let optimized = optimize(plan)?; - insta::assert_snapshot!(optimized.display_tree(), @" + insta::assert_snapshot!(optimized.display_tree(), @r" root: vortex.plan.take(bool, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) "); Ok(()) } @@ -982,14 +938,10 @@ fn multi_field_struct_expression_pushes_into_each_field() -> VortexResult<()> { root: vortex.plan.eval(bool, rows=3) expr=(($.a > 5i32) and ($.b > 7i32)) child: vortex.plan.pack({a=i32, b=i32, c=i32}, rows=3) a: vortex.plan.take(i32, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) - c: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.segment_scan(i32, rows=3) + c: vortex.plan.segment_scan(i32, rows=3) "); let optimized = optimize(plan)?; @@ -997,14 +949,11 @@ fn multi_field_struct_expression_pushes_into_each_field() -> VortexResult<()> { root: vortex.plan.eval(bool, rows=3) expr=($.a and $.b) child: vortex.plan.pack({a=bool, b=bool}, rows=3) a: vortex.plan.take(bool, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) values: vortex.plan.eval(bool, rows=2) expr=($ > 5i32) - child: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) + child: vortex.plan.segment_scan(i32, rows=2) b: vortex.plan.eval(bool, rows=3) expr=($ > 7i32) - child: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) "); let reoptimized = optimize(optimized.clone())?; assert!(PlanRef::ptr_eq(&optimized, &reoptimized)); @@ -1073,12 +1022,9 @@ fn multi_field_struct_expression_keeps_cross_field_refinement() -> VortexResult< root: vortex.plan.eval(bool, rows=3) expr=(($.a + $.b) > 10i32) child: vortex.plan.pack({a=i32, b=i32}, rows=3) a: vortex.plan.take(i32, rows=3) - codes: vortex.plan.filter(u8, rows=3) - child: vortex.plan.segment_scan(u8, rows=3) - values: vortex.plan.filter(i32, rows=2) - child: vortex.plan.segment_scan(i32, rows=2) - b: vortex.plan.filter(i32, rows=3) - child: vortex.plan.segment_scan(i32, rows=3) + codes: vortex.plan.segment_scan(u8, rows=3) + values: vortex.plan.segment_scan(i32, rows=2) + b: vortex.plan.segment_scan(i32, rows=3) "); assert_eq!( optimized.display_tree().to_string(), From 94d7196b948101d34ef8becd2469214bafb76f87 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 19:45:13 +0000 Subject: [PATCH 06/20] refactor(layout): drop mid-run compilation from the pipeline runtime `Cx::spawn` let a source ask for a plan to be compiled into a new inlet while the scan ran. Every plan is now compiled up front: what a stage reads is decided by its plan and its masks before it runs, so the runtime never plans mid-run. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/mod.rs | 20 -------------------- vortex-layout/src/plan/pipeline/scan.rs | 25 ------------------------- 2 files changed, 45 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 47161d379ad..0dfcd93466d 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -26,7 +26,6 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; -use vortex_mask::Mask; use vortex_session::VortexSession; pub use self::compile::Chain; @@ -40,7 +39,6 @@ pub use self::scan::ReadRequest; pub use self::scan::Scan; pub use self::scan::Split; pub use self::scan::Turn; -use crate::plan::PlanRef; use crate::segments::SegmentId; /// Batches an inlet holds before its writer is blocked, unless its reader says otherwise. @@ -119,13 +117,6 @@ pub trait Source: Operator { } } -/// A plan a source asks to have compiled into a new inlet. -pub(crate) struct Spawn { - plan: PlanRef, - rows: std::ops::Range, - mask: Mask, -} - /// What a stage may touch during one call: its pipeline's inlets, the bytes it asked for, and /// the session. pub struct Cx<'a> { @@ -134,7 +125,6 @@ pub struct Cx<'a> { bytes: &'a mut Option, session: &'a VortexSession, exec: &'a mut ExecutionCtx, - spawns: &'a mut Vec, /// The inlets read during this run, so only their writers are checked for room after it. touched: &'a mut SmallVec<[usize; 4]>, } @@ -162,16 +152,6 @@ impl Cx<'_> { pub fn exec(&mut self) -> &mut ExecutionCtx { self.exec } - - /// Asks for `plan` over `rows`, restricted to `mask`, to be compiled into a new inlet of this - /// source once the call returns, and returns the inlet's index. - /// - /// For a source that learns what it must read only from what it has read, as a list learns - /// its elements' range from its offsets. - pub fn spawn(&mut self, plan: PlanRef, rows: std::ops::Range, mask: Mask) -> usize { - self.spawns.push(Spawn { plan, rows, mask }); - self.inlets.len() + self.spawns.len() - 1 - } } #[cfg(test)] diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index bb6adc45ec5..9541f6b7d2d 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -26,7 +26,6 @@ use super::Cx; use super::Input; use super::Operator; use super::Source; -use super::Spawn; use super::Step; use super::compile::Chain; use super::port::Arena; @@ -96,7 +95,6 @@ pub(crate) struct Pipeline { inlets: Vec, outlets: SmallVec<[PortId; 1]>, bytes: Option, - spawns: Vec, /// The inlets the last run read. touched: SmallVec<[usize; 4]>, state: State, @@ -128,7 +126,6 @@ fn drive( inlets, outlets, bytes, - spawns, touched, .. } = pipe; @@ -138,7 +135,6 @@ fn drive( bytes, session, exec, - spawns, touched, }; loop { @@ -276,7 +272,6 @@ impl Core { inlets: chain.inlets, outlets, bytes: None, - spawns: Vec::new(), touched: SmallVec::new(), state: State::Runnable, queued: true, @@ -346,11 +341,6 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; - if !pipe.spawns.is_empty() { - for spawn in std::mem::take(&mut pipe.spawns) { - self.spawn_inlet(id, spawn)?; - } - } match progress { Progress::Read(segment_id) => { let read = ReadId(self.next_read); @@ -375,21 +365,6 @@ impl Core { Ok(()) } - /// Compiles what a source asked for into a new inlet of it. - fn spawn_inlet(&mut self, id: PipelineId, spawn: Spawn) -> VortexResult<()> { - let index = self.pipeline(id).inlets.len(); - let capacity = self.pipeline(id).source.capacity(index); - let port = self.arena.create(capacity, None, Reader::Pipeline(id)); - self.pipeline_mut(id).inlets.push(port); - match self.compile(&spawn.plan, spawn.rows, &spawn.mask)? { - Some(chain) => { - self.add_pipeline(chain, smallvec::smallvec![port]); - } - None => self.arena.close(port), - } - Ok(()) - } - /// Wakes the readers of what pipeline `id` wrote and the writers of what it took. fn wake_neighbours(&mut self, id: PipelineId) { let mut wake: SmallVec<[PipelineId; 4]> = SmallVec::new(); From 82fd9e5b10bb2626524b9bf9795e44a2b1f83bde Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 22:05:09 +0000 Subject: [PATCH 07/20] perf(layout): bound and preallocate pipeline ports, wake only the inlet a reader waits on Ports allocate their buffer once, at their capacity, and keep it when a freed slot is reused. The default capacity rises from 2 to 8 batches, which halves scheduler round trips on narrow pipelines. A stage's output port is bounded like any other, so it needs no growth either; only a share's port, whose reader may be compiled later, can grow past the default. A reader blocked on inlet i is woken only by a batch or close on that inlet, and the inlets a run touched are tracked with a flag per inlet, so a source reading thousands of inlets pays per inlet it read, not per wake. Pipelines and ports live in a slab that never moves its entries. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/mod.rs | 12 ++-- vortex-layout/src/plan/pipeline/port.rs | 93 ++++++++++++++++++++++--- vortex-layout/src/plan/pipeline/scan.rs | 92 +++++++++++++----------- 3 files changed, 142 insertions(+), 55 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 0dfcd93466d..35762f443bf 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -21,7 +21,6 @@ mod port; mod scan; pub mod synthetic; -use smallvec::SmallVec; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::buffer::BufferHandle; @@ -42,7 +41,7 @@ pub use self::scan::Turn; use crate::segments::SegmentId; /// Batches an inlet holds before its writer is blocked, unless its reader says otherwise. -pub const DEFAULT_CAPACITY: usize = 2; +pub const DEFAULT_CAPACITY: usize = 8; /// What a pipeline stage is handed on each call. #[derive(Debug)] @@ -125,14 +124,17 @@ pub struct Cx<'a> { bytes: &'a mut Option, session: &'a VortexSession, exec: &'a mut ExecutionCtx, - /// The inlets read during this run, so only their writers are checked for room after it. - touched: &'a mut SmallVec<[usize; 4]>, + /// The inlets read during this run, each once, so only their writers are checked for room + /// after it, and whether each inlet is listed. + touched: &'a mut Vec, + listed: &'a mut [bool], } impl Cx<'_> { /// The reading side of inlet `index` of this pipeline's source. pub fn inlet(&mut self, index: usize) -> Inlet<'_> { - if self.touched.last() != Some(&index) { + if !self.listed[index] { + self.listed[index] = true; self.touched.push(index); } self.arena.inlet(self.inlets[index]) diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/pipeline/port.rs index 5172e719dc2..098754b3d2b 100644 --- a/vortex-layout/src/plan/pipeline/port.rs +++ b/vortex-layout/src/plan/pipeline/port.rs @@ -8,6 +8,8 @@ use std::collections::VecDeque; use vortex_array::ArrayRef; +use crate::plan::pipeline::DEFAULT_CAPACITY; + /// Identifies a port in the scan's arena. pub(crate) type PortId = usize; @@ -27,6 +29,9 @@ pub(crate) enum Reader { } /// A queue of batches with one writer and one reader. +/// +/// Its buffer is allocated once, for the capacity. Only an unbounded port, a share's waiting for a +/// reader compiled later, grows past the default. pub(crate) struct Queue { batches: VecDeque, /// Batches the queue may hold before its writer is blocked. @@ -44,7 +49,7 @@ pub(crate) struct Queue { impl Queue { fn new(capacity: usize, writer: Option, reader: Reader) -> Self { Self { - batches: VecDeque::new(), + batches: VecDeque::with_capacity(capacity.clamp(1, DEFAULT_CAPACITY)), capacity: capacity.max(1), closed: false, writer, @@ -103,17 +108,82 @@ impl Inlet<'_> { self.queue.batches.len() } + /// Whether the writer can push no more until the reader takes a batch. An unbounded inlet + /// counts as full at the default capacity, so a reader waiting for it to fill still runs. + pub fn full(&self) -> bool { + self.queue.batches.len() >= self.queue.capacity.min(DEFAULT_CAPACITY) + } + + /// Rows queued, over every batch. + pub fn rows(&self) -> usize { + self.queue.batches.iter().map(|batch| batch.len()).sum() + } + /// Whether the inlet has ended: closed with nothing queued. pub fn finished(&self) -> bool { self.queue.closed && self.queue.batches.is_empty() } } +/// A growable table that never moves its entries: entries live in blocks of fixed size, so +/// adding one never copies the others, as a vector's growth does. +pub(crate) struct Slab { + blocks: Vec>, + len: usize, +} + +const BLOCK: usize = 64; + +impl Default for Slab { + fn default() -> Self { + Self { + blocks: Vec::new(), + len: 0, + } + } +} + +impl Slab { + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.len + } + + /// Adds `value`, returning its index. + pub(crate) fn push(&mut self, value: T) -> usize { + if self.len.is_multiple_of(BLOCK) { + self.blocks.push(Vec::with_capacity(BLOCK)); + } + let index = self.len; + self.blocks[index / BLOCK].push(value); + self.len += 1; + index + } + + pub(crate) fn iter(&self) -> impl Iterator { + self.blocks.iter().flatten() + } +} + +impl std::ops::Index for Slab { + type Output = T; + + fn index(&self, index: usize) -> &T { + &self.blocks[index / BLOCK][index % BLOCK] + } +} + +impl std::ops::IndexMut for Slab { + fn index_mut(&mut self, index: usize) -> &mut T { + &mut self.blocks[index / BLOCK][index % BLOCK] + } +} + /// Every port of a scan. Slots are reused once a port's writer has closed it and its reader has /// gone, so no pipeline still refers to it. #[derive(Default)] pub(crate) struct Arena { - queues: Vec, + queues: Slab, free: Vec, } @@ -124,16 +194,21 @@ impl Arena { writer: Option, reader: Reader, ) -> PortId { - let queue = Queue::new(capacity, writer, reader); match self.free.pop() { Some(id) => { - self.queues[id] = queue; + // A freed slot keeps its buffer, so reusing it allocates only for a larger port. + let queue = &mut self.queues[id]; + let capacity = capacity.max(1); + queue.batches.reserve_exact(capacity.min(DEFAULT_CAPACITY)); + queue.capacity = capacity; + queue.closed = false; + queue.writer = writer; + queue.reader = reader; + queue.reader_gone = false; + queue.free = false; id } - None => { - self.queues.push(queue); - self.queues.len() - 1 - } + None => self.queues.push(Queue::new(capacity, writer, reader)), } } @@ -184,7 +259,7 @@ impl Arena { let queue = &mut self.queues[id]; if queue.closed && queue.reader_gone && !queue.free { queue.free = true; - queue.batches = VecDeque::new(); + queue.batches.clear(); self.free.push(id); } } diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index 9541f6b7d2d..fe51bc355a0 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -23,6 +23,7 @@ use vortex_session::VortexSession; use super::Blocked; use super::Cx; +use super::DEFAULT_CAPACITY; use super::Input; use super::Operator; use super::Source; @@ -32,6 +33,7 @@ use super::port::Arena; use super::port::PipelineId; use super::port::PortId; use super::port::Reader; +use super::port::Slab; use crate::plan::PlanRef; use crate::segments::SegmentId; @@ -95,8 +97,10 @@ pub(crate) struct Pipeline { inlets: Vec, outlets: SmallVec<[PortId; 1]>, bytes: Option, - /// The inlets the last run read. - touched: SmallVec<[usize; 4]>, + /// The inlets the last run read, each once. + touched: Vec, + /// For each inlet, whether it is in `touched`. + listed: Vec, state: State, queued: bool, } @@ -127,6 +131,7 @@ fn drive( outlets, bytes, touched, + listed, .. } = pipe; let mut cx = Cx { @@ -136,6 +141,7 @@ fn drive( session, exec, touched, + listed, }; loop { if !outlets.iter().all(|&outlet| cx.arena.has_room(outlet)) { @@ -228,7 +234,7 @@ pub(crate) struct Core { /// The global row index of plan row zero. pub(crate) row_offset: u64, pub(crate) arena: Arena, - pipelines: Vec>, + pipelines: Slab>, free_pipelines: Vec, /// Runnable pipelines, last in first out, so data a pipeline just produced is consumed /// before more is produced. @@ -249,7 +255,7 @@ impl Core { session, row_offset, arena: Arena::default(), - pipelines: Vec::new(), + pipelines: Slab::default(), free_pipelines: Vec::new(), ready: Vec::new(), reads: FxHashMap::default(), @@ -266,22 +272,22 @@ impl Core { chain: Chain, outlets: SmallVec<[PortId; 1]>, ) -> PipelineId { + let listed = vec![false; chain.inlets.len()]; + let touched = Vec::with_capacity(chain.inlets.len()); let pipeline = Pipeline { source: chain.source, stages: chain.stages, inlets: chain.inlets, outlets, bytes: None, - touched: SmallVec::new(), + touched, + listed, state: State::Runnable, queued: true, }; let id = match self.free_pipelines.pop() { Some(id) => id, - None => { - self.pipelines.push(None); - self.pipelines.len() - 1 - } + None => self.pipelines.push(None), }; for &inlet in &pipeline.inlets { self.arena.get_mut(inlet).reader = Reader::Pipeline(id); @@ -308,7 +314,9 @@ impl Core { let Some(chain) = self.compile(plan, rows, mask)? else { return Ok(None); }; - let port = self.arena.create(usize::MAX, None, Reader::Split(slot)); + let port = self + .arena + .create(DEFAULT_CAPACITY, None, Reader::Split(slot)); self.add_pipeline(chain, smallvec::smallvec![port]); Ok(Some(port)) } @@ -367,52 +375,51 @@ impl Core { /// Wakes the readers of what pipeline `id` wrote and the writers of what it took. fn wake_neighbours(&mut self, id: PipelineId) { - let mut wake: SmallVec<[PipelineId; 4]> = SmallVec::new(); - let mut dirty: SmallVec<[usize; 1]> = SmallVec::new(); - let pipe = self.pipeline(id); - for &outlet in &pipe.outlets { - let queue = self.arena.get(outlet); + for index in 0..self.pipeline(id).outlets.len() { + let queue = self.arena.get(self.pipeline(id).outlets[index]); + if !queue.readable() { + continue; + } + let port = self.pipeline(id).outlets[index]; match queue.reader { Reader::Pipeline(reader) => { - if queue.readable() - && matches!( - self.pipeline(reader).state, - State::Blocked(Blocked::Inlet(_)) - ) + // A reader waits on one inlet, so a batch on any other cannot unblock it. + let reader_pipe = self.pipeline(reader); + if let State::Blocked(Blocked::Inlet(waiting)) = reader_pipe.state + && reader_pipe.inlets[waiting] == port { - wake.push(reader); - } - } - Reader::Split(slot) => { - if queue.readable() { - dirty.push(slot); + self.make_runnable(reader); } } + Reader::Split(slot) => self.mark_dirty(slot), Reader::Unclaimed => {} } } // Only an inlet the run read can have gained room, so a source with many inlets, as a // concatenation of many chunks, pays for the inlets it read, not for all of them. - for &index in &pipe.touched { - if let Some(writer) = self.arena.get(pipe.inlets[index]).writer + let touched = std::mem::take(&mut self.pipeline_mut(id).touched); + for &index in &touched { + if let Some(writer) = self.arena.get(self.pipeline(id).inlets[index]).writer && writer != id && self.writer_has_room(writer) { - wake.push(writer); + self.make_runnable(writer); } } - for id in wake { - self.make_runnable(id); + let mut touched = touched; + let pipe = self.pipeline_mut(id); + for &index in &touched { + pipe.listed[index] = false; } - self.pipeline_mut(id).touched.clear(); - for slot in dirty { - if slot >= self.dirty_flags.len() { - self.dirty_flags.resize(slot + 1, false); - } - if !self.dirty_flags[slot] { - self.dirty_flags[slot] = true; - self.dirty.push(slot); - } + touched.clear(); + pipe.touched = touched; + } + + /// Queues split slot `slot` for the scan to read its stage's output. + fn mark_dirty(&mut self, slot: usize) { + if !self.dirty_flags[slot] { + self.dirty_flags[slot] = true; + self.dirty.push(slot); } } @@ -529,7 +536,8 @@ impl Scan { ); } let root = Root::Plan(plan); - let core = Core::new(session, 0); + let mut core = Core::new(session, 0); + core.dirty_flags = vec![false; 1]; Ok(Self { core, root, @@ -544,6 +552,8 @@ impl Scan { /// Runs up to `max_active` splits at once. pub fn with_max_active(mut self, max_active: usize) -> Self { self.active.resize_with(max_active.max(1), || None); + self.core.dirty_flags = vec![false; self.active.len()]; + self.core.dirty = Vec::with_capacity(self.active.len()); self } From fd843a9a5f9d0a5b004b235d903a279606301e75 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 22:07:04 +0000 Subject: [PATCH 08/20] perf(layout): pack joins misaligned fields instead of cutting at every boundary Pack used to emit a struct as soon as every field had a front batch, cut to the shortest one. Fields chunked differently were sliced at every boundary of every other field: 10,000 misaligned columns made 8,192 structs and 80 million slices per scan. Now, when every field's front batch has the same length, Pack takes those batches whole, so aligned fields still stream one struct per chunk with no slicing. Otherwise it waits on the one field with the fewest queued rows until that field is full or closed, then takes that many rows from every field. Whole batches are joined as a ChunkedArray, not copied, and at most one batch per field is sliced. Memory stays bounded by the inlet capacities. Waiting costs O(1) per wake. Queued rows only grow until a struct is built, so Pack resumes at the first inlet it found empty instead of rescanning, and while filling it checks only the inlet it waits on. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/ops/mod.rs | 36 +++-- vortex-layout/src/plan/pipeline/ops/pack.rs | 140 +++++++++++++++----- vortex-layout/src/plan/pipeline/tests.rs | 64 +++++---- vortex-layout/src/plan/plans/concat.rs | 2 +- vortex-layout/src/plan/plans/pack.rs | 4 +- 5 files changed, 177 insertions(+), 69 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/ops/mod.rs b/vortex-layout/src/plan/pipeline/ops/mod.rs index 7ed5d8943a0..f5506551a3d 100644 --- a/vortex-layout/src/plan/pipeline/ops/mod.rs +++ b/vortex-layout/src/plan/pipeline/ops/mod.rs @@ -17,23 +17,37 @@ pub(crate) use filter::*; pub(crate) use pack::*; pub(crate) use port::*; pub(crate) use scan::*; +use smallvec::SmallVec; use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; use vortex_error::VortexResult; use vortex_error::vortex_err; use crate::plan::pipeline::Inlet; -/// Takes the first `len` rows of the inlet's front batch: the batch itself when it is that -/// long, otherwise a slice, leaving the rest in place. +/// Takes the first `len` rows of the inlet, across as many batches as hold them, slicing the last +/// and leaving its rest in place. Rows spanning batches come back chunked, not copied. pub(crate) fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult { - let front = inlet - .peek_mut() - .ok_or_else(|| vortex_err!("Taking rows from an empty inlet"))?; - if front.len() == len { - return inlet - .take() - .ok_or_else(|| vortex_err!("Taking rows from an empty inlet")); + let mut pieces: SmallVec<[ArrayRef; 2]> = SmallVec::new(); + let mut left = len; + while left > 0 { + let front = inlet + .peek_mut() + .ok_or_else(|| vortex_err!("Taking {len} rows from an inlet holding fewer"))?; + if front.len() <= left { + left -= front.len(); + pieces.extend(inlet.take()); + } else { + let rest = front.slice(left..front.len())?; + pieces.push(std::mem::replace(front, rest).slice(0..left)?); + left = 0; + } } - let rest = front.slice(len..front.len())?; - std::mem::replace(front, rest).slice(0..len) + if pieces.len() == 1 { + return Ok(pieces.swap_remove(0)); + } + let dtype = pieces[0].dtype().clone(); + // SAFETY: the pieces come from one inlet, whose batches share its writer's dtype. + Ok(unsafe { ChunkedArray::new_unchecked(pieces, dtype) }.into_array()) } diff --git a/vortex-layout/src/plan/pipeline/ops/pack.rs b/vortex-layout/src/plan/pipeline/ops/pack.rs index a35cb50fbeb..3021d41dce8 100644 --- a/vortex-layout/src/plan/pipeline/ops/pack.rs +++ b/vortex-layout/src/plan/pipeline/ops/pack.rs @@ -6,7 +6,6 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::StructArray; -use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::validity::Validity; @@ -22,12 +21,31 @@ use crate::plan::pipeline::Operator; use crate::plan::pipeline::Source; use crate::plan::pipeline::Step; -/// Zips its inlets, one per field and the validity last when nullable, into structs, taking as -/// many rows from each as the shortest front batch holds. +/// Zips its inlets, one per field and the validity last when nullable, into structs. +/// +/// When every inlet's front batch has the same length, the fields are chunked alike there and the +/// struct takes those batches whole. Otherwise the struct takes as many rows as every inlet has +/// queued, once the inlet with the fewest can queue no more, so fields chunked differently are +/// joined, not cut at every boundary of every field. +/// +/// Waiting costs nothing per wake. Queued rows only grow until the struct is built, so the source +/// resumes at the first inlet it found empty instead of rescanning the others, and then waits on +/// the one inlet with the fewest rows. pub(crate) struct PackSource { fields: StructFields, nullable: bool, inlets: usize, + /// Inlets before this one hold rows. + cursor: usize, + /// The fewest rows queued on an inlet before `cursor` when the scan saw it, and that inlet. + fewest: (usize, usize), + /// The length of the front batch of every inlet before `cursor`, while they agree. + front: Option, + /// Inlets before `cursor` that had ended. + ended: usize, + /// Every inlet holds rows, and the source waits for this one, which holds the fewest, to + /// fill or close. + filling: Option, } impl PackSource { @@ -36,27 +54,77 @@ impl PackSource { fields, nullable, inlets, + cursor: 0, + fewest: (usize::MAX, 0), + front: None, + ended: 0, + filling: None, } } + + /// Finds whether every inlet holds rows, resuming at `cursor`. + fn scan(&mut self, cx: &mut Cx<'_>) -> Option { + while self.cursor < self.inlets { + let mut inlet = cx.inlet(self.cursor); + match inlet.peek_mut().map(|front| front.len()) { + None if inlet.closed() => self.ended += 1, + None => return Some(Step::Blocked(Blocked::Inlet(self.cursor))), + Some(front) => { + self.front = match self.cursor { + 0 => Some(front), + _ => self.front.filter(|&agreed| agreed == front), + }; + let rows = inlet.rows(); + if rows < self.fewest.0 { + self.fewest = (rows, self.cursor); + } + } + } + self.cursor += 1; + } + None + } } impl Operator for PackSource { fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { - let mut len = usize::MAX; - let mut ended = 0; - for index in 0..self.inlets { - let mut inlet = cx.inlet(index); - let closed = inlet.closed(); - match inlet.peek_mut().map(|front| front.len()) { - Some(front) => len = len.min(front), - None if closed => ended += 1, - None => return Ok(Step::Blocked(Blocked::Inlet(index))), + if let Some(index) = self.filling { + let inlet = cx.inlet(index); + if !inlet.full() && !inlet.closed() { + return Ok(Step::Blocked(Blocked::Inlet(index))); } + // Every inlet may have gained rows while the source waited. + self.filling = None; + self.cursor = 0; + self.fewest = (usize::MAX, 0); + self.front = None; } - if ended == self.inlets { + if let Some(blocked) = self.scan(cx) { + return Ok(blocked); + } + if self.ended == self.inlets { return Ok(Step::Finished); } - vortex_ensure!(ended == 0, "Pack fields ended at different rows"); + vortex_ensure!(self.ended == 0, "Pack fields ended at different rows"); + let fewest = self.fewest.1; + let len = match self.front { + Some(front) => front, + None => { + let inlet = cx.inlet(fewest); + if !inlet.full() && !inlet.closed() { + self.filling = Some(fewest); + return Ok(Step::Blocked(Blocked::Inlet(fewest))); + } + // The counts the scan saw are lower bounds by now, so count again. + (0..self.inlets) + .map(|index| cx.inlet(index).rows()) + .min() + .unwrap_or(0) + } + }; + self.cursor = 0; + self.fewest = (usize::MAX, 0); + self.front = None; let mut arrays = Vec::with_capacity(self.inlets); let mut more = true; for index in 0..self.inlets { @@ -73,8 +141,7 @@ impl Operator for PackSource { } else { Validity::NonNullable }; - let array = StructArray::try_new_with_dtype(arrays, self.fields.clone(), len, validity)? - .into_array(); + let array = assemble(&self.fields, arrays, validity).into_array(); Ok(if more { Step::More(array) } else { @@ -91,36 +158,47 @@ impl Source for PackSource { /// Wraps each batch of a pack's only field into a struct. pub(crate) struct WrapStage { - names: FieldNames, + fields: StructFields, } impl WrapStage { - pub(crate) fn new(names: FieldNames) -> Self { - Self { names } + pub(crate) fn new(fields: StructFields) -> Self { + Self { fields } } } impl Operator for WrapStage { fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult { match input { - Input::Chunk(batch) => { - let len = batch.len(); - Ok(Step::Last( - StructArray::try_new( - self.names.clone(), - vec![batch], - len, - Validity::NonNullable, - )? - .into_array(), - )) - } + Input::Chunk(batch) => Ok(Step::Last( + assemble(&self.fields, vec![batch], Validity::NonNullable).into_array(), + )), Input::End => Ok(Step::Finished), Input::None => Ok(Step::Consumed), } } } +/// A struct of `fields` over arrays of one length. +/// +/// The type is not derived or validated per batch: a `Pack` plan checks when it is built that +/// each field's plan produces the field's dtype over the struct's rows, and a pack source takes +/// the same rows from every field. +fn assemble(fields: &StructFields, arrays: Vec, validity: Validity) -> StructArray { + let len = arrays.first().map_or(0, |array| array.len()); + debug_assert!( + arrays.iter().all(|array| array.len() == len) + && fields + .fields() + .zip(&arrays) + .all(|(dtype, array)| &dtype == array.dtype()), + "pack fields do not match the struct" + ); + // SAFETY: the plan validated every field's dtype against `fields`, and every array holds the + // same rows, checked above in debug builds. + unsafe { StructArray::new_unchecked(arrays, fields.clone(), len, validity) } +} + /// The struct of no fields over `len` rows. pub(crate) fn empty_struct( fields: StructFields, diff --git a/vortex-layout/src/plan/pipeline/tests.rs b/vortex-layout/src/plan/pipeline/tests.rs index 41d44dcc3a1..3d8d4845bd1 100644 --- a/vortex-layout/src/plan/pipeline/tests.rs +++ b/vortex-layout/src/plan/pipeline/tests.rs @@ -456,37 +456,55 @@ fn two_columns(store: &mut Store) -> VortexResult<(PlanRef, ArrayRef)> { Ok((lower(&layout)?, expected)) } -/// Pack emits a struct as soon as every field has rows, however the reads complete, and the -/// structs come out in row order, one per chunk of `a`. +/// Fields chunked differently are joined, not cut at each other's boundaries: Pack waits for the +/// field with the fewest rows to fill or close, then emits every row all fields hold, in row +/// order, however the reads complete. /// /// Segments: a0=0, a1=1, a2=2, b=3. #[rstest] -// `b` arrives last: nothing can be emitted before it. -#[case::b_last(&[1, 0, 2, 3], 4)] -// `b` and `a0` first: rows 0..7 go out before `a2` and `a1` arrive. -#[case::a_chunk_last(&[3, 0, 2, 1], 2)] -// Chunks in order: each chunk's struct goes out as it lands. -#[case::in_order(&[3, 0, 1, 2], 2)] -fn pack_streams_in_row_order( - #[case] order: &[u32], - #[case] deliveries_before_first_piece: usize, -) -> VortexResult<()> { +#[case::b_last(&[1, 0, 2, 3])] +#[case::a_chunk_last(&[3, 0, 2, 1])] +#[case::in_order(&[3, 0, 1, 2])] +fn pack_joins_misaligned_fields(#[case] order: &[u32]) -> VortexResult<()> { let mut store = Store::default(); let (plan, expected) = two_columns(&mut store)?; let mask = Mask::new_true(ROWS as usize); let run = run(&store, &plan, 0..ROWS, mask.clone(), scripted(order))?; - assert_eq!(pieces(&run.events), [7, 5, 8]); - let first_piece = run - .events - .iter() - .position(|e| matches!(e, Event::Piece(_))) - .vortex_expect("a piece"); - let delivered = run.events[..first_piece] - .iter() - .filter(|e| matches!(e, Event::Delivered(_))) - .count(); - assert_eq!(delivered, deliveries_before_first_piece); + assert_eq!(pieces(&run.events), [ROWS as usize]); + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} + +/// A misaligned field that fills its inlet makes Pack emit what every field holds, so a struct +/// never waits for more rows than the inlets can queue. +#[test] +fn pack_emits_when_the_shortest_field_fills() -> VortexResult<()> { + let mut store = Store::default(); + let a = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let b = PrimitiveArray::from_iter((0..ROWS as i64).map(|v| v * 10)).into_array(); + let expected = StructArray::from_fields(&[("a", a.clone()), ("b", b.clone())])?.into_array(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![store.chunked(&a, &[1; ROWS as usize])?, store.flat(&b)?], + ) + .into_layout(); + let plan = lower(&layout)?; + let mask = Mask::new_true(ROWS as usize); + + let run = run( + &store, + &plan, + 0..ROWS, + mask.clone(), + delivery(Delivery::Fifo), + )?; + let pieces = pieces(&run.events); + assert!(pieces.len() > 1, "{pieces:?}"); + assert!( + pieces.iter().all(|&piece| piece <= DEFAULT_CAPACITY), + "{pieces:?}" + ); assert_view(&expected, &(0..ROWS), &mask, run.arrays) } diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 852e58fda4b..d2562a6001f 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -157,7 +157,7 @@ impl PlanVTable for Concat { mask: &Mask, compiler: &mut Compiler<'_>, ) -> VortexResult> { - let mut chains = Vec::new(); + let mut chains = Vec::with_capacity(overlapping(plan, &rows).size_hint().1.unwrap_or(0)); for (index, start, local) in overlapping(plan, &rows) { let local_mask = mask.slice( usize::try_from(local.start - rows.start)? diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 991bd9e03c9..89f786ba7fb 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -246,9 +246,7 @@ impl PlanVTable for Pack { if count == 1 && !nullable { // One field needs no zip: each of its batches is wrapped as it passes. let chain = chains.remove(0); - return Ok(Some( - chain.with(WrapStage::new(plan.fields().names().clone())), - )); + return Ok(Some(chain.with(WrapStage::new(plan.fields().clone())))); } let source = PackSource::new(plan.fields().clone(), nullable, count); Ok(Some(compiler.join(chains, source))) From 06f2233c8fc3d204c4ba9be2962fc6eaa43210ff Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Oct 2026 23:02:45 +0000 Subject: [PATCH 09/20] feat(layout): let a pipeline source have several reads in flight A source could ask for one segment at a time: the driver issued its read and did not compute the source until the bytes came back, so a source reading several segments waited for each in turn. Now the driver asks the source for segments until it has none to ask for and issues every read at once. The source is computed as each arrives and takes the bytes tagged with their segment through Cx::take_segment, oldest arrival first. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/mod.rs | 20 ++++-- vortex-layout/src/plan/pipeline/scan.rs | 46 +++++++------ .../src/plan/pipeline/scheduling_tests.rs | 67 +++++++++++++++++++ 3 files changed, 108 insertions(+), 25 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 35762f443bf..1beed39fad1 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -21,6 +21,8 @@ mod port; mod scan; pub mod synthetic; +use std::collections::VecDeque; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::buffer::BufferHandle; @@ -108,9 +110,10 @@ pub trait Source: Operator { DEFAULT_CAPACITY } - /// The segment this source wants read next, asked before every - /// [`compute`](Operator::compute). A source that returns one is not computed until its - /// bytes have arrived through [`Cx::take_bytes`]. + /// A segment this source wants read, asked until it answers `None` before every + /// [`compute`](Operator::compute), so a source may have several reads in flight. A source + /// that asked for any is not computed again until one of them arrives; the bytes come + /// through [`Cx::take_segment`] in the order they arrive. fn request(&mut self) -> Option { None } @@ -121,7 +124,7 @@ pub trait Source: Operator { pub struct Cx<'a> { arena: &'a mut Arena, inlets: &'a [PortId], - bytes: &'a mut Option, + bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, session: &'a VortexSession, exec: &'a mut ExecutionCtx, /// The inlets read during this run, each once, so only their writers are checked for room @@ -140,9 +143,14 @@ impl Cx<'_> { self.arena.inlet(self.inlets[index]) } - /// The bytes of the segment the source requested, once they have arrived. + /// The bytes of a segment the source asked for that has arrived, oldest arrival first. pub fn take_bytes(&mut self) -> Option { - self.bytes.take() + self.take_segment().map(|(_, bytes)| bytes) + } + + /// Like [`take_bytes`](Self::take_bytes), with the segment the bytes are of. + pub fn take_segment(&mut self) -> Option<(SegmentId, BufferHandle)> { + self.bytes.pop_front() } /// The session used for decoding and evaluation. diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index fe51bc355a0..2433fe1c23f 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -96,7 +96,8 @@ pub(crate) struct Pipeline { stages: Vec>, inlets: Vec, outlets: SmallVec<[PortId; 1]>, - bytes: Option, + /// The segments it asked for that have arrived, in order of arrival. + bytes: VecDeque<(SegmentId, BufferHandle)>, /// The inlets the last run read, each once. touched: Vec, /// For each inlet, whether it is in `touched`. @@ -107,8 +108,8 @@ pub(crate) struct Pipeline { /// What one run of a pipeline came to. enum Progress { - /// The source wants this segment read. - Read(SegmentId), + /// The source wants these segments read. + Read(SmallVec<[SegmentId; 1]>), /// The pipeline did work and may do more. Ran, /// The pipeline cannot run until the condition clears. @@ -147,8 +148,13 @@ fn drive( if !outlets.iter().all(|&outlet| cx.arena.has_room(outlet)) { return Ok(Progress::Blocked(Blocked::Outlet)); } - if let Some(segment) = source.request() { - return Ok(Progress::Read(segment)); + if let Some(first) = source.request() { + let mut segments = SmallVec::new(); + segments.push(first); + while let Some(segment) = source.request() { + segments.push(segment); + } + return Ok(Progress::Read(segments)); } match source.compute(Input::None, &mut cx)? { // A source's input is its inlets and reads, which nothing changes while it runs, so @@ -239,7 +245,7 @@ pub(crate) struct Core { /// Runnable pipelines, last in first out, so data a pipeline just produced is consumed /// before more is produced. ready: Vec, - reads: FxHashMap, + reads: FxHashMap, next_read: u64, new_reads: VecDeque, /// Split slots whose stage output has something new, with a flag per slot so a slot is @@ -279,7 +285,7 @@ impl Core { stages: chain.stages, inlets: chain.inlets, outlets, - bytes: None, + bytes: VecDeque::new(), touched, listed, state: State::Runnable, @@ -349,21 +355,23 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; - match progress { - Progress::Read(segment_id) => { - let read = ReadId(self.next_read); - self.next_read += 1; - self.reads.insert(read, id); - self.new_reads.push_back(ReadRequest { - id: read, - segment_id, - }); + match &progress { + Progress::Read(segments) => { + for &segment_id in segments { + let read = ReadId(self.next_read); + self.next_read += 1; + self.reads.insert(read, (id, segment_id)); + self.new_reads.push_back(ReadRequest { + id: read, + segment_id, + }); + } self.pipeline_mut(id).state = State::Blocked(Blocked::Io); } Progress::Ran => self.make_runnable(id), // Nothing changes a port while its reader computes, so the condition the source // reported holds until a neighbour wakes it. - Progress::Blocked(blocked) => self.pipeline_mut(id).state = State::Blocked(blocked), + Progress::Blocked(blocked) => self.pipeline_mut(id).state = State::Blocked(*blocked), Progress::Done => {} } self.wake_neighbours(id); @@ -453,11 +461,11 @@ impl Core { } fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { - let id = self + let (id, segment) = self .reads .remove(&read) .ok_or_else(|| vortex_err!("Unknown read {read:?}"))?; - self.pipeline_mut(id).bytes = Some(bytes); + self.pipeline_mut(id).bytes.push_back((segment, bytes)); self.make_runnable(id); Ok(()) } diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs index 02a1d54aef7..81de957b126 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -21,6 +21,7 @@ use vortex_mask::Mask; use super::synthetic::Probe; use super::synthetic::ProbeRun; use super::synthetic::RowSource; +use super::synthetic::SyntheticPlan; use super::synthetic::empty_segment; use super::synthetic::row_dtype; use super::*; @@ -196,3 +197,69 @@ fn splits_produce_their_own_rows() -> VortexResult<()> { } Ok(()) } + +/// Asks for every segment at once and emits each segment's id as it arrives. +struct ManyReads { + wanted: Vec, + left: usize, +} + +impl Operator for ManyReads { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + if let Some((segment, _bytes)) = cx.take_segment() { + self.left -= 1; + return Ok(Step::More(indices([u64::from(*segment)]))); + } + Ok(if self.left == 0 { + Step::Finished + } else { + Step::Blocked(Blocked::Io) + }) + } +} + +impl Source for ManyReads { + fn request(&mut self) -> Option { + self.wanted.pop() + } +} + +/// A source has all its reads in flight at once, and takes each segment's bytes as they arrive, +/// tagged with the segment. +#[test] +fn a_source_has_several_reads_in_flight() -> VortexResult<()> { + let plan = SyntheticPlan::new(row_dtype(), 3, Vec::new(), |_, _, _| { + Ok(Some(Chain::new(ManyReads { + wanted: [3_u32, 2, 1].into_iter().map(SegmentId::from).collect(), + left: 3, + }))) + }) + .into_plan(); + let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..3)])?; + let mut inflight = Vec::new(); + let mut arrays = Vec::new(); + let mut first_wait = true; + loop { + match scan.step()? { + Turn::Read(read) => inflight.push(read), + Turn::Output(_, array) => arrays.push(array), + Turn::Waiting => { + if first_wait { + assert_eq!( + inflight.len(), + 3, + "every read is asked for before any arrives" + ); + first_wait = false; + } + // Deliver the last asked for first. + let read = inflight + .pop() + .ok_or_else(|| vortex_err!("no read in flight"))?; + scan.deliver(read.id, empty_segment())?; + } + Turn::Done => break, + } + } + assert_rows(indices([3, 2, 1]), arrays) +} From b646fdb62431903495b199769efecfbd64c0da2e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 09:12:49 +0000 Subject: [PATCH 10/20] docs(layout): describe the scan's stages as this stage builds them The scan's docs linked QueryPlan, which a later change adds, and failed the docs build with broken intra-doc links denied. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/scan.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index 2433fe1c23f..f1e19686902 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -496,9 +496,7 @@ struct Active { /// Runs a plan over a list of splits. /// -/// A [`QueryPlan`] at the root is run as stages: zone pruning, then each conjunct under the rows -/// the earlier ones kept, then the projection under the rows they all kept. Any other plan runs -/// as one stage under each split's selection. +/// The plan runs as one stage under each split's selection. /// /// Up to [`with_max_active`](Self::with_max_active) splits are compiled at once, in order, so /// reads of later splits are in flight while earlier ones compute. From 339f74fb35f80ead274c4faeaa114f35a94ee3f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 12:47:37 +0000 Subject: [PATCH 11/20] refactor(layout): source/sink ports, read cancellation, debug validation Ports are named by their ends: a pipeline's sink pushes into its outlets and another pipeline's source reads them as inlets (SourceEnd replaces Reader). Source and Stage replace Operator and Input: a source returns Step::{Batch, Blocked, Finished}, a stage pushes one batch and flushes on finish, and the driver owns full outlets. ConcatSource::new(1) replaces PortSource, and OnceSource moves to its own module. The driver now: - cancels the reads a retiring pipeline left in flight, so their bytes are dropped instead of reaching a pipeline that reuses its id; - reports Done once no split is live, whatever reads are left over; - finishes stages one at a time behind the outlets' room check, so a flush never overfills a port, and closes at once when nothing is left. Debug builds validate the graph after every compile (port ends agree, sources read the inlets they declare, every port has a reader, fan-out is unbounded so no diamond can deadlock) and each run's report (a source waits only on an open inlet with room, or on a read in flight, and finishes only once its inlets close). "Scan is stuck" names the port and sink each waiting source depends on. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/compile.rs | 16 +- vortex-layout/src/plan/pipeline/mod.rs | 110 ++++--- vortex-layout/src/plan/pipeline/port.rs | 122 ++++---- vortex-layout/src/plan/pipeline/scan.rs | 272 +++++++++++------- .../src/plan/pipeline/scan/validate.rs | 105 +++++++ .../src/plan/pipeline/scheduling_tests.rs | 136 ++++++++- vortex-layout/src/plan/pipeline/synthetic.rs | 56 ++-- 7 files changed, 548 insertions(+), 269 deletions(-) create mode 100644 vortex-layout/src/plan/pipeline/scan/validate.rs diff --git a/vortex-layout/src/plan/pipeline/compile.rs b/vortex-layout/src/plan/pipeline/compile.rs index 0f5735b86aa..fa7cd33acb4 100644 --- a/vortex-layout/src/plan/pipeline/compile.rs +++ b/vortex-layout/src/plan/pipeline/compile.rs @@ -10,10 +10,10 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use super::Operator; use super::Source; +use super::Stage; use super::port::PortId; -use super::port::Reader; +use super::port::SourceEnd; use super::scan::Core; use crate::plan::PlanRef; @@ -21,7 +21,7 @@ use crate::plan::PlanRef; /// parent plan can add stages before it becomes a pipeline. pub struct Chain { pub(crate) source: Box, - pub(crate) stages: Vec>, + pub(crate) stages: Vec>, pub(crate) inlets: Vec, } @@ -36,7 +36,7 @@ impl Chain { } /// This chain with `stage` after its last stage. - pub fn with(mut self, stage: impl Operator + 'static) -> Self { + pub fn with(mut self, stage: impl Stage + 'static) -> Self { self.stages.push(Box::new(stage)); self } @@ -82,10 +82,10 @@ impl Compiler<'_> { .into_iter() .enumerate() .map(|(index, chain)| { - let port = self - .core - .arena - .create(source.capacity(index), None, Reader::Unclaimed); + let port = + self.core + .arena + .create(source.capacity(index), None, SourceEnd::Unclaimed); self.core.add_pipeline(chain, smallvec::smallvec![port]); port }) diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 1beed39fad1..5acb7bf878e 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -5,13 +5,14 @@ //! //! A [`Scan`] runs a [`PlanRef`](crate::plan::PlanRef) over a list of splits, row ranges of the //! plan's domain. Each split is compiled into pipelines: a pipeline is one [`Source`], zero or -//! more [`Operator`] stages, and the ports its last stage writes. The driver asks the source for -//! a batch and pushes it through the stages into the ports. Every port has exactly one writer and -//! one reader, and the reader's source declares how many batches it may hold. +//! more [`Stage`]s, and the ports its last stage writes. The driver asks the source for a batch +//! and pushes it through the stages into the ports. Every port runs from one pipeline's sink, +//! the end that pushes into its outlets, to another's source, which reads it as an inlet, and +//! that source declares how many batches it may hold. //! -//! Only a source blocks, and only on one of three things: a segment read it requested, an empty -//! inlet, or a full outlet, which the driver detects. The scheduler re-runs a blocked pipeline -//! when that one condition clears. +//! Only a source blocks, and only on a segment read it requested or an inlet it waits on, which +//! it reports as a [`Blocked`]; a pipeline whose outlet is full is held back by the driver before +//! its source runs. The scheduler re-runs a pipeline when the one condition it waits on clears. //! //! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and //! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. @@ -42,92 +43,82 @@ pub use self::scan::Split; pub use self::scan::Turn; use crate::segments::SegmentId; -/// Batches an inlet holds before its writer is blocked, unless its reader says otherwise. +/// Batches an inlet holds before its sink is blocked, unless its source says otherwise. pub const DEFAULT_CAPACITY: usize = 8; -/// What a pipeline stage is handed on each call. -#[derive(Debug)] -pub enum Input { - /// One batch from the stage below. - Chunk(ArrayRef), - /// The stage below has finished. Flush whatever is held. - End, - /// Nothing new. Sent to a source on every call, and to a stage that returned - /// [`Step::More`]. - None, -} - -/// What a stage says back. +/// What a source says back when asked for its next batch. #[derive(Debug)] pub enum Step { - /// A batch, and more is ready now without further input. The driver pushes the batch on - /// and calls again with [`Input::None`]. - More(ArrayRef), - /// A batch, and the input is spent. - Last(ArrayRef), - /// The input was absorbed and nothing is ready. - Consumed, - /// Nothing can happen until the condition clears. Sources only. + /// A batch for the stage above. The source is asked again at once. + Batch(ArrayRef), + /// No batch until the condition clears. The source is not asked again before it does. Blocked(Blocked), - /// No more output will ever come. The driver sends [`Input::End`] to the next stage. + /// No batch will ever come again. The source is not asked again. Finished, } -/// Why a pipeline cannot run. +/// What a source waits for. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Blocked { - /// A read was requested and has not been delivered. + /// A read it asked for has not arrived. Io, - /// The inlet is empty and its writer has not closed it. + /// Inlet `inlet` is open and has room, so its sink can still push into it or close it, either + /// of which wakes the source. It may already hold batches, as when a source waits for it to + /// fill. Waiting on a closed or full inlet waits forever; debug builds catch it. Inlet(usize), - /// An outlet is full. Produced by the driver, never by a stage. - Outlet, } -/// A stage of a pipeline: takes one input at a time and pushes what it produces on. -/// -/// A stage never blocks. A stage returning [`Step::More`] must eventually return something -/// else when called again with [`Input::None`]. A stage that has returned [`Step::Finished`] is -/// not called again. A stage receiving [`Input::End`] flushes, returning [`Step::More`] while it -/// has more to flush. -pub trait Operator: Send { - /// Takes `input` and says what happens next. - fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult; -} - -/// The start of a pipeline. Reads its inlets and requests segments rather than taking input -/// from a stage below. Always called with [`Input::None`]. -pub trait Source: Operator { +/// The start of a pipeline: produces batches from its inlets and the segments it asks to have +/// read. +pub trait Source: Send { /// The inlets this source reads, numbered from zero. fn inlet_count(&self) -> usize { 0 } - /// How many batches inlet `inlet` may hold before its writer is blocked. Asked once, when + /// How many batches inlet `inlet` may hold before its sink is blocked. Asked once, when /// the port is made. An empty port always has room. fn capacity(&self, inlet: usize) -> usize { let _ = inlet; DEFAULT_CAPACITY } - /// A segment this source wants read, asked until it answers `None` before every - /// [`compute`](Operator::compute), so a source may have several reads in flight. A source - /// that asked for any is not computed again until one of them arrives; the bytes come - /// through [`Cx::take_segment`] in the order they arrive. + /// A segment this source wants read. Asked until it answers `None` before every + /// [`next`](Self::next), so a source may have several reads in flight. A source that asked + /// for any is not asked for a batch again until one of them arrives; the bytes come through + /// [`Cx::take_segment`] in the order they arrive, and a source still waiting for a read it + /// asked for answers [`Blocked::Io`]. fn request(&mut self) -> Option { None } + + /// The next batch, or why there is none. + fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult; } -/// What a stage may touch during one call: its pipeline's inlets, the bytes it asked for, and -/// the session. +/// A stage of a pipeline: transforms the batches pushed into it, one at a time, and never +/// blocks. +pub trait Stage: Send { + /// Takes `batch` from the stage below and returns what to push on, if anything. + fn push(&mut self, batch: ArrayRef, cx: &mut Cx<'_>) -> VortexResult>; + + /// The stage below has finished. Returns whatever the stage still holds, if anything; it is + /// not called again after. + fn finish(&mut self, cx: &mut Cx<'_>) -> VortexResult> { + let _ = cx; + Ok(None) + } +} + +/// What a source or stage may touch during one call: its pipeline's inlets, the bytes the +/// source asked for, and the session. pub struct Cx<'a> { arena: &'a mut Arena, inlets: &'a [PortId], bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, session: &'a VortexSession, exec: &'a mut ExecutionCtx, - /// The inlets read during this run, each once, so only their writers are checked for room + /// The inlets read during this run, each once, so only their sinks are checked for room /// after it, and whether each inlet is listed. touched: &'a mut Vec, listed: &'a mut [bool], @@ -143,12 +134,7 @@ impl Cx<'_> { self.arena.inlet(self.inlets[index]) } - /// The bytes of a segment the source asked for that has arrived, oldest arrival first. - pub fn take_bytes(&mut self) -> Option { - self.take_segment().map(|(_, bytes)| bytes) - } - - /// Like [`take_bytes`](Self::take_bytes), with the segment the bytes are of. + /// A segment the source asked for that has arrived, oldest arrival first, with its bytes. pub fn take_segment(&mut self) -> Option<(SegmentId, BufferHandle)> { self.bytes.pop_front() } diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/pipeline/port.rs index 098754b3d2b..825a1f94e75 100644 --- a/vortex-layout/src/plan/pipeline/port.rs +++ b/vortex-layout/src/plan/pipeline/port.rs @@ -1,8 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Ports: single-writer, single-reader queues of batches, kept in an arena the scan owns so a -//! port can outlive the stage or split that wrote it. +//! Ports: queues of batches with one sink and one source, kept in an arena the scan owns so a +//! port can outlive the pipeline or split at either end. +//! +//! A port has a consumer for as long as anything can be pushed into it: its source goes only +//! once the sink has closed the port and the source has drained it, or once the port is +//! detached from its sink's outlets, so no push ever reaches a port nobody will read. use std::collections::VecDeque; @@ -16,58 +20,73 @@ pub(crate) type PortId = usize; /// Identifies a pipeline in the scan's table. pub(crate) type PipelineId = usize; -/// Who reads a port. +/// The source end of a port: who takes its batches. #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum Reader { - /// Nobody: a share's port waiting for the reader it was made for, or a port whose reader - /// has gone. +pub(crate) enum SourceEnd { + /// Nobody yet: a share's port waiting for the pipeline it was made for, or a port whose + /// source has gone. Unclaimed, - /// A pipeline's source. + /// The source of this pipeline, which reads the port as one of its inlets. Pipeline(PipelineId), - /// The scan, for the output of a split's stage. - Split(usize), + /// The scan, reading the output of the split's stage running in this slot. + Scan(usize), } -/// A queue of batches with one writer and one reader. +/// A queue of batches with one sink and one source. /// /// Its buffer is allocated once, for the capacity. Only an unbounded port, a share's waiting for a -/// reader compiled later, grows past the default. +/// source compiled later, grows past the default. pub(crate) struct Queue { batches: VecDeque, - /// Batches the queue may hold before its writer is blocked. + /// Batches the queue may hold before its sink is blocked. capacity: usize, - /// The writer will push nothing more. + /// The sink will push nothing more. closed: bool, - pub(crate) writer: Option, - pub(crate) reader: Reader, - /// The reader will take nothing more, so pushes are dropped. - reader_gone: bool, + /// The pipeline whose sink pushes into this port; none once closed. + pub(crate) sink: Option, + pub(crate) source: SourceEnd, + /// The source will take nothing more. Set only once nothing more can be pushed, so the + /// slot is freed as soon as the port is also closed. + source_gone: bool, /// The slot is free for reuse. free: bool, } impl Queue { - fn new(capacity: usize, writer: Option, reader: Reader) -> Self { + fn new(capacity: usize, sink: Option, source: SourceEnd) -> Self { Self { batches: VecDeque::with_capacity(capacity.clamp(1, DEFAULT_CAPACITY)), capacity: capacity.max(1), closed: false, - writer, - reader, - reader_gone: false, + sink, + source, + source_gone: false, free: false, } } pub(crate) fn has_room(&self) -> bool { - self.reader_gone || self.batches.len() < self.capacity + self.batches.len() < self.capacity } pub(crate) fn closed(&self) -> bool { self.closed } - /// Whether the reader would find something new: a batch, or the close. + /// Whether the sink is never blocked by it. + pub(crate) fn unbounded(&self) -> bool { + self.capacity == usize::MAX + } + + pub(crate) fn len(&self) -> usize { + self.batches.len() + } + + pub(crate) fn is_free(&self) -> bool { + self.free + } + + /// Whether the source would find something new: a batch, or the close. pub(crate) fn readable(&self) -> bool { !self.batches.is_empty() || self.closed } @@ -83,7 +102,7 @@ pub struct Inlet<'a> { } impl Inlet<'_> { - /// The oldest batch, in place, so the reader can slice it down without removing it. + /// The oldest batch, in place, so the source can slice it down without removing it. pub fn peek_mut(&mut self) -> Option<&mut ArrayRef> { self.queue.batches.front_mut() } @@ -93,7 +112,7 @@ impl Inlet<'_> { self.queue.batches.pop_front() } - /// Whether the writer has finished. Closed and empty is the end of the inlet. + /// Whether the sink has finished. Closed and empty is the end of the inlet. pub fn closed(&self) -> bool { self.queue.closed } @@ -108,8 +127,8 @@ impl Inlet<'_> { self.queue.batches.len() } - /// Whether the writer can push no more until the reader takes a batch. An unbounded inlet - /// counts as full at the default capacity, so a reader waiting for it to fill still runs. + /// Whether the sink can push no more until the source takes a batch. An unbounded inlet + /// counts as full at the default capacity, so a source waiting for it to fill still runs. pub fn full(&self) -> bool { self.queue.batches.len() >= self.queue.capacity.min(DEFAULT_CAPACITY) } @@ -118,11 +137,6 @@ impl Inlet<'_> { pub fn rows(&self) -> usize { self.queue.batches.iter().map(|batch| batch.len()).sum() } - - /// Whether the inlet has ended: closed with nothing queued. - pub fn finished(&self) -> bool { - self.queue.closed && self.queue.batches.is_empty() - } } /// A growable table that never moves its entries: entries live in blocks of fixed size, so @@ -179,7 +193,7 @@ impl std::ops::IndexMut for Slab { } } -/// Every port of a scan. Slots are reused once a port's writer has closed it and its reader has +/// Every port of a scan. Slots are reused once a port's sink has closed it and its source has /// gone, so no pipeline still refers to it. #[derive(Default)] pub(crate) struct Arena { @@ -191,8 +205,8 @@ impl Arena { pub(crate) fn create( &mut self, capacity: usize, - writer: Option, - reader: Reader, + sink: Option, + source: SourceEnd, ) -> PortId { match self.free.pop() { Some(id) => { @@ -202,13 +216,13 @@ impl Arena { queue.batches.reserve_exact(capacity.min(DEFAULT_CAPACITY)); queue.capacity = capacity; queue.closed = false; - queue.writer = writer; - queue.reader = reader; - queue.reader_gone = false; + queue.sink = sink; + queue.source = source; + queue.source_gone = false; queue.free = false; id } - None => self.queues.push(Queue::new(capacity, writer, reader)), + None => self.queues.push(Queue::new(capacity, sink, source)), } } @@ -230,40 +244,52 @@ impl Arena { self.queues[id].has_room() } - /// Appends a batch, unless the reader has gone. Empty batches are dropped. + /// Appends a batch. Empty batches are dropped. pub(crate) fn push(&mut self, id: PortId, batch: ArrayRef) { let queue = &mut self.queues[id]; - if !queue.reader_gone && !batch.is_empty() { + debug_assert!( + !queue.source_gone && !queue.closed, + "push into port {id} with no consumer" + ); + if !batch.is_empty() { queue.batches.push_back(batch); } } - /// The writer will push nothing more. + /// The sink will push nothing more. pub(crate) fn close(&mut self, id: PortId) { let queue = &mut self.queues[id]; queue.closed = true; - queue.writer = None; + queue.sink = None; self.maybe_free(id); } - /// The reader will take nothing more: what is queued is dropped, and so is what is pushed. - pub(crate) fn drop_reader(&mut self, id: PortId) { + /// The source will take nothing more: what is queued is dropped. Nothing more may be + /// pushed, so the sink must have closed the port or be about to, as a retiring pipeline's + /// inlets' sinks have. + pub(crate) fn drop_source(&mut self, id: PortId) { let queue = &mut self.queues[id]; - queue.reader_gone = true; - queue.reader = Reader::Unclaimed; + queue.source_gone = true; + queue.source = SourceEnd::Unclaimed; queue.batches.clear(); self.maybe_free(id); } fn maybe_free(&mut self, id: PortId) { let queue = &mut self.queues[id]; - if queue.closed && queue.reader_gone && !queue.free { + if queue.closed && queue.source_gone && !queue.free { queue.free = true; queue.batches.clear(); self.free.push(id); } } + /// Every port slot, free or not, with its id. + #[cfg(debug_assertions)] + pub(crate) fn iter(&self) -> impl Iterator { + self.queues.iter().enumerate() + } + /// Ports in use. #[cfg(test)] pub(crate) fn live(&self) -> usize { diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index f1e19686902..07abf20dd90 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -24,19 +24,21 @@ use vortex_session::VortexSession; use super::Blocked; use super::Cx; use super::DEFAULT_CAPACITY; -use super::Input; -use super::Operator; use super::Source; +use super::Stage; use super::Step; use super::compile::Chain; use super::port::Arena; use super::port::PipelineId; use super::port::PortId; -use super::port::Reader; use super::port::Slab; +use super::port::SourceEnd; use crate::plan::PlanRef; use crate::segments::SegmentId; +#[cfg(debug_assertions)] +mod validate; + /// Identifies a read a scan asked its owner for. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub struct ReadId(u64); @@ -83,21 +85,30 @@ pub enum Turn { Done, } -/// Where a pipeline is in its life. +/// Whether a pipeline can run, and if not, what it waits for. #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum State { Runnable, + /// What its source said it waits for. Blocked(Blocked), + /// Room in a full outlet, found by the driver before the source ran. + Outlet, } /// A source, its stages, and the ports they read and write. pub(crate) struct Pipeline { source: Box, - stages: Vec>, + stages: Vec>, inlets: Vec, outlets: SmallVec<[PortId; 1]>, /// The segments it asked for that have arrived, in order of arrival. bytes: VecDeque<(SegmentId, BufferHandle)>, + /// Reads it asked for that have not arrived, so retiring it cancels them only when any + /// are left. + pending: usize, + /// Once its source has finished, the next stage to finish: stages are finished one at a + /// time, each when the outlets have room for what it flushes. + finishing: Option, /// The inlets the last run read, each once. touched: Vec, /// For each inlet, whether it is in `touched`. @@ -110,10 +121,10 @@ pub(crate) struct Pipeline { enum Progress { /// The source wants these segments read. Read(SmallVec<[SegmentId; 1]>), - /// The pipeline did work and may do more. - Ran, - /// The pipeline cannot run until the condition clears. + /// The source waits for this. Blocked(Blocked), + /// An outlet is full. + Outlet, /// The pipeline has finished and closed its outlets. Done, } @@ -131,6 +142,7 @@ fn drive( inlets, outlets, bytes, + finishing, touched, listed, .. @@ -145,8 +157,25 @@ fn drive( listed, }; loop { + // Closing needs no room, so a pipeline with nothing left to flush closes at once. + if *finishing == Some(stages.len()) { + for &outlet in outlets.iter() { + cx.arena.close(outlet); + } + return Ok(Progress::Done); + } if !outlets.iter().all(|&outlet| cx.arena.has_room(outlet)) { - return Ok(Progress::Blocked(Blocked::Outlet)); + return Ok(Progress::Outlet); + } + if let Some(stage) = *finishing { + let (this, rest) = stages[stage..] + .split_first_mut() + .unwrap_or_else(|| unreachable!("stage {stage} finished twice")); + *finishing = Some(stage + 1); + if let Some(batch) = this.finish(&mut cx)? { + push(rest, batch, outlets, &mut cx)?; + } + continue; } if let Some(first) = source.request() { let mut segments = SmallVec::new(); @@ -156,73 +185,35 @@ fn drive( } return Ok(Progress::Read(segments)); } - match source.compute(Input::None, &mut cx)? { + match source.next(&mut cx)? { // A source's input is its inlets and reads, which nothing changes while it runs, so // it is asked again at once: it blocks or finishes without a trip through the // scheduler, and the outlets' room still bounds what it produces. - Step::More(batch) | Step::Last(batch) => push(stages, batch, outlets, &mut cx)?, - Step::Consumed => return Ok(Progress::Ran), + Step::Batch(batch) => push(stages, batch, outlets, &mut cx)?, Step::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), - Step::Finished => { - push_end(stages, outlets, &mut cx)?; - return Ok(Progress::Done); - } + Step::Finished => *finishing = Some(0), } } } /// Pushes `batch` through `stages` into `outlets`. fn push( - stages: &mut [Box], - batch: ArrayRef, + stages: &mut [Box], + mut batch: ArrayRef, outlets: &[PortId], cx: &mut Cx<'_>, ) -> VortexResult<()> { - let Some((stage, rest)) = stages.split_first_mut() else { - push_out(batch, outlets, cx.arena); - return Ok(()); - }; - let mut input = Input::Chunk(batch); - loop { - match stage.compute(input, cx)? { - Step::More(batch) => { - push(rest, batch, outlets, cx)?; - input = Input::None; - } - Step::Last(batch) => return push(rest, batch, outlets, cx), - Step::Consumed => return Ok(()), - Step::Finished => return push_end(rest, outlets, cx), - Step::Blocked(_) => vortex_bail!("A pipeline stage blocked; only sources block"), - } - } -} - -/// Sends the end through `stages`, then closes `outlets`. -fn push_end( - stages: &mut [Box], - outlets: &[PortId], - cx: &mut Cx<'_>, -) -> VortexResult<()> { - let Some((stage, rest)) = stages.split_first_mut() else { - for &outlet in outlets { - cx.arena.close(outlet); - } - return Ok(()); - }; - loop { - match stage.compute(Input::End, cx)? { - Step::More(batch) => push(rest, batch, outlets, cx)?, - Step::Last(batch) => { - push(rest, batch, outlets, cx)?; - return push_end(rest, outlets, cx); - } - Step::Consumed | Step::Finished => return push_end(rest, outlets, cx), - Step::Blocked(_) => vortex_bail!("A pipeline stage blocked; only sources block"), - } + for stage in stages { + batch = match stage.push(batch, cx)? { + Some(batch) => batch, + None => return Ok(()), + }; } + push_out(batch, outlets, cx.arena); + Ok(()) } -/// Writes `batch` into every outlet: one for a plain pipeline, one per reader for a share. +/// Writes `batch` into every outlet: one for a plain pipeline, one per claimant for a share. fn push_out(batch: ArrayRef, outlets: &[PortId], arena: &mut Arena) { if let Some((&last, rest)) = outlets.split_last() { for &outlet in rest { @@ -232,8 +223,8 @@ fn push_out(batch: ArrayRef, outlets: &[PortId], arena: &mut Arena) { } } -/// Everything a scan's pipelines share: the ports, the pipeline table, the scheduler, the reads -/// in flight, and the shares. +/// Everything a scan's pipelines share: the ports, the pipeline table, the scheduler, and the +/// reads in flight. pub(crate) struct Core { pub(crate) session: VortexSession, pub(crate) exec: ExecutionCtx, @@ -245,7 +236,9 @@ pub(crate) struct Core { /// Runnable pipelines, last in first out, so data a pipeline just produced is consumed /// before more is produced. ready: Vec, - reads: FxHashMap, + /// Reads in flight: the pipeline that asked, or `None` once that pipeline has finished and + /// the read is cancelled, and the segment. + reads: FxHashMap, SegmentId)>, next_read: u64, new_reads: VecDeque, /// Split slots whose stage output has something new, with a flag per slot so a slot is @@ -286,6 +279,8 @@ impl Core { inlets: chain.inlets, outlets, bytes: VecDeque::new(), + pending: 0, + finishing: None, touched, listed, state: State::Runnable, @@ -296,10 +291,10 @@ impl Core { None => self.pipelines.push(None), }; for &inlet in &pipeline.inlets { - self.arena.get_mut(inlet).reader = Reader::Pipeline(id); + self.arena.get_mut(inlet).source = SourceEnd::Pipeline(id); } for &outlet in &pipeline.outlets { - self.arena.get_mut(outlet).writer = Some(id); + self.arena.get_mut(outlet).sink = Some(id); } self.pipelines[id] = Some(pipeline); self.ready.push(id); @@ -322,8 +317,10 @@ impl Core { }; let port = self .arena - .create(DEFAULT_CAPACITY, None, Reader::Split(slot)); + .create(DEFAULT_CAPACITY, None, SourceEnd::Scan(slot)); self.add_pipeline(chain, smallvec::smallvec![port]); + #[cfg(debug_assertions)] + self.validate(); Ok(Some(port)) } @@ -355,23 +352,21 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; + #[cfg(debug_assertions)] + self.validate_progress(id, &progress); match &progress { Progress::Read(segments) => { for &segment_id in segments { - let read = ReadId(self.next_read); - self.next_read += 1; - self.reads.insert(read, (id, segment_id)); - self.new_reads.push_back(ReadRequest { - id: read, - segment_id, - }); + self.read(id, segment_id); } - self.pipeline_mut(id).state = State::Blocked(Blocked::Io); + let pipe = self.pipeline_mut(id); + pipe.pending += segments.len(); + pipe.state = State::Blocked(Blocked::Io); } - Progress::Ran => self.make_runnable(id), - // Nothing changes a port while its reader computes, so the condition the source - // reported holds until a neighbour wakes it. + // Nothing changes a port while its source's pipeline runs, so the condition the run reported + // holds until a neighbour wakes it. Progress::Blocked(blocked) => self.pipeline_mut(id).state = State::Blocked(*blocked), + Progress::Outlet => self.pipeline_mut(id).state = State::Outlet, Progress::Done => {} } self.wake_neighbours(id); @@ -381,7 +376,7 @@ impl Core { Ok(()) } - /// Wakes the readers of what pipeline `id` wrote and the writers of what it took. + /// Wakes the sources of what pipeline `id` pushed and the sinks of what it took. fn wake_neighbours(&mut self, id: PipelineId) { for index in 0..self.pipeline(id).outlets.len() { let queue = self.arena.get(self.pipeline(id).outlets[index]); @@ -389,29 +384,29 @@ impl Core { continue; } let port = self.pipeline(id).outlets[index]; - match queue.reader { - Reader::Pipeline(reader) => { - // A reader waits on one inlet, so a batch on any other cannot unblock it. - let reader_pipe = self.pipeline(reader); - if let State::Blocked(Blocked::Inlet(waiting)) = reader_pipe.state - && reader_pipe.inlets[waiting] == port + match queue.source { + SourceEnd::Pipeline(source) => { + // A source waits on one inlet, so a batch on any other cannot unblock it. + let source_pipe = self.pipeline(source); + if let State::Blocked(Blocked::Inlet(waiting)) = source_pipe.state + && source_pipe.inlets[waiting] == port { - self.make_runnable(reader); + self.make_runnable(source); } } - Reader::Split(slot) => self.mark_dirty(slot), - Reader::Unclaimed => {} + SourceEnd::Scan(slot) => self.mark_dirty(slot), + SourceEnd::Unclaimed => {} } } // Only an inlet the run read can have gained room, so a source with many inlets, as a // concatenation of many chunks, pays for the inlets it read, not for all of them. let touched = std::mem::take(&mut self.pipeline_mut(id).touched); for &index in &touched { - if let Some(writer) = self.arena.get(self.pipeline(id).inlets[index]).writer - && writer != id - && self.writer_has_room(writer) + if let Some(sink) = self.arena.get(self.pipeline(id).inlets[index]).sink + && sink != id + && self.sink_has_room(sink) { - self.make_runnable(writer); + self.make_runnable(sink); } } let mut touched = touched; @@ -431,50 +426,97 @@ impl Core { } } - /// Whether `writer` is blocked on its outlets and every one of them has room now. - fn writer_has_room(&self, writer: PipelineId) -> bool { - let pipe = self.pipeline(writer); - pipe.state == State::Blocked(Blocked::Outlet) + /// Whether pipeline `sink` is blocked on its outlets and every one of them has room now. + fn sink_has_room(&self, sink: PipelineId) -> bool { + let pipe = self.pipeline(sink); + pipe.state == State::Outlet && pipe .outlets .iter() .all(|&outlet| self.arena.has_room(outlet)) } - /// Drops a finished pipeline. Its outlets are closed; its inlets lose their reader. + /// Drops a finished pipeline. Its outlets are closed; its inlets lose their source; the reads + /// it asked for and never took are cancelled, so their bytes are dropped on arrival instead + /// of reaching a pipeline that reuses its id. fn retire(&mut self, id: PipelineId) { if let Some(pipe) = self.pipelines[id].take() { for inlet in pipe.inlets { - self.arena.drop_reader(inlet); + self.arena.drop_source(inlet); + } + if pipe.pending > 0 { + self.cancel_reads(id); } self.free_pipelines.push(id); } } - /// The scan has read a stage's output: wakes its writer if the output was full. + /// Cancels the reads pipeline `id` waits for: they become reads for no pipeline, whose + /// bytes are dropped on arrival. + fn cancel_reads(&mut self, id: PipelineId) { + for (owner, _) in self.reads.values_mut() { + if *owner == Some(id) { + *owner = None; + } + } + } + + /// The scan has read a stage's output: wakes its sink if the output was full. fn read_output(&mut self, port: PortId) { - if let Some(writer) = self.arena.get(port).writer - && self.writer_has_room(writer) + if let Some(sink) = self.arena.get(port).sink + && self.sink_has_room(sink) { - self.make_runnable(writer); + self.make_runnable(sink); } } fn deliver(&mut self, read: ReadId, bytes: BufferHandle) -> VortexResult<()> { - let (id, segment) = self + let (owner, segment) = self .reads .remove(&read) .ok_or_else(|| vortex_err!("Unknown read {read:?}"))?; - self.pipeline_mut(id).bytes.push_back((segment, bytes)); - self.make_runnable(id); + if let Some(id) = owner { + let pipe = self.pipeline_mut(id); + pipe.pending -= 1; + pipe.bytes.push_back((segment, bytes)); + self.make_runnable(id); + } Ok(()) } + /// Issues a read of `segment` for pipeline `owner`. + fn read(&mut self, owner: PipelineId, segment_id: SegmentId) { + let read = ReadId(self.next_read); + self.next_read += 1; + self.reads.insert(read, (Some(owner), segment_id)); + self.new_reads.push_back(ReadRequest { + id: read, + segment_id, + }); + } + + /// Every live pipeline's state, and for one waiting on an inlet, that port and its sink. fn describe_blocked(&self) -> String { self.pipelines .iter() .enumerate() - .filter_map(|(id, pipe)| pipe.as_ref().map(|pipe| format!("{id}: {:?}", pipe.state))) + .filter_map(|(id, pipe)| { + let pipe = pipe.as_ref()?; + Some(match pipe.state { + State::Blocked(Blocked::Inlet(index)) => { + let port = pipe.inlets[index]; + let queue = self.arena.get(port); + format!( + "{id}: {:?} (port {port}: {} queued, closed {}, sink {:?})", + pipe.state, + queue.len(), + queue.closed(), + queue.sink.map(|sink| (sink, self.pipeline(sink).state)), + ) + } + state => format!("{id}: {state:?}"), + }) + }) .collect::>() .join(", ") } @@ -591,12 +633,14 @@ impl Scan { self.core.run(id)?; continue; } - if !self.core.reads.is_empty() { - return Ok(Turn::Waiting); - } + // Reads still in flight once every split has finished were cancelled; nothing waits + // for their bytes. if self.live == 0 { return Ok(Turn::Done); } + if !self.core.reads.is_empty() { + return Ok(Turn::Waiting); + } vortex_bail!( "Scan is stuck with nothing runnable and no read in flight: {}", self.core.describe_blocked() @@ -646,7 +690,7 @@ impl Scan { if !self.core.arena.get(port).closed() { return Ok(()); } - self.core.arena.drop_reader(port); + self.core.arena.drop_source(port); active.output = None; } self.active[slot] = None; @@ -659,4 +703,14 @@ impl Scan { pub(crate) fn live_ports(&self) -> usize { self.core.arena.live() } + + /// Reads in flight that a pipeline waits for. + #[cfg(test)] + pub(crate) fn awaited_reads(&self) -> usize { + self.core + .reads + .values() + .filter(|(owner, _)| owner.is_some()) + .count() + } } diff --git a/vortex-layout/src/plan/pipeline/scan/validate.rs b/vortex-layout/src/plan/pipeline/scan/validate.rs new file mode 100644 index 00000000000..133769990e5 --- /dev/null +++ b/vortex-layout/src/plan/pipeline/scan/validate.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Checks of the pipeline graph and of what each run reports, compiled into debug builds only. +//! +//! The scheduler trusts sources to report truthfully and the compiler to wire ports so no cycle +//! of full ports can form. A broken promise shows up much later as "Scan is stuck" or as a +//! wrong result; these checks fail at the pipeline that broke it. + +use super::Blocked; +use super::Core; +use super::PipelineId; +use super::Progress; +use super::SourceEnd; + +impl Core { + /// Checks the graph after a compile: ports and pipelines agree on their ends, sources read + /// as many inlets as they declare, every port has a reader, and fan-out is unbounded. + pub(super) fn validate(&self) { + for (id, pipe) in self.pipelines.iter().enumerate() { + let Some(pipe) = pipe else { continue }; + assert_eq!( + pipe.inlets.len(), + pipe.source.inlet_count(), + "pipeline {id} reads {} inlets but its source declares {}", + pipe.inlets.len(), + pipe.source.inlet_count() + ); + for &port in &pipe.inlets { + let source = self.arena.get(port).source; + assert!( + source == SourceEnd::Pipeline(id) || self.arena.get(port).closed(), + "port {port} is an inlet of pipeline {id} but its source is {source:?}" + ); + } + for &port in &pipe.outlets { + let queue = self.arena.get(port); + assert!( + queue.sink == Some(id) || queue.closed(), + "port {port} is an outlet of pipeline {id} but its sink is {:?}", + queue.sink + ); + // A pipeline blocked on one full outlet starves its other readers. If two of + // them meet again downstream, the join waits on the starved one while the fork + // waits on the join: a diamond deadlock. Unbounded fan-out never blocks. + assert!( + pipe.outlets.len() == 1 || queue.unbounded(), + "pipeline {id} fans out into bounded port {port}, which can deadlock a \ + diamond" + ); + } + } + for (port, queue) in self.arena.iter() { + assert!( + queue.is_free() || queue.source != SourceEnd::Unclaimed, + "port {port} has no reader, so its sink {:?} pushes into it for nothing", + queue.sink + ); + } + } + + /// Checks what one run of pipeline `id` reported: a source waits only on a condition some + /// other event will clear. + pub(super) fn validate_progress(&self, id: PipelineId, progress: &Progress) { + let pipe = self.pipeline(id); + match progress { + Progress::Blocked(Blocked::Inlet(index)) => { + assert!( + *index < pipe.inlets.len(), + "pipeline {id} waits on inlet {index} of {}", + pipe.inlets.len() + ); + let port = pipe.inlets[*index]; + let queue = self.arena.get(port); + // Only the sink pushing or closing wakes the source, and it can do neither + // into a closed or full port. + assert!( + !queue.closed(), + "pipeline {id} waits on inlet {index} (port {port}), which is closed" + ); + assert!( + queue.has_room(), + "pipeline {id} waits on inlet {index} (port {port}), which is full, so its \ + sink {:?} can never push to wake it", + queue.sink + ); + } + Progress::Blocked(Blocked::Io) => assert!( + pipe.pending > 0, + "pipeline {id} waits for a read with none in flight" + ), + Progress::Done => { + for (index, &port) in pipe.inlets.iter().enumerate() { + assert!( + self.arena.get(port).closed(), + "pipeline {id} finished before inlet {index} (port {port}) closed, so \ + its sink {:?} would push into a port nobody reads", + self.arena.get(port).sink + ); + } + } + Progress::Read(_) | Progress::Outlet => {} + } + } +} diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs index 81de957b126..2bf02010d24 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -131,7 +131,7 @@ fn masks_are_applied_by_the_source(#[values(false, true)] dense: bool) -> Vortex assert_rows(expected, run.arrays) } -/// A reader runs only when an inlet has changed: every run of the probe after its first finds +/// A source runs only when an inlet has changed: every run of the probe after its first finds /// a batch it has not seen, or an inlet newly closed. #[rstest] fn a_reader_runs_only_when_an_inlet_changed( @@ -155,8 +155,8 @@ fn a_reader_runs_only_when_an_inlet_changed( Ok(()) } -/// A writer is blocked once its reader's inlet holds as many batches as the reader allows, so -/// a fast source never runs far ahead of a slow reader. +/// A sink is blocked once its source's inlet holds as many batches as the source allows, so a +/// fast pipeline never runs far ahead of a slow source. #[test] fn a_full_inlet_blocks_its_writer() -> VortexResult<()> { let runs: Arc>> = Arc::default(); @@ -204,11 +204,15 @@ struct ManyReads { left: usize, } -impl Operator for ManyReads { - fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { +impl Source for ManyReads { + fn request(&mut self) -> Option { + self.wanted.pop() + } + + fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { if let Some((segment, _bytes)) = cx.take_segment() { self.left -= 1; - return Ok(Step::More(indices([u64::from(*segment)]))); + return Ok(Step::Batch(indices([u64::from(*segment)]))); } Ok(if self.left == 0 { Step::Finished @@ -218,12 +222,6 @@ impl Operator for ManyReads { } } -impl Source for ManyReads { - fn request(&mut self) -> Option { - self.wanted.pop() - } -} - /// A source has all its reads in flight at once, and takes each segment's bytes as they arrive, /// tagged with the segment. #[test] @@ -263,3 +261,117 @@ fn a_source_has_several_reads_in_flight() -> VortexResult<()> { } assert_rows(indices([3, 2, 1]), arrays) } + +/// Asks for every segment at once, emits the id of the first to arrive, and finishes without +/// waiting for the rest. +struct FirstOf { + wanted: Vec, +} + +impl Source for FirstOf { + fn request(&mut self) -> Option { + self.wanted.pop() + } + + fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { + Ok(match cx.take_segment() { + Some((segment, _bytes)) => { + self.wanted.clear(); + Step::Batch(indices([u64::from(*segment)])) + } + None if self.wanted.is_empty() => Step::Finished, + None => Step::Blocked(Blocked::Io), + }) + } +} + +/// A source that finishes with reads in flight has them cancelled: their bytes reach neither it +/// nor the next split's pipeline, which reuses its id, and the scan is done without them. +#[test] +fn finishing_cancels_reads_in_flight() -> VortexResult<()> { + let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |rows, _, _| { + let base = 10 * (u32::try_from(rows.start)? + 1); + Ok(Some(Chain::new(FirstOf { + wanted: (base..base + 3).map(SegmentId::from).collect(), + }))) + }) + .into_plan(); + let splits = vec![Split::all(0..1), Split::all(1..2)]; + let mut scan = Scan::try_new(SESSION.clone(), plan, splits)?; + let mut by_split: Vec> = vec![Vec::new(); 2]; + let mut inflight = Vec::new(); + loop { + match scan.step()? { + Turn::Read(read) => inflight.push(read), + Turn::Output(split, array) => by_split[split].push(array), + Turn::Waiting => { + let read = inflight.remove(0); + scan.deliver(read.id, empty_segment())?; + } + Turn::Done => break, + } + } + assert!( + !inflight.is_empty(), + "the second split's other reads are never delivered" + ); + assert_eq!( + scan.awaited_reads(), + 0, + "no pipeline waits for the reads left over" + ); + assert_rows(indices([12]), std::mem::take(&mut by_split[0]))?; + assert_rows(indices([22]), std::mem::take(&mut by_split[1]))?; + assert_eq!(scan.live_ports(), 0); + Ok(()) +} + +/// Emits `left` one-row batches. +struct Repeat { + left: usize, +} + +impl Source for Repeat { + fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(match self.left.checked_sub(1) { + Some(left) => { + self.left = left; + Step::Batch(indices([0])) + } + None => Step::Finished, + }) + } +} + +/// Waits on its one-batch inlet whatever it holds, so once the inlet is full nothing wakes it. +struct WaitsWhenFull; + +impl Source for WaitsWhenFull { + fn inlet_count(&self) -> usize { + 1 + } + + fn capacity(&self, _inlet: usize) -> usize { + 1 + } + + fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(Step::Blocked(Blocked::Inlet(0))) + } +} + +/// Debug builds catch a source waiting on an inlet its sink can no longer wake it through, at +/// that source, instead of reporting the scan stuck. +#[cfg(debug_assertions)] +#[test] +#[should_panic(expected = "which is full")] +fn waiting_on_a_full_inlet_is_caught() { + let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |_, _, compiler| { + Ok(Some( + compiler.join(vec![Chain::new(Repeat { left: 2 })], WaitsWhenFull), + )) + }) + .into_plan(); + let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..2)]).unwrap(); + while !matches!(scan.step().unwrap(), Turn::Done) {} +} diff --git a/vortex-layout/src/plan/pipeline/synthetic.rs b/vortex-layout/src/plan/pipeline/synthetic.rs index c8f3be38f62..f59f973c005 100644 --- a/vortex-layout/src/plan/pipeline/synthetic.rs +++ b/vortex-layout/src/plan/pipeline/synthetic.rs @@ -29,8 +29,6 @@ use super::Blocked; use super::Chain; use super::Compiler; use super::Cx; -use super::Input; -use super::Operator; use super::Source; use super::Step; use crate::plan::Plan; @@ -125,7 +123,7 @@ pub fn row_dtype() -> DType { } /// A source that emits its rows' indices, in order, cut into pieces of the given lengths, one -/// piece per call, optionally after a round trip through the IO service. +/// piece per call, optionally after a read of a segment it does not look at. /// /// A selected source emits only the rows its mask selects, as every plan does. A dense source /// ignores its mask and emits every row of its range, as a bare segment scan does. Piece @@ -135,8 +133,10 @@ pub struct RowSource { remaining: Vec, /// Lengths of the pieces still to emit, last first. pieces: Vec, - /// A segment to read before emitting anything, standing in for a real read. + /// A segment to read before emitting anything, standing in for a real read, until it has + /// been asked for. segment: Option, + /// The segment has been asked for and its bytes not yet taken. waiting: bool, } @@ -172,11 +172,17 @@ impl RowSource { } } -impl Operator for RowSource { - fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { +impl Source for RowSource { + fn request(&mut self) -> Option { + let segment = self.segment.take()?; + self.waiting = true; + Some(segment) + } + + fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { if self.waiting { self.waiting = false; - cx.take_bytes() + cx.take_segment() .ok_or_else(|| vortex_err!("RowSource ran without its bytes"))?; } if self.remaining.is_empty() { @@ -188,20 +194,12 @@ impl Operator for RowSource { .unwrap_or(self.remaining.len()) .min(self.remaining.len()); let piece: Vec = self.remaining.drain(..len).collect(); - Ok(Step::More(Buffer::from(piece).into_array())) - } -} - -impl Source for RowSource { - fn request(&mut self) -> Option { - let segment = self.segment.take()?; - self.waiting = true; - Some(segment) + Ok(Step::Batch(Buffer::from(piece).into_array())) } } /// What a [`Probe`] saw of its inlets on one call: for each, the batches queued and whether -/// its writer had closed it. +/// its sink had closed it. pub type ProbeRun = Vec<(usize, bool)>; /// A source over one inlet per child that records what it sees each time it runs, then passes @@ -245,8 +243,16 @@ impl Probe { } } -impl Operator for Probe { - fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { +impl Source for Probe { + fn inlet_count(&self) -> usize { + self.inlets + } + + fn capacity(&self, _inlet: usize) -> usize { + self.capacity + } + + fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { let run: ProbeRun = (0..self.inlets) .map(|index| { let inlet = cx.inlet(index); @@ -260,7 +266,7 @@ impl Operator for Probe { let mut first = cx.inlet(0); let batch: Option = first.take(); match batch { - Some(batch) => Ok(Step::Last(batch)), + Some(batch) => Ok(Step::Batch(batch)), None if first.closed() => { for index in 1..self.inlets { if !cx.inlet(index).closed() { @@ -274,16 +280,6 @@ impl Operator for Probe { } } -impl Source for Probe { - fn inlet_count(&self) -> usize { - self.inlets - } - - fn capacity(&self, _inlet: usize) -> usize { - self.capacity - } -} - /// An empty segment, for sources whose read is only a scheduling event. pub fn empty_segment() -> BufferHandle { BufferHandle::new_host(Buffer::::empty().into_byte_buffer()) From 9fdb3966e4bf658ed9eec0c26c83bc94b15cfc17 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 12:59:20 +0000 Subject: [PATCH 12/20] fix(layout): compile what only debug validation uses in debug builds only Queue::unbounded and Queue::is_free, and the test sources that break the scheduling contract, are used only under debug_assertions, so release builds (cargo check --release in CI) rejected them as dead code. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/port.rs | 2 + .../src/plan/pipeline/scheduling_tests.rs | 85 ++++++++++--------- 2 files changed, 47 insertions(+), 40 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/pipeline/port.rs index 825a1f94e75..304d431de00 100644 --- a/vortex-layout/src/plan/pipeline/port.rs +++ b/vortex-layout/src/plan/pipeline/port.rs @@ -74,6 +74,7 @@ impl Queue { } /// Whether the sink is never blocked by it. + #[cfg(debug_assertions)] pub(crate) fn unbounded(&self) -> bool { self.capacity == usize::MAX } @@ -82,6 +83,7 @@ impl Queue { self.batches.len() } + #[cfg(debug_assertions)] pub(crate) fn is_free(&self) -> bool { self.free } diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs index 2bf02010d24..af130c6249f 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -326,52 +326,57 @@ fn finishing_cancels_reads_in_flight() -> VortexResult<()> { Ok(()) } -/// Emits `left` one-row batches. -struct Repeat { - left: usize, -} +/// Sources that break the scheduling contract, which debug builds catch. +#[cfg(debug_assertions)] +mod contract { + use super::*; -impl Source for Repeat { - fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { - Ok(match self.left.checked_sub(1) { - Some(left) => { - self.left = left; - Step::Batch(indices([0])) - } - None => Step::Finished, - }) + /// Emits `left` one-row batches. + struct Repeat { + left: usize, } -} -/// Waits on its one-batch inlet whatever it holds, so once the inlet is full nothing wakes it. -struct WaitsWhenFull; - -impl Source for WaitsWhenFull { - fn inlet_count(&self) -> usize { - 1 + impl Source for Repeat { + fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(match self.left.checked_sub(1) { + Some(left) => { + self.left = left; + Step::Batch(indices([0])) + } + None => Step::Finished, + }) + } } - fn capacity(&self, _inlet: usize) -> usize { - 1 - } + /// Waits on its one-batch inlet whatever it holds, so once the inlet is full nothing wakes it. + struct WaitsWhenFull; + + impl Source for WaitsWhenFull { + fn inlet_count(&self) -> usize { + 1 + } + + fn capacity(&self, _inlet: usize) -> usize { + 1 + } - fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { - Ok(Step::Blocked(Blocked::Inlet(0))) + fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(Step::Blocked(Blocked::Inlet(0))) + } } -} -/// Debug builds catch a source waiting on an inlet its sink can no longer wake it through, at -/// that source, instead of reporting the scan stuck. -#[cfg(debug_assertions)] -#[test] -#[should_panic(expected = "which is full")] -fn waiting_on_a_full_inlet_is_caught() { - let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |_, _, compiler| { - Ok(Some( - compiler.join(vec![Chain::new(Repeat { left: 2 })], WaitsWhenFull), - )) - }) - .into_plan(); - let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..2)]).unwrap(); - while !matches!(scan.step().unwrap(), Turn::Done) {} + /// Debug builds catch a source waiting on an inlet its sink can no longer wake it through, at + /// that source, instead of reporting the scan stuck. + #[test] + #[should_panic(expected = "which is full")] + fn waiting_on_a_full_inlet_is_caught() { + let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |_, _, compiler| { + Ok(Some( + compiler.join(vec![Chain::new(Repeat { left: 2 })], WaitsWhenFull), + )) + }) + .into_plan(); + let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..2)]).unwrap(); + while !matches!(scan.step().unwrap(), Turn::Done) {} + } } From b80abc43ac79dd9751af01fbd292a6932be0f7bb Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 13:13:42 +0000 Subject: [PATCH 13/20] refactor(layout): name a pipeline's stages Operator, computed over an Input Stage::push and Stage::finish become one Operator::compute, handed Input::Batch for each batch and Input::End once, last. The driver still hands End one operator at a time behind the outlets' room check, and an operator still never blocks. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/compile.rs | 6 ++--- vortex-layout/src/plan/pipeline/mod.rs | 29 ++++++++++++---------- vortex-layout/src/plan/pipeline/scan.rs | 11 ++++---- 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/compile.rs b/vortex-layout/src/plan/pipeline/compile.rs index fa7cd33acb4..4a9fc550bc4 100644 --- a/vortex-layout/src/plan/pipeline/compile.rs +++ b/vortex-layout/src/plan/pipeline/compile.rs @@ -10,8 +10,8 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; +use super::Operator; use super::Source; -use super::Stage; use super::port::PortId; use super::port::SourceEnd; use super::scan::Core; @@ -21,7 +21,7 @@ use crate::plan::PlanRef; /// parent plan can add stages before it becomes a pipeline. pub struct Chain { pub(crate) source: Box, - pub(crate) stages: Vec>, + pub(crate) stages: Vec>, pub(crate) inlets: Vec, } @@ -36,7 +36,7 @@ impl Chain { } /// This chain with `stage` after its last stage. - pub fn with(mut self, stage: impl Stage + 'static) -> Self { + pub fn with(mut self, stage: impl Operator + 'static) -> Self { self.stages.push(Box::new(stage)); self } diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 5acb7bf878e..19b472e1c1c 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -5,7 +5,7 @@ //! //! A [`Scan`] runs a [`PlanRef`](crate::plan::PlanRef) over a list of splits, row ranges of the //! plan's domain. Each split is compiled into pipelines: a pipeline is one [`Source`], zero or -//! more [`Stage`]s, and the ports its last stage writes. The driver asks the source for a batch +//! more [`Operator`]s, and the ports its last stage writes. The driver asks the source for a batch //! and pushes it through the stages into the ports. Every port runs from one pipeline's sink, //! the end that pushes into its outlets, to another's source, which reads it as an inlet, and //! that source declares how many batches it may hold. @@ -96,18 +96,21 @@ pub trait Source: Send { fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult; } -/// A stage of a pipeline: transforms the batches pushed into it, one at a time, and never -/// blocks. -pub trait Stage: Send { - /// Takes `batch` from the stage below and returns what to push on, if anything. - fn push(&mut self, batch: ArrayRef, cx: &mut Cx<'_>) -> VortexResult>; - - /// The stage below has finished. Returns whatever the stage still holds, if anything; it is - /// not called again after. - fn finish(&mut self, cx: &mut Cx<'_>) -> VortexResult> { - let _ = cx; - Ok(None) - } +/// What an operator is handed by the stage below it. +#[derive(Debug)] +pub enum Input { + /// The next batch. + Batch(ArrayRef), + /// The stage below has finished: nothing more comes. Handed over once, last. + End, +} + +/// A stage of a pipeline after its source: transforms what the stage below hands it, one input +/// at a time, and never blocks. +pub trait Operator: Send { + /// Computes what to push on from `input`, if anything. At [`Input::End`] it returns whatever + /// it still holds; it is not called again after. + fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult>; } /// What a source or stage may touch during one call: its pipeline's inlets, the bytes the diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index 07abf20dd90..4b44483e91c 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -24,8 +24,9 @@ use vortex_session::VortexSession; use super::Blocked; use super::Cx; use super::DEFAULT_CAPACITY; +use super::Input; +use super::Operator; use super::Source; -use super::Stage; use super::Step; use super::compile::Chain; use super::port::Arena; @@ -98,7 +99,7 @@ enum State { /// A source, its stages, and the ports they read and write. pub(crate) struct Pipeline { source: Box, - stages: Vec>, + stages: Vec>, inlets: Vec, outlets: SmallVec<[PortId; 1]>, /// The segments it asked for that have arrived, in order of arrival. @@ -172,7 +173,7 @@ fn drive( .split_first_mut() .unwrap_or_else(|| unreachable!("stage {stage} finished twice")); *finishing = Some(stage + 1); - if let Some(batch) = this.finish(&mut cx)? { + if let Some(batch) = this.compute(Input::End, &mut cx)? { push(rest, batch, outlets, &mut cx)?; } continue; @@ -198,13 +199,13 @@ fn drive( /// Pushes `batch` through `stages` into `outlets`. fn push( - stages: &mut [Box], + stages: &mut [Box], mut batch: ArrayRef, outlets: &[PortId], cx: &mut Cx<'_>, ) -> VortexResult<()> { for stage in stages { - batch = match stage.push(batch, cx)? { + batch = match stage.compute(Input::Batch(batch), cx)? { Some(batch) => batch, None => return Ok(()), }; From 9343ec83297f42e641ddcc62411ce40de43bcf14 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 13:23:34 +0000 Subject: [PATCH 14/20] refactor(layout): a source is one ready(), asking for reads in its Step Source is now one method, ready(), returning Step::{Batch, Read, Blocked, Finished}. A source asks for a segment by returning Step::Read; the driver issues the read once the run ends and asks the source again at once, so it can still have several reads in flight before it answers Blocked::Io. Bytes already there, as of a segment read ahead, run the source again at once. The inlet count comes from the chain, and capacities are the joining plan's: Compiler::join holds DEFAULT_CAPACITY batches per inlet, and Compiler::join_with takes a capacity per inlet. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- vortex-layout/src/plan/pipeline/compile.rs | 23 ++++++++-- vortex-layout/src/plan/pipeline/mod.rs | 33 ++++--------- vortex-layout/src/plan/pipeline/scan.rs | 39 ++++++++-------- .../src/plan/pipeline/scan/validate.rs | 13 ++---- .../src/plan/pipeline/scheduling_tests.rs | 46 +++++++++---------- vortex-layout/src/plan/pipeline/synthetic.rs | 32 ++++--------- 6 files changed, 81 insertions(+), 105 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/compile.rs b/vortex-layout/src/plan/pipeline/compile.rs index 4a9fc550bc4..ed39688a109 100644 --- a/vortex-layout/src/plan/pipeline/compile.rs +++ b/vortex-layout/src/plan/pipeline/compile.rs @@ -10,6 +10,7 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; +use super::DEFAULT_CAPACITY; use super::Operator; use super::Source; use super::port::PortId; @@ -76,16 +77,28 @@ impl Compiler<'_> { } /// A chain reading `chains` through `source`: each chain becomes a pipeline writing a new - /// port, inlet `i` of `source` reading chain `i`, with the capacity `source` asks for. + /// port, inlet `i` of `source` reading chain `i`, holding [`DEFAULT_CAPACITY`] batches. pub fn join(&mut self, chains: Vec, source: impl Source + 'static) -> Chain { + self.join_with(chains, |_| DEFAULT_CAPACITY, source) + } + + /// [`join`](Self::join), with inlet `i` holding `capacity(i)` batches before its sink is + /// blocked: `usize::MAX` for an inlet the source needs whole. An empty inlet always has + /// room. + pub fn join_with( + &mut self, + chains: Vec, + capacity: impl Fn(usize) -> usize, + source: impl Source + 'static, + ) -> Chain { let inlets = chains .into_iter() .enumerate() .map(|(index, chain)| { - let port = - self.core - .arena - .create(source.capacity(index), None, SourceEnd::Unclaimed); + let port = self + .core + .arena + .create(capacity(index), None, SourceEnd::Unclaimed); self.core.add_pipeline(chain, smallvec::smallvec![port]); port }) diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 19b472e1c1c..2b39bfbe811 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -43,7 +43,8 @@ pub use self::scan::Split; pub use self::scan::Turn; use crate::segments::SegmentId; -/// Batches an inlet holds before its sink is blocked, unless its source says otherwise. +/// Batches an inlet holds before its sink is blocked, unless the plan joining it says +/// otherwise; see [`Compiler::join_with`]. pub const DEFAULT_CAPACITY: usize = 8; /// What a source says back when asked for its next batch. @@ -51,6 +52,10 @@ pub const DEFAULT_CAPACITY: usize = 8; pub enum Step { /// A batch for the stage above. The source is asked again at once. Batch(ArrayRef), + /// A segment to read. The source is asked again at once, so it may have several reads in + /// flight; the bytes come through [`Cx::take_segment`] in the order they arrive, and while + /// none has, the source answers [`Blocked::Io`]. + Read(SegmentId), /// No batch until the condition clears. The source is not asked again before it does. Blocked(Blocked), /// No batch will ever come again. The source is not asked again. @@ -71,29 +76,9 @@ pub enum Blocked { /// The start of a pipeline: produces batches from its inlets and the segments it asks to have /// read. pub trait Source: Send { - /// The inlets this source reads, numbered from zero. - fn inlet_count(&self) -> usize { - 0 - } - - /// How many batches inlet `inlet` may hold before its sink is blocked. Asked once, when - /// the port is made. An empty port always has room. - fn capacity(&self, inlet: usize) -> usize { - let _ = inlet; - DEFAULT_CAPACITY - } - - /// A segment this source wants read. Asked until it answers `None` before every - /// [`next`](Self::next), so a source may have several reads in flight. A source that asked - /// for any is not asked for a batch again until one of them arrives; the bytes come through - /// [`Cx::take_segment`] in the order they arrive, and a source still waiting for a read it - /// asked for answers [`Blocked::Io`]. - fn request(&mut self) -> Option { - None - } - - /// The next batch, or why there is none. - fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult; + /// What the source has ready: a batch, a segment to read, why it has neither, or that it + /// has finished. + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult; } /// What an operator is handed by the stage below it. diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index 4b44483e91c..f8f3e3e24a1 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -104,6 +104,8 @@ pub(crate) struct Pipeline { outlets: SmallVec<[PortId; 1]>, /// The segments it asked for that have arrived, in order of arrival. bytes: VecDeque<(SegmentId, BufferHandle)>, + /// Segments its source asked for during the last run, issued after it. + asked: SmallVec<[SegmentId; 1]>, /// Reads it asked for that have not arrived, so retiring it cancels them only when any /// are left. pending: usize, @@ -120,8 +122,6 @@ pub(crate) struct Pipeline { /// What one run of a pipeline came to. enum Progress { - /// The source wants these segments read. - Read(SmallVec<[SegmentId; 1]>), /// The source waits for this. Blocked(Blocked), /// An outlet is full. @@ -130,7 +130,7 @@ enum Progress { Done, } -/// Runs `pipe` until it blocks, finishes, or requests a read. +/// Runs `pipe` until it blocks or finishes, collecting the reads its source asks for. fn drive( pipe: &mut Pipeline, arena: &mut Arena, @@ -143,6 +143,7 @@ fn drive( inlets, outlets, bytes, + asked, finishing, touched, listed, @@ -178,19 +179,12 @@ fn drive( } continue; } - if let Some(first) = source.request() { - let mut segments = SmallVec::new(); - segments.push(first); - while let Some(segment) = source.request() { - segments.push(segment); - } - return Ok(Progress::Read(segments)); - } - match source.next(&mut cx)? { + match source.ready(&mut cx)? { // A source's input is its inlets and reads, which nothing changes while it runs, so // it is asked again at once: it blocks or finishes without a trip through the // scheduler, and the outlets' room still bounds what it produces. Step::Batch(batch) => push(stages, batch, outlets, &mut cx)?, + Step::Read(segment) => asked.push(segment), Step::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), Step::Finished => *finishing = Some(0), } @@ -280,6 +274,7 @@ impl Core { inlets: chain.inlets, outlets, bytes: VecDeque::new(), + asked: SmallVec::new(), pending: 0, finishing: None, touched, @@ -353,16 +348,16 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; + for segment_id in std::mem::take(&mut self.pipeline_mut(id).asked) { + self.ask(id, segment_id); + } #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { - Progress::Read(segments) => { - for &segment_id in segments { - self.read(id, segment_id); - } - let pipe = self.pipeline_mut(id); - pipe.pending += segments.len(); - pipe.state = State::Blocked(Blocked::Io); + // Bytes already here, as of a segment read ahead, are taken at once: the source runs + // again. + Progress::Blocked(Blocked::Io) if !self.pipeline(id).bytes.is_empty() => { + self.make_runnable(id); } // Nothing changes a port while its source's pipeline runs, so the condition the run reported // holds until a neighbour wakes it. @@ -377,6 +372,12 @@ impl Core { Ok(()) } + /// Issues a read of `segment_id`, which pipeline `id`'s source asked for. + fn ask(&mut self, id: PipelineId, segment_id: SegmentId) { + self.read(id, segment_id); + self.pipeline_mut(id).pending += 1; + } + /// Wakes the sources of what pipeline `id` pushed and the sinks of what it took. fn wake_neighbours(&mut self, id: PipelineId) { for index in 0..self.pipeline(id).outlets.len() { diff --git a/vortex-layout/src/plan/pipeline/scan/validate.rs b/vortex-layout/src/plan/pipeline/scan/validate.rs index 133769990e5..8d8fb2b16c6 100644 --- a/vortex-layout/src/plan/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/pipeline/scan/validate.rs @@ -14,18 +14,11 @@ use super::Progress; use super::SourceEnd; impl Core { - /// Checks the graph after a compile: ports and pipelines agree on their ends, sources read - /// as many inlets as they declare, every port has a reader, and fan-out is unbounded. + /// Checks the graph after a compile: ports and pipelines agree on their ends, every port has + /// a reader, and fan-out is unbounded. pub(super) fn validate(&self) { for (id, pipe) in self.pipelines.iter().enumerate() { let Some(pipe) = pipe else { continue }; - assert_eq!( - pipe.inlets.len(), - pipe.source.inlet_count(), - "pipeline {id} reads {} inlets but its source declares {}", - pipe.inlets.len(), - pipe.source.inlet_count() - ); for &port in &pipe.inlets { let source = self.arena.get(port).source; assert!( @@ -99,7 +92,7 @@ impl Core { ); } } - Progress::Read(_) | Progress::Outlet => {} + Progress::Outlet => {} } } } diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs index af130c6249f..c5451a74f7f 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -205,11 +205,10 @@ struct ManyReads { } impl Source for ManyReads { - fn request(&mut self) -> Option { - self.wanted.pop() - } - - fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if let Some(segment) = self.wanted.pop() { + return Ok(Step::Read(segment)); + } if let Some((segment, _bytes)) = cx.take_segment() { self.left -= 1; return Ok(Step::Batch(indices([u64::from(*segment)]))); @@ -266,20 +265,22 @@ fn a_source_has_several_reads_in_flight() -> VortexResult<()> { /// waiting for the rest. struct FirstOf { wanted: Vec, + done: bool, } impl Source for FirstOf { - fn request(&mut self) -> Option { - self.wanted.pop() - } - - fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if self.done { + return Ok(Step::Finished); + } + if let Some(segment) = self.wanted.pop() { + return Ok(Step::Read(segment)); + } Ok(match cx.take_segment() { Some((segment, _bytes)) => { - self.wanted.clear(); + self.done = true; Step::Batch(indices([u64::from(*segment)])) } - None if self.wanted.is_empty() => Step::Finished, None => Step::Blocked(Blocked::Io), }) } @@ -293,6 +294,7 @@ fn finishing_cancels_reads_in_flight() -> VortexResult<()> { let base = 10 * (u32::try_from(rows.start)? + 1); Ok(Some(Chain::new(FirstOf { wanted: (base..base + 3).map(SegmentId::from).collect(), + done: false, }))) }) .into_plan(); @@ -337,7 +339,7 @@ mod contract { } impl Source for Repeat { - fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { Ok(match self.left.checked_sub(1) { Some(left) => { self.left = left; @@ -352,15 +354,7 @@ mod contract { struct WaitsWhenFull; impl Source for WaitsWhenFull { - fn inlet_count(&self) -> usize { - 1 - } - - fn capacity(&self, _inlet: usize) -> usize { - 1 - } - - fn next(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { Ok(Step::Blocked(Blocked::Inlet(0))) } } @@ -371,9 +365,11 @@ mod contract { #[should_panic(expected = "which is full")] fn waiting_on_a_full_inlet_is_caught() { let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |_, _, compiler| { - Ok(Some( - compiler.join(vec![Chain::new(Repeat { left: 2 })], WaitsWhenFull), - )) + Ok(Some(compiler.join_with( + vec![Chain::new(Repeat { left: 2 })], + |_| 1, + WaitsWhenFull, + ))) }) .into_plan(); let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..2)]).unwrap(); diff --git a/vortex-layout/src/plan/pipeline/synthetic.rs b/vortex-layout/src/plan/pipeline/synthetic.rs index f59f973c005..ab0f650a062 100644 --- a/vortex-layout/src/plan/pipeline/synthetic.rs +++ b/vortex-layout/src/plan/pipeline/synthetic.rs @@ -21,7 +21,6 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_buffer::Buffer; use vortex_error::VortexResult; -use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::CachedId; @@ -173,17 +172,16 @@ impl RowSource { } impl Source for RowSource { - fn request(&mut self) -> Option { - let segment = self.segment.take()?; - self.waiting = true; - Some(segment) - } - - fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if let Some(segment) = self.segment.take() { + self.waiting = true; + return Ok(Step::Read(segment)); + } if self.waiting { + if cx.take_segment().is_none() { + return Ok(Step::Blocked(Blocked::Io)); + } self.waiting = false; - cx.take_segment() - .ok_or_else(|| vortex_err!("RowSource ran without its bytes"))?; } if self.remaining.is_empty() { return Ok(Step::Finished); @@ -206,7 +204,6 @@ pub type ProbeRun = Vec<(usize, bool)>; /// its first inlet through, one batch per call, and drops what the others hold. pub struct Probe { inlets: usize, - capacity: usize, runs: Arc>>, } @@ -233,10 +230,9 @@ impl Probe { } let probe = Probe { inlets: chains.len(), - capacity, runs: Arc::clone(&runs), }; - Ok(Some(compiler.join(chains, probe))) + Ok(Some(compiler.join_with(chains, |_| capacity, probe))) }, ) .into_plan() @@ -244,15 +240,7 @@ impl Probe { } impl Source for Probe { - fn inlet_count(&self) -> usize { - self.inlets - } - - fn capacity(&self, _inlet: usize) -> usize { - self.capacity - } - - fn next(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { let run: ProbeRun = (0..self.inlets) .map(|index| { let inlet = cx.inlet(index); From b126427115e8572f98a35ee7eff05c6e2f194517 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 12 Oct 2026 15:16:27 +0100 Subject: [PATCH 15/20] refactor(layout): separate source readiness and drain operator continuations Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/pipeline/mod.rs | 175 ++++++--- vortex-layout/src/plan/pipeline/port.rs | 13 +- vortex-layout/src/plan/pipeline/scan.rs | 170 +++++---- .../src/plan/pipeline/scan/validate.rs | 13 +- .../src/plan/pipeline/scheduling_tests.rs | 332 ++++++++++++++++-- vortex-layout/src/plan/pipeline/synthetic.rs | 67 ++-- 6 files changed, 596 insertions(+), 174 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/pipeline/mod.rs index 2b39bfbe811..3d3c51e6c00 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/pipeline/mod.rs @@ -8,11 +8,11 @@ //! more [`Operator`]s, and the ports its last stage writes. The driver asks the source for a batch //! and pushes it through the stages into the ports. Every port runs from one pipeline's sink, //! the end that pushes into its outlets, to another's source, which reads it as an inlet, and -//! that source declares how many batches it may hold. +//! the plan joining them declares how many batches each inlet may hold. //! -//! Only a source blocks, and only on a segment read it requested or an inlet it waits on, which -//! it reports as a [`Blocked`]; a pipeline whose outlet is full is held back by the driver before -//! its source runs. The scheduler re-runs a pipeline when the one condition it waits on clears. +//! A source checks its readiness before computing, waiting for a segment read or an inlet. +//! The driver also blocks a pipeline when an outlet is full. The scheduler re-runs a blocked +//! pipeline when that one condition clears. //! //! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and //! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. @@ -24,6 +24,7 @@ pub mod synthetic; use std::collections::VecDeque; +use smallvec::SmallVec; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::buffer::BufferHandle; @@ -47,55 +48,134 @@ use crate::segments::SegmentId; /// otherwise; see [`Compiler::join_with`]. pub const DEFAULT_CAPACITY: usize = 8; -/// What a source says back when asked for its next batch. +/// The output of one non-blocking compute call. #[derive(Debug)] pub enum Step { - /// A batch for the stage above. The source is asked again at once. - Batch(ArrayRef), - /// A segment to read. The source is asked again at once, so it may have several reads in - /// flight; the bytes come through [`Cx::take_segment`] in the order they arrive, and while - /// none has, the source answers [`Blocked::Io`]. - Read(SegmentId), - /// No batch until the condition clears. The source is not asked again before it does. - Blocked(Blocked), - /// No batch will ever come again. The source is not asked again. - Finished, + /// No output remains for this input. A source has finished entirely. + End, + /// A batch with more computation available without another input or readiness check. + More(ArrayRef), + /// The last batch for this input. A stage needs another input; a source needs readiness + /// checked again. + Final(ArrayRef), +} + +/// Whether a source can compute or needs the scheduler to wait. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SourceReady { + /// The next compute call can produce output or finish without blocking. + Ready, + /// New reads were queued through [`Cx::request_read`]. Run again when bytes are available. + Read, + /// Nothing can happen until the condition clears. + Blocked(SourceBlocked), } -/// What a source waits for. +/// Why a source cannot compute yet. #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Blocked { - /// A read it asked for has not arrived. +pub enum SourceBlocked { + /// A read was requested and has not been delivered. Io, - /// Inlet `inlet` is open and has room, so its sink can still push into it or close it, either - /// of which wakes the source. It may already hold batches, as when a source waits for it to - /// fill. Waiting on a closed or full inlet waits forever; debug builds catch it. + /// The inlet needs more batches or its sink to close it. It must be open and have room + /// for its sink to push; debug builds check that the sink can still wake this source. Inlet(usize), } /// The start of a pipeline: produces batches from its inlets and the segments it asks to have /// read. -pub trait Source: Send { - /// What the source has ready: a batch, a segment to read, why it has neither, or that it - /// has finished. - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult; +pub trait Source: Operator { + /// Prepares the next [`compute`](Operator::compute) call, after the driver checks outlets. + /// Returning [`SourceReady::Ready`] guarantees that call can proceed without blocking. + /// Any input or choice made here can be retained by the source for that call. + /// + /// Queue one or more reads with [`Cx::request_read`] and return [`SourceReady::Read`] to + /// issue them. Readiness is checked again when bytes are available, including bytes already + /// delivered; they are available through [`Cx::take_segment`] in arrival order. + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(SourceReady::Ready) + } } /// What an operator is handed by the stage below it. #[derive(Debug)] pub enum Input { + /// No new input: a source is ready, or an operator is continuing after [`Step::More`]. + None, /// The next batch. Batch(ArrayRef), - /// The stage below has finished: nothing more comes. Handed over once, last. + /// The stage below has finished: no more batches come. Handed over once; [`Input::None`] + /// may follow to drain any [`Step::More`] outputs. End, } -/// A stage of a pipeline after its source: transforms what the stage below hands it, one input -/// at a time, and never blocks. +/// Computes batches without blocking, either as a source or as a stage transforming input. pub trait Operator: Send { - /// Computes what to push on from `input`, if anything. At [`Input::End`] it returns whatever - /// it still holds; it is not called again after. - fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult>; + /// Computes one output. [`Step::More`] requests another call with [`Input::None`] before + /// any new input is supplied. [`Step::End`] and [`Step::Final`] finish the current input; + /// after [`Input::End`], they finish the stage as well. + /// + /// Sources receive only [`Input::None`]. Their first call and each call after + /// [`Step::Final`] follow [`Source::ready`]; [`Step::End`] finishes the source. + fn compute(&mut self, input: Input, cx: &mut Cx<'_>) -> VortexResult; +} + +/// Distinct inlet indices, with constant-time insertion and removal. +#[derive(Default)] +struct InletSet { + indices: Vec, + positions: Vec, +} + +impl InletSet { + fn new(inlets: usize) -> Self { + Self { + indices: Vec::with_capacity(inlets), + positions: vec![usize::MAX; inlets], + } + } + + fn all(inlets: usize) -> Self { + Self { + indices: (0..inlets).collect(), + positions: (0..inlets).collect(), + } + } + + fn insert(&mut self, index: usize) { + if self.positions[index] == usize::MAX { + self.positions[index] = self.indices.len(); + self.indices.push(index); + } + } + + fn remove(&mut self, index: usize) { + let position = self.positions[index]; + if position == usize::MAX { + return; + } + self.indices.swap_remove(position); + if let Some(&moved) = self.indices.get(position) { + self.positions[moved] = position; + } + self.positions[index] = usize::MAX; + } + + fn pop(&mut self) -> Option { + let index = *self.indices.last()?; + self.remove(index); + Some(index) + } + + fn indices(&self) -> &[usize] { + &self.indices + } + + fn clear(&mut self) { + for &index in &self.indices { + self.positions[index] = usize::MAX; + } + self.indices.clear(); + } } /// What a source or stage may touch during one call: its pipeline's inlets, the bytes the @@ -104,25 +184,40 @@ pub struct Cx<'a> { arena: &'a mut Arena, inlets: &'a [PortId], bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, + reads: SmallVec<[SegmentId; 1]>, session: &'a VortexSession, exec: &'a mut ExecutionCtx, - /// The inlets read during this run, each once, so only their sinks are checked for room - /// after it, and whether each inlet is listed. - touched: &'a mut Vec, - listed: &'a mut [bool], + touched: &'a mut InletSet, + changed: &'a mut InletSet, } impl Cx<'_> { + /// Takes an inlet index whose sink appended data or closed since the last readiness + /// check. Each index is reported once, and all inlets are reported on the first check. + /// Remaining notifications are cleared after [`Source::ready`] returns. A source must + /// account for its own consumption when caching inlet state. + pub fn changed_inlet(&mut self) -> Option { + self.changed.pop() + } + + /// Queues a segment read during [`Source::ready`]. Return [`SourceReady::Read`] after + /// queuing all the reads needed for this call. + pub fn request_read(&mut self, segment: SegmentId) { + self.reads.push(segment); + } + + /// Whether a requested segment has arrived and can be taken. + pub fn has_bytes(&self) -> bool { + !self.bytes.is_empty() + } + /// The reading side of inlet `index` of this pipeline's source. pub fn inlet(&mut self, index: usize) -> Inlet<'_> { - if !self.listed[index] { - self.listed[index] = true; - self.touched.push(index); - } + self.touched.insert(index); self.arena.inlet(self.inlets[index]) } - /// A segment the source asked for that has arrived, oldest arrival first, with its bytes. + /// Takes an arrived segment's ID and bytes, oldest arrival first. pub fn take_segment(&mut self) -> Option<(SegmentId, BufferHandle)> { self.bytes.pop_front() } diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/pipeline/port.rs index 304d431de00..38f9ae143f1 100644 --- a/vortex-layout/src/plan/pipeline/port.rs +++ b/vortex-layout/src/plan/pipeline/port.rs @@ -27,7 +27,7 @@ pub(crate) enum SourceEnd { /// source has gone. Unclaimed, /// The source of this pipeline, which reads the port as one of its inlets. - Pipeline(PipelineId), + Pipeline { id: PipelineId, inlet: usize }, /// The scan, reading the output of the split's stage running in this slot. Scan(usize), } @@ -42,6 +42,8 @@ pub(crate) struct Queue { capacity: usize, /// The sink will push nothing more. closed: bool, + /// The sink appended data or closed this port since its last run was settled. + changed: bool, /// The pipeline whose sink pushes into this port; none once closed. pub(crate) sink: Option, pub(crate) source: SourceEnd, @@ -58,6 +60,7 @@ impl Queue { batches: VecDeque::with_capacity(capacity.clamp(1, DEFAULT_CAPACITY)), capacity: capacity.max(1), closed: false, + changed: false, sink, source, source_gone: false, @@ -88,9 +91,8 @@ impl Queue { self.free } - /// Whether the source would find something new: a batch, or the close. - pub(crate) fn readable(&self) -> bool { - !self.batches.is_empty() || self.closed + pub(crate) fn take_changed(&mut self) -> bool { + std::mem::take(&mut self.changed) } pub(crate) fn pop(&mut self) -> Option { @@ -218,6 +220,7 @@ impl Arena { queue.batches.reserve_exact(capacity.min(DEFAULT_CAPACITY)); queue.capacity = capacity; queue.closed = false; + queue.changed = false; queue.sink = sink; queue.source = source; queue.source_gone = false; @@ -255,12 +258,14 @@ impl Arena { ); if !batch.is_empty() { queue.batches.push_back(batch); + queue.changed = true; } } /// The sink will push nothing more. pub(crate) fn close(&mut self, id: PortId) { let queue = &mut self.queues[id]; + queue.changed |= !queue.closed; queue.closed = true; queue.sink = None; self.maybe_free(id); diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/pipeline/scan.rs index f8f3e3e24a1..5651eac6430 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/pipeline/scan.rs @@ -21,12 +21,14 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use super::Blocked; use super::Cx; use super::DEFAULT_CAPACITY; +use super::InletSet; use super::Input; use super::Operator; use super::Source; +use super::SourceBlocked; +use super::SourceReady; use super::Step; use super::compile::Chain; use super::port::Arena; @@ -91,7 +93,7 @@ pub enum Turn { enum State { Runnable, /// What its source said it waits for. - Blocked(Blocked), + Blocked(SourceBlocked), /// Room in a full outlet, found by the driver before the source ran. Outlet, } @@ -104,33 +106,35 @@ pub(crate) struct Pipeline { outlets: SmallVec<[PortId; 1]>, /// The segments it asked for that have arrived, in order of arrival. bytes: VecDeque<(SegmentId, BufferHandle)>, - /// Segments its source asked for during the last run, issued after it. - asked: SmallVec<[SegmentId; 1]>, /// Reads it asked for that have not arrived, so retiring it cancels them only when any /// are left. pending: usize, /// Once its source has finished, the next stage to finish: stages are finished one at a /// time, each when the outlets have room for what it flushes. finishing: Option, - /// The inlets the last run read, each once. - touched: Vec, - /// For each inlet, whether it is in `touched`. - listed: Vec, + /// Stages with more output for their current input, resumed deepest first. + continuing: Vec, + /// The source can compute again without checking readiness. + source_more: bool, + touched: InletSet, + changed: InletSet, state: State, queued: bool, } /// What one run of a pipeline came to. enum Progress { + /// The source wants these segments read. + Read(SmallVec<[SegmentId; 1]>), /// The source waits for this. - Blocked(Blocked), + Blocked(SourceBlocked), /// An outlet is full. Outlet, /// The pipeline has finished and closed its outlets. Done, } -/// Runs `pipe` until it blocks or finishes, collecting the reads its source asks for. +/// Runs `pipe` until it blocks, finishes, or requests a read. fn drive( pipe: &mut Pipeline, arena: &mut Arena, @@ -143,24 +147,26 @@ fn drive( inlets, outlets, bytes, - asked, finishing, + continuing, + source_more, touched, - listed, + changed, .. } = pipe; let mut cx = Cx { arena, inlets, bytes, + reads: SmallVec::new(), session, exec, touched, - listed, + changed, }; loop { // Closing needs no room, so a pipeline with nothing left to flush closes at once. - if *finishing == Some(stages.len()) { + if *finishing == Some(stages.len()) && continuing.is_empty() { for &outlet in outlets.iter() { cx.arena.close(outlet); } @@ -169,39 +175,78 @@ fn drive( if !outlets.iter().all(|&outlet| cx.arena.has_room(outlet)) { return Ok(Progress::Outlet); } + if let Some(stage) = continuing.pop() { + let step = stages[stage].compute(Input::None, &mut cx)?; + push(stages, stage, step, continuing, outlets, &mut cx)?; + continue; + } if let Some(stage) = *finishing { - let (this, rest) = stages[stage..] - .split_first_mut() - .unwrap_or_else(|| unreachable!("stage {stage} finished twice")); *finishing = Some(stage + 1); - if let Some(batch) = this.compute(Input::End, &mut cx)? { - push(rest, batch, outlets, &mut cx)?; - } + let step = stages[stage].compute(Input::End, &mut cx)?; + push(stages, stage, step, continuing, outlets, &mut cx)?; continue; } - match source.ready(&mut cx)? { - // A source's input is its inlets and reads, which nothing changes while it runs, so - // it is asked again at once: it blocks or finishes without a trip through the - // scheduler, and the outlets' room still bounds what it produces. - Step::Batch(batch) => push(stages, batch, outlets, &mut cx)?, - Step::Read(segment) => asked.push(segment), - Step::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), - Step::Finished => *finishing = Some(0), + if !*source_more { + let ready = source.ready(&mut cx)?; + cx.changed.clear(); + vortex_ensure!( + cx.reads.is_empty() == (ready != SourceReady::Read), + "A source must return Read exactly when it queues reads" + ); + match ready { + SourceReady::Ready => {} + SourceReady::Read => return Ok(Progress::Read(cx.reads)), + SourceReady::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), + } + } + let step = source.compute(Input::None, &mut cx)?; + *source_more = matches!(&step, Step::More(_)); + match step { + Step::More(batch) | Step::Final(batch) => { + push_batch(stages, 0, batch, continuing, outlets, &mut cx)?; + } + Step::End => *finishing = Some(0), } } } -/// Pushes `batch` through `stages` into `outlets`. +/// Records a stage's continuation before pushing its output through the remaining stages. fn push( stages: &mut [Box], + stage: usize, + step: Step, + continuing: &mut Vec, + outlets: &[PortId], + cx: &mut Cx<'_>, +) -> VortexResult<()> { + let batch = match step { + Step::End => return Ok(()), + Step::More(batch) => { + continuing.push(stage); + batch + } + Step::Final(batch) => batch, + }; + push_batch(stages, stage + 1, batch, continuing, outlets, cx) +} + +/// Pushes one batch, retaining continuations so outlet capacity bounds every output. +fn push_batch( + stages: &mut [Box], + start: usize, mut batch: ArrayRef, + continuing: &mut Vec, outlets: &[PortId], cx: &mut Cx<'_>, ) -> VortexResult<()> { - for stage in stages { + for (index, stage) in stages.iter_mut().enumerate().skip(start) { batch = match stage.compute(Input::Batch(batch), cx)? { - Some(batch) => batch, - None => return Ok(()), + Step::More(batch) => { + continuing.push(index); + batch + } + Step::Final(batch) => batch, + Step::End => return Ok(()), }; } push_out(batch, outlets, cx.arena); @@ -266,19 +311,20 @@ impl Core { chain: Chain, outlets: SmallVec<[PortId; 1]>, ) -> PipelineId { - let listed = vec![false; chain.inlets.len()]; - let touched = Vec::with_capacity(chain.inlets.len()); + let touched = InletSet::new(chain.inlets.len()); + let changed = InletSet::all(chain.inlets.len()); let pipeline = Pipeline { + continuing: Vec::new(), + source_more: false, source: chain.source, stages: chain.stages, inlets: chain.inlets, outlets, bytes: VecDeque::new(), - asked: SmallVec::new(), pending: 0, finishing: None, touched, - listed, + changed, state: State::Runnable, queued: true, }; @@ -286,8 +332,8 @@ impl Core { Some(id) => id, None => self.pipelines.push(None), }; - for &inlet in &pipeline.inlets { - self.arena.get_mut(inlet).source = SourceEnd::Pipeline(id); + for (index, &inlet) in pipeline.inlets.iter().enumerate() { + self.arena.get_mut(inlet).source = SourceEnd::Pipeline { id, inlet: index }; } for &outlet in &pipeline.outlets { self.arena.get_mut(outlet).sink = Some(id); @@ -348,15 +394,21 @@ impl Core { .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; - for segment_id in std::mem::take(&mut self.pipeline_mut(id).asked) { - self.ask(id, segment_id); - } #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { - // Bytes already here, as of a segment read ahead, are taken at once: the source runs - // again. - Progress::Blocked(Blocked::Io) if !self.pipeline(id).bytes.is_empty() => { + Progress::Read(segments) => { + for &segment_id in segments { + self.ask(id, segment_id); + } + if self.pipeline(id).bytes.is_empty() { + self.pipeline_mut(id).state = State::Blocked(SourceBlocked::Io); + } else { + self.make_runnable(id); + } + } + // A read-ahead source can still have delivered bytes to consume. + Progress::Blocked(SourceBlocked::Io) if !self.pipeline(id).bytes.is_empty() => { self.make_runnable(id); } // Nothing changes a port while its source's pipeline runs, so the condition the run reported @@ -381,18 +433,17 @@ impl Core { /// Wakes the sources of what pipeline `id` pushed and the sinks of what it took. fn wake_neighbours(&mut self, id: PipelineId) { for index in 0..self.pipeline(id).outlets.len() { - let queue = self.arena.get(self.pipeline(id).outlets[index]); - if !queue.readable() { + let port = self.pipeline(id).outlets[index]; + let queue = self.arena.get_mut(port); + if !queue.take_changed() { continue; } - let port = self.pipeline(id).outlets[index]; match queue.source { - SourceEnd::Pipeline(source) => { - // A source waits on one inlet, so a batch on any other cannot unblock it. - let source_pipe = self.pipeline(source); - if let State::Blocked(Blocked::Inlet(waiting)) = source_pipe.state - && source_pipe.inlets[waiting] == port - { + SourceEnd::Pipeline { id: source, inlet } => { + // Record every change, including inlets the source is not waiting on. + let source_pipe = self.pipeline_mut(source); + source_pipe.changed.insert(inlet); + if source_pipe.state == State::Blocked(SourceBlocked::Inlet(inlet)) { self.make_runnable(source); } } @@ -402,8 +453,8 @@ impl Core { } // Only an inlet the run read can have gained room, so a source with many inlets, as a // concatenation of many chunks, pays for the inlets it read, not for all of them. - let touched = std::mem::take(&mut self.pipeline_mut(id).touched); - for &index in &touched { + let mut touched = std::mem::take(&mut self.pipeline_mut(id).touched); + for &index in touched.indices() { if let Some(sink) = self.arena.get(self.pipeline(id).inlets[index]).sink && sink != id && self.sink_has_room(sink) @@ -411,13 +462,8 @@ impl Core { self.make_runnable(sink); } } - let mut touched = touched; - let pipe = self.pipeline_mut(id); - for &index in &touched { - pipe.listed[index] = false; - } touched.clear(); - pipe.touched = touched; + self.pipeline_mut(id).touched = touched; } /// Queues split slot `slot` for the scan to read its stage's output. @@ -505,7 +551,7 @@ impl Core { .filter_map(|(id, pipe)| { let pipe = pipe.as_ref()?; Some(match pipe.state { - State::Blocked(Blocked::Inlet(index)) => { + State::Blocked(SourceBlocked::Inlet(index)) => { let port = pipe.inlets[index]; let queue = self.arena.get(port); format!( diff --git a/vortex-layout/src/plan/pipeline/scan/validate.rs b/vortex-layout/src/plan/pipeline/scan/validate.rs index 8d8fb2b16c6..d1147c4012b 100644 --- a/vortex-layout/src/plan/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/pipeline/scan/validate.rs @@ -7,10 +7,10 @@ //! of full ports can form. A broken promise shows up much later as "Scan is stuck" or as a //! wrong result; these checks fail at the pipeline that broke it. -use super::Blocked; use super::Core; use super::PipelineId; use super::Progress; +use super::SourceBlocked; use super::SourceEnd; impl Core { @@ -19,10 +19,11 @@ impl Core { pub(super) fn validate(&self) { for (id, pipe) in self.pipelines.iter().enumerate() { let Some(pipe) = pipe else { continue }; - for &port in &pipe.inlets { + for (index, &port) in pipe.inlets.iter().enumerate() { let source = self.arena.get(port).source; assert!( - source == SourceEnd::Pipeline(id) || self.arena.get(port).closed(), + source == (SourceEnd::Pipeline { id, inlet: index }) + || self.arena.get(port).closed(), "port {port} is an inlet of pipeline {id} but its source is {source:?}" ); } @@ -57,7 +58,7 @@ impl Core { pub(super) fn validate_progress(&self, id: PipelineId, progress: &Progress) { let pipe = self.pipeline(id); match progress { - Progress::Blocked(Blocked::Inlet(index)) => { + Progress::Blocked(SourceBlocked::Inlet(index)) => { assert!( *index < pipe.inlets.len(), "pipeline {id} waits on inlet {index} of {}", @@ -78,7 +79,7 @@ impl Core { queue.sink ); } - Progress::Blocked(Blocked::Io) => assert!( + Progress::Blocked(SourceBlocked::Io) => assert!( pipe.pending > 0, "pipeline {id} waits for a read with none in flight" ), @@ -92,7 +93,7 @@ impl Core { ); } } - Progress::Outlet => {} + Progress::Read(_) | Progress::Outlet => {} } } } diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs index c5451a74f7f..5b241ff7613 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -3,6 +3,7 @@ //! The runtime over hand-built sources: order, blocking, backpressure, and wakeups. +use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; @@ -131,7 +132,7 @@ fn masks_are_applied_by_the_source(#[values(false, true)] dense: bool) -> Vortex assert_rows(expected, run.arrays) } -/// A source runs only when an inlet has changed: every run of the probe after its first finds +/// A reader runs only when an inlet has changed: every run of the probe after its first finds /// a batch it has not seen, or an inlet newly closed. #[rstest] fn a_reader_runs_only_when_an_inlet_changed( @@ -155,8 +156,8 @@ fn a_reader_runs_only_when_an_inlet_changed( Ok(()) } -/// A sink is blocked once its source's inlet holds as many batches as the source allows, so a -/// fast pipeline never runs far ahead of a slow source. +/// A writer is blocked once its reader's inlet holds as many batches as the reader allows, so +/// a fast source never runs far ahead of a slow reader. #[test] fn a_full_inlet_blocks_its_writer() -> VortexResult<()> { let runs: Arc>> = Arc::default(); @@ -205,26 +206,36 @@ struct ManyReads { } impl Source for ManyReads { - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { - if let Some(segment) = self.wanted.pop() { - return Ok(Step::Read(segment)); + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if !self.wanted.is_empty() { + while let Some(segment) = self.wanted.pop() { + cx.request_read(segment); + } + return Ok(SourceReady::Read); } + Ok(if self.left == 0 || cx.has_bytes() { + SourceReady::Ready + } else { + SourceReady::Blocked(SourceBlocked::Io) + }) + } +} + +impl Operator for ManyReads { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { if let Some((segment, _bytes)) = cx.take_segment() { self.left -= 1; - return Ok(Step::Batch(indices([u64::from(*segment)]))); + return Ok(Step::Final(indices([u64::from(*segment)]))); } - Ok(if self.left == 0 { - Step::Finished - } else { - Step::Blocked(Blocked::Io) - }) + assert_eq!(self.left, 0, "computed while waiting for a read"); + Ok(Step::End) } } /// A source has all its reads in flight at once, and takes each segment's bytes as they arrive, /// tagged with the segment. -#[test] -fn a_source_has_several_reads_in_flight() -> VortexResult<()> { +#[rstest] +fn a_source_has_several_reads_in_flight(#[values(1, 3)] deliveries: usize) -> VortexResult<()> { let plan = SyntheticPlan::new(row_dtype(), 3, Vec::new(), |_, _, _| { Ok(Some(Chain::new(ManyReads { wanted: [3_u32, 2, 1].into_iter().map(SegmentId::from).collect(), @@ -250,10 +261,12 @@ fn a_source_has_several_reads_in_flight() -> VortexResult<()> { first_wait = false; } // Deliver the last asked for first. - let read = inflight - .pop() - .ok_or_else(|| vortex_err!("no read in flight"))?; - scan.deliver(read.id, empty_segment())?; + for _ in 0..deliveries { + let read = inflight + .pop() + .ok_or_else(|| vortex_err!("no read in flight"))?; + scan.deliver(read.id, empty_segment())?; + } } Turn::Done => break, } @@ -261,6 +274,77 @@ fn a_source_has_several_reads_in_flight() -> VortexResult<()> { assert_rows(indices([3, 2, 1]), arrays) } +/// Holds an inlet batch until a read chosen from that batch has arrived. +struct ReadAfterInlet { + batch: Option, +} + +impl Source for ReadAfterInlet { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + if self.batch.is_some() { + return Ok(if cx.has_bytes() { + SourceReady::Ready + } else { + SourceReady::Blocked(SourceBlocked::Io) + }); + } + let mut inlet = cx.inlet(0); + let Some(batch) = inlet.take() else { + return Ok(if inlet.closed() { + SourceReady::Ready + } else { + SourceReady::Blocked(SourceBlocked::Inlet(0)) + }); + }; + cx.request_read(SegmentId::from(u32::try_from(batch.len())?)); + self.batch = Some(batch); + Ok(SourceReady::Read) + } +} + +impl Operator for ReadAfterInlet { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + let Some(batch) = self.batch.take() else { + let inlet = cx.inlet(0); + assert!( + inlet.closed() && inlet.is_empty(), + "computed before the inlet was ready" + ); + return Ok(Step::End); + }; + let (segment, _) = cx + .take_segment() + .ok_or_else(|| vortex_err!("computed before the read arrived"))?; + assert_eq!(*segment, u32::try_from(batch.len())?); + Ok(Step::Final(batch)) + } +} + +#[test] +fn readiness_can_take_an_inlet_and_request_a_read() -> VortexResult<()> { + let child = RowSource::plan(6, vec![1, 2, 3], segment(9), false); + let plan = SyntheticPlan::new( + row_dtype(), + 6, + vec![child.clone()], + move |rows, mask, compiler| { + let chain = compiler + .compile(&child, rows, mask)? + .ok_or_else(|| vortex_err!("child has no rows"))?; + Ok(Some( + compiler.join_with(vec![chain], |_| 1, ReadAfterInlet { batch: None }), + )) + }, + ) + .into_plan(); + let run = run(&plan, 0..6, Mask::new_true(6), fifo)?; + assert_eq!( + run.events, + [Some(9), Some(1), None, Some(2), None, Some(3), None] + ); + assert_rows(indices(0..6), run.arrays) +} + /// Asks for every segment at once, emits the id of the first to arrive, and finishes without /// waiting for the rest. struct FirstOf { @@ -269,23 +353,37 @@ struct FirstOf { } impl Source for FirstOf { - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { if self.done { - return Ok(Step::Finished); - } - if let Some(segment) = self.wanted.pop() { - return Ok(Step::Read(segment)); + return Ok(SourceReady::Ready); } - Ok(match cx.take_segment() { - Some((segment, _bytes)) => { - self.done = true; - Step::Batch(indices([u64::from(*segment)])) + if !self.wanted.is_empty() { + while let Some(segment) = self.wanted.pop() { + cx.request_read(segment); } - None => Step::Blocked(Blocked::Io), + return Ok(SourceReady::Read); + } + Ok(if cx.has_bytes() { + SourceReady::Ready + } else { + SourceReady::Blocked(SourceBlocked::Io) }) } } +impl Operator for FirstOf { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + if self.done { + return Ok(Step::End); + } + let (segment, _) = cx + .take_segment() + .ok_or_else(|| vortex_err!("FirstOf computed before a read arrived"))?; + self.done = true; + Ok(Step::Final(indices([u64::from(*segment)]))) + } +} + /// A source that finishes with reads in flight has them cancelled: their bytes reach neither it /// nor the next split's pipeline, which reuses its id, and the scan is done without them. #[test] @@ -328,6 +426,156 @@ fn finishing_cancels_reads_in_flight() -> VortexResult<()> { Ok(()) } +/// Emits pairs of batches, requiring readiness only between pairs. +#[derive(Default)] +struct Pairs { + next: u64, + ready: bool, +} + +impl Source for Pairs { + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + assert!(!self.ready, "More must skip readiness"); + self.ready = true; + Ok(SourceReady::Ready) + } +} + +impl Operator for Pairs { + fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult { + assert!(matches!(input, Input::None)); + assert!(self.ready, "Final must check readiness again"); + if self.next == 6 { + return Ok(Step::End); + } + let batch = indices([self.next]); + self.next += 1; + if self.next % 2 == 1 { + Ok(Step::More(batch)) + } else { + self.ready = false; + Ok(Step::Final(batch)) + } + } +} + +struct Copies { + copies: usize, + left: usize, + batch: Option, + ended: bool, +} + +impl Copies { + fn new(copies: usize) -> Self { + Self { + copies, + left: 0, + batch: None, + ended: false, + } + } +} + +impl Operator for Copies { + fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult { + assert!(!self.ended, "a finished stage cannot run again"); + match input { + Input::Batch(batch) => { + assert!(self.batch.is_none(), "new input arrived before draining More"); + self.batch = Some(batch); + self.left = self.copies; + } + Input::None => assert!(self.batch.is_some(), "continued after Final"), + Input::End => { + assert!(self.batch.is_none(), "End arrived before draining More"); + self.ended = true; + return Ok(Step::End); + } + } + self.left -= 1; + if self.left == 0 { + Ok(Step::Final(self.batch.take().expect("batch is present"))) + } else { + Ok(Step::More( + self.batch.as_ref().expect("batch is present").clone(), + )) + } + } +} + +/// Buffers every input, then flushes with either Final or End terminating the continuation. +struct Buffered { + batches: VecDeque, + flushing: bool, + final_batch: bool, +} + +impl Operator for Buffered { + fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult { + match input { + Input::Batch(batch) => { + assert!(!self.flushing); + self.batches.push_back(batch); + return Ok(Step::End); + } + Input::End => { + assert!(!self.flushing, "End must be supplied exactly once"); + self.flushing = true; + } + Input::None => assert!(self.flushing), + } + Ok(match self.batches.pop_front() { + Some(batch) if self.final_batch && self.batches.is_empty() => Step::Final(batch), + Some(batch) => Step::More(batch), + None => Step::End, + }) + } +} + +#[rstest] +fn more_drains_before_new_input(#[values(1, 2)] capacity: usize) -> VortexResult<()> { + let plan = SyntheticPlan::new(row_dtype(), 36, Vec::new(), move |_, _, _| { + let chain = Chain::new(Pairs::default()) + .with(Copies::new(2)) + .with(Copies::new(3)); + Ok(Some(chain)) + }) + .into_plan(); + let plan = Probe::plan(vec![plan], capacity, Arc::default()); + let result = run(&plan, 0..36, Mask::new_true(36), fifo)?; + assert_eq!(result.arrays.len(), 36); + assert_rows(indices((0..6).flat_map(|row| [row; 6])), result.arrays) +} + +#[rstest] +fn more_flushes_before_ending_downstream( + #[values(1, 2)] capacity: usize, + #[values(false, true)] final_batch: bool, +) -> VortexResult<()> { + let plan = SyntheticPlan::new(row_dtype(), 36, Vec::new(), move |_, _, _| { + let chain = Chain::new(Pairs::default()) + .with(Buffered { + batches: VecDeque::new(), + flushing: false, + final_batch, + }) + .with(Copies::new(2)) + .with(Copies::new(3)) + .with(Buffered { + batches: VecDeque::new(), + flushing: false, + final_batch, + }); + Ok(Some(chain)) + }) + .into_plan(); + let plan = Probe::plan(vec![plan], capacity, Arc::default()); + let result = run(&plan, 0..36, Mask::new_true(36), fifo)?; + assert_eq!(result.arrays.len(), 36); + assert_rows(indices((0..6).flat_map(|row| [row; 6])), result.arrays) +} + /// Sources that break the scheduling contract, which debug builds catch. #[cfg(debug_assertions)] mod contract { @@ -338,14 +586,16 @@ mod contract { left: usize, } - impl Source for Repeat { - fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + impl Source for Repeat {} + + impl Operator for Repeat { + fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { Ok(match self.left.checked_sub(1) { Some(left) => { self.left = left; - Step::Batch(indices([0])) + Step::Final(indices([0])) } - None => Step::Finished, + None => Step::End, }) } } @@ -354,8 +604,14 @@ mod contract { struct WaitsWhenFull; impl Source for WaitsWhenFull { - fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { - Ok(Step::Blocked(Blocked::Inlet(0))) + fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { + Ok(SourceReady::Blocked(SourceBlocked::Inlet(0))) + } + } + + impl Operator for WaitsWhenFull { + fn compute(&mut self, _input: Input, _cx: &mut Cx<'_>) -> VortexResult { + Err(vortex_err!("a blocked source must not compute")) } } @@ -365,11 +621,9 @@ mod contract { #[should_panic(expected = "which is full")] fn waiting_on_a_full_inlet_is_caught() { let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |_, _, compiler| { - Ok(Some(compiler.join_with( - vec![Chain::new(Repeat { left: 2 })], - |_| 1, - WaitsWhenFull, - ))) + Ok(Some( + compiler.join_with(vec![Chain::new(Repeat { left: 2 })], |_| 1, WaitsWhenFull), + )) }) .into_plan(); let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..2)]).unwrap(); diff --git a/vortex-layout/src/plan/pipeline/synthetic.rs b/vortex-layout/src/plan/pipeline/synthetic.rs index ab0f650a062..6fd70667678 100644 --- a/vortex-layout/src/plan/pipeline/synthetic.rs +++ b/vortex-layout/src/plan/pipeline/synthetic.rs @@ -12,7 +12,6 @@ use std::ops::Range; use std::sync::Arc; use parking_lot::Mutex; -use vortex_array::ArrayRef; use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::buffer::BufferHandle; @@ -21,14 +20,18 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::CachedId; -use super::Blocked; use super::Chain; use super::Compiler; use super::Cx; +use super::Input; +use super::Operator; use super::Source; +use super::SourceBlocked; +use super::SourceReady; use super::Step; use crate::plan::Plan; use crate::plan::PlanChildren; @@ -172,19 +175,29 @@ impl RowSource { } impl Source for RowSource { - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { if let Some(segment) = self.segment.take() { self.waiting = true; - return Ok(Step::Read(segment)); + cx.request_read(segment); + return Ok(SourceReady::Read); } + Ok(if self.waiting && !cx.has_bytes() { + SourceReady::Blocked(SourceBlocked::Io) + } else { + SourceReady::Ready + }) + } +} + +impl Operator for RowSource { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { if self.waiting { - if cx.take_segment().is_none() { - return Ok(Step::Blocked(Blocked::Io)); - } self.waiting = false; + cx.take_segment() + .ok_or_else(|| vortex_err!("RowSource ran without its bytes"))?; } if self.remaining.is_empty() { - return Ok(Step::Finished); + return Ok(Step::End); } let len = self .pieces @@ -192,7 +205,7 @@ impl Source for RowSource { .unwrap_or(self.remaining.len()) .min(self.remaining.len()); let piece: Vec = self.remaining.drain(..len).collect(); - Ok(Step::Batch(Buffer::from(piece).into_array())) + Ok(Step::Final(Buffer::from(piece).into_array())) } } @@ -200,8 +213,8 @@ impl Source for RowSource { /// its sink had closed it. pub type ProbeRun = Vec<(usize, bool)>; -/// A source over one inlet per child that records what it sees each time it runs, then passes -/// its first inlet through, one batch per call, and drops what the others hold. +/// A source over one inlet per child that records what it sees on each readiness check, then +/// passes its first inlet through, one batch per call, and drops what the others hold. pub struct Probe { inlets: usize, runs: Arc>>, @@ -240,7 +253,7 @@ impl Probe { } impl Source for Probe { - fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { + fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { let run: ProbeRun = (0..self.inlets) .map(|index| { let inlet = cx.inlet(index); @@ -251,20 +264,28 @@ impl Source for Probe { for index in 1..self.inlets { while cx.inlet(index).take().is_some() {} } - let mut first = cx.inlet(0); - let batch: Option = first.take(); - match batch { - Some(batch) => Ok(Step::Batch(batch)), - None if first.closed() => { - for index in 1..self.inlets { - if !cx.inlet(index).closed() { - return Ok(Step::Blocked(Blocked::Inlet(index))); - } + let first = cx.inlet(0); + if !first.is_empty() { + return Ok(SourceReady::Ready); + } + if first.closed() { + for index in 1..self.inlets { + if !cx.inlet(index).closed() { + return Ok(SourceReady::Blocked(SourceBlocked::Inlet(index))); } - Ok(Step::Finished) } - None => Ok(Step::Blocked(Blocked::Inlet(0))), + return Ok(SourceReady::Ready); } + Ok(SourceReady::Blocked(SourceBlocked::Inlet(0))) + } +} + +impl Operator for Probe { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + Ok(match cx.inlet(0).take() { + Some(batch) => Step::Final(batch), + None => Step::End, + }) } } From 0e6559a54112176cb463b6be8a080b69eff04aaf Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 14:26:34 +0000 Subject: [PATCH 16/20] style(layout): cargo fmt the pipeline Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 --- .../src/plan/pipeline/scheduling_tests.rs | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/pipeline/scheduling_tests.rs index 5b241ff7613..5c6597df460 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/pipeline/scheduling_tests.rs @@ -331,9 +331,11 @@ fn readiness_can_take_an_inlet_and_request_a_read() -> VortexResult<()> { let chain = compiler .compile(&child, rows, mask)? .ok_or_else(|| vortex_err!("child has no rows"))?; - Ok(Some( - compiler.join_with(vec![chain], |_| 1, ReadAfterInlet { batch: None }), - )) + Ok(Some(compiler.join_with( + vec![chain], + |_| 1, + ReadAfterInlet { batch: None }, + ))) }, ) .into_plan(); @@ -482,7 +484,10 @@ impl Operator for Copies { assert!(!self.ended, "a finished stage cannot run again"); match input { Input::Batch(batch) => { - assert!(self.batch.is_none(), "new input arrived before draining More"); + assert!( + self.batch.is_none(), + "new input arrived before draining More" + ); self.batch = Some(batch); self.left = self.copies; } @@ -621,9 +626,11 @@ mod contract { #[should_panic(expected = "which is full")] fn waiting_on_a_full_inlet_is_caught() { let plan = SyntheticPlan::new(row_dtype(), 2, Vec::new(), |_, _, compiler| { - Ok(Some( - compiler.join_with(vec![Chain::new(Repeat { left: 2 })], |_| 1, WaitsWhenFull), - )) + Ok(Some(compiler.join_with( + vec![Chain::new(Repeat { left: 2 })], + |_| 1, + WaitsWhenFull, + ))) }) .into_plan(); let mut scan = Scan::try_new(SESSION.clone(), plan, vec![Split::all(0..2)]).unwrap(); From 07cc6f9044e3f5cc2abad5342822e5a008f8b233 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 12 Oct 2026 15:37:40 +0100 Subject: [PATCH 17/20] refactor(layout): organize physical plans and execution modules Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/mod.rs | 76 +-------------- .../src/plan/{ => physical}/children.rs | 2 +- .../src/plan/{ => physical}/display.rs | 0 .../src/plan/{ => physical}/lower.rs | 14 +-- vortex-layout/src/plan/physical/mod.rs | 83 ++++++++++++++++ .../{pipeline/mod.rs => physical/operator.rs} | 96 ++++--------------- .../ops => physical/operators}/concat.rs | 14 +-- .../ops => physical/operators}/filter.rs | 8 +- .../ops => physical/operators}/mod.rs | 2 +- .../ops => physical/operators}/once.rs | 10 +- .../ops => physical/operators}/pack.rs | 18 ++-- .../ops => physical/operators}/scan.rs | 16 ++-- .../src/plan/{ => physical}/optimize.rs | 6 +- .../src/plan/{ => physical}/optimizer/mod.rs | 0 .../plan/{ => physical}/optimizer/rules.rs | 6 +- .../plan/{ => physical}/pipeline/compile.rs | 18 ++-- .../src/plan/physical/pipeline/mod.rs | 66 +++++++++++++ .../src/plan/{ => physical}/pipeline/port.rs | 2 +- .../src/plan/{ => physical}/pipeline/scan.rs | 18 ++-- .../{ => physical}/pipeline/scan/validate.rs | 0 .../pipeline/scheduling_tests.rs | 9 +- .../plan/{ => physical}/pipeline/synthetic.rs | 26 ++--- .../src/plan/{ => physical}/pipeline/tests.rs | 21 ++-- .../src/plan/{ => physical}/plans/concat.rs | 28 +++--- .../src/plan/{ => physical}/plans/eval.rs | 16 ++-- .../src/plan/{ => physical}/plans/filter.rs | 24 ++--- .../plan/{ => physical}/plans/list_pack.rs | 12 +-- .../src/plan/{ => physical}/plans/mod.rs | 0 .../src/plan/{ => physical}/plans/pack.rs | 32 +++---- .../src/plan/{ => physical}/plans/row_idx.rs | 20 ++-- .../plan/{ => physical}/plans/segment_scan.rs | 18 ++-- .../src/plan/{ => physical}/plans/take.rs | 20 ++-- .../src/plan/{ => physical}/tests.rs | 0 .../src/plan/{ => physical}/typed.rs | 14 +-- .../src/plan/{ => physical}/vtable.rs | 12 +-- 35 files changed, 371 insertions(+), 336 deletions(-) rename vortex-layout/src/plan/{ => physical}/children.rs (98%) rename vortex-layout/src/plan/{ => physical}/display.rs (100%) rename vortex-layout/src/plan/{ => physical}/lower.rs (95%) create mode 100644 vortex-layout/src/plan/physical/mod.rs rename vortex-layout/src/plan/{pipeline/mod.rs => physical/operator.rs} (64%) rename vortex-layout/src/plan/{pipeline/ops => physical/operators}/concat.rs (81%) rename vortex-layout/src/plan/{pipeline/ops => physical/operators}/filter.rs (91%) rename vortex-layout/src/plan/{pipeline/ops => physical/operators}/mod.rs (97%) rename vortex-layout/src/plan/{pipeline/ops => physical/operators}/once.rs (74%) rename vortex-layout/src/plan/{pipeline/ops => physical/operators}/pack.rs (93%) rename vortex-layout/src/plan/{pipeline/ops => physical/operators}/scan.rs (90%) rename vortex-layout/src/plan/{ => physical}/optimize.rs (90%) rename vortex-layout/src/plan/{ => physical}/optimizer/mod.rs (100%) rename vortex-layout/src/plan/{ => physical}/optimizer/rules.rs (98%) rename vortex-layout/src/plan/{ => physical}/pipeline/compile.rs (97%) create mode 100644 vortex-layout/src/plan/physical/pipeline/mod.rs rename vortex-layout/src/plan/{ => physical}/pipeline/port.rs (99%) rename vortex-layout/src/plan/{ => physical}/pipeline/scan.rs (98%) rename vortex-layout/src/plan/{ => physical}/pipeline/scan/validate.rs (100%) rename vortex-layout/src/plan/{ => physical}/pipeline/scheduling_tests.rs (98%) rename vortex-layout/src/plan/{ => physical}/pipeline/synthetic.rs (94%) rename vortex-layout/src/plan/{ => physical}/pipeline/tests.rs (96%) rename vortex-layout/src/plan/{ => physical}/plans/concat.rs (91%) rename vortex-layout/src/plan/{ => physical}/plans/eval.rs (92%) rename vortex-layout/src/plan/{ => physical}/plans/filter.rs (88%) rename vortex-layout/src/plan/{ => physical}/plans/list_pack.rs (96%) rename vortex-layout/src/plan/{ => physical}/plans/mod.rs (100%) rename vortex-layout/src/plan/{ => physical}/plans/pack.rs (97%) rename vortex-layout/src/plan/{ => physical}/plans/row_idx.rs (95%) rename vortex-layout/src/plan/{ => physical}/plans/segment_scan.rs (89%) rename vortex-layout/src/plan/{ => physical}/plans/take.rs (92%) rename vortex-layout/src/plan/{ => physical}/tests.rs (100%) rename vortex-layout/src/plan/{ => physical}/typed.rs (97%) rename vortex-layout/src/plan/{ => physical}/vtable.rs (94%) diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 36cb5047149..88c104f81d3 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -1,78 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Physical plans for scans. -//! -//! A plan is a tree of physical operators over a row domain. Operator identity and -//! operator-specific state do not depend on the source layout kind, so rewrites can reason about a -//! plan's shape alone. The common child container can initialize individual slots lazily. +//! Scan planning. -mod children; -mod display; -mod lower; -mod optimize; -pub mod optimizer; -pub mod pipeline; -mod plans; -mod typed; -mod vtable; - -pub use children::PlanChildren; -pub use display::PlanIndentedFormatter; -pub use display::PlanSummaryExtractor; -pub use display::PlanTreeContext; -pub use display::PlanTreeDisplay; -pub use display::PlanTreeExtractor; -pub use lower::lower; -pub use optimize::optimize; -pub use plans::Concat; -pub use plans::ConcatData; -pub use plans::ConcatPlan; -pub use plans::Eval; -pub use plans::EvalData; -pub use plans::EvalPlan; -pub use plans::Filter; -pub use plans::FilterData; -pub use plans::FilterPlan; -pub use plans::ListPack; -pub use plans::ListPackData; -pub use plans::ListPackPlan; -pub use plans::Pack; -pub use plans::PackData; -pub use plans::PackPlan; -pub use plans::RowIdx; -pub use plans::RowIdxData; -pub use plans::RowIdxPlan; -pub use plans::SegmentScan; -pub use plans::SegmentScanData; -pub use plans::SegmentScanPlan; -pub use plans::Take; -pub use plans::TakePlan; -pub use plans::plan_row_idx_expression; -pub use plans::row_idx_dtype; -pub use typed::DynPlan; -pub use typed::Plan; -pub use typed::PlanParts; -pub use typed::PlanRef; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -pub use vtable::PlanId; -pub use vtable::PlanVTable; - -/// Returns an error when `children` does not have exactly `expected` entries. -pub(crate) fn check_child_count( - name: &str, - children: &PlanChildren, - expected: usize, -) -> VortexResult<()> { - if children.len() != expected { - vortex_bail!( - "{name} expects {expected} children but got {}", - children.len() - ); - } - Ok(()) -} - -#[cfg(test)] -mod tests; +pub mod physical; diff --git a/vortex-layout/src/plan/children.rs b/vortex-layout/src/plan/physical/children.rs similarity index 98% rename from vortex-layout/src/plan/children.rs rename to vortex-layout/src/plan/physical/children.rs index b045c1be8d2..8b1475e2a98 100644 --- a/vortex-layout/src/plan/children.rs +++ b/vortex-layout/src/plan/physical/children.rs @@ -8,7 +8,7 @@ use once_cell::sync::OnceCell; use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; type ChildInitializer = dyn Fn(usize) -> VortexResult + 'static + Send + Sync; diff --git a/vortex-layout/src/plan/display.rs b/vortex-layout/src/plan/physical/display.rs similarity index 100% rename from vortex-layout/src/plan/display.rs rename to vortex-layout/src/plan/physical/display.rs diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/physical/lower.rs similarity index 95% rename from vortex-layout/src/plan/lower.rs rename to vortex-layout/src/plan/physical/lower.rs index e8911734549..bade3544384 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -24,13 +24,13 @@ use crate::layouts::list::OFFSETS_CHILD_INDEX; use crate::layouts::list::VALIDITY_CHILD_INDEX; use crate::layouts::struct_::Struct; use crate::layouts::struct_::StructLayout; -use crate::plan::ConcatPlan; -use crate::plan::ListPackPlan; -use crate::plan::PackPlan; -use crate::plan::PlanChildren; -use crate::plan::PlanRef; -use crate::plan::SegmentScanPlan; -use crate::plan::TakePlan; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::ListPackPlan; +use crate::plan::physical::PackPlan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::TakePlan; /// Constructs a physical-plan fixture from `layout` for tests. /// diff --git a/vortex-layout/src/plan/physical/mod.rs b/vortex-layout/src/plan/physical/mod.rs new file mode 100644 index 00000000000..8e604ff7cf3 --- /dev/null +++ b/vortex-layout/src/plan/physical/mod.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Physical plans for scans. +//! +//! A plan is a tree of physical operators over a row domain. Operator identity and +//! operator-specific state do not depend on the source layout kind, so rewrites can reason about a +//! plan's shape alone. The common child container can initialize individual slots lazily. +//! +//! [`operator`] defines the execution contracts, `operators` implements the concrete sources +//! and stages, and [`pipeline`] compiles and runs them. + +mod children; +mod display; +mod lower; +pub mod operator; +pub(crate) mod operators; +mod optimize; +pub mod optimizer; +pub mod pipeline; +mod plans; +mod typed; +mod vtable; + +pub use children::PlanChildren; +pub use display::PlanIndentedFormatter; +pub use display::PlanSummaryExtractor; +pub use display::PlanTreeContext; +pub use display::PlanTreeDisplay; +pub use display::PlanTreeExtractor; +pub use lower::lower; +pub use optimize::optimize; +pub use plans::Concat; +pub use plans::ConcatData; +pub use plans::ConcatPlan; +pub use plans::Eval; +pub use plans::EvalData; +pub use plans::EvalPlan; +pub use plans::Filter; +pub use plans::FilterData; +pub use plans::FilterPlan; +pub use plans::ListPack; +pub use plans::ListPackData; +pub use plans::ListPackPlan; +pub use plans::Pack; +pub use plans::PackData; +pub use plans::PackPlan; +pub use plans::RowIdx; +pub use plans::RowIdxData; +pub use plans::RowIdxPlan; +pub use plans::SegmentScan; +pub use plans::SegmentScanData; +pub use plans::SegmentScanPlan; +pub use plans::Take; +pub use plans::TakePlan; +pub use plans::plan_row_idx_expression; +pub use plans::row_idx_dtype; +pub use typed::DynPlan; +pub use typed::Plan; +pub use typed::PlanParts; +pub use typed::PlanRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +pub use vtable::PlanId; +pub use vtable::PlanVTable; + +/// Returns an error when `children` does not have exactly `expected` entries. +pub(crate) fn check_child_count( + name: &str, + children: &PlanChildren, + expected: usize, +) -> VortexResult<()> { + if children.len() != expected { + vortex_bail!( + "{name} expects {expected} children but got {}", + children.len() + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/physical/operator.rs similarity index 64% rename from vortex-layout/src/plan/pipeline/mod.rs rename to vortex-layout/src/plan/physical/operator.rs index 80eb6462c4e..83d7c81fefd 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -1,47 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Push-based execution of physical plans as pipelines. -//! -//! A [`Scan`] runs a [`PlanRef`](crate::plan::PlanRef) over a list of splits, row ranges of the -//! plan's domain. Each split is compiled into pipelines: a pipeline is one [`Source`], zero or -//! more [`Operator`]s, and the ports its last stage writes. The driver asks the source for a batch -//! and pushes it through the stages into the ports. Every port runs from one pipeline's sink, -//! the end that pushes into its outlets, to another's source, which reads it as an inlet, and -//! the plan joining them declares how many batches each inlet may hold. -//! -//! A source checks its readiness before computing, waiting for a segment read or an inlet. -//! The driver also blocks a pipeline when an outlet is full. The scheduler re-runs a blocked -//! pipeline when that one condition clears. -//! -//! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and -//! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. -//! -//! # Masks -//! -//! Every mask a scan applies is known when the pipelines that apply it are compiled: the split's -//! selection, and the rows the earlier conjuncts kept. The -//! compiler places the mask in the stage that reads the rows, so chunks a mask selects nothing -//! of are never built and never read. -//! -//! # Sharing -//! -//! A segment several readers decode is decoded once. Before the scan starts, one pass over the -//! plans every split's stages run records the rows each segment is read for, so the splits that -//! will read it are found by binary search over the splits. The first reader of a segment read -//! more than once builds one pipeline that decodes it and fans the decoded array out into one -//! single-use port per reader, tagged with the reader's split: a reader in the same stage, a -//! later stage of a query, or a later split claims its own port. A split's ports it never -//! claims, because its rows were pruned or a stage never ran, leave the sharing pipeline's -//! outlets and are freed when it finishes, so nothing is ever pushed into a port nobody will -//! read, and a claimed port is dropped once its source has drained it, so a decoded segment -//! lives only as long as a split that may still read it. - -mod compile; -pub(crate) mod ops; -mod port; -mod scan; -pub mod synthetic; +//! Contracts and execution context for pipeline sources and operators. use std::collections::VecDeque; @@ -52,24 +12,11 @@ use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; use vortex_session::VortexSession; -pub use self::compile::Chain; -pub use self::compile::Compiler; -pub use self::compile::Reach; -pub(crate) use self::compile::overlapping; -use self::port::Arena; -pub use self::port::Inlet; -use self::port::PortId; -pub use self::scan::ReadId; -pub use self::scan::ReadRequest; -pub use self::scan::Scan; -pub use self::scan::Split; -pub use self::scan::Turn; +use super::pipeline::Arena; +use super::pipeline::Inlet; +use super::pipeline::PortId; use crate::segments::SegmentId; -/// Batches an inlet holds before its sink is blocked, unless the plan joining it says -/// otherwise; see [`Compiler::join_with`]. -pub const DEFAULT_CAPACITY: usize = 8; - /// The output of one non-blocking compute call. #[derive(Debug)] pub enum Step { @@ -143,34 +90,34 @@ pub trait Operator: Send { /// Distinct inlet indices, with constant-time insertion and removal. #[derive(Default)] -struct InletSet { +pub(super) struct InletSet { indices: Vec, positions: Vec, } impl InletSet { - fn new(inlets: usize) -> Self { + pub(super) fn new(inlets: usize) -> Self { Self { indices: Vec::with_capacity(inlets), positions: vec![usize::MAX; inlets], } } - fn all(inlets: usize) -> Self { + pub(super) fn all(inlets: usize) -> Self { Self { indices: (0..inlets).collect(), positions: (0..inlets).collect(), } } - fn insert(&mut self, index: usize) { + pub(super) fn insert(&mut self, index: usize) { if self.positions[index] == usize::MAX { self.positions[index] = self.indices.len(); self.indices.push(index); } } - fn remove(&mut self, index: usize) { + pub(super) fn remove(&mut self, index: usize) { let position = self.positions[index]; if position == usize::MAX { return; @@ -188,11 +135,11 @@ impl InletSet { Some(index) } - fn indices(&self) -> &[usize] { + pub(super) fn indices(&self) -> &[usize] { &self.indices } - fn clear(&mut self) { + pub(super) fn clear(&mut self) { for &index in &self.indices { self.positions[index] = usize::MAX; } @@ -203,14 +150,14 @@ impl InletSet { /// What a source or stage may touch during one call: its pipeline's inlets, the bytes the /// source asked for, and the session. pub struct Cx<'a> { - arena: &'a mut Arena, - inlets: &'a [PortId], - bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, - reads: SmallVec<[SegmentId; 1]>, - session: &'a VortexSession, - exec: &'a mut ExecutionCtx, - touched: &'a mut InletSet, - changed: &'a mut InletSet, + pub(super) arena: &'a mut Arena, + pub(super) inlets: &'a [PortId], + pub(super) bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, + pub(super) reads: SmallVec<[SegmentId; 1]>, + pub(super) session: &'a VortexSession, + pub(super) exec: &'a mut ExecutionCtx, + pub(super) touched: &'a mut InletSet, + pub(super) changed: &'a mut InletSet, } impl Cx<'_> { @@ -254,8 +201,3 @@ impl Cx<'_> { self.exec } } - -#[cfg(test)] -mod scheduling_tests; -#[cfg(test)] -mod tests; diff --git a/vortex-layout/src/plan/pipeline/ops/concat.rs b/vortex-layout/src/plan/physical/operators/concat.rs similarity index 81% rename from vortex-layout/src/plan/pipeline/ops/concat.rs rename to vortex-layout/src/plan/physical/operators/concat.rs index e981f1b65e3..57110556280 100644 --- a/vortex-layout/src/plan/pipeline/ops/concat.rs +++ b/vortex-layout/src/plan/physical/operators/concat.rs @@ -6,13 +6,13 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::plan::pipeline::Cx; -use crate::plan::pipeline::Input; -use crate::plan::pipeline::Operator; -use crate::plan::pipeline::Source; -use crate::plan::pipeline::SourceBlocked; -use crate::plan::pipeline::SourceReady; -use crate::plan::pipeline::Step; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; /// Emits its inlets one after another, each to its end. pub(crate) struct ConcatSource { diff --git a/vortex-layout/src/plan/pipeline/ops/filter.rs b/vortex-layout/src/plan/physical/operators/filter.rs similarity index 91% rename from vortex-layout/src/plan/pipeline/ops/filter.rs rename to vortex-layout/src/plan/physical/operators/filter.rs index 690c654a922..ecd0943f97c 100644 --- a/vortex-layout/src/plan/pipeline/ops/filter.rs +++ b/vortex-layout/src/plan/physical/operators/filter.rs @@ -9,10 +9,10 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; -use crate::plan::pipeline::Cx; -use crate::plan::pipeline::Input; -use crate::plan::pipeline::Operator; -use crate::plan::pipeline::Step; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Step; /// The selected fraction of a predicate's batch at or above which it is executed whole and then /// filtered, as the default scan's flat reader does. diff --git a/vortex-layout/src/plan/pipeline/ops/mod.rs b/vortex-layout/src/plan/physical/operators/mod.rs similarity index 97% rename from vortex-layout/src/plan/pipeline/ops/mod.rs rename to vortex-layout/src/plan/physical/operators/mod.rs index 11f45398809..e53f856f973 100644 --- a/vortex-layout/src/plan/pipeline/ops/mod.rs +++ b/vortex-layout/src/plan/physical/operators/mod.rs @@ -24,7 +24,7 @@ use vortex_array::arrays::ChunkedArray; use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::plan::pipeline::Inlet; +use crate::plan::physical::pipeline::Inlet; /// Takes the first `len` rows of the inlet, across as many batches as hold them, slicing the last /// and leaving its rest in place. Rows spanning batches come back chunked, not copied. diff --git a/vortex-layout/src/plan/pipeline/ops/once.rs b/vortex-layout/src/plan/physical/operators/once.rs similarity index 74% rename from vortex-layout/src/plan/pipeline/ops/once.rs rename to vortex-layout/src/plan/physical/operators/once.rs index 4075be8be7f..745f00fb1bf 100644 --- a/vortex-layout/src/plan/pipeline/ops/once.rs +++ b/vortex-layout/src/plan/physical/operators/once.rs @@ -6,11 +6,11 @@ use vortex_array::ArrayRef; use vortex_error::VortexResult; -use crate::plan::pipeline::Cx; -use crate::plan::pipeline::Input; -use crate::plan::pipeline::Operator; -use crate::plan::pipeline::Source; -use crate::plan::pipeline::Step; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::Step; /// Emits one batch it was built with. pub(crate) struct OnceSource(Option); diff --git a/vortex-layout/src/plan/pipeline/ops/pack.rs b/vortex-layout/src/plan/physical/operators/pack.rs similarity index 93% rename from vortex-layout/src/plan/pipeline/ops/pack.rs rename to vortex-layout/src/plan/physical/operators/pack.rs index 5b2af7fa06e..4fefa368f05 100644 --- a/vortex-layout/src/plan/pipeline/ops/pack.rs +++ b/vortex-layout/src/plan/physical/operators/pack.rs @@ -14,15 +14,15 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use super::take_rows; -use crate::plan::pipeline::Cx; -use crate::plan::pipeline::Inlet; -use crate::plan::pipeline::InletSet; -use crate::plan::pipeline::Input; -use crate::plan::pipeline::Operator; -use crate::plan::pipeline::Source; -use crate::plan::pipeline::SourceBlocked; -use crate::plan::pipeline::SourceReady; -use crate::plan::pipeline::Step; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::InletSet; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; +use crate::plan::physical::pipeline::Inlet; #[derive(Clone, Copy)] enum InletState { diff --git a/vortex-layout/src/plan/pipeline/ops/scan.rs b/vortex-layout/src/plan/physical/operators/scan.rs similarity index 90% rename from vortex-layout/src/plan/pipeline/ops/scan.rs rename to vortex-layout/src/plan/physical/operators/scan.rs index 227da9fca9c..d5189bdbd89 100644 --- a/vortex-layout/src/plan/pipeline/ops/scan.rs +++ b/vortex-layout/src/plan/physical/operators/scan.rs @@ -13,14 +13,14 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::plan::SegmentScanPlan; -use crate::plan::pipeline::Cx; -use crate::plan::pipeline::Input; -use crate::plan::pipeline::Operator; -use crate::plan::pipeline::Source; -use crate::plan::pipeline::SourceBlocked; -use crate::plan::pipeline::SourceReady; -use crate::plan::pipeline::Step; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; /// Decodes the whole of a scan's segment. diff --git a/vortex-layout/src/plan/optimize.rs b/vortex-layout/src/plan/physical/optimize.rs similarity index 90% rename from vortex-layout/src/plan/optimize.rs rename to vortex-layout/src/plan/physical/optimize.rs index f471b7744f9..32023f0ee05 100644 --- a/vortex-layout/src/plan/optimize.rs +++ b/vortex-layout/src/plan/physical/optimize.rs @@ -8,9 +8,9 @@ use vortex_error::VortexResult; -use crate::plan::PlanRef; -use crate::plan::optimizer::reduce_parent; -use crate::plan::optimizer::reduce_plan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::optimizer::reduce_parent; +use crate::plan::physical::optimizer::reduce_plan; fn reduce(plan: &PlanRef) -> VortexResult> { if let Some(rewritten) = reduce_plan(plan)? { diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs similarity index 100% rename from vortex-layout/src/plan/optimizer/mod.rs rename to vortex-layout/src/plan/physical/optimizer/mod.rs diff --git a/vortex-layout/src/plan/optimizer/rules.rs b/vortex-layout/src/plan/physical/optimizer/rules.rs similarity index 98% rename from vortex-layout/src/plan/optimizer/rules.rs rename to vortex-layout/src/plan/physical/optimizer/rules.rs index 12490201ee5..1c32e9c6f60 100644 --- a/vortex-layout/src/plan/optimizer/rules.rs +++ b/vortex-layout/src/plan/physical/optimizer/rules.rs @@ -11,9 +11,9 @@ use vortex_error::VortexResult; #[cfg(debug_assertions)] use vortex_error::vortex_ensure_eq; -use crate::plan::Plan; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; /// A rewrite over one concrete plan operator. /// diff --git a/vortex-layout/src/plan/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs similarity index 97% rename from vortex-layout/src/plan/pipeline/compile.rs rename to vortex-layout/src/plan/physical/pipeline/compile.rs index 7131a25cb52..998996b93f9 100644 --- a/vortex-layout/src/plan/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -13,17 +13,17 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use super::DEFAULT_CAPACITY; -use super::Operator; -use super::Source; -use super::ops::ConcatSource; -use super::ops::ScanSource; -use super::ops::SelectStage; use super::port::PortId; use super::port::SourceEnd; use super::scan::Core; -use crate::plan::ConcatPlan; -use crate::plan::PlanRef; -use crate::plan::SegmentScanPlan; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operators::ConcatSource; +use crate::plan::physical::operators::ScanSource; +use crate::plan::physical::operators::SelectStage; use crate::segments::SegmentId; /// A pipeline under construction: a source and the stages after it, not yet given outlets, so a @@ -73,7 +73,7 @@ impl Core { } /// What a plan compiles its children through. See -/// [`PlanVTable::compile`](crate::plan::PlanVTable::compile). +/// [`PlanVTable::compile`](crate::plan::physical::PlanVTable::compile). pub struct Compiler<'a> { pub(crate) core: &'a mut Core, } diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs new file mode 100644 index 00000000000..ab74d6db56b --- /dev/null +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Push-based execution of physical plans as pipelines. +//! +//! A [`Scan`] runs a [`PlanRef`](crate::plan::physical::PlanRef) over a list of splits, row ranges +//! of the plan's domain. Each split is compiled into pipelines: a pipeline is one +//! [`Source`](super::operator::Source), zero or more [`Operator`](super::operator::Operator)s, +//! and the ports its last stage writes. The driver asks the source for a batch and pushes it +//! through the stages into the ports. Every port runs from one pipeline's sink, +//! the end that pushes into its outlets, to another's source, which reads it as an inlet, and +//! the plan joining them declares how many batches each inlet may hold. +//! +//! A source checks its readiness before computing, waiting for a segment read or an inlet. +//! The driver also blocks a pipeline when an outlet is full. The scheduler re-runs a blocked +//! pipeline when that one condition clears. +//! +//! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and +//! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. +//! +//! # Masks +//! +//! Every mask a scan applies is known when the pipelines that apply it are compiled: the split's +//! selection, and the rows the earlier conjuncts kept. The +//! compiler places the mask in the stage that reads the rows, so chunks a mask selects nothing +//! of are never built and never read. +//! +//! # Sharing +//! +//! A segment several readers decode is decoded once. Before the scan starts, one pass over the +//! plans every split's stages run records the rows each segment is read for, so the splits that +//! will read it are found by binary search over the splits. The first reader of a segment read +//! more than once builds one pipeline that decodes it and fans the decoded array out into one +//! single-use port per reader, tagged with the reader's split: a reader in the same stage, a +//! later stage of a query, or a later split claims its own port. A split's ports it never +//! claims, because its rows were pruned or a stage never ran, leave the sharing pipeline's +//! outlets and are freed when it finishes, so nothing is ever pushed into a port nobody will +//! read, and a claimed port is dropped once its source has drained it, so a decoded segment +//! lives only as long as a split that may still read it. + +mod compile; +mod port; +mod scan; +pub mod synthetic; + +pub use self::compile::Chain; +pub use self::compile::Compiler; +pub use self::compile::Reach; +pub(crate) use self::compile::overlapping; +pub(super) use self::port::Arena; +pub use self::port::Inlet; +pub(super) use self::port::PortId; +pub use self::scan::ReadId; +pub use self::scan::ReadRequest; +pub use self::scan::Scan; +pub use self::scan::Split; +pub use self::scan::Turn; + +/// Batches an inlet holds before its sink is blocked, unless the plan joining it says +/// otherwise; see [`Compiler::join_with`]. +pub const DEFAULT_CAPACITY: usize = 8; + +#[cfg(test)] +mod scheduling_tests; +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/physical/pipeline/port.rs similarity index 99% rename from vortex-layout/src/plan/pipeline/port.rs rename to vortex-layout/src/plan/physical/pipeline/port.rs index 72911840a66..a5c4c38a6e4 100644 --- a/vortex-layout/src/plan/pipeline/port.rs +++ b/vortex-layout/src/plan/physical/pipeline/port.rs @@ -12,7 +12,7 @@ use std::collections::VecDeque; use vortex_array::ArrayRef; -use crate::plan::pipeline::DEFAULT_CAPACITY; +use crate::plan::physical::pipeline::DEFAULT_CAPACITY; /// Identifies a port in the scan's arena. pub(crate) type PortId = usize; diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs similarity index 98% rename from vortex-layout/src/plan/pipeline/scan.rs rename to vortex-layout/src/plan/physical/pipeline/scan.rs index a2694b279db..88a3ca13677 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -21,15 +21,7 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use super::Cx; use super::DEFAULT_CAPACITY; -use super::InletSet; -use super::Input; -use super::Operator; -use super::Source; -use super::SourceBlocked; -use super::SourceReady; -use super::Step; use super::compile::Chain; use super::compile::Shares; use super::port::Arena; @@ -37,7 +29,15 @@ use super::port::PipelineId; use super::port::PortId; use super::port::Slab; use super::port::SourceEnd; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::InletSet; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; #[cfg(debug_assertions)] diff --git a/vortex-layout/src/plan/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs similarity index 100% rename from vortex-layout/src/plan/pipeline/scan/validate.rs rename to vortex-layout/src/plan/physical/pipeline/scan/validate.rs diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs similarity index 98% rename from vortex-layout/src/plan/pipeline/scheduling_tests.rs rename to vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index 5b241ff7613..e255689fd94 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -26,7 +26,14 @@ use super::synthetic::SyntheticPlan; use super::synthetic::empty_segment; use super::synthetic::row_dtype; use super::*; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; use crate::test::SESSION; diff --git a/vortex-layout/src/plan/pipeline/synthetic.rs b/vortex-layout/src/plan/physical/pipeline/synthetic.rs similarity index 94% rename from vortex-layout/src/plan/pipeline/synthetic.rs rename to vortex-layout/src/plan/physical/pipeline/synthetic.rs index 6fd70667678..b6046c16c21 100644 --- a/vortex-layout/src/plan/pipeline/synthetic.rs +++ b/vortex-layout/src/plan/physical/pipeline/synthetic.rs @@ -26,19 +26,19 @@ use vortex_session::registry::CachedId; use super::Chain; use super::Compiler; -use super::Cx; -use super::Input; -use super::Operator; -use super::Source; -use super::SourceBlocked; -use super::SourceReady; -use super::Step; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; /// Compiles a [`SyntheticPlan`] over some rows, restricted to a mask. diff --git a/vortex-layout/src/plan/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs similarity index 96% rename from vortex-layout/src/plan/pipeline/tests.rs rename to vortex-layout/src/plan/physical/pipeline/tests.rs index 2a2392633cd..12941035cd5 100644 --- a/vortex-layout/src/plan/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -4,6 +4,7 @@ // Fixtures use a 20-row domain and name columns after their struct fields. #![allow(clippy::cast_possible_truncation, clippy::many_single_char_names)] +use std::collections::VecDeque; use std::ops::Range; use rstest::rstest; @@ -29,7 +30,6 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::ReadContext; -use super::ops::PackSource; use super::port::SourceEnd; use super::synthetic::row_dtype; use super::*; @@ -38,9 +38,18 @@ use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::struct_::StructLayout; -use crate::plan::PlanRef; -use crate::plan::SegmentScan; -use crate::plan::lower; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::lower; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::InletSet; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; +use crate::plan::physical::operators::PackSource; use crate::segments::SegmentId; use crate::test::SESSION; @@ -550,7 +559,7 @@ fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> Vortex let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); let scan = lower(&store.flat(&values)?)?; assert!(scan.is::()); - let filtered = crate::plan::FilterPlan::new(scan.clone()).into_plan(); + let filtered = crate::plan::physical::FilterPlan::new(scan.clone()).into_plan(); let rows = 3..9; let mask = sel.mask(6); @@ -582,7 +591,7 @@ fn a_segment_two_readers_need_is_read_once() -> VortexResult<()> { .as_struct_fields_opt() .vortex_expect("struct") .clone(); - let plan = crate::plan::PackPlan::try_new( + let plan = crate::plan::physical::PackPlan::try_new( fields, vortex_array::dtype::Nullability::NonNullable, ROWS, diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs similarity index 91% rename from vortex-layout/src/plan/plans/concat.rs rename to vortex-layout/src/plan/physical/plans/concat.rs index d2562a6001f..b2e1ffea0f7 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -12,20 +12,20 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use vortex_session::registry::CachedId; -use crate::plan::Eval; -use crate::plan::EvalPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::optimizer::PlanParentReduceRule; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; -use crate::plan::pipeline::ops::ConcatSource; -use crate::plan::pipeline::overlapping; +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::operators::ConcatSource; +use crate::plan::physical::optimizer::PlanParentReduceRule; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; +use crate::plan::physical::pipeline::overlapping; use crate::segments::SegmentId; /// Concatenates its children row-wise. diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs similarity index 92% rename from vortex-layout/src/plan/plans/eval.rs rename to vortex-layout/src/plan/physical/plans/eval.rs index f1edb967e41..04281bdfefc 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -10,14 +10,14 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::optimizer::PlanReduceRule; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::optimizer::PlanReduceRule; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] diff --git a/vortex-layout/src/plan/plans/filter.rs b/vortex-layout/src/plan/physical/plans/filter.rs similarity index 88% rename from vortex-layout/src/plan/plans/filter.rs rename to vortex-layout/src/plan/physical/plans/filter.rs index 8a38a2c06d6..4850674c67b 100644 --- a/vortex-layout/src/plan/plans/filter.rs +++ b/vortex-layout/src/plan/physical/plans/filter.rs @@ -11,18 +11,18 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::CachedId; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::SegmentScan; -use crate::plan::check_child_count; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; -use crate::plan::pipeline::ops::MaskStage; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::SegmentScan; +use crate::plan::physical::check_child_count; +use crate::plan::physical::operators::MaskStage; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; use crate::segments::SegmentId; /// Keeps only the selected rows of its child. diff --git a/vortex-layout/src/plan/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs similarity index 96% rename from vortex-layout/src/plan/plans/list_pack.rs rename to vortex-layout/src/plan/physical/plans/list_pack.rs index d9a23de818b..bb075c7f4e2 100644 --- a/vortex-layout/src/plan/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -12,12 +12,12 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::registry::CachedId; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; const ELEMENTS: usize = 0; const OFFSETS: usize = 1; diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs similarity index 100% rename from vortex-layout/src/plan/plans/mod.rs rename to vortex-layout/src/plan/physical/plans/mod.rs diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs similarity index 97% rename from vortex-layout/src/plan/plans/pack.rs rename to vortex-layout/src/plan/physical/plans/pack.rs index 89f786ba7fb..825897e33ee 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -31,22 +31,22 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::CachedId; -use crate::plan::Eval; -use crate::plan::EvalPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::optimizer::PlanParentReduceRule; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; -use crate::plan::pipeline::ops::OnceSource; -use crate::plan::pipeline::ops::PackSource; -use crate::plan::pipeline::ops::WrapStage; -use crate::plan::pipeline::ops::empty_struct; +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::operators::OnceSource; +use crate::plan::physical::operators::PackSource; +use crate::plan::physical::operators::WrapStage; +use crate::plan::physical::operators::empty_struct; +use crate::plan::physical::optimizer::PlanParentReduceRule; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; use crate::segments::SegmentId; /// Assembles a struct from one child per field, plus an optional trailing validity child. diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs similarity index 95% rename from vortex-layout/src/plan/plans/row_idx.rs rename to vortex-layout/src/plan/physical/plans/row_idx.rs index 6b1edbbc433..90fb7a58a85 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -22,16 +22,16 @@ use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::layouts::row_idx::RowIdx as RowIdxFn; -use crate::plan::EvalPlan; -use crate::plan::PackPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::plans::pack::rewrite_partition_root; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::PackPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::plans::pack::rewrite_partition_root; const ROW_IDX_PARTITION_NAME: &str = "row_idx"; const CHILD_PARTITION_NAME: &str = "child"; diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/physical/plans/segment_scan.rs similarity index 89% rename from vortex-layout/src/plan/plans/segment_scan.rs rename to vortex-layout/src/plan/physical/plans/segment_scan.rs index 9ace7e43ff9..e75181f53cc 100644 --- a/vortex-layout/src/plan/plans/segment_scan.rs +++ b/vortex-layout/src/plan/physical/plans/segment_scan.rs @@ -11,15 +11,15 @@ use vortex_mask::Mask; use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; use crate::segments::SegmentId; /// Reads one serialized array segment. diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/physical/plans/take.rs similarity index 92% rename from vortex-layout/src/plan/plans/take.rs rename to vortex-layout/src/plan/physical/plans/take.rs index 95919090f3c..01ddd4adfd6 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/physical/plans/take.rs @@ -10,16 +10,16 @@ use vortex_array::expr::label_bound_tree; use vortex_error::VortexResult; use vortex_session::registry::CachedId; -use crate::plan::Eval; -use crate::plan::EvalPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::optimizer::PlanParentReduceRule; +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::optimizer::PlanParentReduceRule; const CODES: usize = 0; const VALUES: usize = 1; diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/physical/tests.rs similarity index 100% rename from vortex-layout/src/plan/tests.rs rename to vortex-layout/src/plan/physical/tests.rs diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/physical/typed.rs similarity index 97% rename from vortex-layout/src/plan/typed.rs rename to vortex-layout/src/plan/physical/typed.rs index bd00ebffc9c..20c2f76102a 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/physical/typed.rs @@ -20,13 +20,13 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::Mask; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanVTable; -use crate::plan::display::PlanTreeDisplay; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::display::PlanTreeDisplay; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; use crate::segments::SegmentId; /// The combined allocation behind [`PlanRef`]. diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/physical/vtable.rs similarity index 94% rename from vortex-layout/src/plan/vtable.rs rename to vortex-layout/src/plan/physical/vtable.rs index 74cf44bfcd4..328c7436a6e 100644 --- a/vortex-layout/src/plan/vtable.rs +++ b/vortex-layout/src/plan/physical/vtable.rs @@ -13,11 +13,11 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use vortex_session::registry::Id; -use crate::plan::PlanChildren; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; -use crate::plan::typed::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; +use crate::plan::physical::typed::Plan; use crate::segments::SegmentId; /// A unique identifier for a plan operator. @@ -35,7 +35,7 @@ pub type PlanId = Id; pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { /// Operator-specific data, excluding children. /// - /// Children belong in [`PlanParts::children`](crate::plan::PlanParts::children) so that + /// Children belong in [`PlanParts::children`](crate::plan::physical::PlanParts::children) so that /// traversal and rewriting can discover them generically. type PlanData: 'static + Send + Sync + Clone + Debug; From 128482d2c8fbfa490e39b19339ac462a5a7c5b62 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 12 Oct 2026 15:40:37 +0100 Subject: [PATCH 18/20] refactor(layout): organize physical plans and execution modules Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/mod.rs | 73 +---------------- .../src/plan/{ => physical}/children.rs | 2 +- .../src/plan/{ => physical}/display.rs | 0 .../src/plan/{ => physical}/lower.rs | 14 ++-- vortex-layout/src/plan/physical/mod.rs | 80 +++++++++++++++++++ .../{pipeline/mod.rs => physical/operator.rs} | 72 +++++------------ .../src/plan/physical/operators/mod.rs | 4 + .../src/plan/{ => physical}/optimize.rs | 6 +- .../src/plan/{ => physical}/optimizer/mod.rs | 0 .../plan/{ => physical}/optimizer/rules.rs | 6 +- .../plan/{ => physical}/pipeline/compile.rs | 8 +- .../src/plan/physical/pipeline/mod.rs | 43 ++++++++++ .../src/plan/{ => physical}/pipeline/port.rs | 2 +- .../src/plan/{ => physical}/pipeline/scan.rs | 18 ++--- .../{ => physical}/pipeline/scan/validate.rs | 0 .../pipeline/scheduling_tests.rs | 9 ++- .../plan/{ => physical}/pipeline/synthetic.rs | 26 +++--- .../src/plan/{ => physical}/plans/concat.rs | 18 ++--- .../src/plan/{ => physical}/plans/eval.rs | 16 ++-- .../plan/{ => physical}/plans/list_pack.rs | 12 +-- .../src/plan/{ => physical}/plans/mod.rs | 0 .../src/plan/{ => physical}/plans/pack.rs | 18 ++--- .../src/plan/{ => physical}/plans/row_idx.rs | 20 ++--- .../plan/{ => physical}/plans/segment_scan.rs | 12 +-- .../src/plan/{ => physical}/plans/take.rs | 20 ++--- .../src/plan/{ => physical}/tests.rs | 0 .../src/plan/{ => physical}/typed.rs | 14 ++-- .../src/plan/{ => physical}/vtable.rs | 12 +-- 28 files changed, 268 insertions(+), 237 deletions(-) rename vortex-layout/src/plan/{ => physical}/children.rs (98%) rename vortex-layout/src/plan/{ => physical}/display.rs (100%) rename vortex-layout/src/plan/{ => physical}/lower.rs (95%) create mode 100644 vortex-layout/src/plan/physical/mod.rs rename vortex-layout/src/plan/{pipeline/mod.rs => physical/operator.rs} (74%) create mode 100644 vortex-layout/src/plan/physical/operators/mod.rs rename vortex-layout/src/plan/{ => physical}/optimize.rs (90%) rename vortex-layout/src/plan/{ => physical}/optimizer/mod.rs (100%) rename vortex-layout/src/plan/{ => physical}/optimizer/rules.rs (98%) rename vortex-layout/src/plan/{ => physical}/pipeline/compile.rs (95%) create mode 100644 vortex-layout/src/plan/physical/pipeline/mod.rs rename vortex-layout/src/plan/{ => physical}/pipeline/port.rs (99%) rename vortex-layout/src/plan/{ => physical}/pipeline/scan.rs (98%) rename vortex-layout/src/plan/{ => physical}/pipeline/scan/validate.rs (100%) rename vortex-layout/src/plan/{ => physical}/pipeline/scheduling_tests.rs (98%) rename vortex-layout/src/plan/{ => physical}/pipeline/synthetic.rs (94%) rename vortex-layout/src/plan/{ => physical}/plans/concat.rs (93%) rename vortex-layout/src/plan/{ => physical}/plans/eval.rs (92%) rename vortex-layout/src/plan/{ => physical}/plans/list_pack.rs (96%) rename vortex-layout/src/plan/{ => physical}/plans/mod.rs (100%) rename vortex-layout/src/plan/{ => physical}/plans/pack.rs (98%) rename vortex-layout/src/plan/{ => physical}/plans/row_idx.rs (95%) rename vortex-layout/src/plan/{ => physical}/plans/segment_scan.rs (91%) rename vortex-layout/src/plan/{ => physical}/plans/take.rs (92%) rename vortex-layout/src/plan/{ => physical}/tests.rs (100%) rename vortex-layout/src/plan/{ => physical}/typed.rs (97%) rename vortex-layout/src/plan/{ => physical}/vtable.rs (94%) diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 92ece0c38dd..88c104f81d3 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -1,75 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Physical plans for scans. -//! -//! A plan is a tree of physical operators over a row domain. Operator identity and -//! operator-specific state do not depend on the source layout kind, so rewrites can reason about a -//! plan's shape alone. The common child container can initialize individual slots lazily. +//! Scan planning. -mod children; -mod display; -mod lower; -mod optimize; -pub mod optimizer; -pub mod pipeline; -mod plans; -mod typed; -mod vtable; - -pub use children::PlanChildren; -pub use display::PlanIndentedFormatter; -pub use display::PlanSummaryExtractor; -pub use display::PlanTreeContext; -pub use display::PlanTreeDisplay; -pub use display::PlanTreeExtractor; -pub use lower::lower; -pub use optimize::optimize; -pub use plans::Concat; -pub use plans::ConcatData; -pub use plans::ConcatPlan; -pub use plans::Eval; -pub use plans::EvalData; -pub use plans::EvalPlan; -pub use plans::ListPack; -pub use plans::ListPackData; -pub use plans::ListPackPlan; -pub use plans::Pack; -pub use plans::PackData; -pub use plans::PackPlan; -pub use plans::RowIdx; -pub use plans::RowIdxData; -pub use plans::RowIdxPlan; -pub use plans::SegmentScan; -pub use plans::SegmentScanData; -pub use plans::SegmentScanPlan; -pub use plans::Take; -pub use plans::TakePlan; -pub use plans::plan_row_idx_expression; -pub use plans::row_idx_dtype; -pub use typed::DynPlan; -pub use typed::Plan; -pub use typed::PlanParts; -pub use typed::PlanRef; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -pub use vtable::PlanId; -pub use vtable::PlanVTable; - -/// Returns an error when `children` does not have exactly `expected` entries. -pub(crate) fn check_child_count( - name: &str, - children: &PlanChildren, - expected: usize, -) -> VortexResult<()> { - if children.len() != expected { - vortex_bail!( - "{name} expects {expected} children but got {}", - children.len() - ); - } - Ok(()) -} - -#[cfg(test)] -mod tests; +pub mod physical; diff --git a/vortex-layout/src/plan/children.rs b/vortex-layout/src/plan/physical/children.rs similarity index 98% rename from vortex-layout/src/plan/children.rs rename to vortex-layout/src/plan/physical/children.rs index b045c1be8d2..8b1475e2a98 100644 --- a/vortex-layout/src/plan/children.rs +++ b/vortex-layout/src/plan/physical/children.rs @@ -8,7 +8,7 @@ use once_cell::sync::OnceCell; use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; type ChildInitializer = dyn Fn(usize) -> VortexResult + 'static + Send + Sync; diff --git a/vortex-layout/src/plan/display.rs b/vortex-layout/src/plan/physical/display.rs similarity index 100% rename from vortex-layout/src/plan/display.rs rename to vortex-layout/src/plan/physical/display.rs diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/physical/lower.rs similarity index 95% rename from vortex-layout/src/plan/lower.rs rename to vortex-layout/src/plan/physical/lower.rs index e8911734549..bade3544384 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/physical/lower.rs @@ -24,13 +24,13 @@ use crate::layouts::list::OFFSETS_CHILD_INDEX; use crate::layouts::list::VALIDITY_CHILD_INDEX; use crate::layouts::struct_::Struct; use crate::layouts::struct_::StructLayout; -use crate::plan::ConcatPlan; -use crate::plan::ListPackPlan; -use crate::plan::PackPlan; -use crate::plan::PlanChildren; -use crate::plan::PlanRef; -use crate::plan::SegmentScanPlan; -use crate::plan::TakePlan; +use crate::plan::physical::ConcatPlan; +use crate::plan::physical::ListPackPlan; +use crate::plan::physical::PackPlan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScanPlan; +use crate::plan::physical::TakePlan; /// Constructs a physical-plan fixture from `layout` for tests. /// diff --git a/vortex-layout/src/plan/physical/mod.rs b/vortex-layout/src/plan/physical/mod.rs new file mode 100644 index 00000000000..5378d50d59f --- /dev/null +++ b/vortex-layout/src/plan/physical/mod.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Physical plans for scans. +//! +//! A plan is a tree of physical operators over a row domain. Operator identity and +//! operator-specific state do not depend on the source layout kind, so rewrites can reason about a +//! plan's shape alone. The common child container can initialize individual slots lazily. +//! +//! [`operator`] defines the execution contracts, `operators` implements the concrete sources +//! and stages, and [`pipeline`] compiles and runs them. + +mod children; +mod display; +mod lower; +pub mod operator; +pub(crate) mod operators; +mod optimize; +pub mod optimizer; +pub mod pipeline; +mod plans; +mod typed; +mod vtable; + +pub use children::PlanChildren; +pub use display::PlanIndentedFormatter; +pub use display::PlanSummaryExtractor; +pub use display::PlanTreeContext; +pub use display::PlanTreeDisplay; +pub use display::PlanTreeExtractor; +pub use lower::lower; +pub use optimize::optimize; +pub use plans::Concat; +pub use plans::ConcatData; +pub use plans::ConcatPlan; +pub use plans::Eval; +pub use plans::EvalData; +pub use plans::EvalPlan; +pub use plans::ListPack; +pub use plans::ListPackData; +pub use plans::ListPackPlan; +pub use plans::Pack; +pub use plans::PackData; +pub use plans::PackPlan; +pub use plans::RowIdx; +pub use plans::RowIdxData; +pub use plans::RowIdxPlan; +pub use plans::SegmentScan; +pub use plans::SegmentScanData; +pub use plans::SegmentScanPlan; +pub use plans::Take; +pub use plans::TakePlan; +pub use plans::plan_row_idx_expression; +pub use plans::row_idx_dtype; +pub use typed::DynPlan; +pub use typed::Plan; +pub use typed::PlanParts; +pub use typed::PlanRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +pub use vtable::PlanId; +pub use vtable::PlanVTable; + +/// Returns an error when `children` does not have exactly `expected` entries. +pub(crate) fn check_child_count( + name: &str, + children: &PlanChildren, + expected: usize, +) -> VortexResult<()> { + if children.len() != expected { + vortex_bail!( + "{name} expects {expected} children but got {}", + children.len() + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-layout/src/plan/pipeline/mod.rs b/vortex-layout/src/plan/physical/operator.rs similarity index 74% rename from vortex-layout/src/plan/pipeline/mod.rs rename to vortex-layout/src/plan/physical/operator.rs index 3d3c51e6c00..83d7c81fefd 100644 --- a/vortex-layout/src/plan/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -1,26 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Push-based execution of physical plans as pipelines. -//! -//! A [`Scan`] runs a [`PlanRef`](crate::plan::PlanRef) over a list of splits, row ranges of the -//! plan's domain. Each split is compiled into pipelines: a pipeline is one [`Source`], zero or -//! more [`Operator`]s, and the ports its last stage writes. The driver asks the source for a batch -//! and pushes it through the stages into the ports. Every port runs from one pipeline's sink, -//! the end that pushes into its outlets, to another's source, which reads it as an inlet, and -//! the plan joining them declares how many batches each inlet may hold. -//! -//! A source checks its readiness before computing, waiting for a segment read or an inlet. -//! The driver also blocks a pipeline when an outlet is full. The scheduler re-runs a blocked -//! pipeline when that one condition clears. -//! -//! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and -//! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. - -mod compile; -mod port; -mod scan; -pub mod synthetic; +//! Contracts and execution context for pipeline sources and operators. use std::collections::VecDeque; @@ -31,23 +12,11 @@ use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; use vortex_session::VortexSession; -pub use self::compile::Chain; -pub use self::compile::Compiler; -pub use self::compile::Reach; -use self::port::Arena; -pub use self::port::Inlet; -use self::port::PortId; -pub use self::scan::ReadId; -pub use self::scan::ReadRequest; -pub use self::scan::Scan; -pub use self::scan::Split; -pub use self::scan::Turn; +use super::pipeline::Arena; +use super::pipeline::Inlet; +use super::pipeline::PortId; use crate::segments::SegmentId; -/// Batches an inlet holds before its sink is blocked, unless the plan joining it says -/// otherwise; see [`Compiler::join_with`]. -pub const DEFAULT_CAPACITY: usize = 8; - /// The output of one non-blocking compute call. #[derive(Debug)] pub enum Step { @@ -121,34 +90,34 @@ pub trait Operator: Send { /// Distinct inlet indices, with constant-time insertion and removal. #[derive(Default)] -struct InletSet { +pub(super) struct InletSet { indices: Vec, positions: Vec, } impl InletSet { - fn new(inlets: usize) -> Self { + pub(super) fn new(inlets: usize) -> Self { Self { indices: Vec::with_capacity(inlets), positions: vec![usize::MAX; inlets], } } - fn all(inlets: usize) -> Self { + pub(super) fn all(inlets: usize) -> Self { Self { indices: (0..inlets).collect(), positions: (0..inlets).collect(), } } - fn insert(&mut self, index: usize) { + pub(super) fn insert(&mut self, index: usize) { if self.positions[index] == usize::MAX { self.positions[index] = self.indices.len(); self.indices.push(index); } } - fn remove(&mut self, index: usize) { + pub(super) fn remove(&mut self, index: usize) { let position = self.positions[index]; if position == usize::MAX { return; @@ -166,11 +135,11 @@ impl InletSet { Some(index) } - fn indices(&self) -> &[usize] { + pub(super) fn indices(&self) -> &[usize] { &self.indices } - fn clear(&mut self) { + pub(super) fn clear(&mut self) { for &index in &self.indices { self.positions[index] = usize::MAX; } @@ -181,14 +150,14 @@ impl InletSet { /// What a source or stage may touch during one call: its pipeline's inlets, the bytes the /// source asked for, and the session. pub struct Cx<'a> { - arena: &'a mut Arena, - inlets: &'a [PortId], - bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, - reads: SmallVec<[SegmentId; 1]>, - session: &'a VortexSession, - exec: &'a mut ExecutionCtx, - touched: &'a mut InletSet, - changed: &'a mut InletSet, + pub(super) arena: &'a mut Arena, + pub(super) inlets: &'a [PortId], + pub(super) bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, + pub(super) reads: SmallVec<[SegmentId; 1]>, + pub(super) session: &'a VortexSession, + pub(super) exec: &'a mut ExecutionCtx, + pub(super) touched: &'a mut InletSet, + pub(super) changed: &'a mut InletSet, } impl Cx<'_> { @@ -232,6 +201,3 @@ impl Cx<'_> { self.exec } } - -#[cfg(test)] -mod scheduling_tests; diff --git a/vortex-layout/src/plan/physical/operators/mod.rs b/vortex-layout/src/plan/physical/operators/mod.rs new file mode 100644 index 00000000000..bda7959899b --- /dev/null +++ b/vortex-layout/src/plan/physical/operators/mod.rs @@ -0,0 +1,4 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Concrete sources and stages for physical plans. diff --git a/vortex-layout/src/plan/optimize.rs b/vortex-layout/src/plan/physical/optimize.rs similarity index 90% rename from vortex-layout/src/plan/optimize.rs rename to vortex-layout/src/plan/physical/optimize.rs index f471b7744f9..32023f0ee05 100644 --- a/vortex-layout/src/plan/optimize.rs +++ b/vortex-layout/src/plan/physical/optimize.rs @@ -8,9 +8,9 @@ use vortex_error::VortexResult; -use crate::plan::PlanRef; -use crate::plan::optimizer::reduce_parent; -use crate::plan::optimizer::reduce_plan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::optimizer::reduce_parent; +use crate::plan::physical::optimizer::reduce_plan; fn reduce(plan: &PlanRef) -> VortexResult> { if let Some(rewritten) = reduce_plan(plan)? { diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs similarity index 100% rename from vortex-layout/src/plan/optimizer/mod.rs rename to vortex-layout/src/plan/physical/optimizer/mod.rs diff --git a/vortex-layout/src/plan/optimizer/rules.rs b/vortex-layout/src/plan/physical/optimizer/rules.rs similarity index 98% rename from vortex-layout/src/plan/optimizer/rules.rs rename to vortex-layout/src/plan/physical/optimizer/rules.rs index 12490201ee5..1c32e9c6f60 100644 --- a/vortex-layout/src/plan/optimizer/rules.rs +++ b/vortex-layout/src/plan/physical/optimizer/rules.rs @@ -11,9 +11,9 @@ use vortex_error::VortexResult; #[cfg(debug_assertions)] use vortex_error::vortex_ensure_eq; -use crate::plan::Plan; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; /// A rewrite over one concrete plan operator. /// diff --git a/vortex-layout/src/plan/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs similarity index 95% rename from vortex-layout/src/plan/pipeline/compile.rs rename to vortex-layout/src/plan/physical/pipeline/compile.rs index ed39688a109..21c68d9e34f 100644 --- a/vortex-layout/src/plan/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -11,12 +11,12 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use super::DEFAULT_CAPACITY; -use super::Operator; -use super::Source; use super::port::PortId; use super::port::SourceEnd; use super::scan::Core; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; /// A pipeline under construction: a source and the stages after it, not yet given outlets, so a /// parent plan can add stages before it becomes a pipeline. @@ -56,7 +56,7 @@ impl Core { } /// What a plan compiles its children through. See -/// [`PlanVTable::compile`](crate::plan::PlanVTable::compile). +/// [`PlanVTable::compile`](crate::plan::physical::PlanVTable::compile). pub struct Compiler<'a> { pub(crate) core: &'a mut Core, } diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs new file mode 100644 index 00000000000..ff6565ee2f6 --- /dev/null +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Push-based execution of physical plans as pipelines. +//! +//! A [`Scan`] runs a [`PlanRef`](crate::plan::physical::PlanRef) over a list of splits, row ranges +//! of the plan's domain. Each split is compiled into pipelines: a pipeline is one +//! [`Source`](super::operator::Source), zero or more [`Operator`](super::operator::Operator)s, +//! and the ports its last stage writes. The driver asks the source for a batch and pushes it +//! through the stages into the ports. Every port runs from one pipeline's sink, +//! the end that pushes into its outlets, to another's source, which reads it as an inlet, and +//! the plan joining them declares how many batches each inlet may hold. +//! +//! A source checks its readiness before computing, waiting for a segment read or an inlet. +//! The driver also blocks a pipeline when an outlet is full. The scheduler re-runs a blocked +//! pipeline when that one condition clears. +//! +//! The owner of a scan answers reads: [`Scan::step`] returns each read as a [`Turn::Read`], and +//! [`Scan::deliver`] hands its bytes back. Nothing in a scan awaits. + +mod compile; +mod port; +mod scan; +pub mod synthetic; + +pub use self::compile::Chain; +pub use self::compile::Compiler; +pub use self::compile::Reach; +pub(super) use self::port::Arena; +pub use self::port::Inlet; +pub(super) use self::port::PortId; +pub use self::scan::ReadId; +pub use self::scan::ReadRequest; +pub use self::scan::Scan; +pub use self::scan::Split; +pub use self::scan::Turn; + +/// Batches an inlet holds before its sink is blocked, unless the plan joining it says +/// otherwise; see [`Compiler::join_with`]. +pub const DEFAULT_CAPACITY: usize = 8; + +#[cfg(test)] +mod scheduling_tests; diff --git a/vortex-layout/src/plan/pipeline/port.rs b/vortex-layout/src/plan/physical/pipeline/port.rs similarity index 99% rename from vortex-layout/src/plan/pipeline/port.rs rename to vortex-layout/src/plan/physical/pipeline/port.rs index 38f9ae143f1..039136cf858 100644 --- a/vortex-layout/src/plan/pipeline/port.rs +++ b/vortex-layout/src/plan/physical/pipeline/port.rs @@ -12,7 +12,7 @@ use std::collections::VecDeque; use vortex_array::ArrayRef; -use crate::plan::pipeline::DEFAULT_CAPACITY; +use crate::plan::physical::pipeline::DEFAULT_CAPACITY; /// Identifies a port in the scan's arena. pub(crate) type PortId = usize; diff --git a/vortex-layout/src/plan/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs similarity index 98% rename from vortex-layout/src/plan/pipeline/scan.rs rename to vortex-layout/src/plan/physical/pipeline/scan.rs index 5651eac6430..5ba165fdd38 100644 --- a/vortex-layout/src/plan/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -21,22 +21,22 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use super::Cx; use super::DEFAULT_CAPACITY; -use super::InletSet; -use super::Input; -use super::Operator; -use super::Source; -use super::SourceBlocked; -use super::SourceReady; -use super::Step; use super::compile::Chain; use super::port::Arena; use super::port::PipelineId; use super::port::PortId; use super::port::Slab; use super::port::SourceEnd; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::InletSet; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; #[cfg(debug_assertions)] diff --git a/vortex-layout/src/plan/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs similarity index 100% rename from vortex-layout/src/plan/pipeline/scan/validate.rs rename to vortex-layout/src/plan/physical/pipeline/scan/validate.rs diff --git a/vortex-layout/src/plan/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs similarity index 98% rename from vortex-layout/src/plan/pipeline/scheduling_tests.rs rename to vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index 5c6597df460..28de94f2209 100644 --- a/vortex-layout/src/plan/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -26,7 +26,14 @@ use super::synthetic::SyntheticPlan; use super::synthetic::empty_segment; use super::synthetic::row_dtype; use super::*; -use crate::plan::PlanRef; +use crate::plan::physical::PlanRef; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; use crate::test::SESSION; diff --git a/vortex-layout/src/plan/pipeline/synthetic.rs b/vortex-layout/src/plan/physical/pipeline/synthetic.rs similarity index 94% rename from vortex-layout/src/plan/pipeline/synthetic.rs rename to vortex-layout/src/plan/physical/pipeline/synthetic.rs index 6fd70667678..b6046c16c21 100644 --- a/vortex-layout/src/plan/pipeline/synthetic.rs +++ b/vortex-layout/src/plan/physical/pipeline/synthetic.rs @@ -26,19 +26,19 @@ use vortex_session::registry::CachedId; use super::Chain; use super::Compiler; -use super::Cx; -use super::Input; -use super::Operator; -use super::Source; -use super::SourceBlocked; -use super::SourceReady; -use super::Step; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::operator::Cx; +use crate::plan::physical::operator::Input; +use crate::plan::physical::operator::Operator; +use crate::plan::physical::operator::Source; +use crate::plan::physical::operator::SourceBlocked; +use crate::plan::physical::operator::SourceReady; +use crate::plan::physical::operator::Step; use crate::segments::SegmentId; /// Compiles a [`SyntheticPlan`] over some rows, restricted to a mask. diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs similarity index 93% rename from vortex-layout/src/plan/plans/concat.rs rename to vortex-layout/src/plan/physical/plans/concat.rs index 2e814156309..37b002bcc93 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -10,15 +10,15 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::plan::Eval; -use crate::plan::EvalPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::optimizer::PlanParentReduceRule; +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::optimizer::PlanParentReduceRule; /// Concatenates its children row-wise. #[derive(Clone, Debug)] diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs similarity index 92% rename from vortex-layout/src/plan/plans/eval.rs rename to vortex-layout/src/plan/physical/plans/eval.rs index f1edb967e41..04281bdfefc 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -10,14 +10,14 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::optimizer::PlanReduceRule; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::optimizer::PlanReduceRule; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] diff --git a/vortex-layout/src/plan/plans/list_pack.rs b/vortex-layout/src/plan/physical/plans/list_pack.rs similarity index 96% rename from vortex-layout/src/plan/plans/list_pack.rs rename to vortex-layout/src/plan/physical/plans/list_pack.rs index d9a23de818b..bb075c7f4e2 100644 --- a/vortex-layout/src/plan/plans/list_pack.rs +++ b/vortex-layout/src/plan/physical/plans/list_pack.rs @@ -12,12 +12,12 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::registry::CachedId; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; const ELEMENTS: usize = 0; const OFFSETS: usize = 1; diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs similarity index 100% rename from vortex-layout/src/plan/plans/mod.rs rename to vortex-layout/src/plan/physical/plans/mod.rs diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs similarity index 98% rename from vortex-layout/src/plan/plans/pack.rs rename to vortex-layout/src/plan/physical/plans/pack.rs index 6497cc63d45..20b55874e01 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -29,15 +29,15 @@ 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::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::optimizer::PlanParentReduceRule; +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::optimizer::PlanParentReduceRule; /// Assembles a struct from one child per field, plus an optional trailing validity child. #[derive(Clone, Debug)] diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs similarity index 95% rename from vortex-layout/src/plan/plans/row_idx.rs rename to vortex-layout/src/plan/physical/plans/row_idx.rs index 6b1edbbc433..90fb7a58a85 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -22,16 +22,16 @@ use vortex_error::vortex_err; use vortex_session::registry::CachedId; use crate::layouts::row_idx::RowIdx as RowIdxFn; -use crate::plan::EvalPlan; -use crate::plan::PackPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::plans::pack::rewrite_partition_root; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::PackPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::plans::pack::rewrite_partition_root; const ROW_IDX_PARTITION_NAME: &str = "row_idx"; const CHILD_PARTITION_NAME: &str = "child"; diff --git a/vortex-layout/src/plan/plans/segment_scan.rs b/vortex-layout/src/plan/physical/plans/segment_scan.rs similarity index 91% rename from vortex-layout/src/plan/plans/segment_scan.rs rename to vortex-layout/src/plan/physical/plans/segment_scan.rs index d2b20df89ad..37048ff4ea9 100644 --- a/vortex-layout/src/plan/plans/segment_scan.rs +++ b/vortex-layout/src/plan/physical/plans/segment_scan.rs @@ -8,12 +8,12 @@ use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; use crate::segments::SegmentId; /// Reads one serialized array segment. diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/physical/plans/take.rs similarity index 92% rename from vortex-layout/src/plan/plans/take.rs rename to vortex-layout/src/plan/physical/plans/take.rs index 95919090f3c..01ddd4adfd6 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/physical/plans/take.rs @@ -10,16 +10,16 @@ use vortex_array::expr::label_bound_tree; use vortex_error::VortexResult; use vortex_session::registry::CachedId; -use crate::plan::Eval; -use crate::plan::EvalPlan; -use crate::plan::Plan; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanParts; -use crate::plan::PlanRef; -use crate::plan::PlanVTable; -use crate::plan::check_child_count; -use crate::plan::optimizer::PlanParentReduceRule; +use crate::plan::physical::Eval; +use crate::plan::physical::EvalPlan; +use crate::plan::physical::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanParts; +use crate::plan::physical::PlanRef; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::check_child_count; +use crate::plan::physical::optimizer::PlanParentReduceRule; const CODES: usize = 0; const VALUES: usize = 1; diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/physical/tests.rs similarity index 100% rename from vortex-layout/src/plan/tests.rs rename to vortex-layout/src/plan/physical/tests.rs diff --git a/vortex-layout/src/plan/typed.rs b/vortex-layout/src/plan/physical/typed.rs similarity index 97% rename from vortex-layout/src/plan/typed.rs rename to vortex-layout/src/plan/physical/typed.rs index bd00ebffc9c..20c2f76102a 100644 --- a/vortex-layout/src/plan/typed.rs +++ b/vortex-layout/src/plan/physical/typed.rs @@ -20,13 +20,13 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::Mask; -use crate::plan::PlanChildren; -use crate::plan::PlanId; -use crate::plan::PlanVTable; -use crate::plan::display::PlanTreeDisplay; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::PlanId; +use crate::plan::physical::PlanVTable; +use crate::plan::physical::display::PlanTreeDisplay; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; use crate::segments::SegmentId; /// The combined allocation behind [`PlanRef`]. diff --git a/vortex-layout/src/plan/vtable.rs b/vortex-layout/src/plan/physical/vtable.rs similarity index 94% rename from vortex-layout/src/plan/vtable.rs rename to vortex-layout/src/plan/physical/vtable.rs index 74cf44bfcd4..328c7436a6e 100644 --- a/vortex-layout/src/plan/vtable.rs +++ b/vortex-layout/src/plan/physical/vtable.rs @@ -13,11 +13,11 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use vortex_session::registry::Id; -use crate::plan::PlanChildren; -use crate::plan::pipeline::Chain; -use crate::plan::pipeline::Compiler; -use crate::plan::pipeline::Reach; -use crate::plan::typed::Plan; +use crate::plan::physical::PlanChildren; +use crate::plan::physical::pipeline::Chain; +use crate::plan::physical::pipeline::Compiler; +use crate::plan::physical::pipeline::Reach; +use crate::plan::physical::typed::Plan; use crate::segments::SegmentId; /// A unique identifier for a plan operator. @@ -35,7 +35,7 @@ pub type PlanId = Id; pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { /// Operator-specific data, excluding children. /// - /// Children belong in [`PlanParts::children`](crate::plan::PlanParts::children) so that + /// Children belong in [`PlanParts::children`](crate::plan::physical::PlanParts::children) so that /// traversal and rewriting can discover them generically. type PlanData: 'static + Send + Sync + Clone + Debug; From d0cde1c30c9f79de785f4d0ed355e44678dbd8ba Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 12 Oct 2026 15:42:47 +0100 Subject: [PATCH 19/20] fix(layout): import physical plan types in pipeline tests Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/physical/pipeline/tests.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 1b2e5dca722..2949edb2d12 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -38,6 +38,8 @@ use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::struct_::StructLayout; +use crate::plan::physical::FilterPlan; +use crate::plan::physical::PackPlan; use crate::plan::physical::PlanRef; use crate::plan::physical::SegmentScan; use crate::plan::physical::lower; @@ -559,7 +561,7 @@ fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> Vortex let values = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); let scan = lower(&store.flat(&values)?)?; assert!(scan.is::()); - let filtered = crate::plan::physical::FilterPlan::new(scan.clone()).into_plan(); + let filtered = FilterPlan::new(scan.clone()).into_plan(); let rows = 3..9; let mask = sel.mask(6); @@ -591,7 +593,7 @@ fn a_segment_two_readers_need_is_read_once() -> VortexResult<()> { .as_struct_fields_opt() .vortex_expect("struct") .clone(); - let plan = crate::plan::physical::PackPlan::try_new( + let plan = PackPlan::try_new( fields, vortex_array::dtype::Nullability::NonNullable, ROWS, From b966c279524c1fbef673bf0ae2a5759bd9b17840 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 12 Oct 2026 16:03:15 +0100 Subject: [PATCH 20/20] refactor(layout): use blocked IO for queued and outstanding reads Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/physical/operator.rs | 18 +++++----- .../src/plan/physical/pipeline/scan.rs | 34 +++++++++---------- .../plan/physical/pipeline/scan/validate.rs | 6 ++-- .../physical/pipeline/scheduling_tests.rs | 10 ++++-- .../src/plan/physical/pipeline/synthetic.rs | 2 +- 5 files changed, 36 insertions(+), 34 deletions(-) diff --git a/vortex-layout/src/plan/physical/operator.rs b/vortex-layout/src/plan/physical/operator.rs index 83d7c81fefd..96a8d0cfbd7 100644 --- a/vortex-layout/src/plan/physical/operator.rs +++ b/vortex-layout/src/plan/physical/operator.rs @@ -34,8 +34,6 @@ pub enum Step { pub enum SourceReady { /// The next compute call can produce output or finish without blocking. Ready, - /// New reads were queued through [`Cx::request_read`]. Run again when bytes are available. - Read, /// Nothing can happen until the condition clears. Blocked(SourceBlocked), } @@ -43,7 +41,7 @@ pub enum SourceReady { /// Why a source cannot compute yet. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum SourceBlocked { - /// A read was requested and has not been delivered. + /// Wait for bytes from newly queued or outstanding reads. Io, /// The inlet needs more batches or its sink to close it. It must be open and have room /// for its sink to push; debug builds check that the sink can still wake this source. @@ -57,9 +55,10 @@ pub trait Source: Operator { /// Returning [`SourceReady::Ready`] guarantees that call can proceed without blocking. /// Any input or choice made here can be retained by the source for that call. /// - /// Queue one or more reads with [`Cx::request_read`] and return [`SourceReady::Read`] to - /// issue them. Readiness is checked again when bytes are available, including bytes already - /// delivered; they are available through [`Cx::take_segment`] in arrival order. + /// Queue reads with [`Cx::request_read`] and return + /// `SourceReady::Blocked(SourceBlocked::Io)`. The driver issues newly queued reads + /// before waiting. Readiness is checked again when bytes are available, including bytes + /// already delivered; they are available through [`Cx::take_segment`] in arrival order. fn ready(&mut self, _cx: &mut Cx<'_>) -> VortexResult { Ok(SourceReady::Ready) } @@ -153,7 +152,7 @@ pub struct Cx<'a> { pub(super) arena: &'a mut Arena, pub(super) inlets: &'a [PortId], pub(super) bytes: &'a mut VecDeque<(SegmentId, BufferHandle)>, - pub(super) reads: SmallVec<[SegmentId; 1]>, + pub(super) reads: &'a mut SmallVec<[SegmentId; 1]>, pub(super) session: &'a VortexSession, pub(super) exec: &'a mut ExecutionCtx, pub(super) touched: &'a mut InletSet, @@ -169,8 +168,9 @@ impl Cx<'_> { self.changed.pop() } - /// Queues a segment read during [`Source::ready`]. Return [`SourceReady::Read`] after - /// queuing all the reads needed for this call. + /// Queues a segment read during [`Source::ready`]. Return + /// `SourceReady::Blocked(SourceBlocked::Io)` after queuing the reads needed for + /// this call. pub fn request_read(&mut self, segment: SegmentId) { self.reads.push(segment); } diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 5ba165fdd38..3191df85d69 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -124,8 +124,6 @@ pub(crate) struct Pipeline { /// What one run of a pipeline came to. enum Progress { - /// The source wants these segments read. - Read(SmallVec<[SegmentId; 1]>), /// The source waits for this. Blocked(SourceBlocked), /// An outlet is full. @@ -134,12 +132,13 @@ enum Progress { Done, } -/// Runs `pipe` until it blocks, finishes, or requests a read. +/// Runs `pipe` until it blocks or finishes, collecting any reads its source requests. fn drive( pipe: &mut Pipeline, arena: &mut Arena, session: &VortexSession, exec: &mut ExecutionCtx, + reads: &mut SmallVec<[SegmentId; 1]>, ) -> VortexResult { let Pipeline { source, @@ -158,7 +157,7 @@ fn drive( arena, inlets, bytes, - reads: SmallVec::new(), + reads, session, exec, touched, @@ -190,12 +189,11 @@ fn drive( let ready = source.ready(&mut cx)?; cx.changed.clear(); vortex_ensure!( - cx.reads.is_empty() == (ready != SourceReady::Read), - "A source must return Read exactly when it queues reads" + cx.reads.is_empty() || ready == SourceReady::Blocked(SourceBlocked::Io), + "A source that queues reads must block on IO" ); match ready { SourceReady::Ready => {} - SourceReady::Read => return Ok(Progress::Read(cx.reads)), SourceReady::Blocked(blocked) => return Ok(Progress::Blocked(blocked)), } } @@ -393,20 +391,20 @@ impl Core { .as_mut() .ok_or_else(|| vortex_err!("Pipeline {id} is not live"))?; pipe.queued = false; - let progress = drive(pipe, &mut self.arena, &self.session, &mut self.exec)?; + let mut reads = SmallVec::new(); + let progress = drive( + pipe, + &mut self.arena, + &self.session, + &mut self.exec, + &mut reads, + )?; + for segment_id in reads { + self.ask(id, segment_id); + } #[cfg(debug_assertions)] self.validate_progress(id, &progress); match &progress { - Progress::Read(segments) => { - for &segment_id in segments { - self.ask(id, segment_id); - } - if self.pipeline(id).bytes.is_empty() { - self.pipeline_mut(id).state = State::Blocked(SourceBlocked::Io); - } else { - self.make_runnable(id); - } - } // A read-ahead source can still have delivered bytes to consume. Progress::Blocked(SourceBlocked::Io) if !self.pipeline(id).bytes.is_empty() => { self.make_runnable(id); diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index d1147c4012b..3abe13381d0 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -80,8 +80,8 @@ impl Core { ); } Progress::Blocked(SourceBlocked::Io) => assert!( - pipe.pending > 0, - "pipeline {id} waits for a read with none in flight" + pipe.pending > 0 || !pipe.bytes.is_empty(), + "pipeline {id} waits for a read with none in flight or delivered" ), Progress::Done => { for (index, &port) in pipe.inlets.iter().enumerate() { @@ -93,7 +93,7 @@ impl Core { ); } } - Progress::Read(_) | Progress::Outlet => {} + Progress::Outlet => {} } } } diff --git a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs index 28de94f2209..19bcb5504db 100644 --- a/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/scheduling_tests.rs @@ -218,7 +218,7 @@ impl Source for ManyReads { while let Some(segment) = self.wanted.pop() { cx.request_read(segment); } - return Ok(SourceReady::Read); + return Ok(SourceReady::Blocked(SourceBlocked::Io)); } Ok(if self.left == 0 || cx.has_bytes() { SourceReady::Ready @@ -267,6 +267,10 @@ fn a_source_has_several_reads_in_flight(#[values(1, 3)] deliveries: usize) -> Vo ); first_wait = false; } + assert!( + matches!(scan.step()?, Turn::Waiting), + "waiting for IO must not reissue reads or compute before bytes arrive" + ); // Deliver the last asked for first. for _ in 0..deliveries { let read = inflight @@ -305,7 +309,7 @@ impl Source for ReadAfterInlet { }; cx.request_read(SegmentId::from(u32::try_from(batch.len())?)); self.batch = Some(batch); - Ok(SourceReady::Read) + Ok(SourceReady::Blocked(SourceBlocked::Io)) } } @@ -370,7 +374,7 @@ impl Source for FirstOf { while let Some(segment) = self.wanted.pop() { cx.request_read(segment); } - return Ok(SourceReady::Read); + return Ok(SourceReady::Blocked(SourceBlocked::Io)); } Ok(if cx.has_bytes() { SourceReady::Ready diff --git a/vortex-layout/src/plan/physical/pipeline/synthetic.rs b/vortex-layout/src/plan/physical/pipeline/synthetic.rs index b6046c16c21..82c4f5bc3f6 100644 --- a/vortex-layout/src/plan/physical/pipeline/synthetic.rs +++ b/vortex-layout/src/plan/physical/pipeline/synthetic.rs @@ -179,7 +179,7 @@ impl Source for RowSource { if let Some(segment) = self.segment.take() { self.waiting = true; cx.request_read(segment); - return Ok(SourceReady::Read); + return Ok(SourceReady::Blocked(SourceBlocked::Io)); } Ok(if self.waiting && !cx.has_bytes() { SourceReady::Blocked(SourceBlocked::Io)