lance_datafusion/
planner.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Exec plan planner
5
6use std::borrow::Cow;
7use std::collections::{BTreeSet, VecDeque};
8use std::sync::Arc;
9
10use crate::expr::safe_coerce_scalar;
11use crate::logical_expr::{coerce_filter_type_to_boolean, get_as_string_scalar_opt, resolve_expr};
12use crate::sql::{parse_sql_expr, parse_sql_filter};
13use arrow::compute::CastOptions;
14use arrow_array::ListArray;
15use arrow_buffer::OffsetBuffer;
16use arrow_schema::{DataType as ArrowDataType, Field, SchemaRef, TimeUnit};
17use arrow_select::concat::concat;
18use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
19use datafusion::common::DFSchema;
20use datafusion::config::ConfigOptions;
21use datafusion::error::Result as DFResult;
22use datafusion::execution::config::SessionConfig;
23use datafusion::execution::context::SessionState;
24use datafusion::execution::runtime_env::RuntimeEnvBuilder;
25use datafusion::execution::session_state::SessionStateBuilder;
26use datafusion::logical_expr::expr::ScalarFunction;
27use datafusion::logical_expr::planner::{ExprPlanner, PlannerResult, RawFieldAccessExpr};
28use datafusion::logical_expr::{
29    AggregateUDF, ColumnarValue, GetFieldAccess, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
30    Signature, Volatility, WindowUDF,
31};
32use datafusion::optimizer::simplify_expressions::SimplifyContext;
33use datafusion::sql::planner::{ContextProvider, ParserOptions, PlannerContext, SqlToRel};
34use datafusion::sql::sqlparser::ast::{
35    AccessExpr, Array as SQLArray, BinaryOperator, DataType as SQLDataType, ExactNumberInfo,
36    Expr as SQLExpr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Ident,
37    ObjectNamePart, Subscript, TimezoneInfo, UnaryOperator, Value, ValueWithSpan,
38};
39use datafusion::{
40    common::Column,
41    logical_expr::{col, Between, BinaryExpr, Like, Operator},
42    physical_expr::execution_props::ExecutionProps,
43    physical_plan::PhysicalExpr,
44    prelude::Expr,
45    scalar::ScalarValue,
46};
47use datafusion_functions::core::getfield::GetFieldFunc;
48use lance_arrow::cast::cast_with_options;
49use lance_core::datatypes::Schema;
50use lance_core::error::LanceOptionExt;
51use snafu::location;
52
53use lance_core::{Error, Result};
54
55#[derive(Debug, Clone)]
56struct CastListF16Udf {
57    signature: Signature,
58}
59
60impl CastListF16Udf {
61    pub fn new() -> Self {
62        Self {
63            signature: Signature::any(1, Volatility::Immutable),
64        }
65    }
66}
67
68impl ScalarUDFImpl for CastListF16Udf {
69    fn as_any(&self) -> &dyn std::any::Any {
70        self
71    }
72
73    fn name(&self) -> &str {
74        "_cast_list_f16"
75    }
76
77    fn signature(&self) -> &Signature {
78        &self.signature
79    }
80
81    fn return_type(&self, arg_types: &[ArrowDataType]) -> DFResult<ArrowDataType> {
82        let input = &arg_types[0];
83        match input {
84            ArrowDataType::FixedSizeList(field, size) => {
85                if field.data_type() != &ArrowDataType::Float32
86                    && field.data_type() != &ArrowDataType::Float16
87                {
88                    return Err(datafusion::error::DataFusionError::Execution(
89                        "cast_list_f16 only supports list of float32 or float16".to_string(),
90                    ));
91                }
92                Ok(ArrowDataType::FixedSizeList(
93                    Arc::new(Field::new(
94                        field.name(),
95                        ArrowDataType::Float16,
96                        field.is_nullable(),
97                    )),
98                    *size,
99                ))
100            }
101            ArrowDataType::List(field) => {
102                if field.data_type() != &ArrowDataType::Float32
103                    && field.data_type() != &ArrowDataType::Float16
104                {
105                    return Err(datafusion::error::DataFusionError::Execution(
106                        "cast_list_f16 only supports list of float32 or float16".to_string(),
107                    ));
108                }
109                Ok(ArrowDataType::List(Arc::new(Field::new(
110                    field.name(),
111                    ArrowDataType::Float16,
112                    field.is_nullable(),
113                ))))
114            }
115            _ => Err(datafusion::error::DataFusionError::Execution(
116                "cast_list_f16 only supports FixedSizeList/List arguments".to_string(),
117            )),
118        }
119    }
120
121    fn invoke_with_args(&self, func_args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
122        let ColumnarValue::Array(arr) = &func_args.args[0] else {
123            return Err(datafusion::error::DataFusionError::Execution(
124                "cast_list_f16 only supports array arguments".to_string(),
125            ));
126        };
127
128        let to_type = match arr.data_type() {
129            ArrowDataType::FixedSizeList(field, size) => ArrowDataType::FixedSizeList(
130                Arc::new(Field::new(
131                    field.name(),
132                    ArrowDataType::Float16,
133                    field.is_nullable(),
134                )),
135                *size,
136            ),
137            ArrowDataType::List(field) => ArrowDataType::List(Arc::new(Field::new(
138                field.name(),
139                ArrowDataType::Float16,
140                field.is_nullable(),
141            ))),
142            _ => {
143                return Err(datafusion::error::DataFusionError::Execution(
144                    "cast_list_f16 only supports array arguments".to_string(),
145                ));
146            }
147        };
148
149        let res = cast_with_options(arr.as_ref(), &to_type, &CastOptions::default())?;
150        Ok(ColumnarValue::Array(res))
151    }
152}
153
154// Adapter that instructs datafusion how lance expects expressions to be interpreted
155struct LanceContextProvider {
156    options: datafusion::config::ConfigOptions,
157    state: SessionState,
158    expr_planners: Vec<Arc<dyn ExprPlanner>>,
159}
160
161impl Default for LanceContextProvider {
162    fn default() -> Self {
163        let config = SessionConfig::new();
164        let runtime = RuntimeEnvBuilder::new().build_arc().unwrap();
165        let mut state_builder = SessionStateBuilder::new()
166            .with_config(config)
167            .with_runtime_env(runtime)
168            .with_default_features();
169
170        // SessionState does not expose expr_planners, so we need to get the default ones from
171        // the builder and store them to return from get_expr_planners
172
173        // unwrap safe because with_default_features sets expr_planners
174        let expr_planners = state_builder.expr_planners().as_ref().unwrap().clone();
175
176        Self {
177            options: ConfigOptions::default(),
178            state: state_builder.build(),
179            expr_planners,
180        }
181    }
182}
183
184impl ContextProvider for LanceContextProvider {
185    fn get_table_source(
186        &self,
187        name: datafusion::sql::TableReference,
188    ) -> DFResult<Arc<dyn datafusion::logical_expr::TableSource>> {
189        Err(datafusion::error::DataFusionError::NotImplemented(format!(
190            "Attempt to reference inner table {} not supported",
191            name
192        )))
193    }
194
195    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
196        self.state.aggregate_functions().get(name).cloned()
197    }
198
199    fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
200        self.state.window_functions().get(name).cloned()
201    }
202
203    fn get_function_meta(&self, f: &str) -> Option<Arc<ScalarUDF>> {
204        match f {
205            // TODO: cast should go thru CAST syntax instead of UDF
206            // Going thru UDF makes it hard for the optimizer to find no-ops
207            "_cast_list_f16" => Some(Arc::new(ScalarUDF::new_from_impl(CastListF16Udf::new()))),
208            _ => self.state.scalar_functions().get(f).cloned(),
209        }
210    }
211
212    fn get_variable_type(&self, _: &[String]) -> Option<ArrowDataType> {
213        // Variables (things like @@LANGUAGE) not supported
214        None
215    }
216
217    fn options(&self) -> &datafusion::config::ConfigOptions {
218        &self.options
219    }
220
221    fn udf_names(&self) -> Vec<String> {
222        self.state.scalar_functions().keys().cloned().collect()
223    }
224
225    fn udaf_names(&self) -> Vec<String> {
226        self.state.aggregate_functions().keys().cloned().collect()
227    }
228
229    fn udwf_names(&self) -> Vec<String> {
230        self.state.window_functions().keys().cloned().collect()
231    }
232
233    fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
234        &self.expr_planners
235    }
236}
237
238pub struct Planner {
239    schema: SchemaRef,
240    context_provider: LanceContextProvider,
241}
242
243impl Planner {
244    pub fn new(schema: SchemaRef) -> Self {
245        Self {
246            schema,
247            context_provider: LanceContextProvider::default(),
248        }
249    }
250
251    fn column(idents: &[Ident]) -> Expr {
252        let mut column = col(&idents[0].value);
253        for ident in &idents[1..] {
254            column = Expr::ScalarFunction(ScalarFunction {
255                args: vec![
256                    column,
257                    Expr::Literal(ScalarValue::Utf8(Some(ident.value.clone()))),
258                ],
259                func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
260            });
261        }
262        column
263    }
264
265    fn binary_op(&self, op: &BinaryOperator) -> Result<Operator> {
266        Ok(match op {
267            BinaryOperator::Plus => Operator::Plus,
268            BinaryOperator::Minus => Operator::Minus,
269            BinaryOperator::Multiply => Operator::Multiply,
270            BinaryOperator::Divide => Operator::Divide,
271            BinaryOperator::Modulo => Operator::Modulo,
272            BinaryOperator::StringConcat => Operator::StringConcat,
273            BinaryOperator::Gt => Operator::Gt,
274            BinaryOperator::Lt => Operator::Lt,
275            BinaryOperator::GtEq => Operator::GtEq,
276            BinaryOperator::LtEq => Operator::LtEq,
277            BinaryOperator::Eq => Operator::Eq,
278            BinaryOperator::NotEq => Operator::NotEq,
279            BinaryOperator::And => Operator::And,
280            BinaryOperator::Or => Operator::Or,
281            _ => {
282                return Err(Error::invalid_input(
283                    format!("Operator {op} is not supported"),
284                    location!(),
285                ));
286            }
287        })
288    }
289
290    fn binary_expr(&self, left: &SQLExpr, op: &BinaryOperator, right: &SQLExpr) -> Result<Expr> {
291        Ok(Expr::BinaryExpr(BinaryExpr::new(
292            Box::new(self.parse_sql_expr(left)?),
293            self.binary_op(op)?,
294            Box::new(self.parse_sql_expr(right)?),
295        )))
296    }
297
298    fn unary_expr(&self, op: &UnaryOperator, expr: &SQLExpr) -> Result<Expr> {
299        Ok(match op {
300            UnaryOperator::Not | UnaryOperator::PGBitwiseNot => {
301                Expr::Not(Box::new(self.parse_sql_expr(expr)?))
302            }
303
304            UnaryOperator::Minus => {
305                use datafusion::logical_expr::lit;
306                match expr {
307                    SQLExpr::Value(ValueWithSpan { value: Value::Number(n, _), ..}) => match n.parse::<i64>() {
308                        Ok(n) => lit(-n),
309                        Err(_) => lit(-n
310                            .parse::<f64>()
311                            .map_err(|_e| {
312                                Error::invalid_input(
313                                    format!("negative operator can be only applied to integer and float operands, got: {n}"),
314                                    location!(),
315                                )
316                            })?),
317                    },
318                    _ => {
319                        Expr::Negative(Box::new(self.parse_sql_expr(expr)?))
320                    }
321                }
322            }
323
324            _ => {
325                return Err(Error::invalid_input(
326                    format!("Unary operator '{:?}' is not supported", op),
327                    location!(),
328                ));
329            }
330        })
331    }
332
333    // See datafusion `sqlToRel::parse_sql_number()`
334    fn number(&self, value: &str, negative: bool) -> Result<Expr> {
335        use datafusion::logical_expr::lit;
336        let value: Cow<str> = if negative {
337            Cow::Owned(format!("-{}", value))
338        } else {
339            Cow::Borrowed(value)
340        };
341        if let Ok(n) = value.parse::<i64>() {
342            Ok(lit(n))
343        } else {
344            value.parse::<f64>().map(lit).map_err(|_| {
345                Error::invalid_input(
346                    format!("'{value}' is not supported number value."),
347                    location!(),
348                )
349            })
350        }
351    }
352
353    fn value(&self, value: &Value) -> Result<Expr> {
354        Ok(match value {
355            Value::Number(v, _) => self.number(v.as_str(), false)?,
356            Value::SingleQuotedString(s) => Expr::Literal(ScalarValue::Utf8(Some(s.clone()))),
357            Value::HexStringLiteral(hsl) => {
358                Expr::Literal(ScalarValue::Binary(Self::try_decode_hex_literal(hsl)))
359            }
360            Value::DoubleQuotedString(s) => Expr::Literal(ScalarValue::Utf8(Some(s.clone()))),
361            Value::Boolean(v) => Expr::Literal(ScalarValue::Boolean(Some(*v))),
362            Value::Null => Expr::Literal(ScalarValue::Null),
363            _ => todo!(),
364        })
365    }
366
367    fn parse_function_args(&self, func_args: &FunctionArg) -> Result<Expr> {
368        match func_args {
369            FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => self.parse_sql_expr(expr),
370            _ => Err(Error::invalid_input(
371                format!("Unsupported function args: {:?}", func_args),
372                location!(),
373            )),
374        }
375    }
376
377    // We now use datafusion to parse functions.  This allows us to use datafusion's
378    // entire collection of functions (previously we had just hard-coded support for two functions).
379    //
380    // Unfortunately, one of those two functions was is_valid and the reason we needed it was because
381    // this is a function that comes from duckdb.  Datafusion does not consider is_valid to be a function
382    // but rather an AST node (Expr::IsNotNull) and so we need to handle this case specially.
383    fn legacy_parse_function(&self, func: &Function) -> Result<Expr> {
384        match &func.args {
385            FunctionArguments::List(args) => {
386                if func.name.0.len() != 1 {
387                    return Err(Error::invalid_input(
388                        format!("Function name must have 1 part, got: {:?}", func.name.0),
389                        location!(),
390                    ));
391                }
392                Ok(Expr::IsNotNull(Box::new(
393                    self.parse_function_args(&args.args[0])?,
394                )))
395            }
396            _ => Err(Error::invalid_input(
397                format!("Unsupported function args: {:?}", &func.args),
398                location!(),
399            )),
400        }
401    }
402
403    fn parse_function(&self, function: SQLExpr) -> Result<Expr> {
404        if let SQLExpr::Function(function) = &function {
405            if let Some(ObjectNamePart::Identifier(name)) = &function.name.0.first() {
406                if &name.value == "is_valid" {
407                    return self.legacy_parse_function(function);
408                }
409            }
410        }
411        let sql_to_rel = SqlToRel::new_with_options(
412            &self.context_provider,
413            ParserOptions {
414                parse_float_as_decimal: false,
415                enable_ident_normalization: false,
416                support_varchar_with_length: false,
417                enable_options_value_normalization: false,
418                collect_spans: false,
419                map_varchar_to_utf8view: false,
420            },
421        );
422
423        let mut planner_context = PlannerContext::default();
424        let schema = DFSchema::try_from(self.schema.as_ref().clone())?;
425        Ok(sql_to_rel.sql_to_expr(function, &schema, &mut planner_context)?)
426    }
427
428    fn parse_type(&self, data_type: &SQLDataType) -> Result<ArrowDataType> {
429        const SUPPORTED_TYPES: [&str; 13] = [
430            "int [unsigned]",
431            "tinyint [unsigned]",
432            "smallint [unsigned]",
433            "bigint [unsigned]",
434            "float",
435            "double",
436            "string",
437            "binary",
438            "date",
439            "timestamp(precision)",
440            "datetime(precision)",
441            "decimal(precision,scale)",
442            "boolean",
443        ];
444        match data_type {
445            SQLDataType::String(_) => Ok(ArrowDataType::Utf8),
446            SQLDataType::Binary(_) => Ok(ArrowDataType::Binary),
447            SQLDataType::Float(_) => Ok(ArrowDataType::Float32),
448            SQLDataType::Double(_) => Ok(ArrowDataType::Float64),
449            SQLDataType::Boolean => Ok(ArrowDataType::Boolean),
450            SQLDataType::TinyInt(_) => Ok(ArrowDataType::Int8),
451            SQLDataType::SmallInt(_) => Ok(ArrowDataType::Int16),
452            SQLDataType::Int(_) | SQLDataType::Integer(_) => Ok(ArrowDataType::Int32),
453            SQLDataType::BigInt(_) => Ok(ArrowDataType::Int64),
454            SQLDataType::TinyIntUnsigned(_) => Ok(ArrowDataType::UInt8),
455            SQLDataType::SmallIntUnsigned(_) => Ok(ArrowDataType::UInt16),
456            SQLDataType::IntUnsigned(_) | SQLDataType::IntegerUnsigned(_) => {
457                Ok(ArrowDataType::UInt32)
458            }
459            SQLDataType::BigIntUnsigned(_) => Ok(ArrowDataType::UInt64),
460            SQLDataType::Date => Ok(ArrowDataType::Date32),
461            SQLDataType::Timestamp(resolution, tz) => {
462                match tz {
463                    TimezoneInfo::None => {}
464                    _ => {
465                        return Err(Error::invalid_input(
466                            "Timezone not supported in timestamp".to_string(),
467                            location!(),
468                        ));
469                    }
470                };
471                let time_unit = match resolution {
472                    // Default to microsecond to match PyArrow
473                    None => TimeUnit::Microsecond,
474                    Some(0) => TimeUnit::Second,
475                    Some(3) => TimeUnit::Millisecond,
476                    Some(6) => TimeUnit::Microsecond,
477                    Some(9) => TimeUnit::Nanosecond,
478                    _ => {
479                        return Err(Error::invalid_input(
480                            format!("Unsupported datetime resolution: {:?}", resolution),
481                            location!(),
482                        ));
483                    }
484                };
485                Ok(ArrowDataType::Timestamp(time_unit, None))
486            }
487            SQLDataType::Datetime(resolution) => {
488                let time_unit = match resolution {
489                    None => TimeUnit::Microsecond,
490                    Some(0) => TimeUnit::Second,
491                    Some(3) => TimeUnit::Millisecond,
492                    Some(6) => TimeUnit::Microsecond,
493                    Some(9) => TimeUnit::Nanosecond,
494                    _ => {
495                        return Err(Error::invalid_input(
496                            format!("Unsupported datetime resolution: {:?}", resolution),
497                            location!(),
498                        ));
499                    }
500                };
501                Ok(ArrowDataType::Timestamp(time_unit, None))
502            }
503            SQLDataType::Decimal(number_info) => match number_info {
504                ExactNumberInfo::PrecisionAndScale(precision, scale) => {
505                    Ok(ArrowDataType::Decimal128(*precision as u8, *scale as i8))
506                }
507                _ => Err(Error::invalid_input(
508                    format!(
509                        "Must provide precision and scale for decimal: {:?}",
510                        number_info
511                    ),
512                    location!(),
513                )),
514            },
515            _ => Err(Error::invalid_input(
516                format!(
517                    "Unsupported data type: {:?}. Supported types: {:?}",
518                    data_type, SUPPORTED_TYPES
519                ),
520                location!(),
521            )),
522        }
523    }
524
525    fn plan_field_access(&self, mut field_access_expr: RawFieldAccessExpr) -> Result<Expr> {
526        let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
527        for planner in self.context_provider.get_expr_planners() {
528            match planner.plan_field_access(field_access_expr, &df_schema)? {
529                PlannerResult::Planned(expr) => return Ok(expr),
530                PlannerResult::Original(expr) => {
531                    field_access_expr = expr;
532                }
533            }
534        }
535        Err(Error::invalid_input(
536            "Field access could not be planned",
537            location!(),
538        ))
539    }
540
541    fn parse_sql_expr(&self, expr: &SQLExpr) -> Result<Expr> {
542        match expr {
543            SQLExpr::Identifier(id) => {
544                // Users can pass string literals wrapped in `"`.
545                // (Normally SQL only allows single quotes.)
546                if id.quote_style == Some('"') {
547                    Ok(Expr::Literal(ScalarValue::Utf8(Some(id.value.clone()))))
548                // Users can wrap identifiers with ` to reference non-standard
549                // names, such as uppercase or spaces.
550                } else if id.quote_style == Some('`') {
551                    Ok(Expr::Column(Column::from_name(id.value.clone())))
552                } else {
553                    Ok(Self::column(vec![id.clone()].as_slice()))
554                }
555            }
556            SQLExpr::CompoundIdentifier(ids) => Ok(Self::column(ids.as_slice())),
557            SQLExpr::BinaryOp { left, op, right } => self.binary_expr(left, op, right),
558            SQLExpr::UnaryOp { op, expr } => self.unary_expr(op, expr),
559            SQLExpr::Value(value) => self.value(&value.value),
560            SQLExpr::Array(SQLArray { elem, .. }) => {
561                let mut values = vec![];
562
563                let array_literal_error = |pos: usize, value: &_| {
564                    Err(Error::invalid_input(
565                        format!(
566                            "Expected a literal value in array, instead got {} at position {}",
567                            value, pos
568                        ),
569                        location!(),
570                    ))
571                };
572
573                for (pos, expr) in elem.iter().enumerate() {
574                    match expr {
575                        SQLExpr::Value(value) => {
576                            if let Expr::Literal(value) = self.value(&value.value)? {
577                                values.push(value);
578                            } else {
579                                return array_literal_error(pos, expr);
580                            }
581                        }
582                        SQLExpr::UnaryOp {
583                            op: UnaryOperator::Minus,
584                            expr,
585                        } => {
586                            if let SQLExpr::Value(ValueWithSpan {
587                                value: Value::Number(number, _),
588                                ..
589                            }) = expr.as_ref()
590                            {
591                                if let Expr::Literal(value) = self.number(number, true)? {
592                                    values.push(value);
593                                } else {
594                                    return array_literal_error(pos, expr);
595                                }
596                            } else {
597                                return array_literal_error(pos, expr);
598                            }
599                        }
600                        _ => {
601                            return array_literal_error(pos, expr);
602                        }
603                    }
604                }
605
606                let field = if !values.is_empty() {
607                    let data_type = values[0].data_type();
608
609                    for value in &mut values {
610                        if value.data_type() != data_type {
611                            *value = safe_coerce_scalar(value, &data_type).ok_or_else(|| Error::invalid_input(
612                                format!("Array expressions must have a consistent datatype. Expected: {}, got: {}", data_type, value.data_type()),
613                                location!()
614                            ))?;
615                        }
616                    }
617                    Field::new("item", data_type, true)
618                } else {
619                    Field::new("item", ArrowDataType::Null, true)
620                };
621
622                let values = values
623                    .into_iter()
624                    .map(|v| v.to_array().map_err(Error::from))
625                    .collect::<Result<Vec<_>>>()?;
626                let array_refs = values.iter().map(|v| v.as_ref()).collect::<Vec<_>>();
627                let values = concat(&array_refs)?;
628                let values = ListArray::try_new(
629                    field.into(),
630                    OffsetBuffer::from_lengths([values.len()]),
631                    values,
632                    None,
633                )?;
634
635                Ok(Expr::Literal(ScalarValue::List(Arc::new(values))))
636            }
637            // For example, DATE '2020-01-01'
638            SQLExpr::TypedString { data_type, value } => {
639                let value = value.clone().into_string().expect_ok()?;
640                Ok(Expr::Cast(datafusion::logical_expr::Cast {
641                    expr: Box::new(Expr::Literal(ScalarValue::Utf8(Some(value)))),
642                    data_type: self.parse_type(data_type)?,
643                }))
644            }
645            SQLExpr::IsFalse(expr) => Ok(Expr::IsFalse(Box::new(self.parse_sql_expr(expr)?))),
646            SQLExpr::IsNotFalse(expr) => Ok(Expr::IsNotFalse(Box::new(self.parse_sql_expr(expr)?))),
647            SQLExpr::IsTrue(expr) => Ok(Expr::IsTrue(Box::new(self.parse_sql_expr(expr)?))),
648            SQLExpr::IsNotTrue(expr) => Ok(Expr::IsNotTrue(Box::new(self.parse_sql_expr(expr)?))),
649            SQLExpr::IsNull(expr) => Ok(Expr::IsNull(Box::new(self.parse_sql_expr(expr)?))),
650            SQLExpr::IsNotNull(expr) => Ok(Expr::IsNotNull(Box::new(self.parse_sql_expr(expr)?))),
651            SQLExpr::InList {
652                expr,
653                list,
654                negated,
655            } => {
656                let value_expr = self.parse_sql_expr(expr)?;
657                let list_exprs = list
658                    .iter()
659                    .map(|e| self.parse_sql_expr(e))
660                    .collect::<Result<Vec<_>>>()?;
661                Ok(value_expr.in_list(list_exprs, *negated))
662            }
663            SQLExpr::Nested(inner) => self.parse_sql_expr(inner.as_ref()),
664            SQLExpr::Function(_) => self.parse_function(expr.clone()),
665            SQLExpr::ILike {
666                negated,
667                expr,
668                pattern,
669                escape_char,
670                any: _,
671            } => Ok(Expr::Like(Like::new(
672                *negated,
673                Box::new(self.parse_sql_expr(expr)?),
674                Box::new(self.parse_sql_expr(pattern)?),
675                escape_char.as_ref().and_then(|c| c.chars().next()),
676                true,
677            ))),
678            SQLExpr::Like {
679                negated,
680                expr,
681                pattern,
682                escape_char,
683                any: _,
684            } => Ok(Expr::Like(Like::new(
685                *negated,
686                Box::new(self.parse_sql_expr(expr)?),
687                Box::new(self.parse_sql_expr(pattern)?),
688                escape_char.as_ref().and_then(|c| c.chars().next()),
689                false,
690            ))),
691            SQLExpr::Cast {
692                expr, data_type, ..
693            } => Ok(Expr::Cast(datafusion::logical_expr::Cast {
694                expr: Box::new(self.parse_sql_expr(expr)?),
695                data_type: self.parse_type(data_type)?,
696            })),
697            SQLExpr::JsonAccess { .. } => Err(Error::invalid_input(
698                "JSON access is not supported",
699                location!(),
700            )),
701            SQLExpr::CompoundFieldAccess { root, access_chain } => {
702                let mut expr = self.parse_sql_expr(root)?;
703
704                for access in access_chain {
705                    let field_access = match access {
706                        // x.y or x['y']
707                        AccessExpr::Dot(SQLExpr::Identifier(Ident { value: s, .. }))
708                        | AccessExpr::Subscript(Subscript::Index {
709                            index:
710                                SQLExpr::Value(ValueWithSpan {
711                                    value:
712                                        Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
713                                    ..
714                                }),
715                        }) => GetFieldAccess::NamedStructField {
716                            name: ScalarValue::from(s.as_str()),
717                        },
718                        AccessExpr::Subscript(Subscript::Index { index }) => {
719                            let key = Box::new(self.parse_sql_expr(index)?);
720                            GetFieldAccess::ListIndex { key }
721                        }
722                        AccessExpr::Subscript(Subscript::Slice { .. }) => {
723                            return Err(Error::invalid_input(
724                                "Slice subscript is not supported",
725                                location!(),
726                            ));
727                        }
728                        _ => {
729                            // Handle other cases like JSON access
730                            // Note: JSON access is not supported in lance
731                            return Err(Error::invalid_input(
732                                "Only dot notation or index access is supported for field access",
733                                location!(),
734                            ));
735                        }
736                    };
737
738                    let field_access_expr = RawFieldAccessExpr { expr, field_access };
739                    expr = self.plan_field_access(field_access_expr)?;
740                }
741
742                Ok(expr)
743            }
744            SQLExpr::Between {
745                expr,
746                negated,
747                low,
748                high,
749            } => {
750                // Parse the main expression and bounds
751                let expr = self.parse_sql_expr(expr)?;
752                let low = self.parse_sql_expr(low)?;
753                let high = self.parse_sql_expr(high)?;
754
755                let between = Expr::Between(Between::new(
756                    Box::new(expr),
757                    *negated,
758                    Box::new(low),
759                    Box::new(high),
760                ));
761                Ok(between)
762            }
763            _ => Err(Error::invalid_input(
764                format!("Expression '{expr}' is not supported SQL in lance"),
765                location!(),
766            )),
767        }
768    }
769
770    /// Create Logical [Expr] from a SQL filter clause.
771    ///
772    /// Note: the returned expression must be passed through [optimize_expr()]
773    /// before being passed to [create_physical_expr()].
774    pub fn parse_filter(&self, filter: &str) -> Result<Expr> {
775        // Allow sqlparser to parse filter as part of ONE SQL statement.
776        let ast_expr = parse_sql_filter(filter)?;
777        let expr = self.parse_sql_expr(&ast_expr)?;
778        let schema = Schema::try_from(self.schema.as_ref())?;
779        let resolved = resolve_expr(&expr, &schema)?;
780        coerce_filter_type_to_boolean(resolved)
781    }
782
783    /// Create Logical [Expr] from a SQL expression.
784    ///
785    /// Note: the returned expression must be passed through [optimize_filter()]
786    /// before being passed to [create_physical_expr()].
787    pub fn parse_expr(&self, expr: &str) -> Result<Expr> {
788        let ast_expr = parse_sql_expr(expr)?;
789        let expr = self.parse_sql_expr(&ast_expr)?;
790        let schema = Schema::try_from(self.schema.as_ref())?;
791        let resolved = resolve_expr(&expr, &schema)?;
792        Ok(resolved)
793    }
794
795    /// Try to decode bytes from hex literal string.
796    ///
797    /// Copied from datafusion because this is not public.
798    ///
799    /// TODO: use SqlToRel from Datafusion directly?
800    fn try_decode_hex_literal(s: &str) -> Option<Vec<u8>> {
801        let hex_bytes = s.as_bytes();
802        let mut decoded_bytes = Vec::with_capacity(hex_bytes.len().div_ceil(2));
803
804        let start_idx = hex_bytes.len() % 2;
805        if start_idx > 0 {
806            // The first byte is formed of only one char.
807            decoded_bytes.push(Self::try_decode_hex_char(hex_bytes[0])?);
808        }
809
810        for i in (start_idx..hex_bytes.len()).step_by(2) {
811            let high = Self::try_decode_hex_char(hex_bytes[i])?;
812            let low = Self::try_decode_hex_char(hex_bytes[i + 1])?;
813            decoded_bytes.push((high << 4) | low);
814        }
815
816        Some(decoded_bytes)
817    }
818
819    /// Try to decode a byte from a hex char.
820    ///
821    /// None will be returned if the input char is hex-invalid.
822    const fn try_decode_hex_char(c: u8) -> Option<u8> {
823        match c {
824            b'A'..=b'F' => Some(c - b'A' + 10),
825            b'a'..=b'f' => Some(c - b'a' + 10),
826            b'0'..=b'9' => Some(c - b'0'),
827            _ => None,
828        }
829    }
830
831    /// Optimize the filter expression and coerce data types.
832    pub fn optimize_expr(&self, expr: Expr) -> Result<Expr> {
833        let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?);
834
835        // DataFusion needs the simplify and coerce passes to be applied before
836        // expressions can be handled by the physical planner.
837        let props = ExecutionProps::default();
838        let simplify_context = SimplifyContext::new(&props).with_schema(df_schema.clone());
839        let simplifier =
840            datafusion::optimizer::simplify_expressions::ExprSimplifier::new(simplify_context);
841
842        let expr = simplifier.simplify(expr)?;
843        let expr = simplifier.coerce(expr, &df_schema)?;
844
845        Ok(expr)
846    }
847
848    /// Create the [`PhysicalExpr`] from a logical [`Expr`]
849    pub fn create_physical_expr(&self, expr: &Expr) -> Result<Arc<dyn PhysicalExpr>> {
850        let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?);
851
852        Ok(datafusion::physical_expr::create_physical_expr(
853            expr,
854            df_schema.as_ref(),
855            &Default::default(),
856        )?)
857    }
858
859    /// Collect the columns in the expression.
860    ///
861    /// The columns are returned in sorted order.
862    pub fn column_names_in_expr(expr: &Expr) -> Vec<String> {
863        let mut visitor = ColumnCapturingVisitor {
864            current_path: VecDeque::new(),
865            columns: BTreeSet::new(),
866        };
867        expr.visit(&mut visitor).unwrap();
868        visitor.columns.into_iter().collect()
869    }
870}
871
872struct ColumnCapturingVisitor {
873    // Current column path. If this is empty, we are not in a column expression.
874    current_path: VecDeque<String>,
875    columns: BTreeSet<String>,
876}
877
878impl TreeNodeVisitor<'_> for ColumnCapturingVisitor {
879    type Node = Expr;
880
881    fn f_down(&mut self, node: &Self::Node) -> DFResult<TreeNodeRecursion> {
882        match node {
883            Expr::Column(Column { name, .. }) => {
884                let mut path = name.clone();
885                for part in self.current_path.drain(..) {
886                    path.push('.');
887                    path.push_str(&part);
888                }
889                self.columns.insert(path);
890                self.current_path.clear();
891            }
892            Expr::ScalarFunction(udf) => {
893                if udf.name() == GetFieldFunc::default().name() {
894                    if let Some(name) = get_as_string_scalar_opt(&udf.args[1]) {
895                        self.current_path.push_front(name.to_string())
896                    } else {
897                        self.current_path.clear();
898                    }
899                } else {
900                    self.current_path.clear();
901                }
902            }
903            _ => {
904                self.current_path.clear();
905            }
906        }
907
908        Ok(TreeNodeRecursion::Continue)
909    }
910}
911
912#[cfg(test)]
913mod tests {
914
915    use crate::logical_expr::ExprExt;
916
917    use super::*;
918
919    use arrow::datatypes::Float64Type;
920    use arrow_array::{
921        ArrayRef, BooleanArray, Float32Array, Int32Array, Int64Array, RecordBatch, StringArray,
922        StructArray, TimestampMicrosecondArray, TimestampMillisecondArray,
923        TimestampNanosecondArray, TimestampSecondArray,
924    };
925    use arrow_schema::{DataType, Fields, Schema};
926    use datafusion::{
927        logical_expr::{lit, Cast},
928        prelude::{array_element, get_field},
929    };
930    use datafusion_functions::core::expr_ext::FieldAccessor;
931
932    #[test]
933    fn test_parse_filter_simple() {
934        let schema = Arc::new(Schema::new(vec![
935            Field::new("i", DataType::Int32, false),
936            Field::new("s", DataType::Utf8, true),
937            Field::new(
938                "st",
939                DataType::Struct(Fields::from(vec![
940                    Field::new("x", DataType::Float32, false),
941                    Field::new("y", DataType::Float32, false),
942                ])),
943                true,
944            ),
945        ]));
946
947        let planner = Planner::new(schema.clone());
948
949        let expected = col("i")
950            .gt(lit(3_i32))
951            .and(col("st").field_newstyle("x").lt_eq(lit(5.0_f32)))
952            .and(
953                col("s")
954                    .eq(lit("str-4"))
955                    .or(col("s").in_list(vec![lit("str-4"), lit("str-5")], false)),
956            );
957
958        // double quotes
959        let expr = planner
960            .parse_filter("i > 3 AND st.x <= 5.0 AND (s == 'str-4' OR s in ('str-4', 'str-5'))")
961            .unwrap();
962        assert_eq!(expr, expected);
963
964        // single quote
965        let expr = planner
966            .parse_filter("i > 3 AND st.x <= 5.0 AND (s = 'str-4' OR s in ('str-4', 'str-5'))")
967            .unwrap();
968
969        let physical_expr = planner.create_physical_expr(&expr).unwrap();
970
971        let batch = RecordBatch::try_new(
972            schema,
973            vec![
974                Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef,
975                Arc::new(StringArray::from_iter_values(
976                    (0..10).map(|v| format!("str-{}", v)),
977                )),
978                Arc::new(StructArray::from(vec![
979                    (
980                        Arc::new(Field::new("x", DataType::Float32, false)),
981                        Arc::new(Float32Array::from_iter_values((0..10).map(|v| v as f32)))
982                            as ArrayRef,
983                    ),
984                    (
985                        Arc::new(Field::new("y", DataType::Float32, false)),
986                        Arc::new(Float32Array::from_iter_values(
987                            (0..10).map(|v| (v * 10) as f32),
988                        )),
989                    ),
990                ])),
991            ],
992        )
993        .unwrap();
994        let predicates = physical_expr.evaluate(&batch).unwrap();
995        assert_eq!(
996            predicates.into_array(0).unwrap().as_ref(),
997            &BooleanArray::from(vec![
998                false, false, false, false, true, true, false, false, false, false
999            ])
1000        );
1001    }
1002
1003    #[test]
1004    fn test_nested_col_refs() {
1005        let schema = Arc::new(Schema::new(vec![
1006            Field::new("s0", DataType::Utf8, true),
1007            Field::new(
1008                "st",
1009                DataType::Struct(Fields::from(vec![
1010                    Field::new("s1", DataType::Utf8, true),
1011                    Field::new(
1012                        "st",
1013                        DataType::Struct(Fields::from(vec![Field::new(
1014                            "s2",
1015                            DataType::Utf8,
1016                            true,
1017                        )])),
1018                        true,
1019                    ),
1020                ])),
1021                true,
1022            ),
1023        ]));
1024
1025        let planner = Planner::new(schema);
1026
1027        fn assert_column_eq(planner: &Planner, expr: &str, expected: &Expr) {
1028            let expr = planner.parse_filter(&format!("{expr} = 'val'")).unwrap();
1029            assert!(matches!(
1030                expr,
1031                Expr::BinaryExpr(BinaryExpr {
1032                    left: _,
1033                    op: Operator::Eq,
1034                    right: _
1035                })
1036            ));
1037            if let Expr::BinaryExpr(BinaryExpr { left, .. }) = expr {
1038                assert_eq!(left.as_ref(), expected);
1039            }
1040        }
1041
1042        let expected = Expr::Column(Column::new_unqualified("s0"));
1043        assert_column_eq(&planner, "s0", &expected);
1044        assert_column_eq(&planner, "`s0`", &expected);
1045
1046        let expected = Expr::ScalarFunction(ScalarFunction {
1047            func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1048            args: vec![
1049                Expr::Column(Column::new_unqualified("st")),
1050                Expr::Literal(ScalarValue::Utf8(Some("s1".to_string()))),
1051            ],
1052        });
1053        assert_column_eq(&planner, "st.s1", &expected);
1054        assert_column_eq(&planner, "`st`.`s1`", &expected);
1055        assert_column_eq(&planner, "st.`s1`", &expected);
1056
1057        let expected = Expr::ScalarFunction(ScalarFunction {
1058            func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1059            args: vec![
1060                Expr::ScalarFunction(ScalarFunction {
1061                    func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1062                    args: vec![
1063                        Expr::Column(Column::new_unqualified("st")),
1064                        Expr::Literal(ScalarValue::Utf8(Some("st".to_string()))),
1065                    ],
1066                }),
1067                Expr::Literal(ScalarValue::Utf8(Some("s2".to_string()))),
1068            ],
1069        });
1070
1071        assert_column_eq(&planner, "st.st.s2", &expected);
1072        assert_column_eq(&planner, "`st`.`st`.`s2`", &expected);
1073        assert_column_eq(&planner, "st.st.`s2`", &expected);
1074        assert_column_eq(&planner, "st['st'][\"s2\"]", &expected);
1075    }
1076
1077    #[test]
1078    fn test_nested_list_refs() {
1079        let schema = Arc::new(Schema::new(vec![Field::new(
1080            "l",
1081            DataType::List(Arc::new(Field::new(
1082                "item",
1083                DataType::Struct(Fields::from(vec![Field::new("f1", DataType::Utf8, true)])),
1084                true,
1085            ))),
1086            true,
1087        )]));
1088
1089        let planner = Planner::new(schema);
1090
1091        let expected = array_element(col("l"), lit(0_i64));
1092        let expr = planner.parse_expr("l[0]").unwrap();
1093        assert_eq!(expr, expected);
1094
1095        let expected = get_field(array_element(col("l"), lit(0_i64)), "f1");
1096        let expr = planner.parse_expr("l[0]['f1']").unwrap();
1097        assert_eq!(expr, expected);
1098
1099        // FIXME: This should work, but sqlparser doesn't recognize anything
1100        // after the period for some reason.
1101        // let expr = planner.parse_expr("l[0].f1").unwrap();
1102        // assert_eq!(expr, expected);
1103    }
1104
1105    #[test]
1106    fn test_negative_expressions() {
1107        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
1108
1109        let planner = Planner::new(schema.clone());
1110
1111        let expected = col("x")
1112            .gt(lit(-3_i64))
1113            .and(col("x").lt(-(lit(-5_i64) + lit(3_i64))));
1114
1115        let expr = planner.parse_filter("x > -3 AND x < -(-5 + 3)").unwrap();
1116
1117        assert_eq!(expr, expected);
1118
1119        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1120
1121        let batch = RecordBatch::try_new(
1122            schema,
1123            vec![Arc::new(Int64Array::from_iter_values(-5..5)) as ArrayRef],
1124        )
1125        .unwrap();
1126        let predicates = physical_expr.evaluate(&batch).unwrap();
1127        assert_eq!(
1128            predicates.into_array(0).unwrap().as_ref(),
1129            &BooleanArray::from(vec![
1130                false, false, false, true, true, true, true, false, false, false
1131            ])
1132        );
1133    }
1134
1135    #[test]
1136    fn test_negative_array_expressions() {
1137        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
1138
1139        let planner = Planner::new(schema);
1140
1141        let expected = Expr::Literal(ScalarValue::List(Arc::new(
1142            ListArray::from_iter_primitive::<Float64Type, _, _>(vec![Some(
1143                [-1_f64, -2.0, -3.0, -4.0, -5.0].map(Some),
1144            )]),
1145        )));
1146
1147        let expr = planner
1148            .parse_expr("[-1.0, -2.0, -3.0, -4.0, -5.0]")
1149            .unwrap();
1150
1151        assert_eq!(expr, expected);
1152    }
1153
1154    #[test]
1155    fn test_sql_like() {
1156        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1157
1158        let planner = Planner::new(schema.clone());
1159
1160        let expected = col("s").like(lit("str-4"));
1161        // single quote
1162        let expr = planner.parse_filter("s LIKE 'str-4'").unwrap();
1163        assert_eq!(expr, expected);
1164        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1165
1166        let batch = RecordBatch::try_new(
1167            schema,
1168            vec![Arc::new(StringArray::from_iter_values(
1169                (0..10).map(|v| format!("str-{}", v)),
1170            ))],
1171        )
1172        .unwrap();
1173        let predicates = physical_expr.evaluate(&batch).unwrap();
1174        assert_eq!(
1175            predicates.into_array(0).unwrap().as_ref(),
1176            &BooleanArray::from(vec![
1177                false, false, false, false, true, false, false, false, false, false
1178            ])
1179        );
1180    }
1181
1182    #[test]
1183    fn test_not_like() {
1184        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1185
1186        let planner = Planner::new(schema.clone());
1187
1188        let expected = col("s").not_like(lit("str-4"));
1189        // single quote
1190        let expr = planner.parse_filter("s NOT LIKE 'str-4'").unwrap();
1191        assert_eq!(expr, expected);
1192        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1193
1194        let batch = RecordBatch::try_new(
1195            schema,
1196            vec![Arc::new(StringArray::from_iter_values(
1197                (0..10).map(|v| format!("str-{}", v)),
1198            ))],
1199        )
1200        .unwrap();
1201        let predicates = physical_expr.evaluate(&batch).unwrap();
1202        assert_eq!(
1203            predicates.into_array(0).unwrap().as_ref(),
1204            &BooleanArray::from(vec![
1205                true, true, true, true, false, true, true, true, true, true
1206            ])
1207        );
1208    }
1209
1210    #[test]
1211    fn test_sql_is_in() {
1212        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1213
1214        let planner = Planner::new(schema.clone());
1215
1216        let expected = col("s").in_list(vec![lit("str-4"), lit("str-5")], false);
1217        // single quote
1218        let expr = planner.parse_filter("s IN ('str-4', 'str-5')").unwrap();
1219        assert_eq!(expr, expected);
1220        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1221
1222        let batch = RecordBatch::try_new(
1223            schema,
1224            vec![Arc::new(StringArray::from_iter_values(
1225                (0..10).map(|v| format!("str-{}", v)),
1226            ))],
1227        )
1228        .unwrap();
1229        let predicates = physical_expr.evaluate(&batch).unwrap();
1230        assert_eq!(
1231            predicates.into_array(0).unwrap().as_ref(),
1232            &BooleanArray::from(vec![
1233                false, false, false, false, true, true, false, false, false, false
1234            ])
1235        );
1236    }
1237
1238    #[test]
1239    fn test_sql_is_null() {
1240        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1241
1242        let planner = Planner::new(schema.clone());
1243
1244        let expected = col("s").is_null();
1245        let expr = planner.parse_filter("s IS NULL").unwrap();
1246        assert_eq!(expr, expected);
1247        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1248
1249        let batch = RecordBatch::try_new(
1250            schema,
1251            vec![Arc::new(StringArray::from_iter((0..10).map(|v| {
1252                if v % 3 == 0 {
1253                    Some(format!("str-{}", v))
1254                } else {
1255                    None
1256                }
1257            })))],
1258        )
1259        .unwrap();
1260        let predicates = physical_expr.evaluate(&batch).unwrap();
1261        assert_eq!(
1262            predicates.into_array(0).unwrap().as_ref(),
1263            &BooleanArray::from(vec![
1264                false, true, true, false, true, true, false, true, true, false
1265            ])
1266        );
1267
1268        let expr = planner.parse_filter("s IS NOT NULL").unwrap();
1269        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1270        let predicates = physical_expr.evaluate(&batch).unwrap();
1271        assert_eq!(
1272            predicates.into_array(0).unwrap().as_ref(),
1273            &BooleanArray::from(vec![
1274                true, false, false, true, false, false, true, false, false, true,
1275            ])
1276        );
1277    }
1278
1279    #[test]
1280    fn test_sql_invert() {
1281        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Boolean, true)]));
1282
1283        let planner = Planner::new(schema.clone());
1284
1285        let expr = planner.parse_filter("NOT s").unwrap();
1286        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1287
1288        let batch = RecordBatch::try_new(
1289            schema,
1290            vec![Arc::new(BooleanArray::from_iter(
1291                (0..10).map(|v| Some(v % 3 == 0)),
1292            ))],
1293        )
1294        .unwrap();
1295        let predicates = physical_expr.evaluate(&batch).unwrap();
1296        assert_eq!(
1297            predicates.into_array(0).unwrap().as_ref(),
1298            &BooleanArray::from(vec![
1299                false, true, true, false, true, true, false, true, true, false
1300            ])
1301        );
1302    }
1303
1304    #[test]
1305    fn test_sql_cast() {
1306        let cases = &[
1307            (
1308                "x = cast('2021-01-01 00:00:00' as timestamp)",
1309                ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1310            ),
1311            (
1312                "x = cast('2021-01-01 00:00:00' as timestamp(0))",
1313                ArrowDataType::Timestamp(TimeUnit::Second, None),
1314            ),
1315            (
1316                "x = cast('2021-01-01 00:00:00.123' as timestamp(9))",
1317                ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1318            ),
1319            (
1320                "x = cast('2021-01-01 00:00:00.123' as datetime(9))",
1321                ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1322            ),
1323            ("x = cast('2021-01-01' as date)", ArrowDataType::Date32),
1324            (
1325                "x = cast('1.238' as decimal(9,3))",
1326                ArrowDataType::Decimal128(9, 3),
1327            ),
1328            ("x = cast(1 as float)", ArrowDataType::Float32),
1329            ("x = cast(1 as double)", ArrowDataType::Float64),
1330            ("x = cast(1 as tinyint)", ArrowDataType::Int8),
1331            ("x = cast(1 as smallint)", ArrowDataType::Int16),
1332            ("x = cast(1 as int)", ArrowDataType::Int32),
1333            ("x = cast(1 as integer)", ArrowDataType::Int32),
1334            ("x = cast(1 as bigint)", ArrowDataType::Int64),
1335            ("x = cast(1 as tinyint unsigned)", ArrowDataType::UInt8),
1336            ("x = cast(1 as smallint unsigned)", ArrowDataType::UInt16),
1337            ("x = cast(1 as int unsigned)", ArrowDataType::UInt32),
1338            ("x = cast(1 as integer unsigned)", ArrowDataType::UInt32),
1339            ("x = cast(1 as bigint unsigned)", ArrowDataType::UInt64),
1340            ("x = cast(1 as boolean)", ArrowDataType::Boolean),
1341            ("x = cast(1 as string)", ArrowDataType::Utf8),
1342        ];
1343
1344        for (sql, expected_data_type) in cases {
1345            let schema = Arc::new(Schema::new(vec![Field::new(
1346                "x",
1347                expected_data_type.clone(),
1348                true,
1349            )]));
1350            let planner = Planner::new(schema.clone());
1351            let expr = planner.parse_filter(sql).unwrap();
1352
1353            // Get the thing after 'cast(` but before ' as'.
1354            let expected_value_str = sql
1355                .split("cast(")
1356                .nth(1)
1357                .unwrap()
1358                .split(" as")
1359                .next()
1360                .unwrap();
1361            // Remove any quote marks
1362            let expected_value_str = expected_value_str.trim_matches('\'');
1363
1364            match expr {
1365                Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
1366                    Expr::Cast(Cast { expr, data_type }) => {
1367                        match expr.as_ref() {
1368                            Expr::Literal(ScalarValue::Utf8(Some(value_str))) => {
1369                                assert_eq!(value_str, expected_value_str);
1370                            }
1371                            Expr::Literal(ScalarValue::Int64(Some(value))) => {
1372                                assert_eq!(*value, 1);
1373                            }
1374                            _ => panic!("Expected cast to be applied to literal"),
1375                        }
1376                        assert_eq!(data_type, expected_data_type);
1377                    }
1378                    _ => panic!("Expected right to be a cast"),
1379                },
1380                _ => panic!("Expected binary expression"),
1381            }
1382        }
1383    }
1384
1385    #[test]
1386    fn test_sql_literals() {
1387        let cases = &[
1388            (
1389                "x = timestamp '2021-01-01 00:00:00'",
1390                ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1391            ),
1392            (
1393                "x = timestamp(0) '2021-01-01 00:00:00'",
1394                ArrowDataType::Timestamp(TimeUnit::Second, None),
1395            ),
1396            (
1397                "x = timestamp(9) '2021-01-01 00:00:00.123'",
1398                ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1399            ),
1400            ("x = date '2021-01-01'", ArrowDataType::Date32),
1401            ("x = decimal(9,3) '1.238'", ArrowDataType::Decimal128(9, 3)),
1402        ];
1403
1404        for (sql, expected_data_type) in cases {
1405            let schema = Arc::new(Schema::new(vec![Field::new(
1406                "x",
1407                expected_data_type.clone(),
1408                true,
1409            )]));
1410            let planner = Planner::new(schema.clone());
1411            let expr = planner.parse_filter(sql).unwrap();
1412
1413            let expected_value_str = sql.split('\'').nth(1).unwrap();
1414
1415            match expr {
1416                Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
1417                    Expr::Cast(Cast { expr, data_type }) => {
1418                        match expr.as_ref() {
1419                            Expr::Literal(ScalarValue::Utf8(Some(value_str))) => {
1420                                assert_eq!(value_str, expected_value_str);
1421                            }
1422                            _ => panic!("Expected cast to be applied to literal"),
1423                        }
1424                        assert_eq!(data_type, expected_data_type);
1425                    }
1426                    _ => panic!("Expected right to be a cast"),
1427                },
1428                _ => panic!("Expected binary expression"),
1429            }
1430        }
1431    }
1432
1433    #[test]
1434    fn test_sql_array_literals() {
1435        let cases = [
1436            (
1437                "x = [1, 2, 3]",
1438                ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int64, true))),
1439            ),
1440            (
1441                "x = [1, 2, 3]",
1442                ArrowDataType::FixedSizeList(
1443                    Arc::new(Field::new("item", ArrowDataType::Int64, true)),
1444                    3,
1445                ),
1446            ),
1447        ];
1448
1449        for (sql, expected_data_type) in cases {
1450            let schema = Arc::new(Schema::new(vec![Field::new(
1451                "x",
1452                expected_data_type.clone(),
1453                true,
1454            )]));
1455            let planner = Planner::new(schema.clone());
1456            let expr = planner.parse_filter(sql).unwrap();
1457            let expr = planner.optimize_expr(expr).unwrap();
1458
1459            match expr {
1460                Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
1461                    Expr::Literal(value) => {
1462                        assert_eq!(&value.data_type(), &expected_data_type);
1463                    }
1464                    _ => panic!("Expected right to be a literal"),
1465                },
1466                _ => panic!("Expected binary expression"),
1467            }
1468        }
1469    }
1470
1471    #[test]
1472    fn test_sql_between() {
1473        use arrow_array::{Float64Array, Int32Array, TimestampMicrosecondArray};
1474        use arrow_schema::{DataType, Field, Schema, TimeUnit};
1475        use std::sync::Arc;
1476
1477        let schema = Arc::new(Schema::new(vec![
1478            Field::new("x", DataType::Int32, false),
1479            Field::new("y", DataType::Float64, false),
1480            Field::new(
1481                "ts",
1482                DataType::Timestamp(TimeUnit::Microsecond, None),
1483                false,
1484            ),
1485        ]));
1486
1487        let planner = Planner::new(schema.clone());
1488
1489        // Test integer BETWEEN
1490        let expr = planner
1491            .parse_filter("x BETWEEN CAST(3 AS INT) AND CAST(7 AS INT)")
1492            .unwrap();
1493        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1494
1495        // Create timestamp array with values representing:
1496        // 2024-01-01 00:00:00 to 2024-01-01 00:00:09 (in microseconds)
1497        let base_ts = 1704067200000000_i64; // 2024-01-01 00:00:00
1498        let ts_array = TimestampMicrosecondArray::from_iter_values(
1499            (0..10).map(|i| base_ts + i * 1_000_000), // Each value is 1 second apart
1500        );
1501
1502        let batch = RecordBatch::try_new(
1503            schema,
1504            vec![
1505                Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef,
1506                Arc::new(Float64Array::from_iter_values((0..10).map(|v| v as f64))),
1507                Arc::new(ts_array),
1508            ],
1509        )
1510        .unwrap();
1511
1512        let predicates = physical_expr.evaluate(&batch).unwrap();
1513        assert_eq!(
1514            predicates.into_array(0).unwrap().as_ref(),
1515            &BooleanArray::from(vec![
1516                false, false, false, true, true, true, true, true, false, false
1517            ])
1518        );
1519
1520        // Test NOT BETWEEN
1521        let expr = planner
1522            .parse_filter("x NOT BETWEEN CAST(3 AS INT) AND CAST(7 AS INT)")
1523            .unwrap();
1524        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1525
1526        let predicates = physical_expr.evaluate(&batch).unwrap();
1527        assert_eq!(
1528            predicates.into_array(0).unwrap().as_ref(),
1529            &BooleanArray::from(vec![
1530                true, true, true, false, false, false, false, false, true, true
1531            ])
1532        );
1533
1534        // Test floating point BETWEEN
1535        let expr = planner.parse_filter("y BETWEEN 2.5 AND 6.5").unwrap();
1536        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1537
1538        let predicates = physical_expr.evaluate(&batch).unwrap();
1539        assert_eq!(
1540            predicates.into_array(0).unwrap().as_ref(),
1541            &BooleanArray::from(vec![
1542                false, false, false, true, true, true, true, false, false, false
1543            ])
1544        );
1545
1546        // Test timestamp BETWEEN
1547        let expr = planner
1548            .parse_filter(
1549                "ts BETWEEN timestamp '2024-01-01 00:00:03' AND timestamp '2024-01-01 00:00:07'",
1550            )
1551            .unwrap();
1552        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1553
1554        let predicates = physical_expr.evaluate(&batch).unwrap();
1555        assert_eq!(
1556            predicates.into_array(0).unwrap().as_ref(),
1557            &BooleanArray::from(vec![
1558                false, false, false, true, true, true, true, true, false, false
1559            ])
1560        );
1561    }
1562
1563    #[test]
1564    fn test_sql_comparison() {
1565        // Create a batch with all data types
1566        let batch: Vec<(&str, ArrayRef)> = vec![
1567            (
1568                "timestamp_s",
1569                Arc::new(TimestampSecondArray::from_iter_values(0..10)),
1570            ),
1571            (
1572                "timestamp_ms",
1573                Arc::new(TimestampMillisecondArray::from_iter_values(0..10)),
1574            ),
1575            (
1576                "timestamp_us",
1577                Arc::new(TimestampMicrosecondArray::from_iter_values(0..10)),
1578            ),
1579            (
1580                "timestamp_ns",
1581                Arc::new(TimestampNanosecondArray::from_iter_values(4995..5005)),
1582            ),
1583        ];
1584        let batch = RecordBatch::try_from_iter(batch).unwrap();
1585
1586        let planner = Planner::new(batch.schema());
1587
1588        // Each expression is meant to select the final 5 rows
1589        let expressions = &[
1590            "timestamp_s >= TIMESTAMP '1970-01-01 00:00:05'",
1591            "timestamp_ms >= TIMESTAMP '1970-01-01 00:00:00.005'",
1592            "timestamp_us >= TIMESTAMP '1970-01-01 00:00:00.000005'",
1593            "timestamp_ns >= TIMESTAMP '1970-01-01 00:00:00.000005'",
1594        ];
1595
1596        let expected: ArrayRef = Arc::new(BooleanArray::from_iter(
1597            std::iter::repeat_n(Some(false), 5).chain(std::iter::repeat_n(Some(true), 5)),
1598        ));
1599        for expression in expressions {
1600            // convert to physical expression
1601            let logical_expr = planner.parse_filter(expression).unwrap();
1602            let logical_expr = planner.optimize_expr(logical_expr).unwrap();
1603            let physical_expr = planner.create_physical_expr(&logical_expr).unwrap();
1604
1605            // Evaluate and assert they have correct results
1606            let result = physical_expr.evaluate(&batch).unwrap();
1607            let result = result.into_array(batch.num_rows()).unwrap();
1608            assert_eq!(&expected, &result, "unexpected result for {}", expression);
1609        }
1610    }
1611
1612    #[test]
1613    fn test_columns_in_expr() {
1614        let expr = col("s0").gt(lit("value")).and(
1615            col("st")
1616                .field("st")
1617                .field("s2")
1618                .eq(lit("value"))
1619                .or(col("st")
1620                    .field("s1")
1621                    .in_list(vec![lit("value 1"), lit("value 2")], false)),
1622        );
1623
1624        let columns = Planner::column_names_in_expr(&expr);
1625        assert_eq!(columns, vec!["s0", "st.s1", "st.st.s2"]);
1626    }
1627
1628    #[test]
1629    fn test_parse_binary_expr() {
1630        let bin_str = "x'616263'";
1631
1632        let schema = Arc::new(Schema::new(vec![Field::new(
1633            "binary",
1634            DataType::Binary,
1635            true,
1636        )]));
1637        let planner = Planner::new(schema);
1638        let expr = planner.parse_expr(bin_str).unwrap();
1639        assert_eq!(
1640            expr,
1641            Expr::Literal(ScalarValue::Binary(Some(vec![b'a', b'b', b'c'])))
1642        );
1643    }
1644
1645    #[test]
1646    fn test_lance_context_provider_expr_planners() {
1647        let ctx_provider = LanceContextProvider::default();
1648        assert!(!ctx_provider.get_expr_planners().is_empty());
1649    }
1650}
lance_datafusion/planner.rs

lance_datafusion/
planner.rs
