Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions rust/lance-datafusion/src/logical_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
87 changes: 68 additions & 19 deletions rust/lance-datafusion/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
Loading