less

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

Author: robert3005

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

@@ -21,7 +21,8 @@ use vortex_array::dtype::PType;
 use vortex_array::expr::Expression;
 use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
-use vortex_array::scalar_fn::fns::stats_expression::{contains_stats_fn_array, substitute_stats_fn_array};
+use vortex_array::scalar_fn::fns::stats_expression::contains_stats_fn_array;
+use vortex_array::scalar_fn::fns::stats_expression::substitute_stats_fn_array;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::Validity;
 use vortex_buffer::buffer;
@@ -44,7 +45,7 @@ pub struct ZoneMap {
     // The statistics that are included in the table.
     stats: Arc<[Stat]>,
     // The length of each zone in the zone map.
-    zone_len: usize,
+    zone_len: u64,
     // Number of rows that the zone map covers
     row_count: u64,
 }
@@ -56,7 +57,7 @@ impl ZoneMap {
         column_dtype: DType,
         array: StructArray,
         stats: Arc<[Stat]>,
-        zone_len: usize,
+        zone_len: u64,
         row_count: u64,
     ) -> VortexResult<Self> {
         let expected_dtype = stats_table_dtype(&column_dtype, &stats);
@@ -76,7 +77,7 @@ impl ZoneMap {
     pub unsafe fn new_unchecked(
         array: StructArray,
         stats: Arc<[Stat]>,
-        zone_len: usize,
+        zone_len: u64,
         row_count: u64,
     ) -> Self {
         Self {
@@ -156,11 +157,7 @@ impl ZoneMap {
     /// [`ConstantArray`]; a short final zone uses a run-end encoded array.
     ///
     /// [`checked_pruning_expr`]: vortex_array::expr::pruning::checked_pruning_expr
-    pub fn prune(
-        &self,
-        predicate: &Expression,
-        session: &VortexSession,
-    ) -> VortexResult<Mask> {
+    pub fn prune(&self, predicate: &Expression, session: &VortexSession) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
         let num_zones = self.array.len();
 
@@ -170,7 +167,8 @@ impl ZoneMap {
             return applied.execute::<Mask>(&mut ctx);
         }
 
-        let row_count_array = row_count_array(self.zone_len, self.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)
     }
@@ -182,7 +180,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: usize,
+    zone_len: u64,
     row_count: u64,
     num_zones: usize,
     ctx: &mut ExecutionCtx,
@@ -276,6 +274,7 @@ mod tests {
             ])
             .unwrap(),
             Arc::new([Stat::Max, Stat::Min]),
+            3,
             10,
         )
         .unwrap();
@@ -284,7 +283,7 @@ mod tests {
         // => A.max < 6
         let expr = gt_eq(root(), lit(6i32));
         let (pruning_expr, _) = checked_pruning_expr(&expr, &stats).unwrap();
-        let mask = zone_map.prune(&pruning_expr, 1, 3, &SESSION).unwrap();
+        let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
             BoolArray::from_iter([true, false, false])
@@ -294,7 +293,7 @@ mod tests {
         // => A.max <= 5
         let expr = gt(root(), lit(5i32));
         let (pruning_expr, _) = checked_pruning_expr(&expr, &stats).unwrap();
-        let mask = zone_map.prune(&pruning_expr, 1, 3, &SESSION).unwrap();
+        let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
             BoolArray::from_iter([true, false, false])
@@ -304,7 +303,7 @@ mod tests {
         // => A.min >= 2
         let expr = lt(root(), lit(2i32));
         let (pruning_expr, _) = checked_pruning_expr(&expr, &stats).unwrap();
-        let mask = zone_map.prune(&pruning_expr, 1, 3, &SESSION).unwrap();
+        let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true]));
     }
 
@@ -318,6 +317,8 @@ mod tests {
             )])
             .unwrap(),
             Arc::new([Stat::NullCount]),
+            4,
+            10,
         )
         .unwrap();
 
@@ -326,7 +327,7 @@ mod tests {
         let expr = is_not_null(root());
         let (pruning_expr, _) = checked_pruning_expr(&expr, &available_stats).unwrap();
 
-        let mask = zone_map.prune(&pruning_expr, 4, 10, &SESSION).unwrap();
+        let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
             BoolArray::from_iter([false, false, true])
@@ -343,6 +344,8 @@ mod tests {
             )])
             .unwrap(),
             Arc::new([Stat::NullCount]),
+            4,
+            12,
         )
         .unwrap();
 
@@ -352,7 +355,7 @@ mod tests {
         let (pruning_expr, _) = checked_pruning_expr(&expr, &available_stats).unwrap();
 
         // All three zones have length 4 (total rows = 12).
-        let mask = zone_map.prune(&pruning_expr, 4, 12, &SESSION).unwrap();
+        let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
             BoolArray::from_iter([false, true, false])