Add zoned layout scan plan

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

Author: joseph-isaacs

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<dyn Plan>;
@@ -103,6 +106,9 @@ pub fn new_plan(layout: &LayoutRef) -> VortexResult<PlanRef> {
     if let Some(layout) = layout.as_opt::<Struct>() {
         return Ok(Arc::new(StructPlan::new(layout)));
     }
+    if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
+        return Ok(Arc::new(ZonedPlan::try_new(layout)?));
+    }
     vortex_bail!(
         "No physical plan implementation for layout '{}'",
         layout.encoding_id()

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;

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

@@ -0,0 +1,99 @@
+// 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<Self> {
+        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 name(&self) -> &'static str {
+        "ZonedPlan"
+    }
+
+    fn optimize(&self) -> VortexResult<PlanRef> {
+        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<Option<PlanRef>> {
+        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}]")),
+        }
+    }
+}

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<PlanRef> {
     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.is::<ZonedPlan>());
+    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.is::<ZonedPlan>());
+    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);