1use std::sync::Arc;
5
6use arrow_schema::DataType;
7use arrow_schema::Field;
8use arrow_schema::Schema;
9use datafusion_common::Result as DFResult;
10use datafusion_common::ScalarValue;
11use datafusion_common::exec_datafusion_err;
12use datafusion_common::tree_node::TreeNode;
13use datafusion_common::tree_node::TreeNodeRecursion;
14use datafusion_expr::Operator as DFOperator;
15use datafusion_functions::core::getfield::GetFieldFunc;
16use datafusion_functions::string::octet_length::OctetLengthFunc;
17use datafusion_functions_nested::length::ArrayLength;
18use datafusion_physical_expr::PhysicalExpr;
19use datafusion_physical_expr::ScalarFunctionExpr;
20use datafusion_physical_expr::projection::ProjectionExpr;
21use datafusion_physical_expr::projection::ProjectionExprs;
22use datafusion_physical_expr::utils::collect_columns;
23use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr;
24use datafusion_physical_plan::expressions as df_expr;
25use itertools::Itertools;
26use vortex::VortexSessionDefault;
27use vortex::array::arrow::ArrowSessionExt;
28use vortex::dtype::Nullability;
29use vortex::expr::Expression;
30use vortex::expr::and_collect;
31use vortex::expr::byte_length;
32use vortex::expr::cast;
33use vortex::expr::get_item;
34use vortex::expr::is_not_null;
35use vortex::expr::is_null;
36use vortex::expr::list_contains;
37use vortex::expr::list_length;
38use vortex::expr::lit;
39use vortex::expr::nested_case_when;
40use vortex::expr::not;
41use vortex::expr::pack;
42use vortex::expr::root;
43use vortex::scalar::Scalar;
44use vortex::scalar_fn::ScalarFnVTableExt;
45use vortex::scalar_fn::fns::binary::Binary;
46use vortex::scalar_fn::fns::like::Like;
47use vortex::scalar_fn::fns::like::LikeOptions;
48use vortex::scalar_fn::fns::operators::Operator;
49use vortex::session::VortexSession;
50
51use crate::convert::FromDataFusion;
52
53pub struct ProcessedProjection {
55 pub scan_projection: Expression,
56 pub leftover_projection: ProjectionExprs,
57}
58
59pub(crate) fn make_vortex_predicate(
61 expr_convertor: &dyn ExpressionConvertor,
62 predicate: &[Arc<dyn PhysicalExpr>],
63) -> DFResult<Option<Expression>> {
64 let exprs = predicate
65 .iter()
66 .map(|e| expr_convertor.convert(e.as_ref()))
67 .collect::<DFResult<Vec<_>>>()?;
68
69 Ok(and_collect(exprs))
70}
71
72pub trait ExpressionConvertor: Send + Sync {
74 fn can_be_pushed_down(&self, expr: &Arc<dyn PhysicalExpr>, schema: &Schema) -> bool;
76
77 fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult<Expression>;
79
80 fn split_projection(
83 &self,
84 source_projection: ProjectionExprs,
85 input_schema: &Schema,
86 output_schema: &Schema,
87 ) -> DFResult<ProcessedProjection>;
88
89 fn no_pushdown_projection(
92 &self,
93 source_projection: ProjectionExprs,
94 input_schema: &Schema,
95 ) -> DFResult<ProcessedProjection> {
96 let column_indices = source_projection.column_indices();
98
99 let scan_columns: Vec<(String, Expression)> = column_indices
101 .into_iter()
102 .map(|idx| {
103 let field = input_schema.field(idx);
104 let name = field.name().clone();
105 (name.clone(), get_item(name, root()))
106 })
107 .collect();
108
109 Ok(ProcessedProjection {
110 scan_projection: pack(scan_columns, Nullability::NonNullable),
111 leftover_projection: source_projection,
112 })
113 }
114}
115
116pub struct DefaultExpressionConvertor {
118 session: VortexSession,
123}
124
125impl Default for DefaultExpressionConvertor {
126 fn default() -> Self {
127 Self {
128 session: VortexSession::default(),
129 }
130 }
131}
132
133impl DefaultExpressionConvertor {
134 pub fn new(session: VortexSession) -> Self {
137 Self { session }
138 }
139
140 fn try_convert_octet_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult<Expression> {
142 let [input] = scalar_fn.args() else {
143 return Err(exec_datafusion_err!(
144 "octet_length requires exactly one argument"
145 ));
146 };
147
148 let input = self.convert(input.as_ref())?;
149 let return_dtype = self
150 .session
151 .arrow()
152 .from_arrow_field(&Field::new(
153 "",
154 scalar_fn.return_type().clone(),
155 scalar_fn.nullable(),
156 ))
157 .map_err(|e| exec_datafusion_err!("Failed to convert return type to dtype: {e}"))?;
158 Ok(cast(byte_length(input), return_dtype))
159 }
160
161 fn try_convert_array_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult<Expression> {
167 let Some(input) = array_length_input(scalar_fn) else {
168 return Err(exec_datafusion_err!(
169 "array_length pushdown supports only the one-argument form or an explicit first \
170 dimension"
171 ));
172 };
173
174 let input = self.convert(input.as_ref())?;
175 let return_dtype = self
176 .session
177 .arrow()
178 .from_arrow_field(&Field::new(
179 "",
180 scalar_fn.return_type().clone(),
181 scalar_fn.nullable(),
182 ))
183 .map_err(|e| exec_datafusion_err!("Failed to convert return type to dtype: {e}"))?;
184 Ok(cast(list_length(input), return_dtype))
185 }
186
187 fn try_convert_scalar_function(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult<Expression> {
189 if let Some(octet_length_fn) =
190 ScalarFunctionExpr::try_downcast_func::<OctetLengthFunc>(scalar_fn)
191 {
192 return self.try_convert_octet_length(octet_length_fn);
193 }
194
195 if let Some(array_length_fn) =
196 ScalarFunctionExpr::try_downcast_func::<ArrayLength>(scalar_fn)
197 {
198 return self.try_convert_array_length(array_length_fn);
199 }
200
201 if let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn)
202 {
203 let (source_expr, field_names) = get_field_fn
208 .args()
209 .split_first()
210 .ok_or_else(|| exec_datafusion_err!("get_field missing source expression"))?;
211
212 let mut result = self.convert(source_expr.as_ref())?;
213 for expr in field_names {
214 let field_name = expr
215 .downcast_ref::<df_expr::Literal>()
216 .ok_or_else(|| exec_datafusion_err!("get_field field name must be a literal"))?
217 .value()
218 .try_as_str()
219 .flatten()
220 .ok_or_else(|| {
221 exec_datafusion_err!("get_field field name must be a UTF-8 string")
222 })?;
223 result = get_item(field_name.to_string(), result);
224 }
225 return Ok(result);
226 }
227
228 Err(exec_datafusion_err!(
229 "Unsupported ScalarFunctionExpr: {}",
230 scalar_fn.name()
231 ))
232 }
233
234 fn try_convert_case_expr(&self, case_expr: &df_expr::CaseExpr) -> DFResult<Expression> {
236 if case_expr.expr().is_some() {
243 return Err(exec_datafusion_err!(
244 "CASE expr WHEN form is not yet supported, only searched CASE is supported"
245 ));
246 }
247
248 let when_then_pairs = case_expr.when_then_expr();
249 if when_then_pairs.is_empty() {
250 return Err(exec_datafusion_err!(
251 "CASE expression must have at least one WHEN clause"
252 ));
253 }
254
255 let mut pairs = Vec::with_capacity(when_then_pairs.len());
257 for (when_expr, then_expr) in when_then_pairs {
258 let condition = self.convert(when_expr.as_ref())?;
259 let value = self.convert(then_expr.as_ref())?;
260 pairs.push((condition, value));
261 }
262
263 let else_value = case_expr
265 .else_expr()
266 .map(|e| self.convert(e.as_ref()))
267 .transpose()?;
268
269 Ok(nested_case_when(pairs, else_value))
271 }
272}
273
274impl ExpressionConvertor for DefaultExpressionConvertor {
275 fn can_be_pushed_down(&self, expr: &Arc<dyn PhysicalExpr>, schema: &Schema) -> bool {
276 can_be_pushed_down_impl(expr, schema)
277 }
278
279 fn convert(&self, df: &dyn PhysicalExpr) -> DFResult<Expression> {
280 if let Some(binary_expr) = df.downcast_ref::<df_expr::BinaryExpr>() {
283 let left = self.convert(binary_expr.left().as_ref())?;
284 let right = self.convert(binary_expr.right().as_ref())?;
285 let operator = try_operator_from_df(binary_expr.op())?;
286
287 return Ok(Binary.new_expr(operator, [left, right]));
288 }
289
290 if let Some(col_expr) = df.downcast_ref::<df_expr::Column>() {
291 return Ok(get_item(col_expr.name().to_owned(), root()));
292 }
293
294 if let Some(like) = df.downcast_ref::<df_expr::LikeExpr>() {
295 let child = self.convert(like.expr().as_ref())?;
296 let pattern = self.convert(like.pattern().as_ref())?;
297 return Ok(Like.new_expr(
298 LikeOptions {
299 negated: like.negated(),
300 case_insensitive: like.case_insensitive(),
301 },
302 [child, pattern],
303 ));
304 }
305
306 if let Some(literal) = df.downcast_ref::<df_expr::Literal>() {
307 let value = Scalar::from_df(literal.value());
308 return Ok(lit(value));
309 }
310
311 if let Some(cast_expr) = df.downcast_ref::<df_expr::CastExpr>() {
312 let cast_dtype = self
313 .session
314 .arrow()
315 .from_arrow_field(cast_expr.target_field().as_ref())
316 .map_err(|e| exec_datafusion_err!("Failed to convert cast target to dtype: {e}"))?;
317 let child = self.convert(cast_expr.expr().as_ref())?;
318 return Ok(cast(child, cast_dtype));
319 }
320
321 if let Some(is_null_expr) = df.downcast_ref::<df_expr::IsNullExpr>() {
322 let arg = self.convert(is_null_expr.arg().as_ref())?;
323 return Ok(is_null(arg));
324 }
325
326 if let Some(is_not_null_expr) = df.downcast_ref::<df_expr::IsNotNullExpr>() {
327 let arg = self.convert(is_not_null_expr.arg().as_ref())?;
328 return Ok(is_not_null(arg));
329 }
330
331 if let Some(in_list) = df.downcast_ref::<df_expr::InListExpr>() {
332 let value = self.convert(in_list.expr().as_ref())?;
333 let list_elements: Vec<_> = in_list
334 .list()
335 .iter()
336 .map(|e| {
337 if let Some(lit) = e.downcast_ref::<df_expr::Literal>() {
338 Ok(Scalar::from_df(lit.value()))
339 } else {
340 Err(exec_datafusion_err!("Failed to cast sub-expression"))
341 }
342 })
343 .try_collect()?;
344
345 let list = Scalar::list(
346 list_elements[0].dtype().clone(),
347 list_elements,
348 Nullability::Nullable,
349 );
350 let expr = list_contains(lit(list), value);
351
352 return Ok(if in_list.negated() { not(expr) } else { expr });
353 }
354
355 if let Some(scalar_fn) = df.downcast_ref::<ScalarFunctionExpr>() {
356 return self.try_convert_scalar_function(scalar_fn);
357 }
358
359 if let Some(case_expr) = df.downcast_ref::<df_expr::CaseExpr>() {
360 return self.try_convert_case_expr(case_expr);
361 }
362
363 Err(exec_datafusion_err!(
364 "Couldn't convert DataFusion physical {df} expression to a vortex expression"
365 ))
366 }
367
368 fn split_projection(
369 &self,
370 source_projection: ProjectionExprs,
371 input_schema: &Schema,
372 output_schema: &Schema,
373 ) -> DFResult<ProcessedProjection> {
374 let mut scan_projection = vec![];
375 let mut leftover_projection: Vec<ProjectionExpr> = vec![];
376
377 for projection_expr in source_projection.iter() {
378 let r = projection_expr.expr.apply(|node| {
379 if let Some(scalar_fn_expr) = node.downcast_ref::<ScalarFunctionExpr>()
381 && !can_scalar_fn_be_pushed_down(scalar_fn_expr, input_schema)
382 {
383 scan_projection.extend(
384 collect_columns(node)
385 .into_iter()
386 .map(|c| (c.name().to_string(), get_item(c.name(), root()))),
387 );
388
389 leftover_projection.push(projection_expr.clone());
390 return Ok(TreeNodeRecursion::Stop);
391 }
392
393 if let Some(binary_expr) = node.downcast_ref::<df_expr::BinaryExpr>()
396 && binary_expr.op().is_numerical_operators()
397 && binary_expr.left().data_type(input_schema)?.is_decimal()
398 && binary_expr.right().data_type(input_schema)?.is_decimal()
399 {
400 scan_projection.extend(
401 collect_columns(node)
402 .into_iter()
403 .map(|c| (c.name().to_string(), get_item(c.name(), root()))),
404 );
405
406 leftover_projection.push(projection_expr.clone());
407 return Ok(TreeNodeRecursion::Stop);
408 }
409
410 Ok(TreeNodeRecursion::Continue)
411 })?;
412
413 if matches!(r, TreeNodeRecursion::Continue) {
415 scan_projection.push((
416 projection_expr.alias.clone(),
417 self.convert(projection_expr.expr.as_ref())?,
418 ));
419 leftover_projection.push(ProjectionExpr {
420 expr: Arc::new(df_expr::Column::new_with_schema(
421 projection_expr.alias.as_str(),
422 output_schema,
423 )?),
424 alias: projection_expr.alias.clone(),
425 });
426 }
427 }
428
429 Ok(ProcessedProjection {
430 scan_projection: pack(scan_projection, Nullability::NonNullable),
431 leftover_projection: leftover_projection.into(),
432 })
433 }
434}
435
436fn try_operator_from_df(value: &DFOperator) -> DFResult<Operator> {
437 match value {
438 DFOperator::Eq => Ok(Operator::Eq),
439 DFOperator::NotEq => Ok(Operator::NotEq),
440 DFOperator::Lt => Ok(Operator::Lt),
441 DFOperator::LtEq => Ok(Operator::Lte),
442 DFOperator::Gt => Ok(Operator::Gt),
443 DFOperator::GtEq => Ok(Operator::Gte),
444 DFOperator::And => Ok(Operator::And),
445 DFOperator::Or => Ok(Operator::Or),
446 DFOperator::Plus => Ok(Operator::Add),
447 DFOperator::Minus => Ok(Operator::Sub),
448 DFOperator::Multiply => Ok(Operator::Mul),
449 DFOperator::Divide => Ok(Operator::Div),
450 DFOperator::IsDistinctFrom
451 | DFOperator::IsNotDistinctFrom
452 | DFOperator::RegexMatch
453 | DFOperator::RegexIMatch
454 | DFOperator::RegexNotMatch
455 | DFOperator::RegexNotIMatch
456 | DFOperator::LikeMatch
457 | DFOperator::ILikeMatch
458 | DFOperator::NotLikeMatch
459 | DFOperator::NotILikeMatch
460 | DFOperator::BitwiseAnd
461 | DFOperator::BitwiseOr
462 | DFOperator::BitwiseXor
463 | DFOperator::BitwiseShiftRight
464 | DFOperator::BitwiseShiftLeft
465 | DFOperator::StringConcat
466 | DFOperator::AtArrow
467 | DFOperator::ArrowAt
468 | DFOperator::Modulo
469 | DFOperator::Arrow
470 | DFOperator::LongArrow
471 | DFOperator::HashArrow
472 | DFOperator::HashLongArrow
473 | DFOperator::AtAt
474 | DFOperator::IntegerDivide
475 | DFOperator::HashMinus
476 | DFOperator::AtQuestion
477 | DFOperator::Question
478 | DFOperator::QuestionAnd
479 | DFOperator::QuestionPipe
480 | DFOperator::Colon => {
481 tracing::debug!(operator = %value, "Can't pushdown binary_operator operator");
482 Err(exec_datafusion_err!(
483 "Unsupported datafusion operator {value}"
484 ))
485 }
486 }
487}
488
489fn can_be_pushed_down_impl(expr: &Arc<dyn PhysicalExpr>, schema: &Schema) -> bool {
490 if is_dynamic_physical_expr(expr) {
493 return false;
494 }
495
496 if let Some(binary) = expr.downcast_ref::<df_expr::BinaryExpr>() {
497 can_binary_be_pushed_down(binary, schema)
498 } else if let Some(col) = expr.downcast_ref::<df_expr::Column>() {
499 schema
500 .field_with_name(col.name())
501 .ok()
502 .is_some_and(|field| supported_data_types(field.data_type()))
503 } else if let Some(like) = expr.downcast_ref::<df_expr::LikeExpr>() {
504 can_be_pushed_down_impl(like.expr(), schema)
505 && can_be_pushed_down_impl(like.pattern(), schema)
506 } else if let Some(lit) = expr.downcast_ref::<df_expr::Literal>() {
507 supported_data_types(&lit.value().data_type())
508 } else if let Some(cast_expr) = expr.downcast_ref::<df_expr::CastExpr>() {
509 is_convertible_expr(cast_expr.expr())
511 } else if let Some(is_null) = expr.downcast_ref::<df_expr::IsNullExpr>() {
512 can_be_pushed_down_impl(is_null.arg(), schema)
513 } else if let Some(is_not_null) = expr.downcast_ref::<df_expr::IsNotNullExpr>() {
514 can_be_pushed_down_impl(is_not_null.arg(), schema)
515 } else if let Some(in_list) = expr.downcast_ref::<df_expr::InListExpr>() {
516 can_be_pushed_down_impl(in_list.expr(), schema)
517 && in_list
518 .list()
519 .iter()
520 .all(|e| can_be_pushed_down_impl(e, schema))
521 } else if let Some(scalar_fn) = expr.downcast_ref::<ScalarFunctionExpr>() {
522 can_scalar_fn_be_pushed_down(scalar_fn, schema)
523 } else if let Some(case_expr) = expr.downcast_ref::<df_expr::CaseExpr>() {
524 can_case_be_pushed_down(case_expr, schema)
525 } else {
526 tracing::debug!(%expr, "DataFusion expression can't be pushed down");
527 false
528 }
529}
530
531fn is_convertible_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
535 expr.downcast_ref::<df_expr::BinaryExpr>().is_some()
537 || expr.downcast_ref::<df_expr::Column>().is_some()
538 || expr.downcast_ref::<df_expr::LikeExpr>().is_some()
539 || expr.downcast_ref::<df_expr::Literal>().is_some()
540 || expr
541 .downcast_ref::<df_expr::CastExpr>()
542 .is_some_and(|e| is_convertible_expr(e.expr()))
543 || expr.downcast_ref::<df_expr::IsNullExpr>().is_some()
544 || expr.downcast_ref::<df_expr::IsNotNullExpr>().is_some()
545 || expr.downcast_ref::<df_expr::InListExpr>().is_some()
546 || expr.downcast_ref::<ScalarFunctionExpr>().is_some_and(|sf| {
547 ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(sf).is_some()
548 || ScalarFunctionExpr::try_downcast_func::<OctetLengthFunc>(sf).is_some()
549 || ScalarFunctionExpr::try_downcast_func::<ArrayLength>(sf).is_some()
550 })
551}
552
553fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool {
554 let is_op_supported = try_operator_from_df(binary.op()).is_ok();
555 is_op_supported
556 && can_be_pushed_down_impl(binary.left(), schema)
557 && can_be_pushed_down_impl(binary.right(), schema)
558}
559
560fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool {
561 if case_expr.expr().is_some() {
564 return false;
565 }
566
567 for (when_expr, then_expr) in case_expr.when_then_expr() {
569 if !can_be_pushed_down_impl(when_expr, schema)
570 || !can_be_pushed_down_impl(then_expr, schema)
571 {
572 return false;
573 }
574 }
575
576 if let Some(else_expr) = case_expr.else_expr()
578 && !can_be_pushed_down_impl(else_expr, schema)
579 {
580 return false;
581 }
582
583 true
584}
585
586fn supported_data_types(dt: &DataType) -> bool {
587 use DataType::*;
588
589 if let Dictionary(_, value_type) = dt {
591 return supported_data_types(value_type.as_ref());
592 }
593
594 let is_supported = dt.is_null()
595 || dt.is_numeric()
596 || dt.is_binary()
597 || dt.is_string()
598 || matches!(
599 dt,
600 Boolean | Date32 | Date64 | Timestamp(_, _) | Time32(_) | Time64(_)
601 );
602
603 if !is_supported {
604 tracing::debug!("DataFusion data type {dt:?} is not supported");
605 }
606
607 is_supported
608}
609
610fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool {
613 if ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn).is_some() {
614 return true;
615 }
616
617 if ScalarFunctionExpr::try_downcast_func::<OctetLengthFunc>(scalar_fn)
618 .is_some_and(|octet_length| can_octet_length_be_pushed_down(octet_length, schema))
619 {
620 return true;
621 }
622
623 ScalarFunctionExpr::try_downcast_func::<ArrayLength>(scalar_fn)
624 .is_some_and(|array_length| can_array_length_be_pushed_down(array_length, schema))
625}
626
627fn can_octet_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool {
628 let [input] = scalar_fn.args() else {
629 return false;
630 };
631
632 input.data_type(schema).as_ref().is_ok_and(|data_type| {
633 let dt = if let DataType::Dictionary(_, value_type) = data_type {
634 value_type.as_ref()
635 } else {
636 data_type
637 };
638
639 dt.is_binary() || dt.is_string()
640 }) && can_be_pushed_down_impl(input, schema)
641}
642
643fn can_array_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool {
644 let Some(input) = array_length_input(scalar_fn) else {
645 return false;
646 };
647
648 input.data_type(schema).as_ref().is_ok_and(|data_type| {
652 matches!(
653 data_type,
654 DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _)
655 )
656 }) && is_convertible_expr(input)
657}
658
659fn array_length_input(scalar_fn: &ScalarFunctionExpr) -> Option<&Arc<dyn PhysicalExpr>> {
664 match scalar_fn.args() {
665 [input] => Some(input),
666 [input, dimension] if is_dimension_one(dimension) => Some(input),
667 _ => None,
668 }
669}
670
671fn is_dimension_one(expr: &Arc<dyn PhysicalExpr>) -> bool {
675 expr.downcast_ref::<df_expr::Literal>()
676 .is_some_and(|literal| matches!(literal.value(), ScalarValue::Int64(Some(1))))
677}
678
679#[cfg(test)]
680mod tests {
681 use std::sync::Arc;
682
683 use arrow_schema::DataType;
684 use arrow_schema::Field;
685 use arrow_schema::Schema;
686 use arrow_schema::TimeUnit as ArrowTimeUnit;
687 use datafusion::arrow::array::AsArray;
688 use datafusion::arrow::datatypes::Int32Type;
689 use datafusion_common::ScalarValue;
690 use datafusion_common::config::ConfigOptions;
691 use datafusion_expr::Operator as DFOperator;
692 use datafusion_expr::ScalarUDF;
693 use datafusion_physical_expr::PhysicalExpr;
694 use datafusion_physical_plan::expressions as df_expr;
695 use insta::assert_snapshot;
696 use rstest::rstest;
697
698 use super::*;
699 use crate::common_tests::TestSessionContext;
700
701 #[rstest::fixture]
702 fn test_schema() -> Schema {
703 Schema::new(vec![
704 Field::new("id", DataType::Int32, false),
705 Field::new("name", DataType::Utf8, true),
706 Field::new("score", DataType::Float64, true),
707 Field::new("active", DataType::Boolean, false),
708 Field::new(
709 "created_at",
710 DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
711 true,
712 ),
713 Field::new(
714 "tags",
715 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
716 true,
717 ),
718 ])
719 }
720
721 fn octet_length_expr(input: Arc<dyn PhysicalExpr>, schema: &Schema) -> Arc<dyn PhysicalExpr> {
722 Arc::new(
723 ScalarFunctionExpr::try_new(
724 Arc::new(ScalarUDF::from(OctetLengthFunc::new())),
725 vec![input],
726 schema,
727 Arc::new(ConfigOptions::new()),
728 )
729 .unwrap(),
730 )
731 }
732
733 fn array_length_expr(
734 args: Vec<Arc<dyn PhysicalExpr>>,
735 schema: &Schema,
736 ) -> Arc<dyn PhysicalExpr> {
737 Arc::new(
738 ScalarFunctionExpr::try_new(
739 Arc::new(ScalarUDF::from(ArrayLength::new())),
740 args,
741 schema,
742 Arc::new(ConfigOptions::new()),
743 )
744 .unwrap(),
745 )
746 }
747
748 #[test]
749 fn test_make_vortex_predicate_empty() {
750 let expr_convertor = DefaultExpressionConvertor::default();
751 let result = make_vortex_predicate(&expr_convertor, &[]).unwrap();
752 assert!(result.is_none());
753 }
754
755 #[test]
756 fn test_make_vortex_predicate_single() {
757 let expr_convertor = DefaultExpressionConvertor::default();
758 let col_expr = Arc::new(df_expr::Column::new("test", 0)) as Arc<dyn PhysicalExpr>;
759 let result = make_vortex_predicate(&expr_convertor, &[col_expr]).unwrap();
760 assert!(result.is_some());
761 }
762
763 #[test]
764 fn test_make_vortex_predicate_multiple() {
765 let expr_convertor = DefaultExpressionConvertor::default();
766 let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc<dyn PhysicalExpr>;
767 let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc<dyn PhysicalExpr>;
768 let result = make_vortex_predicate(&expr_convertor, &[col1, col2]).unwrap();
769 assert!(result.is_some());
770 }
772
773 #[rstest]
774 #[case::eq(DFOperator::Eq, Operator::Eq)]
775 #[case::not_eq(DFOperator::NotEq, Operator::NotEq)]
776 #[case::lt(DFOperator::Lt, Operator::Lt)]
777 #[case::lte(DFOperator::LtEq, Operator::Lte)]
778 #[case::gt(DFOperator::Gt, Operator::Gt)]
779 #[case::gte(DFOperator::GtEq, Operator::Gte)]
780 #[case::and(DFOperator::And, Operator::And)]
781 #[case::or(DFOperator::Or, Operator::Or)]
782 #[case::plus(DFOperator::Plus, Operator::Add)]
783 #[case::plus(DFOperator::Minus, Operator::Sub)]
784 #[case::plus(DFOperator::Multiply, Operator::Mul)]
785 #[case::plus(DFOperator::Divide, Operator::Div)]
786 fn test_operator_conversion_supported(
787 #[case] df_op: DFOperator,
788 #[case] expected_vortex_op: Operator,
789 ) {
790 let result = try_operator_from_df(&df_op).unwrap();
791 assert_eq!(result, expected_vortex_op);
792 }
793
794 #[rstest]
795 #[case::modulo(DFOperator::Modulo)]
796 #[case::bitwise_and(DFOperator::BitwiseAnd)]
797 #[case::regex_match(DFOperator::RegexMatch)]
798 #[case::like_match(DFOperator::LikeMatch)]
799 fn test_operator_conversion_unsupported(#[case] df_op: DFOperator) {
800 let result = try_operator_from_df(&df_op);
801 assert!(result.is_err());
802 assert!(
803 result
804 .unwrap_err()
805 .to_string()
806 .contains("Unsupported datafusion operator")
807 );
808 }
809
810 #[test]
811 fn test_expr_from_df_column() {
812 let col_expr = df_expr::Column::new("test_column", 0);
813 let result = DefaultExpressionConvertor::default()
814 .convert(&col_expr)
815 .unwrap();
816
817 assert_snapshot!(result.display_tree().to_string(), @r"
818 vortex.get_item(test_column)
819 └── input: vortex.root()
820 ");
821 }
822
823 #[test]
824 fn test_expr_from_df_literal() {
825 let literal_expr = df_expr::Literal::new(ScalarValue::Int32(Some(42)));
826 let result = DefaultExpressionConvertor::default()
827 .convert(&literal_expr)
828 .unwrap();
829
830 assert_snapshot!(result.display_tree().to_string(), @"vortex.literal(42i32)");
831 }
832
833 #[test]
834 fn test_expr_from_df_binary() {
835 let left = Arc::new(df_expr::Column::new("left", 0)) as Arc<dyn PhysicalExpr>;
836 let right =
837 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
838 let binary_expr = df_expr::BinaryExpr::new(left, DFOperator::Eq, right);
839
840 let result = DefaultExpressionConvertor::default()
841 .convert(&binary_expr)
842 .unwrap();
843
844 assert_snapshot!(result.display_tree().to_string(), @r"
845 vortex.binary(=)
846 ├── lhs: vortex.get_item(left)
847 │ └── input: vortex.root()
848 └── rhs: vortex.literal(42i32)
849 ");
850 }
851
852 #[rstest]
853 #[case::like_normal(false, false)]
854 #[case::like_negated(true, false)]
855 #[case::like_case_insensitive(false, true)]
856 #[case::like_negated_case_insensitive(true, true)]
857 fn test_expr_from_df_like(#[case] negated: bool, #[case] case_insensitive: bool) {
858 let expr = Arc::new(df_expr::Column::new("text_col", 0)) as Arc<dyn PhysicalExpr>;
859 let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some(
860 "test%".to_string(),
861 )))) as Arc<dyn PhysicalExpr>;
862 let like_expr = df_expr::LikeExpr::new(negated, case_insensitive, expr, pattern);
863
864 let result = DefaultExpressionConvertor::default()
865 .convert(&like_expr)
866 .unwrap();
867 let like_opts = result.as_::<Like>();
868 assert_eq!(
869 like_opts,
870 &LikeOptions {
871 negated,
872 case_insensitive
873 }
874 );
875 }
876
877 #[rstest]
878 fn test_expr_from_df_octet_length(test_schema: Schema) {
879 let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc<dyn PhysicalExpr>;
880 let octet_length = octet_length_expr(expr, &test_schema);
881
882 let result = DefaultExpressionConvertor::default()
883 .convert(octet_length.as_ref())
884 .unwrap();
885
886 assert_snapshot!(result.display_tree().to_string(), @r"
887 vortex.cast(i32?)
888 └── input: vortex.byte_length()
889 └── input: vortex.get_item(name)
890 └── input: vortex.root()
891 ");
892 }
893
894 #[rstest]
895 fn test_expr_from_df_array_length(test_schema: Schema) {
896 let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
897 let array_length = array_length_expr(vec![expr], &test_schema);
898
899 let result = DefaultExpressionConvertor::default()
900 .convert(array_length.as_ref())
901 .unwrap();
902
903 assert_snapshot!(result.display_tree().to_string(), @r"
904 vortex.cast(u64?)
905 └── input: vortex.list.length()
906 └── input: vortex.get_item(tags)
907 └── input: vortex.root()
908 ");
909 }
910
911 #[rstest]
912 #[case::null(DataType::Null, true)]
914 #[case::boolean(DataType::Boolean, true)]
915 #[case::int8(DataType::Int8, true)]
916 #[case::int16(DataType::Int16, true)]
917 #[case::int32(DataType::Int32, true)]
918 #[case::int64(DataType::Int64, true)]
919 #[case::uint8(DataType::UInt8, true)]
920 #[case::uint16(DataType::UInt16, true)]
921 #[case::uint32(DataType::UInt32, true)]
922 #[case::uint64(DataType::UInt64, true)]
923 #[case::float32(DataType::Float32, true)]
924 #[case::float64(DataType::Float64, true)]
925 #[case::utf8(DataType::Utf8, true)]
926 #[case::utf8_view(DataType::Utf8View, true)]
927 #[case::binary(DataType::Binary, true)]
928 #[case::binary_view(DataType::BinaryView, true)]
929 #[case::date32(DataType::Date32, true)]
930 #[case::date64(DataType::Date64, true)]
931 #[case::timestamp_ms(DataType::Timestamp(ArrowTimeUnit::Millisecond, None), true)]
932 #[case::timestamp_us(
933 DataType::Timestamp(ArrowTimeUnit::Microsecond, Some(Arc::from("UTC"))),
934 true
935 )]
936 #[case::time32_s(DataType::Time32(ArrowTimeUnit::Second), true)]
937 #[case::time64_ns(DataType::Time64(ArrowTimeUnit::Nanosecond), true)]
938 #[case::list(
940 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
941 false
942 )]
943 #[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into()
944 ), false)]
945 #[case::dict_utf8(
947 DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
948 true
949 )]
950 #[case::dict_int32(
951 DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Int32)),
952 true
953 )]
954 #[case::dict_unsupported(
955 DataType::Dictionary(
956 Box::new(DataType::UInt32),
957 Box::new(DataType::List(Arc::new(Field::new("item", DataType::Int32, true))))
958 ),
959 false
960 )]
961 fn test_supported_data_types(#[case] data_type: DataType, #[case] expected: bool) {
962 assert_eq!(supported_data_types(&data_type), expected);
963 }
964
965 #[rstest]
966 fn test_can_be_pushed_down_column_supported(test_schema: Schema) {
967 let col_expr = Arc::new(df_expr::Column::new("id", 0)) as Arc<dyn PhysicalExpr>;
968
969 assert!(can_be_pushed_down_impl(&col_expr, &test_schema));
970 }
971
972 #[rstest]
973 fn test_can_be_pushed_down_column_unsupported_type(test_schema: Schema) {
974 let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
975
976 assert!(!can_be_pushed_down_impl(&col_expr, &test_schema));
977 }
978
979 #[rstest]
980 fn test_can_be_pushed_down_column_not_found(test_schema: Schema) {
981 let col_expr = Arc::new(df_expr::Column::new("nonexistent", 99)) as Arc<dyn PhysicalExpr>;
982
983 assert!(!can_be_pushed_down_impl(&col_expr, &test_schema));
984 }
985
986 #[rstest]
987 fn test_can_be_pushed_down_literal_supported(test_schema: Schema) {
988 let lit_expr =
989 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
990
991 assert!(can_be_pushed_down_impl(&lit_expr, &test_schema));
992 }
993
994 #[rstest]
995 fn test_can_be_pushed_down_literal_unsupported(test_schema: Schema) {
996 let unsupported_literal = ScalarValue::DurationSecond(Some(42));
998 let lit_expr =
999 Arc::new(df_expr::Literal::new(unsupported_literal)) as Arc<dyn PhysicalExpr>;
1000
1001 assert!(!can_be_pushed_down_impl(&lit_expr, &test_schema));
1002 }
1003
1004 #[rstest]
1005 fn test_can_be_pushed_down_binary_supported(test_schema: Schema) {
1006 let left = Arc::new(df_expr::Column::new("id", 0)) as Arc<dyn PhysicalExpr>;
1007 let right =
1008 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
1009 let binary_expr = Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right))
1010 as Arc<dyn PhysicalExpr>;
1011
1012 assert!(can_be_pushed_down_impl(&binary_expr, &test_schema));
1013 }
1014
1015 #[rstest]
1016 fn test_can_be_pushed_down_binary_unsupported_operator(test_schema: Schema) {
1017 let left = Arc::new(df_expr::Column::new("id", 0)) as Arc<dyn PhysicalExpr>;
1018 let right =
1019 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
1020 let binary_expr = Arc::new(df_expr::BinaryExpr::new(
1021 left,
1022 DFOperator::AtQuestion,
1023 right,
1024 )) as Arc<dyn PhysicalExpr>;
1025
1026 assert!(!can_be_pushed_down_impl(&binary_expr, &test_schema));
1027 }
1028
1029 #[rstest]
1030 fn test_can_be_pushed_down_binary_unsupported_operand(test_schema: Schema) {
1031 let left = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
1032 let right =
1033 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
1034 let binary_expr = Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right))
1035 as Arc<dyn PhysicalExpr>;
1036
1037 assert!(!can_be_pushed_down_impl(&binary_expr, &test_schema));
1038 }
1039
1040 #[rstest]
1041 fn test_can_be_pushed_down_like_supported(test_schema: Schema) {
1042 let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc<dyn PhysicalExpr>;
1043 let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some(
1044 "test%".to_string(),
1045 )))) as Arc<dyn PhysicalExpr>;
1046 let like_expr =
1047 Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc<dyn PhysicalExpr>;
1048
1049 assert!(can_be_pushed_down_impl(&like_expr, &test_schema));
1050 }
1051
1052 #[rstest]
1053 fn test_can_be_pushed_down_like_unsupported_operand(test_schema: Schema) {
1054 let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
1055 let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some(
1056 "test%".to_string(),
1057 )))) as Arc<dyn PhysicalExpr>;
1058 let like_expr =
1059 Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc<dyn PhysicalExpr>;
1060
1061 assert!(!can_be_pushed_down_impl(&like_expr, &test_schema));
1062 }
1063
1064 #[rstest]
1065 fn test_can_be_pushed_down_octet_length_supported(test_schema: Schema) {
1066 let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc<dyn PhysicalExpr>;
1067 let octet_length = octet_length_expr(expr, &test_schema);
1068
1069 assert!(can_be_pushed_down_impl(&octet_length, &test_schema));
1070 }
1071
1072 #[rstest]
1073 fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) {
1074 let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
1075 let octet_length = Arc::new(ScalarFunctionExpr::new(
1076 "octet_length",
1077 Arc::new(ScalarUDF::from(OctetLengthFunc::new())),
1078 vec![expr],
1079 Arc::new(Field::new("octet_length", DataType::Int32, true)),
1080 Arc::new(ConfigOptions::new()),
1081 )) as Arc<dyn PhysicalExpr>;
1082
1083 assert!(!can_be_pushed_down_impl(&octet_length, &test_schema));
1084 }
1085
1086 #[rstest]
1087 fn test_can_be_pushed_down_array_length_supported(test_schema: Schema) {
1088 let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
1089 let array_length = array_length_expr(vec![expr], &test_schema);
1090
1091 assert!(can_be_pushed_down_impl(&array_length, &test_schema));
1092 }
1093
1094 #[rstest]
1095 fn test_can_be_pushed_down_array_length_unsupported_operand(test_schema: Schema) {
1096 let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc<dyn PhysicalExpr>;
1098 let array_length = Arc::new(ScalarFunctionExpr::new(
1099 "array_length",
1100 Arc::new(ScalarUDF::from(ArrayLength::new())),
1101 vec![expr],
1102 Arc::new(Field::new("array_length", DataType::UInt64, true)),
1103 Arc::new(ConfigOptions::new()),
1104 )) as Arc<dyn PhysicalExpr>;
1105
1106 assert!(!can_be_pushed_down_impl(&array_length, &test_schema));
1107 }
1108
1109 #[rstest]
1110 fn test_can_be_pushed_down_array_length_dimension_one_supported(test_schema: Schema) {
1111 let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
1113 let dimension =
1114 Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(1)))) as Arc<dyn PhysicalExpr>;
1115 let array_length = array_length_expr(vec![list, dimension], &test_schema);
1116
1117 assert!(can_be_pushed_down_impl(&array_length, &test_schema));
1118 }
1119
1120 #[rstest]
1121 fn test_can_be_pushed_down_array_length_higher_dimension_not_supported(test_schema: Schema) {
1122 let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc<dyn PhysicalExpr>;
1125 let dimension =
1126 Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(2)))) as Arc<dyn PhysicalExpr>;
1127 let array_length = array_length_expr(vec![list, dimension], &test_schema);
1128
1129 assert!(!can_be_pushed_down_impl(&array_length, &test_schema));
1130 }
1131
1132 #[tokio::test]
1134 async fn test_cast_int_to_string() -> anyhow::Result<()> {
1135 let ctx = TestSessionContext::default();
1136
1137 ctx.session
1138 .sql(r#"copy (select 1 as id) to 'example.vortex'"#)
1139 .await?
1140 .show()
1141 .await?;
1142
1143 ctx.session
1144 .sql(r#"select cast(id as string) as sid from 'example.vortex' where id > 0"#)
1145 .await?
1146 .show()
1147 .await?;
1148
1149 ctx.session
1150 .sql(r#"select id from 'example.vortex' where cast (id as string) == '1'"#)
1151 .await?
1152 .show()
1153 .await?;
1154
1155 ctx.session
1157 .sql(r#"select cast(id as string) from 'example.vortex'"#)
1158 .await?
1159 .collect()
1160 .await?;
1161
1162 Ok(())
1163 }
1164
1165 #[test]
1170 fn test_cast_to_uuid_resolves_via_registry() -> anyhow::Result<()> {
1171 use arrow_schema::extension::Uuid;
1172
1173 let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), true);
1174 uuid_field.try_with_extension_type(Uuid)?;
1175
1176 let child = Arc::new(df_expr::Column::new("id", 0)) as Arc<dyn PhysicalExpr>;
1177 let cast = df_expr::CastExpr::new_with_target_field(child, Arc::new(uuid_field), None);
1178
1179 DefaultExpressionConvertor::default().convert(&cast)?;
1181 Ok(())
1182 }
1183
1184 #[test]
1187 fn test_case_when_datafusion_vortex_equivalence() {
1188 use datafusion::arrow::array::Int32Array;
1189 use datafusion::arrow::array::RecordBatch;
1190 use datafusion_physical_expr::expressions::CaseExpr;
1191 use vortex::VortexSessionDefault;
1192 use vortex::array::ArrayRef;
1193 use vortex::array::Canonical;
1194 use vortex::array::VortexSessionExecute as _;
1195 use vortex::array::arrow::FromArrowArray;
1196 use vortex::session::VortexSession;
1197
1198 let values = Arc::new(Int32Array::from(vec![1, 5, 10, 15, 20]));
1200 let schema = Arc::new(Schema::new(vec![Field::new(
1201 "value",
1202 DataType::Int32,
1203 false,
1204 )]));
1205 let batch = RecordBatch::try_new(schema, vec![values]).unwrap();
1206
1207 let col_value = Arc::new(df_expr::Column::new("value", 0)) as Arc<dyn PhysicalExpr>;
1210 let lit_10 =
1211 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(10)))) as Arc<dyn PhysicalExpr>;
1212 let lit_5 =
1213 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(5)))) as Arc<dyn PhysicalExpr>;
1214 let lit_100 =
1215 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(100)))) as Arc<dyn PhysicalExpr>;
1216 let lit_50 =
1217 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(50)))) as Arc<dyn PhysicalExpr>;
1218 let lit_0 =
1219 Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))) as Arc<dyn PhysicalExpr>;
1220
1221 let when1 = Arc::new(df_expr::BinaryExpr::new(
1223 Arc::clone(&col_value),
1224 DFOperator::Gt,
1225 lit_10,
1226 )) as Arc<dyn PhysicalExpr>;
1227 let when2 = Arc::new(df_expr::BinaryExpr::new(col_value, DFOperator::Gt, lit_5))
1229 as Arc<dyn PhysicalExpr>;
1230
1231 let case_expr =
1232 CaseExpr::try_new(None, vec![(when1, lit_100), (when2, lit_50)], Some(lit_0)).unwrap();
1233
1234 let df_result = case_expr.evaluate(&batch).unwrap();
1236 let df_array = df_result.into_array(batch.num_rows()).unwrap();
1237
1238 let expr_convertor = DefaultExpressionConvertor::default();
1240 let vortex_expr = expr_convertor.try_convert_case_expr(&case_expr).unwrap();
1241
1242 let vortex_array: ArrayRef = ArrayRef::from_arrow(&batch, false).unwrap();
1244
1245 let session = VortexSession::default();
1247 let mut ctx = session.create_execution_ctx();
1248 let vortex_result = vortex_array
1249 .apply(&vortex_expr)
1250 .unwrap()
1251 .execute::<Canonical>(&mut ctx)
1252 .unwrap();
1253
1254 let vortex_as_arrow = vortex_result.into_primitive().as_slice::<i32>().to_vec();
1256
1257 let df_as_arrow: Vec<i32> = df_array.as_primitive::<Int32Type>().values().to_vec();
1259
1260 assert_eq!(df_as_arrow, vec![0, 0, 50, 100, 100]);
1268 assert_eq!(vortex_as_arrow, df_as_arrow);
1269 }
1270}