Skip to main content

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::exec::{LanceExecutionOptions, get_session_context};
11use crate::expr::safe_coerce_scalar;
12use crate::logical_expr::{coerce_filter_type_to_boolean, get_as_string_scalar_opt, resolve_expr};
13use crate::signed_zero::{normalize_zero_comparisons, rewrite_signed_zero_comparisons};
14use crate::sql::{parse_sql_expr, parse_sql_filter};
15use arrow::compute::CastOptions;
16use arrow_array::ListArray;
17use arrow_buffer::OffsetBuffer;
18use arrow_cast::cast_with_options;
19use arrow_schema::{DataType as ArrowDataType, Field, SchemaRef, TimeUnit};
20use arrow_select::concat::concat;
21use datafusion::catalog::Session;
22use datafusion::common::DFSchema;
23use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor};
24use datafusion::config::ConfigOptions;
25use datafusion::error::Result as DFResult;
26use datafusion::execution::context::{SessionContext, SessionState};
27use datafusion::logical_expr::expr::ScalarFunction;
28use datafusion::logical_expr::planner::{ExprPlanner, PlannerResult, RawFieldAccessExpr};
29use datafusion::logical_expr::{
30    AggregateUDF, ColumnarValue, GetFieldAccess, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
31    Signature, Volatility, WindowUDF,
32};
33use datafusion::optimizer::simplify_expressions::SimplifyContext;
34use datafusion::sql::planner::{
35    ContextProvider, NullOrdering, ParserOptions, PlannerContext, SqlToRel,
36};
37use datafusion::sql::sqlparser::ast::{
38    AccessExpr, Array as SQLArray, BinaryOperator, DataType as SQLDataType, ExactNumberInfo,
39    Expr as SQLExpr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Ident,
40    ObjectNamePart, Subscript, TimezoneInfo, TypedString, UnaryOperator, Value, ValueWithSpan,
41};
42use datafusion::{
43    common::Column,
44    logical_expr::{Between, BinaryExpr, Like, Operator},
45    physical_plan::PhysicalExpr,
46    prelude::Expr,
47    scalar::ScalarValue,
48};
49use datafusion_functions::core::getfield::GetFieldFunc;
50use lance_core::datatypes::Schema;
51use lance_core::error::LanceOptionExt;
52
53use chrono::Utc;
54use lance_core::{Error, Result};
55
56/// Encode a JSON string into a JSONB `LargeBinary` literal expression.
57fn encode_jsonb(json_str: &str) -> Result<Expr> {
58    let bytes = lance_arrow::json::encode_json(json_str)
59        .map_err(|e| Error::invalid_input(format!("Failed to encode JSONB: {e}")))?;
60    Ok(Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), None))
61}
62
63// The escape in `LIKE/ILIKE ... ESCAPE '<char>'` must be exactly one character.
64// Reject empty or multi-character escape strings rather than silently treating
65// them as "no escape" or truncating to the first character.
66fn parse_like_escape_char(escape_char: &Option<ValueWithSpan>) -> Result<Option<char>> {
67    let Some(value) = escape_char else {
68        return Ok(None);
69    };
70    let ValueWithSpan {
71        value: Value::SingleQuotedString(escape),
72        ..
73    } = value
74    else {
75        return Err(Error::invalid_input(format!(
76            "Invalid escape character in LIKE expression. Expected a single character wrapped with single quotes, got {value}"
77        )));
78    };
79    let mut chars = escape.chars();
80    match (chars.next(), chars.next()) {
81        (Some(c), None) => Ok(Some(c)),
82        _ => Err(Error::invalid_input(format!(
83            "Invalid escape character in LIKE expression. Expected a single character, got '{escape}'"
84        ))),
85    }
86}
87
88#[derive(Debug, Clone, Eq, PartialEq, Hash)]
89struct CastListF16Udf {
90    signature: Signature,
91}
92
93impl CastListF16Udf {
94    pub fn new() -> Self {
95        Self {
96            signature: Signature::any(1, Volatility::Immutable),
97        }
98    }
99}
100
101impl ScalarUDFImpl for CastListF16Udf {
102    fn name(&self) -> &str {
103        "_cast_list_f16"
104    }
105
106    fn signature(&self) -> &Signature {
107        &self.signature
108    }
109
110    fn return_type(&self, arg_types: &[ArrowDataType]) -> DFResult<ArrowDataType> {
111        let input = &arg_types[0];
112        match input {
113            ArrowDataType::FixedSizeList(field, size) => {
114                if field.data_type() != &ArrowDataType::Float32
115                    && field.data_type() != &ArrowDataType::Float16
116                {
117                    return Err(datafusion::error::DataFusionError::Execution(
118                        "cast_list_f16 only supports list of float32 or float16".to_string(),
119                    ));
120                }
121                Ok(ArrowDataType::FixedSizeList(
122                    Arc::new(Field::new(
123                        field.name(),
124                        ArrowDataType::Float16,
125                        field.is_nullable(),
126                    )),
127                    *size,
128                ))
129            }
130            ArrowDataType::List(field) => {
131                if field.data_type() != &ArrowDataType::Float32
132                    && field.data_type() != &ArrowDataType::Float16
133                {
134                    return Err(datafusion::error::DataFusionError::Execution(
135                        "cast_list_f16 only supports list of float32 or float16".to_string(),
136                    ));
137                }
138                Ok(ArrowDataType::List(Arc::new(Field::new(
139                    field.name(),
140                    ArrowDataType::Float16,
141                    field.is_nullable(),
142                ))))
143            }
144            _ => Err(datafusion::error::DataFusionError::Execution(
145                "cast_list_f16 only supports FixedSizeList/List arguments".to_string(),
146            )),
147        }
148    }
149
150    fn invoke_with_args(&self, func_args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
151        let ColumnarValue::Array(arr) = &func_args.args[0] else {
152            return Err(datafusion::error::DataFusionError::Execution(
153                "cast_list_f16 only supports array arguments".to_string(),
154            ));
155        };
156
157        let to_type = match arr.data_type() {
158            ArrowDataType::FixedSizeList(field, size) => ArrowDataType::FixedSizeList(
159                Arc::new(Field::new(
160                    field.name(),
161                    ArrowDataType::Float16,
162                    field.is_nullable(),
163                )),
164                *size,
165            ),
166            ArrowDataType::List(field) => ArrowDataType::List(Arc::new(Field::new(
167                field.name(),
168                ArrowDataType::Float16,
169                field.is_nullable(),
170            ))),
171            _ => {
172                return Err(datafusion::error::DataFusionError::Execution(
173                    "cast_list_f16 only supports array arguments".to_string(),
174                ));
175            }
176        };
177
178        let res = cast_with_options(arr.as_ref(), &to_type, &CastOptions::default())?;
179        Ok(ColumnarValue::Array(res))
180    }
181}
182
183// Adapter that instructs datafusion how lance expects expressions to be interpreted
184struct LanceContextProvider {
185    options: datafusion::config::ConfigOptions,
186    // Functions are looked up in the session's shared state rather than in a
187    // clone of it: planners are created several times per query, and cloning
188    // every function map and rule list for each one is a measurable share of
189    // CPU for short queries.
190    session: SessionContext,
191    expr_planners: Vec<Arc<dyn ExprPlanner>>,
192}
193
194impl LanceContextProvider {
195    fn new(session: SessionContext) -> Self {
196        let expr_planners = session.state_ref().read().expr_planners().to_vec();
197        Self {
198            options: ConfigOptions::default(),
199            session,
200            expr_planners,
201        }
202    }
203
204    fn with_state<T>(&self, read: impl FnOnce(&SessionState) -> T) -> T {
205        read(&self.session.state_ref().read())
206    }
207}
208
209impl Default for LanceContextProvider {
210    fn default() -> Self {
211        Self::new(get_session_context(&LanceExecutionOptions::default()))
212    }
213}
214
215impl ContextProvider for LanceContextProvider {
216    fn get_table_source(
217        &self,
218        name: datafusion::sql::TableReference,
219    ) -> DFResult<Arc<dyn datafusion::logical_expr::TableSource>> {
220        Err(datafusion::error::DataFusionError::NotImplemented(format!(
221            "Attempt to reference inner table {} not supported",
222            name
223        )))
224    }
225
226    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
227        self.with_state(|state| state.aggregate_functions().get(name).cloned())
228    }
229
230    fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
231        self.with_state(|state| state.window_functions().get(name).cloned())
232    }
233
234    fn get_higher_order_meta(
235        &self,
236        name: &str,
237    ) -> Option<Arc<datafusion::logical_expr::HigherOrderUDF>> {
238        self.with_state(|state| state.higher_order_functions().get(name).cloned())
239    }
240
241    fn get_function_meta(&self, f: &str) -> Option<Arc<ScalarUDF>> {
242        match f {
243            // TODO: cast should go thru CAST syntax instead of UDF
244            // Going thru UDF makes it hard for the optimizer to find no-ops
245            "_cast_list_f16" => Some(Arc::new(ScalarUDF::new_from_impl(CastListF16Udf::new()))),
246            _ => self.with_state(|state| state.scalar_functions().get(f).cloned()),
247        }
248    }
249
250    fn get_variable_type(&self, _: &[String]) -> Option<ArrowDataType> {
251        // Variables (things like @@LANGUAGE) not supported
252        None
253    }
254
255    fn options(&self) -> &datafusion::config::ConfigOptions {
256        &self.options
257    }
258
259    fn udf_names(&self) -> Vec<String> {
260        self.with_state(|state| state.scalar_functions().keys().cloned().collect())
261    }
262
263    fn udaf_names(&self) -> Vec<String> {
264        self.with_state(|state| state.aggregate_functions().keys().cloned().collect())
265    }
266
267    fn udwf_names(&self) -> Vec<String> {
268        self.with_state(|state| state.window_functions().keys().cloned().collect())
269    }
270
271    fn higher_order_function_names(&self) -> Vec<String> {
272        self.with_state(|state| state.higher_order_functions().keys().cloned().collect())
273    }
274
275    fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
276        &self.expr_planners
277    }
278}
279
280pub struct Planner {
281    schema: SchemaRef,
282    context_provider: LanceContextProvider,
283    enable_relations: bool,
284}
285
286impl Planner {
287    pub fn new(schema: SchemaRef) -> Self {
288        Self {
289            schema,
290            context_provider: LanceContextProvider::default(),
291            enable_relations: false,
292        }
293    }
294
295    /// If passed with `true`, then the first identifier in column reference
296    /// is parsed as the relation. For example, `table.field.inner` will be
297    /// read as the nested field `field.inner` (`inner` on struct field `field`)
298    /// on the `table` relation. If `false` (the default), then no relations
299    /// are used and all identifiers are assumed to be a nested column path.
300    pub fn with_enable_relations(mut self, enable_relations: bool) -> Self {
301        self.enable_relations = enable_relations;
302        self
303    }
304
305    /// Resolve a column name using case-insensitive matching against the schema.
306    /// Returns the actual field name if found, otherwise returns the original name.
307    fn resolve_column_name(&self, name: &str) -> String {
308        // Try exact match first
309        if self.schema.field_with_name(name).is_ok() {
310            return name.to_string();
311        }
312        // Fall back to case-insensitive match
313        for field in self.schema.fields() {
314            if field.name().eq_ignore_ascii_case(name) {
315                return field.name().clone();
316            }
317        }
318        // Not found in schema - return original (might be computed column, system column, etc.)
319        name.to_string()
320    }
321
322    fn column(&self, idents: &[Ident]) -> Expr {
323        fn handle_remaining_idents(expr: &mut Expr, idents: &[Ident]) {
324            for ident in idents {
325                *expr = Expr::ScalarFunction(ScalarFunction {
326                    args: vec![
327                        std::mem::take(expr),
328                        Expr::Literal(ScalarValue::Utf8(Some(ident.value.clone())), None),
329                    ],
330                    func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
331                });
332            }
333        }
334
335        if self.enable_relations && idents.len() > 1 {
336            // Create qualified column reference (relation.column)
337            let relation = &idents[0].value;
338            let column_name = self.resolve_column_name(&idents[1].value);
339            let column = Expr::Column(Column::new(Some(relation.clone()), column_name));
340            let mut result = column;
341            handle_remaining_idents(&mut result, &idents[2..]);
342            result
343        } else {
344            // Default behavior - treat as struct field access
345            // Use resolved column name to handle case-insensitive matching
346            let resolved_name = self.resolve_column_name(&idents[0].value);
347            let mut column = Expr::Column(Column::from_name(resolved_name));
348            handle_remaining_idents(&mut column, &idents[1..]);
349            column
350        }
351    }
352
353    fn binary_op(&self, op: &BinaryOperator) -> Result<Operator> {
354        Ok(match op {
355            BinaryOperator::Plus => Operator::Plus,
356            BinaryOperator::Minus => Operator::Minus,
357            BinaryOperator::Multiply => Operator::Multiply,
358            BinaryOperator::Divide => Operator::Divide,
359            BinaryOperator::Modulo => Operator::Modulo,
360            BinaryOperator::StringConcat => Operator::StringConcat,
361            BinaryOperator::Gt => Operator::Gt,
362            BinaryOperator::Lt => Operator::Lt,
363            BinaryOperator::GtEq => Operator::GtEq,
364            BinaryOperator::LtEq => Operator::LtEq,
365            BinaryOperator::Eq => Operator::Eq,
366            BinaryOperator::NotEq => Operator::NotEq,
367            BinaryOperator::And => Operator::And,
368            BinaryOperator::Or => Operator::Or,
369            BinaryOperator::PGBitwiseShiftLeft => Operator::BitwiseShiftLeft,
370            BinaryOperator::PGBitwiseShiftRight => Operator::BitwiseShiftRight,
371            _ => {
372                return Err(Error::invalid_input(format!(
373                    "Operator {op} is not supported"
374                )));
375            }
376        })
377    }
378
379    fn is_logical_binary_op(op: &BinaryOperator) -> bool {
380        matches!(op, BinaryOperator::And | BinaryOperator::Or)
381    }
382
383    fn is_same_logical_binary_op(left: &BinaryOperator, right: &BinaryOperator) -> bool {
384        matches!(
385            (left, right),
386            (BinaryOperator::And, BinaryOperator::And) | (BinaryOperator::Or, BinaryOperator::Or)
387        )
388    }
389
390    fn flatten_logical_binary_exprs<'a>(
391        left: &'a SQLExpr,
392        op: &BinaryOperator,
393        right: &'a SQLExpr,
394    ) -> Vec<&'a SQLExpr> {
395        let mut leaves = Vec::new();
396        let mut stack = vec![right, left];
397
398        while let Some(expr) = stack.pop() {
399            match expr {
400                SQLExpr::BinaryOp {
401                    left,
402                    op: child_op,
403                    right,
404                } if Self::is_same_logical_binary_op(op, child_op) => {
405                    stack.push(right.as_ref());
406                    stack.push(left.as_ref());
407                }
408                _ => leaves.push(expr),
409            }
410        }
411
412        leaves
413    }
414
415    fn balanced_binary_expr(mut exprs: VecDeque<Expr>, op: Operator) -> Result<Expr> {
416        if exprs.is_empty() {
417            return Err(Error::invalid_input("Binary expression has no operands"));
418        }
419
420        while exprs.len() > 1 {
421            let mut next = VecDeque::with_capacity(exprs.len().div_ceil(2));
422            while let Some(left) = exprs.pop_front() {
423                if let Some(right) = exprs.pop_front() {
424                    next.push_back(Expr::BinaryExpr(BinaryExpr::new(
425                        Box::new(left),
426                        op,
427                        Box::new(right),
428                    )));
429                } else {
430                    next.push_back(left);
431                }
432            }
433            exprs = next;
434        }
435
436        exprs
437            .pop_front()
438            .ok_or_else(|| Error::invalid_input("Binary expression has no operands"))
439    }
440
441    fn binary_expr(&self, left: &SQLExpr, op: &BinaryOperator, right: &SQLExpr) -> Result<Expr> {
442        let df_op = self.binary_op(op)?;
443        if Self::is_logical_binary_op(op) {
444            let leaves = Self::flatten_logical_binary_exprs(left, op, right);
445            let mut exprs = VecDeque::with_capacity(leaves.len());
446            for leaf in leaves {
447                exprs.push_back(self.parse_sql_expr(leaf)?);
448            }
449            return Self::balanced_binary_expr(exprs, df_op);
450        }
451
452        Ok(Expr::BinaryExpr(BinaryExpr::new(
453            Box::new(self.parse_sql_expr(left)?),
454            df_op,
455            Box::new(self.parse_sql_expr(right)?),
456        )))
457    }
458
459    fn unary_expr(&self, op: &UnaryOperator, expr: &SQLExpr) -> Result<Expr> {
460        Ok(match op {
461            UnaryOperator::Not | UnaryOperator::BitwiseNot => {
462                Expr::Not(Box::new(self.parse_sql_expr(expr)?))
463            }
464
465            UnaryOperator::Minus => {
466                use datafusion::logical_expr::lit;
467                match expr {
468                    SQLExpr::Value(ValueWithSpan { value: Value::Number(n, _), ..}) => match n.parse::<i64>() {
469                        Ok(n) => lit(-n),
470                        Err(_) => lit(-n
471                            .parse::<f64>()
472                            .map_err(|_e| {
473                                Error::invalid_input(format!("negative operator can be only applied to integer and float operands, got: {n}"))
474                            })?),
475                    },
476                    _ => {
477                        Expr::Negative(Box::new(self.parse_sql_expr(expr)?))
478                    }
479                }
480            }
481
482            _ => {
483                return Err(Error::invalid_input(format!(
484                    "Unary operator '{:?}' is not supported",
485                    op
486                )));
487            }
488        })
489    }
490
491    // See datafusion `sqlToRel::parse_sql_number()`
492    fn number(&self, value: &str, negative: bool) -> Result<Expr> {
493        use datafusion::logical_expr::lit;
494        let value: Cow<str> = if negative {
495            Cow::Owned(format!("-{}", value))
496        } else {
497            Cow::Borrowed(value)
498        };
499        if let Ok(n) = value.parse::<i64>() {
500            Ok(lit(n))
501        } else if let Ok(n) = value.parse::<u64>() {
502            Ok(lit(n))
503        } else {
504            value.parse::<f64>().map(lit).map_err(|_| {
505                Error::invalid_input(format!("'{value}' is not supported number value."))
506            })
507        }
508    }
509
510    fn value(&self, value: &Value) -> Result<Expr> {
511        Ok(match value {
512            Value::Number(v, _) => self.number(v.as_str(), false)?,
513            Value::SingleQuotedString(s) => Expr::Literal(ScalarValue::Utf8(Some(s.clone())), None),
514            Value::HexStringLiteral(hsl) => {
515                Expr::Literal(ScalarValue::Binary(Self::try_decode_hex_literal(hsl)), None)
516            }
517            Value::DoubleQuotedString(s) => Expr::Literal(ScalarValue::Utf8(Some(s.clone())), None),
518            Value::Boolean(v) => Expr::Literal(ScalarValue::Boolean(Some(*v)), None),
519            Value::Null => Expr::Literal(ScalarValue::Null, None),
520            _ => todo!(),
521        })
522    }
523
524    fn parse_function_args(&self, func_args: &FunctionArg) -> Result<Expr> {
525        match func_args {
526            FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => self.parse_sql_expr(expr),
527            _ => Err(Error::invalid_input(format!(
528                "Unsupported function args: {:?}",
529                func_args
530            ))),
531        }
532    }
533
534    // We now use datafusion to parse functions.  This allows us to use datafusion's
535    // entire collection of functions (previously we had just hard-coded support for two functions).
536    //
537    // Unfortunately, one of those two functions was is_valid and the reason we needed it was because
538    // this is a function that comes from duckdb.  Datafusion does not consider is_valid to be a function
539    // but rather an AST node (Expr::IsNotNull) and so we need to handle this case specially.
540    fn legacy_parse_function(&self, func: &Function) -> Result<Expr> {
541        match &func.args {
542            FunctionArguments::List(args) => {
543                if func.name.0.len() != 1 {
544                    return Err(Error::invalid_input(format!(
545                        "Function name must have 1 part, got: {:?}",
546                        func.name.0
547                    )));
548                }
549                Ok(Expr::IsNotNull(Box::new(
550                    self.parse_function_args(&args.args[0])?,
551                )))
552            }
553            _ => Err(Error::invalid_input(format!(
554                "Unsupported function args: {:?}",
555                func.args
556            ))),
557        }
558    }
559
560    fn parse_function(&self, function: SQLExpr) -> Result<Expr> {
561        if let SQLExpr::Function(function) = &function
562            && let Some(ObjectNamePart::Identifier(name)) = &function.name.0.first()
563            && &name.value == "is_valid"
564        {
565            return self.legacy_parse_function(function);
566        }
567        let sql_to_rel = SqlToRel::new_with_options(
568            &self.context_provider,
569            ParserOptions {
570                parse_float_as_decimal: false,
571                enable_ident_normalization: false,
572                support_varchar_with_length: false,
573                enable_options_value_normalization: false,
574                collect_spans: false,
575                map_string_types_to_utf8view: false,
576                default_null_ordering: NullOrdering::NullsMax,
577            },
578        );
579
580        let mut planner_context = PlannerContext::default();
581        let schema = DFSchema::try_from(self.schema.as_ref().clone())?;
582        sql_to_rel
583            .sql_to_expr(function, &schema, &mut planner_context)
584            .map_err(|e| Error::invalid_input(format!("Error parsing function: {e}")))
585    }
586
587    fn parse_type(&self, data_type: &SQLDataType) -> Result<ArrowDataType> {
588        const SUPPORTED_TYPES: [&str; 13] = [
589            "int [unsigned]",
590            "tinyint [unsigned]",
591            "smallint [unsigned]",
592            "bigint [unsigned]",
593            "float",
594            "double",
595            "string",
596            "binary",
597            "date",
598            "timestamp(precision)",
599            "datetime(precision)",
600            "decimal(precision,scale)",
601            "boolean",
602        ];
603        match data_type {
604            SQLDataType::String(_) => Ok(ArrowDataType::Utf8),
605            SQLDataType::Binary(_) => Ok(ArrowDataType::Binary),
606            SQLDataType::Float(_) => Ok(ArrowDataType::Float32),
607            SQLDataType::Double(_) => Ok(ArrowDataType::Float64),
608            SQLDataType::Boolean => Ok(ArrowDataType::Boolean),
609            SQLDataType::TinyInt(_) => Ok(ArrowDataType::Int8),
610            SQLDataType::SmallInt(_) => Ok(ArrowDataType::Int16),
611            SQLDataType::Int(_) | SQLDataType::Integer(_) => Ok(ArrowDataType::Int32),
612            SQLDataType::BigInt(_) => Ok(ArrowDataType::Int64),
613            SQLDataType::TinyIntUnsigned(_) => Ok(ArrowDataType::UInt8),
614            SQLDataType::SmallIntUnsigned(_) => Ok(ArrowDataType::UInt16),
615            SQLDataType::IntUnsigned(_) | SQLDataType::IntegerUnsigned(_) => {
616                Ok(ArrowDataType::UInt32)
617            }
618            SQLDataType::BigIntUnsigned(_) => Ok(ArrowDataType::UInt64),
619            SQLDataType::Date => Ok(ArrowDataType::Date32),
620            SQLDataType::Timestamp(resolution, tz) => {
621                match tz {
622                    TimezoneInfo::None => {}
623                    _ => {
624                        return Err(Error::invalid_input(
625                            "Timezone not supported in timestamp".to_string(),
626                        ));
627                    }
628                };
629                let time_unit = match resolution {
630                    // Default to microsecond to match PyArrow
631                    None => TimeUnit::Microsecond,
632                    Some(0) => TimeUnit::Second,
633                    Some(3) => TimeUnit::Millisecond,
634                    Some(6) => TimeUnit::Microsecond,
635                    Some(9) => TimeUnit::Nanosecond,
636                    _ => {
637                        return Err(Error::invalid_input(format!(
638                            "Unsupported datetime resolution: {:?}",
639                            resolution
640                        )));
641                    }
642                };
643                Ok(ArrowDataType::Timestamp(time_unit, None))
644            }
645            SQLDataType::Datetime(resolution) => {
646                let time_unit = match resolution {
647                    None => TimeUnit::Microsecond,
648                    Some(0) => TimeUnit::Second,
649                    Some(3) => TimeUnit::Millisecond,
650                    Some(6) => TimeUnit::Microsecond,
651                    Some(9) => TimeUnit::Nanosecond,
652                    _ => {
653                        return Err(Error::invalid_input(format!(
654                            "Unsupported datetime resolution: {:?}",
655                            resolution
656                        )));
657                    }
658                };
659                Ok(ArrowDataType::Timestamp(time_unit, None))
660            }
661            SQLDataType::Decimal(number_info) => match number_info {
662                ExactNumberInfo::PrecisionAndScale(precision, scale) => {
663                    Ok(ArrowDataType::Decimal128(*precision as u8, *scale as i8))
664                }
665                _ => Err(Error::invalid_input(format!(
666                    "Must provide precision and scale for decimal: {:?}",
667                    number_info
668                ))),
669            },
670            _ => Err(Error::invalid_input(format!(
671                "Unsupported data type: {:?}. Supported types: {:?}",
672                data_type, SUPPORTED_TYPES
673            ))),
674        }
675    }
676
677    fn plan_field_access(&self, mut field_access_expr: RawFieldAccessExpr) -> Result<Expr> {
678        let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
679        for planner in self.context_provider.get_expr_planners() {
680            match planner.plan_field_access(field_access_expr, &df_schema)? {
681                PlannerResult::Planned(expr) => return Ok(expr),
682                PlannerResult::Original(expr) => {
683                    field_access_expr = expr;
684                }
685            }
686        }
687        Err(Error::invalid_input("Field access could not be planned"))
688    }
689
690    fn parse_sql_expr(&self, expr: &SQLExpr) -> Result<Expr> {
691        match expr {
692            SQLExpr::Identifier(id) => {
693                // Users can pass string literals wrapped in `"`.
694                // (Normally SQL only allows single quotes.)
695                if id.quote_style == Some('"') {
696                    Ok(Expr::Literal(
697                        ScalarValue::Utf8(Some(id.value.clone())),
698                        None,
699                    ))
700                // Users can wrap identifiers with ` to reference non-standard
701                // names, such as uppercase or spaces.
702                } else if id.quote_style == Some('`') {
703                    Ok(Expr::Column(Column::from_name(id.value.clone())))
704                } else {
705                    Ok(self.column(vec![id.clone()].as_slice()))
706                }
707            }
708            SQLExpr::CompoundIdentifier(ids) => Ok(self.column(ids.as_slice())),
709            SQLExpr::BinaryOp { left, op, right } => self.binary_expr(left, op, right),
710            SQLExpr::UnaryOp { op, expr } => self.unary_expr(op, expr),
711            SQLExpr::Value(value) => self.value(&value.value),
712            SQLExpr::Array(SQLArray { elem, .. }) => {
713                let mut values = vec![];
714
715                let array_literal_error = |pos: usize, value: &_| {
716                    Err(Error::invalid_input(format!(
717                        "Expected a literal value in array, instead got {} at position {}",
718                        value, pos
719                    )))
720                };
721
722                for (pos, expr) in elem.iter().enumerate() {
723                    match expr {
724                        SQLExpr::Value(value) => {
725                            if let Expr::Literal(value, _) = self.value(&value.value)? {
726                                values.push(value);
727                            } else {
728                                return array_literal_error(pos, expr);
729                            }
730                        }
731                        SQLExpr::UnaryOp {
732                            op: UnaryOperator::Minus,
733                            expr,
734                        } => {
735                            if let SQLExpr::Value(ValueWithSpan {
736                                value: Value::Number(number, _),
737                                ..
738                            }) = expr.as_ref()
739                            {
740                                if let Expr::Literal(value, _) = self.number(number, true)? {
741                                    values.push(value);
742                                } else {
743                                    return array_literal_error(pos, expr);
744                                }
745                            } else {
746                                return array_literal_error(pos, expr);
747                            }
748                        }
749                        _ => {
750                            return array_literal_error(pos, expr);
751                        }
752                    }
753                }
754
755                let field = if !values.is_empty() {
756                    let data_type = values[0].data_type();
757
758                    for value in &mut values {
759                        if value.data_type() != data_type {
760                            *value = safe_coerce_scalar(value, &data_type).ok_or_else(|| Error::invalid_input(format!("Array expressions must have a consistent datatype. Expected: {}, got: {}", data_type, value.data_type())))?;
761                        }
762                    }
763                    Field::new("item", data_type, true)
764                } else {
765                    Field::new("item", ArrowDataType::Null, true)
766                };
767
768                let values = values
769                    .into_iter()
770                    .map(|v| v.to_array().map_err(Error::from))
771                    .collect::<Result<Vec<_>>>()?;
772                let array_refs = values.iter().map(|v| v.as_ref()).collect::<Vec<_>>();
773                let values = concat(&array_refs)?;
774                let values = ListArray::try_new(
775                    field.into(),
776                    OffsetBuffer::from_lengths([values.len()]),
777                    values,
778                    None,
779                )?;
780
781                Ok(Expr::Literal(ScalarValue::List(Arc::new(values)), None))
782            }
783            // JSONB literal: jsonb '{"key": "value"}'
784            SQLExpr::TypedString(TypedString {
785                data_type: SQLDataType::JSONB,
786                value,
787                ..
788            }) => match &value.value {
789                Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => encode_jsonb(s),
790                _ => Err(Error::invalid_input(
791                    "Expected a string value for JSONB literal",
792                )),
793            },
794            // For example, DATE '2020-01-01'
795            SQLExpr::TypedString(TypedString {
796                data_type, value, ..
797            }) => {
798                let value = value.clone().into_string().expect_ok()?;
799                Ok(Expr::Cast(datafusion::logical_expr::Cast::new(
800                    Box::new(Expr::Literal(ScalarValue::Utf8(Some(value)), None)),
801                    self.parse_type(data_type)?,
802                )))
803            }
804            SQLExpr::IsFalse(expr) => Ok(Expr::IsFalse(Box::new(self.parse_sql_expr(expr)?))),
805            SQLExpr::IsNotFalse(expr) => Ok(Expr::IsNotFalse(Box::new(self.parse_sql_expr(expr)?))),
806            SQLExpr::IsTrue(expr) => Ok(Expr::IsTrue(Box::new(self.parse_sql_expr(expr)?))),
807            SQLExpr::IsNotTrue(expr) => Ok(Expr::IsNotTrue(Box::new(self.parse_sql_expr(expr)?))),
808            SQLExpr::IsNull(expr) => Ok(Expr::IsNull(Box::new(self.parse_sql_expr(expr)?))),
809            SQLExpr::IsNotNull(expr) => Ok(Expr::IsNotNull(Box::new(self.parse_sql_expr(expr)?))),
810            SQLExpr::InList {
811                expr,
812                list,
813                negated,
814            } => {
815                let value_expr = self.parse_sql_expr(expr)?;
816                let list_exprs = list
817                    .iter()
818                    .map(|e| self.parse_sql_expr(e))
819                    .collect::<Result<Vec<_>>>()?;
820                Ok(value_expr.in_list(list_exprs, *negated))
821            }
822            SQLExpr::Nested(inner) => self.parse_sql_expr(inner.as_ref()),
823            SQLExpr::Function(_) => self.parse_function(expr.clone()),
824            SQLExpr::ILike {
825                negated,
826                expr,
827                pattern,
828                escape_char,
829                any: _,
830            } => Ok(Expr::Like(Like::new(
831                *negated,
832                Box::new(self.parse_sql_expr(expr)?),
833                Box::new(self.parse_sql_expr(pattern)?),
834                parse_like_escape_char(escape_char)?,
835                true,
836            ))),
837            SQLExpr::Like {
838                negated,
839                expr,
840                pattern,
841                escape_char,
842                any: _,
843            } => Ok(Expr::Like(Like::new(
844                *negated,
845                Box::new(self.parse_sql_expr(expr)?),
846                Box::new(self.parse_sql_expr(pattern)?),
847                parse_like_escape_char(escape_char)?,
848                false,
849            ))),
850            // JSONB cast: CAST('...' AS JSONB) or '...'::jsonb
851            SQLExpr::Cast {
852                data_type: SQLDataType::JSONB,
853                expr: inner,
854                ..
855            } => match inner.as_ref() {
856                SQLExpr::Value(ValueWithSpan {
857                    value: Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
858                    ..
859                }) => encode_jsonb(s),
860                _ => Err(Error::invalid_input(
861                    "CAST to JSONB only supports string literals",
862                )),
863            },
864            SQLExpr::Cast {
865                expr,
866                data_type,
867                kind,
868                ..
869            } => match kind {
870                datafusion::sql::sqlparser::ast::CastKind::TryCast
871                | datafusion::sql::sqlparser::ast::CastKind::SafeCast => {
872                    Ok(Expr::TryCast(datafusion::logical_expr::TryCast::new(
873                        Box::new(self.parse_sql_expr(expr)?),
874                        self.parse_type(data_type)?,
875                    )))
876                }
877                _ => Ok(Expr::Cast(datafusion::logical_expr::Cast::new(
878                    Box::new(self.parse_sql_expr(expr)?),
879                    self.parse_type(data_type)?,
880                ))),
881            },
882            SQLExpr::JsonAccess { .. } => Err(Error::invalid_input("JSON access is not supported")),
883            SQLExpr::CompoundFieldAccess { root, access_chain } => {
884                let mut expr = self.parse_sql_expr(root)?;
885
886                for access in access_chain {
887                    let field_access = match access {
888                        // x.y or x['y']
889                        AccessExpr::Dot(SQLExpr::Identifier(Ident { value: s, .. }))
890                        | AccessExpr::Subscript(Subscript::Index {
891                            index:
892                                SQLExpr::Value(ValueWithSpan {
893                                    value:
894                                        Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
895                                    ..
896                                }),
897                        }) => GetFieldAccess::NamedStructField {
898                            name: ScalarValue::from(s.as_str()),
899                        },
900                        AccessExpr::Subscript(Subscript::Index { index }) => {
901                            let key = Box::new(self.parse_sql_expr(index)?);
902                            GetFieldAccess::ListIndex { key }
903                        }
904                        AccessExpr::Subscript(Subscript::Slice { .. }) => {
905                            return Err(Error::invalid_input("Slice subscript is not supported"));
906                        }
907                        _ => {
908                            // Handle other cases like JSON access
909                            // Note: JSON access is not supported in lance
910                            return Err(Error::invalid_input(
911                                "Only dot notation or index access is supported for field access",
912                            ));
913                        }
914                    };
915
916                    let field_access_expr = RawFieldAccessExpr { expr, field_access };
917                    expr = self.plan_field_access(field_access_expr)?;
918                }
919
920                Ok(expr)
921            }
922            SQLExpr::Between {
923                expr,
924                negated,
925                low,
926                high,
927            } => {
928                // Parse the main expression and bounds
929                let expr = self.parse_sql_expr(expr)?;
930                let low = self.parse_sql_expr(low)?;
931                let high = self.parse_sql_expr(high)?;
932
933                let between = Expr::Between(Between::new(
934                    Box::new(expr),
935                    *negated,
936                    Box::new(low),
937                    Box::new(high),
938                ));
939                Ok(between)
940            }
941            _ => Err(Error::invalid_input(format!(
942                "Expression '{expr}' is not supported SQL in lance"
943            ))),
944        }
945    }
946
947    /// Create Logical [Expr] from a SQL filter clause.
948    ///
949    /// Note: the returned expression must be passed through `optimize_expr()`
950    /// before being passed to `create_physical_expr()`.
951    pub fn parse_filter(&self, filter: &str) -> Result<Expr> {
952        // Allow sqlparser to parse filter as part of ONE SQL statement.
953        let ast_expr = parse_sql_filter(filter)?;
954        let expr = self.parse_sql_expr(&ast_expr)?;
955        let schema = Schema::try_from(self.schema.as_ref())?;
956        let resolved = resolve_expr(&expr, &schema).map_err(|e| {
957            Error::invalid_input(format!("Error resolving filter expression {filter}: {e}"))
958        })?;
959
960        Ok(coerce_filter_type_to_boolean(resolved))
961    }
962
963    /// Create Logical [Expr] from a SQL expression.
964    ///
965    /// Note: the returned expression must be passed through `optimize_filter()`
966    /// before being passed to `create_physical_expr()`.
967    pub fn parse_expr(&self, expr: &str) -> Result<Expr> {
968        // First check if it's a simple column reference (no operators, functions, etc.)
969        // resolve_column_name tries exact match first, then falls back to case-insensitive
970        let resolved_name = self.resolve_column_name(expr);
971        if self.schema.field_with_name(&resolved_name).is_ok() {
972            return Ok(Expr::Column(Column::from_name(resolved_name)));
973        }
974
975        // Parse as SQL expression
976        let ast_expr = parse_sql_expr(expr)?;
977        let expr = self.parse_sql_expr(&ast_expr)?;
978        let schema = Schema::try_from(self.schema.as_ref())?;
979        let resolved = resolve_expr(&expr, &schema)?;
980        Ok(resolved)
981    }
982
983    /// Try to decode bytes from hex literal string.
984    ///
985    /// Copied from datafusion because this is not public.
986    ///
987    /// TODO: use SqlToRel from Datafusion directly?
988    fn try_decode_hex_literal(s: &str) -> Option<Vec<u8>> {
989        let hex_bytes = s.as_bytes();
990        let mut decoded_bytes = Vec::with_capacity(hex_bytes.len().div_ceil(2));
991
992        let start_idx = hex_bytes.len() % 2;
993        if start_idx > 0 {
994            // The first byte is formed of only one char.
995            decoded_bytes.push(Self::try_decode_hex_char(hex_bytes[0])?);
996        }
997
998        for i in (start_idx..hex_bytes.len()).step_by(2) {
999            let high = Self::try_decode_hex_char(hex_bytes[i])?;
1000            let low = Self::try_decode_hex_char(hex_bytes[i + 1])?;
1001            decoded_bytes.push((high << 4) | low);
1002        }
1003
1004        Some(decoded_bytes)
1005    }
1006
1007    /// Try to decode a byte from a hex char.
1008    ///
1009    /// None will be returned if the input char is hex-invalid.
1010    const fn try_decode_hex_char(c: u8) -> Option<u8> {
1011        match c {
1012            b'A'..=b'F' => Some(c - b'A' + 10),
1013            b'a'..=b'f' => Some(c - b'a' + 10),
1014            b'0'..=b'9' => Some(c - b'0'),
1015            _ => None,
1016        }
1017    }
1018
1019    /// Optimize the filter expression and coerce data types.
1020    pub fn optimize_expr(&self, expr: Expr) -> Result<Expr> {
1021        let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?);
1022
1023        // DataFusion rewrites arrow_cast to Expr::Cast, whose Arrow kernel does not support
1024        // integer-to-Time32 casts. Convert literal values with Lance's scalar coercion first.
1025        let expr = expr
1026            .transform_up(|expr| {
1027                let coerced = match &expr {
1028                    Expr::ScalarFunction(ScalarFunction { func, args })
1029                        if func.name() == "arrow_cast" =>
1030                    {
1031                        match args.as_slice() {
1032                            [
1033                                Expr::Literal(value, metadata),
1034                                Expr::Literal(ScalarValue::Utf8(Some(data_type)), _),
1035                            ] => data_type
1036                                .parse::<ArrowDataType>()
1037                                .ok()
1038                                .filter(|data_type| matches!(data_type, ArrowDataType::Time32(_)))
1039                                .and_then(|data_type| {
1040                                    if matches!(value, ScalarValue::Null) {
1041                                        ScalarValue::try_new_null(&data_type).ok()
1042                                    } else {
1043                                        safe_coerce_scalar(value, &data_type)
1044                                    }
1045                                })
1046                                .map(|value| Expr::Literal(value, metadata.clone())),
1047                            _ => None,
1048                        }
1049                    }
1050                    _ => None,
1051                };
1052
1053                Ok(match coerced {
1054                    Some(coerced) => Transformed::yes(coerced),
1055                    None => Transformed::no(expr),
1056                })
1057            })?
1058            .data;
1059
1060        // DataFusion needs the coerce and simplify passes to be applied before
1061        // expressions can be handled by the physical planner.
1062        let simplify_context = SimplifyContext::builder()
1063            .with_schema(df_schema.clone())
1064            .with_query_execution_start_time(Some(Utc::now()))
1065            .build();
1066        let simplifier =
1067            datafusion::optimizer::simplify_expressions::ExprSimplifier::new(simplify_context);
1068
1069        // Coerce before simplify to match DataFusion's analyzer-before-optimizer pipeline.
1070        let expr = simplifier.coerce(expr, &df_schema)?;
1071
1072        // Fold each comparison's own operands and rewrite it before anything above
1073        // it folds. `simplify` folds an operand and everything above it in one
1074        // pass, so a fully constant predicate whose zero appears only as a result
1075        // of folding never presents a zero literal to the rewrite:
1076        // `-1.0 * 0.0 < (1.0 - 1.0)` answered `true` where IEEE says false, and a
1077        // wrapper such as `IS TRUE` or a `CAST` did the same to the comparison's
1078        // own result.
1079        let expr = normalize_zero_comparisons(expr, &|operand| simplifier.simplify(operand))?;
1080
1081        // Again after simplify, which is what expands `BETWEEN` into two
1082        // comparisons and folds the casts `coerce` inserts, so those forms only
1083        // become visible on this pass.
1084        //
1085        // Running the rewrite more than once is safe because its output is a fixed
1086        // point of `optimize_expr`; `optimizing_twice_changes_nothing` pins that.
1087        let expr = simplifier.simplify(expr)?;
1088        let expr = rewrite_signed_zero_comparisons(expr)?;
1089
1090        Ok(expr)
1091    }
1092
1093    /// Create the [`PhysicalExpr`] from a logical [`Expr`]
1094    pub fn create_physical_expr(&self, expr: &Expr) -> Result<Arc<dyn PhysicalExpr>> {
1095        let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?);
1096        Ok(datafusion::physical_expr::create_physical_expr(
1097            expr,
1098            df_schema.as_ref(),
1099            &Default::default(),
1100        )?)
1101    }
1102
1103    /// Create a [`PhysicalExpr`] using the caller's DataFusion session.
1104    pub fn create_physical_expr_with_session(
1105        &self,
1106        expr: &Expr,
1107        session: &dyn Session,
1108    ) -> Result<Arc<dyn PhysicalExpr>> {
1109        let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
1110        Ok(session.create_physical_expr(expr.clone(), &df_schema)?)
1111    }
1112
1113    /// Collect the columns in the expression.
1114    ///
1115    /// The columns are returned in sorted order.
1116    ///
1117    /// If the expr refers to nested columns these will be returned
1118    /// as dotted paths (x.y.z)
1119    pub fn column_names_in_expr(expr: &Expr) -> Vec<String> {
1120        let mut visitor = ColumnCapturingVisitor {
1121            current_path: VecDeque::new(),
1122            columns: BTreeSet::new(),
1123        };
1124        expr.visit(&mut visitor).unwrap();
1125        visitor.columns.into_iter().collect()
1126    }
1127}
1128
1129struct ColumnCapturingVisitor {
1130    // Current column path. If this is empty, we are not in a column expression.
1131    current_path: VecDeque<String>,
1132    columns: BTreeSet<String>,
1133}
1134
1135impl TreeNodeVisitor<'_> for ColumnCapturingVisitor {
1136    type Node = Expr;
1137
1138    fn f_down(&mut self, node: &Self::Node) -> DFResult<TreeNodeRecursion> {
1139        match node {
1140            Expr::Column(Column { name, .. }) => {
1141                // Build the field path from the column name and any nested fields
1142                // The nested field names from get_field already come as literal strings,
1143                // so we just need to concatenate them properly
1144                let mut path = name.clone();
1145                for part in self.current_path.drain(..) {
1146                    path.push('.');
1147                    // Check if the part needs quoting (contains dots)
1148                    if part.contains('.') || part.contains('`') {
1149                        // Quote the field name with backticks and escape any existing backticks
1150                        let escaped = part.replace('`', "``");
1151                        path.push('`');
1152                        path.push_str(&escaped);
1153                        path.push('`');
1154                    } else {
1155                        path.push_str(&part);
1156                    }
1157                }
1158                self.columns.insert(path);
1159                self.current_path.clear();
1160            }
1161            Expr::ScalarFunction(udf) if udf.name() == GetFieldFunc::default().name() => {
1162                if let Some(name) = get_as_string_scalar_opt(&udf.args[1]) {
1163                    self.current_path.push_front(name.to_string())
1164                } else {
1165                    self.current_path.clear();
1166                }
1167            }
1168            _ => {
1169                self.current_path.clear();
1170            }
1171        }
1172
1173        Ok(TreeNodeRecursion::Continue)
1174    }
1175}
1176
1177#[cfg(test)]
1178mod tests {
1179
1180    use super::*;
1181
1182    use arrow::datatypes::Float64Type;
1183    use arrow_array::{
1184        ArrayRef, BooleanArray, Float32Array, Int32Array, Int64Array, RecordBatch, StringArray,
1185        StructArray, Time32SecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
1186        TimestampNanosecondArray, TimestampSecondArray, UInt64Array,
1187    };
1188    use arrow_schema::{DataType, Fields, Schema};
1189    use datafusion::{
1190        logical_expr::{Cast, col, lit},
1191        prelude::{array_element, get_field},
1192    };
1193    use datafusion_functions::core::expr_ext::FieldAccessor;
1194    use rstest::rstest;
1195
1196    #[test]
1197    fn test_parse_filter_simple() {
1198        let schema = Arc::new(Schema::new(vec![
1199            Field::new("i", DataType::Int32, false),
1200            Field::new("s", DataType::Utf8, true),
1201            Field::new(
1202                "st",
1203                DataType::Struct(Fields::from(vec![
1204                    Field::new("x", DataType::Float32, false),
1205                    Field::new("y", DataType::Float32, false),
1206                ])),
1207                true,
1208            ),
1209        ]));
1210
1211        let planner = Planner::new(schema.clone());
1212
1213        let expected = col("i")
1214            .gt(lit(3_i32))
1215            .and(col("st").field("x").lt_eq(lit(5.0_f32)))
1216            .and(
1217                col("s")
1218                    .eq(lit("str-4"))
1219                    .or(col("s").in_list(vec![lit("str-4"), lit("str-5")], false)),
1220            );
1221
1222        // double quotes
1223        let expr = planner
1224            .parse_filter("i > 3 AND st.x <= 5.0 AND (s == 'str-4' OR s in ('str-4', 'str-5'))")
1225            .unwrap();
1226        assert_eq!(expr, expected);
1227
1228        // single quote
1229        let expr = planner
1230            .parse_filter("i > 3 AND st.x <= 5.0 AND (s = 'str-4' OR s in ('str-4', 'str-5'))")
1231            .unwrap();
1232
1233        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1234
1235        let batch = RecordBatch::try_new(
1236            schema,
1237            vec![
1238                Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef,
1239                Arc::new(StringArray::from_iter_values(
1240                    (0..10).map(|v| format!("str-{}", v)),
1241                )),
1242                Arc::new(StructArray::from(vec![
1243                    (
1244                        Arc::new(Field::new("x", DataType::Float32, false)),
1245                        Arc::new(Float32Array::from_iter_values((0..10).map(|v| v as f32)))
1246                            as ArrayRef,
1247                    ),
1248                    (
1249                        Arc::new(Field::new("y", DataType::Float32, false)),
1250                        Arc::new(Float32Array::from_iter_values(
1251                            (0..10).map(|v| (v * 10) as f32),
1252                        )),
1253                    ),
1254                ])),
1255            ],
1256        )
1257        .unwrap();
1258        let predicates = physical_expr.evaluate(&batch).unwrap();
1259        assert_eq!(
1260            predicates.into_array(0).unwrap().as_ref(),
1261            &BooleanArray::from(vec![
1262                false, false, false, false, true, true, false, false, false, false
1263            ])
1264        );
1265    }
1266
1267    #[test]
1268    fn test_parse_filter_uint64_literal_above_i64_max() {
1269        let value = u64::MAX - 1;
1270        let batch = arrow_array::record_batch!(("id", UInt64, [1, value])).unwrap();
1271        let planner = Planner::new(batch.schema());
1272
1273        let expr = planner.parse_filter(&format!("id = {value}")).unwrap();
1274        assert_eq!(expr, col("id").eq(lit(value)));
1275
1276        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1277        let predicates = physical_expr.evaluate(&batch).unwrap();
1278        assert_eq!(
1279            predicates.into_array(0).unwrap().as_ref(),
1280            &BooleanArray::from(vec![false, true])
1281        );
1282
1283        let expr = planner
1284            .parse_expr("arrow_cast(NULL, 'Time32(Second)')")
1285            .unwrap();
1286        let expr = planner.optimize_expr(expr).unwrap();
1287        assert_eq!(expr, Expr::Literal(ScalarValue::Time32Second(None), None));
1288    }
1289
1290    #[test]
1291    fn test_parse_deep_logical_filter() {
1292        let planner = Planner::new(Arc::new(Schema::empty()));
1293
1294        for op in ["AND", "OR"] {
1295            let filter = std::iter::repeat_n("true", 1000)
1296                .collect::<Vec<_>>()
1297                .join(&format!(" {op} "));
1298
1299            let expr = planner.parse_filter(&filter).unwrap();
1300            let optimized = planner.optimize_expr(expr).unwrap();
1301
1302            assert_eq!(optimized, lit(true));
1303        }
1304    }
1305
1306    #[derive(Debug, Eq, PartialEq, Hash)]
1307    struct StrictFloat64Udf {
1308        signature: Signature,
1309    }
1310
1311    impl StrictFloat64Udf {
1312        fn new() -> Self {
1313            Self {
1314                signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
1315            }
1316        }
1317    }
1318
1319    impl ScalarUDFImpl for StrictFloat64Udf {
1320        fn name(&self) -> &str {
1321            "strict_float64"
1322        }
1323
1324        fn signature(&self) -> &Signature {
1325            &self.signature
1326        }
1327
1328        fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
1329            Ok(DataType::Float64)
1330        }
1331
1332        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
1333            let data_type = args.args[0].data_type();
1334            assert_eq!(
1335                data_type,
1336                DataType::Float64,
1337                "strict_float64 expected Float64, got {data_type}"
1338            );
1339            Ok(ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))))
1340        }
1341    }
1342
1343    #[test]
1344    fn test_coerce_before_simplify() {
1345        let planner = Planner::new(Arc::new(Schema::empty()));
1346        let strict_float64 = Arc::new(ScalarUDF::new_from_impl(StrictFloat64Udf::new()));
1347        let expr = Expr::ScalarFunction(ScalarFunction::new_udf(strict_float64, vec![lit(0_i64)]))
1348            .eq(lit(0.0_f64));
1349
1350        let optimized = planner.optimize_expr(expr).unwrap();
1351
1352        planner.create_physical_expr(&optimized).unwrap();
1353    }
1354
1355    #[test]
1356    fn test_nested_col_refs() {
1357        let schema = Arc::new(Schema::new(vec![
1358            Field::new("s0", DataType::Utf8, true),
1359            Field::new(
1360                "st",
1361                DataType::Struct(Fields::from(vec![
1362                    Field::new("s1", DataType::Utf8, true),
1363                    Field::new(
1364                        "st",
1365                        DataType::Struct(Fields::from(vec![Field::new(
1366                            "s2",
1367                            DataType::Utf8,
1368                            true,
1369                        )])),
1370                        true,
1371                    ),
1372                ])),
1373                true,
1374            ),
1375        ]));
1376
1377        let planner = Planner::new(schema);
1378
1379        fn assert_column_eq(planner: &Planner, expr: &str, expected: &Expr) {
1380            let expr = planner.parse_filter(&format!("{expr} = 'val'")).unwrap();
1381            assert!(matches!(
1382                expr,
1383                Expr::BinaryExpr(BinaryExpr {
1384                    left: _,
1385                    op: Operator::Eq,
1386                    right: _
1387                })
1388            ));
1389            if let Expr::BinaryExpr(BinaryExpr { left, .. }) = expr {
1390                assert_eq!(left.as_ref(), expected);
1391            }
1392        }
1393
1394        let expected = Expr::Column(Column::new_unqualified("s0"));
1395        assert_column_eq(&planner, "s0", &expected);
1396        assert_column_eq(&planner, "`s0`", &expected);
1397
1398        let expected = Expr::ScalarFunction(ScalarFunction {
1399            func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1400            args: vec![
1401                Expr::Column(Column::new_unqualified("st")),
1402                Expr::Literal(ScalarValue::Utf8(Some("s1".to_string())), None),
1403            ],
1404        });
1405        assert_column_eq(&planner, "st.s1", &expected);
1406        assert_column_eq(&planner, "`st`.`s1`", &expected);
1407        assert_column_eq(&planner, "st.`s1`", &expected);
1408
1409        let expected = Expr::ScalarFunction(ScalarFunction {
1410            func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1411            args: vec![
1412                Expr::ScalarFunction(ScalarFunction {
1413                    func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1414                    args: vec![
1415                        Expr::Column(Column::new_unqualified("st")),
1416                        Expr::Literal(ScalarValue::Utf8(Some("st".to_string())), None),
1417                    ],
1418                }),
1419                Expr::Literal(ScalarValue::Utf8(Some("s2".to_string())), None),
1420            ],
1421        });
1422
1423        assert_column_eq(&planner, "st.st.s2", &expected);
1424        assert_column_eq(&planner, "`st`.`st`.`s2`", &expected);
1425        assert_column_eq(&planner, "st.st.`s2`", &expected);
1426        assert_column_eq(&planner, "st['st'][\"s2\"]", &expected);
1427    }
1428
1429    #[test]
1430    fn test_nested_list_refs() {
1431        let schema = Arc::new(Schema::new(vec![Field::new(
1432            "l",
1433            DataType::List(Arc::new(Field::new(
1434                "item",
1435                DataType::Struct(Fields::from(vec![Field::new("f1", DataType::Utf8, true)])),
1436                true,
1437            ))),
1438            true,
1439        )]));
1440
1441        let planner = Planner::new(schema);
1442
1443        let expected = array_element(col("l"), lit(0_i64));
1444        let expr = planner.parse_expr("l[0]").unwrap();
1445        assert_eq!(expr, expected);
1446
1447        let expected = get_field(array_element(col("l"), lit(0_i64)), "f1");
1448        let expr = planner.parse_expr("l[0]['f1']").unwrap();
1449        assert_eq!(expr, expected);
1450
1451        // FIXME: This should work, but sqlparser doesn't recognize anything
1452        // after the period for some reason.
1453        // let expr = planner.parse_expr("l[0].f1").unwrap();
1454        // assert_eq!(expr, expected);
1455    }
1456
1457    #[test]
1458    fn test_negative_expressions() {
1459        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
1460
1461        let planner = Planner::new(schema.clone());
1462
1463        let expected = col("x")
1464            .gt(lit(-3_i64))
1465            .and(col("x").lt(-(lit(-5_i64) + lit(3_i64))));
1466
1467        let expr = planner.parse_filter("x > -3 AND x < -(-5 + 3)").unwrap();
1468
1469        assert_eq!(expr, expected);
1470
1471        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1472
1473        let batch = RecordBatch::try_new(
1474            schema,
1475            vec![Arc::new(Int64Array::from_iter_values(-5..5)) as ArrayRef],
1476        )
1477        .unwrap();
1478        let predicates = physical_expr.evaluate(&batch).unwrap();
1479        assert_eq!(
1480            predicates.into_array(0).unwrap().as_ref(),
1481            &BooleanArray::from(vec![
1482                false, false, false, true, true, true, true, false, false, false
1483            ])
1484        );
1485    }
1486
1487    #[rstest]
1488    #[case::right("value >> 32", Operator::BitwiseShiftRight, vec![0, 1, 3])]
1489    #[case::left(
1490        "value << 1",
1491        Operator::BitwiseShiftLeft,
1492        vec![0, 2_u64 << 32, ((3_u64 << 32) + 7) << 1]
1493    )]
1494    fn test_bitwise_shift_expressions(
1495        #[case] sql: &str,
1496        #[case] expected_op: Operator,
1497        #[case] expected: Vec<u64>,
1498    ) {
1499        let input = vec![0, 1_u64 << 32, (3_u64 << 32) + 7];
1500        let batch =
1501            RecordBatch::try_from_iter([("value", Arc::new(UInt64Array::from(input)) as ArrayRef)])
1502                .unwrap();
1503        let planner = Planner::new(batch.schema());
1504
1505        let expr = planner.parse_expr(sql).unwrap();
1506        let Expr::BinaryExpr(binary_expr) = &expr else {
1507            panic!("expected binary expression for {sql}, got {expr}");
1508        };
1509        assert_eq!(binary_expr.op, expected_op);
1510
1511        let expr = planner.optimize_expr(expr).unwrap();
1512        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1513        let values = physical_expr
1514            .evaluate(&batch)
1515            .unwrap()
1516            .into_array(batch.num_rows())
1517            .unwrap();
1518        assert_eq!(values.as_ref(), &UInt64Array::from(expected));
1519    }
1520
1521    #[test]
1522    fn test_negative_array_expressions() {
1523        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
1524
1525        let planner = Planner::new(schema);
1526
1527        let expected = Expr::Literal(
1528            ScalarValue::List(Arc::new(
1529                ListArray::from_iter_primitive::<Float64Type, _, _>(vec![Some(
1530                    [-1_f64, -2.0, -3.0, -4.0, -5.0].map(Some),
1531                )]),
1532            )),
1533            None,
1534        );
1535
1536        let expr = planner
1537            .parse_expr("[-1.0, -2.0, -3.0, -4.0, -5.0]")
1538            .unwrap();
1539
1540        assert_eq!(expr, expected);
1541    }
1542
1543    #[test]
1544    fn test_sql_like() {
1545        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1546
1547        let planner = Planner::new(schema.clone());
1548
1549        let expected = col("s").like(lit("str-4"));
1550        // single quote
1551        let expr = planner.parse_filter("s LIKE 'str-4'").unwrap();
1552        assert_eq!(expr, expected);
1553        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1554
1555        let batch = RecordBatch::try_new(
1556            schema,
1557            vec![Arc::new(StringArray::from_iter_values(
1558                (0..10).map(|v| format!("str-{}", v)),
1559            ))],
1560        )
1561        .unwrap();
1562        let predicates = physical_expr.evaluate(&batch).unwrap();
1563        assert_eq!(
1564            predicates.into_array(0).unwrap().as_ref(),
1565            &BooleanArray::from(vec![
1566                false, false, false, false, true, false, false, false, false, false
1567            ])
1568        );
1569    }
1570
1571    #[test]
1572    fn test_not_like() {
1573        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1574
1575        let planner = Planner::new(schema.clone());
1576
1577        let expected = col("s").not_like(lit("str-4"));
1578        // single quote
1579        let expr = planner.parse_filter("s NOT LIKE 'str-4'").unwrap();
1580        assert_eq!(expr, expected);
1581        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1582
1583        let batch = RecordBatch::try_new(
1584            schema,
1585            vec![Arc::new(StringArray::from_iter_values(
1586                (0..10).map(|v| format!("str-{}", v)),
1587            ))],
1588        )
1589        .unwrap();
1590        let predicates = physical_expr.evaluate(&batch).unwrap();
1591        assert_eq!(
1592            predicates.into_array(0).unwrap().as_ref(),
1593            &BooleanArray::from(vec![
1594                true, true, true, true, false, true, true, true, true, true
1595            ])
1596        );
1597    }
1598
1599    #[test]
1600    fn test_like_escape_char() {
1601        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1602        let planner = Planner::new(schema);
1603
1604        // A valid single-character escape is captured for both LIKE and ILIKE.
1605        for filter in ["s LIKE 'a!%' ESCAPE '!'", "s ILIKE 'a!%' ESCAPE '!'"] {
1606            match planner.parse_filter(filter).unwrap() {
1607                Expr::Like(like) => assert_eq!(like.escape_char, Some('!'), "{filter}"),
1608                other => panic!("expected a LIKE expression for `{filter}`, got {other:?}"),
1609            }
1610        }
1611
1612        // Empty and multi-character escapes are rejected rather than silently
1613        // dropped or truncated to the first character.
1614        for filter in [
1615            "s LIKE 'x' ESCAPE ''",
1616            "s LIKE 'x' ESCAPE 'ab'",
1617            "s ILIKE 'x' ESCAPE ''",
1618            "s ILIKE 'x' ESCAPE 'ab'",
1619        ] {
1620            let err = planner.parse_filter(filter).unwrap_err();
1621            assert!(
1622                err.to_string()
1623                    .contains("Invalid escape character in LIKE expression"),
1624                "unexpected error for `{filter}`: {err}"
1625            );
1626        }
1627    }
1628
1629    #[test]
1630    fn test_sql_is_in() {
1631        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1632
1633        let planner = Planner::new(schema.clone());
1634
1635        let expected = col("s").in_list(vec![lit("str-4"), lit("str-5")], false);
1636        // single quote
1637        let expr = planner.parse_filter("s IN ('str-4', 'str-5')").unwrap();
1638        assert_eq!(expr, expected);
1639        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1640
1641        let batch = RecordBatch::try_new(
1642            schema,
1643            vec![Arc::new(StringArray::from_iter_values(
1644                (0..10).map(|v| format!("str-{}", v)),
1645            ))],
1646        )
1647        .unwrap();
1648        let predicates = physical_expr.evaluate(&batch).unwrap();
1649        assert_eq!(
1650            predicates.into_array(0).unwrap().as_ref(),
1651            &BooleanArray::from(vec![
1652                false, false, false, false, true, true, false, false, false, false
1653            ])
1654        );
1655    }
1656
1657    #[test]
1658    fn test_sql_is_null() {
1659        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1660
1661        let planner = Planner::new(schema.clone());
1662
1663        let expected = col("s").is_null();
1664        let expr = planner.parse_filter("s IS NULL").unwrap();
1665        assert_eq!(expr, expected);
1666        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1667
1668        let batch = RecordBatch::try_new(
1669            schema,
1670            vec![Arc::new(StringArray::from_iter((0..10).map(|v| {
1671                if v % 3 == 0 {
1672                    Some(format!("str-{}", v))
1673                } else {
1674                    None
1675                }
1676            })))],
1677        )
1678        .unwrap();
1679        let predicates = physical_expr.evaluate(&batch).unwrap();
1680        assert_eq!(
1681            predicates.into_array(0).unwrap().as_ref(),
1682            &BooleanArray::from(vec![
1683                false, true, true, false, true, true, false, true, true, false
1684            ])
1685        );
1686
1687        let expr = planner.parse_filter("s IS NOT NULL").unwrap();
1688        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1689        let predicates = physical_expr.evaluate(&batch).unwrap();
1690        assert_eq!(
1691            predicates.into_array(0).unwrap().as_ref(),
1692            &BooleanArray::from(vec![
1693                true, false, false, true, false, false, true, false, false, true,
1694            ])
1695        );
1696    }
1697
1698    #[test]
1699    fn test_sql_invert() {
1700        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Boolean, true)]));
1701
1702        let planner = Planner::new(schema.clone());
1703
1704        let expr = planner.parse_filter("NOT s").unwrap();
1705        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1706
1707        let batch = RecordBatch::try_new(
1708            schema,
1709            vec![Arc::new(BooleanArray::from_iter(
1710                (0..10).map(|v| Some(v % 3 == 0)),
1711            ))],
1712        )
1713        .unwrap();
1714        let predicates = physical_expr.evaluate(&batch).unwrap();
1715        assert_eq!(
1716            predicates.into_array(0).unwrap().as_ref(),
1717            &BooleanArray::from(vec![
1718                false, true, true, false, true, true, false, true, true, false
1719            ])
1720        );
1721    }
1722
1723    #[test]
1724    fn test_sql_cast() {
1725        let cases = &[
1726            (
1727                "x = cast('2021-01-01 00:00:00' as timestamp)",
1728                ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1729            ),
1730            (
1731                "x = cast('2021-01-01 00:00:00' as timestamp(0))",
1732                ArrowDataType::Timestamp(TimeUnit::Second, None),
1733            ),
1734            (
1735                "x = cast('2021-01-01 00:00:00.123' as timestamp(9))",
1736                ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1737            ),
1738            (
1739                "x = cast('2021-01-01 00:00:00.123' as datetime(9))",
1740                ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1741            ),
1742            ("x = cast('2021-01-01' as date)", ArrowDataType::Date32),
1743            (
1744                "x = cast('1.238' as decimal(9,3))",
1745                ArrowDataType::Decimal128(9, 3),
1746            ),
1747            ("x = cast(1 as float)", ArrowDataType::Float32),
1748            ("x = cast(1 as double)", ArrowDataType::Float64),
1749            ("x = cast(1 as tinyint)", ArrowDataType::Int8),
1750            ("x = cast(1 as smallint)", ArrowDataType::Int16),
1751            ("x = cast(1 as int)", ArrowDataType::Int32),
1752            ("x = cast(1 as integer)", ArrowDataType::Int32),
1753            ("x = cast(1 as bigint)", ArrowDataType::Int64),
1754            ("x = cast(1 as tinyint unsigned)", ArrowDataType::UInt8),
1755            ("x = cast(1 as smallint unsigned)", ArrowDataType::UInt16),
1756            ("x = cast(1 as int unsigned)", ArrowDataType::UInt32),
1757            ("x = cast(1 as integer unsigned)", ArrowDataType::UInt32),
1758            ("x = cast(1 as bigint unsigned)", ArrowDataType::UInt64),
1759            ("x = cast(1 as boolean)", ArrowDataType::Boolean),
1760            ("x = cast(1 as string)", ArrowDataType::Utf8),
1761        ];
1762
1763        for (sql, expected_data_type) in cases {
1764            let schema = Arc::new(Schema::new(vec![Field::new(
1765                "x",
1766                expected_data_type.clone(),
1767                true,
1768            )]));
1769            let planner = Planner::new(schema.clone());
1770            let expr = planner.parse_filter(sql).unwrap();
1771
1772            // Get the thing after 'cast(` but before ' as'.
1773            let expected_value_str = sql
1774                .split("cast(")
1775                .nth(1)
1776                .unwrap()
1777                .split(" as")
1778                .next()
1779                .unwrap();
1780            // Remove any quote marks
1781            let expected_value_str = expected_value_str.trim_matches('\'');
1782
1783            match expr {
1784                Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
1785                    Expr::Cast(Cast { expr, field }) => {
1786                        match expr.as_ref() {
1787                            Expr::Literal(ScalarValue::Utf8(Some(value_str)), _) => {
1788                                assert_eq!(value_str, expected_value_str);
1789                            }
1790                            Expr::Literal(ScalarValue::Int64(Some(value)), _) => {
1791                                assert_eq!(*value, 1);
1792                            }
1793                            _ => panic!("Expected cast to be applied to literal"),
1794                        }
1795                        assert_eq!(field.data_type(), expected_data_type);
1796                    }
1797                    _ => panic!("Expected right to be a cast"),
1798                },
1799                _ => panic!("Expected binary expression"),
1800            }
1801        }
1802    }
1803
1804    #[test]
1805    fn test_arrow_cast_int_literal_to_time32() {
1806        let batch = RecordBatch::try_from_iter([(
1807            "v",
1808            Arc::new(Time32SecondArray::from(vec![3725, 3726])) as ArrayRef,
1809        )])
1810        .unwrap();
1811        let planner = Planner::new(batch.schema());
1812
1813        let expr = planner
1814            .parse_filter("v = arrow_cast(3726, 'Time32(Second)')")
1815            .unwrap();
1816        let expr = planner.optimize_expr(expr).unwrap();
1817        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1818        let predicates = physical_expr.evaluate(&batch).unwrap();
1819
1820        assert_eq!(
1821            predicates.into_array(0).unwrap().as_ref(),
1822            &BooleanArray::from(vec![false, true])
1823        );
1824    }
1825
1826    #[test]
1827    fn test_sql_literals() {
1828        let cases = &[
1829            (
1830                "x = timestamp '2021-01-01 00:00:00'",
1831                ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1832            ),
1833            (
1834                "x = timestamp(0) '2021-01-01 00:00:00'",
1835                ArrowDataType::Timestamp(TimeUnit::Second, None),
1836            ),
1837            (
1838                "x = timestamp(9) '2021-01-01 00:00:00.123'",
1839                ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1840            ),
1841            ("x = date '2021-01-01'", ArrowDataType::Date32),
1842            ("x = decimal(9,3) '1.238'", ArrowDataType::Decimal128(9, 3)),
1843        ];
1844
1845        for (sql, expected_data_type) in cases {
1846            let schema = Arc::new(Schema::new(vec![Field::new(
1847                "x",
1848                expected_data_type.clone(),
1849                true,
1850            )]));
1851            let planner = Planner::new(schema.clone());
1852            let expr = planner.parse_filter(sql).unwrap();
1853
1854            let expected_value_str = sql.split('\'').nth(1).unwrap();
1855
1856            match expr {
1857                Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
1858                    Expr::Cast(Cast { expr, field }) => {
1859                        match expr.as_ref() {
1860                            Expr::Literal(ScalarValue::Utf8(Some(value_str)), _) => {
1861                                assert_eq!(value_str, expected_value_str);
1862                            }
1863                            _ => panic!("Expected cast to be applied to literal"),
1864                        }
1865                        assert_eq!(field.data_type(), expected_data_type);
1866                    }
1867                    _ => panic!("Expected right to be a cast"),
1868                },
1869                _ => panic!("Expected binary expression"),
1870            }
1871        }
1872    }
1873
1874    #[test]
1875    fn test_sql_array_literals() {
1876        let cases = [
1877            (
1878                "x = [1, 2, 3]",
1879                ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int64, true))),
1880            ),
1881            (
1882                "x = [1, 2, 3]",
1883                ArrowDataType::FixedSizeList(
1884                    Arc::new(Field::new("item", ArrowDataType::Int64, true)),
1885                    3,
1886                ),
1887            ),
1888        ];
1889
1890        for (sql, expected_data_type) in cases {
1891            let schema = Arc::new(Schema::new(vec![Field::new(
1892                "x",
1893                expected_data_type.clone(),
1894                true,
1895            )]));
1896            let planner = Planner::new(schema.clone());
1897            let expr = planner.parse_filter(sql).unwrap();
1898            let expr = planner.optimize_expr(expr).unwrap();
1899
1900            match expr {
1901                Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
1902                    Expr::Literal(value, _) => {
1903                        assert_eq!(&value.data_type(), &expected_data_type);
1904                    }
1905                    _ => panic!("Expected right to be a literal"),
1906                },
1907                _ => panic!("Expected binary expression"),
1908            }
1909        }
1910    }
1911
1912    #[test]
1913    fn test_sql_between() {
1914        use arrow_array::{Float64Array, Int32Array, TimestampMicrosecondArray};
1915        use arrow_schema::{DataType, Field, Schema, TimeUnit};
1916        use std::sync::Arc;
1917
1918        let schema = Arc::new(Schema::new(vec![
1919            Field::new("x", DataType::Int32, false),
1920            Field::new("y", DataType::Float64, false),
1921            Field::new(
1922                "ts",
1923                DataType::Timestamp(TimeUnit::Microsecond, None),
1924                false,
1925            ),
1926        ]));
1927
1928        let planner = Planner::new(schema.clone());
1929
1930        // Test integer BETWEEN
1931        let expr = planner
1932            .parse_filter("x BETWEEN CAST(3 AS INT) AND CAST(7 AS INT)")
1933            .unwrap();
1934        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1935
1936        // Create timestamp array with values representing:
1937        // 2024-01-01 00:00:00 to 2024-01-01 00:00:09 (in microseconds)
1938        let base_ts = 1704067200000000_i64; // 2024-01-01 00:00:00
1939        let ts_array = TimestampMicrosecondArray::from_iter_values(
1940            (0..10).map(|i| base_ts + i * 1_000_000), // Each value is 1 second apart
1941        );
1942
1943        let batch = RecordBatch::try_new(
1944            schema,
1945            vec![
1946                Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef,
1947                Arc::new(Float64Array::from_iter_values((0..10).map(|v| v as f64))),
1948                Arc::new(ts_array),
1949            ],
1950        )
1951        .unwrap();
1952
1953        let predicates = physical_expr.evaluate(&batch).unwrap();
1954        assert_eq!(
1955            predicates.into_array(0).unwrap().as_ref(),
1956            &BooleanArray::from(vec![
1957                false, false, false, true, true, true, true, true, false, false
1958            ])
1959        );
1960
1961        // Test NOT BETWEEN
1962        let expr = planner
1963            .parse_filter("x NOT BETWEEN CAST(3 AS INT) AND CAST(7 AS INT)")
1964            .unwrap();
1965        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1966
1967        let predicates = physical_expr.evaluate(&batch).unwrap();
1968        assert_eq!(
1969            predicates.into_array(0).unwrap().as_ref(),
1970            &BooleanArray::from(vec![
1971                true, true, true, false, false, false, false, false, true, true
1972            ])
1973        );
1974
1975        // Test floating point BETWEEN
1976        let expr = planner.parse_filter("y BETWEEN 2.5 AND 6.5").unwrap();
1977        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1978
1979        let predicates = physical_expr.evaluate(&batch).unwrap();
1980        assert_eq!(
1981            predicates.into_array(0).unwrap().as_ref(),
1982            &BooleanArray::from(vec![
1983                false, false, false, true, true, true, true, false, false, false
1984            ])
1985        );
1986
1987        // Test timestamp BETWEEN
1988        let expr = planner
1989            .parse_filter(
1990                "ts BETWEEN timestamp '2024-01-01 00:00:03' AND timestamp '2024-01-01 00:00:07'",
1991            )
1992            .unwrap();
1993        let physical_expr = planner.create_physical_expr(&expr).unwrap();
1994
1995        let predicates = physical_expr.evaluate(&batch).unwrap();
1996        assert_eq!(
1997            predicates.into_array(0).unwrap().as_ref(),
1998            &BooleanArray::from(vec![
1999                false, false, false, true, true, true, true, true, false, false
2000            ])
2001        );
2002    }
2003
2004    #[test]
2005    fn test_sql_comparison() {
2006        // Create a batch with all data types
2007        let batch: Vec<(&str, ArrayRef)> = vec![
2008            (
2009                "timestamp_s",
2010                Arc::new(TimestampSecondArray::from_iter_values(0..10)),
2011            ),
2012            (
2013                "timestamp_ms",
2014                Arc::new(TimestampMillisecondArray::from_iter_values(0..10)),
2015            ),
2016            (
2017                "timestamp_us",
2018                Arc::new(TimestampMicrosecondArray::from_iter_values(0..10)),
2019            ),
2020            (
2021                "timestamp_ns",
2022                Arc::new(TimestampNanosecondArray::from_iter_values(4995..5005)),
2023            ),
2024        ];
2025        let batch = RecordBatch::try_from_iter(batch).unwrap();
2026
2027        let planner = Planner::new(batch.schema());
2028
2029        // Each expression is meant to select the final 5 rows
2030        let expressions = &[
2031            "timestamp_s >= TIMESTAMP '1970-01-01 00:00:05'",
2032            "timestamp_ms >= TIMESTAMP '1970-01-01 00:00:00.005'",
2033            "timestamp_us >= TIMESTAMP '1970-01-01 00:00:00.000005'",
2034            "timestamp_ns >= TIMESTAMP '1970-01-01 00:00:00.000005'",
2035        ];
2036
2037        let expected: ArrayRef = Arc::new(BooleanArray::from_iter(
2038            std::iter::repeat_n(Some(false), 5).chain(std::iter::repeat_n(Some(true), 5)),
2039        ));
2040        for expression in expressions {
2041            // convert to physical expression
2042            let logical_expr = planner.parse_filter(expression).unwrap();
2043            let logical_expr = planner.optimize_expr(logical_expr).unwrap();
2044            let physical_expr = planner.create_physical_expr(&logical_expr).unwrap();
2045
2046            // Evaluate and assert they have correct results
2047            let result = physical_expr.evaluate(&batch).unwrap();
2048            let result = result.into_array(batch.num_rows()).unwrap();
2049            assert_eq!(&expected, &result, "unexpected result for {}", expression);
2050        }
2051    }
2052
2053    #[test]
2054    fn test_columns_in_expr() {
2055        let expr = col("s0").gt(lit("value")).and(
2056            col("st")
2057                .field("st")
2058                .field("s2")
2059                .eq(lit("value"))
2060                .or(col("st")
2061                    .field("s1")
2062                    .in_list(vec![lit("value 1"), lit("value 2")], false)),
2063        );
2064
2065        let columns = Planner::column_names_in_expr(&expr);
2066        assert_eq!(columns, vec!["s0", "st.s1", "st.st.s2"]);
2067    }
2068
2069    #[test]
2070    fn test_parse_binary_expr() {
2071        let bin_str = "x'616263'";
2072
2073        let schema = Arc::new(Schema::new(vec![Field::new(
2074            "binary",
2075            DataType::Binary,
2076            true,
2077        )]));
2078        let planner = Planner::new(schema);
2079        let expr = planner.parse_expr(bin_str).unwrap();
2080        assert_eq!(
2081            expr,
2082            Expr::Literal(ScalarValue::Binary(Some(vec![b'a', b'b', b'c'])), None)
2083        );
2084    }
2085
2086    #[test]
2087    fn test_lance_context_provider_expr_planners() {
2088        let ctx_provider = LanceContextProvider::default();
2089        assert!(!ctx_provider.get_expr_planners().is_empty());
2090    }
2091
2092    #[test]
2093    fn test_lance_context_provider_resolves_session_functions() {
2094        let session = crate::exec::new_session_context(&LanceExecutionOptions::default());
2095        let ctx_provider = LanceContextProvider::new(session.clone());
2096        assert!(ctx_provider.get_function_meta("contains_tokens").is_some());
2097        assert!(ctx_provider.get_function_meta("lower").is_some());
2098        assert!(ctx_provider.get_aggregate_meta("sum").is_some());
2099        assert!(ctx_provider.get_window_meta("row_number").is_some());
2100        assert!(ctx_provider.get_function_meta("registered_later").is_none());
2101
2102        // The provider reads the session's functions instead of a copy, so a
2103        // function registered on the session afterwards resolves as well.
2104        session.register_udf(datafusion::logical_expr::create_udf(
2105            "registered_later",
2106            vec![DataType::Utf8],
2107            DataType::Utf8,
2108            Volatility::Immutable,
2109            Arc::new(|args: &[ColumnarValue]| Ok(args[0].clone())),
2110        ));
2111        assert!(ctx_provider.get_function_meta("registered_later").is_some());
2112        assert!(
2113            ctx_provider
2114                .udf_names()
2115                .contains(&"registered_later".to_string())
2116        );
2117    }
2118
2119    #[test]
2120    fn test_regexp_match_and_non_empty_captions() {
2121        // Repro for a bug where regexp_match inside an AND chain wasn't coerced to boolean,
2122        // causing planning/evaluation failures. This should evaluate successfully.
2123        let schema = Arc::new(Schema::new(vec![
2124            Field::new("keywords", DataType::Utf8, true),
2125            Field::new("natural_caption", DataType::Utf8, true),
2126            Field::new("poetic_caption", DataType::Utf8, true),
2127        ]));
2128
2129        let planner = Planner::new(schema.clone());
2130
2131        let expr = planner
2132            .parse_filter(
2133                "regexp_match(keywords, 'Liberty|revolution') AND \
2134                 (natural_caption IS NOT NULL AND natural_caption <> '' AND \
2135                  poetic_caption IS NOT NULL AND poetic_caption <> '')",
2136            )
2137            .unwrap();
2138
2139        let physical_expr = planner.create_physical_expr(&expr).unwrap();
2140
2141        let batch = RecordBatch::try_new(
2142            schema,
2143            vec![
2144                Arc::new(StringArray::from(vec![
2145                    Some("Liberty for all"),
2146                    Some("peace"),
2147                    Some("revolution now"),
2148                    Some("Liberty"),
2149                    Some("revolutionary"),
2150                    Some("none"),
2151                ])) as ArrayRef,
2152                Arc::new(StringArray::from(vec![
2153                    Some("a"),
2154                    Some("b"),
2155                    None,
2156                    Some(""),
2157                    Some("c"),
2158                    Some("d"),
2159                ])) as ArrayRef,
2160                Arc::new(StringArray::from(vec![
2161                    Some("x"),
2162                    Some(""),
2163                    Some("y"),
2164                    Some("z"),
2165                    None,
2166                    Some("w"),
2167                ])) as ArrayRef,
2168            ],
2169        )
2170        .unwrap();
2171
2172        let result = physical_expr.evaluate(&batch).unwrap();
2173        assert_eq!(
2174            result.into_array(0).unwrap().as_ref(),
2175            &BooleanArray::from(vec![true, false, false, false, false, false])
2176        );
2177    }
2178
2179    #[test]
2180    fn test_regexp_match_infer_error_without_boolean_coercion() {
2181        // With the fix applied, using parse_filter should coerce regexp_match to boolean
2182        // even when nested in a larger AND expression, so this should plan successfully.
2183        let schema = Arc::new(Schema::new(vec![
2184            Field::new("keywords", DataType::Utf8, true),
2185            Field::new("natural_caption", DataType::Utf8, true),
2186            Field::new("poetic_caption", DataType::Utf8, true),
2187        ]));
2188
2189        let planner = Planner::new(schema);
2190
2191        let expr = planner
2192            .parse_filter(
2193                "regexp_match(keywords, 'Liberty|revolution') AND \
2194                 (natural_caption IS NOT NULL AND natural_caption <> '' AND \
2195                  poetic_caption IS NOT NULL AND poetic_caption <> '')",
2196            )
2197            .unwrap();
2198
2199        // Should not panic
2200        let _physical = planner.create_physical_expr(&expr).unwrap();
2201    }
2202
2203    #[test]
2204    fn test_jsonb_literals() {
2205        let schema = Arc::new(Schema::new(vec![Field::new(
2206            "j",
2207            DataType::LargeBinary,
2208            true,
2209        )]));
2210        let planner = Planner::new(schema);
2211
2212        let cases = [
2213            ("jsonb '{\"key\": \"value\"}'", r#"{"key":"value"}"#),
2214            ("cast('{\"a\": 1}' as jsonb)", r#"{"a":1}"#),
2215            ("'{\"a\": 1}'::jsonb", r#"{"a":1}"#),
2216        ];
2217        for (sql, expected) in cases {
2218            let ast = parse_sql_expr(sql).unwrap();
2219            let expr = planner.parse_sql_expr(&ast).unwrap();
2220            match expr {
2221                Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), _) => {
2222                    assert_eq!(
2223                        lance_arrow::json::decode_json(&bytes),
2224                        expected,
2225                        "failed for: {sql}"
2226                    );
2227                }
2228                other => panic!("Expected LargeBinary literal for '{sql}', got: {other:?}"),
2229            }
2230        }
2231    }
2232
2233    #[test]
2234    fn test_jsonb_literal_errors() {
2235        let schema = Arc::new(Schema::new(vec![Field::new(
2236            "j",
2237            DataType::LargeBinary,
2238            true,
2239        )]));
2240        let planner = Planner::new(schema);
2241
2242        // Invalid JSON content
2243        let ast = parse_sql_expr("jsonb 'not valid json'").unwrap();
2244        let err = planner.parse_sql_expr(&ast).unwrap_err();
2245        assert!(
2246            err.to_string().contains("Failed to encode JSONB"),
2247            "expected JSONB encoding error, got: {err}"
2248        );
2249
2250        // CAST with non-literal expression
2251        let ast = parse_sql_expr("cast(j as jsonb)").unwrap();
2252        let err = planner.parse_sql_expr(&ast).unwrap_err();
2253        assert!(
2254            err.to_string()
2255                .contains("CAST to JSONB only supports string literals"),
2256            "got: {err}"
2257        );
2258    }
2259}