more

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

Author: robert3005

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

@@ -155,7 +155,7 @@ impl LayoutStrategy for ZonedStrategy {
             )
             .await?;
 
-        let Some(stats_table) = stats_accumulator.lock().as_stats_table()? else {
+        let Some((stats_table, stats)) = stats_accumulator.lock().as_stats_table()? 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);
@@ -174,13 +174,7 @@ impl LayoutStrategy for ZonedStrategy {
             .write_stream(ctx, Arc::clone(&segment_sink), stats_stream, eof, &session)
             .await?;
 
-        Ok(ZonedLayout::new(
-            data_layout,
-            zones_layout,
-            block_size,
-            Arc::clone(stats_table.present_stats()),
-        )
-        .into_layout())
+        Ok(ZonedLayout::new(data_layout, zones_layout, block_size, stats).into_layout())
     }
 
     fn buffered_bytes(&self) -> u64 {

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

@@ -43,6 +43,10 @@ pub struct ZoneMap {
     array: StructArray,
     // The statistics that are included in the table.
     stats: Arc<[Stat]>,
+    // The length of each zone in the zone map.
+    zone_len: usize,
+    // Number of rows that the zone map covers
+    row_count: u64,
 }
 
 impl ZoneMap {
@@ -52,24 +56,35 @@ impl ZoneMap {
         column_dtype: DType,
         array: StructArray,
         stats: Arc<[Stat]>,
+        zone_len: usize,
+        row_count: u64,
     ) -> VortexResult<Self> {
         let expected_dtype = stats_table_dtype(&column_dtype, &stats);
         if &expected_dtype != array.dtype() {
             vortex_bail!("Array dtype does not match expected zone map dtype: {expected_dtype}");
         }
 
         // SAFETY: We checked that the array matches the expected stats-table schema.
-        Ok(unsafe { Self::new_unchecked(array, stats) })
+        Ok(unsafe { Self::new_unchecked(array, stats, zone_len, row_count) })
     }
 
     /// Creates [`ZoneMap`] without validating return array against expected stats.
     ///
     /// # Safety
     ///
     /// Assumes that the input struct array has the correct statistics as fields. Or in other words,
-    /// the [`DType`] of the input array is equal to the result of `stats_table_dtype`.
-    pub unsafe fn new_unchecked(array: StructArray, stats: Arc<[Stat]>) -> Self {
-        Self { array, stats }
+    pub unsafe fn new_unchecked(
+        array: StructArray,
+        stats: Arc<[Stat]>,
+        zone_len: usize,
+        row_count: u64,
+    ) -> Self {
+        Self {
+            array,
+            stats,
+            zone_len,
+            row_count,
+        }
     }
 
     /// Returns the [`DType`] of the statistics table given a set of statistics and column [`DType`].
@@ -144,8 +159,6 @@ impl ZoneMap {
     pub fn prune(
         &self,
         predicate: &Expression,
-        zone_len: u64,
-        row_count: u64,
         session: &VortexSession,
     ) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
@@ -157,7 +170,7 @@ impl ZoneMap {
             return applied.execute::<Mask>(&mut ctx);
         }
 
-        let row_count_array = row_count_array(zone_len, row_count, num_zones, &mut ctx)?;
+        let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones, &mut ctx)?;
         let substituted = substitute_stats_fn_array::<RowCount>(applied, &row_count_array)?;
         substituted.execute::<Mask>(&mut ctx)
     }
@@ -169,7 +182,7 @@ impl ZoneMap {
 /// result is a [`ConstantArray`] for uniform zone sizes, otherwise a two-run
 /// run-end encoded array whose trailing run carries the final zone length.
 fn row_count_array(
-    zone_len: u64,
+    zone_len: usize,
     row_count: u64,
     num_zones: usize,
     ctx: &mut ExecutionCtx,
@@ -263,6 +276,7 @@ mod tests {
             ])
             .unwrap(),
             Arc::new([Stat::Max, Stat::Min]),
+            10,
         )
         .unwrap();