fix[vortex-datafusion]: check field exists in get_field pushdown (#5295)

Previously, any get field expression was pushed down, which resulted in
incorrectly pushing down get fields on non-existent fields, causing an
error at execution time rather than plan time.

This was a mistake on my part that should've been included in #5024

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

Author: asubiotto

vortex-datafusion/src/convert/exprs.rs

@@ -229,8 +229,7 @@ pub(crate) fn can_be_pushed_down(df_expr: &PhysicalExprRef, schema: &Schema) ->
         can_be_pushed_down(in_list.expr(), schema)
             && in_list.list().iter().all(|e| can_be_pushed_down(e, schema))
     } else if let Some(scalar_fn) = expr.downcast_ref::<ScalarFunctionExpr>() {
-        // Only get_field pushdown is supported.
-        ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn).is_some()
+        can_scalar_fn_be_pushed_down(scalar_fn, schema)
     } else {
         tracing::debug!(%df_expr, "DataFusion expression can't be pushed down");
         false
@@ -277,6 +276,53 @@ fn supported_data_types(dt: &DataType) -> bool {
     is_supported
 }
 
+/// Checks if a GetField scalar function can be pushed down.
+fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool {
+    let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn)
+    else {
+        // Only get_field pushdown is supported.
+        return false;
+    };
+
+    let args = get_field_fn.args();
+    if args.len() != 2 {
+        tracing::debug!(
+            "Expected 2 arguments for GetField, not pushing down {} arguments",
+            args.len()
+        );
+        return false;
+    }
+    let source_expr = &args[0];
+    let field_name_expr = &args[1];
+    let Some(field_name) = field_name_expr
+        .as_any()
+        .downcast_ref::<df_expr::Literal>()
+        .and_then(|lit| lit.value().try_as_str().flatten())
+    else {
+        return false;
+    };
+
+    let Ok(source_dt) = source_expr.data_type(schema) else {
+        tracing::debug!(
+            field_name = field_name,
+            schema = ?schema,
+            source_expr = ?source_expr,
+            "Failed to get source type for GetField, not pushing down"
+        );
+        return false;
+    };
+    let DataType::Struct(fields) = source_dt else {
+        tracing::debug!(
+            field_name = field_name,
+            schema = ?schema,
+            source_expr = ?source_expr,
+            "Failed to get source type as struct for GetField, not pushing down"
+        );
+        return false;
+    };
+    fields.find(field_name).is_some()
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
@@ -606,8 +652,10 @@ mod tests {
         "#);
     }
 
-    #[test]
-    fn test_can_be_pushed_down_get_field() {
+    #[rstest]
+    #[case::valid_field("field1", true)]
+    #[case::missing_field("nonexistent_field", false)]
+    fn test_can_be_pushed_down_get_field(#[case] field_name: &str, #[case] expected: bool) {
         let struct_fields = Fields::from(vec![
             Field::new("field1", DataType::Utf8, true),
             Field::new("field2", DataType::Int32, true),
@@ -619,18 +667,18 @@ mod tests {
         )]);
 
         let struct_col = Arc::new(df_expr::Column::new("my_struct", 0)) as Arc<dyn PhysicalExpr>;
-        let field_name = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some(
-            "field1".to_string(),
+        let field_name_lit = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some(
+            field_name.to_string(),
         )))) as Arc<dyn PhysicalExpr>;
 
         let get_field_expr = Arc::new(ScalarFunctionExpr::new(
             "get_field",
             Arc::new(ScalarUDF::from(GetFieldFunc::new())),
-            vec![struct_col, field_name],
-            Arc::new(Field::new("field1", DataType::Utf8, true)),
+            vec![struct_col, field_name_lit],
+            Arc::new(Field::new(field_name, DataType::Utf8, true)),
             Arc::new(ConfigOptions::new()),
         )) as Arc<dyn PhysicalExpr>;
 
-        assert!(can_be_pushed_down(&get_field_expr, &schema));
+        assert_eq!(can_be_pushed_down(&get_field_expr, &schema), expected);
     }
 }