Use aggregate descriptors for zoned stats

vortex-array/src/aggregate_fn/fns/all_nan/mod.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -37,7 +38,15 @@ impl AggregateFnVTable for AllNan {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -37,7 +38,15 @@ impl AggregateFnVTable for AllNonNan {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -29,7 +30,15 @@ impl AggregateFnVTable for AllNonNull {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/all_null/mod.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -29,7 +30,15 @@ impl AggregateFnVTable for AllNull {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/max/mod.rs

@@ -3,6 +3,7 @@
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -52,7 +53,15 @@ impl AggregateFnVTable for Max {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/min/mod.rs

@@ -3,6 +3,7 @@
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -52,7 +53,15 @@ impl AggregateFnVTable for Min {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/nan_count/mod.rs

@@ -7,6 +7,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
+use vortex_session::VortexSession;
 
 use self::primitive::accumulate_primitive;
 use crate::ArrayRef;
@@ -87,7 +88,15 @@ impl AggregateFnVTable for NanCount {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        unimplemented!("NanCount is not yet serializable");
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/null_count/mod.rs

@@ -4,6 +4,7 @@
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Columnar;
@@ -62,7 +63,15 @@ impl AggregateFnVTable for NullCount {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        Ok(None)
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/sum/mod.rs

@@ -12,6 +12,7 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
+use vortex_session::VortexSession;
 
 use self::bool::accumulate_bool;
 use self::constant::multiply_constant;
@@ -76,7 +77,15 @@ impl AggregateFnVTable for Sum {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        unimplemented!("Sum is not yet serializable");
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs

@@ -27,6 +27,7 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 use vortex_mask::Mask;
+use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::Canonical;
@@ -104,7 +105,15 @@ impl AggregateFnVTable for UncompressedSizeInBytes {
     }
 
     fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        unimplemented!("UncompressedSizeInBytes is not yet serializable");
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {

vortex-cuda/src/layout.rs

@@ -40,6 +40,7 @@ use vortex::error::VortexResult;
 use vortex::error::vortex_bail;
 use vortex::error::vortex_panic;
 use vortex::layout::IntoLayout;
+use vortex::layout::LayoutBuildContext;
 use vortex::layout::LayoutChildType;
 use vortex::layout::LayoutChildren;
 use vortex::layout::LayoutEncodingRef;
@@ -198,7 +199,7 @@ impl VTable for CudaFlat {
         metadata: &<Self::Metadata as DeserializeMetadata>::Output,
         segment_ids: Vec<SegmentId>,
         _children: &dyn LayoutChildren,
-        ctx: &ReadContext,
+        build_ctx: &LayoutBuildContext<'_>,
     ) -> VortexResult<Self::Layout> {
         if segment_ids.len() != 1 {
             vortex_bail!("CudaFlatLayout must have exactly one segment ID");
@@ -212,7 +213,7 @@ impl VTable for CudaFlat {
             row_count,
             dtype: dtype.clone(),
             segment_id: segment_ids[0],
-            ctx: ctx.clone(),
+            ctx: build_ctx.array_read_ctx.clone(),
             array_tree: ByteBuffer::from(metadata.array_encoding_tree.clone()),
             host_buffers: Arc::new(host_buffers),
         })

vortex-file/src/footer/mod.rs

@@ -35,7 +35,6 @@ use vortex_flatbuffers::footer as fb;
 use vortex_layout::LayoutEncodingId;
 use vortex_layout::LayoutRef;
 use vortex_layout::layout_from_flatbuffer_with_options;
-use vortex_layout::session::LayoutSessionExt;
 use vortex_session::VortexSession;
 use vortex_session::registry::ReadContext;
 
@@ -106,7 +105,7 @@ impl Footer {
             &dtype,
             &layout_read_ctx,
             &array_read_ctx,
-            session.layouts().registry(),
+            session,
             session.allows_unknown(),
         )?;
 

vortex-file/src/tests.rs

@@ -1886,7 +1886,10 @@ async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
 
     // Find all zoned layouts and verify data segments come before zone map segments.
     fn check_zoned_ordering(layout: &dyn Layout, segment_specs: &[SegmentSpec]) {
-        if layout.encoding_id().as_ref() == "vortex.stats" {
+        if matches!(
+            layout.encoding_id().as_ref(),
+            "vortex.zoned" | "vortex.stats"
+        ) {
             // child 0 = data, child 1 = zones
             let data_offsets =
                 collect_segment_offsets(layout.child(0).unwrap().as_ref(), segment_specs);
@@ -1918,7 +1921,10 @@ async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
         all_data: &mut Vec<u64>,
         all_zones: &mut Vec<u64>,
     ) {
-        if layout.encoding_id().as_ref() == "vortex.stats" {
+        if matches!(
+            layout.encoding_id().as_ref(),
+            "vortex.zoned" | "vortex.stats"
+        ) {
             // child 0 = data, child 1 = zones
             all_data.extend(collect_segment_offsets(
                 layout.child(0).unwrap().as_ref(),

vortex-layout/src/children.rs

@@ -14,8 +14,10 @@ use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 use vortex_flatbuffers::FlatBuffer;
 use vortex_flatbuffers::layout as fbl;
+use vortex_session::VortexSession;
 use vortex_session::registry::ReadContext;
 
+use crate::LayoutBuildContext;
 use crate::LayoutRef;
 use crate::layouts::foreign::new_foreign_layout;
 use crate::segments::SegmentId;
@@ -105,6 +107,7 @@ pub(crate) struct ViewedLayoutChildren {
     layout_read_ctx: ReadContext,
     layouts: LayoutRegistry,
     allow_unknown: bool,
+    session: VortexSession,
     cache: Arc<[OnceCell<LayoutRef>]>,
 }
 
@@ -121,6 +124,7 @@ impl ViewedLayoutChildren {
         layout_read_ctx: ReadContext,
         layouts: LayoutRegistry,
         allow_unknown: bool,
+        session: VortexSession,
     ) -> Self {
         // SAFETY: guaranteed by caller
         let nchildren = unsafe { fbl::Layout::follow(flatbuffer.as_ref(), flatbuffer_loc) }
@@ -135,6 +139,7 @@ impl ViewedLayoutChildren {
             layout_read_ctx,
             layouts,
             allow_unknown,
+            session,
             cache,
         }
     }
@@ -206,6 +211,7 @@ impl LayoutChildren for ViewedLayoutChildren {
                     self.layout_read_ctx.clone(),
                     self.layouts.clone(),
                     self.allow_unknown,
+                    self.session.clone(),
                 )
             };
 
@@ -223,6 +229,10 @@ impl LayoutChildren for ViewedLayoutChildren {
                 ));
             };
 
+            let build_ctx = LayoutBuildContext {
+                session: &self.session,
+                array_read_ctx: &self.array_read_ctx,
+            };
             encoding.build(
                 dtype,
                 fb_child.row_count(),
@@ -237,7 +247,7 @@ impl LayoutChildren for ViewedLayoutChildren {
                     .map(SegmentId::from)
                     .collect_vec(),
                 &viewed_children,
-                &self.array_read_ctx,
+                &build_ctx,
             )
         })?;
         Ok(Arc::clone(layout_ref))