From d67da2107885c24fbc3fcc7ba3f3d303302577ad Mon Sep 17 00:00:00 2001
From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com>
Date: Sat, 3 Oct 2026 01:54:05 +0000
Subject: [PATCH] fix(datafusion): preserve regexp_match null checks

---
 rust/lance-datafusion/src/logical_expr.rs | 19 ++---
 rust/lance-datafusion/src/planner.rs      | 87 ++++++++++++++++++-----
 2 files changed, 79 insertions(+), 27 deletions(-)

diff --git a/rust/lance-datafusion/src/logical_expr.rs b/rust/lance-datafusion/src/logical_expr.rs
index 2f724eb296a..a5c7b531fa9 100644
--- a/rust/lance-datafusion/src/logical_expr.rs
+++ b/rust/lance-datafusion/src/logical_expr.rs
@@ -195,15 +195,18 @@ pub fn coerce_filter_type_to_boolean(expr: Expr) -> Expr {
             Expr::IsNotNull(Box::new(Expr::ScalarFunction(sf)))
         }
 
-        // Recurse into boolean contexts so nested regexp_match terms are also coerced
-        Expr::BinaryExpr(BinaryExpr { left, op, right }) => Expr::BinaryExpr(BinaryExpr {
-            left: Box::new(coerce_filter_type_to_boolean(*left)),
-            op,
-            right: Box::new(coerce_filter_type_to_boolean(*right)),
-        }),
+        // Only boolean operands need coercion. Null checks and comparisons must
+        // preserve regexp_match's nullable list result.
+        Expr::BinaryExpr(BinaryExpr { left, op, right })
+            if matches!(op, Operator::And | Operator::Or) =>
+        {
+            Expr::BinaryExpr(BinaryExpr {
+                left: Box::new(coerce_filter_type_to_boolean(*left)),
+                op,
+                right: Box::new(coerce_filter_type_to_boolean(*right)),
+            })
+        }
         Expr::Not(inner) => Expr::Not(Box::new(coerce_filter_type_to_boolean(*inner))),
-        Expr::IsNull(inner) => Expr::IsNull(Box::new(coerce_filter_type_to_boolean(*inner))),
-        Expr::IsNotNull(inner) => Expr::IsNotNull(Box::new(coerce_filter_type_to_boolean(*inner))),
 
         // Pass-through for all other nodes
         other => other,
diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs
index 065a6b204b6..6638c9d77e9 100644
--- a/rust/lance-datafusion/src/planner.rs
+++ b/rust/lance-datafusion/src/planner.rs
@@ -2176,28 +2176,77 @@ mod tests {
         );
     }
 
-    #[test]
-    fn test_regexp_match_infer_error_without_boolean_coercion() {
-        // With the fix applied, using parse_filter should coerce regexp_match to boolean
-        // even when nested in a larger AND expression, so this should plan successfully.
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("keywords", DataType::Utf8, true),
-            Field::new("natural_caption", DataType::Utf8, true),
-            Field::new("poetic_caption", DataType::Utf8, true),
-        ]));
+    #[rstest]
+    #[case::bare("regexp_match(name, 'e[12]')", [false, true, true, false, false, false])]
+    #[case::is_not_null(
+        "regexp_match(name, 'e[12]') IS NOT NULL",
+        [false, true, true, false, false, false]
+    )]
+    #[case::is_null(
+        "regexp_match(name, 'e[12]') IS NULL",
+        [true, false, false, true, true, true]
+    )]
+    #[case::not_bare(
+        "NOT regexp_match(name, 'e[12]')",
+        [true, false, false, true, true, true]
+    )]
+    #[case::and_bare(
+        "regexp_match(name, 'e[12]') AND name <> 'name2'",
+        [false, true, false, false, false, false]
+    )]
+    #[case::or_bare(
+        "regexp_match(name, 'e[12]') OR name IS NULL",
+        [false, true, true, true, false, false]
+    )]
+    #[case::not_is_not_null(
+        "NOT (regexp_match(name, 'e[12]') IS NOT NULL)",
+        [true, false, false, true, true, true]
+    )]
+    #[case::and_is_null(
+        "regexp_match(name, 'e[12]') IS NULL AND name IS NOT NULL",
+        [true, false, false, false, true, true]
+    )]
+    #[case::or_is_not_null(
+        "regexp_match(name, 'e[12]') IS NOT NULL OR name IS NULL",
+        [false, true, true, true, false, false]
+    )]
+    fn test_regexp_match_filter_coercion(#[case] filter: &str, #[case] expected: [bool; 6]) {
+        let batch = arrow_array::record_batch!((
+            "name",
+            Utf8,
+            [
+                Some("name0"),
+                Some("name1"),
+                Some("name2"),
+                None,
+                Some("name4"),
+                Some("name5")
+            ]
+        ))
+        .unwrap();
+        let planner = Planner::new(batch.schema());
+        let expr = planner.parse_filter(filter).unwrap();
+        let physical_expr = planner.create_physical_expr(&expr).unwrap();
+        let result = physical_expr.evaluate(&batch).unwrap();
 
-        let planner = Planner::new(schema);
+        assert_eq!(
+            result.into_array(batch.num_rows()).unwrap().as_ref(),
+            &BooleanArray::from(expected.to_vec())
+        );
+    }
 
-        let expr = planner
-            .parse_filter(
-                "regexp_match(keywords, 'Liberty|revolution') AND \
-                 (natural_caption IS NOT NULL AND natural_caption <> '' AND \
-                  poetic_caption IS NOT NULL AND poetic_caption <> '')",
-            )
-            .unwrap();
+    #[rstest]
+    #[case::is_not_null("regexp_match(name, 'e[12]') IS NOT NULL")]
+    #[case::is_null("regexp_match(name, 'e[12]') IS NULL")]
+    #[case::comparison("regexp_match(name, 'e[12]') = regexp_match(name, 'e[12]')")]
+    fn test_regexp_match_preserves_value_contexts(#[case] filter: &str) {
+        let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Utf8, true)]));
+        let planner = Planner::new(schema);
 
-        // Should not panic
-        let _physical = planner.create_physical_expr(&expr).unwrap();
+        assert_eq!(
+            planner.parse_filter(filter).unwrap(),
+            planner.parse_expr(filter).unwrap()
+        );
     }
 
     #[test]