lance_datafusion/
logical_expr.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Extends logical expression.
5
6use arrow_schema::DataType;
7
8use crate::expr::safe_coerce_scalar;
9use datafusion::logical_expr::{Between, ScalarUDFImpl};
10use datafusion::logical_expr::{BinaryExpr, Operator};
11use datafusion::prelude::*;
12use datafusion::scalar::ScalarValue;
13use datafusion_functions::core::expr_ext::FieldAccessor;
14use datafusion_functions::core::getfield::GetFieldFunc;
15use lance_arrow::DataTypeExt;
16
17use lance_core::datatypes::Schema;
18use lance_core::{Error, Result};
19/// Resolve a Value
20fn resolve_value(expr: &Expr, data_type: &DataType) -> Result<Expr> {
21    match expr {
22        Expr::Literal(scalar_value, metadata) => {
23            Ok(Expr::Literal(safe_coerce_scalar(scalar_value, data_type).ok_or_else(|| Error::invalid_input(format!("Received literal {expr} and could not convert to literal of type '{data_type:?}'")))?, metadata.clone()))
24        }
25        _ => Err(Error::invalid_input(format!("Expected a literal of type '{data_type:?}' but received: {expr}"))),
26    }
27}
28
29/// A simple helper function that interprets an Expr as a string scalar
30/// or returns None if it is not.
31pub fn get_as_string_scalar_opt(expr: &Expr) -> Option<&str> {
32    match expr {
33        Expr::Literal(ScalarValue::Utf8(Some(s)), _) => Some(s),
34        _ => None,
35    }
36}
37
38/// Given a column or nested `get_field` expression, get the data type of the referenced
39/// field in the schema.
40///
41/// If the column is not found in the schema, return None. If the expression is
42/// not a field reference, also returns None.
43pub fn resolve_column_type(expr: &Expr, schema: &Schema) -> Option<DataType> {
44    let mut field_path = Vec::new();
45    let mut current_expr = expr;
46    // We are looping from outer-most reference to inner-most.
47    loop {
48        match current_expr {
49            Expr::Column(c) => {
50                field_path.push(c.name.as_str());
51                break;
52            }
53            Expr::ScalarFunction(udf) if udf.name() == GetFieldFunc::default().name() => {
54                let name = get_as_string_scalar_opt(&udf.args[1])?;
55                field_path.push(name);
56                current_expr = &udf.args[0];
57            }
58            _ => return None,
59        }
60    }
61
62    let mut path_iter = field_path.iter().rev();
63    let mut field = schema.field(path_iter.next()?)?;
64    for name in path_iter {
65        if field.data_type().is_struct() {
66            field = field.children.iter().find(|f| &f.name == name)?;
67        } else {
68            return None;
69        }
70    }
71    Some(field.data_type())
72}
73
74fn is_literal_only_same_type_expr(expr: &Expr) -> bool {
75    match expr {
76        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
77            (op.is_numerical_operators()
78                || matches!(
79                    op,
80                    Operator::BitwiseShiftLeft
81                        | Operator::BitwiseShiftRight
82                        | Operator::StringConcat
83                ))
84                && is_literal_only_same_type_expr(left)
85                && is_literal_only_same_type_expr(right)
86        }
87        Expr::Literal(..) => true,
88        _ => false,
89    }
90}
91
92/// Resolve logical expression `expr`.
93///
94/// Parameters
95///
96/// - *expr*: a datafusion logical expression
97/// - *schema*: lance schema.
98pub fn resolve_expr(expr: &Expr, schema: &Schema) -> Result<Expr> {
99    match expr {
100        Expr::Between(Between {
101            expr: inner_expr,
102            low,
103            high,
104            negated,
105        }) => {
106            if let Some(inner_expr_type) = resolve_column_type(inner_expr.as_ref(), schema) {
107                Ok(Expr::Between(Between {
108                    expr: inner_expr.clone(),
109                    low: Box::new(coerce_expr(low.as_ref(), &inner_expr_type)?),
110                    high: Box::new(coerce_expr(high.as_ref(), &inner_expr_type)?),
111                    negated: *negated,
112                }))
113            } else {
114                Ok(expr.clone())
115            }
116        }
117        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
118            if matches!(op, Operator::And | Operator::Or) {
119                Ok(Expr::BinaryExpr(BinaryExpr {
120                    left: Box::new(resolve_expr(left.as_ref(), schema)?),
121                    op: *op,
122                    right: Box::new(resolve_expr(right.as_ref(), schema)?),
123                }))
124            } else if let Some(left_type) = resolve_column_type(left.as_ref(), schema) {
125                match right.as_ref() {
126                    Expr::Literal(..) | Expr::Cast(_) => Ok(Expr::BinaryExpr(BinaryExpr {
127                        left: left.clone(),
128                        op: *op,
129                        right: Box::new(coerce_expr(right.as_ref(), &left_type)?),
130                    })),
131                    // Constant expressions need the column type applied to all of their literals.
132                    Expr::BinaryExpr(_) if is_literal_only_same_type_expr(right) => {
133                        Ok(Expr::BinaryExpr(BinaryExpr {
134                            left: left.clone(),
135                            op: *op,
136                            right: Box::new(coerce_expr(right, &left_type)?),
137                        }))
138                    }
139                    // Let DataFusion coerce other expressions. Recursively resolving them
140                    // can narrow literals to an inner column's type before arithmetic promotes
141                    // its operands, for example the 0.5 in `float_col = int_col + 0.5`.
142                    _ => Ok(expr.clone()),
143                }
144            } else if let Some(right_type) = resolve_column_type(right.as_ref(), schema) {
145                match left.as_ref() {
146                    Expr::Literal(..) | Expr::Cast(_) => Ok(Expr::BinaryExpr(BinaryExpr {
147                        left: Box::new(coerce_expr(left.as_ref(), &right_type)?),
148                        op: *op,
149                        right: right.clone(),
150                    })),
151                    _ => Ok(expr.clone()),
152                }
153            } else {
154                Ok(expr.clone())
155            }
156        }
157        Expr::InList(in_list) => {
158            if matches!(in_list.expr.as_ref(), Expr::Column(_)) {
159                if let Some(resolved_type) = resolve_column_type(in_list.expr.as_ref(), schema) {
160                    let resolved_values = in_list
161                        .list
162                        .iter()
163                        .map(|val| coerce_expr(val, &resolved_type))
164                        .collect::<Result<Vec<_>>>()?;
165                    Ok(Expr::in_list(
166                        in_list.expr.as_ref().clone(),
167                        resolved_values,
168                        in_list.negated,
169                    ))
170                } else {
171                    Ok(expr.clone())
172                }
173            } else {
174                Ok(expr.clone())
175            }
176        }
177        _ => {
178            // Passthrough
179            Ok(expr.clone())
180        }
181    }
182}
183
184/// Coerce expression of literals to column type.
185///
186/// Parameters
187///
188/// - *expr*: a datafusion logical expression
189/// - *dtype*: a lance data type
190pub fn coerce_expr(expr: &Expr, dtype: &DataType) -> Result<Expr> {
191    match expr {
192        Expr::BinaryExpr(BinaryExpr { left, op, right }) => Ok(Expr::BinaryExpr(BinaryExpr {
193            left: Box::new(coerce_expr(left, dtype)?),
194            op: *op,
195            right: Box::new(coerce_expr(right, dtype)?),
196        })),
197        literal_expr @ Expr::Literal(..) => Ok(resolve_value(literal_expr, dtype)?),
198        Expr::Cast(_) => Ok(narrow_double_literal(expr, dtype).unwrap_or_else(|| expr.clone())),
199        _ => Ok(expr.clone()),
200    }
201}
202
203/// Narrows a double literal such as `CAST(0.5 AS double)` to a `Float32` column's
204/// type when the value converts exactly, so the column is not cast and stays
205/// indexable.
206fn narrow_double_literal(expr: &Expr, dtype: &DataType) -> Option<Expr> {
207    let Expr::Cast(cast) = expr else {
208        return None;
209    };
210    let Expr::Literal(value, _) = cast.expr.as_ref() else {
211        return None;
212    };
213    if *dtype != DataType::Float32 || *cast.field.data_type() != DataType::Float64 {
214        return None;
215    }
216    let ScalarValue::Float64(value) = value.cast_to(&DataType::Float64).ok()? else {
217        return None;
218    };
219    let narrowed = value.map(|value| value as f32);
220    // `total_cmp` also tells NaN signs and signed zeros apart.
221    let is_exact = value
222        .zip(narrowed)
223        .is_none_or(|(value, narrowed)| f64::from(narrowed).total_cmp(&value).is_eq());
224    is_exact.then(|| lit(ScalarValue::Float32(narrowed)))
225}
226
227/// Coerce logical expression for filters to boolean.
228///
229/// Parameters
230///
231/// - *expr*: a datafusion logical expression
232pub fn coerce_filter_type_to_boolean(expr: Expr) -> Expr {
233    match expr {
234        // Coerce regexp_match to boolean by checking for non-null
235        Expr::ScalarFunction(sf) if sf.func.name() == "regexp_match" => {
236            log::warn!(
237                "regexp_match now is coerced to boolean, this may be changed in the future, please use `regexp_like` instead"
238            );
239            Expr::IsNotNull(Box::new(Expr::ScalarFunction(sf)))
240        }
241
242        // Only boolean operands need coercion. Null checks and comparisons must
243        // preserve regexp_match's nullable list result.
244        Expr::BinaryExpr(BinaryExpr { left, op, right })
245            if matches!(op, Operator::And | Operator::Or) =>
246        {
247            Expr::BinaryExpr(BinaryExpr {
248                left: Box::new(coerce_filter_type_to_boolean(*left)),
249                op,
250                right: Box::new(coerce_filter_type_to_boolean(*right)),
251            })
252        }
253        Expr::Not(inner) => Expr::Not(Box::new(coerce_filter_type_to_boolean(*inner))),
254
255        // Pass-through for all other nodes
256        other => other,
257    }
258}
259
260/// Convert a field path string into a DataFusion expression.
261///
262/// This function handles:
263/// - Simple column names: "column"
264/// - Nested paths: "parent.child" or "parent.child.grandchild"
265/// - Backtick-escaped field names: "parent.`field.with.dots`"
266///
267/// # Arguments
268///
269/// * `field_path` - The field path to convert. Supports simple columns, nested paths,
270///   and backtick-escaped field names.
271///
272/// # Returns
273///
274/// Returns `Result<Expr>` - Ok with the DataFusion expression, or Err if the path
275/// could not be parsed.
276///
277/// # Example
278///
279/// ```
280/// use lance_datafusion::logical_expr::field_path_to_expr;
281///
282/// // Simple column
283/// let expr = field_path_to_expr("column_name").unwrap();
284///
285/// // Nested field
286/// let expr = field_path_to_expr("parent.child").unwrap();
287///
288/// // Backtick-escaped field with dots
289/// let expr = field_path_to_expr("parent.`field.with.dots`").unwrap();
290/// ```
291pub fn field_path_to_expr(field_path: &str) -> Result<Expr> {
292    // Parse the field path to handle nested fields and backtick-escaped names
293    let parts = lance_core::datatypes::parse_field_path(field_path)?;
294
295    if parts.is_empty() {
296        return Err(Error::invalid_input(format!(
297            "Invalid empty field path: {}",
298            field_path
299        )));
300    }
301
302    // Build the column expression, handling nested fields.
303    let mut expr = Expr::Column(datafusion::common::Column::new_unqualified(
304        parts[0].clone(),
305    ));
306    for part in &parts[1..] {
307        expr = expr.field(part.as_str());
308    }
309
310    Ok(expr)
311}
312
313#[cfg(test)]
314mod tests {
315    use std::sync::Arc;
316
317    use super::*;
318
319    use arrow_schema::{Field, Schema as ArrowSchema};
320    use datafusion::common::Column;
321
322    #[test]
323    fn test_field_path_to_expr_preserves_case_sensitive_root_column() {
324        let expr = field_path_to_expr("VECTOR").unwrap();
325
326        assert_eq!(expr, Expr::Column(Column::new_unqualified("VECTOR")));
327    }
328
329    #[rstest::rstest]
330    #[case::nested("Parent.Child", "Parent", vec!["Child"])]
331    #[case::deeply_nested("Parent.Child.Leaf", "Parent", vec!["Child", "Leaf"])]
332    #[case::escaped_child("Parent.`Child.With.Dot`", "Parent", vec!["Child.With.Dot"])]
333    #[case::escaped_root("`Parent.With.Dot`.Child", "Parent.With.Dot", vec!["Child"])]
334    #[case::escaped_nested(
335        "`Parent.With.Dot`.`Child.With.Dot`.Leaf",
336        "Parent.With.Dot",
337        vec!["Child.With.Dot", "Leaf"]
338    )]
339    fn test_field_path_to_expr_preserves_nested_names(
340        #[case] path: &str,
341        #[case] root: &str,
342        #[case] children: Vec<&str>,
343    ) {
344        let expr = field_path_to_expr(path).unwrap();
345        let mut expected = Expr::Column(Column::new_unqualified(root));
346        for name in children {
347            expected = get_field(expected, name);
348        }
349
350        assert_eq!(expr, expected);
351    }
352
353    #[test]
354    fn test_resolve_large_utf8() {
355        let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::LargeUtf8, false)]);
356        let expr = Expr::BinaryExpr(BinaryExpr {
357            left: Box::new(Expr::Column("a".to_string().into())),
358            op: Operator::Eq,
359            right: Box::new(Expr::Literal(
360                ScalarValue::Utf8(Some("a".to_string())),
361                None,
362            )),
363        });
364
365        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
366        match resolved {
367            Expr::BinaryExpr(be) => {
368                assert_eq!(
369                    be.right.as_ref(),
370                    &Expr::Literal(ScalarValue::LargeUtf8(Some("a".to_string())), None)
371                )
372            }
373            _ => unreachable!("Expected BinaryExpr"),
374        };
375    }
376
377    #[test]
378    fn test_resolve_binary_expr_on_right() {
379        let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::Float64, false)]);
380        let expr = Expr::BinaryExpr(BinaryExpr {
381            left: Box::new(Expr::Column("a".to_string().into())),
382            op: Operator::Eq,
383            right: Box::new(Expr::BinaryExpr(BinaryExpr {
384                left: Box::new(Expr::Literal(ScalarValue::Int64(Some(2)), None)),
385                op: Operator::Minus,
386                right: Box::new(Expr::Literal(ScalarValue::Int64(Some(-1)), None)),
387            })),
388        });
389        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
390
391        match resolved {
392            Expr::BinaryExpr(be) => match be.right.as_ref() {
393                Expr::BinaryExpr(r_be) => {
394                    assert_eq!(
395                        r_be.left.as_ref(),
396                        &Expr::Literal(ScalarValue::Float64(Some(2.0)), None)
397                    );
398                    assert_eq!(
399                        r_be.right.as_ref(),
400                        &Expr::Literal(ScalarValue::Float64(Some(-1.0)), None)
401                    );
402                }
403                _ => panic!("Expected BinaryExpr"),
404            },
405            _ => panic!("Expected BinaryExpr"),
406        }
407    }
408
409    #[test]
410    fn test_resolve_in_expr() {
411        // Type coercion should apply for `A IN (0)` or `A NOT IN (0)`
412        let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::Float32, false)]);
413        let expr = Expr::in_list(
414            Expr::Column("a".to_string().into()),
415            vec![Expr::Literal(ScalarValue::Float64(Some(0.0)), None)],
416            false,
417        );
418        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
419        let expected = Expr::in_list(
420            Expr::Column("a".to_string().into()),
421            vec![Expr::Literal(ScalarValue::Float32(Some(0.0)), None)],
422            false,
423        );
424        assert_eq!(resolved, expected);
425
426        let expr = Expr::in_list(
427            Expr::Column("a".to_string().into()),
428            vec![Expr::Literal(ScalarValue::Float64(Some(0.0)), None)],
429            true,
430        );
431        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
432        let expected = Expr::in_list(
433            Expr::Column("a".to_string().into()),
434            vec![Expr::Literal(ScalarValue::Float32(Some(0.0)), None)],
435            true,
436        );
437        assert_eq!(resolved, expected);
438    }
439
440    #[rstest::rstest]
441    #[case::half("x > CAST(0.5 AS DOUBLE)", col("x").gt(lit(0.5_f32)))]
442    #[case::reversed("CAST(0.5 AS DOUBLE) < x", lit(0.5_f32).lt(col("x")))]
443    #[case::negative_infinity("x < CAST('-inf' AS DOUBLE)", col("x").lt(lit(f32::NEG_INFINITY)))]
444    #[case::null("x > CAST(NULL AS DOUBLE)", col("x").gt(lit(ScalarValue::Float32(None))))]
445    #[case::between(
446        "x BETWEEN 0.0 AND CAST(0.5 AS DOUBLE)",
447        col("x").between(lit(0.0_f32), lit(0.5_f32))
448    )]
449    #[case::in_list("x IN (CAST(0.5 AS DOUBLE))", col("x").in_list(vec![lit(0.5_f32)], false))]
450    #[case::inexact(
451        "x = CAST(0.1 AS DOUBLE)",
452        col("x").eq(cast(lit(0.1), DataType::Float64))
453    )]
454    #[case::overflow(
455        "x > CAST(1e300 AS DOUBLE)",
456        col("x").gt(cast(lit(1e300), DataType::Float64))
457    )]
458    #[case::underflow(
459        "x > CAST(1e-300 AS DOUBLE)",
460        col("x").gt(cast(lit(1e-300), DataType::Float64))
461    )]
462    fn test_resolve_double_literal_on_float32(#[case] sql: &str, #[case] expected: Expr) {
463        let schema = ArrowSchema::new(vec![Field::new("x", DataType::Float32, true)]);
464        let planner = crate::planner::Planner::new(Arc::new(schema));
465        assert_eq!(planner.parse_filter(sql).unwrap(), expected, "{sql}");
466    }
467
468    #[test]
469    fn test_resolve_column_type() {
470        let schema = Arc::new(ArrowSchema::new(vec![
471            Field::new("int", DataType::Int32, true),
472            Field::new(
473                "st",
474                DataType::Struct(
475                    vec![
476                        Field::new("str", DataType::Utf8, true),
477                        Field::new("Child.With.Dot", DataType::Int32, true),
478                        Field::new(
479                            "st",
480                            DataType::Struct(
481                                vec![Field::new("float", DataType::Float64, true)].into(),
482                            ),
483                            true,
484                        ),
485                    ]
486                    .into(),
487                ),
488                true,
489            ),
490        ]));
491        let schema = Schema::try_from(schema.as_ref()).unwrap();
492
493        assert_eq!(
494            resolve_column_type(&col("int"), &schema),
495            Some(DataType::Int32)
496        );
497        assert_eq!(
498            resolve_column_type(&col("st").field("str"), &schema),
499            Some(DataType::Utf8)
500        );
501        assert_eq!(
502            resolve_column_type(&col("st").field("st").field("float"), &schema),
503            Some(DataType::Float64)
504        );
505
506        assert_eq!(
507            resolve_column_type(&field_path_to_expr("st.`Child.With.Dot`").unwrap(), &schema),
508            Some(DataType::Int32)
509        );
510
511        assert_eq!(resolve_column_type(&col("x"), &schema), None);
512        assert_eq!(resolve_column_type(&col("str"), &schema), None);
513        assert_eq!(resolve_column_type(&col("float"), &schema), None);
514        assert_eq!(
515            resolve_column_type(&col("st").field("str").eq(lit("x")), &schema),
516            None
517        );
518    }
519
520    #[test]
521    fn test_resolve_utf8view_literal_against_utf8_column() {
522        // Simulates DataFusion 43+ producing a Utf8View literal (e.g. from md5())
523        // being compared against a Utf8 column stored in Lance.
524        let arrow_schema = ArrowSchema::new(vec![Field::new("hash", DataType::Utf8, false)]);
525        let schema = Schema::try_from(&arrow_schema).unwrap();
526
527        let expr = Expr::BinaryExpr(BinaryExpr {
528            left: Box::new(Expr::Column("hash".to_string().into())),
529            op: Operator::Eq,
530            right: Box::new(Expr::Literal(
531                ScalarValue::Utf8View(Some("abc".to_string())),
532                None,
533            )),
534        });
535
536        let resolved = resolve_expr(&expr, &schema).unwrap();
537        match resolved {
538            Expr::BinaryExpr(be) => {
539                assert_eq!(
540                    be.right.as_ref(),
541                    &Expr::Literal(ScalarValue::Utf8(Some("abc".to_string())), None)
542                )
543            }
544            _ => unreachable!("Expected BinaryExpr"),
545        }
546    }
547
548    #[test]
549    fn test_resolve_typed_null_against_dictionary_column() {
550        // A dictionary-encoded string column, e.g. a categorical field.
551        let dict_ty = DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8));
552        let arrow_schema = ArrowSchema::new(vec![Field::new("etld", dict_ty, true)]);
553        let schema = Schema::try_from(&arrow_schema).unwrap();
554
555        // A typed null must be wrapped in the dictionary type, not left as a bare
556        // `Utf8(None)` literal sitting next to a `Dictionary(...)` column.
557        let expected_null = Expr::Literal(
558            ScalarValue::Dictionary(Box::new(DataType::Int16), Box::new(ScalarValue::Utf8(None))),
559            None,
560        );
561
562        // `etld = <typed null>` built directly via the API, as opposed to coming
563        // through SQL parsing.
564        let expr = Expr::BinaryExpr(BinaryExpr {
565            left: Box::new(Expr::Column("etld".to_string().into())),
566            op: Operator::Eq,
567            right: Box::new(Expr::Literal(ScalarValue::Utf8(None), None)),
568        });
569        match resolve_expr(&expr, &schema).unwrap() {
570            Expr::BinaryExpr(be) => assert_eq!(be.right.as_ref(), &expected_null),
571            other => unreachable!("Expected BinaryExpr, got {other:?}"),
572        }
573
574        // `etld IN ('a', <typed null>)` — a typed value mixed with a typed null,
575        // both already typed as Utf8. Every list element is wrapped in the
576        // dictionary type.
577        let expr = Expr::in_list(
578            Expr::Column("etld".to_string().into()),
579            vec![
580                Expr::Literal(ScalarValue::Utf8(Some("a".to_string())), None),
581                Expr::Literal(ScalarValue::Utf8(None), None),
582            ],
583            false,
584        );
585        let expected = Expr::in_list(
586            Expr::Column("etld".to_string().into()),
587            vec![
588                Expr::Literal(
589                    ScalarValue::Dictionary(
590                        Box::new(DataType::Int16),
591                        Box::new(ScalarValue::Utf8(Some("a".to_string()))),
592                    ),
593                    None,
594                ),
595                expected_null,
596            ],
597            false,
598        );
599        assert_eq!(resolve_expr(&expr, &schema).unwrap(), expected);
600    }
601}
lance_datafusion/logical_expr.rs

lance_datafusion/
logical_expr.rs
