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
74/// Resolve logical expression `expr`.
75///
76/// Parameters
77///
78/// - *expr*: a datafusion logical expression
79/// - *schema*: lance schema.
80pub fn resolve_expr(expr: &Expr, schema: &Schema) -> Result<Expr> {
81    match expr {
82        Expr::Between(Between {
83            expr: inner_expr,
84            low,
85            high,
86            negated,
87        }) => {
88            if let Some(inner_expr_type) = resolve_column_type(inner_expr.as_ref(), schema) {
89                Ok(Expr::Between(Between {
90                    expr: inner_expr.clone(),
91                    low: Box::new(coerce_expr(low.as_ref(), &inner_expr_type)?),
92                    high: Box::new(coerce_expr(high.as_ref(), &inner_expr_type)?),
93                    negated: *negated,
94                }))
95            } else {
96                Ok(expr.clone())
97            }
98        }
99        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
100            if matches!(op, Operator::And | Operator::Or) {
101                Ok(Expr::BinaryExpr(BinaryExpr {
102                    left: Box::new(resolve_expr(left.as_ref(), schema)?),
103                    op: *op,
104                    right: Box::new(resolve_expr(right.as_ref(), schema)?),
105                }))
106            } else if let Some(left_type) = resolve_column_type(left.as_ref(), schema) {
107                match right.as_ref() {
108                    Expr::Literal(..) => Ok(Expr::BinaryExpr(BinaryExpr {
109                        left: left.clone(),
110                        op: *op,
111                        right: Box::new(resolve_value(right.as_ref(), &left_type)?),
112                    })),
113                    // For cases complex expressions (not just literals) on right hand side like x = 1 + 1 + -2*2
114                    Expr::BinaryExpr(r) => Ok(Expr::BinaryExpr(BinaryExpr {
115                        left: left.clone(),
116                        op: *op,
117                        right: Box::new(Expr::BinaryExpr(BinaryExpr {
118                            left: coerce_expr(&r.left, &left_type).map(Box::new)?,
119                            op: r.op,
120                            right: coerce_expr(&r.right, &left_type).map(Box::new)?,
121                        })),
122                    })),
123                    _ => Ok(expr.clone()),
124                }
125            } else if let Some(right_type) = resolve_column_type(right.as_ref(), schema) {
126                match left.as_ref() {
127                    Expr::Literal(..) => Ok(Expr::BinaryExpr(BinaryExpr {
128                        left: Box::new(resolve_value(left.as_ref(), &right_type)?),
129                        op: *op,
130                        right: right.clone(),
131                    })),
132                    _ => Ok(expr.clone()),
133                }
134            } else {
135                Ok(expr.clone())
136            }
137        }
138        Expr::InList(in_list) => {
139            if matches!(in_list.expr.as_ref(), Expr::Column(_)) {
140                if let Some(resolved_type) = resolve_column_type(in_list.expr.as_ref(), schema) {
141                    let resolved_values = in_list
142                        .list
143                        .iter()
144                        .map(|val| coerce_expr(val, &resolved_type))
145                        .collect::<Result<Vec<_>>>()?;
146                    Ok(Expr::in_list(
147                        in_list.expr.as_ref().clone(),
148                        resolved_values,
149                        in_list.negated,
150                    ))
151                } else {
152                    Ok(expr.clone())
153                }
154            } else {
155                Ok(expr.clone())
156            }
157        }
158        _ => {
159            // Passthrough
160            Ok(expr.clone())
161        }
162    }
163}
164
165/// Coerce expression of literals to column type.
166///
167/// Parameters
168///
169/// - *expr*: a datafusion logical expression
170/// - *dtype*: a lance data type
171pub fn coerce_expr(expr: &Expr, dtype: &DataType) -> Result<Expr> {
172    match expr {
173        Expr::BinaryExpr(BinaryExpr { left, op, right }) => Ok(Expr::BinaryExpr(BinaryExpr {
174            left: Box::new(coerce_expr(left, dtype)?),
175            op: *op,
176            right: Box::new(coerce_expr(right, dtype)?),
177        })),
178        literal_expr @ Expr::Literal(..) => Ok(resolve_value(literal_expr, dtype)?),
179        _ => Ok(expr.clone()),
180    }
181}
182
183/// Coerce logical expression for filters to boolean.
184///
185/// Parameters
186///
187/// - *expr*: a datafusion logical expression
188pub fn coerce_filter_type_to_boolean(expr: Expr) -> Expr {
189    match expr {
190        // Coerce regexp_match to boolean by checking for non-null
191        Expr::ScalarFunction(sf) if sf.func.name() == "regexp_match" => {
192            log::warn!(
193                "regexp_match now is coerced to boolean, this may be changed in the future, please use `regexp_like` instead"
194            );
195            Expr::IsNotNull(Box::new(Expr::ScalarFunction(sf)))
196        }
197
198        // Recurse into boolean contexts so nested regexp_match terms are also coerced
199        Expr::BinaryExpr(BinaryExpr { left, op, right }) => Expr::BinaryExpr(BinaryExpr {
200            left: Box::new(coerce_filter_type_to_boolean(*left)),
201            op,
202            right: Box::new(coerce_filter_type_to_boolean(*right)),
203        }),
204        Expr::Not(inner) => Expr::Not(Box::new(coerce_filter_type_to_boolean(*inner))),
205        Expr::IsNull(inner) => Expr::IsNull(Box::new(coerce_filter_type_to_boolean(*inner))),
206        Expr::IsNotNull(inner) => Expr::IsNotNull(Box::new(coerce_filter_type_to_boolean(*inner))),
207
208        // Pass-through for all other nodes
209        other => other,
210    }
211}
212
213/// Convert a field path string into a DataFusion expression.
214///
215/// This function handles:
216/// - Simple column names: "column"
217/// - Nested paths: "parent.child" or "parent.child.grandchild"
218/// - Backtick-escaped field names: "parent.`field.with.dots`"
219///
220/// # Arguments
221///
222/// * `field_path` - The field path to convert. Supports simple columns, nested paths,
223///   and backtick-escaped field names.
224///
225/// # Returns
226///
227/// Returns `Result<Expr>` - Ok with the DataFusion expression, or Err if the path
228/// could not be parsed.
229///
230/// # Example
231///
232/// ```
233/// use lance_datafusion::logical_expr::field_path_to_expr;
234///
235/// // Simple column
236/// let expr = field_path_to_expr("column_name").unwrap();
237///
238/// // Nested field
239/// let expr = field_path_to_expr("parent.child").unwrap();
240///
241/// // Backtick-escaped field with dots
242/// let expr = field_path_to_expr("parent.`field.with.dots`").unwrap();
243/// ```
244pub fn field_path_to_expr(field_path: &str) -> Result<Expr> {
245    // Parse the field path to handle nested fields and backtick-escaped names
246    let parts = lance_core::datatypes::parse_field_path(field_path)?;
247
248    if parts.is_empty() {
249        return Err(Error::invalid_input(format!(
250            "Invalid empty field path: {}",
251            field_path
252        )));
253    }
254
255    // Build the column expression, handling nested fields.
256    let mut expr = Expr::Column(datafusion::common::Column::new_unqualified(
257        parts[0].clone(),
258    ));
259    for part in &parts[1..] {
260        expr = expr.field(part.as_str());
261    }
262
263    Ok(expr)
264}
265
266#[cfg(test)]
267mod tests {
268    use std::sync::Arc;
269
270    use super::*;
271
272    use arrow_schema::{Field, Schema as ArrowSchema};
273    use datafusion::common::Column;
274
275    #[test]
276    fn test_field_path_to_expr_preserves_case_sensitive_root_column() {
277        let expr = field_path_to_expr("VECTOR").unwrap();
278
279        assert_eq!(expr, Expr::Column(Column::new_unqualified("VECTOR")));
280    }
281
282    #[rstest::rstest]
283    #[case::nested("Parent.Child", "Parent", vec!["Child"])]
284    #[case::deeply_nested("Parent.Child.Leaf", "Parent", vec!["Child", "Leaf"])]
285    #[case::escaped_child("Parent.`Child.With.Dot`", "Parent", vec!["Child.With.Dot"])]
286    #[case::escaped_root("`Parent.With.Dot`.Child", "Parent.With.Dot", vec!["Child"])]
287    #[case::escaped_nested(
288        "`Parent.With.Dot`.`Child.With.Dot`.Leaf",
289        "Parent.With.Dot",
290        vec!["Child.With.Dot", "Leaf"]
291    )]
292    fn test_field_path_to_expr_preserves_nested_names(
293        #[case] path: &str,
294        #[case] root: &str,
295        #[case] children: Vec<&str>,
296    ) {
297        let expr = field_path_to_expr(path).unwrap();
298        let mut expected = Expr::Column(Column::new_unqualified(root));
299        for name in children {
300            expected = get_field(expected, name);
301        }
302
303        assert_eq!(expr, expected);
304    }
305
306    #[test]
307    fn test_resolve_large_utf8() {
308        let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::LargeUtf8, false)]);
309        let expr = Expr::BinaryExpr(BinaryExpr {
310            left: Box::new(Expr::Column("a".to_string().into())),
311            op: Operator::Eq,
312            right: Box::new(Expr::Literal(
313                ScalarValue::Utf8(Some("a".to_string())),
314                None,
315            )),
316        });
317
318        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
319        match resolved {
320            Expr::BinaryExpr(be) => {
321                assert_eq!(
322                    be.right.as_ref(),
323                    &Expr::Literal(ScalarValue::LargeUtf8(Some("a".to_string())), None)
324                )
325            }
326            _ => unreachable!("Expected BinaryExpr"),
327        };
328    }
329
330    #[test]
331    fn test_resolve_binary_expr_on_right() {
332        let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::Float64, false)]);
333        let expr = Expr::BinaryExpr(BinaryExpr {
334            left: Box::new(Expr::Column("a".to_string().into())),
335            op: Operator::Eq,
336            right: Box::new(Expr::BinaryExpr(BinaryExpr {
337                left: Box::new(Expr::Literal(ScalarValue::Int64(Some(2)), None)),
338                op: Operator::Minus,
339                right: Box::new(Expr::Literal(ScalarValue::Int64(Some(-1)), None)),
340            })),
341        });
342        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
343
344        match resolved {
345            Expr::BinaryExpr(be) => match be.right.as_ref() {
346                Expr::BinaryExpr(r_be) => {
347                    assert_eq!(
348                        r_be.left.as_ref(),
349                        &Expr::Literal(ScalarValue::Float64(Some(2.0)), None)
350                    );
351                    assert_eq!(
352                        r_be.right.as_ref(),
353                        &Expr::Literal(ScalarValue::Float64(Some(-1.0)), None)
354                    );
355                }
356                _ => panic!("Expected BinaryExpr"),
357            },
358            _ => panic!("Expected BinaryExpr"),
359        }
360    }
361
362    #[test]
363    fn test_resolve_in_expr() {
364        // Type coercion should apply for `A IN (0)` or `A NOT IN (0)`
365        let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::Float32, false)]);
366        let expr = Expr::in_list(
367            Expr::Column("a".to_string().into()),
368            vec![Expr::Literal(ScalarValue::Float64(Some(0.0)), None)],
369            false,
370        );
371        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
372        let expected = Expr::in_list(
373            Expr::Column("a".to_string().into()),
374            vec![Expr::Literal(ScalarValue::Float32(Some(0.0)), None)],
375            false,
376        );
377        assert_eq!(resolved, expected);
378
379        let expr = Expr::in_list(
380            Expr::Column("a".to_string().into()),
381            vec![Expr::Literal(ScalarValue::Float64(Some(0.0)), None)],
382            true,
383        );
384        let resolved = resolve_expr(&expr, &Schema::try_from(&arrow_schema).unwrap()).unwrap();
385        let expected = Expr::in_list(
386            Expr::Column("a".to_string().into()),
387            vec![Expr::Literal(ScalarValue::Float32(Some(0.0)), None)],
388            true,
389        );
390        assert_eq!(resolved, expected);
391    }
392
393    #[test]
394    fn test_resolve_column_type() {
395        let schema = Arc::new(ArrowSchema::new(vec![
396            Field::new("int", DataType::Int32, true),
397            Field::new(
398                "st",
399                DataType::Struct(
400                    vec![
401                        Field::new("str", DataType::Utf8, true),
402                        Field::new("Child.With.Dot", DataType::Int32, true),
403                        Field::new(
404                            "st",
405                            DataType::Struct(
406                                vec![Field::new("float", DataType::Float64, true)].into(),
407                            ),
408                            true,
409                        ),
410                    ]
411                    .into(),
412                ),
413                true,
414            ),
415        ]));
416        let schema = Schema::try_from(schema.as_ref()).unwrap();
417
418        assert_eq!(
419            resolve_column_type(&col("int"), &schema),
420            Some(DataType::Int32)
421        );
422        assert_eq!(
423            resolve_column_type(&col("st").field("str"), &schema),
424            Some(DataType::Utf8)
425        );
426        assert_eq!(
427            resolve_column_type(&col("st").field("st").field("float"), &schema),
428            Some(DataType::Float64)
429        );
430
431        assert_eq!(
432            resolve_column_type(&field_path_to_expr("st.`Child.With.Dot`").unwrap(), &schema),
433            Some(DataType::Int32)
434        );
435
436        assert_eq!(resolve_column_type(&col("x"), &schema), None);
437        assert_eq!(resolve_column_type(&col("str"), &schema), None);
438        assert_eq!(resolve_column_type(&col("float"), &schema), None);
439        assert_eq!(
440            resolve_column_type(&col("st").field("str").eq(lit("x")), &schema),
441            None
442        );
443    }
444
445    #[test]
446    fn test_resolve_utf8view_literal_against_utf8_column() {
447        // Simulates DataFusion 43+ producing a Utf8View literal (e.g. from md5())
448        // being compared against a Utf8 column stored in Lance.
449        let arrow_schema = ArrowSchema::new(vec![Field::new("hash", DataType::Utf8, false)]);
450        let schema = Schema::try_from(&arrow_schema).unwrap();
451
452        let expr = Expr::BinaryExpr(BinaryExpr {
453            left: Box::new(Expr::Column("hash".to_string().into())),
454            op: Operator::Eq,
455            right: Box::new(Expr::Literal(
456                ScalarValue::Utf8View(Some("abc".to_string())),
457                None,
458            )),
459        });
460
461        let resolved = resolve_expr(&expr, &schema).unwrap();
462        match resolved {
463            Expr::BinaryExpr(be) => {
464                assert_eq!(
465                    be.right.as_ref(),
466                    &Expr::Literal(ScalarValue::Utf8(Some("abc".to_string())), None)
467                )
468            }
469            _ => unreachable!("Expected BinaryExpr"),
470        }
471    }
472
473    #[test]
474    fn test_resolve_typed_null_against_dictionary_column() {
475        // A dictionary-encoded string column, e.g. a categorical field.
476        let dict_ty = DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8));
477        let arrow_schema = ArrowSchema::new(vec![Field::new("etld", dict_ty, true)]);
478        let schema = Schema::try_from(&arrow_schema).unwrap();
479
480        // A typed null must be wrapped in the dictionary type, not left as a bare
481        // `Utf8(None)` literal sitting next to a `Dictionary(...)` column.
482        let expected_null = Expr::Literal(
483            ScalarValue::Dictionary(Box::new(DataType::Int16), Box::new(ScalarValue::Utf8(None))),
484            None,
485        );
486
487        // `etld = <typed null>` built directly via the API, as opposed to coming
488        // through SQL parsing.
489        let expr = Expr::BinaryExpr(BinaryExpr {
490            left: Box::new(Expr::Column("etld".to_string().into())),
491            op: Operator::Eq,
492            right: Box::new(Expr::Literal(ScalarValue::Utf8(None), None)),
493        });
494        match resolve_expr(&expr, &schema).unwrap() {
495            Expr::BinaryExpr(be) => assert_eq!(be.right.as_ref(), &expected_null),
496            other => unreachable!("Expected BinaryExpr, got {other:?}"),
497        }
498
499        // `etld IN ('a', <typed null>)` — a typed value mixed with a typed null,
500        // both already typed as Utf8. Every list element is wrapped in the
501        // dictionary type.
502        let expr = Expr::in_list(
503            Expr::Column("etld".to_string().into()),
504            vec![
505                Expr::Literal(ScalarValue::Utf8(Some("a".to_string())), None),
506                Expr::Literal(ScalarValue::Utf8(None), None),
507            ],
508            false,
509        );
510        let expected = Expr::in_list(
511            Expr::Column("etld".to_string().into()),
512            vec![
513                Expr::Literal(
514                    ScalarValue::Dictionary(
515                        Box::new(DataType::Int16),
516                        Box::new(ScalarValue::Utf8(Some("a".to_string()))),
517                    ),
518                    None,
519                ),
520                expected_null,
521            ],
522            false,
523        );
524        assert_eq!(resolve_expr(&expr, &schema).unwrap(), expected);
525    }
526}
lance_datafusion/logical_expr.rs

lance_datafusion/
logical_expr.rs
