Add zoned layout scan plan

Add a `Zoned` operator pairing data with the zone statistics summarising
it, and lower both `vortex.zoned` and legacy `vortex.stats` layouts into
it, since the two share a child shape.

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012obBhJ8oPZoBbKyeS79yMv
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-layout/src/plan/lower.rs

@@ -24,12 +24,15 @@ 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::layouts::zoned::LegacyStats;
+use crate::layouts::zoned::Zoned;
 use crate::plan::ConcatPlan;
 use crate::plan::ListPackPlan;
 use crate::plan::PackPlan;
 use crate::plan::PlanRef;
 use crate::plan::SegmentScanPlan;
 use crate::plan::TakePlan;
+use crate::plan::ZonedPlan;
 
 /// Lowers `layout` into a physical plan.
 ///
@@ -51,6 +54,9 @@ pub fn lower(layout: &LayoutRef) -> VortexResult<PlanRef> {
     if let Some(layout) = layout.as_opt::<List>() {
         return Ok(lower_list(layout)?.into_plan());
     }
+    if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
+        return Ok(lower_zoned(layout)?.into_plan());
+    }
     vortex_bail!(
         "No physical plan implementation for layout '{}'",
         layout.encoding_id()
@@ -143,3 +149,18 @@ fn lower_list(layout: &ListLayout) -> VortexResult<ListPackPlan> {
         validity,
     )
 }
+
+fn lower_zoned(layout: &LayoutRef) -> VortexResult<ZonedPlan> {
+    // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones.
+    let data = lower(
+        &layout
+            .slot(0)?
+            .ok_or_else(|| vortex_err!("Zoned data child is absent"))?,
+    )?;
+    let zones = lower(
+        &layout
+            .slot(1)?
+            .ok_or_else(|| vortex_err!("Zoned zones child is absent"))?,
+    )?;
+    Ok(ZonedPlan::new(data, zones))
+}

vortex-layout/src/plan/mod.rs

@@ -51,6 +51,8 @@ pub use plans::SegmentScanData;
 pub use plans::SegmentScanPlan;
 pub use plans::Take;
 pub use plans::TakePlan;
+pub use plans::Zoned;
+pub use plans::ZonedPlan;
 pub use plans::row_idx_dtype;
 pub use typed::DynPlan;
 pub use typed::Plan;

vortex-layout/src/plan/plans/mod.rs

@@ -10,6 +10,7 @@ mod row_idx_partition;
 mod row_idx_values;
 mod segment_scan;
 mod take;
+mod zoned;
 
 pub use concat::Concat;
 pub use concat::ConcatData;
@@ -43,3 +44,5 @@ pub use segment_scan::SegmentScanPlan;
 pub(crate) use take::ExpressionTakeRule;
 pub use take::Take;
 pub use take::TakePlan;
+pub use zoned::Zoned;
+pub use zoned::ZonedPlan;

vortex-layout/src/plan/plans/zoned.rs

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::borrow::Cow;
+
+use vortex_array::EmptyMetadata;
+use vortex_array::dtype::DType;
+use vortex_error::VortexResult;
+use vortex_session::registry::CachedId;
+
+use crate::plan::Plan;
+use crate::plan::PlanId;
+use crate::plan::PlanParts;
+use crate::plan::PlanRef;
+use crate::plan::PlanVTable;
+use crate::plan::check_child_count;
+
+pub(crate) const DATA: usize = 0;
+pub(crate) const ZONES: usize = 1;
+
+/// Reads data alongside the zone statistics summarising it.
+///
+/// This operator covers both `vortex.zoned` layouts and legacy `vortex.stats` layouts, which have
+/// the same physical child shape.
+#[derive(Clone, Debug)]
+pub struct Zoned;
+
+/// A plan that pairs data with its zone statistics.
+pub type ZonedPlan = Plan<Zoned>;
+
+impl ZonedPlan {
+    /// Creates a zoned plan over `data` summarised by `zones`.
+    pub fn new(data: PlanRef, zones: PlanRef) -> Self {
+        let dtype: DType = data.dtype().clone();
+        let row_count = data.row_count();
+        PlanParts {
+            vtable: Zoned,
+            dtype,
+            row_count,
+            children: vec![data, zones],
+            data: (),
+        }
+        .into_typed()
+    }
+
+    /// Returns the plan producing the summarised data.
+    pub fn data_plan(&self) -> &PlanRef {
+        &self.children()[DATA]
+    }
+
+    /// Returns the plan producing the zone statistics.
+    pub fn zones_plan(&self) -> &PlanRef {
+        &self.children()[ZONES]
+    }
+}
+
+impl PlanVTable for Zoned {
+    type PlanData = ();
+    type Metadata = EmptyMetadata;
+
+    fn id(&self) -> PlanId {
+        static ID: CachedId = CachedId::new("vortex.plan.zoned");
+        *ID
+    }
+
+    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
+        Some(EmptyMetadata)
+    }
+
+    fn with_children(_plan: &Plan<Self>, mut children: Vec<PlanRef>) -> VortexResult<Plan<Self>> {
+        check_child_count("Zoned", &children, 2)?;
+        let zones = children.remove(ZONES);
+        let data = children.remove(DATA);
+        Ok(ZonedPlan::new(data, zones))
+    }
+
+    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
+        match index {
+            DATA => Cow::Borrowed("data"),
+            ZONES => Cow::Borrowed("zones"),
+            _ => Cow::Owned(format!("child[{index}]")),
+        }
+    }
+}

vortex-layout/src/plan/tests.rs

@@ -2,8 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::fmt;
+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;
@@ -22,6 +24,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;
@@ -31,6 +35,8 @@ use crate::layouts::foreign::new_foreign_layout;
 use crate::layouts::list::ListLayout;
 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 {
@@ -852,3 +858,61 @@ fn nullable_struct_keeps_expression_above_parent_validity() -> VortexResult<()>
     assert!(eval.child_plan().is::<Pack>());
     Ok(())
 }
+
+#[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.is::<Zoned>());
+    insta::assert_snapshot!(plan.display_tree(), @"
+    root: vortex.plan.zoned(i32, rows=5)
+      data: vortex.plan.segment_scan(i32, rows=5)
+      zones: vortex.plan.segment_scan({}, 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.is::<Zoned>());
+    insta::assert_snapshot!(plan.display_tree(), @"
+    root: vortex.plan.zoned(i32, rows=5)
+      data: vortex.plan.segment_scan(i32, rows=5)
+      zones: vortex.plan.segment_scan({}, rows=2)
+    ");
+    Ok(())
+}