From 71fd6e56fbf2d005b2598ae904850b421653446c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 5 Aug 2026 10:06:04 +0100 Subject: [PATCH] Add zoned layout scan plan Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/mod.rs | 6 ++ vortex-layout/src/plan/plans/mod.rs | 2 + vortex-layout/src/plan/plans/zoned.rs | 103 ++++++++++++++++++++++++++ vortex-layout/src/plan/tests.rs | 64 ++++++++++++++++ 4 files changed, 175 insertions(+) create mode 100644 vortex-layout/src/plan/plans/zoned.rs diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 64473c3d95f..75ba259b98f 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -28,6 +28,7 @@ pub use plans::RowIdxPartitionPlan; pub use plans::RowIdxPlan; pub use plans::RowIdxValuesPlan; pub use plans::StructPlan; +pub use plans::ZonedPlan; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -38,6 +39,8 @@ use crate::layouts::dict::Dict; use crate::layouts::flat::Flat; use crate::layouts::list::List; use crate::layouts::struct_::Struct; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; /// Shared handle to a heap-allocated physical plan. pub type PlanRef = Arc; @@ -106,6 +109,9 @@ pub fn new_plan(layout: &LayoutRef) -> VortexResult { if let Some(layout) = layout.as_opt::() { return Ok(Arc::new(StructPlan::new(layout))); } + if layout.is::() || layout.is::() { + return Ok(Arc::new(ZonedPlan::try_new(layout)?)); + } vortex_bail!( "No physical plan implementation for layout '{}'", layout.encoding_id() diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index ae1e7a96124..84c924b1dee 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -8,6 +8,7 @@ mod flat; mod list; mod row_idx; mod struct_; +mod zoned; pub use chunked::ChunkedPlan; pub(crate) use chunked::ExpressionChunkedRule; @@ -22,3 +23,4 @@ pub use row_idx::RowIdxPlan; pub use row_idx::RowIdxValuesPlan; pub(crate) use struct_::ExpressionStructRule; pub use struct_::StructPlan; +pub use zoned::ZonedPlan; diff --git a/vortex-layout/src/plan/plans/zoned.rs b/vortex-layout/src/plan/plans/zoned.rs new file mode 100644 index 00000000000..da686a021ea --- /dev/null +++ b/vortex-layout/src/plan/plans/zoned.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::sync::Arc; + +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::LayoutRef; +use crate::plan::Plan; +use crate::plan::PlanRef; +use crate::plan::new_plan; + +const DATA_CHILD_INDEX: usize = 0; +const ZONES_CHILD_INDEX: usize = 1; + +/// A physical zoned plan with a transparent data child and an auxiliary zones child. +/// +/// This plan represents both current `vortex.zoned` layouts and legacy `vortex.stats` layouts, +/// which have the same physical child shape. +pub struct ZonedPlan { + layout: LayoutRef, + dtype: DType, + data: PlanRef, + zones: PlanRef, +} + +impl ZonedPlan { + pub(crate) fn try_new(layout: &LayoutRef) -> VortexResult { + let data = new_plan( + &layout + .slot(DATA_CHILD_INDEX)? + .ok_or_else(|| vortex_err!("Zoned data child is absent"))?, + )?; + let zones = new_plan( + &layout + .slot(ZONES_CHILD_INDEX)? + .ok_or_else(|| vortex_err!("Zoned zones child is absent"))?, + )?; + Ok(Self { + layout: Arc::clone(layout), + dtype: layout.dtype().clone(), + data, + zones, + }) + } + + fn with_children(&self, data: PlanRef, zones: PlanRef) -> Self { + Self { + layout: Arc::clone(&self.layout), + dtype: self.dtype.clone(), + data, + zones, + } + } +} + +impl Plan for ZonedPlan { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn name(&self) -> &'static str { + "ZonedPlan" + } + + fn optimize(&self) -> VortexResult { + let data = self.data.optimize()?; + let zones = self.zones.optimize()?; + Ok(Arc::new(self.with_children(data, zones))) + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.layout.row_count() + } + + fn child_count(&self) -> usize { + 2 + } + + fn child(&self, index: usize) -> VortexResult> { + match index { + DATA_CHILD_INDEX => Ok(Some(Arc::clone(&self.data))), + ZONES_CHILD_INDEX => Ok(Some(Arc::clone(&self.zones))), + _ => vortex_bail!("Zoned plan has no child {index}"), + } + } + + fn child_name(&self, index: usize) -> Cow<'_, str> { + match index { + DATA_CHILD_INDEX => Cow::Borrowed("data"), + ZONES_CHILD_INDEX => Cow::Borrowed("zones"), + _ => Cow::Owned(format!("child[{index}]")), + } + } +} diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 01fa7fcecfb..3dc0b453a54 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::num::NonZeroUsize; use std::sync::Arc; +use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -20,6 +22,8 @@ use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; use super::*; +use crate::LayoutBuildContext; +use crate::LayoutEncoding; use crate::LayoutRef; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; @@ -28,6 +32,8 @@ use crate::layouts::flat::FlatLayout; use crate::layouts::foreign::new_foreign_layout; use crate::layouts::row_idx::row_idx; use crate::layouts::struct_::StructLayout; +use crate::layouts::zoned::LegacyStatsLayoutEncoding; +use crate::layouts::zoned::ZonedLayout; use crate::segments::SegmentId; fn primitive(ptype: PType, nullability: Nullability) -> DType { @@ -53,6 +59,64 @@ fn make_plan(layout: LayoutRef) -> VortexResult { new_plan(&layout) } +#[test] +fn zoned_plan_exposes_data_and_zones() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let zones_dtype = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let aggregate_fns: Arc<[AggregateFnRef]> = Vec::new().into(); + let layout = ZonedLayout::try_new( + flat(5, dtype, 0), + flat(2, zones_dtype, 1), + zone_len, + aggregate_fns, + )? + .into_layout(); + + let plan = make_plan(layout)?; + assert!(plan.as_any().is::()); + insta::assert_snapshot!(plan.tree_display(), @r" + root: ZonedPlan(i32, rows=5) + data: FlatPlan(i32, rows=5) + zones: FlatPlan({}, rows=2) + "); + Ok(()) +} + +#[test] +fn legacy_stats_layout_uses_zoned_plan() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let zones_dtype = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let children = OwnedLayoutChildren::layout_children(vec![ + flat(5, dtype.clone(), 0), + flat(2, zones_dtype, 1), + ]); + let session = vortex_array::array_session(); + let read_ctx = ReadContext::new([]); + let build_ctx = LayoutBuildContext { + session: &session, + array_read_ctx: &read_ctx, + }; + let layout = LayoutEncoding::build( + &LegacyStatsLayoutEncoding, + &dtype, + 5, + &3_u32.to_le_bytes(), + Vec::new(), + children.as_ref(), + &build_ctx, + )?; + + let plan = make_plan(layout)?; + assert!(plan.as_any().is::()); + insta::assert_snapshot!(plan.tree_display(), @r" + root: ZonedPlan(i32, rows=5) + data: FlatPlan(i32, rows=5) + zones: FlatPlan({}, rows=2) + "); + Ok(()) +} + #[test] fn struct_plan_optimization_visits_all_fields() -> VortexResult<()> { let field_dtype = primitive(PType::I32, Nullability::NonNullable);