Thread scope dtype through zoned pruning

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

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

@@ -116,7 +116,7 @@ impl PruningState {
         self.pruning_predicates
             .entry(expr.clone())
             .or_default()
-            .get_or_init(move || match expr.falsify(&self.session) {
+            .get_or_init(move || match expr.falsify(&self.dtype, &self.session) {
                 Ok(predicate) => predicate,
                 Err(error) => {
                     trace!(%expr, %error, "failed to construct stats rewrite predicate");

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

@@ -269,8 +269,10 @@ mod tests {
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::StructArray;
     use vortex_array::assert_arrays_eq;
+    use vortex_array::dtype::DType;
     use vortex_array::dtype::FieldNames;
     use vortex_array::dtype::PType;
+    use vortex_array::expr::Expression;
     use vortex_array::expr::gt;
     use vortex_array::expr::gt_eq;
     use vortex_array::expr::is_not_null;
@@ -288,6 +290,10 @@ mod tests {
     use crate::layouts::zoned::zone_map::ZoneMap;
     use crate::test::SESSION;
 
+    fn falsify(expr: &Expression, dtype: DType) -> Expression {
+        expr.falsify(&dtype, &SESSION).unwrap().unwrap()
+    }
+
     #[test]
     fn test_zone_map_prunes() {
         // Construct a zone map with 3 zones:
@@ -331,7 +337,7 @@ mod tests {
         // A >= 6
         // => A.max < 6
         let expr = gt_eq(root(), lit(6i32));
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::I32.into());
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
@@ -341,7 +347,7 @@ mod tests {
         // A > 5
         // => A.max <= 5
         let expr = gt(root(), lit(5i32));
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::I32.into());
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
@@ -351,7 +357,7 @@ mod tests {
         // A < 2
         // => A.min >= 2
         let expr = lt(root(), lit(2i32));
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::I32.into());
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true]));
     }
@@ -372,7 +378,7 @@ mod tests {
         .unwrap();
 
         let expr = is_not_null(root());
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::U64.into());
 
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
@@ -458,7 +464,7 @@ mod tests {
         );
 
         let expr = gt(root(), lit(5u64));
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::U64.into());
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
@@ -488,7 +494,7 @@ mod tests {
         .unwrap();
 
         let expr = gt(root(), lit(5.0f32));
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::F32.into());
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
         assert_arrays_eq!(
             mask.into_array(),
@@ -541,7 +547,7 @@ mod tests {
         .unwrap();
 
         let expr = is_not_null(root());
-        let pruning_expr = expr.falsify(&SESSION).unwrap().unwrap();
+        let pruning_expr = falsify(&expr, PType::U64.into());
 
         // All three zones have length 4 (total rows = 12).
         let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();