1mod required_indices;
21
22use crate::optimizer::ApplyOrder;
23use crate::{OptimizerConfig, OptimizerRule};
24use std::sync::Arc;
25
26use datafusion_common::{
27 Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err,
28 get_required_group_by_exprs_indices, internal_datafusion_err, internal_err,
29};
30use datafusion_expr::expr::Alias;
31use datafusion_expr::{
32 Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScanBuilder, Unnest,
33 Window, logical_plan::LogicalPlan,
34};
35
36use crate::optimize_projections::required_indices::RequiredIndices;
37use crate::utils::NamePreserver;
38use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeContainer};
39
40#[derive(Default, Debug)]
74pub struct OptimizeProjections {}
75
76impl OptimizeProjections {
77 #[expect(missing_docs)]
78 pub fn new() -> Self {
79 Self {}
80 }
81}
82
83impl OptimizerRule for OptimizeProjections {
84 fn name(&self) -> &str {
85 "optimize_projections"
86 }
87
88 fn apply_order(&self) -> Option<ApplyOrder> {
89 None
90 }
91
92 fn supports_rewrite(&self) -> bool {
93 true
94 }
95
96 fn rewrite(
97 &self,
98 plan: LogicalPlan,
99 config: &dyn OptimizerConfig,
100 ) -> Result<Transformed<LogicalPlan>> {
101 let indices = RequiredIndices::new_for_all_exprs(&plan);
103 optimize_projections(plan, config, indices)
104 }
105}
106
107#[cfg_attr(feature = "recursive_protection", recursive::recursive)]
127fn optimize_projections(
128 plan: LogicalPlan,
129 config: &dyn OptimizerConfig,
130 indices: RequiredIndices,
131) -> Result<Transformed<LogicalPlan>> {
132 match plan {
135 LogicalPlan::Projection(proj) => {
136 return merge_consecutive_projections(proj)?
137 .transform_data(|proj| {
138 rewrite_projection_given_requirements(proj, config, &indices)
139 })?
140 .transform_data(|plan| optimize_subqueries(plan, config));
141 }
142 LogicalPlan::Aggregate(aggregate) => {
143 let n_group_exprs = aggregate.group_expr_len()?;
145 let (group_by_reqs, aggregate_reqs) = indices.split_off(n_group_exprs);
148
149 let new_group_bys = if aggregate
154 .input
155 .schema()
156 .functional_dependencies()
157 .is_empty()
158 {
159 aggregate.group_expr
160 } else {
161 let group_by_expr_existing = aggregate
162 .group_expr
163 .iter()
164 .map(|group_by_expr| group_by_expr.schema_name().to_string())
165 .collect::<Vec<_>>();
166
167 if let Some(simplest_groupby_indices) =
168 get_required_group_by_exprs_indices(
169 aggregate.input.schema(),
170 &group_by_expr_existing,
171 )
172 {
173 group_by_reqs
177 .append(&simplest_groupby_indices)
178 .get_at_indices(&aggregate.group_expr)
179 } else {
180 aggregate.group_expr
181 }
182 };
183
184 let new_aggr_expr = aggregate_reqs.get_at_indices(&aggregate.aggr_expr);
187
188 if new_group_bys.is_empty() && new_aggr_expr.is_empty() {
189 return Ok(Transformed::yes(LogicalPlan::EmptyRelation(
191 EmptyRelation {
192 produce_one_row: true,
193 schema: Arc::new(DFSchema::empty()),
194 },
195 )));
196 }
197
198 let all_exprs_iter = new_group_bys.iter().chain(new_aggr_expr.iter());
199 let schema = aggregate.input.schema();
200 let necessary_indices =
201 RequiredIndices::new().with_exprs(schema, all_exprs_iter);
202 let necessary_exprs = necessary_indices.get_required_exprs(schema);
203
204 return optimize_projections(
205 Arc::unwrap_or_clone(aggregate.input),
206 config,
207 necessary_indices,
208 )?
209 .transform_data(|aggregate_input| {
210 add_projection_on_top_if_helpful(aggregate_input, necessary_exprs)
215 })?
216 .map_data(|aggregate_input| {
217 Aggregate::try_new(
220 Arc::new(aggregate_input),
221 new_group_bys,
222 new_aggr_expr,
223 )
224 .map(LogicalPlan::Aggregate)
225 })?
226 .transform_data(|plan| optimize_subqueries(plan, config));
227 }
228 LogicalPlan::Window(window) => {
229 let input_schema = Arc::clone(window.input.schema());
230 let n_input_fields = input_schema.fields().len();
232 let (child_reqs, window_reqs) = indices.split_off(n_input_fields);
235
236 let new_window_expr = window_reqs.get_at_indices(&window.window_expr);
239
240 let required_indices = child_reqs.with_exprs(&input_schema, &new_window_expr);
243
244 return optimize_projections(
245 Arc::unwrap_or_clone(window.input),
246 config,
247 required_indices.clone(),
248 )?
249 .transform_data(|window_child| {
250 if new_window_expr.is_empty() {
251 Ok(Transformed::no(window_child))
253 } else {
254 let required_exprs =
258 required_indices.get_required_exprs(&input_schema);
259 let window_child =
260 add_projection_on_top_if_helpful(window_child, required_exprs)?
261 .data;
262 Window::try_new(new_window_expr, Arc::new(window_child))
263 .map(LogicalPlan::Window)
264 .map(Transformed::yes)
265 }
266 })?
267 .transform_data(|plan| optimize_subqueries(plan, config));
268 }
269 LogicalPlan::TableScan(table_scan) => {
270 let projection = match &table_scan.projection {
273 Some(projection) => indices.into_mapped_indices(|idx| projection[idx]),
274 None => indices.into_inner(),
275 };
276 let new_scan = TableScanBuilder::from(table_scan)
277 .with_projection(Some(projection))
278 .build()?;
279
280 return Transformed::yes(LogicalPlan::TableScan(new_scan))
281 .transform_data(|plan| optimize_subqueries(plan, config));
282 }
283 _ => {}
285 };
286
287 let mut child_required_indices: Vec<RequiredIndices> = match &plan {
290 LogicalPlan::Sort(_)
291 | LogicalPlan::Filter(_)
292 | LogicalPlan::Repartition(_)
293 | LogicalPlan::Union(_)
294 | LogicalPlan::SubqueryAlias(_)
295 | LogicalPlan::Distinct(Distinct::On(_)) => {
296 plan.inputs()
301 .into_iter()
302 .map(|input| {
303 indices
304 .clone()
305 .with_projection_beneficial()
306 .with_plan_exprs(&plan, input.schema())
307 })
308 .collect::<Result<_>>()?
309 }
310 LogicalPlan::Limit(_) => {
311 plan.inputs()
316 .into_iter()
317 .map(|input| indices.clone().with_plan_exprs(&plan, input.schema()))
318 .collect::<Result<_>>()?
319 }
320 LogicalPlan::Copy(_)
321 | LogicalPlan::Ddl(_)
322 | LogicalPlan::Dml(_)
323 | LogicalPlan::Explain(_)
324 | LogicalPlan::Analyze(_)
325 | LogicalPlan::Subquery(_)
326 | LogicalPlan::Statement(_)
327 | LogicalPlan::Distinct(Distinct::All(_)) => {
328 plan.inputs()
334 .into_iter()
335 .map(RequiredIndices::new_for_all_exprs)
336 .collect()
337 }
338 LogicalPlan::Extension(extension) => {
339 if let Some(necessary_children_indices) =
340 extension.node.necessary_children_exprs(indices.indices())
341 {
342 let children = extension.node.inputs();
343 assert_eq_or_internal_err!(
344 children.len(),
345 necessary_children_indices.len(),
346 "Inconsistent length between children and necessary children indices. \
347 Make sure `.necessary_children_exprs` implementation of the \
348 `UserDefinedLogicalNode` is consistent with actual children length \
349 for the node."
350 );
351 children
352 .into_iter()
353 .zip(necessary_children_indices)
354 .map(|(child, necessary_indices)| {
355 RequiredIndices::new_from_indices(necessary_indices)
356 .with_plan_exprs(&plan, child.schema())
357 })
358 .collect::<Result<Vec<_>>>()?
359 } else {
360 plan.inputs()
363 .into_iter()
364 .map(RequiredIndices::new_for_all_exprs)
365 .collect()
366 }
367 }
368 LogicalPlan::EmptyRelation(_)
369 | LogicalPlan::Values(_)
370 | LogicalPlan::DescribeTable(_) => {
371 return Ok(Transformed::no(plan));
373 }
374 LogicalPlan::RecursiveQuery(_) => {
375 return plan.map_children(|c| {
380 let indices = RequiredIndices::new_for_all_exprs(&c);
381 optimize_projections(c, config, indices)
382 });
383 }
384 LogicalPlan::Join(join) => {
385 let left_len = join.left.schema().fields().len();
386 let right_len = join.right.schema().fields().len();
387 let (left_req_indices, right_req_indices) =
388 split_join_requirements(left_len, right_len, indices, &join.join_type);
389 let mut left_indices =
390 left_req_indices.with_plan_exprs(&plan, join.left.schema())?;
391 let mut right_indices =
392 right_req_indices.with_plan_exprs(&plan, join.right.schema())?;
393 match join.join_type {
395 JoinType::LeftMark if right_indices.indices().is_empty() => {
396 right_indices = right_indices.append(&[0]);
397 }
398 JoinType::RightMark if left_indices.indices().is_empty() => {
399 left_indices = left_indices.append(&[0]);
400 }
401 _ => {}
402 }
403 vec![
406 left_indices.with_projection_beneficial(),
407 right_indices.with_projection_beneficial(),
408 ]
409 }
410 LogicalPlan::Projection(_)
412 | LogicalPlan::Aggregate(_)
413 | LogicalPlan::Window(_)
414 | LogicalPlan::TableScan(_) => {
415 return internal_err!(
416 "OptimizeProjection: should have handled in the match statement above"
417 );
418 }
419 LogicalPlan::Unnest(Unnest {
420 input,
421 dependency_indices,
422 ..
423 }) => {
424 let required_indices =
426 RequiredIndices::new().with_plan_exprs(&plan, input.schema())?;
427
428 let mut additional_necessary_child_indices = Vec::new();
430 indices.indices().iter().for_each(|idx| {
431 if let Some(index) = dependency_indices.get(*idx) {
432 additional_necessary_child_indices.push(*index);
433 }
434 });
435 vec![required_indices.append(&additional_necessary_child_indices)]
436 }
437 };
438
439 child_required_indices.reverse();
442 assert_eq_or_internal_err!(
443 child_required_indices.len(),
444 plan.inputs().len(),
445 "OptimizeProjection: child_required_indices length mismatch with plan inputs"
446 );
447
448 let transformed_plan = plan.map_children(|child| {
450 let required_indices = child_required_indices.pop().ok_or_else(|| {
451 internal_datafusion_err!(
452 "Unexpected number of required_indices in OptimizeProjections rule"
453 )
454 })?;
455
456 let projection_beneficial = required_indices.projection_beneficial();
457 let project_exprs = required_indices.get_required_exprs(child.schema());
458
459 optimize_projections(child, config, required_indices)?.transform_data(
460 |new_input| {
461 if projection_beneficial {
462 add_projection_on_top_if_helpful(new_input, project_exprs)
463 } else {
464 Ok(Transformed::no(new_input))
465 }
466 },
467 )
468 })?;
469
470 let transformed_plan =
471 transformed_plan.transform_data(|plan| optimize_subqueries(plan, config))?;
472
473 if transformed_plan.transformed {
475 transformed_plan.map_data(|plan| plan.recompute_schema())
476 } else {
477 Ok(transformed_plan)
478 }
479}
480
481fn optimize_subqueries(
485 plan: LogicalPlan,
486 config: &dyn OptimizerConfig,
487) -> Result<Transformed<LogicalPlan>> {
488 plan.map_uncorrelated_subqueries(|subquery_plan| {
489 let indices = RequiredIndices::new_for_all_exprs(&subquery_plan);
490 optimize_projections(subquery_plan, config, indices)
491 })
492}
493
494fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Projection>> {
525 let mut current = proj;
528 let mut transformed_any = false;
529 loop {
530 let Transformed {
531 data, transformed, ..
532 } = merge_consecutive_projections_one_level(current)?;
533 current = data;
534 if !transformed {
535 break;
536 }
537 transformed_any = true;
538 }
539 Ok(if transformed_any {
540 Transformed::yes(current)
541 } else {
542 Transformed::no(current)
543 })
544}
545
546fn merge_consecutive_projections_one_level(
547 proj: Projection,
548) -> Result<Transformed<Projection>> {
549 let Projection {
550 expr,
551 input,
552 schema,
553 ..
554 } = proj;
555 let LogicalPlan::Projection(prev_projection) = input.as_ref() else {
556 return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
557 };
558
559 if prev_projection.expr == expr {
562 return Projection::try_new_with_schema(
563 expr,
564 Arc::clone(&prev_projection.input),
565 schema,
566 )
567 .map(Transformed::yes);
568 }
569
570 let mut column_referral_map = HashMap::<&Column, usize>::new();
572 expr.iter()
573 .for_each(|expr| expr.add_column_ref_counts(&mut column_referral_map));
574
575 if column_referral_map.into_iter().any(|(col, usage)| {
579 usage > 1
580 && !prev_projection.expr[prev_projection.schema.index_of_column(col).unwrap()]
581 .placement()
582 .should_push_to_leaves()
583 }) {
584 return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
586 }
587
588 let LogicalPlan::Projection(prev_projection) = Arc::unwrap_or_clone(input) else {
589 unreachable!();
591 };
592
593 let name_preserver = NamePreserver::new_for_projection();
596 let mut original_names = vec![];
597 let new_exprs = expr.map_elements(|expr| {
598 original_names.push(name_preserver.save(&expr));
599
600 match expr {
602 Expr::Alias(Alias {
603 expr,
604 relation,
605 name,
606 metadata,
607 }) => rewrite_expr(*expr, &prev_projection).map(|result| {
608 result.update_data(|expr| {
609 if metadata.is_none() && expr.schema_name().to_string() == name {
616 expr
617 } else {
618 Expr::Alias(Alias {
619 expr: Box::new(expr),
620 relation,
621 name,
622 metadata,
623 })
624 }
625 })
626 }),
627 e => rewrite_expr(e, &prev_projection),
628 }
629 })?;
630
631 if new_exprs.transformed {
634 let new_exprs = new_exprs
636 .data
637 .into_iter()
638 .zip(original_names)
639 .map(|(expr, original_name)| original_name.restore(expr))
640 .collect::<Vec<_>>();
641 Projection::try_new(new_exprs, prev_projection.input).map(Transformed::yes)
642 } else {
643 let input = Arc::new(LogicalPlan::Projection(prev_projection));
645 Projection::try_new_with_schema(new_exprs.data, input, schema)
646 .map(Transformed::no)
647 }
648}
649
650fn rewrite_expr(expr: Expr, input: &Projection) -> Result<Transformed<Expr>> {
696 expr.transform_up(|expr| {
697 match expr {
698 Expr::Alias(alias) => {
700 match alias
701 .metadata
702 .as_ref()
703 .map(|h| h.is_empty())
704 .unwrap_or(true)
705 {
706 true => Ok(Transformed::yes(*alias.expr)),
707 false => Ok(Transformed::no(Expr::Alias(alias))),
708 }
709 }
710 Expr::Column(col) => {
711 let idx = input.schema.index_of_column(&col)?;
713 let input_expr = input.expr[idx].clone().unalias_nested().data;
721 Ok(Transformed::yes(input_expr))
722 }
723 _ => Ok(Transformed::no(expr)),
725 }
726 })
727}
728
729fn split_join_requirements(
760 left_len: usize,
761 right_len: usize,
762 indices: RequiredIndices,
763 join_type: &JoinType,
764) -> (RequiredIndices, RequiredIndices) {
765 match join_type {
766 JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {
768 indices.split_off(left_len)
771 }
772 JoinType::LeftMark => {
773 let (left_indices, _mark) = indices.split_off(left_len);
777 (left_indices, RequiredIndices::new())
778 }
779 JoinType::RightMark => {
780 let (right_indices, _mark) = indices.split_off(right_len);
782 (RequiredIndices::new(), right_indices)
783 }
784 JoinType::LeftAnti | JoinType::LeftSemi => (indices, RequiredIndices::new()),
786 JoinType::RightSemi | JoinType::RightAnti => (RequiredIndices::new(), indices),
789 }
790}
791
792fn add_projection_on_top_if_helpful(
810 plan: LogicalPlan,
811 project_exprs: Vec<Expr>,
812) -> Result<Transformed<LogicalPlan>> {
813 if project_exprs.len() >= plan.schema().fields().len() {
815 Ok(Transformed::no(plan))
816 } else {
817 Projection::try_new(project_exprs, Arc::new(plan))
818 .map(LogicalPlan::Projection)
819 .map(Transformed::yes)
820 }
821}
822
823fn rewrite_projection_given_requirements(
841 proj: Projection,
842 config: &dyn OptimizerConfig,
843 indices: &RequiredIndices,
844) -> Result<Transformed<LogicalPlan>> {
845 let Projection { expr, input, .. } = proj;
846
847 let exprs_used = indices.get_at_indices(&expr);
848
849 let required_indices =
850 RequiredIndices::new().with_exprs(input.schema(), exprs_used.iter());
851
852 optimize_projections(Arc::unwrap_or_clone(input), config, required_indices)?
855 .transform_data(|input| {
856 if is_projection_unnecessary(&input, &exprs_used)? {
857 Ok(Transformed::yes(input))
858 } else {
859 Projection::try_new(exprs_used, Arc::new(input))
860 .map(LogicalPlan::Projection)
861 .map(Transformed::yes)
862 }
863 })
864}
865
866pub fn is_projection_unnecessary(
870 input: &LogicalPlan,
871 proj_exprs: &[Expr],
872) -> Result<bool> {
873 if proj_exprs.len() != input.schema().fields().len() {
875 return Ok(false);
876 }
877 Ok(input.schema().iter().zip(proj_exprs.iter()).all(
878 |((field_relation, field_name), expr)| {
879 if let Expr::Column(col) = expr {
881 col.relation.as_ref() == field_relation && col.name.eq(field_name.name())
882 } else {
883 false
884 }
885 },
886 ))
887}
888
889#[cfg(test)]
890mod tests {
891 use std::cmp::Ordering;
892 use std::collections::HashMap;
893 use std::fmt::Formatter;
894 use std::ops::Add;
895 use std::sync::Arc;
896 use std::vec;
897
898 use crate::optimize_projections::OptimizeProjections;
899 use crate::optimizer::Optimizer;
900 use crate::test::{
901 assert_fields_eq, scan_empty, test_table_scan, test_table_scan_fields,
902 test_table_scan_with_name,
903 };
904 use crate::{OptimizerContext, OptimizerRule};
905 use arrow::datatypes::{DataType, Field, Schema};
906 use datafusion_common::{
907 Column, DFSchema, DFSchemaRef, JoinType, Result, TableReference,
908 };
909 use datafusion_expr::ExprFunctionExt;
910 use datafusion_expr::{
911 BinaryExpr, Expr, Extension, Like, LogicalPlan, Operator, Projection,
912 UserDefinedLogicalNodeCore, WindowFunctionDefinition, binary_expr,
913 build_join_schema,
914 builder::table_scan_with_filters,
915 col,
916 expr::{self, Cast},
917 lit,
918 logical_plan::{builder::LogicalPlanBuilder, table_scan},
919 not, try_cast, when,
920 };
921 use insta::assert_snapshot;
922
923 use crate::assert_optimized_plan_eq_snapshot;
924 use datafusion_functions_aggregate::count::count_udaf;
925 use datafusion_functions_aggregate::expr_fn::{count, max, min};
926 use datafusion_functions_aggregate::min_max::max_udaf;
927
928 macro_rules! assert_optimized_plan_equal {
929 (
930 $plan:expr,
931 @ $expected:literal $(,)?
932 ) => {{
933 let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
934 let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(OptimizeProjections::new())];
935 assert_optimized_plan_eq_snapshot!(
936 optimizer_ctx,
937 rules,
938 $plan,
939 @ $expected,
940 )
941 }};
942 }
943
944 #[derive(Debug, Hash, PartialEq, Eq)]
945 struct NoOpUserDefined {
946 exprs: Vec<Expr>,
947 schema: DFSchemaRef,
948 input: Arc<LogicalPlan>,
949 }
950
951 impl NoOpUserDefined {
952 fn new(schema: DFSchemaRef, input: Arc<LogicalPlan>) -> Self {
953 Self {
954 exprs: vec![],
955 schema,
956 input,
957 }
958 }
959
960 fn with_exprs(mut self, exprs: Vec<Expr>) -> Self {
961 self.exprs = exprs;
962 self
963 }
964 }
965
966 impl PartialOrd for NoOpUserDefined {
968 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
969 match self.exprs.partial_cmp(&other.exprs) {
970 Some(Ordering::Equal) => self.input.partial_cmp(&other.input),
971 cmp => cmp,
972 }
973 .filter(|cmp| *cmp != Ordering::Equal || self == other)
975 }
976 }
977
978 impl UserDefinedLogicalNodeCore for NoOpUserDefined {
979 fn name(&self) -> &str {
980 "NoOpUserDefined"
981 }
982
983 fn inputs(&self) -> Vec<&LogicalPlan> {
984 vec![&self.input]
985 }
986
987 fn schema(&self) -> &DFSchemaRef {
988 &self.schema
989 }
990
991 fn expressions(&self) -> Vec<Expr> {
992 self.exprs.clone()
993 }
994
995 fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
996 write!(f, "NoOpUserDefined")
997 }
998
999 fn with_exprs_and_inputs(
1000 &self,
1001 exprs: Vec<Expr>,
1002 mut inputs: Vec<LogicalPlan>,
1003 ) -> Result<Self> {
1004 Ok(Self {
1005 exprs,
1006 input: Arc::new(inputs.swap_remove(0)),
1007 schema: Arc::clone(&self.schema),
1008 })
1009 }
1010
1011 fn necessary_children_exprs(
1012 &self,
1013 output_columns: &[usize],
1014 ) -> Option<Vec<Vec<usize>>> {
1015 Some(vec![output_columns.to_vec()])
1017 }
1018
1019 fn supports_limit_pushdown(&self) -> bool {
1020 false }
1022 }
1023
1024 #[derive(Debug, Hash, PartialEq, Eq)]
1025 struct UserDefinedCrossJoin {
1026 exprs: Vec<Expr>,
1027 schema: DFSchemaRef,
1028 left_child: Arc<LogicalPlan>,
1029 right_child: Arc<LogicalPlan>,
1030 }
1031
1032 impl UserDefinedCrossJoin {
1033 fn new(left_child: Arc<LogicalPlan>, right_child: Arc<LogicalPlan>) -> Self {
1034 let left_schema = left_child.schema();
1035 let right_schema = right_child.schema();
1036 let schema = Arc::new(
1037 build_join_schema(left_schema, right_schema, &JoinType::Inner).unwrap(),
1038 );
1039 Self {
1040 exprs: vec![],
1041 schema,
1042 left_child,
1043 right_child,
1044 }
1045 }
1046 }
1047
1048 impl PartialOrd for UserDefinedCrossJoin {
1050 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1051 match self.exprs.partial_cmp(&other.exprs) {
1052 Some(Ordering::Equal) => {
1053 match self.left_child.partial_cmp(&other.left_child) {
1054 Some(Ordering::Equal) => {
1055 self.right_child.partial_cmp(&other.right_child)
1056 }
1057 cmp => cmp,
1058 }
1059 }
1060 cmp => cmp,
1061 }
1062 .filter(|cmp| *cmp != Ordering::Equal || self == other)
1064 }
1065 }
1066
1067 impl UserDefinedLogicalNodeCore for UserDefinedCrossJoin {
1068 fn name(&self) -> &str {
1069 "UserDefinedCrossJoin"
1070 }
1071
1072 fn inputs(&self) -> Vec<&LogicalPlan> {
1073 vec![&self.left_child, &self.right_child]
1074 }
1075
1076 fn schema(&self) -> &DFSchemaRef {
1077 &self.schema
1078 }
1079
1080 fn expressions(&self) -> Vec<Expr> {
1081 self.exprs.clone()
1082 }
1083
1084 fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
1085 write!(f, "UserDefinedCrossJoin")
1086 }
1087
1088 fn with_exprs_and_inputs(
1089 &self,
1090 exprs: Vec<Expr>,
1091 mut inputs: Vec<LogicalPlan>,
1092 ) -> Result<Self> {
1093 assert_eq!(inputs.len(), 2);
1094 Ok(Self {
1095 exprs,
1096 left_child: Arc::new(inputs.remove(0)),
1097 right_child: Arc::new(inputs.remove(0)),
1098 schema: Arc::clone(&self.schema),
1099 })
1100 }
1101
1102 fn necessary_children_exprs(
1103 &self,
1104 output_columns: &[usize],
1105 ) -> Option<Vec<Vec<usize>>> {
1106 let left_child_len = self.left_child.schema().fields().len();
1107 let mut left_reqs = vec![];
1108 let mut right_reqs = vec![];
1109 for &out_idx in output_columns {
1110 if out_idx < left_child_len {
1111 left_reqs.push(out_idx);
1112 } else {
1113 right_reqs.push(out_idx - left_child_len)
1116 }
1117 }
1118 Some(vec![left_reqs, right_reqs])
1119 }
1120
1121 fn supports_limit_pushdown(&self) -> bool {
1122 false }
1124 }
1125
1126 #[derive(Debug, Hash, PartialEq, Eq)]
1130 struct OpaqueRequirementsUserDefined {
1131 input: Arc<LogicalPlan>,
1132 schema: DFSchemaRef,
1133 }
1134
1135 impl PartialOrd for OpaqueRequirementsUserDefined {
1137 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1138 self.input
1139 .partial_cmp(&other.input)
1140 .filter(|cmp| *cmp != Ordering::Equal || self == other)
1141 }
1142 }
1143
1144 impl UserDefinedLogicalNodeCore for OpaqueRequirementsUserDefined {
1145 fn name(&self) -> &str {
1146 "OpaqueRequirementsUserDefined"
1147 }
1148
1149 fn inputs(&self) -> Vec<&LogicalPlan> {
1150 vec![&self.input]
1151 }
1152
1153 fn schema(&self) -> &DFSchemaRef {
1154 &self.schema
1155 }
1156
1157 fn expressions(&self) -> Vec<Expr> {
1158 vec![]
1159 }
1160
1161 fn with_exprs_and_inputs(
1162 &self,
1163 _exprs: Vec<Expr>,
1164 mut inputs: Vec<LogicalPlan>,
1165 ) -> Result<Self> {
1166 Ok(Self {
1167 input: Arc::new(inputs.swap_remove(0)),
1168 schema: Arc::clone(&self.schema),
1169 })
1170 }
1171
1172 fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
1173 write!(f, "OpaqueRequirementsUserDefined")
1174 }
1175 }
1176
1177 #[test]
1178 fn merge_two_projection() -> Result<()> {
1179 let table_scan = test_table_scan()?;
1180 let plan = LogicalPlanBuilder::from(table_scan)
1181 .project(vec![col("a")])?
1182 .project(vec![binary_expr(lit(1), Operator::Plus, col("a"))])?
1183 .build()?;
1184
1185 assert_optimized_plan_equal!(
1186 plan,
1187 @r"
1188 Projection: Int32(1) + test.a
1189 TableScan: test projection=[a]
1190 "
1191 )
1192 }
1193
1194 #[test]
1195 fn merge_three_projection() -> Result<()> {
1196 let table_scan = test_table_scan()?;
1197 let plan = LogicalPlanBuilder::from(table_scan)
1198 .project(vec![col("a"), col("b")])?
1199 .project(vec![col("a")])?
1200 .project(vec![binary_expr(lit(1), Operator::Plus, col("a"))])?
1201 .build()?;
1202
1203 assert_optimized_plan_equal!(
1204 plan,
1205 @r"
1206 Projection: Int32(1) + test.a
1207 TableScan: test projection=[a]
1208 "
1209 )
1210 }
1211
1212 #[test]
1213 fn merge_alias() -> Result<()> {
1214 let table_scan = test_table_scan()?;
1215 let plan = LogicalPlanBuilder::from(table_scan)
1216 .project(vec![col("a")])?
1217 .project(vec![col("a").alias("alias")])?
1218 .build()?;
1219
1220 assert_optimized_plan_equal!(
1221 plan,
1222 @r"
1223 Projection: test.a AS alias
1224 TableScan: test projection=[a]
1225 "
1226 )
1227 }
1228
1229 #[test]
1230 fn merge_nested_alias() -> Result<()> {
1231 let table_scan = test_table_scan()?;
1232 let plan = LogicalPlanBuilder::from(table_scan)
1233 .project(vec![col("a").alias("alias1").alias("alias2")])?
1234 .project(vec![col("alias2").alias("alias")])?
1235 .build()?;
1236
1237 assert_optimized_plan_equal!(
1238 plan,
1239 @r"
1240 Projection: test.a AS alias
1241 TableScan: test projection=[a]
1242 "
1243 )
1244 }
1245
1246 #[test]
1247 fn test_nested_count() -> Result<()> {
1248 let schema = Schema::new(vec![Field::new("foo", DataType::Int32, false)]);
1249
1250 let groups: Vec<Expr> = vec![];
1251
1252 let plan = table_scan(TableReference::none(), &schema, None)
1253 .unwrap()
1254 .aggregate(groups.clone(), vec![count(lit(1))])
1255 .unwrap()
1256 .aggregate(groups, vec![count(lit(1))])
1257 .unwrap()
1258 .build()
1259 .unwrap();
1260
1261 assert_optimized_plan_equal!(
1262 plan,
1263 @r"
1264 Aggregate: groupBy=[[]], aggr=[[count(Int32(1))]]
1265 EmptyRelation: rows=1
1266 "
1267 )
1268 }
1269
1270 #[test]
1271 fn test_neg_push_down() -> Result<()> {
1272 let table_scan = test_table_scan()?;
1273 let plan = LogicalPlanBuilder::from(table_scan)
1274 .project(vec![-col("a")])?
1275 .build()?;
1276
1277 assert_optimized_plan_equal!(
1278 plan,
1279 @r"
1280 Projection: (- test.a)
1281 TableScan: test projection=[a]
1282 "
1283 )
1284 }
1285
1286 #[test]
1287 fn test_is_null() -> Result<()> {
1288 let table_scan = test_table_scan()?;
1289 let plan = LogicalPlanBuilder::from(table_scan)
1290 .project(vec![col("a").is_null()])?
1291 .build()?;
1292
1293 assert_optimized_plan_equal!(
1294 plan,
1295 @r"
1296 Projection: test.a IS NULL
1297 TableScan: test projection=[a]
1298 "
1299 )
1300 }
1301
1302 #[test]
1303 fn test_is_not_null() -> Result<()> {
1304 let table_scan = test_table_scan()?;
1305 let plan = LogicalPlanBuilder::from(table_scan)
1306 .project(vec![col("a").is_not_null()])?
1307 .build()?;
1308
1309 assert_optimized_plan_equal!(
1310 plan,
1311 @r"
1312 Projection: test.a IS NOT NULL
1313 TableScan: test projection=[a]
1314 "
1315 )
1316 }
1317
1318 #[test]
1319 fn test_is_true() -> Result<()> {
1320 let table_scan = test_table_scan()?;
1321 let plan = LogicalPlanBuilder::from(table_scan)
1322 .project(vec![col("a").is_true()])?
1323 .build()?;
1324
1325 assert_optimized_plan_equal!(
1326 plan,
1327 @r"
1328 Projection: test.a IS TRUE
1329 TableScan: test projection=[a]
1330 "
1331 )
1332 }
1333
1334 #[test]
1335 fn test_is_not_true() -> Result<()> {
1336 let table_scan = test_table_scan()?;
1337 let plan = LogicalPlanBuilder::from(table_scan)
1338 .project(vec![col("a").is_not_true()])?
1339 .build()?;
1340
1341 assert_optimized_plan_equal!(
1342 plan,
1343 @r"
1344 Projection: test.a IS NOT TRUE
1345 TableScan: test projection=[a]
1346 "
1347 )
1348 }
1349
1350 #[test]
1351 fn test_is_false() -> Result<()> {
1352 let table_scan = test_table_scan()?;
1353 let plan = LogicalPlanBuilder::from(table_scan)
1354 .project(vec![col("a").is_false()])?
1355 .build()?;
1356
1357 assert_optimized_plan_equal!(
1358 plan,
1359 @r"
1360 Projection: test.a IS FALSE
1361 TableScan: test projection=[a]
1362 "
1363 )
1364 }
1365
1366 #[test]
1367 fn test_is_not_false() -> Result<()> {
1368 let table_scan = test_table_scan()?;
1369 let plan = LogicalPlanBuilder::from(table_scan)
1370 .project(vec![col("a").is_not_false()])?
1371 .build()?;
1372
1373 assert_optimized_plan_equal!(
1374 plan,
1375 @r"
1376 Projection: test.a IS NOT FALSE
1377 TableScan: test projection=[a]
1378 "
1379 )
1380 }
1381
1382 #[test]
1383 fn test_is_unknown() -> Result<()> {
1384 let table_scan = test_table_scan()?;
1385 let plan = LogicalPlanBuilder::from(table_scan)
1386 .project(vec![col("a").is_unknown()])?
1387 .build()?;
1388
1389 assert_optimized_plan_equal!(
1390 plan,
1391 @r"
1392 Projection: test.a IS UNKNOWN
1393 TableScan: test projection=[a]
1394 "
1395 )
1396 }
1397
1398 #[test]
1399 fn test_is_not_unknown() -> Result<()> {
1400 let table_scan = test_table_scan()?;
1401 let plan = LogicalPlanBuilder::from(table_scan)
1402 .project(vec![col("a").is_not_unknown()])?
1403 .build()?;
1404
1405 assert_optimized_plan_equal!(
1406 plan,
1407 @r"
1408 Projection: test.a IS NOT UNKNOWN
1409 TableScan: test projection=[a]
1410 "
1411 )
1412 }
1413
1414 #[test]
1415 fn test_not() -> Result<()> {
1416 let table_scan = test_table_scan()?;
1417 let plan = LogicalPlanBuilder::from(table_scan)
1418 .project(vec![not(col("a"))])?
1419 .build()?;
1420
1421 assert_optimized_plan_equal!(
1422 plan,
1423 @r"
1424 Projection: NOT test.a
1425 TableScan: test projection=[a]
1426 "
1427 )
1428 }
1429
1430 #[test]
1431 fn test_try_cast() -> Result<()> {
1432 let table_scan = test_table_scan()?;
1433 let plan = LogicalPlanBuilder::from(table_scan)
1434 .project(vec![try_cast(col("a"), DataType::Float64)])?
1435 .build()?;
1436
1437 assert_optimized_plan_equal!(
1438 plan,
1439 @r"
1440 Projection: TRY_CAST(test.a AS Float64)
1441 TableScan: test projection=[a]
1442 "
1443 )
1444 }
1445
1446 #[test]
1447 fn test_similar_to() -> Result<()> {
1448 let table_scan = test_table_scan()?;
1449 let expr = Box::new(col("a"));
1450 let pattern = Box::new(lit("[0-9]"));
1451 let similar_to_expr =
1452 Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
1453 let plan = LogicalPlanBuilder::from(table_scan)
1454 .project(vec![similar_to_expr])?
1455 .build()?;
1456
1457 assert_optimized_plan_equal!(
1458 plan,
1459 @r#"
1460 Projection: test.a SIMILAR TO Utf8("[0-9]")
1461 TableScan: test projection=[a]
1462 "#
1463 )
1464 }
1465
1466 #[test]
1467 fn test_between() -> Result<()> {
1468 let table_scan = test_table_scan()?;
1469 let plan = LogicalPlanBuilder::from(table_scan)
1470 .project(vec![col("a").between(lit(1), lit(3))])?
1471 .build()?;
1472
1473 assert_optimized_plan_equal!(
1474 plan,
1475 @r"
1476 Projection: test.a BETWEEN Int32(1) AND Int32(3)
1477 TableScan: test projection=[a]
1478 "
1479 )
1480 }
1481
1482 #[test]
1484 fn test_case_merged() -> Result<()> {
1485 let table_scan = test_table_scan()?;
1486 let plan = LogicalPlanBuilder::from(table_scan)
1487 .project(vec![col("a"), lit(0).alias("d")])?
1488 .project(vec![
1489 col("a"),
1490 when(col("a").eq(lit(1)), lit(10))
1491 .otherwise(col("d"))?
1492 .alias("d"),
1493 ])?
1494 .build()?;
1495
1496 assert_optimized_plan_equal!(
1497 plan,
1498 @r"
1499 Projection: test.a, CASE WHEN test.a = Int32(1) THEN Int32(10) ELSE Int32(0) END AS d
1500 TableScan: test projection=[a]
1501 "
1502 )
1503 }
1504
1505 #[test]
1508 fn test_derived_column() -> Result<()> {
1509 let table_scan = test_table_scan()?;
1510 let plan = LogicalPlanBuilder::from(table_scan)
1511 .project(vec![col("a").add(lit(1)).alias("a"), lit(0).alias("d")])?
1512 .project(vec![
1513 col("a"),
1514 when(col("a").eq(lit(1)), lit(10))
1515 .otherwise(col("d"))?
1516 .alias("d"),
1517 ])?
1518 .build()?;
1519
1520 assert_optimized_plan_equal!(
1521 plan,
1522 @r"
1523 Projection: a, CASE WHEN a = Int32(1) THEN Int32(10) ELSE d END AS d
1524 Projection: test.a + Int32(1) AS a, Int32(0) AS d
1525 TableScan: test projection=[a]
1526 "
1527 )
1528 }
1529
1530 #[test]
1533 fn test_user_defined_logical_plan_node() -> Result<()> {
1534 let table_scan = test_table_scan()?;
1535 let custom_plan = LogicalPlan::Extension(Extension {
1536 node: Arc::new(NoOpUserDefined::new(
1537 Arc::clone(table_scan.schema()),
1538 Arc::new(table_scan.clone()),
1539 )),
1540 });
1541 let plan = LogicalPlanBuilder::from(custom_plan)
1542 .project(vec![col("a"), lit(0).alias("d")])?
1543 .build()?;
1544
1545 assert_optimized_plan_equal!(
1546 plan,
1547 @r"
1548 Projection: test.a, Int32(0) AS d
1549 NoOpUserDefined
1550 TableScan: test projection=[a]
1551 "
1552 )
1553 }
1554
1555 #[test]
1560 fn test_user_defined_logical_plan_node2() -> Result<()> {
1561 let table_scan = test_table_scan()?;
1562 let exprs = vec![Expr::Column(Column::from_qualified_name("b"))];
1563 let custom_plan = LogicalPlan::Extension(Extension {
1564 node: Arc::new(
1565 NoOpUserDefined::new(
1566 Arc::clone(table_scan.schema()),
1567 Arc::new(table_scan.clone()),
1568 )
1569 .with_exprs(exprs),
1570 ),
1571 });
1572 let plan = LogicalPlanBuilder::from(custom_plan)
1573 .project(vec![col("a"), lit(0).alias("d")])?
1574 .build()?;
1575
1576 assert_optimized_plan_equal!(
1577 plan,
1578 @r"
1579 Projection: test.a, Int32(0) AS d
1580 NoOpUserDefined
1581 TableScan: test projection=[a, b]
1582 "
1583 )
1584 }
1585
1586 #[test]
1592 fn test_user_defined_logical_plan_node3() -> Result<()> {
1593 let table_scan = test_table_scan()?;
1594 let left_expr = Expr::Column(Column::from_qualified_name("b"));
1595 let right_expr = Expr::Column(Column::from_qualified_name("c"));
1596 let binary_expr = Expr::BinaryExpr(BinaryExpr::new(
1597 Box::new(left_expr),
1598 Operator::Plus,
1599 Box::new(right_expr),
1600 ));
1601 let exprs = vec![binary_expr];
1602 let custom_plan = LogicalPlan::Extension(Extension {
1603 node: Arc::new(
1604 NoOpUserDefined::new(
1605 Arc::clone(table_scan.schema()),
1606 Arc::new(table_scan.clone()),
1607 )
1608 .with_exprs(exprs),
1609 ),
1610 });
1611 let plan = LogicalPlanBuilder::from(custom_plan)
1612 .project(vec![col("a"), lit(0).alias("d")])?
1613 .build()?;
1614
1615 assert_optimized_plan_equal!(
1616 plan,
1617 @r"
1618 Projection: test.a, Int32(0) AS d
1619 NoOpUserDefined
1620 TableScan: test projection=[a, b, c]
1621 "
1622 )
1623 }
1624
1625 #[test]
1630 fn test_user_defined_logical_plan_node4() -> Result<()> {
1631 let left_table = test_table_scan_with_name("l")?;
1632 let right_table = test_table_scan_with_name("r")?;
1633 let custom_plan = LogicalPlan::Extension(Extension {
1634 node: Arc::new(UserDefinedCrossJoin::new(
1635 Arc::new(left_table),
1636 Arc::new(right_table),
1637 )),
1638 });
1639 let plan = LogicalPlanBuilder::from(custom_plan)
1640 .project(vec![col("l.a"), col("l.c"), col("r.a"), lit(0).alias("d")])?
1641 .build()?;
1642
1643 assert_optimized_plan_equal!(
1644 plan,
1645 @r"
1646 Projection: l.a, l.c, r.a, Int32(0) AS d
1647 UserDefinedCrossJoin
1648 TableScan: l projection=[a, c]
1649 TableScan: r projection=[a]
1650 "
1651 )
1652 }
1653
1654 #[test]
1655 fn aggregate_no_group_by() -> Result<()> {
1656 let table_scan = test_table_scan()?;
1657
1658 let plan = LogicalPlanBuilder::from(table_scan)
1659 .aggregate(Vec::<Expr>::new(), vec![max(col("b"))])?
1660 .build()?;
1661
1662 assert_optimized_plan_equal!(
1663 plan,
1664 @r"
1665 Aggregate: groupBy=[[]], aggr=[[max(test.b)]]
1666 TableScan: test projection=[b]
1667 "
1668 )
1669 }
1670
1671 #[test]
1672 fn aggregate_group_by() -> Result<()> {
1673 let table_scan = test_table_scan()?;
1674
1675 let plan = LogicalPlanBuilder::from(table_scan)
1676 .aggregate(vec![col("c")], vec![max(col("b"))])?
1677 .build()?;
1678
1679 assert_optimized_plan_equal!(
1680 plan,
1681 @r"
1682 Aggregate: groupBy=[[test.c]], aggr=[[max(test.b)]]
1683 TableScan: test projection=[b, c]
1684 "
1685 )
1686 }
1687
1688 #[test]
1689 fn aggregate_group_by_with_table_alias() -> Result<()> {
1690 let table_scan = test_table_scan()?;
1691
1692 let plan = LogicalPlanBuilder::from(table_scan)
1693 .alias("a")?
1694 .aggregate(vec![col("c")], vec![max(col("b"))])?
1695 .build()?;
1696
1697 assert_optimized_plan_equal!(
1698 plan,
1699 @r"
1700 Aggregate: groupBy=[[a.c]], aggr=[[max(a.b)]]
1701 SubqueryAlias: a
1702 TableScan: test projection=[b, c]
1703 "
1704 )
1705 }
1706
1707 #[test]
1708 fn aggregate_no_group_by_with_filter() -> Result<()> {
1709 let table_scan = test_table_scan()?;
1710
1711 let plan = LogicalPlanBuilder::from(table_scan)
1712 .filter(col("c").gt(lit(1)))?
1713 .aggregate(Vec::<Expr>::new(), vec![max(col("b"))])?
1714 .build()?;
1715
1716 assert_optimized_plan_equal!(
1717 plan,
1718 @r"
1719 Aggregate: groupBy=[[]], aggr=[[max(test.b)]]
1720 Projection: test.b
1721 Filter: test.c > Int32(1)
1722 TableScan: test projection=[b, c]
1723 "
1724 )
1725 }
1726
1727 #[test]
1728 fn aggregate_with_periods() -> Result<()> {
1729 let schema = Schema::new(vec![Field::new("tag.one", DataType::Utf8, false)]);
1730
1731 let plan = table_scan(Some("m4"), &schema, None)?
1738 .aggregate(
1739 Vec::<Expr>::new(),
1740 vec![max(col(Column::new_unqualified("tag.one"))).alias("tag.one")],
1741 )?
1742 .project([col(Column::new_unqualified("tag.one"))])?
1743 .build()?;
1744
1745 assert_optimized_plan_equal!(
1746 plan,
1747 @r"
1748 Aggregate: groupBy=[[]], aggr=[[max(m4.tag.one) AS tag.one]]
1749 TableScan: m4 projection=[tag.one]
1750 "
1751 )
1752 }
1753
1754 #[test]
1755 fn redundant_project() -> Result<()> {
1756 let table_scan = test_table_scan()?;
1757
1758 let plan = LogicalPlanBuilder::from(table_scan)
1759 .project(vec![col("a"), col("b"), col("c")])?
1760 .project(vec![col("a"), col("c"), col("b")])?
1761 .build()?;
1762 assert_optimized_plan_equal!(
1763 plan,
1764 @r"
1765 Projection: test.a, test.c, test.b
1766 TableScan: test projection=[a, b, c]
1767 "
1768 )
1769 }
1770
1771 #[test]
1772 fn reorder_scan() -> Result<()> {
1773 let schema = Schema::new(test_table_scan_fields());
1774
1775 let plan = table_scan(Some("test"), &schema, Some(vec![1, 0, 2]))?.build()?;
1776 assert_optimized_plan_equal!(
1777 plan,
1778 @"TableScan: test projection=[b, a, c]"
1779 )
1780 }
1781
1782 #[test]
1783 fn reorder_scan_projection() -> Result<()> {
1784 let schema = Schema::new(test_table_scan_fields());
1785
1786 let plan = table_scan(Some("test"), &schema, Some(vec![1, 0, 2]))?
1787 .project(vec![col("a"), col("b")])?
1788 .build()?;
1789 assert_optimized_plan_equal!(
1790 plan,
1791 @r"
1792 Projection: test.a, test.b
1793 TableScan: test projection=[b, a]
1794 "
1795 )
1796 }
1797
1798 #[test]
1799 fn reorder_projection() -> Result<()> {
1800 let table_scan = test_table_scan()?;
1801
1802 let plan = LogicalPlanBuilder::from(table_scan)
1803 .project(vec![col("c"), col("b"), col("a")])?
1804 .build()?;
1805 assert_optimized_plan_equal!(
1806 plan,
1807 @r"
1808 Projection: test.c, test.b, test.a
1809 TableScan: test projection=[a, b, c]
1810 "
1811 )
1812 }
1813
1814 #[test]
1815 fn noncontinuous_redundant_projection() -> Result<()> {
1816 let table_scan = test_table_scan()?;
1817
1818 let plan = LogicalPlanBuilder::from(table_scan)
1819 .project(vec![col("c"), col("b"), col("a")])?
1820 .filter(col("c").gt(lit(1)))?
1821 .project(vec![col("c"), col("a"), col("b")])?
1822 .filter(col("b").gt(lit(1)))?
1823 .filter(col("a").gt(lit(1)))?
1824 .project(vec![col("a"), col("c"), col("b")])?
1825 .build()?;
1826 assert_optimized_plan_equal!(
1827 plan,
1828 @r"
1829 Projection: test.a, test.c, test.b
1830 Filter: test.a > Int32(1)
1831 Filter: test.b > Int32(1)
1832 Projection: test.c, test.a, test.b
1833 Filter: test.c > Int32(1)
1834 Projection: test.c, test.b, test.a
1835 TableScan: test projection=[a, b, c]
1836 "
1837 )
1838 }
1839
1840 #[test]
1841 fn join_schema_trim_full_join_column_projection() -> Result<()> {
1842 let table_scan = test_table_scan()?;
1843
1844 let schema = Schema::new(vec![Field::new("c1", DataType::UInt32, false)]);
1845 let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
1846
1847 let plan = LogicalPlanBuilder::from(table_scan)
1848 .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
1849 .project(vec![col("a"), col("b"), col("c1")])?
1850 .build()?;
1851
1852 let optimized_plan = optimize(plan)?;
1853
1854 assert_snapshot!(
1856 optimized_plan.clone(),
1857 @r"
1858 Left Join: test.a = test2.c1
1859 TableScan: test projection=[a, b]
1860 TableScan: test2 projection=[c1]
1861 "
1862 );
1863
1864 let optimized_join = optimized_plan;
1866 assert_eq!(
1867 **optimized_join.schema(),
1868 DFSchema::new_with_metadata(
1869 vec![
1870 (
1871 Some("test".into()),
1872 Arc::new(Field::new("a", DataType::UInt32, false))
1873 ),
1874 (
1875 Some("test".into()),
1876 Arc::new(Field::new("b", DataType::UInt32, false))
1877 ),
1878 (
1879 Some("test2".into()),
1880 Arc::new(Field::new("c1", DataType::UInt32, true))
1881 ),
1882 ],
1883 HashMap::new()
1884 )?,
1885 );
1886
1887 Ok(())
1888 }
1889
1890 #[test]
1891 fn join_schema_trim_partial_join_column_projection() -> Result<()> {
1892 let table_scan = test_table_scan()?;
1895
1896 let schema = Schema::new(vec![Field::new("c1", DataType::UInt32, false)]);
1897 let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
1898
1899 let plan = LogicalPlanBuilder::from(table_scan)
1900 .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
1901 .project(vec![col("a"), col("b")])?
1904 .build()?;
1905
1906 let optimized_plan = optimize(plan)?;
1907
1908 assert_snapshot!(
1910 optimized_plan.clone(),
1911 @r"
1912 Projection: test.a, test.b
1913 Left Join: test.a = test2.c1
1914 TableScan: test projection=[a, b]
1915 TableScan: test2 projection=[c1]
1916 "
1917 );
1918
1919 let optimized_join = optimized_plan.inputs()[0];
1921 assert_eq!(
1922 **optimized_join.schema(),
1923 DFSchema::new_with_metadata(
1924 vec![
1925 (
1926 Some("test".into()),
1927 Arc::new(Field::new("a", DataType::UInt32, false))
1928 ),
1929 (
1930 Some("test".into()),
1931 Arc::new(Field::new("b", DataType::UInt32, false))
1932 ),
1933 (
1934 Some("test2".into()),
1935 Arc::new(Field::new("c1", DataType::UInt32, true))
1936 ),
1937 ],
1938 HashMap::new()
1939 )?,
1940 );
1941
1942 Ok(())
1943 }
1944
1945 #[test]
1946 fn join_schema_trim_using_join() -> Result<()> {
1947 let table_scan = test_table_scan()?;
1950
1951 let schema = Schema::new(vec![Field::new("a", DataType::UInt32, false)]);
1952 let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
1953
1954 let plan = LogicalPlanBuilder::from(table_scan)
1955 .join_using(table2_scan, JoinType::Left, vec!["a".into()])?
1956 .project(vec![col("a"), col("b")])?
1957 .build()?;
1958
1959 let optimized_plan = optimize(plan)?;
1960
1961 assert_snapshot!(
1963 optimized_plan.clone(),
1964 @r"
1965 Projection: test.a, test.b
1966 Left Join: Using test.a = test2.a
1967 TableScan: test projection=[a, b]
1968 TableScan: test2 projection=[a]
1969 "
1970 );
1971
1972 let optimized_join = optimized_plan.inputs()[0];
1974 assert_eq!(
1975 **optimized_join.schema(),
1976 DFSchema::new_with_metadata(
1977 vec![
1978 (
1979 Some("test".into()),
1980 Arc::new(Field::new("a", DataType::UInt32, false))
1981 ),
1982 (
1983 Some("test".into()),
1984 Arc::new(Field::new("b", DataType::UInt32, false))
1985 ),
1986 (
1987 Some("test2".into()),
1988 Arc::new(Field::new("a", DataType::UInt32, true))
1989 ),
1990 ],
1991 HashMap::new()
1992 )?,
1993 );
1994
1995 Ok(())
1996 }
1997
1998 #[test]
1999 fn cast() -> Result<()> {
2000 let table_scan = test_table_scan()?;
2001
2002 let plan = LogicalPlanBuilder::from(table_scan)
2003 .project(vec![Expr::Cast(Cast::new(
2004 Box::new(col("c")),
2005 DataType::Float64,
2006 ))])?
2007 .build()?;
2008
2009 assert_optimized_plan_equal!(
2010 plan,
2011 @r"
2012 Projection: CAST(test.c AS Float64)
2013 TableScan: test projection=[c]
2014 "
2015 )
2016 }
2017
2018 #[test]
2019 fn table_scan_projected_schema() -> Result<()> {
2020 let table_scan = test_table_scan()?;
2021 let plan = LogicalPlanBuilder::from(test_table_scan()?)
2022 .project(vec![col("a"), col("b")])?
2023 .build()?;
2024
2025 assert_eq!(3, table_scan.schema().fields().len());
2026 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2027 assert_fields_eq(&plan, vec!["a", "b"]);
2028
2029 assert_optimized_plan_equal!(
2030 plan,
2031 @"TableScan: test projection=[a, b]"
2032 )
2033 }
2034
2035 #[test]
2036 fn table_scan_projected_schema_non_qualified_relation() -> Result<()> {
2037 let table_scan = test_table_scan()?;
2038 let input_schema = table_scan.schema();
2039 assert_eq!(3, input_schema.fields().len());
2040 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2041
2042 let expr = vec![col("test.a"), col("test.b")];
2046 let plan =
2047 LogicalPlan::Projection(Projection::try_new(expr, Arc::new(table_scan))?);
2048
2049 assert_fields_eq(&plan, vec!["a", "b"]);
2050
2051 assert_optimized_plan_equal!(
2052 plan,
2053 @"TableScan: test projection=[a, b]"
2054 )
2055 }
2056
2057 #[test]
2058 fn table_limit() -> Result<()> {
2059 let table_scan = test_table_scan()?;
2060 assert_eq!(3, table_scan.schema().fields().len());
2061 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2062
2063 let plan = LogicalPlanBuilder::from(table_scan)
2064 .project(vec![col("c"), col("a")])?
2065 .limit(0, Some(5))?
2066 .build()?;
2067
2068 assert_fields_eq(&plan, vec!["c", "a"]);
2069
2070 assert_optimized_plan_equal!(
2071 plan,
2072 @r"
2073 Limit: skip=0, fetch=5
2074 Projection: test.c, test.a
2075 TableScan: test projection=[a, c]
2076 "
2077 )
2078 }
2079
2080 #[test]
2081 fn table_scan_without_projection() -> Result<()> {
2082 let table_scan = test_table_scan()?;
2083 let plan = LogicalPlanBuilder::from(table_scan).build()?;
2084 assert_optimized_plan_equal!(
2086 plan,
2087 @"TableScan: test projection=[a, b, c]"
2088 )
2089 }
2090
2091 #[test]
2092 fn table_scan_with_literal_projection() -> Result<()> {
2093 let table_scan = test_table_scan()?;
2094 let plan = LogicalPlanBuilder::from(table_scan)
2095 .project(vec![lit(1_i64), lit(2_i64)])?
2096 .build()?;
2097 assert_optimized_plan_equal!(
2098 plan,
2099 @r"
2100 Projection: Int64(1), Int64(2)
2101 TableScan: test projection=[]
2102 "
2103 )
2104 }
2105
2106 #[test]
2108 fn table_unused_column() -> Result<()> {
2109 let table_scan = test_table_scan()?;
2110 assert_eq!(3, table_scan.schema().fields().len());
2111 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2112
2113 let plan = LogicalPlanBuilder::from(table_scan)
2115 .project(vec![col("c"), col("a"), col("b")])?
2116 .filter(col("c").gt(lit(1)))?
2117 .aggregate(vec![col("c")], vec![max(col("a"))])?
2118 .build()?;
2119
2120 assert_fields_eq(&plan, vec!["c", "max(test.a)"]);
2121
2122 let plan = optimize(plan).expect("failed to optimize plan");
2123 assert_optimized_plan_equal!(
2124 plan,
2125 @r"
2126 Aggregate: groupBy=[[test.c]], aggr=[[max(test.a)]]
2127 Filter: test.c > Int32(1)
2128 Projection: test.c, test.a
2129 TableScan: test projection=[a, c]
2130 "
2131 )
2132 }
2133
2134 #[test]
2136 fn table_unused_projection() -> Result<()> {
2137 let table_scan = test_table_scan()?;
2138 assert_eq!(3, table_scan.schema().fields().len());
2139 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2140
2141 let plan = LogicalPlanBuilder::from(table_scan)
2143 .project(vec![col("b")])?
2144 .project(vec![lit(1).alias("a")])?
2145 .build()?;
2146
2147 assert_fields_eq(&plan, vec!["a"]);
2148
2149 assert_optimized_plan_equal!(
2150 plan,
2151 @r"
2152 Projection: Int32(1) AS a
2153 TableScan: test projection=[]
2154 "
2155 )
2156 }
2157
2158 #[test]
2159 fn table_full_filter_pushdown() -> Result<()> {
2160 let schema = Schema::new(test_table_scan_fields());
2161
2162 let table_scan = table_scan_with_filters(
2163 Some("test"),
2164 &schema,
2165 None,
2166 vec![col("b").eq(lit(1))],
2167 )?
2168 .build()?;
2169 assert_eq!(3, table_scan.schema().fields().len());
2170 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2171
2172 let plan = LogicalPlanBuilder::from(table_scan)
2174 .project(vec![col("b")])?
2175 .project(vec![lit(1).alias("a")])?
2176 .build()?;
2177
2178 assert_fields_eq(&plan, vec!["a"]);
2179
2180 assert_optimized_plan_equal!(
2181 plan,
2182 @r"
2183 Projection: Int32(1) AS a
2184 TableScan: test projection=[], full_filters=[b = Int32(1)]
2185 "
2186 )
2187 }
2188
2189 #[test]
2191 fn test_double_optimization() -> Result<()> {
2192 let table_scan = test_table_scan()?;
2193
2194 let plan = LogicalPlanBuilder::from(table_scan)
2195 .project(vec![col("b")])?
2196 .project(vec![lit(1).alias("a")])?
2197 .build()?;
2198
2199 let optimized_plan1 = optimize(plan).expect("failed to optimize plan");
2200 let optimized_plan2 =
2201 optimize(optimized_plan1.clone()).expect("failed to optimize plan");
2202
2203 let formatted_plan1 = format!("{optimized_plan1:?}");
2204 let formatted_plan2 = format!("{optimized_plan2:?}");
2205 assert_eq!(formatted_plan1, formatted_plan2);
2206 Ok(())
2207 }
2208
2209 #[test]
2210 fn test_continue_processing_through_extension() -> Result<()> {
2211 let table_scan = test_table_scan()?;
2212 let plan = LogicalPlanBuilder::from(table_scan.clone())
2213 .project(vec![col("a")])?
2214 .project(vec![col("a")])?
2215 .build()?;
2216 let plan = LogicalPlan::Extension(Extension {
2217 node: Arc::new(OpaqueRequirementsUserDefined {
2218 input: Arc::new(plan),
2219 schema: Arc::clone(table_scan.schema()),
2220 }),
2221 });
2222 let plan = optimize(plan).expect("failed to optimize plan");
2223 assert_optimized_plan_equal!(
2224 plan,
2225 @r"
2226 OpaqueRequirementsUserDefined
2227 TableScan: test projection=[a]
2228 "
2229 )
2230 }
2231
2232 #[test]
2234 fn table_unused_aggregate() -> Result<()> {
2235 let table_scan = test_table_scan()?;
2236 assert_eq!(3, table_scan.schema().fields().len());
2237 assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2238
2239 let plan = LogicalPlanBuilder::from(table_scan)
2241 .aggregate(vec![col("a"), col("c")], vec![max(col("b")), min(col("b"))])?
2242 .filter(col("c").gt(lit(1)))?
2243 .project(vec![col("c"), col("a"), col("max(test.b)")])?
2244 .build()?;
2245
2246 assert_fields_eq(&plan, vec!["c", "a", "max(test.b)"]);
2247
2248 assert_optimized_plan_equal!(
2249 plan,
2250 @r"
2251 Projection: test.c, test.a, max(test.b)
2252 Filter: test.c > Int32(1)
2253 Aggregate: groupBy=[[test.a, test.c]], aggr=[[max(test.b)]]
2254 TableScan: test projection=[a, b, c]
2255 "
2256 )
2257 }
2258
2259 #[test]
2260 fn aggregate_filter_pushdown() -> Result<()> {
2261 let table_scan = test_table_scan()?;
2262 let aggr_with_filter = count_udaf()
2263 .call(vec![col("b")])
2264 .filter(col("c").gt(lit(42)))
2265 .build()?;
2266 let plan = LogicalPlanBuilder::from(table_scan)
2267 .aggregate(
2268 vec![col("a")],
2269 vec![count(col("b")), aggr_with_filter.alias("count2")],
2270 )?
2271 .build()?;
2272
2273 assert_optimized_plan_equal!(
2274 plan,
2275 @r"
2276 Aggregate: groupBy=[[test.a]], aggr=[[count(test.b), count(test.b) FILTER (WHERE test.c > Int32(42)) AS count2]]
2277 TableScan: test projection=[a, b, c]
2278 "
2279 )
2280 }
2281
2282 #[test]
2283 fn pushdown_through_distinct() -> Result<()> {
2284 let table_scan = test_table_scan()?;
2285
2286 let plan = LogicalPlanBuilder::from(table_scan)
2287 .project(vec![col("a"), col("b")])?
2288 .distinct()?
2289 .project(vec![col("a")])?
2290 .build()?;
2291
2292 assert_optimized_plan_equal!(
2293 plan,
2294 @r"
2295 Projection: test.a
2296 Distinct:
2297 TableScan: test projection=[a, b]
2298 "
2299 )
2300 }
2301
2302 #[test]
2303 fn test_window() -> Result<()> {
2304 let table_scan = test_table_scan()?;
2305
2306 let max1 = Expr::from(expr::WindowFunction::new(
2307 WindowFunctionDefinition::AggregateUDF(max_udaf()),
2308 vec![col("test.a")],
2309 ))
2310 .partition_by(vec![col("test.b")])
2311 .build()
2312 .unwrap();
2313
2314 let max2 = Expr::from(expr::WindowFunction::new(
2315 WindowFunctionDefinition::AggregateUDF(max_udaf()),
2316 vec![col("test.b")],
2317 ));
2318 let col1 = col(max1.schema_name().to_string());
2319 let col2 = col(max2.schema_name().to_string());
2320
2321 let plan = LogicalPlanBuilder::from(table_scan)
2322 .window(vec![max1])?
2323 .window(vec![max2])?
2324 .project(vec![col1, col2])?
2325 .build()?;
2326
2327 assert_optimized_plan_equal!(
2328 plan,
2329 @r"
2330 Projection: max(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, max(test.b) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
2331 WindowAggr: windowExpr=[[max(test.b) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
2332 Projection: test.b, max(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
2333 WindowAggr: windowExpr=[[max(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
2334 TableScan: test projection=[a, b]
2335 "
2336 )
2337 }
2338
2339 #[test]
2343 fn optimize_projections_exists_or_exists_with_outer_join() -> Result<()> {
2344 use datafusion_expr::utils::disjunction;
2345 use datafusion_expr::{exists, out_ref_col};
2346
2347 let table_a = test_table_scan_with_name("a")?;
2348 let table_b = test_table_scan_with_name("b")?;
2349
2350 let sq_a = Arc::new(
2351 LogicalPlanBuilder::from(test_table_scan_with_name("sq_a")?)
2352 .filter(col("sq_a.a").eq(out_ref_col(DataType::UInt32, "a.a")))?
2353 .project(vec![lit(1)])?
2354 .build()?,
2355 );
2356
2357 let sq_b = Arc::new(
2358 LogicalPlanBuilder::from(test_table_scan_with_name("sq_b")?)
2359 .filter(col("sq_b.b").eq(out_ref_col(DataType::UInt32, "a.b")))?
2360 .project(vec![lit(1)])?
2361 .build()?,
2362 );
2363
2364 let plan = LogicalPlanBuilder::from(table_a)
2365 .filter(disjunction(vec![exists(sq_a), exists(sq_b)]).unwrap())?
2366 .join(table_b, JoinType::Left, (vec!["a"], vec!["a"]), None)?
2367 .build()?;
2368
2369 let optimizer = Optimizer::new();
2370 let config = OptimizerContext::new();
2371 optimizer.optimize(plan, &config, observe)?;
2372
2373 Ok(())
2374 }
2375
2376 #[test]
2377 fn optimize_projections_left_mark_join_with_projection() -> Result<()> {
2378 let table_a = test_table_scan_with_name("a")?;
2379 let table_b = test_table_scan_with_name("b")?;
2380 let table_c = test_table_scan_with_name("c")?;
2381
2382 let plan = LogicalPlanBuilder::from(table_a)
2383 .join(table_b, JoinType::LeftMark, (vec!["a"], vec!["a"]), None)?
2384 .project(vec![col("a.a"), col("a.b"), col("a.c")])?
2385 .join(table_c, JoinType::Left, (vec!["a"], vec!["a"]), None)?
2386 .build()?;
2387
2388 assert_optimized_plan_equal!(
2389 plan,
2390 @r"
2391 Left Join: a.a = c.a
2392 Projection: a.a, a.b, a.c
2393 LeftMark Join: a.a = b.a
2394 TableScan: a projection=[a, b, c]
2395 TableScan: b projection=[a]
2396 TableScan: c projection=[a, b, c]
2397 "
2398 )
2399 }
2400
2401 #[test]
2404 fn optimize_projections_stacked_mark_joins_keep_qualified_mark() -> Result<()> {
2405 let person = test_table_scan_with_name("person")?;
2406
2407 let aliased_scan = |table: &str, alias: &str| -> Result<LogicalPlan> {
2408 LogicalPlanBuilder::from(test_table_scan_with_name(table)?)
2409 .project(vec![col(format!("{table}.a"))])?
2410 .alias(alias)?
2411 .build()
2412 };
2413
2414 let plan = LogicalPlanBuilder::from(person)
2415 .join_on(
2416 aliased_scan("s1", "__correlated_sq_1")?,
2417 JoinType::LeftMark,
2418 vec![lit(true)],
2419 )?
2420 .join_on(
2421 aliased_scan("s2", "__correlated_sq_2")?,
2422 JoinType::LeftMark,
2423 vec![lit(true)],
2424 )?
2425 .join_on(
2426 aliased_scan("s3", "__correlated_sq_3")?,
2427 JoinType::LeftMark,
2428 vec![lit(true)],
2429 )?
2430 .filter(
2431 col("__correlated_sq_1.mark")
2432 .or(col("__correlated_sq_2.mark"))
2433 .and(not(col("__correlated_sq_3.mark"))),
2434 )?
2435 .project(vec![col("person.a")])?
2436 .build()?;
2437
2438 assert_optimized_plan_equal!(
2439 plan,
2440 @r"
2441 Projection: person.a
2442 Filter: (__correlated_sq_1.mark OR __correlated_sq_2.mark) AND NOT __correlated_sq_3.mark
2443 LeftMark Join: Filter: Boolean(true)
2444 LeftMark Join: Filter: Boolean(true)
2445 LeftMark Join: Filter: Boolean(true)
2446 TableScan: person projection=[a]
2447 SubqueryAlias: __correlated_sq_1
2448 TableScan: s1 projection=[a]
2449 SubqueryAlias: __correlated_sq_2
2450 TableScan: s2 projection=[a]
2451 SubqueryAlias: __correlated_sq_3
2452 TableScan: s3 projection=[a]
2453 "
2454 )
2455 }
2456
2457 fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
2458
2459 fn optimize(plan: LogicalPlan) -> Result<LogicalPlan> {
2460 let optimizer = Optimizer::with_rules(vec![Arc::new(OptimizeProjections::new())]);
2461 let optimized_plan =
2462 optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
2463 Ok(optimized_plan)
2464 }
2465}