fix: validate substrait filters that need no field pruning

rust/lance-datafusion/src/substrait.rs

@@ -153,8 +153,21 @@ fn remove_extension_types(
     Ok((new_substrait_schema, new_arrow_schema, index_mapping))
 }
 
+/// Substrait's optional message fields are `None` when a producer omits them, so unwrapping one
+/// turns a malformed (or merely terse) filter into a panic. This walker runs on every filter, so
+/// report the missing field instead.
+fn missing_field(what: &str) -> Error {
+    Error::invalid_input(format!(
+        "filter expression was missing a required {what} field"
+    ))
+}
+
 fn remap_expr_references(expr: &mut Expression, mapping: &HashMap<usize, usize>) -> Result<()> {
-    match expr.rex_type.as_mut().unwrap() {
+    match expr
+        .rex_type
+        .as_mut()
+        .ok_or_else(|| missing_field("expression"))?
+    {
         // Simple, no field references possible
         RexType::Literal(_) | RexType::Nested(_) | RexType::DynamicParameter(_) => Ok(()),
         // Enum literals are deprecated in Substrait and should only appear in older plans.
@@ -174,33 +187,61 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap<usize, usize>)
                 remap_expr_references(arg, mapping)?;
             }
             for arg in &mut func.arguments {
-                match arg.arg_type.as_mut().unwrap() {
+                match arg
+                    .arg_type
+                    .as_mut()
+                    .ok_or_else(|| missing_field("function argument"))?
+                {
                     ArgType::Value(expr) => remap_expr_references(expr, mapping)?,
                     ArgType::Enum(_) | ArgType::Type(_) => {}
                 }
             }
             Ok(())
         }
         RexType::IfThen(ifthen) => {
-            for clause in ifthen.ifs.iter_mut() {
-                remap_expr_references(clause.r#if.as_mut().unwrap(), mapping)?;
-                remap_expr_references(clause.then.as_mut().unwrap(), mapping)?;
+            for (i, clause) in ifthen.ifs.iter_mut().enumerate() {
+                remap_expr_references(
+                    clause
+                        .r#if
+                        .as_mut()
+                        .ok_or_else(|| missing_field("if clause condition"))?,
+                    mapping,
+                )?;
+                match clause.then.as_mut() {
+                    Some(then) => remap_expr_references(then, mapping)?,
+                    // Only the leading clause may omit `then`, in which case its condition is
+                    // the case expression being matched against.
+                    None if i == 0 => {}
+                    None => return Err(missing_field("if clause result")),
+                }
+            }
+            if let Some(otherwise) = ifthen.r#else.as_mut() {
+                remap_expr_references(otherwise, mapping)?;
             }
-            remap_expr_references(ifthen.r#else.as_mut().unwrap(), mapping)?;
             Ok(())
         }
         RexType::SwitchExpression(switch) => {
             for clause in switch.ifs.iter_mut() {
-                remap_expr_references(clause.then.as_mut().unwrap(), mapping)?;
+                if let Some(then) = clause.then.as_mut() {
+                    remap_expr_references(then, mapping)?;
+                }
+            }
+            if let Some(otherwise) = switch.r#else.as_mut() {
+                remap_expr_references(otherwise, mapping)?;
             }
-            remap_expr_references(switch.r#else.as_mut().unwrap(), mapping)?;
             Ok(())
         }
         RexType::SingularOrList(orlist) => {
             for opt in orlist.options.iter_mut() {
                 remap_expr_references(opt, mapping)?;
             }
-            remap_expr_references(orlist.value.as_mut().unwrap(), mapping)?;
+            remap_expr_references(
+                orlist
+                    .value
+                    .as_mut()
+                    .ok_or_else(|| missing_field("IN list value"))?,
+                mapping,
+            )?;
             Ok(())
         }
         RexType::MultiOrList(orlist) => {
@@ -215,22 +256,35 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap<usize, usize>)
             Ok(())
         }
         RexType::Cast(cast) => {
-            remap_expr_references(cast.input.as_mut().unwrap(), mapping)?;
+            remap_expr_references(
+                cast.input
+                    .as_mut()
+                    .ok_or_else(|| missing_field("cast input"))?,
+                mapping,
+            )?;
             Ok(())
         }
         RexType::Selection(sel) => {
-            // Finally, the selection, which might actually have field references
-            let root_type = sel.root_type.as_mut().unwrap();
-            // These types of references do not reference input fields so no remap needed
+            // Finally, the selection, which might actually have field references.
+            // An omitted root is a reference into the input, same as RootReference.
             if matches!(
-                root_type,
-                RootType::Expression(_) | RootType::OuterReference(_)
+                sel.root_type.as_mut(),
+                Some(RootType::Expression(_) | RootType::OuterReference(_))
             ) {
+                // These types of references do not reference input fields so no remap needed
                 return Ok(());
             }
-            match sel.reference_type.as_mut().unwrap() {
+            match sel
+                .reference_type
+                .as_mut()
+                .ok_or_else(|| missing_field("field reference"))?
+            {
                 ReferenceType::DirectReference(direct) => {
-                    match direct.reference_type.as_mut().unwrap() {
+                    match direct
+                        .reference_type
+                        .as_mut()
+                        .ok_or_else(|| missing_field("reference segment"))?
+                    {
                         reference_segment::ReferenceType::ListElement(_)
                         | reference_segment::ReferenceType::MapKey(_) => Err(Error::invalid_input(
                             "map/list nested references not supported in pushdown filters",
@@ -301,19 +355,9 @@ pub async fn parse_substrait(
         let (substrait_schema, _, index_mapping) =
             remove_extension_types(envelope.base_schema.as_ref().unwrap(), input_schema.clone())?;
 
-        if substrait_schema.r#struct.as_ref().unwrap().types.len()
-            != envelope
-                .base_schema
-                .as_ref()
-                .unwrap()
-                .r#struct
-                .as_ref()
-                .unwrap()
-                .types
-                .len()
-        {
-            remap_expr_references(&mut expr, &index_mapping)?;
-        }
+        // Always walk the expression: this also rejects operators we cannot push down. When no
+        // fields were removed the mapping is the identity, so the remap itself is a no-op.
+        remap_expr_references(&mut expr, &index_mapping)?;
 
         substrait_schema
     } else {
@@ -539,16 +583,17 @@ mod tests {
     use arrow_schema::{DataType, Field, Schema};
     use datafusion::{
         execution::SessionState,
-        logical_expr::{BinaryExpr, Operator},
+        logical_expr::{BinaryExpr, Case, Operator},
         prelude::{Expr, SessionContext},
     };
     use datafusion_common::{Column, ScalarValue};
     use datafusion_substrait::substrait::proto::{
         Expression, ExpressionReference, ExtendedExpression, FunctionArgument, NamedStruct, Type,
         Version,
         expression::{
-            FieldReference, Literal, ReferenceSegment, RexType, ScalarFunction,
+            FieldReference, IfThen, Literal, ReferenceSegment, RexType, ScalarFunction,
             field_reference::{ReferenceType, RootReference, RootType},
+            if_then::IfClause,
             literal::LiteralType,
             reference_segment::{self, StructField},
         },
@@ -664,6 +709,155 @@ mod tests {
         assert_eq!(df_expr, expected);
     }
 
+    /// A base schema with no extension types needs no field pruning, which is the case that
+    /// used to skip validation entirely.
+    async fn parse_unpruned_expr(rex_type: RexType) -> lance_core::Result<Expr> {
+        let expr = ExtendedExpression {
+            version: Some(Version {
+                major_number: 0,
+                minor_number: 63,
+                patch_number: 1,
+                git_hash: "".to_string(),
+                producer: "unit-test".to_string(),
+            }),
+            extension_urns: vec![],
+            extensions: vec![],
+            referred_expr: vec![ExpressionReference {
+                output_names: vec!["filter_mask".to_string()],
+                expr_type: Some(ExprType::Expression(Expression {
+                    rex_type: Some(rex_type),
+                })),
+            }],
+            base_schema: Some(NamedStruct {
+                names: vec!["x".to_string()],
+                r#struct: Some(Struct {
+                    types: vec![Type {
+                        kind: Some(Kind::I32(I32 {
+                            type_variation_reference: 0,
+                            nullability: Nullability::Nullable as i32,
+                        })),
+                    }],
+                    type_variation_reference: 0,
+                    nullability: Nullability::Required as i32,
+                }),
+            }),
+            advanced_extensions: None,
+            expected_type_urls: vec![],
+        };
+
+        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
+        parse_substrait(expr.encode_to_vec().as_slice(), schema, &session_state()).await
+    }
+
+    #[tokio::test]
+    async fn test_unsupported_operator_rejected_without_pruning() {
+        let err = parse_unpruned_expr(RexType::Subquery(Box::default()))
+            .await
+            .expect_err("subqueries should be rejected in filter expressions");
+        assert!(
+            err.to_string()
+                .contains("Window functions or subqueries not allowed in filter expression"),
+            "unexpected error: {err}"
+        );
+
+        let err = parse_unpruned_expr(RexType::Lambda(Box::default()))
+            .await
+            .expect_err("lambdas should be rejected in filter expressions");
+        assert!(
+            err.to_string()
+                .contains("Lambda expressions not allowed in filter expression"),
+            "unexpected error: {err}"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_unpruned_selection_without_explicit_root() {
+        let expr = parse_unpruned_expr(RexType::Selection(Box::new(FieldReference {
+            reference_type: Some(ReferenceType::DirectReference(ReferenceSegment {
+                reference_type: Some(reference_segment::ReferenceType::StructField(Box::new(
+                    StructField {
+                        field: 0,
+                        child: None,
+                    },
+                ))),
+            })),
+            root_type: None,
+        })))
+        .await
+        .unwrap();
+
+        assert_eq!(expr, Expr::Column(Column::new_unqualified("x")));
+    }
+
+    /// Optional message fields that a producer may legitimately omit must not panic the walker.
+    #[tokio::test]
+    async fn test_unpruned_if_then_with_omitted_optional_fields() {
+        let condition = Expression {
+            rex_type: Some(RexType::Literal(Literal {
+                nullable: false,
+                type_variation_reference: 0,
+                literal_type: Some(LiteralType::Boolean(true)),
+            })),
+        };
+
+        // A leading clause without `then` supplies the case expression, and `else` is optional.
+        let expr = parse_unpruned_expr(RexType::IfThen(Box::new(IfThen {
+            ifs: vec![IfClause {
+                r#if: Some(condition),
+                then: None,
+            }],
+            r#else: None,
+        })))
+        .await
+        .unwrap();
+
+        assert_eq!(
+            expr,
+            Expr::Case(Case {
+                expr: Some(Box::new(Expr::Literal(
+                    ScalarValue::Boolean(Some(true)),
+                    None
+                ))),
+                when_then_expr: vec![],
+                else_expr: None,
+            })
+        );
+    }
+
+    /// DataFusion only tolerates an omitted `then` on the leading clause, so a later clause
+    /// missing one has to be rejected here rather than reaching the consumer.
+    #[tokio::test]
+    async fn test_unpruned_if_then_missing_nonleading_then() {
+        let condition = || Expression {
+            rex_type: Some(RexType::Literal(Literal {
+                nullable: false,
+                type_variation_reference: 0,
+                literal_type: Some(LiteralType::Boolean(true)),
+            })),
+        };
+
+        let err = parse_unpruned_expr(RexType::IfThen(Box::new(IfThen {
+            ifs: vec![
+                IfClause {
+                    r#if: Some(condition()),
+                    then: Some(condition()),
+                },
+                IfClause {
+                    r#if: Some(condition()),
+                    then: None,
+                },
+            ],
+            r#else: None,
+        })))
+        .await
+        .expect_err("a non-leading clause without `then` is not valid");
+
+        assert!(
+            err.to_string().contains("if clause result"),
+            "unexpected error: {err}"
+        );
+    }
+
     #[tokio::test]
     async fn test_expr_substrait_roundtrip() {
         let schema = arrow_schema::Schema::new(vec![Field::new("x", DataType::Int32, true)]);