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

lance_datafusion/
logical_expr.rs
