1use std::collections::{HashMap, HashSet};
21use std::sync::Arc;
22
23use arrow::datatypes::DataType;
24use indexmap::IndexSet;
25use itertools::Itertools;
26use log::{Level, debug, log_enabled};
27
28use datafusion_common::instant::Instant;
29use datafusion_common::tree_node::{
30 Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
31};
32use datafusion_common::{
33 Column, DFSchema, Result, assert_eq_or_internal_err, internal_err, plan_err,
34 qualified_name,
35};
36use datafusion_expr::expr::WindowFunction;
37use datafusion_expr::expr_rewriter::replace_col;
38use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan};
39use datafusion_expr::utils::{
40 conjunction, expr_to_columns, split_conjunction, split_conjunction_owned,
41};
42use datafusion_expr::{
43 BinaryExpr, Distinct, Expr, Filter, Operator, Projection,
44 TableProviderFilterPushDown, and, or,
45};
46
47use crate::optimizer::ApplyOrder;
48use crate::simplify_expressions::{reorder_predicates, simplify_predicates};
49use crate::utils::{
50 ColumnReference, has_all_column_refs, is_restrict_null_predicate, schema_columns,
51};
52use crate::{OptimizerConfig, OptimizerRule};
53use datafusion_expr::ExpressionPlacement;
54
55#[derive(Default, Debug)]
143pub struct PushDownFilter {}
144
145pub(crate) fn lr_is_preserved(join_type: JoinType) -> (bool, bool) {
170 match join_type {
171 JoinType::Inner => (true, true),
172 JoinType::Left => (true, false),
173 JoinType::Right => (false, true),
174 JoinType::Full => (false, false),
175 JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => (true, false),
178 JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => (false, true),
181 }
182}
183
184pub(crate) fn on_lr_is_preserved(join_type: JoinType) -> (bool, bool) {
186 join_type.on_lr_is_preserved()
187}
188
189#[derive(Debug)]
192struct ColumnChecker<'a> {
193 left_schema: &'a DFSchema,
195 left_columns: Option<HashSet<ColumnReference<'a>>>,
197 right_schema: &'a DFSchema,
199 right_columns: Option<HashSet<ColumnReference<'a>>>,
201}
202
203impl<'a> ColumnChecker<'a> {
204 fn new(left_schema: &'a DFSchema, right_schema: &'a DFSchema) -> Self {
205 Self {
206 left_schema,
207 left_columns: None,
208 right_schema,
209 right_columns: None,
210 }
211 }
212
213 fn is_left_only(&mut self, predicate: &Expr) -> bool {
215 if self.left_columns.is_none() {
216 self.left_columns = Some(schema_columns(self.left_schema));
217 }
218 has_all_column_refs(predicate, self.left_columns.as_ref().unwrap())
219 }
220
221 fn is_right_only(&mut self, predicate: &Expr) -> bool {
223 if self.right_columns.is_none() {
224 self.right_columns = Some(schema_columns(self.right_schema));
225 }
226 has_all_column_refs(predicate, self.right_columns.as_ref().unwrap())
227 }
228}
229
230fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
232 let mut is_evaluate = true;
233 predicate.apply(|expr| match expr {
234 Expr::Column(_)
235 | Expr::Literal(_, _)
236 | Expr::Placeholder(_)
237 | Expr::ScalarVariable(_, _) => Ok(TreeNodeRecursion::Jump),
238 Expr::Exists { .. }
239 | Expr::InSubquery(_)
240 | Expr::SetComparison(_)
241 | Expr::ScalarSubquery(_)
242 | Expr::OuterReferenceColumn(_, _)
243 | Expr::Unnest(_) => {
244 is_evaluate = false;
245 Ok(TreeNodeRecursion::Stop)
246 }
247 Expr::Alias(_)
248 | Expr::BinaryExpr(_)
249 | Expr::Like(_)
250 | Expr::SimilarTo(_)
251 | Expr::Not(_)
252 | Expr::IsNotNull(_)
253 | Expr::IsNull(_)
254 | Expr::IsTrue(_)
255 | Expr::IsFalse(_)
256 | Expr::IsUnknown(_)
257 | Expr::IsNotTrue(_)
258 | Expr::IsNotFalse(_)
259 | Expr::IsNotUnknown(_)
260 | Expr::Negative(_)
261 | Expr::Between(_)
262 | Expr::Case(_)
263 | Expr::Cast(_)
264 | Expr::TryCast(_)
265 | Expr::InList { .. }
266 | Expr::ScalarFunction(_)
267 | Expr::HigherOrderFunction(_)
268 | Expr::Lambda(_)
269 | Expr::LambdaVariable(_) => Ok(TreeNodeRecursion::Continue),
270 #[expect(deprecated)]
272 Expr::AggregateFunction(_)
273 | Expr::WindowFunction(_)
274 | Expr::Wildcard { .. }
275 | Expr::GroupingSet(_) => internal_err!("Unsupported predicate type"),
276 })?;
277 Ok(is_evaluate)
278}
279
280fn extract_or_clauses_for_join<'a>(
314 filters: &'a [Expr],
315 schema_cols: &'a HashSet<ColumnReference>,
316) -> impl Iterator<Item = Expr> + 'a {
317 filters.iter().filter_map(move |expr| {
319 if let Expr::BinaryExpr(BinaryExpr {
320 left,
321 op: Operator::Or,
322 right,
323 }) = expr
324 {
325 let left_expr = extract_or_clause(left.as_ref(), schema_cols);
326 let right_expr = extract_or_clause(right.as_ref(), schema_cols);
327
328 if let (Some(left_expr), Some(right_expr)) = (left_expr, right_expr) {
330 return Some(or(left_expr, right_expr));
331 }
332 }
333 None
334 })
335}
336
337fn extract_or_clause(
349 expr: &Expr,
350 schema_columns: &HashSet<ColumnReference>,
351) -> Option<Expr> {
352 let mut predicate = None;
353
354 match expr {
355 Expr::BinaryExpr(BinaryExpr {
356 left: l_expr,
357 op: Operator::Or,
358 right: r_expr,
359 }) => {
360 let l_expr = extract_or_clause(l_expr, schema_columns);
361 let r_expr = extract_or_clause(r_expr, schema_columns);
362
363 if let (Some(l_expr), Some(r_expr)) = (l_expr, r_expr) {
364 predicate = Some(or(l_expr, r_expr));
365 }
366 }
367 Expr::BinaryExpr(BinaryExpr {
368 left: l_expr,
369 op: Operator::And,
370 right: r_expr,
371 }) => {
372 let l_expr = extract_or_clause(l_expr, schema_columns);
373 let r_expr = extract_or_clause(r_expr, schema_columns);
374
375 match (l_expr, r_expr) {
376 (Some(l_expr), Some(r_expr)) => {
377 predicate = Some(and(l_expr, r_expr));
378 }
379 (Some(l_expr), None) => {
380 predicate = Some(l_expr);
381 }
382 (None, Some(r_expr)) => {
383 predicate = Some(r_expr);
384 }
385 (None, None) => {
386 predicate = None;
387 }
388 }
389 }
390 _ => {
391 if has_all_column_refs(expr, schema_columns) {
392 predicate = Some(expr.clone());
393 }
394 }
395 }
396
397 predicate
398}
399
400fn push_down_all_join(
402 predicates: Vec<Expr>,
403 inferred_join_predicates: Vec<Expr>,
404 mut join: Join,
405 on_filter: Vec<Expr>,
406) -> Result<Transformed<LogicalPlan>> {
407 let is_inner_join = join.join_type == JoinType::Inner;
408 let (left_preserved, right_preserved) = lr_is_preserved(join.join_type);
410
411 let left_schema = join.left.schema();
416 let right_schema = join.right.schema();
417
418 let left_schema_columns = schema_columns(left_schema.as_ref());
419 let right_schema_columns = schema_columns(right_schema.as_ref());
420
421 let mut left_push = vec![];
422 let mut right_push = vec![];
423 let mut keep_predicates = vec![];
424 let mut join_conditions = vec![];
425 let mut checker = ColumnChecker::new(left_schema, right_schema);
426 for predicate in predicates {
427 if left_preserved && checker.is_left_only(&predicate) {
428 left_push.push(predicate);
429 } else if right_preserved && checker.is_right_only(&predicate) {
430 right_push.push(predicate);
431 } else if is_inner_join && can_evaluate_as_join_condition(&predicate)? {
432 join_conditions.push(predicate);
435 } else {
436 keep_predicates.push(predicate);
437 }
438 }
439
440 for predicate in inferred_join_predicates {
442 if checker.is_left_only(&predicate) {
443 left_push.push(predicate);
444 } else if checker.is_right_only(&predicate) {
445 right_push.push(predicate);
446 }
447 }
448
449 let mut on_filter_join_conditions = vec![];
450 let (on_left_preserved, on_right_preserved) = on_lr_is_preserved(join.join_type);
451 for on in on_filter {
452 if on_left_preserved && checker.is_left_only(&on) {
453 left_push.push(on)
454 } else if on_right_preserved && checker.is_right_only(&on) {
455 right_push.push(on)
456 } else {
457 on_filter_join_conditions.push(on)
458 }
459 }
460
461 if left_preserved {
464 left_push.extend(extract_or_clauses_for_join(
465 &keep_predicates,
466 &left_schema_columns,
467 ));
468 left_push.extend(extract_or_clauses_for_join(
469 &join_conditions,
470 &left_schema_columns,
471 ));
472 }
473 if right_preserved {
474 right_push.extend(extract_or_clauses_for_join(
475 &keep_predicates,
476 &right_schema_columns,
477 ));
478 right_push.extend(extract_or_clauses_for_join(
479 &join_conditions,
480 &right_schema_columns,
481 ));
482 }
483
484 if on_left_preserved {
487 left_push.extend(extract_or_clauses_for_join(
488 &on_filter_join_conditions,
489 &left_schema_columns,
490 ));
491 }
492 if on_right_preserved {
493 right_push.extend(extract_or_clauses_for_join(
494 &on_filter_join_conditions,
495 &right_schema_columns,
496 ));
497 }
498
499 let join_conditions_empty = join_conditions.is_empty();
501 join_conditions.extend(on_filter_join_conditions);
502 join.filter = conjunction(join_conditions);
503
504 if join_conditions_empty && left_push.is_empty() && right_push.is_empty() {
505 return Ok(Transformed::no(with_filters(
507 keep_predicates,
508 LogicalPlan::Join(join),
509 )));
510 }
511
512 if let Some(predicate) = conjunction(left_push) {
513 join.left = Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.left)));
514 }
515
516 if let Some(predicate) = conjunction(right_push) {
517 join.right = Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.right)));
518 }
519
520 Ok(Transformed::yes(with_filters(
522 keep_predicates,
523 LogicalPlan::Join(join),
524 )))
525}
526
527fn push_down_join(
528 mut join: Join,
529 parent_predicate: Option<Expr>,
530) -> Result<Transformed<LogicalPlan>> {
531 let predicates = parent_predicate.map_or_else(Vec::new, split_conjunction_owned);
533
534 let on_filters = join
536 .filter
537 .take()
538 .map_or_else(Vec::new, split_conjunction_owned);
539
540 let inferred_join_predicates = with_debug_timing("infer_join_predicates", || {
542 infer_join_predicates(&join, &predicates, &on_filters)
543 })?;
544
545 if log_enabled!(Level::Debug) {
546 debug!(
547 "push_down_filter: join_type={:?}, parent_predicates={}, on_filters={}, inferred_join_predicates={}",
548 join.join_type,
549 predicates.len(),
550 on_filters.len(),
551 inferred_join_predicates.len()
552 );
553 }
554
555 if on_filters.is_empty()
556 && predicates.is_empty()
557 && inferred_join_predicates.is_empty()
558 {
559 return Ok(Transformed::no(LogicalPlan::Join(join)));
560 }
561
562 push_down_all_join(predicates, inferred_join_predicates, join, on_filters)
563}
564
565fn infer_join_predicates(
575 join: &Join,
576 predicates: &[Expr],
577 on_filters: &[Expr],
578) -> Result<Vec<Expr>> {
579 if join.null_aware {
587 return Ok(vec![]);
588 }
589
590 let join_col_keys = join
592 .on
593 .iter()
594 .filter_map(|(l, r)| {
595 let left_col = l.try_as_col()?;
596 let right_col = r.try_as_col()?;
597 Some((left_col, right_col))
598 })
599 .collect::<Vec<_>>();
600
601 let join_type = join.join_type;
602
603 let mut inferred_predicates = InferredPredicates::new(join_type);
604
605 infer_join_predicates_from_predicates(
606 &join_col_keys,
607 predicates,
608 &mut inferred_predicates,
609 )?;
610
611 infer_join_predicates_from_on_filters(
612 &join_col_keys,
613 join_type,
614 on_filters,
615 &mut inferred_predicates,
616 )?;
617
618 Ok(inferred_predicates.predicates)
619}
620
621struct InferredPredicates {
630 predicates: Vec<Expr>,
631 is_inner_join: bool,
632}
633
634impl InferredPredicates {
635 fn new(join_type: JoinType) -> Self {
636 Self {
637 predicates: vec![],
638 is_inner_join: join_type == JoinType::Inner,
639 }
640 }
641
642 fn try_build_predicate(
643 &mut self,
644 predicate: Expr,
645 replace_map: &HashMap<&Column, &Column>,
646 ) -> Result<()> {
647 if self.is_inner_join
648 || matches!(
649 is_restrict_null_predicate(
650 predicate.clone(),
651 replace_map.keys().cloned()
652 ),
653 Ok(true)
654 )
655 {
656 self.predicates.push(replace_col(predicate, replace_map)?);
657 }
658
659 Ok(())
660 }
661}
662
663fn infer_join_predicates_from_predicates(
672 join_col_keys: &[(&Column, &Column)],
673 predicates: &[Expr],
674 inferred_predicates: &mut InferredPredicates,
675) -> Result<()> {
676 infer_join_predicates_impl::<true, true>(
677 join_col_keys,
678 predicates,
679 inferred_predicates,
680 )
681}
682
683fn infer_join_predicates_from_on_filters(
695 join_col_keys: &[(&Column, &Column)],
696 join_type: JoinType,
697 on_filters: &[Expr],
698 inferred_predicates: &mut InferredPredicates,
699) -> Result<()> {
700 match join_type {
701 JoinType::Full | JoinType::LeftAnti | JoinType::RightAnti => Ok(()),
702 JoinType::Inner => infer_join_predicates_impl::<true, true>(
703 join_col_keys,
704 on_filters,
705 inferred_predicates,
706 ),
707 JoinType::Left | JoinType::LeftSemi | JoinType::LeftMark => {
708 infer_join_predicates_impl::<true, false>(
709 join_col_keys,
710 on_filters,
711 inferred_predicates,
712 )
713 }
714 JoinType::Right | JoinType::RightSemi | JoinType::RightMark => {
715 infer_join_predicates_impl::<false, true>(
716 join_col_keys,
717 on_filters,
718 inferred_predicates,
719 )
720 }
721 }
722}
723
724fn infer_join_predicates_impl<
740 const ENABLE_LEFT_TO_RIGHT: bool,
741 const ENABLE_RIGHT_TO_LEFT: bool,
742>(
743 join_col_keys: &[(&Column, &Column)],
744 input_predicates: &[Expr],
745 inferred_predicates: &mut InferredPredicates,
746) -> Result<()> {
747 for predicate in input_predicates {
748 let mut join_cols_to_replace = HashMap::new();
749
750 for &col in &predicate.column_refs() {
751 for (l, r) in join_col_keys.iter() {
752 if ENABLE_LEFT_TO_RIGHT && col == *l {
753 join_cols_to_replace.insert(col, *r);
754 break;
755 }
756 if ENABLE_RIGHT_TO_LEFT && col == *r {
757 join_cols_to_replace.insert(col, *l);
758 break;
759 }
760 }
761 }
762 if join_cols_to_replace.is_empty() {
763 continue;
764 }
765
766 inferred_predicates
767 .try_build_predicate(predicate.clone(), &join_cols_to_replace)?;
768 }
769 Ok(())
770}
771
772impl OptimizerRule for PushDownFilter {
773 fn name(&self) -> &str {
774 "push_down_filter"
775 }
776
777 fn apply_order(&self) -> Option<ApplyOrder> {
778 Some(ApplyOrder::TopDown)
779 }
780
781 fn supports_rewrite(&self) -> bool {
782 true
783 }
784
785 fn rewrite(
786 &self,
787 plan: LogicalPlan,
788 config: &dyn OptimizerConfig,
789 ) -> Result<Transformed<LogicalPlan>> {
790 let _ = config;
791 if let LogicalPlan::Join(join) = plan {
792 return push_down_join(join, None);
793 };
794
795 let LogicalPlan::Filter(mut filter) = plan else {
796 return Ok(Transformed::no(plan));
797 };
798
799 let predicate = split_conjunction_owned(filter.predicate.clone());
800 let old_predicate_len = predicate.len();
801 let new_predicates =
802 with_debug_timing("simplify_predicates", || simplify_predicates(predicate))?;
803
804 if log_enabled!(Level::Debug) {
805 debug!(
806 "push_down_filter: simplify_predicates old_count={}, new_count={}",
807 old_predicate_len,
808 new_predicates.len()
809 );
810 }
811
812 let (new_predicates, reorder_changed) = reorder_predicates(new_predicates);
816
817 let count_changed = old_predicate_len != new_predicates.len();
818 if count_changed || reorder_changed {
819 let Some(new_predicate) = conjunction(new_predicates) else {
820 return Ok(Transformed::yes(Arc::unwrap_or_clone(filter.input)));
823 };
824 filter.predicate = new_predicate;
825 }
826
827 if filter.input.fetch()?.is_some() || filter.input.skip()?.is_some() {
831 return Ok(Transformed::no(LogicalPlan::Filter(filter)));
832 }
833
834 match Arc::unwrap_or_clone(filter.input) {
835 LogicalPlan::Filter(mut child_filter) => {
836 let new_predicates: IndexSet<Expr> =
839 split_conjunction_owned(child_filter.predicate)
840 .into_iter()
841 .chain(split_conjunction_owned(filter.predicate))
842 .collect();
843
844 let Some(new_predicate) = conjunction(new_predicates) else {
845 return plan_err!("at least one expression exists");
846 };
847
848 child_filter.predicate = new_predicate;
849 self.rewrite(LogicalPlan::Filter(child_filter), config)
850 }
851 LogicalPlan::Repartition(mut repartition) => {
852 filter.input = repartition.input;
853 repartition.input = Arc::new(LogicalPlan::Filter(filter));
854 Ok(Transformed::yes(LogicalPlan::Repartition(repartition)))
855 }
856 LogicalPlan::Distinct(distinct) => {
857 let distinct = match distinct {
858 Distinct::All(input) => {
859 filter.input = input;
860 Distinct::All(Arc::new(LogicalPlan::Filter(filter)))
861 }
862 Distinct::On(mut distinct) => {
863 filter.input = distinct.input;
864 distinct.input = Arc::new(LogicalPlan::Filter(filter));
865 Distinct::On(distinct)
866 }
867 };
868
869 Ok(Transformed::yes(LogicalPlan::Distinct(distinct)))
870 }
871 LogicalPlan::Sort(mut sort) => {
872 filter.input = sort.input;
873 sort.input = Arc::new(LogicalPlan::Filter(filter));
874 Ok(Transformed::yes(LogicalPlan::Sort(sort)))
875 }
876 LogicalPlan::SubqueryAlias(mut subquery_alias) => {
877 let mut replace_map = HashMap::new();
878 for (i, (qualifier, field)) in
879 subquery_alias.input.schema().iter().enumerate()
880 {
881 let (sub_qualifier, sub_field) =
882 subquery_alias.schema.qualified_field(i);
883 replace_map.insert(
884 qualified_name(sub_qualifier, sub_field.name()),
885 Expr::Column(Column::new(qualifier.cloned(), field.name())),
886 );
887 }
888
889 filter.predicate = replace_cols_by_name(filter.predicate, &replace_map)?;
890 filter.input = subquery_alias.input;
891 subquery_alias.input = Arc::new(LogicalPlan::Filter(filter));
892 Ok(Transformed::yes(LogicalPlan::SubqueryAlias(subquery_alias)))
893 }
894 LogicalPlan::Projection(projection) => {
895 let predicates = split_conjunction_owned(filter.predicate.clone());
896 let (mut result, keep_predicates) =
897 rewrite_projection(predicates, projection)?;
898 if result.transformed {
899 result.data = with_filters(keep_predicates, result.data)
900 } else {
901 filter.input = Arc::new(result.data);
902 result.data = LogicalPlan::Filter(filter)
903 }
904
905 Ok(result)
906 }
907 LogicalPlan::Unnest(mut unnest) => {
908 let predicates = split_conjunction_owned(filter.predicate.clone());
909 let mut non_unnest_predicates = vec![];
910 let mut unnest_predicates = vec![];
911 let mut unnest_struct_columns = vec![];
912
913 for idx in &unnest.struct_type_columns {
914 let (sub_qualifier, field) =
915 unnest.input.schema().qualified_field(*idx);
916 if let DataType::Struct(children) = field.data_type() {
917 let field_name = field.name();
918 for child in children {
919 let child_name = child.name();
920 unnest_struct_columns.push(Column::new(
921 sub_qualifier.cloned(),
922 format!("{field_name}.{child_name}"),
923 ));
924 }
925 }
926 }
927
928 for predicate in predicates {
929 let mut accum: HashSet<Column> = HashSet::new();
931 expr_to_columns(&predicate, &mut accum)?;
932
933 let contains_list_columns =
934 unnest.list_type_columns.iter().any(|(_, unnest_list)| {
935 accum.contains(&unnest_list.output_column)
936 });
937 let contains_struct_columns =
938 unnest_struct_columns.iter().any(|c| accum.contains(c));
939
940 if contains_list_columns || contains_struct_columns {
941 unnest_predicates.push(predicate);
942 } else {
943 non_unnest_predicates.push(predicate);
944 }
945 }
946
947 if non_unnest_predicates.is_empty() {
950 filter.input = Arc::new(LogicalPlan::Unnest(unnest));
951 return Ok(Transformed::no(LogicalPlan::Filter(filter)));
952 }
953
954 filter.predicate = conjunction(non_unnest_predicates).unwrap();
964 filter.input = unnest.input;
965 unnest.input = Arc::new(LogicalPlan::Filter(filter));
969 Ok(Transformed::yes(with_filters(
970 unnest_predicates,
971 LogicalPlan::Unnest(unnest),
972 )))
973 }
974 LogicalPlan::Union(mut union) => {
975 let mut inputs = Vec::with_capacity(union.inputs.len());
976 for input in union.inputs {
977 let mut replace_map = HashMap::new();
978 for (i, (qualifier, field)) in input.schema().iter().enumerate() {
979 let (union_qualifier, union_field) =
980 union.schema.qualified_field(i);
981 replace_map.insert(
982 qualified_name(union_qualifier, union_field.name()),
983 Expr::Column(Column::new(qualifier.cloned(), field.name())),
984 );
985 }
986
987 let push_predicate =
988 replace_cols_by_name(filter.predicate.clone(), &replace_map)?;
989 inputs.push(Arc::new(LogicalPlan::Filter(Filter::new(
990 push_predicate,
991 input,
992 ))))
993 }
994
995 union.inputs = inputs;
996 Ok(Transformed::yes(LogicalPlan::Union(union)))
997 }
998 LogicalPlan::Aggregate(mut agg) => {
999 let group_expr_columns = expr_columns(&agg.group_expr);
1001
1002 let mut replace_map = HashMap::new();
1006 for expr in &agg.group_expr {
1007 replace_map.insert(expr.schema_name().to_string(), unalias(expr));
1008 }
1009
1010 let predicates = split_conjunction_owned(filter.predicate);
1011 let mut keep_predicates = vec![];
1012 let mut push_predicates = vec![];
1013 for expr in predicates {
1014 let cols = expr.column_refs();
1015 if cols.iter().all(|c| group_expr_columns.contains(c)) {
1016 push_predicates.push(replace_cols_by_name(expr, &replace_map)?);
1017 } else {
1018 keep_predicates.push(expr);
1019 }
1020 }
1021
1022 let result = if let Some(predicate) = conjunction(push_predicates) {
1024 filter.predicate = predicate;
1025 filter.input = agg.input;
1026 agg.input = Arc::new(LogicalPlan::Filter(filter));
1027 Transformed::yes(LogicalPlan::Aggregate(agg))
1028 } else {
1029 Transformed::no(LogicalPlan::Aggregate(agg))
1030 };
1031
1032 result.map_data(|plan| Ok(with_filters(keep_predicates, plan)))
1034 }
1035 LogicalPlan::Window(mut window) => {
1046 fn extract_partition_keys(func: &WindowFunction) -> HashSet<Column> {
1052 expr_columns(&func.params.partition_by)
1053 }
1054
1055 let potential_partition_keys = window
1056 .window_expr
1057 .iter()
1058 .map(|e| {
1059 match e {
1060 Expr::WindowFunction(window_func) => {
1061 extract_partition_keys(window_func)
1062 }
1063 Expr::Alias(alias) => {
1064 if let Expr::WindowFunction(window_func) =
1065 alias.expr.as_ref()
1066 {
1067 extract_partition_keys(window_func)
1068 } else {
1069 unreachable!()
1071 }
1072 }
1073 _ => {
1074 unreachable!()
1076 }
1077 }
1078 })
1079 .reduce(|a, b| &a & &b)
1082 .unwrap_or_default();
1083
1084 let predicates = split_conjunction_owned(filter.predicate);
1085 let mut keep_predicates = vec![];
1086 let mut push_predicates = vec![];
1087 for expr in predicates {
1088 let cols = expr.column_refs();
1089 if cols.iter().all(|c| potential_partition_keys.contains(c)) {
1090 push_predicates.push(expr);
1091 } else {
1092 keep_predicates.push(expr);
1093 }
1094 }
1095
1096 let result = if let Some(predicate) = conjunction(push_predicates) {
1106 filter.predicate = predicate;
1107 filter.input = window.input;
1108 window.input = Arc::new(LogicalPlan::Filter(filter));
1109 Transformed::yes(LogicalPlan::Window(window))
1110 } else {
1111 Transformed::no(LogicalPlan::Window(window))
1112 };
1113
1114 result.map_data(|plan| Ok(with_filters(keep_predicates, plan)))
1116 }
1117 LogicalPlan::Join(join) => push_down_join(join, Some(filter.predicate)),
1118 LogicalPlan::TableScan(mut scan) => {
1119 let filter_predicates = split_conjunction(&filter.predicate);
1120 let (subquery_filters, pushdown_candidates): (Vec<&Expr>, Vec<&Expr>) =
1124 filter_predicates
1125 .into_iter()
1126 .partition(|pred| pred.contains_scalar_subquery());
1127
1128 let (volatile_filters, non_volatile_filters): (Vec<&Expr>, Vec<&Expr>) =
1129 pushdown_candidates
1130 .into_iter()
1131 .partition(|pred| pred.is_volatile());
1132
1133 let supported_filters = scan
1135 .source
1136 .supports_filters_pushdown(non_volatile_filters.as_slice())?;
1137 assert_eq_or_internal_err!(
1138 non_volatile_filters.len(),
1139 supported_filters.len(),
1140 "Vec returned length: {} from supports_filters_pushdown is not the same size as the filters passed, which length is: {}",
1141 supported_filters.len(),
1142 non_volatile_filters.len()
1143 );
1144
1145 if supported_filters
1146 .iter()
1147 .all(|res| res == &TableProviderFilterPushDown::Unsupported)
1148 {
1149 filter.input = Arc::new(LogicalPlan::TableScan(scan));
1150 return Ok(Transformed::no(LogicalPlan::Filter(filter)));
1151 }
1152
1153 let zip = non_volatile_filters.iter().zip(supported_filters.iter());
1155
1156 let new_scan_filters = zip
1157 .clone()
1158 .filter(|(_, res)| *res != &TableProviderFilterPushDown::Unsupported)
1159 .map(|(&pred, _)| pred);
1160
1161 let new_scan_filters: Vec<Expr> = scan
1163 .filters
1164 .iter()
1165 .chain(new_scan_filters)
1166 .unique()
1167 .cloned()
1168 .collect();
1169
1170 if supported_filters
1171 .iter()
1172 .all(|res| res == &TableProviderFilterPushDown::Inexact)
1173 && scan.filters == new_scan_filters
1174 {
1175 filter.input = Arc::new(LogicalPlan::TableScan(scan));
1176 return Ok(Transformed::no(LogicalPlan::Filter(filter)));
1177 } else {
1178 scan.filters = new_scan_filters;
1179 }
1180
1181 let new_predicate: Vec<Expr> = zip
1184 .filter(|(_, res)| *res != &TableProviderFilterPushDown::Exact)
1185 .map(|(&pred, _)| pred)
1186 .chain(volatile_filters)
1187 .chain(subquery_filters)
1188 .cloned()
1189 .collect();
1190
1191 Ok(Transformed::yes(with_filters(
1192 new_predicate,
1193 LogicalPlan::TableScan(scan),
1194 )))
1195 }
1196 LogicalPlan::Extension(extension_plan) => {
1197 if extension_plan.node.inputs().is_empty() {
1200 filter.input = Arc::new(LogicalPlan::Extension(extension_plan));
1201 return Ok(Transformed::no(LogicalPlan::Filter(filter)));
1202 }
1203 let prevent_cols =
1204 extension_plan.node.prevent_predicate_push_down_columns();
1205
1206 let predicate_push_or_keep: Vec<bool> =
1210 split_conjunction(&filter.predicate)
1211 .iter()
1212 .map(|expr| {
1213 !expr
1214 .column_refs()
1215 .iter()
1216 .any(|c| prevent_cols.contains(&c.name))
1217 })
1218 .collect();
1219
1220 if predicate_push_or_keep.iter().all(|&x| !x) {
1222 filter.input = Arc::new(LogicalPlan::Extension(extension_plan));
1223 return Ok(Transformed::no(LogicalPlan::Filter(filter)));
1224 }
1225
1226 let mut keep_predicates = vec![];
1228 let mut push_predicates = vec![];
1229 for (push, expr) in predicate_push_or_keep
1230 .into_iter()
1231 .zip(split_conjunction_owned(filter.predicate))
1232 {
1233 if !push {
1234 keep_predicates.push(expr);
1235 } else {
1236 push_predicates.push(expr);
1237 }
1238 }
1239
1240 let predicate = conjunction(push_predicates).unwrap();
1242 let new_children = extension_plan
1243 .node
1244 .inputs()
1245 .into_iter()
1246 .map(|child| {
1247 LogicalPlan::Filter(Filter::new(
1248 predicate.clone(),
1249 Arc::new(child.clone()),
1250 ))
1251 })
1252 .collect();
1253
1254 let extension = LogicalPlan::Extension(extension_plan);
1256 let new_plan =
1257 extension.with_new_exprs(extension.expressions(), new_children)?;
1258 Ok(Transformed::yes(with_filters(keep_predicates, new_plan)))
1259 }
1260 child => {
1261 filter.input = Arc::new(child);
1262 Ok(Transformed::no(LogicalPlan::Filter(filter)))
1263 }
1264 }
1265 }
1266}
1267
1268fn rewrite_projection(
1296 predicates: Vec<Expr>,
1297 mut projection: Projection,
1298) -> Result<(Transformed<LogicalPlan>, Vec<Expr>)> {
1299 let (non_pushable_map, pushable_map) = projection
1306 .schema
1307 .iter()
1308 .zip(projection.expr.iter())
1309 .map(|((qualifier, field), expr)| {
1310 (qualified_name(qualifier, field.name()), unalias(expr))
1311 })
1312 .partition(|(_, value)| {
1313 value.is_volatile()
1314 || value.placement() == ExpressionPlacement::MoveTowardsLeafNodes
1315 });
1316
1317 let mut push_predicates = vec![];
1318 let mut keep_predicates = vec![];
1319 for expr in predicates {
1320 if contain(&expr, &non_pushable_map) {
1321 keep_predicates.push(expr);
1322 } else {
1323 push_predicates.push(expr);
1324 }
1325 }
1326
1327 let projection = if let Some(expr) = conjunction(push_predicates) {
1328 projection.input = Arc::new(LogicalPlan::Filter(Filter::new(
1331 replace_cols_by_name(expr, &pushable_map)?,
1332 projection.input,
1333 )));
1334
1335 Transformed::yes(LogicalPlan::Projection(projection))
1336 } else {
1337 Transformed::no(LogicalPlan::Projection(projection))
1338 };
1339
1340 Ok((projection, keep_predicates))
1341}
1342
1343#[deprecated]
1347pub fn make_filter(predicate: Expr, input: Arc<LogicalPlan>) -> Result<LogicalPlan> {
1348 Filter::try_new(predicate, input).map(LogicalPlan::Filter)
1349}
1350
1351impl PushDownFilter {
1352 #[expect(missing_docs)]
1353 pub fn new() -> Self {
1354 Self {}
1355 }
1356}
1357
1358fn with_debug_timing<T, F>(label: &'static str, f: F) -> Result<T>
1359where
1360 F: FnOnce() -> Result<T>,
1361{
1362 if !log_enabled!(Level::Debug) {
1363 return f();
1364 }
1365 let start = Instant::now();
1366 let result = f();
1367 debug!(
1368 "push_down_filter_timing: section={label}, elapsed_us={}",
1369 start.elapsed().as_micros()
1370 );
1371 result
1372}
1373
1374pub fn replace_cols_by_name(
1376 e: Expr,
1377 replace_map: &HashMap<String, impl AsRef<Expr>>,
1378) -> Result<Expr> {
1379 e.transform_up(|expr| {
1380 if let Expr::Column(c) = &expr
1381 && let Some(new_expr) = replace_map.get(&c.flat_name())
1382 {
1383 Ok(Transformed::yes(new_expr.as_ref().clone()))
1384 } else {
1385 Ok(Transformed::no(expr))
1386 }
1387 })
1388 .data()
1389}
1390
1391fn unalias(expr: &Expr) -> &Expr {
1393 if let Expr::Alias(alias) = expr {
1394 unalias(&alias.expr)
1395 } else {
1396 expr
1397 }
1398}
1399
1400fn contain<T>(e: &Expr, check_map: &HashMap<String, T>) -> bool {
1402 let mut is_contain = false;
1403 e.apply(|expr| {
1404 if let Expr::Column(c) = &expr
1405 && check_map.contains_key(&c.flat_name())
1406 {
1407 is_contain = true;
1408 Ok(TreeNodeRecursion::Stop)
1409 } else {
1410 Ok(TreeNodeRecursion::Continue)
1411 }
1412 })
1413 .unwrap();
1414 is_contain
1415}
1416
1417fn with_filters(predicates: Vec<Expr>, plan: LogicalPlan) -> LogicalPlan {
1418 if let Some(predicate) = conjunction(predicates) {
1419 LogicalPlan::Filter(Filter::new(predicate, Arc::new(plan)))
1420 } else {
1421 plan
1422 }
1423}
1424
1425fn expr_columns(exprs: &[Expr]) -> HashSet<Column> {
1426 exprs
1427 .iter()
1428 .map(|expr| {
1429 let (relation, name) = expr.qualified_name();
1430 Column::new(relation, name)
1431 })
1432 .collect()
1433}
1434
1435#[cfg(test)]
1436mod tests {
1437 use std::cmp::Ordering;
1438 use std::fmt::{Debug, Formatter};
1439
1440 use arrow::datatypes::{Field, Schema, SchemaRef};
1441 use async_trait::async_trait;
1442
1443 use datafusion_common::{DFSchemaRef, DataFusionError, ScalarValue};
1444 use datafusion_expr::expr::ScalarFunction;
1445 use datafusion_expr::logical_plan::table_scan;
1446 use datafusion_expr::{
1447 ColumnarValue, ExprFunctionExt, Extension, LogicalPlanBuilder,
1448 ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TableScan, TableSource,
1449 TableType, UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col,
1450 in_list, in_subquery, lit,
1451 };
1452
1453 use crate::OptimizerContext;
1454 use crate::assert_optimized_plan_eq_snapshot;
1455 use crate::optimizer::Optimizer;
1456 use crate::simplify_expressions::SimplifyExpressions;
1457 use crate::test::udfs::leaf_udf_expr;
1458 use crate::test::*;
1459 use datafusion_expr::test::function_stub::sum;
1460 use insta::assert_snapshot;
1461
1462 use super::*;
1463
1464 fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
1465
1466 macro_rules! assert_optimized_plan_equal {
1467 (
1468 $plan:expr,
1469 @ $expected:literal $(,)?
1470 ) => {{
1471 let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
1472 let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(PushDownFilter::new())];
1473 assert_optimized_plan_eq_snapshot!(
1474 optimizer_ctx,
1475 rules,
1476 $plan,
1477 @ $expected,
1478 )
1479 }};
1480 }
1481
1482 macro_rules! assert_optimized_plan_eq_with_rewrite_predicate {
1483 (
1484 $plan:expr,
1485 @ $expected:literal $(,)?
1486 ) => {{
1487 let optimizer = Optimizer::with_rules(vec![
1488 Arc::new(SimplifyExpressions::new()),
1489 Arc::new(PushDownFilter::new()),
1490 ]);
1491 let optimized_plan = optimizer.optimize($plan, &OptimizerContext::new(), observe)?;
1492 assert_snapshot!(optimized_plan, @ $expected);
1493 Ok::<(), DataFusionError>(())
1494 }};
1495 }
1496
1497 macro_rules! assert_plan_not_transformed {
1500 ($plan:expr) => {{
1501 let transformed = PushDownFilter::new()
1502 .rewrite($plan, &OptimizerContext::new())
1503 .expect("failed to optimize plan");
1504 assert!(!transformed.transformed);
1505 }};
1506 }
1507
1508 #[test]
1509 fn filter_before_projection() -> Result<()> {
1510 let table_scan = test_table_scan()?;
1511 let plan = LogicalPlanBuilder::from(table_scan)
1512 .project(vec![col("a"), col("b")])?
1513 .filter(col("a").eq(lit(1i64)))?
1514 .build()?;
1515 assert_optimized_plan_equal!(
1517 plan,
1518 @r"
1519 Projection: test.a, test.b
1520 TableScan: test, full_filters=[test.a = Int64(1)]
1521 "
1522 )
1523 }
1524
1525 #[test]
1526 fn filter_after_limit() -> Result<()> {
1527 let table_scan = test_table_scan()?;
1528 let plan = LogicalPlanBuilder::from(table_scan)
1529 .project(vec![col("a"), col("b")])?
1530 .limit(0, Some(10))?
1531 .filter(col("a").eq(lit(1i64)))?
1532 .build()?;
1533 assert_optimized_plan_equal!(
1535 plan,
1536 @r"
1537 Filter: test.a = Int64(1)
1538 Limit: skip=0, fetch=10
1539 Projection: test.a, test.b
1540 TableScan: test
1541 "
1542 )
1543 }
1544
1545 #[test]
1546 fn filter_no_columns() -> Result<()> {
1547 let table_scan = test_table_scan()?;
1548 let plan = LogicalPlanBuilder::from(table_scan)
1549 .filter(lit(0i64).eq(lit(1i64)))?
1550 .build()?;
1551 assert_optimized_plan_equal!(
1552 plan,
1553 @"TableScan: test, full_filters=[Int64(0) = Int64(1)]"
1554 )
1555 }
1556
1557 #[test]
1558 fn filter_jump_2_plans() -> Result<()> {
1559 let table_scan = test_table_scan()?;
1560 let plan = LogicalPlanBuilder::from(table_scan)
1561 .project(vec![col("a"), col("b"), col("c")])?
1562 .project(vec![col("c"), col("b")])?
1563 .filter(col("a").eq(lit(1i64)))?
1564 .build()?;
1565 assert_optimized_plan_equal!(
1567 plan,
1568 @r"
1569 Projection: test.c, test.b
1570 Projection: test.a, test.b, test.c
1571 TableScan: test, full_filters=[test.a = Int64(1)]
1572 "
1573 )
1574 }
1575
1576 #[test]
1577 fn filter_move_agg() -> Result<()> {
1578 let table_scan = test_table_scan()?;
1579 let plan = LogicalPlanBuilder::from(table_scan)
1580 .aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
1581 .filter(col("a").gt(lit(10i64)))?
1582 .build()?;
1583 assert_optimized_plan_equal!(
1585 plan,
1586 @r"
1587 Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b) AS total_salary]]
1588 TableScan: test, full_filters=[test.a > Int64(10)]
1589 "
1590 )
1591 }
1592
1593 #[test]
1595 fn filter_move_agg_special() -> Result<()> {
1596 let schema = Schema::new(vec![
1597 Field::new("$a", DataType::UInt32, false),
1598 Field::new("$b", DataType::UInt32, false),
1599 Field::new("$c", DataType::UInt32, false),
1600 ]);
1601 let table_scan = table_scan(Some("test"), &schema, None)?.build()?;
1602
1603 let plan = LogicalPlanBuilder::from(table_scan)
1604 .aggregate(vec![col("$a")], vec![sum(col("$b")).alias("total_salary")])?
1605 .filter(col("$a").gt(lit(10i64)))?
1606 .build()?;
1607 assert_optimized_plan_equal!(
1609 plan,
1610 @r"
1611 Aggregate: groupBy=[[test.$a]], aggr=[[sum(test.$b) AS total_salary]]
1612 TableScan: test, full_filters=[test.$a > Int64(10)]
1613 "
1614 )
1615 }
1616
1617 #[test]
1618 fn filter_complex_group_by() -> Result<()> {
1619 let table_scan = test_table_scan()?;
1620 let plan = LogicalPlanBuilder::from(table_scan)
1621 .aggregate(vec![add(col("b"), col("a"))], vec![sum(col("a")), col("b")])?
1622 .filter(col("b").gt(lit(10i64)))?
1623 .build()?;
1624 assert_optimized_plan_equal!(
1625 plan,
1626 @r"
1627 Filter: test.b > Int64(10)
1628 Aggregate: groupBy=[[test.b + test.a]], aggr=[[sum(test.a), test.b]]
1629 TableScan: test
1630 "
1631 )
1632 }
1633
1634 #[test]
1635 fn push_agg_need_replace_expr() -> Result<()> {
1636 let plan = LogicalPlanBuilder::from(test_table_scan()?)
1637 .aggregate(vec![add(col("b"), col("a"))], vec![sum(col("a")), col("b")])?
1638 .filter(col("test.b + test.a").gt(lit(10i64)))?
1639 .build()?;
1640 assert_optimized_plan_equal!(
1641 plan,
1642 @r"
1643 Aggregate: groupBy=[[test.b + test.a]], aggr=[[sum(test.a), test.b]]
1644 TableScan: test, full_filters=[test.b + test.a > Int64(10)]
1645 "
1646 )
1647 }
1648
1649 #[test]
1650 fn filter_keep_agg() -> Result<()> {
1651 let table_scan = test_table_scan()?;
1652 let plan = LogicalPlanBuilder::from(table_scan)
1653 .aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
1654 .filter(col("b").gt(lit(10i64)))?
1655 .build()?;
1656 assert_plan_not_transformed!(plan.clone());
1657
1658 assert_optimized_plan_equal!(
1660 plan,
1661 @r"
1662 Filter: b > Int64(10)
1663 Aggregate: groupBy=[[test.a]], aggr=[[sum(test.b) AS b]]
1664 TableScan: test
1665 "
1666 )
1667 }
1668
1669 #[test]
1671 fn filter_move_window() -> Result<()> {
1672 let table_scan = test_table_scan()?;
1673
1674 let window = Expr::from(WindowFunction::new(
1675 WindowFunctionDefinition::WindowUDF(
1676 datafusion_functions_window::rank::rank_udwf(),
1677 ),
1678 vec![],
1679 ))
1680 .partition_by(vec![col("a"), col("b")])
1681 .order_by(vec![col("c").sort(true, true)])
1682 .build()
1683 .unwrap();
1684
1685 let plan = LogicalPlanBuilder::from(table_scan)
1686 .window(vec![window])?
1687 .filter(col("b").gt(lit(10i64)))?
1688 .build()?;
1689
1690 assert_optimized_plan_equal!(
1691 plan,
1692 @r"
1693 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a, test.b] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1694 TableScan: test, full_filters=[test.b > Int64(10)]
1695 "
1696 )
1697 }
1698
1699 #[test]
1701 fn filter_window_special_identifier() -> Result<()> {
1702 let schema = Schema::new(vec![
1703 Field::new("$a", DataType::UInt32, false),
1704 Field::new("$b", DataType::UInt32, false),
1705 Field::new("$c", DataType::UInt32, false),
1706 ]);
1707 let table_scan = table_scan(Some("test"), &schema, None)?.build()?;
1708
1709 let window = Expr::from(WindowFunction::new(
1710 WindowFunctionDefinition::WindowUDF(
1711 datafusion_functions_window::rank::rank_udwf(),
1712 ),
1713 vec![],
1714 ))
1715 .partition_by(vec![col("$a"), col("$b")])
1716 .order_by(vec![col("$c").sort(true, true)])
1717 .build()
1718 .unwrap();
1719
1720 let plan = LogicalPlanBuilder::from(table_scan)
1721 .window(vec![window])?
1722 .filter(col("$b").gt(lit(10i64)))?
1723 .build()?;
1724
1725 assert_optimized_plan_equal!(
1726 plan,
1727 @r"
1728 WindowAggr: windowExpr=[[rank() PARTITION BY [test.$a, test.$b] ORDER BY [test.$c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1729 TableScan: test, full_filters=[test.$b > Int64(10)]
1730 "
1731 )
1732 }
1733
1734 #[test]
1737 fn filter_move_complex_window() -> Result<()> {
1738 let table_scan = test_table_scan()?;
1739
1740 let window = Expr::from(WindowFunction::new(
1741 WindowFunctionDefinition::WindowUDF(
1742 datafusion_functions_window::rank::rank_udwf(),
1743 ),
1744 vec![],
1745 ))
1746 .partition_by(vec![col("a"), col("b")])
1747 .order_by(vec![col("c").sort(true, true)])
1748 .build()
1749 .unwrap();
1750
1751 let plan = LogicalPlanBuilder::from(table_scan)
1752 .window(vec![window])?
1753 .filter(and(col("a").gt(lit(10i64)), col("b").eq(lit(1i64))))?
1754 .build()?;
1755
1756 assert_optimized_plan_equal!(
1757 plan,
1758 @r"
1759 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a, test.b] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1760 TableScan: test, full_filters=[test.a > Int64(10), test.b = Int64(1)]
1761 "
1762 )
1763 }
1764
1765 #[test]
1767 fn filter_move_partial_window() -> Result<()> {
1768 let table_scan = test_table_scan()?;
1769
1770 let window = Expr::from(WindowFunction::new(
1771 WindowFunctionDefinition::WindowUDF(
1772 datafusion_functions_window::rank::rank_udwf(),
1773 ),
1774 vec![],
1775 ))
1776 .partition_by(vec![col("a")])
1777 .order_by(vec![col("c").sort(true, true)])
1778 .build()
1779 .unwrap();
1780
1781 let plan = LogicalPlanBuilder::from(table_scan)
1782 .window(vec![window])?
1783 .filter(and(col("a").gt(lit(10i64)), col("b").eq(lit(1i64))))?
1784 .build()?;
1785
1786 assert_optimized_plan_equal!(
1787 plan,
1788 @r"
1789 Filter: test.b = Int64(1)
1790 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1791 TableScan: test, full_filters=[test.a > Int64(10)]
1792 "
1793 )
1794 }
1795
1796 #[test]
1799 fn filter_expression_keep_window() -> Result<()> {
1800 let table_scan = test_table_scan()?;
1801
1802 let window = Expr::from(WindowFunction::new(
1803 WindowFunctionDefinition::WindowUDF(
1804 datafusion_functions_window::rank::rank_udwf(),
1805 ),
1806 vec![],
1807 ))
1808 .partition_by(vec![add(col("a"), col("b"))]) .order_by(vec![col("c").sort(true, true)])
1810 .build()
1811 .unwrap();
1812
1813 let plan = LogicalPlanBuilder::from(table_scan)
1814 .window(vec![window])?
1815 .filter(add(col("a"), col("b")).gt(lit(10i64)))?
1818 .build()?;
1819
1820 assert_optimized_plan_equal!(
1821 plan,
1822 @r"
1823 Filter: test.a + test.b > Int64(10)
1824 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1825 TableScan: test
1826 "
1827 )
1828 }
1829
1830 #[test]
1832 fn filter_order_keep_window() -> Result<()> {
1833 let table_scan = test_table_scan()?;
1834
1835 let window = Expr::from(WindowFunction::new(
1836 WindowFunctionDefinition::WindowUDF(
1837 datafusion_functions_window::rank::rank_udwf(),
1838 ),
1839 vec![],
1840 ))
1841 .partition_by(vec![col("a")])
1842 .order_by(vec![col("c").sort(true, true)])
1843 .build()
1844 .unwrap();
1845
1846 let plan = LogicalPlanBuilder::from(table_scan)
1847 .window(vec![window])?
1848 .filter(col("c").gt(lit(10i64)))?
1849 .build()?;
1850 assert_plan_not_transformed!(plan.clone());
1851
1852 assert_optimized_plan_equal!(
1853 plan,
1854 @r"
1855 Filter: test.c > Int64(10)
1856 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1857 TableScan: test
1858 "
1859 )
1860 }
1861
1862 #[test]
1865 fn filter_multiple_windows_common_partitions() -> Result<()> {
1866 let table_scan = test_table_scan()?;
1867
1868 let window1 = Expr::from(WindowFunction::new(
1869 WindowFunctionDefinition::WindowUDF(
1870 datafusion_functions_window::rank::rank_udwf(),
1871 ),
1872 vec![],
1873 ))
1874 .partition_by(vec![col("a")])
1875 .order_by(vec![col("c").sort(true, true)])
1876 .build()
1877 .unwrap();
1878
1879 let window2 = Expr::from(WindowFunction::new(
1880 WindowFunctionDefinition::WindowUDF(
1881 datafusion_functions_window::rank::rank_udwf(),
1882 ),
1883 vec![],
1884 ))
1885 .partition_by(vec![col("b"), col("a")])
1886 .order_by(vec![col("c").sort(true, true)])
1887 .build()
1888 .unwrap();
1889
1890 let plan = LogicalPlanBuilder::from(table_scan)
1891 .window(vec![window1, window2])?
1892 .filter(col("a").gt(lit(10i64)))? .build()?;
1894
1895 assert_optimized_plan_equal!(
1896 plan,
1897 @r"
1898 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, rank() PARTITION BY [test.b, test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1899 TableScan: test, full_filters=[test.a > Int64(10)]
1900 "
1901 )
1902 }
1903
1904 #[test]
1907 fn filter_multiple_windows_disjoint_partitions() -> Result<()> {
1908 let table_scan = test_table_scan()?;
1909
1910 let window1 = Expr::from(WindowFunction::new(
1911 WindowFunctionDefinition::WindowUDF(
1912 datafusion_functions_window::rank::rank_udwf(),
1913 ),
1914 vec![],
1915 ))
1916 .partition_by(vec![col("a")])
1917 .order_by(vec![col("c").sort(true, true)])
1918 .build()
1919 .unwrap();
1920
1921 let window2 = Expr::from(WindowFunction::new(
1922 WindowFunctionDefinition::WindowUDF(
1923 datafusion_functions_window::rank::rank_udwf(),
1924 ),
1925 vec![],
1926 ))
1927 .partition_by(vec![col("b"), col("a")])
1928 .order_by(vec![col("c").sort(true, true)])
1929 .build()
1930 .unwrap();
1931
1932 let plan = LogicalPlanBuilder::from(table_scan)
1933 .window(vec![window1, window2])?
1934 .filter(col("b").gt(lit(10i64)))? .build()?;
1936
1937 assert_optimized_plan_equal!(
1938 plan,
1939 @r"
1940 Filter: test.b > Int64(10)
1941 WindowAggr: windowExpr=[[rank() PARTITION BY [test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, rank() PARTITION BY [test.b, test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
1942 TableScan: test
1943 "
1944 )
1945 }
1946
1947 #[test]
1949 fn alias() -> Result<()> {
1950 let table_scan = test_table_scan()?;
1951 let plan = LogicalPlanBuilder::from(table_scan)
1952 .project(vec![col("a").alias("b"), col("c")])?
1953 .filter(col("b").eq(lit(1i64)))?
1954 .build()?;
1955 assert_optimized_plan_equal!(
1957 plan,
1958 @r"
1959 Projection: test.a AS b, test.c
1960 TableScan: test, full_filters=[test.a = Int64(1)]
1961 "
1962 )
1963 }
1964
1965 fn add(left: Expr, right: Expr) -> Expr {
1966 Expr::BinaryExpr(BinaryExpr::new(
1967 Box::new(left),
1968 Operator::Plus,
1969 Box::new(right),
1970 ))
1971 }
1972
1973 fn multiply(left: Expr, right: Expr) -> Expr {
1974 Expr::BinaryExpr(BinaryExpr::new(
1975 Box::new(left),
1976 Operator::Multiply,
1977 Box::new(right),
1978 ))
1979 }
1980
1981 #[test]
1983 fn complex_expression() -> Result<()> {
1984 let table_scan = test_table_scan()?;
1985 let plan = LogicalPlanBuilder::from(table_scan)
1986 .project(vec![
1987 add(multiply(col("a"), lit(2)), col("c")).alias("b"),
1988 col("c"),
1989 ])?
1990 .filter(col("b").eq(lit(1i64)))?
1991 .build()?;
1992
1993 assert_snapshot!(plan,
1995 @r"
1996 Filter: b = Int64(1)
1997 Projection: test.a * Int32(2) + test.c AS b, test.c
1998 TableScan: test
1999 ",
2000 );
2001 assert_optimized_plan_equal!(
2003 plan,
2004 @r"
2005 Projection: test.a * Int32(2) + test.c AS b, test.c
2006 TableScan: test, full_filters=[test.a * Int32(2) + test.c = Int64(1)]
2007 "
2008 )
2009 }
2010
2011 #[test]
2013 fn complex_plan() -> Result<()> {
2014 let table_scan = test_table_scan()?;
2015 let plan = LogicalPlanBuilder::from(table_scan)
2016 .project(vec![
2017 add(multiply(col("a"), lit(2)), col("c")).alias("b"),
2018 col("c"),
2019 ])?
2020 .project(vec![multiply(col("b"), lit(3)).alias("a"), col("c")])?
2022 .filter(col("a").eq(lit(1i64)))?
2023 .build()?;
2024
2025 assert_snapshot!(plan,
2027 @r"
2028 Filter: a = Int64(1)
2029 Projection: b * Int32(3) AS a, test.c
2030 Projection: test.a * Int32(2) + test.c AS b, test.c
2031 TableScan: test
2032 ",
2033 );
2034 assert_optimized_plan_equal!(
2036 plan,
2037 @r"
2038 Projection: b * Int32(3) AS a, test.c
2039 Projection: test.a * Int32(2) + test.c AS b, test.c
2040 TableScan: test, full_filters=[(test.a * Int32(2) + test.c) * Int32(3) = Int64(1)]
2041 "
2042 )
2043 }
2044
2045 #[derive(Debug, PartialEq, Eq, Hash)]
2046 struct NoopPlan {
2047 input: Vec<LogicalPlan>,
2048 schema: DFSchemaRef,
2049 }
2050
2051 impl PartialOrd for NoopPlan {
2053 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2054 self.input
2055 .partial_cmp(&other.input)
2056 .filter(|cmp| *cmp != Ordering::Equal || self == other)
2058 }
2059 }
2060
2061 impl UserDefinedLogicalNodeCore for NoopPlan {
2062 fn name(&self) -> &str {
2063 "NoopPlan"
2064 }
2065
2066 fn inputs(&self) -> Vec<&LogicalPlan> {
2067 self.input.iter().collect()
2068 }
2069
2070 fn schema(&self) -> &DFSchemaRef {
2071 &self.schema
2072 }
2073
2074 fn expressions(&self) -> Vec<Expr> {
2075 self.input
2076 .iter()
2077 .flat_map(|child| child.expressions())
2078 .collect()
2079 }
2080
2081 fn prevent_predicate_push_down_columns(&self) -> HashSet<String> {
2082 HashSet::from_iter(vec!["c".to_string()])
2083 }
2084
2085 fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
2086 write!(f, "NoopPlan")
2087 }
2088
2089 fn with_exprs_and_inputs(
2090 &self,
2091 _exprs: Vec<Expr>,
2092 inputs: Vec<LogicalPlan>,
2093 ) -> Result<Self> {
2094 Ok(Self {
2095 input: inputs,
2096 schema: Arc::clone(&self.schema),
2097 })
2098 }
2099
2100 fn supports_limit_pushdown(&self) -> bool {
2101 false }
2103 }
2104
2105 #[test]
2106 fn user_defined_plan() -> Result<()> {
2107 let table_scan = test_table_scan()?;
2108
2109 let custom_plan = LogicalPlan::Extension(Extension {
2110 node: Arc::new(NoopPlan {
2111 input: vec![table_scan.clone()],
2112 schema: Arc::clone(table_scan.schema()),
2113 }),
2114 });
2115 let plan = LogicalPlanBuilder::from(custom_plan)
2116 .filter(col("a").eq(lit(1i64)))?
2117 .build()?;
2118
2119 assert_optimized_plan_equal!(
2121 plan,
2122 @r"
2123 NoopPlan
2124 TableScan: test, full_filters=[test.a = Int64(1)]
2125 "
2126 )?;
2127
2128 let custom_plan = LogicalPlan::Extension(Extension {
2129 node: Arc::new(NoopPlan {
2130 input: vec![table_scan.clone()],
2131 schema: Arc::clone(table_scan.schema()),
2132 }),
2133 });
2134 let plan = LogicalPlanBuilder::from(custom_plan)
2135 .filter(col("a").eq(lit(1i64)).and(col("c").eq(lit(2i64))))?
2136 .build()?;
2137
2138 assert_optimized_plan_equal!(
2140 plan,
2141 @r"
2142 Filter: test.c = Int64(2)
2143 NoopPlan
2144 TableScan: test, full_filters=[test.a = Int64(1)]
2145 "
2146 )?;
2147
2148 let custom_plan = LogicalPlan::Extension(Extension {
2149 node: Arc::new(NoopPlan {
2150 input: vec![table_scan.clone(), table_scan.clone()],
2151 schema: Arc::clone(table_scan.schema()),
2152 }),
2153 });
2154 let plan = LogicalPlanBuilder::from(custom_plan)
2155 .filter(col("a").eq(lit(1i64)))?
2156 .build()?;
2157
2158 assert_optimized_plan_equal!(
2160 plan,
2161 @r"
2162 NoopPlan
2163 TableScan: test, full_filters=[test.a = Int64(1)]
2164 TableScan: test, full_filters=[test.a = Int64(1)]
2165 "
2166 )?;
2167
2168 let custom_plan = LogicalPlan::Extension(Extension {
2169 node: Arc::new(NoopPlan {
2170 input: vec![table_scan.clone(), table_scan.clone()],
2171 schema: Arc::clone(table_scan.schema()),
2172 }),
2173 });
2174 let plan = LogicalPlanBuilder::from(custom_plan)
2175 .filter(col("a").eq(lit(1i64)).and(col("c").eq(lit(2i64))))?
2176 .build()?;
2177
2178 assert_optimized_plan_equal!(
2180 plan,
2181 @r"
2182 Filter: test.c = Int64(2)
2183 NoopPlan
2184 TableScan: test, full_filters=[test.a = Int64(1)]
2185 TableScan: test, full_filters=[test.a = Int64(1)]
2186 "
2187 )
2188 }
2189
2190 #[test]
2193 fn multi_filter() -> Result<()> {
2194 let table_scan = test_table_scan()?;
2196 let plan = LogicalPlanBuilder::from(table_scan)
2197 .project(vec![col("a").alias("b"), col("c")])?
2198 .aggregate(vec![col("b")], vec![sum(col("c"))])?
2199 .filter(col("b").gt(lit(10i64)))?
2200 .filter(col("sum(test.c)").gt(lit(10i64)))?
2201 .build()?;
2202
2203 assert_snapshot!(plan,
2205 @r"
2206 Filter: sum(test.c) > Int64(10)
2207 Filter: b > Int64(10)
2208 Aggregate: groupBy=[[b]], aggr=[[sum(test.c)]]
2209 Projection: test.a AS b, test.c
2210 TableScan: test
2211 ",
2212 );
2213 assert_optimized_plan_equal!(
2215 plan,
2216 @r"
2217 Filter: sum(test.c) > Int64(10)
2218 Aggregate: groupBy=[[b]], aggr=[[sum(test.c)]]
2219 Projection: test.a AS b, test.c
2220 TableScan: test, full_filters=[test.a > Int64(10)]
2221 "
2222 )
2223 }
2224
2225 #[test]
2228 fn split_filter() -> Result<()> {
2229 let table_scan = test_table_scan()?;
2231 let plan = LogicalPlanBuilder::from(table_scan)
2232 .project(vec![col("a").alias("b"), col("c")])?
2233 .aggregate(vec![col("b")], vec![sum(col("c"))])?
2234 .filter(and(
2235 col("sum(test.c)").gt(lit(10i64)),
2236 and(col("b").gt(lit(10i64)), col("sum(test.c)").lt(lit(20i64))),
2237 ))?
2238 .build()?;
2239
2240 assert_snapshot!(plan,
2242 @r"
2243 Filter: sum(test.c) > Int64(10) AND b > Int64(10) AND sum(test.c) < Int64(20)
2244 Aggregate: groupBy=[[b]], aggr=[[sum(test.c)]]
2245 Projection: test.a AS b, test.c
2246 TableScan: test
2247 ",
2248 );
2249 assert_optimized_plan_equal!(
2251 plan,
2252 @r"
2253 Filter: sum(test.c) > Int64(10) AND sum(test.c) < Int64(20)
2254 Aggregate: groupBy=[[b]], aggr=[[sum(test.c)]]
2255 Projection: test.a AS b, test.c
2256 TableScan: test, full_filters=[test.a > Int64(10)]
2257 "
2258 )
2259 }
2260
2261 #[test]
2263 fn double_limit() -> Result<()> {
2264 let table_scan = test_table_scan()?;
2265 let plan = LogicalPlanBuilder::from(table_scan)
2266 .project(vec![col("a"), col("b")])?
2267 .limit(0, Some(20))?
2268 .limit(0, Some(10))?
2269 .project(vec![col("a"), col("b")])?
2270 .filter(col("a").eq(lit(1i64)))?
2271 .build()?;
2272 assert_optimized_plan_equal!(
2274 plan,
2275 @r"
2276 Projection: test.a, test.b
2277 Filter: test.a = Int64(1)
2278 Limit: skip=0, fetch=10
2279 Limit: skip=0, fetch=20
2280 Projection: test.a, test.b
2281 TableScan: test
2282 "
2283 )
2284 }
2285
2286 #[test]
2287 fn union_all() -> Result<()> {
2288 let table_scan = test_table_scan()?;
2289 let table_scan2 = test_table_scan_with_name("test2")?;
2290 let plan = LogicalPlanBuilder::from(table_scan)
2291 .union(LogicalPlanBuilder::from(table_scan2).build()?)?
2292 .filter(col("a").eq(lit(1i64)))?
2293 .build()?;
2294 assert_optimized_plan_equal!(
2296 plan,
2297 @r"
2298 Union
2299 TableScan: test, full_filters=[test.a = Int64(1)]
2300 TableScan: test2, full_filters=[test2.a = Int64(1)]
2301 "
2302 )
2303 }
2304
2305 #[test]
2306 fn union_all_on_projection() -> Result<()> {
2307 let table_scan = test_table_scan()?;
2308 let table = LogicalPlanBuilder::from(table_scan)
2309 .project(vec![col("a").alias("b")])?
2310 .alias("test2")?;
2311
2312 let plan = table
2313 .clone()
2314 .union(table.build()?)?
2315 .filter(col("b").eq(lit(1i64)))?
2316 .build()?;
2317
2318 assert_optimized_plan_equal!(
2320 plan,
2321 @r"
2322 Union
2323 SubqueryAlias: test2
2324 Projection: test.a AS b
2325 TableScan: test, full_filters=[test.a = Int64(1)]
2326 SubqueryAlias: test2
2327 Projection: test.a AS b
2328 TableScan: test, full_filters=[test.a = Int64(1)]
2329 "
2330 )
2331 }
2332
2333 #[test]
2334 fn test_union_different_schema() -> Result<()> {
2335 let left = LogicalPlanBuilder::from(test_table_scan()?)
2336 .project(vec![col("a"), col("b"), col("c")])?
2337 .build()?;
2338
2339 let schema = Schema::new(vec![
2340 Field::new("d", DataType::UInt32, false),
2341 Field::new("e", DataType::UInt32, false),
2342 Field::new("f", DataType::UInt32, false),
2343 ]);
2344 let right = table_scan(Some("test1"), &schema, None)?
2345 .project(vec![col("d"), col("e"), col("f")])?
2346 .build()?;
2347 let filter = and(col("test.a").eq(lit(1)), col("test1.d").gt(lit(2)));
2348 let plan = LogicalPlanBuilder::from(left)
2349 .cross_join(right)?
2350 .project(vec![col("test.a"), col("test1.d")])?
2351 .filter(filter)?
2352 .build()?;
2353
2354 assert_optimized_plan_equal!(
2355 plan,
2356 @r"
2357 Projection: test.a, test1.d
2358 Cross Join:
2359 Projection: test.a, test.b, test.c
2360 TableScan: test, full_filters=[test.a = Int32(1)]
2361 Projection: test1.d, test1.e, test1.f
2362 TableScan: test1, full_filters=[test1.d > Int32(2)]
2363 "
2364 )
2365 }
2366
2367 #[test]
2368 fn test_project_same_name_different_qualifier() -> Result<()> {
2369 let table_scan = test_table_scan()?;
2370 let left = LogicalPlanBuilder::from(table_scan)
2371 .project(vec![col("a"), col("b"), col("c")])?
2372 .build()?;
2373 let right_table_scan = test_table_scan_with_name("test1")?;
2374 let right = LogicalPlanBuilder::from(right_table_scan)
2375 .project(vec![col("a"), col("b"), col("c")])?
2376 .build()?;
2377 let filter = and(col("test.a").eq(lit(1)), col("test1.a").gt(lit(2)));
2378 let plan = LogicalPlanBuilder::from(left)
2379 .cross_join(right)?
2380 .project(vec![col("test.a"), col("test1.a")])?
2381 .filter(filter)?
2382 .build()?;
2383
2384 assert_optimized_plan_equal!(
2385 plan,
2386 @r"
2387 Projection: test.a, test1.a
2388 Cross Join:
2389 Projection: test.a, test.b, test.c
2390 TableScan: test, full_filters=[test.a = Int32(1)]
2391 Projection: test1.a, test1.b, test1.c
2392 TableScan: test1, full_filters=[test1.a > Int32(2)]
2393 "
2394 )
2395 }
2396
2397 #[test]
2399 fn filter_2_breaks_limits() -> Result<()> {
2400 let table_scan = test_table_scan()?;
2401 let plan = LogicalPlanBuilder::from(table_scan)
2402 .project(vec![col("a")])?
2403 .filter(col("a").lt_eq(lit(1i64)))?
2404 .limit(0, Some(1))?
2405 .project(vec![col("a")])?
2406 .filter(col("a").gt_eq(lit(1i64)))?
2407 .build()?;
2408 assert_snapshot!(plan,
2412 @r"
2413 Filter: test.a >= Int64(1)
2414 Projection: test.a
2415 Limit: skip=0, fetch=1
2416 Filter: test.a <= Int64(1)
2417 Projection: test.a
2418 TableScan: test
2419 ",
2420 );
2421 assert_optimized_plan_equal!(
2422 plan,
2423 @r"
2424 Projection: test.a
2425 Filter: test.a >= Int64(1)
2426 Limit: skip=0, fetch=1
2427 Projection: test.a
2428 TableScan: test, full_filters=[test.a <= Int64(1)]
2429 "
2430 )
2431 }
2432
2433 #[test]
2435 fn two_filters_on_same_depth() -> Result<()> {
2436 let table_scan = test_table_scan()?;
2437 let plan = LogicalPlanBuilder::from(table_scan)
2438 .limit(0, Some(1))?
2439 .filter(col("a").lt_eq(lit(1i64)))?
2440 .filter(col("a").gt_eq(lit(1i64)))?
2441 .project(vec![col("a")])?
2442 .build()?;
2443
2444 assert_snapshot!(plan,
2446 @r"
2447 Projection: test.a
2448 Filter: test.a >= Int64(1)
2449 Filter: test.a <= Int64(1)
2450 Limit: skip=0, fetch=1
2451 TableScan: test
2452 ",
2453 );
2454 assert_optimized_plan_equal!(
2455 plan,
2456 @r"
2457 Projection: test.a
2458 Filter: test.a <= Int64(1) AND test.a >= Int64(1)
2459 Limit: skip=0, fetch=1
2460 TableScan: test
2461 "
2462 )
2463 }
2464
2465 #[test]
2468 fn filters_user_defined_node() -> Result<()> {
2469 let table_scan = test_table_scan()?;
2470 let plan = LogicalPlanBuilder::from(table_scan)
2471 .filter(col("a").lt_eq(lit(1i64)))?
2472 .build()?;
2473
2474 let plan = user_defined::new(plan);
2475
2476 assert_snapshot!(plan,
2478 @r"
2479 TestUserDefined
2480 Filter: test.a <= Int64(1)
2481 TableScan: test
2482 ",
2483 );
2484 assert_optimized_plan_equal!(
2485 plan,
2486 @r"
2487 TestUserDefined
2488 TableScan: test, full_filters=[test.a <= Int64(1)]
2489 "
2490 )
2491 }
2492
2493 #[test]
2495 fn filter_on_join_on_common_independent() -> Result<()> {
2496 let table_scan = test_table_scan()?;
2497 let left = LogicalPlanBuilder::from(table_scan).build()?;
2498 let right_table_scan = test_table_scan_with_name("test2")?;
2499 let right = LogicalPlanBuilder::from(right_table_scan)
2500 .project(vec![col("a")])?
2501 .build()?;
2502 let plan = LogicalPlanBuilder::from(left)
2503 .join(
2504 right,
2505 JoinType::Inner,
2506 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
2507 None,
2508 )?
2509 .filter(col("test.a").lt_eq(lit(1i64)))?
2510 .build()?;
2511
2512 assert_snapshot!(plan,
2514 @r"
2515 Filter: test.a <= Int64(1)
2516 Inner Join: test.a = test2.a
2517 TableScan: test
2518 Projection: test2.a
2519 TableScan: test2
2520 ",
2521 );
2522 assert_optimized_plan_equal!(
2524 plan,
2525 @r"
2526 Inner Join: test.a = test2.a
2527 TableScan: test, full_filters=[test.a <= Int64(1)]
2528 Projection: test2.a
2529 TableScan: test2, full_filters=[test2.a <= Int64(1)]
2530 "
2531 )
2532 }
2533
2534 #[test]
2536 fn filter_using_join_on_common_independent() -> Result<()> {
2537 let table_scan = test_table_scan()?;
2538 let left = LogicalPlanBuilder::from(table_scan).build()?;
2539 let right_table_scan = test_table_scan_with_name("test2")?;
2540 let right = LogicalPlanBuilder::from(right_table_scan)
2541 .project(vec![col("a")])?
2542 .build()?;
2543 let plan = LogicalPlanBuilder::from(left)
2544 .join_using(
2545 right,
2546 JoinType::Inner,
2547 vec![Column::from_name("a".to_string())],
2548 )?
2549 .filter(col("a").lt_eq(lit(1i64)))?
2550 .build()?;
2551
2552 assert_snapshot!(plan,
2554 @r"
2555 Filter: test.a <= Int64(1)
2556 Inner Join: Using test.a = test2.a
2557 TableScan: test
2558 Projection: test2.a
2559 TableScan: test2
2560 ",
2561 );
2562 assert_optimized_plan_equal!(
2564 plan,
2565 @r"
2566 Inner Join: Using test.a = test2.a
2567 TableScan: test, full_filters=[test.a <= Int64(1)]
2568 Projection: test2.a
2569 TableScan: test2, full_filters=[test2.a <= Int64(1)]
2570 "
2571 )
2572 }
2573
2574 #[test]
2576 fn filter_join_on_common_dependent() -> Result<()> {
2577 let table_scan = test_table_scan()?;
2578 let left = LogicalPlanBuilder::from(table_scan)
2579 .project(vec![col("a"), col("c")])?
2580 .build()?;
2581 let right_table_scan = test_table_scan_with_name("test2")?;
2582 let right = LogicalPlanBuilder::from(right_table_scan)
2583 .project(vec![col("a"), col("b")])?
2584 .build()?;
2585 let plan = LogicalPlanBuilder::from(left)
2586 .join(
2587 right,
2588 JoinType::Inner,
2589 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
2590 None,
2591 )?
2592 .filter(col("c").lt_eq(col("b")))?
2593 .build()?;
2594
2595 assert_snapshot!(plan,
2597 @r"
2598 Filter: test.c <= test2.b
2599 Inner Join: test.a = test2.a
2600 Projection: test.a, test.c
2601 TableScan: test
2602 Projection: test2.a, test2.b
2603 TableScan: test2
2604 ",
2605 );
2606 assert_optimized_plan_equal!(
2608 plan,
2609 @r"
2610 Inner Join: test.a = test2.a Filter: test.c <= test2.b
2611 Projection: test.a, test.c
2612 TableScan: test
2613 Projection: test2.a, test2.b
2614 TableScan: test2
2615 "
2616 )
2617 }
2618
2619 #[test]
2621 fn filter_join_on_one_side() -> Result<()> {
2622 let table_scan = test_table_scan()?;
2623 let left = LogicalPlanBuilder::from(table_scan)
2624 .project(vec![col("a"), col("b")])?
2625 .build()?;
2626 let table_scan_right = test_table_scan_with_name("test2")?;
2627 let right = LogicalPlanBuilder::from(table_scan_right)
2628 .project(vec![col("a"), col("c")])?
2629 .build()?;
2630
2631 let plan = LogicalPlanBuilder::from(left)
2632 .join(
2633 right,
2634 JoinType::Inner,
2635 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
2636 None,
2637 )?
2638 .filter(col("b").lt_eq(lit(1i64)))?
2639 .build()?;
2640
2641 assert_snapshot!(plan,
2643 @r"
2644 Filter: test.b <= Int64(1)
2645 Inner Join: test.a = test2.a
2646 Projection: test.a, test.b
2647 TableScan: test
2648 Projection: test2.a, test2.c
2649 TableScan: test2
2650 ",
2651 );
2652 assert_optimized_plan_equal!(
2653 plan,
2654 @r"
2655 Inner Join: test.a = test2.a
2656 Projection: test.a, test.b
2657 TableScan: test, full_filters=[test.b <= Int64(1)]
2658 Projection: test2.a, test2.c
2659 TableScan: test2
2660 "
2661 )
2662 }
2663
2664 #[test]
2667 fn filter_using_left_join() -> Result<()> {
2668 let table_scan = test_table_scan()?;
2669 let left = LogicalPlanBuilder::from(table_scan).build()?;
2670 let right_table_scan = test_table_scan_with_name("test2")?;
2671 let right = LogicalPlanBuilder::from(right_table_scan)
2672 .project(vec![col("a")])?
2673 .build()?;
2674 let plan = LogicalPlanBuilder::from(left)
2675 .join_using(
2676 right,
2677 JoinType::Left,
2678 vec![Column::from_name("a".to_string())],
2679 )?
2680 .filter(col("test2.a").lt_eq(lit(1i64)))?
2681 .build()?;
2682
2683 assert_snapshot!(plan,
2685 @r"
2686 Filter: test2.a <= Int64(1)
2687 Left Join: Using test.a = test2.a
2688 TableScan: test
2689 Projection: test2.a
2690 TableScan: test2
2691 ",
2692 );
2693 assert_optimized_plan_equal!(
2695 plan,
2696 @r"
2697 Filter: test2.a <= Int64(1)
2698 Left Join: Using test.a = test2.a
2699 TableScan: test, full_filters=[test.a <= Int64(1)]
2700 Projection: test2.a
2701 TableScan: test2
2702 "
2703 )
2704 }
2705
2706 #[test]
2708 fn filter_using_right_join() -> Result<()> {
2709 let table_scan = test_table_scan()?;
2710 let left = LogicalPlanBuilder::from(table_scan).build()?;
2711 let right_table_scan = test_table_scan_with_name("test2")?;
2712 let right = LogicalPlanBuilder::from(right_table_scan)
2713 .project(vec![col("a")])?
2714 .build()?;
2715 let plan = LogicalPlanBuilder::from(left)
2716 .join_using(
2717 right,
2718 JoinType::Right,
2719 vec![Column::from_name("a".to_string())],
2720 )?
2721 .filter(col("test.a").lt_eq(lit(1i64)))?
2722 .build()?;
2723
2724 assert_snapshot!(plan,
2726 @r"
2727 Filter: test.a <= Int64(1)
2728 Right Join: Using test.a = test2.a
2729 TableScan: test
2730 Projection: test2.a
2731 TableScan: test2
2732 ",
2733 );
2734 assert_optimized_plan_equal!(
2736 plan,
2737 @r"
2738 Filter: test.a <= Int64(1)
2739 Right Join: Using test.a = test2.a
2740 TableScan: test
2741 Projection: test2.a
2742 TableScan: test2, full_filters=[test2.a <= Int64(1)]
2743 "
2744 )
2745 }
2746
2747 #[test]
2749 fn filter_using_left_join_on_common() -> Result<()> {
2750 let table_scan = test_table_scan()?;
2751 let left = LogicalPlanBuilder::from(table_scan).build()?;
2752 let right_table_scan = test_table_scan_with_name("test2")?;
2753 let right = LogicalPlanBuilder::from(right_table_scan)
2754 .project(vec![col("a")])?
2755 .build()?;
2756 let plan = LogicalPlanBuilder::from(left)
2757 .join_using(
2758 right,
2759 JoinType::Left,
2760 vec![Column::from_name("a".to_string())],
2761 )?
2762 .filter(col("a").lt_eq(lit(1i64)))?
2763 .build()?;
2764
2765 assert_snapshot!(plan,
2767 @r"
2768 Filter: test.a <= Int64(1)
2769 Left Join: Using test.a = test2.a
2770 TableScan: test
2771 Projection: test2.a
2772 TableScan: test2
2773 ",
2774 );
2775 assert_optimized_plan_equal!(
2777 plan,
2778 @r"
2779 Left Join: Using test.a = test2.a
2780 TableScan: test, full_filters=[test.a <= Int64(1)]
2781 Projection: test2.a
2782 TableScan: test2, full_filters=[test2.a <= Int64(1)]
2783 "
2784 )
2785 }
2786
2787 #[test]
2789 fn filter_using_right_join_on_common() -> Result<()> {
2790 let table_scan = test_table_scan()?;
2791 let left = LogicalPlanBuilder::from(table_scan).build()?;
2792 let right_table_scan = test_table_scan_with_name("test2")?;
2793 let right = LogicalPlanBuilder::from(right_table_scan)
2794 .project(vec![col("a")])?
2795 .build()?;
2796 let plan = LogicalPlanBuilder::from(left)
2797 .join_using(
2798 right,
2799 JoinType::Right,
2800 vec![Column::from_name("a".to_string())],
2801 )?
2802 .filter(col("test2.a").lt_eq(lit(1i64)))?
2803 .build()?;
2804
2805 assert_snapshot!(plan,
2807 @r"
2808 Filter: test2.a <= Int64(1)
2809 Right Join: Using test.a = test2.a
2810 TableScan: test
2811 Projection: test2.a
2812 TableScan: test2
2813 ",
2814 );
2815 assert_optimized_plan_equal!(
2817 plan,
2818 @r"
2819 Right Join: Using test.a = test2.a
2820 TableScan: test, full_filters=[test.a <= Int64(1)]
2821 Projection: test2.a
2822 TableScan: test2, full_filters=[test2.a <= Int64(1)]
2823 "
2824 )
2825 }
2826
2827 #[test]
2829 fn join_on_with_filter() -> Result<()> {
2830 let table_scan = test_table_scan()?;
2831 let left = LogicalPlanBuilder::from(table_scan)
2832 .project(vec![col("a"), col("b"), col("c")])?
2833 .build()?;
2834 let right_table_scan = test_table_scan_with_name("test2")?;
2835 let right = LogicalPlanBuilder::from(right_table_scan)
2836 .project(vec![col("a"), col("b"), col("c")])?
2837 .build()?;
2838 let filter = col("test.c")
2839 .gt(lit(1u32))
2840 .and(col("test.b").lt(col("test2.b")))
2841 .and(col("test2.c").gt(lit(4u32)));
2842 let plan = LogicalPlanBuilder::from(left)
2843 .join(
2844 right,
2845 JoinType::Inner,
2846 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
2847 Some(filter),
2848 )?
2849 .build()?;
2850
2851 assert_snapshot!(plan,
2853 @r"
2854 Inner Join: test.a = test2.a Filter: test.c > UInt32(1) AND test.b < test2.b AND test2.c > UInt32(4)
2855 Projection: test.a, test.b, test.c
2856 TableScan: test
2857 Projection: test2.a, test2.b, test2.c
2858 TableScan: test2
2859 ",
2860 );
2861 assert_optimized_plan_equal!(
2862 plan,
2863 @r"
2864 Inner Join: test.a = test2.a Filter: test.b < test2.b
2865 Projection: test.a, test.b, test.c
2866 TableScan: test, full_filters=[test.c > UInt32(1)]
2867 Projection: test2.a, test2.b, test2.c
2868 TableScan: test2, full_filters=[test2.c > UInt32(4)]
2869 "
2870 )
2871 }
2872
2873 #[test]
2875 fn join_filter_removed() -> Result<()> {
2876 let table_scan = test_table_scan()?;
2877 let left = LogicalPlanBuilder::from(table_scan)
2878 .project(vec![col("a"), col("b"), col("c")])?
2879 .build()?;
2880 let right_table_scan = test_table_scan_with_name("test2")?;
2881 let right = LogicalPlanBuilder::from(right_table_scan)
2882 .project(vec![col("a"), col("b"), col("c")])?
2883 .build()?;
2884 let filter = col("test.b")
2885 .gt(lit(1u32))
2886 .and(col("test2.c").gt(lit(4u32)));
2887 let plan = LogicalPlanBuilder::from(left)
2888 .join(
2889 right,
2890 JoinType::Inner,
2891 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
2892 Some(filter),
2893 )?
2894 .build()?;
2895
2896 assert_snapshot!(plan,
2898 @r"
2899 Inner Join: test.a = test2.a Filter: test.b > UInt32(1) AND test2.c > UInt32(4)
2900 Projection: test.a, test.b, test.c
2901 TableScan: test
2902 Projection: test2.a, test2.b, test2.c
2903 TableScan: test2
2904 ",
2905 );
2906 assert_optimized_plan_equal!(
2907 plan,
2908 @r"
2909 Inner Join: test.a = test2.a
2910 Projection: test.a, test.b, test.c
2911 TableScan: test, full_filters=[test.b > UInt32(1)]
2912 Projection: test2.a, test2.b, test2.c
2913 TableScan: test2, full_filters=[test2.c > UInt32(4)]
2914 "
2915 )
2916 }
2917
2918 #[test]
2920 fn join_filter_on_common() -> Result<()> {
2921 let table_scan = test_table_scan()?;
2922 let left = LogicalPlanBuilder::from(table_scan)
2923 .project(vec![col("a")])?
2924 .build()?;
2925 let right_table_scan = test_table_scan_with_name("test2")?;
2926 let right = LogicalPlanBuilder::from(right_table_scan)
2927 .project(vec![col("b")])?
2928 .build()?;
2929 let filter = col("test.a").gt(lit(1u32));
2930 let plan = LogicalPlanBuilder::from(left)
2931 .join(
2932 right,
2933 JoinType::Inner,
2934 (vec![Column::from_name("a")], vec![Column::from_name("b")]),
2935 Some(filter),
2936 )?
2937 .build()?;
2938
2939 assert_snapshot!(plan,
2941 @r"
2942 Inner Join: test.a = test2.b Filter: test.a > UInt32(1)
2943 Projection: test.a
2944 TableScan: test
2945 Projection: test2.b
2946 TableScan: test2
2947 ",
2948 );
2949 assert_optimized_plan_equal!(
2950 plan,
2951 @r"
2952 Inner Join: test.a = test2.b
2953 Projection: test.a
2954 TableScan: test, full_filters=[test.a > UInt32(1)]
2955 Projection: test2.b
2956 TableScan: test2, full_filters=[test2.b > UInt32(1)]
2957 "
2958 )
2959 }
2960
2961 #[test]
2963 fn left_join_on_with_filter() -> Result<()> {
2964 let table_scan = test_table_scan()?;
2965 let left = LogicalPlanBuilder::from(table_scan)
2966 .project(vec![col("a"), col("b"), col("c")])?
2967 .build()?;
2968 let right_table_scan = test_table_scan_with_name("test2")?;
2969 let right = LogicalPlanBuilder::from(right_table_scan)
2970 .project(vec![col("a"), col("b"), col("c")])?
2971 .build()?;
2972 let filter = col("test.a")
2973 .gt(lit(1u32))
2974 .and(col("test.b").lt(col("test2.b")))
2975 .and(col("test2.c").gt(lit(4u32)));
2976 let plan = LogicalPlanBuilder::from(left)
2977 .join(
2978 right,
2979 JoinType::Left,
2980 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
2981 Some(filter),
2982 )?
2983 .build()?;
2984
2985 assert_snapshot!(plan,
2987 @r"
2988 Left Join: test.a = test2.a Filter: test.a > UInt32(1) AND test.b < test2.b AND test2.c > UInt32(4)
2989 Projection: test.a, test.b, test.c
2990 TableScan: test
2991 Projection: test2.a, test2.b, test2.c
2992 TableScan: test2
2993 ",
2994 );
2995 assert_optimized_plan_equal!(
2996 plan,
2997 @r"
2998 Left Join: test.a = test2.a Filter: test.a > UInt32(1) AND test.b < test2.b
2999 Projection: test.a, test.b, test.c
3000 TableScan: test
3001 Projection: test2.a, test2.b, test2.c
3002 TableScan: test2, full_filters=[test2.a > UInt32(1), test2.c > UInt32(4)]
3003 "
3004 )
3005 }
3006
3007 #[test]
3009 fn right_join_on_with_filter() -> Result<()> {
3010 let table_scan = test_table_scan()?;
3011 let left = LogicalPlanBuilder::from(table_scan)
3012 .project(vec![col("a"), col("b"), col("c")])?
3013 .build()?;
3014 let right_table_scan = test_table_scan_with_name("test2")?;
3015 let right = LogicalPlanBuilder::from(right_table_scan)
3016 .project(vec![col("a"), col("b"), col("c")])?
3017 .build()?;
3018 let filter = col("test.a")
3019 .gt(lit(1u32))
3020 .and(col("test.b").lt(col("test2.b")))
3021 .and(col("test2.c").gt(lit(4u32)));
3022 let plan = LogicalPlanBuilder::from(left)
3023 .join(
3024 right,
3025 JoinType::Right,
3026 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
3027 Some(filter),
3028 )?
3029 .build()?;
3030
3031 assert_snapshot!(plan,
3033 @r"
3034 Right Join: test.a = test2.a Filter: test.a > UInt32(1) AND test.b < test2.b AND test2.c > UInt32(4)
3035 Projection: test.a, test.b, test.c
3036 TableScan: test
3037 Projection: test2.a, test2.b, test2.c
3038 TableScan: test2
3039 ",
3040 );
3041 assert_optimized_plan_equal!(
3042 plan,
3043 @r"
3044 Right Join: test.a = test2.a Filter: test.b < test2.b AND test2.c > UInt32(4)
3045 Projection: test.a, test.b, test.c
3046 TableScan: test, full_filters=[test.a > UInt32(1)]
3047 Projection: test2.a, test2.b, test2.c
3048 TableScan: test2
3049 "
3050 )
3051 }
3052
3053 #[test]
3055 fn full_join_on_with_filter() -> Result<()> {
3056 let table_scan = test_table_scan()?;
3057 let left = LogicalPlanBuilder::from(table_scan)
3058 .project(vec![col("a"), col("b"), col("c")])?
3059 .build()?;
3060 let right_table_scan = test_table_scan_with_name("test2")?;
3061 let right = LogicalPlanBuilder::from(right_table_scan)
3062 .project(vec![col("a"), col("b"), col("c")])?
3063 .build()?;
3064 let filter = col("test.a")
3065 .gt(lit(1u32))
3066 .and(col("test.b").lt(col("test2.b")))
3067 .and(col("test2.c").gt(lit(4u32)));
3068 let plan = LogicalPlanBuilder::from(left)
3069 .join(
3070 right,
3071 JoinType::Full,
3072 (vec![Column::from_name("a")], vec![Column::from_name("a")]),
3073 Some(filter),
3074 )?
3075 .build()?;
3076 assert_plan_not_transformed!(plan.clone());
3077
3078 assert_snapshot!(plan,
3080 @r"
3081 Full Join: test.a = test2.a Filter: test.a > UInt32(1) AND test.b < test2.b AND test2.c > UInt32(4)
3082 Projection: test.a, test.b, test.c
3083 TableScan: test
3084 Projection: test2.a, test2.b, test2.c
3085 TableScan: test2
3086 ",
3087 );
3088 assert_optimized_plan_equal!(
3089 plan,
3090 @r"
3091 Full Join: test.a = test2.a Filter: test.a > UInt32(1) AND test.b < test2.b AND test2.c > UInt32(4)
3092 Projection: test.a, test.b, test.c
3093 TableScan: test
3094 Projection: test2.a, test2.b, test2.c
3095 TableScan: test2
3096 "
3097 )
3098 }
3099
3100 struct PushDownProvider {
3101 pub filter_support: TableProviderFilterPushDown,
3102 }
3103
3104 #[async_trait]
3105 impl TableSource for PushDownProvider {
3106 fn schema(&self) -> SchemaRef {
3107 Arc::new(Schema::new(vec![
3108 Field::new("a", DataType::Int32, true),
3109 Field::new("b", DataType::Int32, true),
3110 ]))
3111 }
3112
3113 fn table_type(&self) -> TableType {
3114 TableType::Base
3115 }
3116
3117 fn supports_filters_pushdown(
3118 &self,
3119 filters: &[&Expr],
3120 ) -> Result<Vec<TableProviderFilterPushDown>> {
3121 Ok((0..filters.len())
3122 .map(|_| self.filter_support.clone())
3123 .collect())
3124 }
3125 }
3126
3127 fn table_scan_with_pushdown_provider_builder(
3128 filter_support: TableProviderFilterPushDown,
3129 filters: Vec<Expr>,
3130 projection: Option<Vec<usize>>,
3131 ) -> Result<LogicalPlanBuilder> {
3132 let test_provider = PushDownProvider { filter_support };
3133
3134 let table_scan = LogicalPlan::TableScan(TableScan {
3135 table_name: "test".into(),
3136 filters,
3137 projected_schema: Arc::new(DFSchema::try_from(test_provider.schema())?),
3138 projection,
3139 source: Arc::new(test_provider),
3140 fetch: None,
3141 statistics_requests: std::collections::BTreeSet::new(),
3142 });
3143
3144 Ok(LogicalPlanBuilder::from(table_scan))
3145 }
3146
3147 fn table_scan_with_pushdown_provider(
3148 filter_support: TableProviderFilterPushDown,
3149 ) -> Result<LogicalPlan> {
3150 table_scan_with_pushdown_provider_builder(filter_support, vec![], None)?
3151 .filter(col("a").eq(lit(1i64)))?
3152 .build()
3153 }
3154
3155 #[test]
3156 fn filter_with_table_provider_exact() -> Result<()> {
3157 let plan = table_scan_with_pushdown_provider(TableProviderFilterPushDown::Exact)?;
3158
3159 assert_optimized_plan_equal!(
3160 plan,
3161 @"TableScan: test, full_filters=[a = Int64(1)]"
3162 )
3163 }
3164
3165 #[test]
3166 fn filter_with_table_provider_inexact() -> Result<()> {
3167 let plan =
3168 table_scan_with_pushdown_provider(TableProviderFilterPushDown::Inexact)?;
3169
3170 assert_optimized_plan_equal!(
3171 plan,
3172 @r"
3173 Filter: a = Int64(1)
3174 TableScan: test, partial_filters=[a = Int64(1)]
3175 "
3176 )
3177 }
3178
3179 #[test]
3180 fn filter_with_table_provider_multiple_invocations() -> Result<()> {
3181 let plan =
3182 table_scan_with_pushdown_provider(TableProviderFilterPushDown::Inexact)?;
3183
3184 let optimized = PushDownFilter::new()
3185 .rewrite(plan, &OptimizerContext::new())
3186 .expect("failed to optimize plan");
3187 assert!(optimized.transformed);
3188 assert_plan_not_transformed!(optimized.data.clone());
3189
3190 assert_optimized_plan_equal!(
3193 optimized.data,
3194 @r"
3195 Filter: a = Int64(1)
3196 TableScan: test, partial_filters=[a = Int64(1)]
3197 "
3198 )
3199 }
3200
3201 #[test]
3202 fn filter_with_table_provider_unsupported() -> Result<()> {
3203 let plan =
3204 table_scan_with_pushdown_provider(TableProviderFilterPushDown::Unsupported)?;
3205 assert_plan_not_transformed!(plan.clone());
3206
3207 assert_optimized_plan_equal!(
3208 plan,
3209 @r"
3210 Filter: a = Int64(1)
3211 TableScan: test
3212 "
3213 )
3214 }
3215
3216 #[test]
3217 fn multi_combined_filter() -> Result<()> {
3218 let plan = table_scan_with_pushdown_provider_builder(
3219 TableProviderFilterPushDown::Inexact,
3220 vec![col("a").eq(lit(10i64)), col("b").gt(lit(11i64))],
3221 Some(vec![0]),
3222 )?
3223 .filter(and(col("a").eq(lit(10i64)), col("b").gt(lit(11i64))))?
3224 .project(vec![col("a"), col("b")])?
3225 .build()?;
3226
3227 assert_optimized_plan_equal!(
3228 plan,
3229 @r"
3230 Projection: a, b
3231 Filter: a = Int64(10) AND b > Int64(11)
3232 TableScan: test projection=[a], partial_filters=[a = Int64(10), b > Int64(11)]
3233 "
3234 )
3235 }
3236
3237 #[test]
3238 fn multi_combined_two_filters() -> Result<()> {
3239 let plan = table_scan_with_pushdown_provider_builder(
3240 TableProviderFilterPushDown::Inexact,
3241 vec![col("a").eq(lit(10i64)), col("b").gt(lit(11i64))],
3242 Some(vec![0]),
3243 )?
3244 .filter(col("a").eq(lit(10i64)))?
3245 .filter(col("b").gt(lit(11i64)))?
3246 .project(vec![col("a"), col("b")])?
3247 .build()?;
3248
3249 assert_optimized_plan_equal!(
3250 plan,
3251 @r"
3252 Projection: a, b
3253 Filter: a = Int64(10) AND b > Int64(11)
3254 TableScan: test projection=[a], partial_filters=[a = Int64(10), b > Int64(11)]
3255 "
3256 )
3257 }
3258
3259 #[test]
3260 fn multi_combined_filter_exact() -> Result<()> {
3261 let plan = table_scan_with_pushdown_provider_builder(
3262 TableProviderFilterPushDown::Exact,
3263 vec![],
3264 Some(vec![0]),
3265 )?
3266 .filter(and(col("a").eq(lit(10i64)), col("b").gt(lit(11i64))))?
3267 .project(vec![col("a"), col("b")])?
3268 .build()?;
3269
3270 assert_optimized_plan_equal!(
3271 plan,
3272 @r"
3273 Projection: a, b
3274 TableScan: test projection=[a], full_filters=[a = Int64(10), b > Int64(11)]
3275 "
3276 )
3277 }
3278
3279 #[test]
3280 fn multi_combined_two_filters_exact() -> Result<()> {
3281 let plan = table_scan_with_pushdown_provider_builder(
3282 TableProviderFilterPushDown::Exact,
3283 vec![],
3284 Some(vec![0]),
3285 )?
3286 .filter(col("a").eq(lit(10i64)))?
3287 .filter(col("b").gt(lit(11i64)))?
3288 .project(vec![col("a"), col("b")])?
3289 .build()?;
3290
3291 assert_optimized_plan_equal!(
3292 plan,
3293 @r"
3294 Projection: a, b
3295 TableScan: test projection=[a], full_filters=[a = Int64(10), b > Int64(11)]
3296 "
3297 )
3298 }
3299
3300 #[test]
3301 fn test_filter_with_alias() -> Result<()> {
3302 let table_scan = test_table_scan()?;
3306 let plan = LogicalPlanBuilder::from(table_scan)
3307 .project(vec![col("a").alias("b"), col("c")])?
3308 .filter(and(col("b").gt(lit(10i64)), col("c").gt(lit(10i64))))?
3309 .build()?;
3310
3311 assert_snapshot!(plan,
3313 @r"
3314 Filter: b > Int64(10) AND test.c > Int64(10)
3315 Projection: test.a AS b, test.c
3316 TableScan: test
3317 ",
3318 );
3319 assert_optimized_plan_equal!(
3321 plan,
3322 @r"
3323 Projection: test.a AS b, test.c
3324 TableScan: test, full_filters=[test.a > Int64(10), test.c > Int64(10)]
3325 "
3326 )
3327 }
3328
3329 #[test]
3330 fn test_filter_with_alias_2() -> Result<()> {
3331 let table_scan = test_table_scan()?;
3335 let plan = LogicalPlanBuilder::from(table_scan)
3336 .project(vec![col("a").alias("b"), col("c")])?
3337 .project(vec![col("b"), col("c")])?
3338 .filter(and(col("b").gt(lit(10i64)), col("c").gt(lit(10i64))))?
3339 .build()?;
3340
3341 assert_snapshot!(plan,
3343 @r"
3344 Filter: b > Int64(10) AND test.c > Int64(10)
3345 Projection: b, test.c
3346 Projection: test.a AS b, test.c
3347 TableScan: test
3348 ",
3349 );
3350 assert_optimized_plan_equal!(
3352 plan,
3353 @r"
3354 Projection: b, test.c
3355 Projection: test.a AS b, test.c
3356 TableScan: test, full_filters=[test.a > Int64(10), test.c > Int64(10)]
3357 "
3358 )
3359 }
3360
3361 #[test]
3362 fn test_filter_with_multi_alias() -> Result<()> {
3363 let table_scan = test_table_scan()?;
3364 let plan = LogicalPlanBuilder::from(table_scan)
3365 .project(vec![col("a").alias("b"), col("c").alias("d")])?
3366 .filter(and(col("b").gt(lit(10i64)), col("d").gt(lit(10i64))))?
3367 .build()?;
3368
3369 assert_snapshot!(plan,
3371 @r"
3372 Filter: b > Int64(10) AND d > Int64(10)
3373 Projection: test.a AS b, test.c AS d
3374 TableScan: test
3375 ",
3376 );
3377 assert_optimized_plan_equal!(
3379 plan,
3380 @r"
3381 Projection: test.a AS b, test.c AS d
3382 TableScan: test, full_filters=[test.a > Int64(10), test.c > Int64(10)]
3383 "
3384 )
3385 }
3386
3387 #[test]
3389 fn join_filter_with_alias() -> Result<()> {
3390 let table_scan = test_table_scan()?;
3391 let left = LogicalPlanBuilder::from(table_scan)
3392 .project(vec![col("a").alias("c")])?
3393 .build()?;
3394 let right_table_scan = test_table_scan_with_name("test2")?;
3395 let right = LogicalPlanBuilder::from(right_table_scan)
3396 .project(vec![col("b").alias("d")])?
3397 .build()?;
3398 let filter = col("c").gt(lit(1u32));
3399 let plan = LogicalPlanBuilder::from(left)
3400 .join(
3401 right,
3402 JoinType::Inner,
3403 (vec![Column::from_name("c")], vec![Column::from_name("d")]),
3404 Some(filter),
3405 )?
3406 .build()?;
3407
3408 assert_snapshot!(plan,
3409 @r"
3410 Inner Join: c = d Filter: c > UInt32(1)
3411 Projection: test.a AS c
3412 TableScan: test
3413 Projection: test2.b AS d
3414 TableScan: test2
3415 ",
3416 );
3417 assert_optimized_plan_equal!(
3419 plan,
3420 @r"
3421 Inner Join: c = d
3422 Projection: test.a AS c
3423 TableScan: test, full_filters=[test.a > UInt32(1)]
3424 Projection: test2.b AS d
3425 TableScan: test2, full_filters=[test2.b > UInt32(1)]
3426 "
3427 )
3428 }
3429
3430 #[test]
3431 fn test_in_filter_with_alias() -> Result<()> {
3432 let table_scan = test_table_scan()?;
3436 let filter_value = vec![lit(1u32), lit(2u32), lit(3u32), lit(4u32)];
3437 let plan = LogicalPlanBuilder::from(table_scan)
3438 .project(vec![col("a").alias("b"), col("c")])?
3439 .filter(in_list(col("b"), filter_value, false))?
3440 .build()?;
3441
3442 assert_snapshot!(plan,
3444 @r"
3445 Filter: b IN ([UInt32(1), UInt32(2), UInt32(3), UInt32(4)])
3446 Projection: test.a AS b, test.c
3447 TableScan: test
3448 ",
3449 );
3450 assert_optimized_plan_equal!(
3452 plan,
3453 @r"
3454 Projection: test.a AS b, test.c
3455 TableScan: test, full_filters=[test.a IN ([UInt32(1), UInt32(2), UInt32(3), UInt32(4)])]
3456 "
3457 )
3458 }
3459
3460 #[test]
3461 fn test_in_filter_with_alias_2() -> Result<()> {
3462 let table_scan = test_table_scan()?;
3466 let filter_value = vec![lit(1u32), lit(2u32), lit(3u32), lit(4u32)];
3467 let plan = LogicalPlanBuilder::from(table_scan)
3468 .project(vec![col("a").alias("b"), col("c")])?
3469 .project(vec![col("b"), col("c")])?
3470 .filter(in_list(col("b"), filter_value, false))?
3471 .build()?;
3472
3473 assert_snapshot!(plan,
3475 @r"
3476 Filter: b IN ([UInt32(1), UInt32(2), UInt32(3), UInt32(4)])
3477 Projection: b, test.c
3478 Projection: test.a AS b, test.c
3479 TableScan: test
3480 ",
3481 );
3482 assert_optimized_plan_equal!(
3484 plan,
3485 @r"
3486 Projection: b, test.c
3487 Projection: test.a AS b, test.c
3488 TableScan: test, full_filters=[test.a IN ([UInt32(1), UInt32(2), UInt32(3), UInt32(4)])]
3489 "
3490 )
3491 }
3492
3493 #[test]
3494 fn test_in_subquery_with_alias() -> Result<()> {
3495 let table_scan = test_table_scan()?;
3498 let table_scan_sq = test_table_scan_with_name("sq")?;
3499 let subplan = Arc::new(
3500 LogicalPlanBuilder::from(table_scan_sq)
3501 .project(vec![col("c")])?
3502 .build()?,
3503 );
3504 let plan = LogicalPlanBuilder::from(table_scan)
3505 .project(vec![col("a").alias("b"), col("c")])?
3506 .filter(in_subquery(col("b"), subplan))?
3507 .build()?;
3508
3509 assert_snapshot!(plan,
3511 @r"
3512 Filter: b IN (<subquery>)
3513 Subquery:
3514 Projection: sq.c
3515 TableScan: sq
3516 Projection: test.a AS b, test.c
3517 TableScan: test
3518 ",
3519 );
3520 assert_optimized_plan_equal!(
3522 plan,
3523 @r"
3524 Projection: test.a AS b, test.c
3525 TableScan: test, full_filters=[test.a IN (<subquery>)]
3526 Subquery:
3527 Projection: sq.c
3528 TableScan: sq
3529 "
3530 )
3531 }
3532
3533 #[test]
3534 fn test_propagation_of_optimized_inner_filters_with_projections() -> Result<()> {
3535 let plan = LogicalPlanBuilder::empty(true)
3537 .project(vec![lit(0i64).alias("a")])?
3538 .alias("b")?
3539 .project(vec![col("b.a")])?
3540 .alias("b")?
3541 .filter(col("b.a").eq(lit(1i64)))?
3542 .project(vec![col("b.a")])?
3543 .build()?;
3544
3545 assert_snapshot!(plan,
3546 @r"
3547 Projection: b.a
3548 Filter: b.a = Int64(1)
3549 SubqueryAlias: b
3550 Projection: b.a
3551 SubqueryAlias: b
3552 Projection: Int64(0) AS a
3553 EmptyRelation: rows=1
3554 ",
3555 );
3556 assert_optimized_plan_equal!(
3559 plan,
3560 @r"
3561 Projection: b.a
3562 SubqueryAlias: b
3563 Projection: b.a
3564 SubqueryAlias: b
3565 Projection: Int64(0) AS a
3566 Filter: Int64(0) = Int64(1)
3567 EmptyRelation: rows=1
3568 "
3569 )
3570 }
3571
3572 #[test]
3573 fn test_crossjoin_with_or_clause() -> Result<()> {
3574 let table_scan = test_table_scan()?;
3576 let left = LogicalPlanBuilder::from(table_scan)
3577 .project(vec![col("a"), col("b"), col("c")])?
3578 .build()?;
3579 let right_table_scan = test_table_scan_with_name("test1")?;
3580 let right = LogicalPlanBuilder::from(right_table_scan)
3581 .project(vec![col("a").alias("d"), col("a").alias("e")])?
3582 .build()?;
3583 let filter = or(
3584 and(col("a").eq(col("d")), col("b").gt(lit(1u32))),
3585 and(col("b").eq(col("e")), col("c").lt(lit(10u32))),
3586 );
3587 let plan = LogicalPlanBuilder::from(left)
3588 .cross_join(right)?
3589 .filter(filter)?
3590 .build()?;
3591
3592 assert_optimized_plan_eq_with_rewrite_predicate!(plan.clone(), @r"
3593 Inner Join: Filter: test.a = d AND test.b > UInt32(1) OR test.b = e AND test.c < UInt32(10)
3594 Projection: test.a, test.b, test.c
3595 TableScan: test, full_filters=[test.b > UInt32(1) OR test.c < UInt32(10)]
3596 Projection: test1.a AS d, test1.a AS e
3597 TableScan: test1
3598 ")?;
3599
3600 let optimized_plan = PushDownFilter::new()
3603 .rewrite(plan, &OptimizerContext::new())
3604 .expect("failed to optimize plan")
3605 .data;
3606 assert_optimized_plan_equal!(
3607 optimized_plan,
3608 @r"
3609 Inner Join: Filter: test.a = d AND test.b > UInt32(1) OR test.b = e AND test.c < UInt32(10)
3610 Projection: test.a, test.b, test.c
3611 TableScan: test, full_filters=[test.b > UInt32(1) OR test.c < UInt32(10)]
3612 Projection: test1.a AS d, test1.a AS e
3613 TableScan: test1
3614 "
3615 )
3616 }
3617
3618 #[test]
3619 fn left_semi_join() -> Result<()> {
3620 let left = test_table_scan_with_name("test1")?;
3621 let right_table_scan = test_table_scan_with_name("test2")?;
3622 let right = LogicalPlanBuilder::from(right_table_scan)
3623 .project(vec![col("a"), col("b")])?
3624 .build()?;
3625 let plan = LogicalPlanBuilder::from(left)
3626 .join(
3627 right,
3628 JoinType::LeftSemi,
3629 (
3630 vec![Column::from_qualified_name("test1.a")],
3631 vec![Column::from_qualified_name("test2.a")],
3632 ),
3633 None,
3634 )?
3635 .filter(col("test2.a").lt_eq(lit(1i64)))?
3636 .build()?;
3637
3638 assert_snapshot!(plan,
3640 @r"
3641 Filter: test2.a <= Int64(1)
3642 LeftSemi Join: test1.a = test2.a
3643 TableScan: test1
3644 Projection: test2.a, test2.b
3645 TableScan: test2
3646 ",
3647 );
3648 assert_optimized_plan_equal!(
3650 plan,
3651 @r"
3652 Filter: test2.a <= Int64(1)
3653 LeftSemi Join: test1.a = test2.a
3654 TableScan: test1, full_filters=[test1.a <= Int64(1)]
3655 Projection: test2.a, test2.b
3656 TableScan: test2
3657 "
3658 )
3659 }
3660
3661 #[test]
3662 fn left_semi_join_with_filters() -> Result<()> {
3663 let left = test_table_scan_with_name("test1")?;
3664 let right_table_scan = test_table_scan_with_name("test2")?;
3665 let right = LogicalPlanBuilder::from(right_table_scan)
3666 .project(vec![col("a"), col("b")])?
3667 .build()?;
3668 let plan = LogicalPlanBuilder::from(left)
3669 .join(
3670 right,
3671 JoinType::LeftSemi,
3672 (
3673 vec![Column::from_qualified_name("test1.a")],
3674 vec![Column::from_qualified_name("test2.a")],
3675 ),
3676 Some(
3677 col("test1.b")
3678 .gt(lit(1u32))
3679 .and(col("test2.b").gt(lit(2u32))),
3680 ),
3681 )?
3682 .build()?;
3683
3684 assert_snapshot!(plan,
3686 @r"
3687 LeftSemi Join: test1.a = test2.a Filter: test1.b > UInt32(1) AND test2.b > UInt32(2)
3688 TableScan: test1
3689 Projection: test2.a, test2.b
3690 TableScan: test2
3691 ",
3692 );
3693 assert_optimized_plan_equal!(
3695 plan,
3696 @r"
3697 LeftSemi Join: test1.a = test2.a
3698 TableScan: test1, full_filters=[test1.b > UInt32(1)]
3699 Projection: test2.a, test2.b
3700 TableScan: test2, full_filters=[test2.b > UInt32(2)]
3701 "
3702 )
3703 }
3704
3705 #[test]
3706 fn right_semi_join() -> Result<()> {
3707 let left = test_table_scan_with_name("test1")?;
3708 let right_table_scan = test_table_scan_with_name("test2")?;
3709 let right = LogicalPlanBuilder::from(right_table_scan)
3710 .project(vec![col("a"), col("b")])?
3711 .build()?;
3712 let plan = LogicalPlanBuilder::from(left)
3713 .join(
3714 right,
3715 JoinType::RightSemi,
3716 (
3717 vec![Column::from_qualified_name("test1.a")],
3718 vec![Column::from_qualified_name("test2.a")],
3719 ),
3720 None,
3721 )?
3722 .filter(col("test1.a").lt_eq(lit(1i64)))?
3723 .build()?;
3724
3725 assert_snapshot!(plan,
3727 @r"
3728 Filter: test1.a <= Int64(1)
3729 RightSemi Join: test1.a = test2.a
3730 TableScan: test1
3731 Projection: test2.a, test2.b
3732 TableScan: test2
3733 ",
3734 );
3735 assert_optimized_plan_equal!(
3737 plan,
3738 @r"
3739 Filter: test1.a <= Int64(1)
3740 RightSemi Join: test1.a = test2.a
3741 TableScan: test1
3742 Projection: test2.a, test2.b
3743 TableScan: test2, full_filters=[test2.a <= Int64(1)]
3744 "
3745 )
3746 }
3747
3748 #[test]
3749 fn right_semi_join_with_filters() -> Result<()> {
3750 let left = test_table_scan_with_name("test1")?;
3751 let right_table_scan = test_table_scan_with_name("test2")?;
3752 let right = LogicalPlanBuilder::from(right_table_scan)
3753 .project(vec![col("a"), col("b")])?
3754 .build()?;
3755 let plan = LogicalPlanBuilder::from(left)
3756 .join(
3757 right,
3758 JoinType::RightSemi,
3759 (
3760 vec![Column::from_qualified_name("test1.a")],
3761 vec![Column::from_qualified_name("test2.a")],
3762 ),
3763 Some(
3764 col("test1.b")
3765 .gt(lit(1u32))
3766 .and(col("test2.b").gt(lit(2u32))),
3767 ),
3768 )?
3769 .build()?;
3770
3771 assert_snapshot!(plan,
3773 @r"
3774 RightSemi Join: test1.a = test2.a Filter: test1.b > UInt32(1) AND test2.b > UInt32(2)
3775 TableScan: test1
3776 Projection: test2.a, test2.b
3777 TableScan: test2
3778 ",
3779 );
3780 assert_optimized_plan_equal!(
3782 plan,
3783 @r"
3784 RightSemi Join: test1.a = test2.a
3785 TableScan: test1, full_filters=[test1.b > UInt32(1)]
3786 Projection: test2.a, test2.b
3787 TableScan: test2, full_filters=[test2.b > UInt32(2)]
3788 "
3789 )
3790 }
3791
3792 #[test]
3793 fn left_anti_join() -> Result<()> {
3794 let table_scan = test_table_scan_with_name("test1")?;
3795 let left = LogicalPlanBuilder::from(table_scan)
3796 .project(vec![col("a"), col("b")])?
3797 .build()?;
3798 let right_table_scan = test_table_scan_with_name("test2")?;
3799 let right = LogicalPlanBuilder::from(right_table_scan)
3800 .project(vec![col("a"), col("b")])?
3801 .build()?;
3802 let plan = LogicalPlanBuilder::from(left)
3803 .join(
3804 right,
3805 JoinType::LeftAnti,
3806 (
3807 vec![Column::from_qualified_name("test1.a")],
3808 vec![Column::from_qualified_name("test2.a")],
3809 ),
3810 None,
3811 )?
3812 .filter(col("test2.a").gt(lit(2u32)))?
3813 .build()?;
3814
3815 assert_snapshot!(plan,
3817 @r"
3818 Filter: test2.a > UInt32(2)
3819 LeftAnti Join: test1.a = test2.a
3820 Projection: test1.a, test1.b
3821 TableScan: test1
3822 Projection: test2.a, test2.b
3823 TableScan: test2
3824 ",
3825 );
3826 assert_optimized_plan_equal!(
3828 plan,
3829 @r"
3830 Filter: test2.a > UInt32(2)
3831 LeftAnti Join: test1.a = test2.a
3832 Projection: test1.a, test1.b
3833 TableScan: test1, full_filters=[test1.a > UInt32(2)]
3834 Projection: test2.a, test2.b
3835 TableScan: test2
3836 "
3837 )
3838 }
3839
3840 #[test]
3846 fn null_aware_left_anti_join_no_inferred_pushdown() -> Result<()> {
3847 let table_scan = test_table_scan_with_name("test1")?;
3848 let left = LogicalPlanBuilder::from(table_scan)
3849 .project(vec![col("a"), col("b")])?
3850 .build()?;
3851 let right_table_scan = test_table_scan_with_name("test2")?;
3852 let right = LogicalPlanBuilder::from(right_table_scan)
3853 .project(vec![col("a"), col("b")])?
3854 .build()?;
3855 let plan = LogicalPlanBuilder::from(left)
3856 .join_detailed_with_options(
3857 right,
3858 JoinType::LeftAnti,
3859 (
3860 vec![Column::from_qualified_name("test1.a")],
3861 vec![Column::from_qualified_name("test2.a")],
3862 ),
3863 None,
3864 datafusion_common::NullEquality::NullEqualsNothing,
3865 true,
3866 )?
3867 .filter(col("test1.a").gt(lit(2u32)))?
3868 .build()?;
3869
3870 assert_optimized_plan_equal!(
3874 plan,
3875 @r"
3876 LeftAnti Join: test1.a = test2.a null_aware
3877 Projection: test1.a, test1.b
3878 TableScan: test1, full_filters=[test1.a > UInt32(2)]
3879 Projection: test2.a, test2.b
3880 TableScan: test2
3881 "
3882 )
3883 }
3884
3885 #[test]
3886 fn left_anti_join_with_filters() -> Result<()> {
3887 let table_scan = test_table_scan_with_name("test1")?;
3888 let left = LogicalPlanBuilder::from(table_scan)
3889 .project(vec![col("a"), col("b")])?
3890 .build()?;
3891 let right_table_scan = test_table_scan_with_name("test2")?;
3892 let right = LogicalPlanBuilder::from(right_table_scan)
3893 .project(vec![col("a"), col("b")])?
3894 .build()?;
3895 let plan = LogicalPlanBuilder::from(left)
3896 .join(
3897 right,
3898 JoinType::LeftAnti,
3899 (
3900 vec![Column::from_qualified_name("test1.a")],
3901 vec![Column::from_qualified_name("test2.a")],
3902 ),
3903 Some(
3904 col("test1.b")
3905 .gt(lit(1u32))
3906 .and(col("test2.b").gt(lit(2u32))),
3907 ),
3908 )?
3909 .build()?;
3910
3911 assert_snapshot!(plan,
3913 @r"
3914 LeftAnti Join: test1.a = test2.a Filter: test1.b > UInt32(1) AND test2.b > UInt32(2)
3915 Projection: test1.a, test1.b
3916 TableScan: test1
3917 Projection: test2.a, test2.b
3918 TableScan: test2
3919 ",
3920 );
3921 assert_optimized_plan_equal!(
3923 plan,
3924 @r"
3925 LeftAnti Join: test1.a = test2.a Filter: test1.b > UInt32(1)
3926 Projection: test1.a, test1.b
3927 TableScan: test1
3928 Projection: test2.a, test2.b
3929 TableScan: test2, full_filters=[test2.b > UInt32(2)]
3930 "
3931 )
3932 }
3933
3934 #[test]
3935 fn right_anti_join() -> Result<()> {
3936 let table_scan = test_table_scan_with_name("test1")?;
3937 let left = LogicalPlanBuilder::from(table_scan)
3938 .project(vec![col("a"), col("b")])?
3939 .build()?;
3940 let right_table_scan = test_table_scan_with_name("test2")?;
3941 let right = LogicalPlanBuilder::from(right_table_scan)
3942 .project(vec![col("a"), col("b")])?
3943 .build()?;
3944 let plan = LogicalPlanBuilder::from(left)
3945 .join(
3946 right,
3947 JoinType::RightAnti,
3948 (
3949 vec![Column::from_qualified_name("test1.a")],
3950 vec![Column::from_qualified_name("test2.a")],
3951 ),
3952 None,
3953 )?
3954 .filter(col("test1.a").gt(lit(2u32)))?
3955 .build()?;
3956
3957 assert_snapshot!(plan,
3959 @r"
3960 Filter: test1.a > UInt32(2)
3961 RightAnti Join: test1.a = test2.a
3962 Projection: test1.a, test1.b
3963 TableScan: test1
3964 Projection: test2.a, test2.b
3965 TableScan: test2
3966 ",
3967 );
3968 assert_optimized_plan_equal!(
3970 plan,
3971 @r"
3972 Filter: test1.a > UInt32(2)
3973 RightAnti Join: test1.a = test2.a
3974 Projection: test1.a, test1.b
3975 TableScan: test1
3976 Projection: test2.a, test2.b
3977 TableScan: test2, full_filters=[test2.a > UInt32(2)]
3978 "
3979 )
3980 }
3981
3982 #[test]
3983 fn right_anti_join_with_filters() -> Result<()> {
3984 let table_scan = test_table_scan_with_name("test1")?;
3985 let left = LogicalPlanBuilder::from(table_scan)
3986 .project(vec![col("a"), col("b")])?
3987 .build()?;
3988 let right_table_scan = test_table_scan_with_name("test2")?;
3989 let right = LogicalPlanBuilder::from(right_table_scan)
3990 .project(vec![col("a"), col("b")])?
3991 .build()?;
3992 let plan = LogicalPlanBuilder::from(left)
3993 .join(
3994 right,
3995 JoinType::RightAnti,
3996 (
3997 vec![Column::from_qualified_name("test1.a")],
3998 vec![Column::from_qualified_name("test2.a")],
3999 ),
4000 Some(
4001 col("test1.b")
4002 .gt(lit(1u32))
4003 .and(col("test2.b").gt(lit(2u32))),
4004 ),
4005 )?
4006 .build()?;
4007
4008 assert_snapshot!(plan,
4010 @r"
4011 RightAnti Join: test1.a = test2.a Filter: test1.b > UInt32(1) AND test2.b > UInt32(2)
4012 Projection: test1.a, test1.b
4013 TableScan: test1
4014 Projection: test2.a, test2.b
4015 TableScan: test2
4016 ",
4017 );
4018 assert_optimized_plan_equal!(
4020 plan,
4021 @r"
4022 RightAnti Join: test1.a = test2.a Filter: test2.b > UInt32(2)
4023 Projection: test1.a, test1.b
4024 TableScan: test1, full_filters=[test1.b > UInt32(1)]
4025 Projection: test2.a, test2.b
4026 TableScan: test2
4027 "
4028 )
4029 }
4030
4031 #[derive(Debug, PartialEq, Eq, Hash)]
4032 struct TestScalarUDF {
4033 signature: Signature,
4034 }
4035
4036 impl ScalarUDFImpl for TestScalarUDF {
4037 fn name(&self) -> &str {
4038 "TestScalarUDF"
4039 }
4040
4041 fn signature(&self) -> &Signature {
4042 &self.signature
4043 }
4044
4045 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
4046 Ok(DataType::Int32)
4047 }
4048
4049 fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
4050 Ok(ColumnarValue::Scalar(ScalarValue::from(1)))
4051 }
4052 }
4053
4054 #[test]
4055 fn test_push_down_volatile_function_in_aggregate() -> Result<()> {
4056 let table_scan = test_table_scan_with_name("test1")?;
4058 let fun = ScalarUDF::new_from_impl(TestScalarUDF {
4059 signature: Signature::exact(vec![], Volatility::Volatile),
4060 });
4061 let expr = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(fun), vec![]));
4062
4063 let plan = LogicalPlanBuilder::from(table_scan)
4064 .aggregate(vec![col("a")], vec![sum(col("b"))])?
4065 .project(vec![col("a"), sum(col("b")), add(expr, lit(1)).alias("r")])?
4066 .alias("t")?
4067 .filter(col("t.a").gt(lit(5)).and(col("t.r").gt(lit(0.5))))?
4068 .project(vec![col("t.a"), col("t.r")])?
4069 .build()?;
4070
4071 assert_snapshot!(plan,
4072 @r"
4073 Projection: t.a, t.r
4074 Filter: t.a > Int32(5) AND t.r > Float64(0.5)
4075 SubqueryAlias: t
4076 Projection: test1.a, sum(test1.b), TestScalarUDF() + Int32(1) AS r
4077 Aggregate: groupBy=[[test1.a]], aggr=[[sum(test1.b)]]
4078 TableScan: test1
4079 ",
4080 );
4081 assert_optimized_plan_equal!(
4082 plan,
4083 @r"
4084 Projection: t.a, t.r
4085 SubqueryAlias: t
4086 Filter: r > Float64(0.5)
4087 Projection: test1.a, sum(test1.b), TestScalarUDF() + Int32(1) AS r
4088 Aggregate: groupBy=[[test1.a]], aggr=[[sum(test1.b)]]
4089 TableScan: test1, full_filters=[test1.a > Int32(5)]
4090 "
4091 )
4092 }
4093
4094 #[test]
4095 fn test_push_down_volatile_function_in_join() -> Result<()> {
4096 let table_scan = test_table_scan_with_name("test1")?;
4098 let fun = ScalarUDF::new_from_impl(TestScalarUDF {
4099 signature: Signature::exact(vec![], Volatility::Volatile),
4100 });
4101 let expr = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(fun), vec![]));
4102 let left = LogicalPlanBuilder::from(table_scan).build()?;
4103 let right_table_scan = test_table_scan_with_name("test2")?;
4104 let right = LogicalPlanBuilder::from(right_table_scan).build()?;
4105 let plan = LogicalPlanBuilder::from(left)
4106 .join(
4107 right,
4108 JoinType::Inner,
4109 (
4110 vec![Column::from_qualified_name("test1.a")],
4111 vec![Column::from_qualified_name("test2.a")],
4112 ),
4113 None,
4114 )?
4115 .project(vec![col("test1.a").alias("a"), expr.alias("r")])?
4116 .alias("t")?
4117 .filter(col("t.r").gt(lit(0.8)))?
4118 .project(vec![col("t.a"), col("t.r")])?
4119 .build()?;
4120
4121 assert_snapshot!(plan,
4122 @r"
4123 Projection: t.a, t.r
4124 Filter: t.r > Float64(0.8)
4125 SubqueryAlias: t
4126 Projection: test1.a AS a, TestScalarUDF() AS r
4127 Inner Join: test1.a = test2.a
4128 TableScan: test1
4129 TableScan: test2
4130 ",
4131 );
4132 assert_optimized_plan_equal!(
4133 plan,
4134 @r"
4135 Projection: t.a, t.r
4136 SubqueryAlias: t
4137 Filter: r > Float64(0.8)
4138 Projection: test1.a AS a, TestScalarUDF() AS r
4139 Inner Join: test1.a = test2.a
4140 TableScan: test1
4141 TableScan: test2
4142 "
4143 )
4144 }
4145
4146 #[test]
4147 fn test_push_down_volatile_table_scan() -> Result<()> {
4148 let table_scan = test_table_scan()?;
4150 let fun = ScalarUDF::new_from_impl(TestScalarUDF {
4151 signature: Signature::exact(vec![], Volatility::Volatile),
4152 });
4153 let expr = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(fun), vec![]));
4154 let plan = LogicalPlanBuilder::from(table_scan)
4155 .project(vec![col("a"), col("b")])?
4156 .filter(expr.gt(lit(0.1)))?
4157 .build()?;
4158
4159 assert_snapshot!(plan,
4160 @r"
4161 Filter: TestScalarUDF() > Float64(0.1)
4162 Projection: test.a, test.b
4163 TableScan: test
4164 ",
4165 );
4166 assert_optimized_plan_equal!(
4167 plan,
4168 @r"
4169 Projection: test.a, test.b
4170 Filter: TestScalarUDF() > Float64(0.1)
4171 TableScan: test
4172 "
4173 )
4174 }
4175
4176 #[test]
4177 fn test_push_down_volatile_mixed_table_scan() -> Result<()> {
4178 let table_scan = test_table_scan()?;
4180 let fun = ScalarUDF::new_from_impl(TestScalarUDF {
4181 signature: Signature::exact(vec![], Volatility::Volatile),
4182 });
4183 let expr = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(fun), vec![]));
4184 let plan = LogicalPlanBuilder::from(table_scan)
4185 .project(vec![col("a"), col("b")])?
4186 .filter(
4187 expr.gt(lit(0.1))
4188 .and(col("t.a").gt(lit(5)))
4189 .and(col("t.b").gt(lit(10))),
4190 )?
4191 .build()?;
4192
4193 assert_snapshot!(plan,
4194 @r"
4195 Filter: TestScalarUDF() > Float64(0.1) AND t.a > Int32(5) AND t.b > Int32(10)
4196 Projection: test.a, test.b
4197 TableScan: test
4198 ",
4199 );
4200 assert_optimized_plan_equal!(
4201 plan,
4202 @r"
4203 Projection: test.a, test.b
4204 Filter: TestScalarUDF() > Float64(0.1)
4205 TableScan: test, full_filters=[t.a > Int32(5), t.b > Int32(10)]
4206 "
4207 )
4208 }
4209
4210 #[test]
4211 fn test_push_down_volatile_mixed_unsupported_table_scan() -> Result<()> {
4212 let fun = ScalarUDF::new_from_impl(TestScalarUDF {
4214 signature: Signature::exact(vec![], Volatility::Volatile),
4215 });
4216 let expr = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(fun), vec![]));
4217 let plan = table_scan_with_pushdown_provider_builder(
4218 TableProviderFilterPushDown::Unsupported,
4219 vec![],
4220 None,
4221 )?
4222 .project(vec![col("a"), col("b")])?
4223 .filter(
4224 expr.gt(lit(0.1))
4225 .and(col("t.a").gt(lit(5)))
4226 .and(col("t.b").gt(lit(10))),
4227 )?
4228 .build()?;
4229
4230 assert_snapshot!(plan,
4231 @r"
4232 Filter: TestScalarUDF() > Float64(0.1) AND t.a > Int32(5) AND t.b > Int32(10)
4233 Projection: a, b
4234 TableScan: test
4235 ",
4236 );
4237 assert_optimized_plan_equal!(
4238 plan,
4239 @r"
4240 Projection: a, b
4241 Filter: TestScalarUDF() > Float64(0.1) AND t.a > Int32(5) AND t.b > Int32(10)
4242 TableScan: test
4243 "
4244 )
4245 }
4246
4247 #[test]
4248 fn test_push_down_filter_to_user_defined_node() -> Result<()> {
4249 #[derive(Debug, Hash, Eq, PartialEq)]
4251 struct TestUserNode {
4252 schema: DFSchemaRef,
4253 }
4254
4255 impl PartialOrd for TestUserNode {
4256 fn partial_cmp(&self, _other: &Self) -> Option<Ordering> {
4257 None
4258 }
4259 }
4260
4261 impl TestUserNode {
4262 fn new() -> Self {
4263 let schema = Arc::new(
4264 DFSchema::new_with_metadata(
4265 vec![(None, Field::new("a", DataType::Int64, false).into())],
4266 Default::default(),
4267 )
4268 .unwrap(),
4269 );
4270
4271 Self { schema }
4272 }
4273 }
4274
4275 impl UserDefinedLogicalNodeCore for TestUserNode {
4276 fn name(&self) -> &str {
4277 "test_node"
4278 }
4279
4280 fn inputs(&self) -> Vec<&LogicalPlan> {
4281 vec![]
4282 }
4283
4284 fn schema(&self) -> &DFSchemaRef {
4285 &self.schema
4286 }
4287
4288 fn expressions(&self) -> Vec<Expr> {
4289 vec![]
4290 }
4291
4292 fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
4293 write!(f, "TestUserNode")
4294 }
4295
4296 fn with_exprs_and_inputs(
4297 &self,
4298 exprs: Vec<Expr>,
4299 inputs: Vec<LogicalPlan>,
4300 ) -> Result<Self> {
4301 assert!(exprs.is_empty());
4302 assert!(inputs.is_empty());
4303 Ok(Self {
4304 schema: Arc::clone(&self.schema),
4305 })
4306 }
4307 }
4308
4309 let node = LogicalPlan::Extension(Extension {
4311 node: Arc::new(TestUserNode::new()),
4312 });
4313
4314 let plan = LogicalPlanBuilder::from(node).filter(lit(false))?.build()?;
4315
4316 assert_snapshot!(plan,
4318 @r"
4319 Filter: Boolean(false)
4320 TestUserNode
4321 ",
4322 );
4323 assert_optimized_plan_equal!(
4325 plan,
4326 @r"
4327 Filter: Boolean(false)
4328 TestUserNode
4329 "
4330 )
4331 }
4332
4333 #[test]
4339 fn filter_not_pushed_through_move_towards_leaves_projection() -> Result<()> {
4340 let table_scan = test_table_scan()?;
4341
4342 let proj = LogicalPlanBuilder::from(table_scan)
4344 .project(vec![
4345 leaf_udf_expr(col("a")).alias("val"),
4346 col("b"),
4347 col("c"),
4348 ])?
4349 .build()?;
4350
4351 let plan = LogicalPlanBuilder::from(proj)
4353 .filter(col("val").gt(lit(150i64)))?
4354 .build()?;
4355
4356 assert_optimized_plan_equal!(
4358 plan,
4359 @r"
4360 Filter: val > Int64(150)
4361 Projection: leaf_udf(test.a) AS val, test.b, test.c
4362 TableScan: test
4363 "
4364 )
4365 }
4366
4367 #[test]
4369 fn filter_mixed_predicates_partial_push() -> Result<()> {
4370 let table_scan = test_table_scan()?;
4371
4372 let proj = LogicalPlanBuilder::from(table_scan)
4374 .project(vec![
4375 leaf_udf_expr(col("a")).alias("val"),
4376 col("b"),
4377 col("c"),
4378 ])?
4379 .build()?;
4380
4381 let plan = LogicalPlanBuilder::from(proj)
4383 .filter(col("val").gt(lit(150i64)).and(col("b").gt(lit(5i64))))?
4384 .build()?;
4385
4386 assert_optimized_plan_equal!(
4388 plan,
4389 @r"
4390 Filter: val > Int64(150)
4391 Projection: leaf_udf(test.a) AS val, test.b, test.c
4392 TableScan: test, full_filters=[test.b > Int64(5)]
4393 "
4394 )
4395 }
4396
4397 #[test]
4398 fn filter_not_pushed_down_through_table_scan_with_fetch() -> Result<()> {
4399 let scan = test_table_scan()?;
4400 let scan_with_fetch = match scan {
4401 LogicalPlan::TableScan(scan) => LogicalPlan::TableScan(TableScan {
4402 fetch: Some(10),
4403 ..scan
4404 }),
4405 _ => unreachable!(),
4406 };
4407 let plan = LogicalPlanBuilder::from(scan_with_fetch)
4408 .filter(col("a").gt(lit(10i64)))?
4409 .build()?;
4410 assert_optimized_plan_equal!(
4412 plan,
4413 @r"
4414 Filter: test.a > Int64(10)
4415 TableScan: test, fetch=10
4416 "
4417 )
4418 }
4419
4420 #[test]
4421 fn filter_push_down_through_sort_without_fetch() -> Result<()> {
4422 let table_scan = test_table_scan()?;
4423 let plan = LogicalPlanBuilder::from(table_scan)
4424 .sort(vec![col("a").sort(true, true)])?
4425 .filter(col("a").gt(lit(10i64)))?
4426 .build()?;
4427 assert_optimized_plan_equal!(
4429 plan,
4430 @r"
4431 Sort: test.a ASC NULLS FIRST
4432 TableScan: test, full_filters=[test.a > Int64(10)]
4433 "
4434 )
4435 }
4436
4437 #[test]
4438 fn filter_not_pushed_down_through_sort_with_fetch() -> Result<()> {
4439 let table_scan = test_table_scan()?;
4440 let plan = LogicalPlanBuilder::from(table_scan)
4441 .sort_with_limit(vec![col("a").sort(true, true)], Some(5))?
4442 .filter(col("a").gt(lit(10i64)))?
4443 .build()?;
4444 assert_optimized_plan_equal!(
4447 plan,
4448 @r"
4449 Filter: test.a > Int64(10)
4450 Sort: test.a ASC NULLS FIRST, fetch=5
4451 TableScan: test
4452 "
4453 )
4454 }
4455}