fixes

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-layout/src/layouts/zoned/builder.rs

@@ -4,6 +4,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::marker::PhantomData;
+use std::sync::Arc;
 
 use itertools::Itertools;
 use vortex_array::ArrayRef;
@@ -90,9 +91,10 @@ impl StatsAccumulator {
         Ok(())
     }
 
-    pub fn as_array(&mut self) -> VortexResult<Option<StructArray>> {
+    pub fn as_array(&mut self) -> VortexResult<Option<(StructArray, Arc<[Stat]>)>> {
         let mut names = Vec::new();
         let mut fields = Vec::new();
+        let mut stats = Vec::new();
 
         for builder in self
             .builders
@@ -107,6 +109,7 @@ impl StatsAccumulator {
                 continue;
             }
 
+            stats.push(builder.stat());
             names.extend(values.names);
             fields.extend(values.arrays);
         }
@@ -115,7 +118,8 @@ impl StatsAccumulator {
             return Ok(None);
         }
 
-        StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable).map(Some)
+        let array = StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable)?;
+        Ok(Some((array, stats.into())))
     }
 
     /// Returns an aggregated stats set for the table.
@@ -125,7 +129,7 @@ impl StatsAccumulator {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<StatsSet> {
         let mut stats_set = StatsSet::default();
-        let Some(array) = self.as_array()? else {
+        let Some((array, _)) = self.as_array()? else {
             return Ok(stats_set);
         };
 
@@ -415,7 +419,7 @@ mod tests {
             .vortex_expect("push_chunk should succeed for test data");
         acc.push_chunk(&builder2.finish(), &mut ctx)
             .vortex_expect("push_chunk should succeed for test data");
-        let stats_table = acc.as_array().unwrap().expect("Must have stats table");
+        let (stats_table, _) = acc.as_array().unwrap().expect("Must have stats table");
         assert_eq!(
             stats_table.names().as_ref(),
             &[
@@ -452,7 +456,7 @@ mod tests {
         let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12);
         acc.push_chunk(&array, &mut ctx)
             .vortex_expect("push_chunk should succeed for test array");
-        let stats_table = acc.as_array().unwrap().expect("Must have stats table");
+        let (stats_table, _) = acc.as_array().unwrap().expect("Must have stats table");
         assert_eq!(
             stats_table.names().as_ref(),
             &[

vortex-layout/src/layouts/zoned/writer.rs

@@ -7,7 +7,6 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use futures::StreamExt as _;
-use itertools::Itertools;
 use parking_lot::Mutex;
 use vortex_array::ArrayContext;
 use vortex_array::IntoArray;
@@ -94,14 +93,7 @@ impl LayoutStrategy for ZonedStrategy {
             "ZonedStrategy requires block_size > 0 when writing"
         );
 
-        let stats: Arc<[Stat]> = self
-            .options
-            .stats
-            .as_ref()
-            .iter()
-            .cloned()
-            .sorted_unstable()
-            .collect();
+        let stats = Arc::clone(&self.options.stats);
         let session = session.clone();
         let compute_session = session.clone();
         let handle = session.handle();
@@ -113,13 +105,12 @@ impl LayoutStrategy for ZonedStrategy {
             self.options.max_variable_length_statistics_size,
         )));
 
-        let stats2 = Arc::clone(&stats);
         // We can compute per-chunk statistics in parallel, so we spawn tasks for each chunk
         let stream = SequentialStreamAdapter::new(
             stream.dtype().clone(),
             stream
                 .map(move |chunk| {
-                    let stats = Arc::clone(&stats2);
+                    let stats = Arc::clone(&stats);
                     let session = compute_session.clone();
                     handle2.spawn_cpu(move || {
                         let (sequence_id, chunk) = chunk?;
@@ -164,15 +155,15 @@ impl LayoutStrategy for ZonedStrategy {
             )
             .await?;
 
-        let Some(stats_table) = stats_accumulator.lock().as_array()? else {
+        let Some((stats_array, stats)) = stats_accumulator.lock().as_array()? else {
             // If we have no stats (e.g. the DType doesn't support them), then we just return the
             // child layout.
             return Ok(data_layout);
         };
 
         // We must defer creating the stats table LayoutWriter until now, because the DType of
         // the table depends on which stats were successfully computed.
-        let stats_stream = stats_table
+        let stats_stream = stats_array
             .into_array()
             .to_array_stream()
             .sequenced(eof.split_off());