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;
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::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 "_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::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 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 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 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 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 pub fn parse_filter(&self, filter: &str) -> Result<Expr> {
952 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 pub fn parse_expr(&self, expr: &str) -> Result<Expr> {
968 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 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 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 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 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 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 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 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 let expr = simplifier.coerce(expr, &df_schema)?;
1071
1072 let expr = normalize_zero_comparisons(expr, &|operand| simplifier.simplify(operand))?;
1080
1081 let expr = simplifier.simplify(expr)?;
1088 let expr = rewrite_signed_zero_comparisons(expr)?;
1089
1090 Ok(expr)
1091 }
1092
1093 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 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 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_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 let mut path = name.clone();
1145 for part in self.current_path.drain(..) {
1146 path.push('.');
1147 if part.contains('.') || part.contains('`') {
1149 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 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 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 }
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 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 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 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 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 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 let expected_value_str = sql
1774 .split("cast(")
1775 .nth(1)
1776 .unwrap()
1777 .split(" as")
1778 .next()
1779 .unwrap();
1780 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 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 let base_ts = 1704067200000000_i64; let ts_array = TimestampMicrosecondArray::from_iter_values(
1940 (0..10).map(|i| base_ts + i * 1_000_000), );
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 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 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 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 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 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 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 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 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 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 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 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 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 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}