1use 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, parse_field_path};
51use lance_core::error::LanceOptionExt;
52
53use chrono::Utc;
54use lance_core::{Error, Result};
55
56fn 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
63fn 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
183struct LanceContextProvider {
185 options: datafusion::config::ConfigOptions,
186 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::common::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 "_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 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 pub fn with_enable_relations(mut self, enable_relations: bool) -> Self {
301 self.enable_relations = enable_relations;
302 self
303 }
304
305 fn resolve_column_name(&self, name: &str) -> String {
308 if self.schema.field_with_name(name).is_ok() {
310 return name.to_string();
311 }
312 for field in self.schema.fields() {
314 if field.name().eq_ignore_ascii_case(name) {
315 return field.name().clone();
316 }
317 }
318 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 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 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 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 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 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 if id.quote_style == Some('"') {
696 Ok(Expr::Literal(
697 ScalarValue::Utf8(Some(id.value.clone())),
698 None,
699 ))
700 } 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 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 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::IsDistinctFrom(left, right) => Ok(Expr::BinaryExpr(BinaryExpr::new(
811 Box::new(self.parse_sql_expr(left)?),
812 Operator::IsDistinctFrom,
813 Box::new(self.parse_sql_expr(right)?),
814 ))),
815 SQLExpr::IsNotDistinctFrom(left, right) => Ok(Expr::BinaryExpr(BinaryExpr::new(
816 Box::new(self.parse_sql_expr(left)?),
817 Operator::IsNotDistinctFrom,
818 Box::new(self.parse_sql_expr(right)?),
819 ))),
820 SQLExpr::InList {
821 expr,
822 list,
823 negated,
824 } => {
825 let value_expr = self.parse_sql_expr(expr)?;
826 let list_exprs = list
827 .iter()
828 .map(|e| self.parse_sql_expr(e))
829 .collect::<Result<Vec<_>>>()?;
830 Ok(value_expr.in_list(list_exprs, *negated))
831 }
832 SQLExpr::Nested(inner) => self.parse_sql_expr(inner.as_ref()),
833 SQLExpr::Function(_) => self.parse_function(expr.clone()),
834 SQLExpr::ILike {
835 negated,
836 expr,
837 pattern,
838 escape_char,
839 any: _,
840 } => Ok(Expr::Like(Like::new(
841 *negated,
842 Box::new(self.parse_sql_expr(expr)?),
843 Box::new(self.parse_sql_expr(pattern)?),
844 parse_like_escape_char(escape_char)?,
845 true,
846 ))),
847 SQLExpr::Like {
848 negated,
849 expr,
850 pattern,
851 escape_char,
852 any: _,
853 } => Ok(Expr::Like(Like::new(
854 *negated,
855 Box::new(self.parse_sql_expr(expr)?),
856 Box::new(self.parse_sql_expr(pattern)?),
857 parse_like_escape_char(escape_char)?,
858 false,
859 ))),
860 SQLExpr::Cast {
862 data_type: SQLDataType::JSONB,
863 expr: inner,
864 ..
865 } => match inner.as_ref() {
866 SQLExpr::Value(ValueWithSpan {
867 value: Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
868 ..
869 }) => encode_jsonb(s),
870 _ => Err(Error::invalid_input(
871 "CAST to JSONB only supports string literals",
872 )),
873 },
874 SQLExpr::Cast {
875 expr,
876 data_type,
877 kind,
878 ..
879 } => match kind {
880 datafusion::sql::sqlparser::ast::CastKind::TryCast
881 | datafusion::sql::sqlparser::ast::CastKind::SafeCast => {
882 Ok(Expr::TryCast(datafusion::logical_expr::TryCast::new(
883 Box::new(self.parse_sql_expr(expr)?),
884 self.parse_type(data_type)?,
885 )))
886 }
887 _ => Ok(Expr::Cast(datafusion::logical_expr::Cast::new(
888 Box::new(self.parse_sql_expr(expr)?),
889 self.parse_type(data_type)?,
890 ))),
891 },
892 SQLExpr::JsonAccess { .. } => Err(Error::invalid_input("JSON access is not supported")),
893 SQLExpr::CompoundFieldAccess { root, access_chain } => {
894 let mut expr = self.parse_sql_expr(root)?;
895
896 for access in access_chain {
897 let field_access = match access {
898 AccessExpr::Dot(SQLExpr::Identifier(Ident { value: s, .. }))
900 | AccessExpr::Subscript(Subscript::Index {
901 index:
902 SQLExpr::Value(ValueWithSpan {
903 value:
904 Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
905 ..
906 }),
907 }) => GetFieldAccess::NamedStructField {
908 name: ScalarValue::from(s.as_str()),
909 },
910 AccessExpr::Subscript(Subscript::Index { index }) => {
911 let key = Box::new(self.parse_sql_expr(index)?);
912 GetFieldAccess::ListIndex { key }
913 }
914 AccessExpr::Subscript(Subscript::Slice { .. }) => {
915 return Err(Error::invalid_input("Slice subscript is not supported"));
916 }
917 _ => {
918 return Err(Error::invalid_input(
921 "Only dot notation or index access is supported for field access",
922 ));
923 }
924 };
925
926 let field_access_expr = RawFieldAccessExpr { expr, field_access };
927 expr = self.plan_field_access(field_access_expr)?;
928 }
929
930 Ok(expr)
931 }
932 SQLExpr::Between {
933 expr,
934 negated,
935 low,
936 high,
937 } => {
938 let expr = self.parse_sql_expr(expr)?;
940 let low = self.parse_sql_expr(low)?;
941 let high = self.parse_sql_expr(high)?;
942
943 let between = Expr::Between(Between::new(
944 Box::new(expr),
945 *negated,
946 Box::new(low),
947 Box::new(high),
948 ));
949 Ok(between)
950 }
951 _ => Err(Error::invalid_input(format!(
952 "Expression '{expr}' is not supported SQL in lance"
953 ))),
954 }
955 }
956
957 pub fn parse_filter(&self, filter: &str) -> Result<Expr> {
962 let ast_expr = parse_sql_filter(filter)?;
964 let expr = self.parse_sql_expr(&ast_expr)?;
965 let schema = Schema::try_from(self.schema.as_ref())?;
966 let resolved = resolve_expr(&expr, &schema).map_err(|e| {
967 Error::invalid_input(format!("Error resolving filter expression {filter}: {e}"))
968 })?;
969
970 Ok(coerce_filter_type_to_boolean(resolved))
971 }
972
973 pub fn parse_expr(&self, expr: &str) -> Result<Expr> {
985 let resolved_name = self.resolve_column_name(expr);
990 if self.schema.field_with_name(&resolved_name).is_ok() {
991 return Ok(Expr::Column(Column::from_name(resolved_name)));
992 }
993
994 if let Ok(path) = parse_field_path(expr)
999 && let [name] = path.as_slice()
1000 && self.schema.field_with_name(name).is_ok()
1001 {
1002 return Ok(Expr::Column(Column::from_name(name)));
1003 }
1004
1005 let ast_expr = parse_sql_expr(expr)?;
1007 let expr = self.parse_sql_expr(&ast_expr)?;
1008 let schema = Schema::try_from(self.schema.as_ref())?;
1009 let resolved = resolve_expr(&expr, &schema)?;
1010 Ok(resolved)
1011 }
1012
1013 fn try_decode_hex_literal(s: &str) -> Option<Vec<u8>> {
1019 let hex_bytes = s.as_bytes();
1020 let mut decoded_bytes = Vec::with_capacity(hex_bytes.len().div_ceil(2));
1021
1022 let start_idx = hex_bytes.len() % 2;
1023 if start_idx > 0 {
1024 decoded_bytes.push(Self::try_decode_hex_char(hex_bytes[0])?);
1026 }
1027
1028 for i in (start_idx..hex_bytes.len()).step_by(2) {
1029 let high = Self::try_decode_hex_char(hex_bytes[i])?;
1030 let low = Self::try_decode_hex_char(hex_bytes[i + 1])?;
1031 decoded_bytes.push((high << 4) | low);
1032 }
1033
1034 Some(decoded_bytes)
1035 }
1036
1037 const fn try_decode_hex_char(c: u8) -> Option<u8> {
1041 match c {
1042 b'A'..=b'F' => Some(c - b'A' + 10),
1043 b'a'..=b'f' => Some(c - b'a' + 10),
1044 b'0'..=b'9' => Some(c - b'0'),
1045 _ => None,
1046 }
1047 }
1048
1049 pub fn optimize_expr(&self, expr: Expr) -> Result<Expr> {
1051 let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?);
1052
1053 let expr = expr
1056 .transform_up(|expr| {
1057 let coerced = match &expr {
1058 Expr::ScalarFunction(ScalarFunction { func, args })
1059 if func.name() == "arrow_cast" =>
1060 {
1061 match args.as_slice() {
1062 [
1063 Expr::Literal(value, metadata),
1064 Expr::Literal(ScalarValue::Utf8(Some(data_type)), _),
1065 ] => data_type
1066 .parse::<ArrowDataType>()
1067 .ok()
1068 .filter(|data_type| matches!(data_type, ArrowDataType::Time32(_)))
1069 .and_then(|data_type| {
1070 if matches!(value, ScalarValue::Null) {
1071 ScalarValue::try_new_null(&data_type).ok()
1072 } else {
1073 safe_coerce_scalar(value, &data_type)
1074 }
1075 })
1076 .map(|value| Expr::Literal(value, metadata.clone())),
1077 _ => None,
1078 }
1079 }
1080 _ => None,
1081 };
1082
1083 Ok(match coerced {
1084 Some(coerced) => Transformed::yes(coerced),
1085 None => Transformed::no(expr),
1086 })
1087 })?
1088 .data;
1089
1090 let simplify_context = SimplifyContext::builder()
1093 .with_schema(df_schema.clone())
1094 .with_query_execution_start_time(Some(Utc::now()))
1095 .build();
1096 let simplifier =
1097 datafusion::optimizer::simplify_expressions::ExprSimplifier::new(simplify_context);
1098
1099 let expr = simplifier.coerce(expr, &df_schema)?;
1101
1102 let expr = normalize_zero_comparisons(expr, &|operand| simplifier.simplify(operand))?;
1110
1111 let expr = simplifier.simplify(expr)?;
1118 let expr = rewrite_signed_zero_comparisons(expr)?;
1119
1120 Ok(expr)
1121 }
1122
1123 pub fn create_physical_expr(&self, expr: &Expr) -> Result<Arc<dyn PhysicalExpr>> {
1125 let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?);
1126 Ok(datafusion::physical_expr::create_physical_expr(
1127 expr,
1128 df_schema.as_ref(),
1129 &Default::default(),
1130 )?)
1131 }
1132
1133 pub fn create_physical_expr_with_session(
1135 &self,
1136 expr: &Expr,
1137 session: &dyn Session,
1138 ) -> Result<Arc<dyn PhysicalExpr>> {
1139 let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
1140 Ok(session.create_physical_expr(expr.clone(), &df_schema)?)
1141 }
1142
1143 pub fn column_names_in_expr(expr: &Expr) -> Vec<String> {
1150 let mut visitor = ColumnCapturingVisitor {
1151 current_path: VecDeque::new(),
1152 columns: BTreeSet::new(),
1153 };
1154 expr.visit(&mut visitor).unwrap();
1155 visitor.columns.into_iter().collect()
1156 }
1157}
1158
1159struct ColumnCapturingVisitor {
1160 current_path: VecDeque<String>,
1162 columns: BTreeSet<String>,
1163}
1164
1165impl TreeNodeVisitor<'_> for ColumnCapturingVisitor {
1166 type Node = Expr;
1167
1168 fn f_down(&mut self, node: &Self::Node) -> DFResult<TreeNodeRecursion> {
1169 match node {
1170 Expr::Column(Column { name, .. }) => {
1171 let mut path = name.clone();
1175 for part in self.current_path.drain(..) {
1176 path.push('.');
1177 if part.contains('.') || part.contains('`') {
1179 let escaped = part.replace('`', "``");
1181 path.push('`');
1182 path.push_str(&escaped);
1183 path.push('`');
1184 } else {
1185 path.push_str(&part);
1186 }
1187 }
1188 self.columns.insert(path);
1189 self.current_path.clear();
1190 }
1191 Expr::ScalarFunction(udf) if udf.name() == GetFieldFunc::default().name() => {
1192 if let Some(name) = get_as_string_scalar_opt(&udf.args[1]) {
1193 self.current_path.push_front(name.to_string())
1194 } else {
1195 self.current_path.clear();
1196 }
1197 }
1198 _ => {
1199 self.current_path.clear();
1200 }
1201 }
1202
1203 Ok(TreeNodeRecursion::Continue)
1204 }
1205}
1206
1207#[cfg(test)]
1208mod tests {
1209
1210 use super::*;
1211
1212 use arrow::datatypes::Float64Type;
1213 use arrow_array::{
1214 ArrayRef, BooleanArray, Float32Array, Int32Array, Int64Array, RecordBatch, StringArray,
1215 StructArray, Time32SecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
1216 TimestampNanosecondArray, TimestampSecondArray, UInt64Array,
1217 };
1218 use arrow_schema::{DataType, Fields, Schema};
1219 use datafusion::{
1220 logical_expr::{Cast, col, lit},
1221 prelude::{array_element, get_field},
1222 };
1223 use datafusion_functions::core::expr_ext::FieldAccessor;
1224 use rstest::rstest;
1225
1226 #[test]
1227 fn test_parse_filter_simple() {
1228 let schema = Arc::new(Schema::new(vec![
1229 Field::new("i", DataType::Int32, false),
1230 Field::new("s", DataType::Utf8, true),
1231 Field::new(
1232 "st",
1233 DataType::Struct(Fields::from(vec![
1234 Field::new("x", DataType::Float32, false),
1235 Field::new("y", DataType::Float32, false),
1236 ])),
1237 true,
1238 ),
1239 ]));
1240
1241 let planner = Planner::new(schema.clone());
1242
1243 let expected = col("i")
1244 .gt(lit(3_i32))
1245 .and(col("st").field("x").lt_eq(lit(5.0_f32)))
1246 .and(
1247 col("s")
1248 .eq(lit("str-4"))
1249 .or(col("s").in_list(vec![lit("str-4"), lit("str-5")], false)),
1250 );
1251
1252 let expr = planner
1254 .parse_filter("i > 3 AND st.x <= 5.0 AND (s == 'str-4' OR s in ('str-4', 'str-5'))")
1255 .unwrap();
1256 assert_eq!(expr, expected);
1257
1258 let expr = planner
1260 .parse_filter("i > 3 AND st.x <= 5.0 AND (s = 'str-4' OR s in ('str-4', 'str-5'))")
1261 .unwrap();
1262
1263 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1264
1265 let batch = RecordBatch::try_new(
1266 schema,
1267 vec![
1268 Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef,
1269 Arc::new(StringArray::from_iter_values(
1270 (0..10).map(|v| format!("str-{}", v)),
1271 )),
1272 Arc::new(StructArray::from(vec![
1273 (
1274 Arc::new(Field::new("x", DataType::Float32, false)),
1275 Arc::new(Float32Array::from_iter_values((0..10).map(|v| v as f32)))
1276 as ArrayRef,
1277 ),
1278 (
1279 Arc::new(Field::new("y", DataType::Float32, false)),
1280 Arc::new(Float32Array::from_iter_values(
1281 (0..10).map(|v| (v * 10) as f32),
1282 )),
1283 ),
1284 ])),
1285 ],
1286 )
1287 .unwrap();
1288 let predicates = physical_expr.evaluate(&batch).unwrap();
1289 assert_eq!(
1290 predicates.into_array(0).unwrap().as_ref(),
1291 &BooleanArray::from(vec![
1292 false, false, false, false, true, true, false, false, false, false
1293 ])
1294 );
1295 }
1296
1297 #[rstest]
1298 #[case::not_equal_to_integer_expression(
1299 "flag != (id > 0)",
1300 [false, false, true]
1301 )]
1302 #[case::equal_to_float_expression("flag = (x > 1.0)", [false, false, false])]
1303 #[case::expression_on_left("(id > 0) != flag", [false, false, true])]
1304 #[case::literal_expression("flag != (1 > 0)", [false, true, false])]
1305 #[case::boolean_literal("flag != TRUE", [false, true, false])]
1306 fn test_parse_boolean_column_compared_to_expression(
1307 #[case] filter: &str,
1308 #[case] expected: [bool; 3],
1309 ) {
1310 let batch = arrow_array::record_batch!(
1311 ("flag", Boolean, [true, false, true]),
1312 ("id", Int32, [1, 0, -1]),
1313 ("x", Float32, [0.5, 2.0, 1.0])
1314 )
1315 .unwrap();
1316 let planner = Planner::new(batch.schema());
1317
1318 let expr = planner.parse_filter(filter).unwrap();
1319 let expr = planner.optimize_expr(expr).unwrap();
1320 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1321 let predicates = physical_expr.evaluate(&batch).unwrap();
1322
1323 assert_eq!(
1324 predicates.into_array(0).unwrap().as_ref(),
1325 &BooleanArray::from(expected.to_vec())
1326 );
1327 }
1328
1329 #[rstest]
1330 #[case::shift("u = (1 << 63)", [true, false, false])]
1331 #[case::string_concat("s = ('a' || 'b')", [true, false, false])]
1332 fn test_parse_comparison_to_same_type_literal_expression(
1333 #[case] filter: &str,
1334 #[case] expected: [bool; 3],
1335 ) {
1336 let batch = arrow_array::record_batch!(
1337 ("u", UInt64, [1_u64 << 63, 3, 0]),
1338 ("s", LargeUtf8, ["ab", "a", "b"])
1339 )
1340 .unwrap();
1341 let planner = Planner::new(batch.schema());
1342
1343 let expr = planner.parse_filter(filter).unwrap();
1344 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1345 let predicates = physical_expr.evaluate(&batch).unwrap();
1346
1347 assert_eq!(
1348 predicates.into_array(0).unwrap().as_ref(),
1349 &BooleanArray::from(expected.to_vec())
1350 );
1351 }
1352
1353 #[rstest]
1354 #[case::float_literal("f = n + 0.5")]
1355 #[case::float_literal_on_left("f = 0.5 + n")]
1356 #[case::float_expression_on_left("n + 0.5 = f")]
1357 #[case::wide_integer_literal("wide = small + 128")]
1358 #[case::wide_integer_literal_on_left("wide = 128 + small")]
1359 #[case::integer_expression_on_left("small + 128 = wide")]
1360 fn test_parse_comparison_to_mixed_type_expression(#[case] filter: &str) {
1361 let batch = arrow_array::record_batch!(
1362 ("f", Float64, [Some(1.5), Some(1.5), None, Some(1.5)]),
1363 ("n", Int64, [Some(1), Some(2), Some(1), None]),
1364 ("wide", Int64, [Some(129), Some(129), Some(129), None]),
1365 ("small", Int8, [Some(1), Some(2), None, Some(1)])
1366 )
1367 .unwrap();
1368 let planner = Planner::new(batch.schema());
1369
1370 let expr = planner.parse_filter(filter).unwrap();
1371 let expr = planner.optimize_expr(expr).unwrap();
1372 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1373 let predicates = physical_expr.evaluate(&batch).unwrap();
1374
1375 assert_eq!(
1376 predicates.into_array(batch.num_rows()).unwrap().as_ref(),
1377 &BooleanArray::from(vec![Some(true), Some(false), None, None])
1378 );
1379 }
1380
1381 #[test]
1382 fn test_parse_filter_uint64_literal_above_i64_max() {
1383 let value = u64::MAX - 1;
1384 let batch = arrow_array::record_batch!(("id", UInt64, [1, value])).unwrap();
1385 let planner = Planner::new(batch.schema());
1386
1387 let expr = planner.parse_filter(&format!("id = {value}")).unwrap();
1388 assert_eq!(expr, col("id").eq(lit(value)));
1389
1390 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1391 let predicates = physical_expr.evaluate(&batch).unwrap();
1392 assert_eq!(
1393 predicates.into_array(0).unwrap().as_ref(),
1394 &BooleanArray::from(vec![false, true])
1395 );
1396
1397 let expr = planner
1398 .parse_expr("arrow_cast(NULL, 'Time32(Second)')")
1399 .unwrap();
1400 let expr = planner.optimize_expr(expr).unwrap();
1401 assert_eq!(expr, Expr::Literal(ScalarValue::Time32Second(None), None));
1402 }
1403
1404 fn column_resolution_schema() -> SchemaRef {
1405 Arc::new(Schema::new(vec![
1406 Field::new("id", DataType::Int32, false),
1407 Field::new("order", DataType::Utf8, true),
1408 Field::new("a`b", DataType::Int32, false),
1409 Field::new(
1410 "parent",
1411 DataType::Struct(Fields::from(vec![Field::new(
1412 "child",
1413 DataType::Int32,
1414 false,
1415 )])),
1416 false,
1417 ),
1418 ]))
1419 }
1420
1421 fn parse_expr_via_sql(planner: &Planner, input: &str) -> Expr {
1424 let ast_expr = parse_sql_expr(input).unwrap();
1425 let expr = planner.parse_sql_expr(&ast_expr).unwrap();
1426 let schema = lance_core::datatypes::Schema::try_from(planner.schema.as_ref()).unwrap();
1427 resolve_expr(&expr, &schema).unwrap()
1428 }
1429
1430 #[rstest]
1431 #[case::bare("id", col("id"))]
1432 #[case::backtick_quoted("`id`", col("id"))]
1433 #[case::case_insensitive("ID", col("id"))]
1434 #[case::reserved_word("`order`", col("order"))]
1435 #[case::escaped_backtick("`a``b`", col("a`b"))]
1436 fn test_parse_expr_column_resolution(#[case] input: &str, #[case] expected: Expr) {
1437 let planner = Planner::new(column_resolution_schema());
1438
1439 let expr = planner.parse_expr(input).unwrap();
1440 assert_eq!(expr, expected);
1441 assert_eq!(expr, parse_expr_via_sql(&planner, input));
1443 }
1444
1445 #[rstest]
1447 #[case::unknown_column("`missing`")]
1448 #[case::quoted_wrong_case("`ID`")]
1449 #[case::unknown_nested_root("`missing`.`child`")]
1450 #[case::not_a_path("id + 1")]
1451 #[case::nested("`parent`.`child`")]
1452 #[case::nested_unquoted("parent.child")]
1453 #[case::nested_case_insensitive("`PARENT`.`child`")]
1454 #[case::nested_arithmetic("parent.child + 1")]
1455 #[case::nested_comparison("parent.child > 1")]
1456 #[case::nested_null_test("parent.child IS NULL")]
1457 #[case::nested_cast("parent.child::BIGINT")]
1458 #[case::quoted_nested_arithmetic("`parent`.`child` + 1")]
1459 fn test_parse_expr_column_resolution_falls_back_to_sql(#[case] input: &str) {
1460 let planner = Planner::new(column_resolution_schema());
1461
1462 let fast = planner.parse_expr(input);
1463 let sql = parse_sql_expr(input)
1464 .and_then(|ast| planner.parse_sql_expr(&ast))
1465 .and_then(|expr| {
1466 let schema = lance_core::datatypes::Schema::try_from(planner.schema.as_ref())?;
1467 resolve_expr(&expr, &schema)
1468 });
1469 assert_eq!(fast.is_ok(), sql.is_ok(), "{input}: {fast:?} vs {sql:?}");
1470 if let (Ok(fast), Ok(sql)) = (fast, sql) {
1471 assert_eq!(fast, sql);
1472 }
1473 }
1474
1475 #[test]
1478 fn test_parse_expr_literal_dotted_column_beats_nested_path() {
1479 let child = Arc::new(Field::new("child", DataType::Int32, false));
1480 let parent = StructArray::new(
1481 vec![child].into(),
1482 vec![Arc::new(Int32Array::from(vec![7, 10]))],
1483 None,
1484 );
1485 let batch = RecordBatch::try_from_iter([
1486 ("parent", Arc::new(parent) as ArrayRef),
1487 (
1488 "parent.child",
1489 Arc::new(Int32Array::from(vec![100, 200])) as ArrayRef,
1490 ),
1491 ])
1492 .unwrap();
1493 let planner = Planner::new(batch.schema());
1494
1495 let expr = planner.parse_expr("parent.child").unwrap();
1496 assert_eq!(expr, Expr::Column(Column::from_name("parent.child")));
1497 let expr = planner.optimize_expr(expr).unwrap();
1498 let physical = planner.create_physical_expr(&expr).unwrap();
1499 let actual = physical.evaluate(&batch).unwrap().into_array(2).unwrap();
1500 let expected: ArrayRef = Arc::new(Int32Array::from(vec![100, 200]));
1501 assert_eq!(actual.as_ref(), expected.as_ref());
1502 }
1503
1504 #[test]
1505 fn test_parse_expr_with_relations_uses_sql_path() {
1506 let planner = Planner::new(column_resolution_schema()).with_enable_relations(true);
1507
1508 let expr = planner.parse_expr("`t`.`id`").unwrap();
1510 assert_eq!(expr, Expr::Column(Column::new(Some("t"), "id")));
1511 }
1512
1513 #[test]
1514 fn test_parse_expr_sql_fallback() {
1515 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1516 let planner = Planner::new(schema);
1517
1518 let expr = planner.parse_expr("id + 1").unwrap();
1519 assert_eq!(expr, col("id") + lit(1_i32));
1520 }
1521
1522 #[test]
1523 fn test_parse_deep_logical_filter() {
1524 let planner = Planner::new(Arc::new(Schema::empty()));
1525
1526 for op in ["AND", "OR"] {
1527 let filter = std::iter::repeat_n("true", 1000)
1528 .collect::<Vec<_>>()
1529 .join(&format!(" {op} "));
1530
1531 let expr = planner.parse_filter(&filter).unwrap();
1532 let optimized = planner.optimize_expr(expr).unwrap();
1533
1534 assert_eq!(optimized, lit(true));
1535 }
1536 }
1537
1538 #[derive(Debug, Eq, PartialEq, Hash)]
1539 struct StrictFloat64Udf {
1540 signature: Signature,
1541 }
1542
1543 impl StrictFloat64Udf {
1544 fn new() -> Self {
1545 Self {
1546 signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
1547 }
1548 }
1549 }
1550
1551 impl ScalarUDFImpl for StrictFloat64Udf {
1552 fn name(&self) -> &str {
1553 "strict_float64"
1554 }
1555
1556 fn signature(&self) -> &Signature {
1557 &self.signature
1558 }
1559
1560 fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
1561 Ok(DataType::Float64)
1562 }
1563
1564 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
1565 let data_type = args.args[0].data_type();
1566 assert_eq!(
1567 data_type,
1568 DataType::Float64,
1569 "strict_float64 expected Float64, got {data_type}"
1570 );
1571 Ok(ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))))
1572 }
1573 }
1574
1575 #[test]
1576 fn test_coerce_before_simplify() {
1577 let planner = Planner::new(Arc::new(Schema::empty()));
1578 let strict_float64 = Arc::new(ScalarUDF::new_from_impl(StrictFloat64Udf::new()));
1579 let expr = Expr::ScalarFunction(ScalarFunction::new_udf(strict_float64, vec![lit(0_i64)]))
1580 .eq(lit(0.0_f64));
1581
1582 let optimized = planner.optimize_expr(expr).unwrap();
1583
1584 planner.create_physical_expr(&optimized).unwrap();
1585 }
1586
1587 #[test]
1588 fn test_nested_col_refs() {
1589 let schema = Arc::new(Schema::new(vec![
1590 Field::new("s0", DataType::Utf8, true),
1591 Field::new(
1592 "st",
1593 DataType::Struct(Fields::from(vec![
1594 Field::new("s1", DataType::Utf8, true),
1595 Field::new(
1596 "st",
1597 DataType::Struct(Fields::from(vec![Field::new(
1598 "s2",
1599 DataType::Utf8,
1600 true,
1601 )])),
1602 true,
1603 ),
1604 ])),
1605 true,
1606 ),
1607 ]));
1608
1609 let planner = Planner::new(schema);
1610
1611 fn assert_column_eq(planner: &Planner, expr: &str, expected: &Expr) {
1612 let expr = planner.parse_filter(&format!("{expr} = 'val'")).unwrap();
1613 assert!(matches!(
1614 expr,
1615 Expr::BinaryExpr(BinaryExpr {
1616 left: _,
1617 op: Operator::Eq,
1618 right: _
1619 })
1620 ));
1621 if let Expr::BinaryExpr(BinaryExpr { left, .. }) = expr {
1622 assert_eq!(left.as_ref(), expected);
1623 }
1624 }
1625
1626 let expected = Expr::Column(Column::new_unqualified("s0"));
1627 assert_column_eq(&planner, "s0", &expected);
1628 assert_column_eq(&planner, "`s0`", &expected);
1629
1630 let expected = Expr::ScalarFunction(ScalarFunction {
1631 func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1632 args: vec![
1633 Expr::Column(Column::new_unqualified("st")),
1634 Expr::Literal(ScalarValue::Utf8(Some("s1".to_string())), None),
1635 ],
1636 });
1637 assert_column_eq(&planner, "st.s1", &expected);
1638 assert_column_eq(&planner, "`st`.`s1`", &expected);
1639 assert_column_eq(&planner, "st.`s1`", &expected);
1640
1641 let expected = Expr::ScalarFunction(ScalarFunction {
1642 func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1643 args: vec![
1644 Expr::ScalarFunction(ScalarFunction {
1645 func: Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::default())),
1646 args: vec![
1647 Expr::Column(Column::new_unqualified("st")),
1648 Expr::Literal(ScalarValue::Utf8(Some("st".to_string())), None),
1649 ],
1650 }),
1651 Expr::Literal(ScalarValue::Utf8(Some("s2".to_string())), None),
1652 ],
1653 });
1654
1655 assert_column_eq(&planner, "st.st.s2", &expected);
1656 assert_column_eq(&planner, "`st`.`st`.`s2`", &expected);
1657 assert_column_eq(&planner, "st.st.`s2`", &expected);
1658 assert_column_eq(&planner, "st['st'][\"s2\"]", &expected);
1659 }
1660
1661 #[test]
1662 fn test_nested_list_refs() {
1663 let schema = Arc::new(Schema::new(vec![Field::new(
1664 "l",
1665 DataType::List(Arc::new(Field::new(
1666 "item",
1667 DataType::Struct(Fields::from(vec![Field::new("f1", DataType::Utf8, true)])),
1668 true,
1669 ))),
1670 true,
1671 )]));
1672
1673 let planner = Planner::new(schema);
1674
1675 let expected = array_element(col("l"), lit(0_i64));
1676 let expr = planner.parse_expr("l[0]").unwrap();
1677 assert_eq!(expr, expected);
1678
1679 let expected = get_field(array_element(col("l"), lit(0_i64)), "f1");
1680 let expr = planner.parse_expr("l[0]['f1']").unwrap();
1681 assert_eq!(expr, expected);
1682
1683 }
1688
1689 #[test]
1690 fn test_negative_expressions() {
1691 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
1692
1693 let planner = Planner::new(schema.clone());
1694
1695 let expected = col("x")
1696 .gt(lit(-3_i64))
1697 .and(col("x").lt(-(lit(-5_i64) + lit(3_i64))));
1698
1699 let expr = planner.parse_filter("x > -3 AND x < -(-5 + 3)").unwrap();
1700
1701 assert_eq!(expr, expected);
1702
1703 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1704
1705 let batch = RecordBatch::try_new(
1706 schema,
1707 vec![Arc::new(Int64Array::from_iter_values(-5..5)) as ArrayRef],
1708 )
1709 .unwrap();
1710 let predicates = physical_expr.evaluate(&batch).unwrap();
1711 assert_eq!(
1712 predicates.into_array(0).unwrap().as_ref(),
1713 &BooleanArray::from(vec![
1714 false, false, false, true, true, true, true, false, false, false
1715 ])
1716 );
1717 }
1718
1719 #[rstest]
1720 #[case::right("value >> 32", Operator::BitwiseShiftRight, vec![0, 1, 3])]
1721 #[case::left(
1722 "value << 1",
1723 Operator::BitwiseShiftLeft,
1724 vec![0, 2_u64 << 32, ((3_u64 << 32) + 7) << 1]
1725 )]
1726 fn test_bitwise_shift_expressions(
1727 #[case] sql: &str,
1728 #[case] expected_op: Operator,
1729 #[case] expected: Vec<u64>,
1730 ) {
1731 let input = vec![0, 1_u64 << 32, (3_u64 << 32) + 7];
1732 let batch =
1733 RecordBatch::try_from_iter([("value", Arc::new(UInt64Array::from(input)) as ArrayRef)])
1734 .unwrap();
1735 let planner = Planner::new(batch.schema());
1736
1737 let expr = planner.parse_expr(sql).unwrap();
1738 let Expr::BinaryExpr(binary_expr) = &expr else {
1739 panic!("expected binary expression for {sql}, got {expr}");
1740 };
1741 assert_eq!(binary_expr.op, expected_op);
1742
1743 let expr = planner.optimize_expr(expr).unwrap();
1744 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1745 let values = physical_expr
1746 .evaluate(&batch)
1747 .unwrap()
1748 .into_array(batch.num_rows())
1749 .unwrap();
1750 assert_eq!(values.as_ref(), &UInt64Array::from(expected));
1751 }
1752
1753 #[test]
1754 fn test_negative_array_expressions() {
1755 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
1756
1757 let planner = Planner::new(schema);
1758
1759 let expected = Expr::Literal(
1760 ScalarValue::List(Arc::new(
1761 ListArray::from_iter_primitive::<Float64Type, _, _>(vec![Some(
1762 [-1_f64, -2.0, -3.0, -4.0, -5.0].map(Some),
1763 )]),
1764 )),
1765 None,
1766 );
1767
1768 let expr = planner
1769 .parse_expr("[-1.0, -2.0, -3.0, -4.0, -5.0]")
1770 .unwrap();
1771
1772 assert_eq!(expr, expected);
1773 }
1774
1775 #[test]
1776 fn test_sql_like() {
1777 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1778
1779 let planner = Planner::new(schema.clone());
1780
1781 let expected = col("s").like(lit("str-4"));
1782 let expr = planner.parse_filter("s LIKE 'str-4'").unwrap();
1784 assert_eq!(expr, expected);
1785 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1786
1787 let batch = RecordBatch::try_new(
1788 schema,
1789 vec![Arc::new(StringArray::from_iter_values(
1790 (0..10).map(|v| format!("str-{}", v)),
1791 ))],
1792 )
1793 .unwrap();
1794 let predicates = physical_expr.evaluate(&batch).unwrap();
1795 assert_eq!(
1796 predicates.into_array(0).unwrap().as_ref(),
1797 &BooleanArray::from(vec![
1798 false, false, false, false, true, false, false, false, false, false
1799 ])
1800 );
1801 }
1802
1803 #[test]
1804 fn test_not_like() {
1805 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1806
1807 let planner = Planner::new(schema.clone());
1808
1809 let expected = col("s").not_like(lit("str-4"));
1810 let expr = planner.parse_filter("s NOT LIKE 'str-4'").unwrap();
1812 assert_eq!(expr, expected);
1813 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1814
1815 let batch = RecordBatch::try_new(
1816 schema,
1817 vec![Arc::new(StringArray::from_iter_values(
1818 (0..10).map(|v| format!("str-{}", v)),
1819 ))],
1820 )
1821 .unwrap();
1822 let predicates = physical_expr.evaluate(&batch).unwrap();
1823 assert_eq!(
1824 predicates.into_array(0).unwrap().as_ref(),
1825 &BooleanArray::from(vec![
1826 true, true, true, true, false, true, true, true, true, true
1827 ])
1828 );
1829 }
1830
1831 #[test]
1832 fn test_like_escape_char() {
1833 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1834 let planner = Planner::new(schema);
1835
1836 for filter in ["s LIKE 'a!%' ESCAPE '!'", "s ILIKE 'a!%' ESCAPE '!'"] {
1838 match planner.parse_filter(filter).unwrap() {
1839 Expr::Like(like) => assert_eq!(like.escape_char, Some('!'), "{filter}"),
1840 other => panic!("expected a LIKE expression for `{filter}`, got {other:?}"),
1841 }
1842 }
1843
1844 for filter in [
1847 "s LIKE 'x' ESCAPE ''",
1848 "s LIKE 'x' ESCAPE 'ab'",
1849 "s ILIKE 'x' ESCAPE ''",
1850 "s ILIKE 'x' ESCAPE 'ab'",
1851 ] {
1852 let err = planner.parse_filter(filter).unwrap_err();
1853 assert!(
1854 err.to_string()
1855 .contains("Invalid escape character in LIKE expression"),
1856 "unexpected error for `{filter}`: {err}"
1857 );
1858 }
1859 }
1860
1861 #[test]
1862 fn test_sql_is_in() {
1863 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1864
1865 let planner = Planner::new(schema.clone());
1866
1867 let expected = col("s").in_list(vec![lit("str-4"), lit("str-5")], false);
1868 let expr = planner.parse_filter("s IN ('str-4', 'str-5')").unwrap();
1870 assert_eq!(expr, expected);
1871 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1872
1873 let batch = RecordBatch::try_new(
1874 schema,
1875 vec![Arc::new(StringArray::from_iter_values(
1876 (0..10).map(|v| format!("str-{}", v)),
1877 ))],
1878 )
1879 .unwrap();
1880 let predicates = physical_expr.evaluate(&batch).unwrap();
1881 assert_eq!(
1882 predicates.into_array(0).unwrap().as_ref(),
1883 &BooleanArray::from(vec![
1884 false, false, false, false, true, true, false, false, false, false
1885 ])
1886 );
1887 }
1888
1889 #[test]
1890 fn test_sql_is_null() {
1891 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
1892
1893 let planner = Planner::new(schema.clone());
1894
1895 let expected = col("s").is_null();
1896 let expr = planner.parse_filter("s IS NULL").unwrap();
1897 assert_eq!(expr, expected);
1898 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1899
1900 let batch = RecordBatch::try_new(
1901 schema,
1902 vec![Arc::new(StringArray::from_iter((0..10).map(|v| {
1903 if v % 3 == 0 {
1904 Some(format!("str-{}", v))
1905 } else {
1906 None
1907 }
1908 })))],
1909 )
1910 .unwrap();
1911 let predicates = physical_expr.evaluate(&batch).unwrap();
1912 assert_eq!(
1913 predicates.into_array(0).unwrap().as_ref(),
1914 &BooleanArray::from(vec![
1915 false, true, true, false, true, true, false, true, true, false
1916 ])
1917 );
1918
1919 let expr = planner.parse_filter("s IS NOT NULL").unwrap();
1920 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1921 let predicates = physical_expr.evaluate(&batch).unwrap();
1922 assert_eq!(
1923 predicates.into_array(0).unwrap().as_ref(),
1924 &BooleanArray::from(vec![
1925 true, false, false, true, false, false, true, false, false, true,
1926 ])
1927 );
1928 }
1929
1930 #[rstest]
1931 #[case::distinct("a IS DISTINCT FROM b", [false, true, true, false])]
1932 #[case::not_distinct("a IS NOT DISTINCT FROM b", [true, false, false, true])]
1933 fn test_sql_is_distinct_from(#[case] filter: &str, #[case] expected: [bool; 4]) {
1934 let batch = arrow_array::record_batch!(
1935 ("a", Int32, [Some(1), None, Some(2), None]),
1936 ("b", Int32, [Some(1), Some(2), None, None])
1937 )
1938 .unwrap();
1939 let planner = Planner::new(batch.schema());
1940 let expr = planner.optimize_expr(planner.parse_filter(filter).unwrap());
1941 let physical = planner.create_physical_expr(&expr.unwrap()).unwrap();
1942 let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap();
1943 assert_eq!(result.as_ref(), &BooleanArray::from(expected.to_vec()));
1944 }
1945
1946 #[test]
1947 fn test_sql_invert() {
1948 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Boolean, true)]));
1949
1950 let planner = Planner::new(schema.clone());
1951
1952 let expr = planner.parse_filter("NOT s").unwrap();
1953 let physical_expr = planner.create_physical_expr(&expr).unwrap();
1954
1955 let batch = RecordBatch::try_new(
1956 schema,
1957 vec![Arc::new(BooleanArray::from_iter(
1958 (0..10).map(|v| Some(v % 3 == 0)),
1959 ))],
1960 )
1961 .unwrap();
1962 let predicates = physical_expr.evaluate(&batch).unwrap();
1963 assert_eq!(
1964 predicates.into_array(0).unwrap().as_ref(),
1965 &BooleanArray::from(vec![
1966 false, true, true, false, true, true, false, true, true, false
1967 ])
1968 );
1969 }
1970
1971 #[test]
1972 fn test_sql_cast() {
1973 let cases = &[
1974 (
1975 "x = cast('2021-01-01 00:00:00' as timestamp)",
1976 ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1977 ),
1978 (
1979 "x = cast('2021-01-01 00:00:00' as timestamp(0))",
1980 ArrowDataType::Timestamp(TimeUnit::Second, None),
1981 ),
1982 (
1983 "x = cast('2021-01-01 00:00:00.123' as timestamp(9))",
1984 ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1985 ),
1986 (
1987 "x = cast('2021-01-01 00:00:00.123' as datetime(9))",
1988 ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1989 ),
1990 ("x = cast('2021-01-01' as date)", ArrowDataType::Date32),
1991 (
1992 "x = cast('1.238' as decimal(9,3))",
1993 ArrowDataType::Decimal128(9, 3),
1994 ),
1995 ("x = cast(1 as float)", ArrowDataType::Float32),
1996 ("x = cast(1 as double)", ArrowDataType::Float64),
1997 ("x = cast(1 as tinyint)", ArrowDataType::Int8),
1998 ("x = cast(1 as smallint)", ArrowDataType::Int16),
1999 ("x = cast(1 as int)", ArrowDataType::Int32),
2000 ("x = cast(1 as integer)", ArrowDataType::Int32),
2001 ("x = cast(1 as bigint)", ArrowDataType::Int64),
2002 ("x = cast(1 as tinyint unsigned)", ArrowDataType::UInt8),
2003 ("x = cast(1 as smallint unsigned)", ArrowDataType::UInt16),
2004 ("x = cast(1 as int unsigned)", ArrowDataType::UInt32),
2005 ("x = cast(1 as integer unsigned)", ArrowDataType::UInt32),
2006 ("x = cast(1 as bigint unsigned)", ArrowDataType::UInt64),
2007 ("x = cast(1 as boolean)", ArrowDataType::Boolean),
2008 ("x = cast(1 as string)", ArrowDataType::Utf8),
2009 ];
2010
2011 for (sql, expected_data_type) in cases {
2012 let schema = Arc::new(Schema::new(vec![Field::new(
2013 "x",
2014 expected_data_type.clone(),
2015 true,
2016 )]));
2017 let planner = Planner::new(schema.clone());
2018 let expr = planner.parse_filter(sql).unwrap();
2019
2020 let expected_value_str = sql
2022 .split("cast(")
2023 .nth(1)
2024 .unwrap()
2025 .split(" as")
2026 .next()
2027 .unwrap();
2028 let expected_value_str = expected_value_str.trim_matches('\'');
2030
2031 match expr {
2032 Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
2033 Expr::Cast(Cast { expr, field }) => {
2034 match expr.as_ref() {
2035 Expr::Literal(ScalarValue::Utf8(Some(value_str)), _) => {
2036 assert_eq!(value_str, expected_value_str);
2037 }
2038 Expr::Literal(ScalarValue::Int64(Some(value)), _) => {
2039 assert_eq!(*value, 1);
2040 }
2041 _ => panic!("Expected cast to be applied to literal"),
2042 }
2043 assert_eq!(field.data_type(), expected_data_type);
2044 }
2045 _ => panic!("Expected right to be a cast"),
2046 },
2047 _ => panic!("Expected binary expression"),
2048 }
2049 }
2050 }
2051
2052 #[test]
2053 fn test_arrow_cast_int_literal_to_time32() {
2054 let batch = RecordBatch::try_from_iter([(
2055 "v",
2056 Arc::new(Time32SecondArray::from(vec![3725, 3726])) as ArrayRef,
2057 )])
2058 .unwrap();
2059 let planner = Planner::new(batch.schema());
2060
2061 let expr = planner
2062 .parse_filter("v = arrow_cast(3726, 'Time32(Second)')")
2063 .unwrap();
2064 let expr = planner.optimize_expr(expr).unwrap();
2065 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2066 let predicates = physical_expr.evaluate(&batch).unwrap();
2067
2068 assert_eq!(
2069 predicates.into_array(0).unwrap().as_ref(),
2070 &BooleanArray::from(vec![false, true])
2071 );
2072 }
2073
2074 #[test]
2075 fn test_sql_literals() {
2076 let cases = &[
2077 (
2078 "x = timestamp '2021-01-01 00:00:00'",
2079 ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
2080 ),
2081 (
2082 "x = timestamp(0) '2021-01-01 00:00:00'",
2083 ArrowDataType::Timestamp(TimeUnit::Second, None),
2084 ),
2085 (
2086 "x = timestamp(9) '2021-01-01 00:00:00.123'",
2087 ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
2088 ),
2089 ("x = date '2021-01-01'", ArrowDataType::Date32),
2090 ("x = decimal(9,3) '1.238'", ArrowDataType::Decimal128(9, 3)),
2091 ];
2092
2093 for (sql, expected_data_type) in cases {
2094 let schema = Arc::new(Schema::new(vec![Field::new(
2095 "x",
2096 expected_data_type.clone(),
2097 true,
2098 )]));
2099 let planner = Planner::new(schema.clone());
2100 let expr = planner.parse_filter(sql).unwrap();
2101
2102 let expected_value_str = sql.split('\'').nth(1).unwrap();
2103
2104 match expr {
2105 Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
2106 Expr::Cast(Cast { expr, field }) => {
2107 match expr.as_ref() {
2108 Expr::Literal(ScalarValue::Utf8(Some(value_str)), _) => {
2109 assert_eq!(value_str, expected_value_str);
2110 }
2111 _ => panic!("Expected cast to be applied to literal"),
2112 }
2113 assert_eq!(field.data_type(), expected_data_type);
2114 }
2115 _ => panic!("Expected right to be a cast"),
2116 },
2117 _ => panic!("Expected binary expression"),
2118 }
2119 }
2120 }
2121
2122 #[test]
2123 fn test_sql_array_literals() {
2124 let cases = [
2125 (
2126 "x = [1, 2, 3]",
2127 ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int64, true))),
2128 ),
2129 (
2130 "x = [1, 2, 3]",
2131 ArrowDataType::FixedSizeList(
2132 Arc::new(Field::new("item", ArrowDataType::Int64, true)),
2133 3,
2134 ),
2135 ),
2136 ];
2137
2138 for (sql, expected_data_type) in cases {
2139 let schema = Arc::new(Schema::new(vec![Field::new(
2140 "x",
2141 expected_data_type.clone(),
2142 true,
2143 )]));
2144 let planner = Planner::new(schema.clone());
2145 let expr = planner.parse_filter(sql).unwrap();
2146 let expr = planner.optimize_expr(expr).unwrap();
2147
2148 match expr {
2149 Expr::BinaryExpr(BinaryExpr { right, .. }) => match right.as_ref() {
2150 Expr::Literal(value, _) => {
2151 assert_eq!(&value.data_type(), &expected_data_type);
2152 }
2153 _ => panic!("Expected right to be a literal"),
2154 },
2155 _ => panic!("Expected binary expression"),
2156 }
2157 }
2158 }
2159
2160 #[test]
2161 fn test_sql_between() {
2162 use arrow_array::{Float64Array, Int32Array, TimestampMicrosecondArray};
2163 use arrow_schema::{DataType, Field, Schema, TimeUnit};
2164 use std::sync::Arc;
2165
2166 let schema = Arc::new(Schema::new(vec![
2167 Field::new("x", DataType::Int32, false),
2168 Field::new("y", DataType::Float64, false),
2169 Field::new(
2170 "ts",
2171 DataType::Timestamp(TimeUnit::Microsecond, None),
2172 false,
2173 ),
2174 ]));
2175
2176 let planner = Planner::new(schema.clone());
2177
2178 let expr = planner
2180 .parse_filter("x BETWEEN CAST(3 AS INT) AND CAST(7 AS INT)")
2181 .unwrap();
2182 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2183
2184 let base_ts = 1704067200000000_i64; let ts_array = TimestampMicrosecondArray::from_iter_values(
2188 (0..10).map(|i| base_ts + i * 1_000_000), );
2190
2191 let batch = RecordBatch::try_new(
2192 schema,
2193 vec![
2194 Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef,
2195 Arc::new(Float64Array::from_iter_values((0..10).map(|v| v as f64))),
2196 Arc::new(ts_array),
2197 ],
2198 )
2199 .unwrap();
2200
2201 let predicates = physical_expr.evaluate(&batch).unwrap();
2202 assert_eq!(
2203 predicates.into_array(0).unwrap().as_ref(),
2204 &BooleanArray::from(vec![
2205 false, false, false, true, true, true, true, true, false, false
2206 ])
2207 );
2208
2209 let expr = planner
2211 .parse_filter("x NOT BETWEEN CAST(3 AS INT) AND CAST(7 AS INT)")
2212 .unwrap();
2213 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2214
2215 let predicates = physical_expr.evaluate(&batch).unwrap();
2216 assert_eq!(
2217 predicates.into_array(0).unwrap().as_ref(),
2218 &BooleanArray::from(vec![
2219 true, true, true, false, false, false, false, false, true, true
2220 ])
2221 );
2222
2223 let expr = planner.parse_filter("y BETWEEN 2.5 AND 6.5").unwrap();
2225 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2226
2227 let predicates = physical_expr.evaluate(&batch).unwrap();
2228 assert_eq!(
2229 predicates.into_array(0).unwrap().as_ref(),
2230 &BooleanArray::from(vec![
2231 false, false, false, true, true, true, true, false, false, false
2232 ])
2233 );
2234
2235 let expr = planner
2237 .parse_filter(
2238 "ts BETWEEN timestamp '2024-01-01 00:00:03' AND timestamp '2024-01-01 00:00:07'",
2239 )
2240 .unwrap();
2241 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2242
2243 let predicates = physical_expr.evaluate(&batch).unwrap();
2244 assert_eq!(
2245 predicates.into_array(0).unwrap().as_ref(),
2246 &BooleanArray::from(vec![
2247 false, false, false, true, true, true, true, true, false, false
2248 ])
2249 );
2250 }
2251
2252 #[test]
2253 fn test_sql_comparison() {
2254 let batch: Vec<(&str, ArrayRef)> = vec![
2256 (
2257 "timestamp_s",
2258 Arc::new(TimestampSecondArray::from_iter_values(0..10)),
2259 ),
2260 (
2261 "timestamp_ms",
2262 Arc::new(TimestampMillisecondArray::from_iter_values(0..10)),
2263 ),
2264 (
2265 "timestamp_us",
2266 Arc::new(TimestampMicrosecondArray::from_iter_values(0..10)),
2267 ),
2268 (
2269 "timestamp_ns",
2270 Arc::new(TimestampNanosecondArray::from_iter_values(4995..5005)),
2271 ),
2272 ];
2273 let batch = RecordBatch::try_from_iter(batch).unwrap();
2274
2275 let planner = Planner::new(batch.schema());
2276
2277 let expressions = &[
2279 "timestamp_s >= TIMESTAMP '1970-01-01 00:00:05'",
2280 "timestamp_ms >= TIMESTAMP '1970-01-01 00:00:00.005'",
2281 "timestamp_us >= TIMESTAMP '1970-01-01 00:00:00.000005'",
2282 "timestamp_ns >= TIMESTAMP '1970-01-01 00:00:00.000005'",
2283 ];
2284
2285 let expected: ArrayRef = Arc::new(BooleanArray::from_iter(
2286 std::iter::repeat_n(Some(false), 5).chain(std::iter::repeat_n(Some(true), 5)),
2287 ));
2288 for expression in expressions {
2289 let logical_expr = planner.parse_filter(expression).unwrap();
2291 let logical_expr = planner.optimize_expr(logical_expr).unwrap();
2292 let physical_expr = planner.create_physical_expr(&logical_expr).unwrap();
2293
2294 let result = physical_expr.evaluate(&batch).unwrap();
2296 let result = result.into_array(batch.num_rows()).unwrap();
2297 assert_eq!(&expected, &result, "unexpected result for {}", expression);
2298 }
2299 }
2300
2301 #[test]
2302 fn test_columns_in_expr() {
2303 let expr = col("s0").gt(lit("value")).and(
2304 col("st")
2305 .field("st")
2306 .field("s2")
2307 .eq(lit("value"))
2308 .or(col("st")
2309 .field("s1")
2310 .in_list(vec![lit("value 1"), lit("value 2")], false)),
2311 );
2312
2313 let columns = Planner::column_names_in_expr(&expr);
2314 assert_eq!(columns, vec!["s0", "st.s1", "st.st.s2"]);
2315 }
2316
2317 #[test]
2318 fn test_parse_binary_expr() {
2319 let bin_str = "x'616263'";
2320
2321 let schema = Arc::new(Schema::new(vec![Field::new(
2322 "binary",
2323 DataType::Binary,
2324 true,
2325 )]));
2326 let planner = Planner::new(schema);
2327 let expr = planner.parse_expr(bin_str).unwrap();
2328 assert_eq!(
2329 expr,
2330 Expr::Literal(ScalarValue::Binary(Some(vec![b'a', b'b', b'c'])), None)
2331 );
2332 }
2333
2334 #[test]
2335 fn test_lance_context_provider_expr_planners() {
2336 let ctx_provider = LanceContextProvider::default();
2337 assert!(!ctx_provider.get_expr_planners().is_empty());
2338 }
2339
2340 #[test]
2341 fn test_lance_context_provider_resolves_session_functions() {
2342 let session = crate::exec::new_session_context(&LanceExecutionOptions::default());
2343 let ctx_provider = LanceContextProvider::new(session.clone());
2344 assert!(ctx_provider.get_function_meta("contains_tokens").is_some());
2345 assert!(ctx_provider.get_function_meta("lower").is_some());
2346 assert!(ctx_provider.get_aggregate_meta("sum").is_some());
2347 assert!(ctx_provider.get_window_meta("row_number").is_some());
2348 assert!(ctx_provider.get_function_meta("registered_later").is_none());
2349
2350 session.register_udf(datafusion::logical_expr::create_udf(
2353 "registered_later",
2354 vec![DataType::Utf8],
2355 DataType::Utf8,
2356 Volatility::Immutable,
2357 Arc::new(|args: &[ColumnarValue]| Ok(args[0].clone())),
2358 ));
2359 assert!(ctx_provider.get_function_meta("registered_later").is_some());
2360 assert!(
2361 ctx_provider
2362 .udf_names()
2363 .contains(&"registered_later".to_string())
2364 );
2365 }
2366
2367 #[test]
2368 fn test_regexp_match_and_non_empty_captions() {
2369 let schema = Arc::new(Schema::new(vec![
2372 Field::new("keywords", DataType::Utf8, true),
2373 Field::new("natural_caption", DataType::Utf8, true),
2374 Field::new("poetic_caption", DataType::Utf8, true),
2375 ]));
2376
2377 let planner = Planner::new(schema.clone());
2378
2379 let expr = planner
2380 .parse_filter(
2381 "regexp_match(keywords, 'Liberty|revolution') AND \
2382 (natural_caption IS NOT NULL AND natural_caption <> '' AND \
2383 poetic_caption IS NOT NULL AND poetic_caption <> '')",
2384 )
2385 .unwrap();
2386
2387 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2388
2389 let batch = RecordBatch::try_new(
2390 schema,
2391 vec![
2392 Arc::new(StringArray::from(vec![
2393 Some("Liberty for all"),
2394 Some("peace"),
2395 Some("revolution now"),
2396 Some("Liberty"),
2397 Some("revolutionary"),
2398 Some("none"),
2399 ])) as ArrayRef,
2400 Arc::new(StringArray::from(vec![
2401 Some("a"),
2402 Some("b"),
2403 None,
2404 Some(""),
2405 Some("c"),
2406 Some("d"),
2407 ])) as ArrayRef,
2408 Arc::new(StringArray::from(vec![
2409 Some("x"),
2410 Some(""),
2411 Some("y"),
2412 Some("z"),
2413 None,
2414 Some("w"),
2415 ])) as ArrayRef,
2416 ],
2417 )
2418 .unwrap();
2419
2420 let result = physical_expr.evaluate(&batch).unwrap();
2421 assert_eq!(
2422 result.into_array(0).unwrap().as_ref(),
2423 &BooleanArray::from(vec![true, false, false, false, false, false])
2424 );
2425 }
2426
2427 #[rstest]
2428 #[case::bare("regexp_match(name, 'e[12]')", [false, true, true, false, false, false])]
2429 #[case::is_not_null(
2430 "regexp_match(name, 'e[12]') IS NOT NULL",
2431 [false, true, true, false, false, false]
2432 )]
2433 #[case::is_null(
2434 "regexp_match(name, 'e[12]') IS NULL",
2435 [true, false, false, true, true, true]
2436 )]
2437 #[case::not_bare(
2438 "NOT regexp_match(name, 'e[12]')",
2439 [true, false, false, true, true, true]
2440 )]
2441 #[case::and_bare(
2442 "regexp_match(name, 'e[12]') AND name <> 'name2'",
2443 [false, true, false, false, false, false]
2444 )]
2445 #[case::or_bare(
2446 "regexp_match(name, 'e[12]') OR name IS NULL",
2447 [false, true, true, true, false, false]
2448 )]
2449 #[case::not_is_not_null(
2450 "NOT (regexp_match(name, 'e[12]') IS NOT NULL)",
2451 [true, false, false, true, true, true]
2452 )]
2453 #[case::and_is_null(
2454 "regexp_match(name, 'e[12]') IS NULL AND name IS NOT NULL",
2455 [true, false, false, false, true, true]
2456 )]
2457 #[case::or_is_not_null(
2458 "regexp_match(name, 'e[12]') IS NOT NULL OR name IS NULL",
2459 [false, true, true, true, false, false]
2460 )]
2461 fn test_regexp_match_filter_coercion(#[case] filter: &str, #[case] expected: [bool; 6]) {
2462 let batch = arrow_array::record_batch!((
2463 "name",
2464 Utf8,
2465 [
2466 Some("name0"),
2467 Some("name1"),
2468 Some("name2"),
2469 None,
2470 Some("name4"),
2471 Some("name5")
2472 ]
2473 ))
2474 .unwrap();
2475 let planner = Planner::new(batch.schema());
2476 let expr = planner.parse_filter(filter).unwrap();
2477 let physical_expr = planner.create_physical_expr(&expr).unwrap();
2478 let result = physical_expr.evaluate(&batch).unwrap();
2479
2480 assert_eq!(
2481 result.into_array(batch.num_rows()).unwrap().as_ref(),
2482 &BooleanArray::from(expected.to_vec())
2483 );
2484 }
2485
2486 #[rstest]
2487 #[case::is_not_null("regexp_match(name, 'e[12]') IS NOT NULL")]
2488 #[case::is_null("regexp_match(name, 'e[12]') IS NULL")]
2489 #[case::comparison("regexp_match(name, 'e[12]') = regexp_match(name, 'e[12]')")]
2490 fn test_regexp_match_preserves_value_contexts(#[case] filter: &str) {
2491 let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Utf8, true)]));
2492 let planner = Planner::new(schema);
2493
2494 assert_eq!(
2495 planner.parse_filter(filter).unwrap(),
2496 planner.parse_expr(filter).unwrap()
2497 );
2498 }
2499
2500 #[test]
2501 fn test_jsonb_literals() {
2502 let schema = Arc::new(Schema::new(vec![Field::new(
2503 "j",
2504 DataType::LargeBinary,
2505 true,
2506 )]));
2507 let planner = Planner::new(schema);
2508
2509 let cases = [
2510 ("jsonb '{\"key\": \"value\"}'", r#"{"key":"value"}"#),
2511 ("cast('{\"a\": 1}' as jsonb)", r#"{"a":1}"#),
2512 ("'{\"a\": 1}'::jsonb", r#"{"a":1}"#),
2513 ];
2514 for (sql, expected) in cases {
2515 let ast = parse_sql_expr(sql).unwrap();
2516 let expr = planner.parse_sql_expr(&ast).unwrap();
2517 match expr {
2518 Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), _) => {
2519 assert_eq!(
2520 lance_arrow::json::decode_json(&bytes),
2521 expected,
2522 "failed for: {sql}"
2523 );
2524 }
2525 other => panic!("Expected LargeBinary literal for '{sql}', got: {other:?}"),
2526 }
2527 }
2528 }
2529
2530 #[test]
2531 fn test_jsonb_literal_errors() {
2532 let schema = Arc::new(Schema::new(vec![Field::new(
2533 "j",
2534 DataType::LargeBinary,
2535 true,
2536 )]));
2537 let planner = Planner::new(schema);
2538
2539 let ast = parse_sql_expr("jsonb 'not valid json'").unwrap();
2541 let err = planner.parse_sql_expr(&ast).unwrap_err();
2542 assert!(
2543 err.to_string().contains("Failed to encode JSONB"),
2544 "expected JSONB encoding error, got: {err}"
2545 );
2546
2547 let ast = parse_sql_expr("cast(j as jsonb)").unwrap();
2549 let err = planner.parse_sql_expr(&ast).unwrap_err();
2550 assert!(
2551 err.to_string()
2552 .contains("CAST to JSONB only supports string literals"),
2553 "got: {err}"
2554 );
2555 }
2556}