fix[vortex-expr]: delegate AnalysisExpr to CastExpr child (#5179)

checked_pruning_expr would otherwise return None which would result in
no pruning using expressions that contained a CastExpr somewhere. It's
relatively common for e.g. struct fields to be cast to the right type
for comparison with literals by datafusion, and this expression
translates to get_item(cast(struct), "field") = "value". This cast
within a get_item would return None when field_path was called before
this change.

Signed-off-by: Alfonso Subiotto Marques <[email protected]>

Author: asubiotto

vortex-expr/src/exprs/cast.rs

@@ -3,12 +3,14 @@
 
 use vortex_array::compute::cast as compute_cast;
 use vortex_array::{ArrayRef, DeserializeMetadata, ProstMetadata};
-use vortex_dtype::DType;
+use vortex_dtype::{DType, FieldPath};
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr as pb;
 
 use crate::display::{DisplayAs, DisplayFormat};
-use crate::{AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable};
+use crate::{
+    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, StatsCatalog, VTable, vtable,
+};
 
 vtable!(Cast);
 
@@ -118,7 +120,23 @@ impl DisplayAs for CastExpr {
     }
 }
 
-impl AnalysisExpr for CastExpr {}
+impl AnalysisExpr for CastExpr {
+    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
+        self.child.max(catalog)
+    }
+
+    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
+        self.child.min(catalog)
+    }
+
+    fn nan_count(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
+        self.child.nan_count(catalog)
+    }
+
+    fn field_path(&self) -> Option<FieldPath> {
+        self.child.field_path()
+    }
+}
 
 /// Creates an expression that casts values to a target data type.
 ///

vortex-expr/src/pruning/pruning_expr.rs

@@ -104,12 +104,14 @@ mod tests {
     use rstest::{fixture, rstest};
     use vortex_array::compute::{BetweenOptions, StrictComparison};
     use vortex_array::stats::Stat;
-    use vortex_dtype::{FieldName, FieldPath, FieldPathSet};
+    use vortex_dtype::{
+        DType, FieldName, FieldNames, FieldPath, FieldPathSet, Nullability, StructFields,
+    };
 
     use crate::pruning::pruning_expr::HashMap;
     use crate::pruning::{checked_pruning_expr, field_path_stat_field_name};
     use crate::{
-        HashSet, and, between, col, eq, get_item, gt, gt_eq, lit, lt, lt_eq, not_eq, or, root,
+        HashSet, and, between, cast, col, eq, get_item, gt, gt_eq, lit, lt, lt_eq, not_eq, or, root,
     };
 
     // Implement some checked pruning expressions.
@@ -491,4 +493,36 @@ mod tests {
             &or(gt(lit(10), col("a_max")), gt(col("a_min"), lit(50)))
         );
     }
+
+    #[rstest]
+    fn pruning_cast_get_item_eq(available_stats: FieldPathSet) {
+        // This test verifies that cast properly forwards analysis methods to
+        // enable pruning.
+        let struct_dtype = DType::Struct(
+            StructFields::new(
+                FieldNames::from([FieldName::from("a"), FieldName::from("b")]),
+                vec![
+                    DType::Utf8(Nullability::Nullable),
+                    DType::Utf8(Nullability::Nullable),
+                ],
+            ),
+            Nullability::NonNullable,
+        );
+        let expr = eq(get_item("a", cast(root(), struct_dtype)), lit("value"));
+        let (converted, refs) = checked_pruning_expr(&expr, &available_stats).unwrap();
+        assert_eq!(
+            refs.map(),
+            &HashMap::from_iter([(
+                FieldPath::from_name("a"),
+                HashSet::from_iter([Stat::Min, Stat::Max])
+            )])
+        );
+        assert_eq!(
+            &converted,
+            &or(
+                gt(col("a_min"), lit("value")),
+                gt(lit("value"), col("a_max"))
+            )
+        );
+    }
 }