1use std::cmp::Ordering;
21use std::collections::{BTreeSet, HashMap, HashSet};
22use std::fmt::{self, Debug, Display, Formatter};
23use std::hash::{Hash, Hasher};
24use std::sync::{Arc, LazyLock};
25
26use super::DdlStatement;
27use super::dml::CopyTo;
28use super::invariants::{
29 InvariantLevel, assert_always_invariants_at_current_node,
30 assert_executable_invariants,
31};
32use crate::builder::{unique_field_aliases, unnest_with_options};
33use crate::expr::{
34 Alias, Placeholder, Sort as SortExpr, WindowFunction, WindowFunctionParams,
35 intersect_metadata_for_union,
36};
37use crate::expr_rewriter::{
38 NamePreserver, create_col_from_scalar_expr, normalize_cols, normalize_sorts,
39};
40use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor};
41use crate::logical_plan::extension::UserDefinedLogicalNode;
42use crate::logical_plan::{DmlStatement, Statement, WriteOp};
43use crate::utils::{
44 check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields,
45 find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist,
46 merge_schema, split_conjunction,
47};
48use crate::{
49 BinaryExpr, CreateMemoryTable, CreateView, Execute, Expr, ExprSchemable, GroupingSet,
50 LogicalPlanBuilder, Operator, Prepare, TableProviderFilterPushDown, TableSource,
51 WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed,
52};
53
54use crate::statistics::StatisticsRequest;
55use arrow::compute::SortOptions;
56use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
57use datafusion_common::cse::{NormalizeEq, Normalizeable};
58use datafusion_common::format::{ExplainAnalyzeCategories, ExplainFormat, MetricType};
59use datafusion_common::metadata::check_metadata_with_storage_equal;
60use datafusion_common::tree_node::{
61 Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion,
62};
63use datafusion_common::{
64 Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Dependency,
65 FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result,
66 ScalarValue, Spans, SplitPoint, TableReference, UnnestOptions,
67 aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err,
68 internal_err, plan_err, validate_range_split_points,
69};
70use indexmap::IndexSet;
71use itertools::Itertools as _;
72
73use crate::display::PgJsonVisitor;
75pub use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan};
76pub use datafusion_common::{JoinConstraint, JoinType};
77
78#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
211pub enum LogicalPlan {
212 Projection(Projection),
215 Filter(Filter),
224 Window(Window),
230 Aggregate(Aggregate),
236 Sort(Sort),
239 Join(Join),
242 Repartition(Repartition),
246 Union(Union),
250 TableScan(TableScan),
253 EmptyRelation(EmptyRelation),
257 Subquery(Subquery),
260 SubqueryAlias(SubqueryAlias),
262 Limit(Limit),
264 Statement(Statement),
266 Values(Values),
271 Explain(Explain),
274 Analyze(Analyze),
278 Extension(Extension),
281 Distinct(Distinct),
284 Dml(DmlStatement),
286 Ddl(DdlStatement),
288 Copy(CopyTo),
290 DescribeTable(DescribeTable),
293 Unnest(Unnest),
296 RecursiveQuery(RecursiveQuery),
298}
299
300impl Default for LogicalPlan {
301 fn default() -> Self {
302 LogicalPlan::EmptyRelation(EmptyRelation {
306 produce_one_row: false,
307 schema: Arc::clone(DFSchema::empty_ref()),
308 })
309 }
310}
311
312impl<'a> TreeNodeContainer<'a, Self> for LogicalPlan {
313 fn apply_elements<F: FnMut(&'a Self) -> Result<TreeNodeRecursion>>(
314 &'a self,
315 mut f: F,
316 ) -> Result<TreeNodeRecursion> {
317 f(self)
318 }
319
320 fn map_elements<F: FnMut(Self) -> Result<Transformed<Self>>>(
321 self,
322 mut f: F,
323 ) -> Result<Transformed<Self>> {
324 f(self)
325 }
326}
327
328impl LogicalPlan {
329 pub fn schema(&self) -> &DFSchemaRef {
331 match self {
332 LogicalPlan::EmptyRelation(EmptyRelation { schema, .. }) => schema,
333 LogicalPlan::Values(Values { schema, .. }) => schema,
334 LogicalPlan::TableScan(TableScan {
335 projected_schema, ..
336 }) => projected_schema,
337 LogicalPlan::Projection(Projection { schema, .. }) => schema,
338 LogicalPlan::Filter(Filter { input, .. }) => input.schema(),
339 LogicalPlan::Distinct(Distinct::All(input)) => input.schema(),
340 LogicalPlan::Distinct(Distinct::On(DistinctOn { schema, .. })) => schema,
341 LogicalPlan::Window(Window { schema, .. }) => schema,
342 LogicalPlan::Aggregate(Aggregate { schema, .. }) => schema,
343 LogicalPlan::Sort(Sort { input, .. }) => input.schema(),
344 LogicalPlan::Join(Join { schema, .. }) => schema,
345 LogicalPlan::Repartition(Repartition { input, .. }) => input.schema(),
346 LogicalPlan::Limit(Limit { input, .. }) => input.schema(),
347 LogicalPlan::Statement(statement) => statement.schema(),
348 LogicalPlan::Subquery(Subquery { subquery, .. }) => subquery.schema(),
349 LogicalPlan::SubqueryAlias(SubqueryAlias { schema, .. }) => schema,
350 LogicalPlan::Explain(explain) => &explain.schema,
351 LogicalPlan::Analyze(analyze) => &analyze.schema,
352 LogicalPlan::Extension(extension) => extension.node.schema(),
353 LogicalPlan::Union(Union { schema, .. }) => schema,
354 LogicalPlan::DescribeTable(DescribeTable { output_schema, .. }) => {
355 output_schema
356 }
357 LogicalPlan::Dml(DmlStatement { output_schema, .. }) => output_schema,
358 LogicalPlan::Copy(CopyTo { output_schema, .. }) => output_schema,
359 LogicalPlan::Ddl(ddl) => ddl.schema(),
360 LogicalPlan::Unnest(Unnest { schema, .. }) => schema,
361 LogicalPlan::RecursiveQuery(RecursiveQuery { schema, .. }) => schema,
362 }
363 }
364
365 pub fn fallback_normalize_schemas(&self) -> Vec<&DFSchema> {
368 match self {
369 LogicalPlan::Window(_)
370 | LogicalPlan::Projection(_)
371 | LogicalPlan::Aggregate(_)
372 | LogicalPlan::Unnest(_)
373 | LogicalPlan::Join(_) => self
374 .inputs()
375 .iter()
376 .map(|input| input.schema().as_ref())
377 .collect(),
378 _ => vec![],
379 }
380 }
381
382 pub fn explain_schema() -> SchemaRef {
384 SchemaRef::new(Schema::new(vec![
385 Field::new("plan_type", DataType::Utf8, false),
386 Field::new("plan", DataType::Utf8, false),
387 ]))
388 }
389
390 pub fn describe_schema() -> Schema {
392 Schema::new(vec![
393 Field::new("column_name", DataType::Utf8, false),
394 Field::new("data_type", DataType::Utf8, false),
395 Field::new("is_nullable", DataType::Utf8, false),
396 ])
397 }
398
399 pub fn expressions(self: &LogicalPlan) -> Vec<Expr> {
416 let mut exprs = vec![];
417 self.apply_expressions(|e| {
418 exprs.push(e.clone());
419 Ok(TreeNodeRecursion::Continue)
420 })
421 .unwrap();
423 exprs
424 }
425
426 pub fn all_out_ref_exprs(self: &LogicalPlan) -> Vec<Expr> {
429 let mut exprs = vec![];
430 self.apply_expressions(|e| {
431 find_out_reference_exprs(e).into_iter().for_each(|e| {
432 if !exprs.contains(&e) {
433 exprs.push(e)
434 }
435 });
436 Ok(TreeNodeRecursion::Continue)
437 })
438 .unwrap();
440 self.inputs()
441 .into_iter()
442 .flat_map(|child| child.all_out_ref_exprs())
443 .for_each(|e| {
444 if !exprs.contains(&e) {
445 exprs.push(e)
446 }
447 });
448 exprs
449 }
450
451 pub fn inputs(&self) -> Vec<&LogicalPlan> {
455 match self {
456 LogicalPlan::Projection(Projection { input, .. }) => vec![input],
457 LogicalPlan::Filter(Filter { input, .. }) => vec![input],
458 LogicalPlan::Repartition(Repartition { input, .. }) => vec![input],
459 LogicalPlan::Window(Window { input, .. }) => vec![input],
460 LogicalPlan::Aggregate(Aggregate { input, .. }) => vec![input],
461 LogicalPlan::Sort(Sort { input, .. }) => vec![input],
462 LogicalPlan::Join(Join { left, right, .. }) => vec![left, right],
463 LogicalPlan::Limit(Limit { input, .. }) => vec![input],
464 LogicalPlan::Subquery(Subquery { subquery, .. }) => vec![subquery],
465 LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => vec![input],
466 LogicalPlan::Extension(extension) => extension.node.inputs(),
467 LogicalPlan::Union(Union { inputs, .. }) => {
468 inputs.iter().map(|arc| arc.as_ref()).collect()
469 }
470 LogicalPlan::Distinct(
471 Distinct::All(input) | Distinct::On(DistinctOn { input, .. }),
472 ) => vec![input],
473 LogicalPlan::Explain(explain) => vec![&explain.plan],
474 LogicalPlan::Analyze(analyze) => vec![&analyze.input],
475 LogicalPlan::Dml(write) => vec![&write.input],
476 LogicalPlan::Copy(copy) => vec![©.input],
477 LogicalPlan::Ddl(ddl) => ddl.inputs(),
478 LogicalPlan::Unnest(Unnest { input, .. }) => vec![input],
479 LogicalPlan::RecursiveQuery(RecursiveQuery {
480 static_term,
481 recursive_term,
482 ..
483 }) => vec![static_term, recursive_term],
484 LogicalPlan::Statement(stmt) => stmt.inputs(),
485 LogicalPlan::TableScan { .. }
487 | LogicalPlan::EmptyRelation { .. }
488 | LogicalPlan::Values { .. }
489 | LogicalPlan::DescribeTable(_) => vec![],
490 }
491 }
492
493 pub fn using_columns(&self) -> Result<Vec<HashSet<Column>>, DataFusionError> {
495 let mut using_columns: Vec<HashSet<Column>> = vec![];
496
497 self.apply_with_subqueries(|plan| {
498 if let LogicalPlan::Join(Join {
499 join_constraint: JoinConstraint::Using,
500 on,
501 ..
502 }) = plan
503 {
504 let columns =
506 on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| {
507 let Some(l) = l.get_as_join_column() else {
508 return internal_err!(
509 "Invalid join key. Expected column, found {l:?}"
510 );
511 };
512 let Some(r) = r.get_as_join_column() else {
513 return internal_err!(
514 "Invalid join key. Expected column, found {r:?}"
515 );
516 };
517 accumu.insert(l.to_owned());
518 accumu.insert(r.to_owned());
519 Result::<_, DataFusionError>::Ok(accumu)
520 })?;
521 using_columns.push(columns);
522 }
523 Ok(TreeNodeRecursion::Continue)
524 })?;
525
526 Ok(using_columns)
527 }
528
529 pub fn head_output_expr(&self) -> Result<Option<Expr>> {
531 match self {
532 LogicalPlan::Projection(projection) => {
533 Ok(Some(projection.expr.as_slice()[0].clone()))
534 }
535 LogicalPlan::Aggregate(agg) => {
536 if agg.group_expr.is_empty() {
537 Ok(Some(agg.aggr_expr.as_slice()[0].clone()))
538 } else {
539 Ok(Some(agg.group_expr.as_slice()[0].clone()))
540 }
541 }
542 LogicalPlan::Distinct(Distinct::On(DistinctOn { select_expr, .. })) => {
543 Ok(Some(select_expr[0].clone()))
544 }
545 LogicalPlan::Filter(Filter { input, .. })
546 | LogicalPlan::Distinct(Distinct::All(input))
547 | LogicalPlan::Sort(Sort { input, .. })
548 | LogicalPlan::Limit(Limit { input, .. })
549 | LogicalPlan::Repartition(Repartition { input, .. })
550 | LogicalPlan::Window(Window { input, .. }) => input.head_output_expr(),
551 LogicalPlan::Join(Join {
552 left,
553 right,
554 join_type,
555 ..
556 }) => match join_type {
557 JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {
558 if left.schema().fields().is_empty() {
559 right.head_output_expr()
560 } else {
561 left.head_output_expr()
562 }
563 }
564 JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => {
565 left.head_output_expr()
566 }
567 JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => {
568 right.head_output_expr()
569 }
570 },
571 LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => {
572 static_term.head_output_expr()
573 }
574 LogicalPlan::Union(union) => Ok(Some(Expr::Column(Column::from(
575 union.schema.qualified_field(0),
576 )))),
577 LogicalPlan::TableScan(table) => Ok(Some(Expr::Column(Column::from(
578 table.projected_schema.qualified_field(0),
579 )))),
580 LogicalPlan::SubqueryAlias(subquery_alias) => {
581 let expr_opt = subquery_alias.input.head_output_expr()?;
582 expr_opt
583 .map(|expr| {
584 Ok(Expr::Column(create_col_from_scalar_expr(
585 &expr,
586 subquery_alias.alias.to_string(),
587 )?))
588 })
589 .map_or(Ok(None), |v| v.map(Some))
590 }
591 LogicalPlan::Subquery(_) => Ok(None),
592 LogicalPlan::EmptyRelation(_)
593 | LogicalPlan::Statement(_)
594 | LogicalPlan::Values(_)
595 | LogicalPlan::Explain(_)
596 | LogicalPlan::Analyze(_)
597 | LogicalPlan::Extension(_)
598 | LogicalPlan::Dml(_)
599 | LogicalPlan::Copy(_)
600 | LogicalPlan::Ddl(_)
601 | LogicalPlan::DescribeTable(_)
602 | LogicalPlan::Unnest(_) => Ok(None),
603 }
604 }
605
606 pub fn recompute_schema(self) -> Result<Self> {
629 match self {
630 LogicalPlan::Projection(Projection {
633 expr,
634 input,
635 schema: _,
636 }) => Projection::try_new(expr, input).map(LogicalPlan::Projection),
637 LogicalPlan::Dml(_) => Ok(self),
638 LogicalPlan::Copy(_) => Ok(self),
639 LogicalPlan::Values(Values { schema, values }) => {
640 Ok(LogicalPlan::Values(Values { schema, values }))
642 }
643 LogicalPlan::Filter(Filter { predicate, input }) => {
644 Filter::try_new(predicate, input).map(LogicalPlan::Filter)
645 }
646 LogicalPlan::Repartition(_) => Ok(self),
647 LogicalPlan::Window(Window {
648 input,
649 window_expr,
650 schema: _,
651 }) => Window::try_new(window_expr, input).map(LogicalPlan::Window),
652 LogicalPlan::Aggregate(Aggregate {
653 input,
654 group_expr,
655 aggr_expr,
656 schema: _,
657 }) => Aggregate::try_new(input, group_expr, aggr_expr)
658 .map(LogicalPlan::Aggregate),
659 LogicalPlan::Sort(_) => Ok(self),
660 LogicalPlan::Join(Join {
661 left,
662 right,
663 filter,
664 join_type,
665 join_constraint,
666 on,
667 schema: _,
668 null_equality,
669 null_aware,
670 }) => {
671 let schema =
672 build_join_schema(left.schema(), right.schema(), &join_type)?;
673
674 let new_on: Vec<_> = on
675 .into_iter()
676 .map(|equi_expr| {
677 (equi_expr.0.unalias(), equi_expr.1.unalias())
679 })
680 .collect();
681
682 Ok(LogicalPlan::Join(Join {
683 left,
684 right,
685 join_type,
686 join_constraint,
687 on: new_on,
688 filter,
689 schema: DFSchemaRef::new(schema),
690 null_equality,
691 null_aware,
692 }))
693 }
694 LogicalPlan::Subquery(_) => Ok(self),
695 LogicalPlan::SubqueryAlias(SubqueryAlias {
696 input,
697 alias,
698 schema: _,
699 }) => SubqueryAlias::try_new(input, alias).map(LogicalPlan::SubqueryAlias),
700 LogicalPlan::Limit(_) => Ok(self),
701 LogicalPlan::Ddl(_) => Ok(self),
702 LogicalPlan::Extension(Extension { node }) => {
703 let expr = node.expressions();
706 let inputs: Vec<_> = node.inputs().into_iter().cloned().collect();
707 Ok(LogicalPlan::Extension(Extension {
708 node: node.with_exprs_and_inputs(expr, inputs)?,
709 }))
710 }
711 LogicalPlan::Union(Union { inputs, schema }) => {
712 let first_input_schema = inputs[0].schema();
713 if schema.fields().len() == first_input_schema.fields().len() {
714 Ok(LogicalPlan::Union(Union { inputs, schema }))
716 } else {
717 Ok(LogicalPlan::Union(Union::try_new(inputs)?))
725 }
726 }
727 LogicalPlan::Distinct(distinct) => {
728 let distinct = match distinct {
729 Distinct::All(input) => Distinct::All(input),
730 Distinct::On(DistinctOn {
731 on_expr,
732 select_expr,
733 sort_expr,
734 input,
735 schema: _,
736 }) => Distinct::On(DistinctOn::try_new(
737 on_expr,
738 select_expr,
739 sort_expr,
740 input,
741 )?),
742 };
743 Ok(LogicalPlan::Distinct(distinct))
744 }
745 LogicalPlan::RecursiveQuery(RecursiveQuery {
746 name,
747 static_term,
748 recursive_term,
749 is_distinct,
750 schema: _,
751 }) => RecursiveQuery::try_new(name, static_term, recursive_term, is_distinct)
752 .map(LogicalPlan::RecursiveQuery),
753 LogicalPlan::Analyze(_) => Ok(self),
754 LogicalPlan::Explain(_) => Ok(self),
755 LogicalPlan::TableScan(_) => Ok(self),
756 LogicalPlan::EmptyRelation(_) => Ok(self),
757 LogicalPlan::Statement(_) => Ok(self),
758 LogicalPlan::DescribeTable(_) => Ok(self),
759 LogicalPlan::Unnest(Unnest {
760 input,
761 exec_columns,
762 options,
763 ..
764 }) => {
765 unnest_with_options(Arc::unwrap_or_clone(input), exec_columns, options)
767 }
768 }
769 }
770
771 pub fn with_new_exprs(
797 &self,
798 mut expr: Vec<Expr>,
799 inputs: Vec<LogicalPlan>,
800 ) -> Result<LogicalPlan> {
801 match self {
802 LogicalPlan::Projection(Projection { .. }) => {
805 let input = self.only_input(inputs)?;
806 Projection::try_new(expr, Arc::new(input)).map(LogicalPlan::Projection)
807 }
808 LogicalPlan::Dml(DmlStatement {
809 table_name,
810 target,
811 op,
812 ..
813 }) => {
814 let input = self.only_input(inputs)?;
815 let op = match op {
816 WriteOp::MergeInto(merge_op) => {
817 WriteOp::MergeInto(Box::new(merge_op.with_new_exprs(expr)?))
818 }
819 other => {
820 self.assert_no_expressions(expr)?;
821 other.clone()
822 }
823 };
824 Ok(LogicalPlan::Dml(DmlStatement::new(
825 table_name.clone(),
826 Arc::clone(target),
827 op,
828 Arc::new(input),
829 )))
830 }
831 LogicalPlan::Copy(CopyTo {
832 input: _,
833 output_url,
834 file_type,
835 options,
836 partition_by,
837 output_schema: _,
838 }) => {
839 self.assert_no_expressions(expr)?;
840 let input = self.only_input(inputs)?;
841 Ok(LogicalPlan::Copy(CopyTo::new(
842 Arc::new(input),
843 output_url.clone(),
844 partition_by.clone(),
845 Arc::clone(file_type),
846 options.clone(),
847 )))
848 }
849 LogicalPlan::Values(Values { schema, .. }) => {
850 self.assert_no_inputs(inputs)?;
851 Ok(LogicalPlan::Values(Values {
852 schema: Arc::clone(schema),
853 values: expr
854 .chunks_exact(schema.fields().len())
855 .map(|s| s.to_vec())
856 .collect(),
857 }))
858 }
859 LogicalPlan::Filter { .. } => {
860 let predicate = self.only_expr(expr)?;
861 let input = self.only_input(inputs)?;
862
863 Filter::try_new(predicate, Arc::new(input)).map(LogicalPlan::Filter)
864 }
865 LogicalPlan::Repartition(Repartition {
866 partitioning_scheme,
867 ..
868 }) => match partitioning_scheme {
869 Partitioning::RoundRobinBatch(n) => {
870 self.assert_no_expressions(expr)?;
871 let input = self.only_input(inputs)?;
872 Ok(LogicalPlan::Repartition(Repartition {
873 partitioning_scheme: Partitioning::RoundRobinBatch(*n),
874 input: Arc::new(input),
875 }))
876 }
877 Partitioning::Hash(_, n) => {
878 let input = self.only_input(inputs)?;
879 Ok(LogicalPlan::Repartition(Repartition {
880 partitioning_scheme: Partitioning::Hash(expr, *n),
881 input: Arc::new(input),
882 }))
883 }
884 Partitioning::Range(range) => {
885 if expr.len() != range.ordering().len() {
886 return internal_err!(
887 "Incorrect number of expressions for Range partitioning"
888 );
889 }
890 let input = self.only_input(inputs)?;
891 let ordering = range
892 .ordering()
893 .iter()
894 .zip(expr)
895 .map(|(sort_expr, expr)| SortExpr {
896 expr,
897 asc: sort_expr.asc,
898 nulls_first: sort_expr.nulls_first,
899 })
900 .collect();
901 let range = RangePartitioning::try_new(
902 ordering,
903 range.split_points().to_vec(),
904 )?;
905 Ok(LogicalPlan::Repartition(Repartition {
906 partitioning_scheme: Partitioning::Range(range),
907 input: Arc::new(input),
908 }))
909 }
910 Partitioning::DistributeBy(_) => {
911 let input = self.only_input(inputs)?;
912 Ok(LogicalPlan::Repartition(Repartition {
913 partitioning_scheme: Partitioning::DistributeBy(expr),
914 input: Arc::new(input),
915 }))
916 }
917 },
918 LogicalPlan::Window(Window { window_expr, .. }) => {
919 assert_eq!(window_expr.len(), expr.len());
920 let input = self.only_input(inputs)?;
921 Window::try_new(expr, Arc::new(input)).map(LogicalPlan::Window)
922 }
923 LogicalPlan::Aggregate(Aggregate { group_expr, .. }) => {
924 let input = self.only_input(inputs)?;
925 let agg_expr = expr.split_off(group_expr.len());
927
928 Aggregate::try_new(Arc::new(input), expr, agg_expr)
929 .map(LogicalPlan::Aggregate)
930 }
931 LogicalPlan::Sort(Sort {
932 expr: sort_expr,
933 fetch,
934 ..
935 }) => {
936 let input = self.only_input(inputs)?;
937 Ok(LogicalPlan::Sort(Sort {
938 expr: expr
939 .into_iter()
940 .zip(sort_expr.iter())
941 .map(|(expr, sort)| sort.with_expr(expr))
942 .collect(),
943 input: Arc::new(input),
944 fetch: *fetch,
945 }))
946 }
947 LogicalPlan::Join(Join {
948 join_type,
949 join_constraint,
950 on,
951 null_equality,
952 null_aware,
953 ..
954 }) => {
955 let (left, right) = self.only_two_inputs(inputs)?;
956 let schema = build_join_schema(left.schema(), right.schema(), join_type)?;
957
958 let equi_expr_count = on.len() * 2;
959 assert!(expr.len() >= equi_expr_count);
960
961 let filter_expr = if expr.len() > equi_expr_count {
964 expr.pop()
965 } else {
966 None
967 };
968
969 assert_eq!(expr.len(), equi_expr_count);
972 let mut new_on = Vec::with_capacity(on.len());
973 let mut iter = expr.into_iter();
974 while let Some(left) = iter.next() {
975 let Some(right) = iter.next() else {
976 internal_err!(
977 "Expected a pair of expressions to construct the join on expression"
978 )?
979 };
980
981 new_on.push((left.unalias(), right.unalias()));
983 }
984
985 Ok(LogicalPlan::Join(Join {
986 left: Arc::new(left),
987 right: Arc::new(right),
988 join_type: *join_type,
989 join_constraint: *join_constraint,
990 on: new_on,
991 filter: filter_expr,
992 schema: DFSchemaRef::new(schema),
993 null_equality: *null_equality,
994 null_aware: *null_aware,
995 }))
996 }
997 LogicalPlan::Subquery(Subquery {
998 outer_ref_columns,
999 spans,
1000 ..
1001 }) => {
1002 self.assert_no_expressions(expr)?;
1003 let input = self.only_input(inputs)?;
1004 let subquery = LogicalPlanBuilder::from(input).build()?;
1005 Ok(LogicalPlan::Subquery(Subquery {
1006 subquery: Arc::new(subquery),
1007 outer_ref_columns: outer_ref_columns.clone(),
1008 spans: spans.clone(),
1009 }))
1010 }
1011 LogicalPlan::SubqueryAlias(SubqueryAlias { alias, .. }) => {
1012 self.assert_no_expressions(expr)?;
1013 let input = self.only_input(inputs)?;
1014 SubqueryAlias::try_new(Arc::new(input), alias.clone())
1015 .map(LogicalPlan::SubqueryAlias)
1016 }
1017 LogicalPlan::Limit(Limit { skip, fetch, .. }) => {
1018 let old_expr_len = skip.iter().chain(fetch.iter()).count();
1019 assert_eq_or_internal_err!(
1020 old_expr_len,
1021 expr.len(),
1022 "Invalid number of new Limit expressions: expected {}, got {}",
1023 old_expr_len,
1024 expr.len()
1025 );
1026 let new_fetch = fetch.as_ref().and_then(|_| expr.pop());
1028 let new_skip = skip.as_ref().and_then(|_| expr.pop());
1029 let input = self.only_input(inputs)?;
1030 Ok(LogicalPlan::Limit(Limit {
1031 skip: new_skip.map(Box::new),
1032 fetch: new_fetch.map(Box::new),
1033 input: Arc::new(input),
1034 }))
1035 }
1036 LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(CreateMemoryTable {
1037 name,
1038 if_not_exists,
1039 or_replace,
1040 column_defaults,
1041 temporary,
1042 ..
1043 })) => {
1044 self.assert_no_expressions(expr)?;
1045 let input = self.only_input(inputs)?;
1046 Ok(LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(
1047 CreateMemoryTable {
1048 input: Arc::new(input),
1049 constraints: Constraints::default(),
1050 name: name.clone(),
1051 if_not_exists: *if_not_exists,
1052 or_replace: *or_replace,
1053 column_defaults: column_defaults.clone(),
1054 temporary: *temporary,
1055 },
1056 )))
1057 }
1058 LogicalPlan::Ddl(DdlStatement::CreateView(CreateView {
1059 name,
1060 or_replace,
1061 definition,
1062 temporary,
1063 ..
1064 })) => {
1065 self.assert_no_expressions(expr)?;
1066 let input = self.only_input(inputs)?;
1067 Ok(LogicalPlan::Ddl(DdlStatement::CreateView(CreateView {
1068 input: Arc::new(input),
1069 name: name.clone(),
1070 or_replace: *or_replace,
1071 temporary: *temporary,
1072 definition: definition.clone(),
1073 })))
1074 }
1075 LogicalPlan::Extension(e) => Ok(LogicalPlan::Extension(Extension {
1076 node: e.node.with_exprs_and_inputs(expr, inputs)?,
1077 })),
1078 LogicalPlan::Union(Union { schema, .. }) => {
1079 self.assert_no_expressions(expr)?;
1080 let input_schema = inputs[0].schema();
1081 let schema = if schema.fields().len() == input_schema.fields().len() {
1083 Arc::clone(schema)
1084 } else {
1085 Arc::clone(input_schema)
1086 };
1087 Ok(LogicalPlan::Union(Union {
1088 inputs: inputs.into_iter().map(Arc::new).collect(),
1089 schema,
1090 }))
1091 }
1092 LogicalPlan::Distinct(distinct) => {
1093 let distinct = match distinct {
1094 Distinct::All(_) => {
1095 self.assert_no_expressions(expr)?;
1096 let input = self.only_input(inputs)?;
1097 Distinct::All(Arc::new(input))
1098 }
1099 Distinct::On(DistinctOn {
1100 on_expr,
1101 select_expr,
1102 ..
1103 }) => {
1104 let input = self.only_input(inputs)?;
1105 let sort_expr = expr.split_off(on_expr.len() + select_expr.len());
1106 let select_expr = expr.split_off(on_expr.len());
1107 assert!(
1108 sort_expr.is_empty(),
1109 "with_new_exprs for Distinct does not support sort expressions"
1110 );
1111 Distinct::On(DistinctOn::try_new(
1112 expr,
1113 select_expr,
1114 None, Arc::new(input),
1116 )?)
1117 }
1118 };
1119 Ok(LogicalPlan::Distinct(distinct))
1120 }
1121 LogicalPlan::RecursiveQuery(RecursiveQuery {
1122 name, is_distinct, ..
1123 }) => {
1124 self.assert_no_expressions(expr)?;
1125 let (static_term, recursive_term) = self.only_two_inputs(inputs)?;
1126 RecursiveQuery::try_new(
1127 name.clone(),
1128 Arc::new(static_term),
1129 Arc::new(recursive_term),
1130 *is_distinct,
1131 )
1132 .map(LogicalPlan::RecursiveQuery)
1133 }
1134 LogicalPlan::Analyze(a) => {
1135 self.assert_no_expressions(expr)?;
1136 let input = self.only_input(inputs)?;
1137 Ok(LogicalPlan::Analyze(Analyze {
1138 verbose: a.verbose,
1139 format: a.format.clone(),
1140 schema: Arc::clone(&a.schema),
1141 input: Arc::new(input),
1142 analyze_level: a.analyze_level,
1143 analyze_categories: a.analyze_categories.clone(),
1144 }))
1145 }
1146 LogicalPlan::Explain(e) => {
1147 self.assert_no_expressions(expr)?;
1148 let input = self.only_input(inputs)?;
1149 Ok(LogicalPlan::Explain(Explain {
1150 verbose: e.verbose,
1151 plan: Arc::new(input),
1152 explain_format: e.explain_format.clone(),
1153 stringified_plans: e.stringified_plans.clone(),
1154 schema: Arc::clone(&e.schema),
1155 logical_optimization_succeeded: e.logical_optimization_succeeded,
1156 show_statistics: e.show_statistics,
1157 }))
1158 }
1159 LogicalPlan::Statement(Statement::Prepare(Prepare {
1160 name, fields, ..
1161 })) => {
1162 self.assert_no_expressions(expr)?;
1163 let input = self.only_input(inputs)?;
1164 Ok(LogicalPlan::Statement(Statement::Prepare(Prepare {
1165 name: name.clone(),
1166 fields: fields.clone(),
1167 input: Arc::new(input),
1168 })))
1169 }
1170 LogicalPlan::Statement(Statement::Execute(Execute { name, .. })) => {
1171 self.assert_no_inputs(inputs)?;
1172 Ok(LogicalPlan::Statement(Statement::Execute(Execute {
1173 name: name.clone(),
1174 parameters: expr,
1175 })))
1176 }
1177 LogicalPlan::TableScan(ts) => {
1178 self.assert_no_inputs(inputs)?;
1179 Ok(LogicalPlan::TableScan(TableScan {
1180 filters: expr,
1181 ..ts.clone()
1182 }))
1183 }
1184 LogicalPlan::EmptyRelation(_)
1185 | LogicalPlan::Ddl(_)
1186 | LogicalPlan::Statement(_)
1187 | LogicalPlan::DescribeTable(_) => {
1188 self.assert_no_expressions(expr)?;
1190 self.assert_no_inputs(inputs)?;
1191 Ok(self.clone())
1192 }
1193 LogicalPlan::Unnest(Unnest {
1194 exec_columns: columns,
1195 options,
1196 ..
1197 }) => {
1198 let exec_columns = if expr.is_empty() {
1199 columns.clone()
1200 } else {
1201 expr.into_iter()
1202 .map(|e| match e {
1203 Expr::Column(c) => Ok(c),
1204 other => internal_err!(
1205 "Expected Expr::Column for Unnest exec_columns, got {other:?}"
1206 ),
1207 })
1208 .collect::<Result<Vec<_>>>()?
1209 };
1210 let input = self.only_input(inputs)?;
1211 Ok(unnest_with_options(input, exec_columns, options.clone())?)
1212 }
1213 }
1214 }
1215
1216 pub fn check_invariants(&self, check: InvariantLevel) -> Result<()> {
1218 match check {
1219 InvariantLevel::Always => assert_always_invariants_at_current_node(self),
1220 InvariantLevel::Executable => assert_executable_invariants(self),
1221 }
1222 }
1223
1224 #[inline]
1226 #[expect(clippy::needless_pass_by_value)] fn assert_no_expressions(&self, expr: Vec<Expr>) -> Result<()> {
1228 assert_or_internal_err!(
1229 expr.is_empty(),
1230 "{self:?} should have no exprs, got {:?}",
1231 expr
1232 );
1233 Ok(())
1234 }
1235
1236 #[inline]
1238 #[expect(clippy::needless_pass_by_value)] fn assert_no_inputs(&self, inputs: Vec<LogicalPlan>) -> Result<()> {
1240 assert_or_internal_err!(
1241 inputs.is_empty(),
1242 "{self:?} should have no inputs, got: {:?}",
1243 inputs
1244 );
1245 Ok(())
1246 }
1247
1248 #[inline]
1250 fn only_expr(&self, mut expr: Vec<Expr>) -> Result<Expr> {
1251 assert_eq_or_internal_err!(
1252 expr.len(),
1253 1,
1254 "{self:?} should have exactly one expr, got {:?}",
1255 &expr
1256 );
1257 Ok(expr.remove(0))
1258 }
1259
1260 #[inline]
1262 fn only_input(&self, mut inputs: Vec<LogicalPlan>) -> Result<LogicalPlan> {
1263 assert_eq_or_internal_err!(
1264 inputs.len(),
1265 1,
1266 "{self:?} should have exactly one input, got {:?}",
1267 &inputs
1268 );
1269 Ok(inputs.remove(0))
1270 }
1271
1272 #[inline]
1274 fn only_two_inputs(
1275 &self,
1276 mut inputs: Vec<LogicalPlan>,
1277 ) -> Result<(LogicalPlan, LogicalPlan)> {
1278 assert_eq_or_internal_err!(
1279 inputs.len(),
1280 2,
1281 "{self:?} should have exactly two inputs, got {:?}",
1282 &inputs
1283 );
1284 let right = inputs.remove(1);
1285 let left = inputs.remove(0);
1286 Ok((left, right))
1287 }
1288
1289 pub fn with_param_values(
1342 self,
1343 param_values: impl Into<ParamValues>,
1344 ) -> Result<LogicalPlan> {
1345 let param_values = param_values.into();
1346 let plan_with_values = self.replace_params_with_values(¶m_values)?;
1347
1348 Ok(
1350 if let LogicalPlan::Statement(Statement::Prepare(prepare_lp)) =
1351 plan_with_values
1352 {
1353 param_values.verify_fields(&prepare_lp.fields)?;
1354 Arc::unwrap_or_clone(prepare_lp.input)
1356 } else {
1357 plan_with_values
1358 },
1359 )
1360 }
1361
1362 pub fn max_rows(self: &LogicalPlan) -> Option<usize> {
1367 match self {
1368 LogicalPlan::Projection(Projection { input, .. }) => input.max_rows(),
1369 LogicalPlan::Filter(filter) => {
1370 if filter.is_scalar() {
1371 Some(1)
1372 } else {
1373 filter.input.max_rows()
1374 }
1375 }
1376 LogicalPlan::Window(Window { input, .. }) => input.max_rows(),
1377 LogicalPlan::Aggregate(Aggregate {
1378 input, group_expr, ..
1379 }) => {
1380 if group_expr
1382 .iter()
1383 .all(|expr| matches!(expr, Expr::Literal(_, _)))
1384 {
1385 Some(1)
1386 } else {
1387 input.max_rows()
1388 }
1389 }
1390 LogicalPlan::Sort(Sort { input, fetch, .. }) => {
1391 match (fetch, input.max_rows()) {
1392 (Some(fetch_limit), Some(input_max)) => {
1393 Some(input_max.min(*fetch_limit))
1394 }
1395 (Some(fetch_limit), None) => Some(*fetch_limit),
1396 (None, Some(input_max)) => Some(input_max),
1397 (None, None) => None,
1398 }
1399 }
1400 LogicalPlan::Join(Join {
1401 left,
1402 right,
1403 join_type,
1404 ..
1405 }) => match join_type {
1406 JoinType::Inner => Some(left.max_rows()? * right.max_rows()?),
1407 JoinType::Left | JoinType::Right | JoinType::Full => {
1408 match (left.max_rows()?, right.max_rows()?, join_type) {
1409 (0, 0, _) => Some(0),
1410 (max_rows, 0, JoinType::Left | JoinType::Full) => Some(max_rows),
1411 (0, max_rows, JoinType::Right | JoinType::Full) => Some(max_rows),
1412 (left_max, right_max, _) => Some(left_max * right_max),
1413 }
1414 }
1415 JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => {
1416 left.max_rows()
1417 }
1418 JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => {
1419 right.max_rows()
1420 }
1421 },
1422 LogicalPlan::Repartition(Repartition { input, .. }) => input.max_rows(),
1423 LogicalPlan::Union(Union { inputs, .. }) => {
1424 inputs.iter().try_fold(0usize, |mut acc, plan| {
1425 acc += plan.max_rows()?;
1426 Some(acc)
1427 })
1428 }
1429 LogicalPlan::TableScan(TableScan { fetch, .. }) => *fetch,
1430 LogicalPlan::EmptyRelation(_) => Some(0),
1431 LogicalPlan::RecursiveQuery(_) => None,
1432 LogicalPlan::Subquery(_) => None,
1433 LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => input.max_rows(),
1434 LogicalPlan::Limit(limit) => match limit.get_fetch_type() {
1435 Ok(FetchType::Literal(s)) => s,
1436 _ => None,
1437 },
1438 LogicalPlan::Distinct(
1439 Distinct::All(input) | Distinct::On(DistinctOn { input, .. }),
1440 ) => input.max_rows(),
1441 LogicalPlan::Values(v) => Some(v.values.len()),
1442 LogicalPlan::Unnest(_) => None,
1443 LogicalPlan::Ddl(_)
1444 | LogicalPlan::Explain(_)
1445 | LogicalPlan::Analyze(_)
1446 | LogicalPlan::Dml(_)
1447 | LogicalPlan::Copy(_)
1448 | LogicalPlan::DescribeTable(_)
1449 | LogicalPlan::Statement(_)
1450 | LogicalPlan::Extension(_) => None,
1451 }
1452 }
1453
1454 pub fn skip(&self) -> Result<Option<usize>> {
1459 match self {
1460 LogicalPlan::Limit(limit) => match limit.get_skip_type()? {
1461 SkipType::Literal(0) => Ok(None),
1462 SkipType::Literal(n) => Ok(Some(n)),
1463 SkipType::UnsupportedExpr => Ok(None),
1464 },
1465 LogicalPlan::Sort(_) => Ok(None),
1466 LogicalPlan::TableScan(_) => Ok(None),
1467 LogicalPlan::Projection(_) => Ok(None),
1468 LogicalPlan::Filter(_) => Ok(None),
1469 LogicalPlan::Window(_) => Ok(None),
1470 LogicalPlan::Aggregate(_) => Ok(None),
1471 LogicalPlan::Join(_) => Ok(None),
1472 LogicalPlan::Repartition(_) => Ok(None),
1473 LogicalPlan::Union(_) => Ok(None),
1474 LogicalPlan::EmptyRelation(_) => Ok(None),
1475 LogicalPlan::Subquery(_) => Ok(None),
1476 LogicalPlan::SubqueryAlias(_) => Ok(None),
1477 LogicalPlan::Statement(_) => Ok(None),
1478 LogicalPlan::Values(_) => Ok(None),
1479 LogicalPlan::Explain(_) => Ok(None),
1480 LogicalPlan::Analyze(_) => Ok(None),
1481 LogicalPlan::Extension(_) => Ok(None),
1482 LogicalPlan::Distinct(_) => Ok(None),
1483 LogicalPlan::Dml(_) => Ok(None),
1484 LogicalPlan::Ddl(_) => Ok(None),
1485 LogicalPlan::Copy(_) => Ok(None),
1486 LogicalPlan::DescribeTable(_) => Ok(None),
1487 LogicalPlan::Unnest(_) => Ok(None),
1488 LogicalPlan::RecursiveQuery(_) => Ok(None),
1489 }
1490 }
1491
1492 pub fn fetch(&self) -> Result<Option<usize>> {
1498 match self {
1499 LogicalPlan::Sort(Sort { fetch, .. }) => Ok(*fetch),
1500 LogicalPlan::TableScan(TableScan { fetch, .. }) => Ok(*fetch),
1501 LogicalPlan::Limit(limit) => match limit.get_fetch_type()? {
1502 FetchType::Literal(s) => Ok(s),
1503 FetchType::UnsupportedExpr => Ok(None),
1504 },
1505 LogicalPlan::Projection(_) => Ok(None),
1506 LogicalPlan::Filter(_) => Ok(None),
1507 LogicalPlan::Window(_) => Ok(None),
1508 LogicalPlan::Aggregate(_) => Ok(None),
1509 LogicalPlan::Join(_) => Ok(None),
1510 LogicalPlan::Repartition(_) => Ok(None),
1511 LogicalPlan::Union(_) => Ok(None),
1512 LogicalPlan::EmptyRelation(_) => Ok(None),
1513 LogicalPlan::Subquery(_) => Ok(None),
1514 LogicalPlan::SubqueryAlias(_) => Ok(None),
1515 LogicalPlan::Statement(_) => Ok(None),
1516 LogicalPlan::Values(_) => Ok(None),
1517 LogicalPlan::Explain(_) => Ok(None),
1518 LogicalPlan::Analyze(_) => Ok(None),
1519 LogicalPlan::Extension(_) => Ok(None),
1520 LogicalPlan::Distinct(_) => Ok(None),
1521 LogicalPlan::Dml(_) => Ok(None),
1522 LogicalPlan::Ddl(_) => Ok(None),
1523 LogicalPlan::Copy(_) => Ok(None),
1524 LogicalPlan::DescribeTable(_) => Ok(None),
1525 LogicalPlan::Unnest(_) => Ok(None),
1526 LogicalPlan::RecursiveQuery(_) => Ok(None),
1527 }
1528 }
1529
1530 pub fn contains_outer_reference(&self) -> bool {
1532 let mut contains = false;
1533 self.apply_expressions(|expr| {
1534 Ok(if expr.contains_outer() {
1535 contains = true;
1536 TreeNodeRecursion::Stop
1537 } else {
1538 TreeNodeRecursion::Continue
1539 })
1540 })
1541 .unwrap();
1542 contains
1543 }
1544
1545 pub fn columnized_output_exprs(&self) -> Result<Vec<(&Expr, Column)>> {
1553 match self {
1554 LogicalPlan::Aggregate(aggregate) => Ok(aggregate
1555 .output_expressions()?
1556 .into_iter()
1557 .zip(self.schema().columns())
1558 .collect()),
1559 LogicalPlan::Window(Window {
1560 window_expr,
1561 input,
1562 schema,
1563 }) => {
1564 let mut output_exprs = input.columnized_output_exprs()?;
1572 let input_len = input.schema().fields().len();
1573 output_exprs.extend(
1574 window_expr
1575 .iter()
1576 .zip(schema.columns().into_iter().skip(input_len)),
1577 );
1578 Ok(output_exprs)
1579 }
1580 _ => Ok(vec![]),
1581 }
1582 }
1583}
1584
1585impl LogicalPlan {
1586 pub fn replace_params_with_values(
1593 self,
1594 param_values: &ParamValues,
1595 ) -> Result<LogicalPlan> {
1596 self.transform_up_with_subqueries(|plan| {
1597 let schema = Arc::clone(plan.schema());
1598 let name_preserver = NamePreserver::new(&plan);
1599 plan.map_expressions(|e| {
1600 let (e, has_placeholder) = e.infer_placeholder_types(&schema)?;
1601 if !has_placeholder {
1602 Ok(Transformed::no(e))
1606 } else {
1607 let original_name = name_preserver.save(&e);
1608 let transformed_expr = e.transform_up(|e| {
1609 if let Expr::Placeholder(Placeholder { id, .. }) = e {
1610 let (value, metadata) = param_values
1611 .get_placeholders_with_values(&id)?
1612 .into_inner();
1613 Ok(Transformed::yes(Expr::Literal(value, metadata)))
1614 } else {
1615 Ok(Transformed::no(e))
1616 }
1617 })?;
1618 Ok(transformed_expr.update_data(|expr| original_name.restore(expr)))
1620 }
1621 })?
1622 .map_data(|plan| plan.update_schema_data_type())
1623 })
1624 .map(|res| res.data)
1625 }
1626
1627 fn update_schema_data_type(self) -> Result<LogicalPlan> {
1633 match self {
1634 LogicalPlan::Values(Values { values, schema: _ }) => {
1638 LogicalPlanBuilder::values(values)?.build()
1639 }
1640 plan => plan.recompute_schema(),
1642 }
1643 }
1644
1645 pub fn get_parameter_names(&self) -> Result<HashSet<String>> {
1647 let mut param_names = HashSet::new();
1648 self.apply_with_subqueries(|plan| {
1649 plan.apply_expressions(|expr| {
1650 expr.apply(|expr| {
1651 if let Expr::Placeholder(Placeholder { id, .. }) = expr {
1652 param_names.insert(id.clone());
1653 }
1654 Ok(TreeNodeRecursion::Continue)
1655 })
1656 })
1657 })
1658 .map(|_| param_names)
1659 }
1660
1661 pub fn get_parameter_types(
1666 &self,
1667 ) -> Result<HashMap<String, Option<DataType>>, DataFusionError> {
1668 let mut parameter_fields = self.get_parameter_fields()?;
1669 Ok(parameter_fields
1670 .drain()
1671 .map(|(name, maybe_field)| {
1672 (name, maybe_field.map(|field| field.data_type().clone()))
1673 })
1674 .collect())
1675 }
1676
1677 pub fn get_parameter_fields(
1679 &self,
1680 ) -> Result<HashMap<String, Option<FieldRef>>, DataFusionError> {
1681 let mut param_types: HashMap<String, Option<FieldRef>> = HashMap::new();
1682
1683 self.apply_with_subqueries(|plan| {
1684 plan.apply_expressions(|expr| {
1685 expr.apply(|expr| {
1686 if let Expr::Placeholder(Placeholder { id, field }) = expr {
1687 let prev = param_types.get(id);
1688 match (prev, field) {
1689 (Some(Some(prev)), Some(field)) => {
1690 check_metadata_with_storage_equal(
1691 (field.data_type(), Some(field.metadata())),
1692 (prev.data_type(), Some(prev.metadata())),
1693 "parameter",
1694 &format!(": Conflicting types for id {id}"),
1695 )?;
1696 }
1697 (_, Some(field)) => {
1698 param_types.insert(id.clone(), Some(Arc::clone(field)));
1699 }
1700 _ => {
1701 param_types.insert(id.clone(), None);
1702 }
1703 }
1704 }
1705 Ok(TreeNodeRecursion::Continue)
1706 })
1707 })
1708 })
1709 .map(|_| param_types)
1710 }
1711
1712 pub fn display_indent(&self) -> impl Display + '_ {
1744 struct Wrapper<'a>(&'a LogicalPlan);
1747 impl Display for Wrapper<'_> {
1748 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
1749 let with_schema = false;
1750 let mut visitor = IndentVisitor::new(f, with_schema);
1751 match self.0.visit_with_subqueries(&mut visitor) {
1752 Ok(_) => Ok(()),
1753 Err(_) => Err(fmt::Error),
1754 }
1755 }
1756 }
1757 Wrapper(self)
1758 }
1759
1760 pub fn display_indent_schema(&self) -> impl Display + '_ {
1790 struct Wrapper<'a>(&'a LogicalPlan);
1793 impl Display for Wrapper<'_> {
1794 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
1795 let with_schema = true;
1796 let mut visitor = IndentVisitor::new(f, with_schema);
1797 match self.0.visit_with_subqueries(&mut visitor) {
1798 Ok(_) => Ok(()),
1799 Err(_) => Err(fmt::Error),
1800 }
1801 }
1802 }
1803 Wrapper(self)
1804 }
1805
1806 pub fn display_pg_json(&self) -> impl Display + '_ {
1810 struct Wrapper<'a>(&'a LogicalPlan);
1813 impl Display for Wrapper<'_> {
1814 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
1815 let mut visitor = PgJsonVisitor::new(f);
1816 visitor.with_schema(true);
1817 match self.0.visit_with_subqueries(&mut visitor) {
1818 Ok(_) => Ok(()),
1819 Err(_) => Err(fmt::Error),
1820 }
1821 }
1822 }
1823 Wrapper(self)
1824 }
1825
1826 pub fn display_graphviz(&self) -> impl Display + '_ {
1856 struct Wrapper<'a>(&'a LogicalPlan);
1859 impl Display for Wrapper<'_> {
1860 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
1861 let mut visitor = GraphvizVisitor::new(f);
1862
1863 visitor.start_graph()?;
1864
1865 visitor.pre_visit_plan("LogicalPlan")?;
1866 self.0
1867 .visit_with_subqueries(&mut visitor)
1868 .map_err(|_| fmt::Error)?;
1869 visitor.post_visit_plan()?;
1870
1871 visitor.set_with_schema(true);
1872 visitor.pre_visit_plan("Detailed LogicalPlan")?;
1873 self.0
1874 .visit_with_subqueries(&mut visitor)
1875 .map_err(|_| fmt::Error)?;
1876 visitor.post_visit_plan()?;
1877
1878 visitor.end_graph()?;
1879 Ok(())
1880 }
1881 }
1882 Wrapper(self)
1883 }
1884
1885 pub fn display(&self) -> impl Display + '_ {
1907 struct Wrapper<'a>(&'a LogicalPlan);
1910 impl Display for Wrapper<'_> {
1911 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
1912 match self.0 {
1913 LogicalPlan::EmptyRelation(EmptyRelation {
1914 produce_one_row,
1915 schema: _,
1916 }) => {
1917 let rows = if *produce_one_row { 1 } else { 0 };
1918 write!(f, "EmptyRelation: rows={rows}")
1919 }
1920 LogicalPlan::RecursiveQuery(RecursiveQuery {
1921 is_distinct, ..
1922 }) => {
1923 write!(f, "RecursiveQuery: is_distinct={is_distinct}")
1924 }
1925 LogicalPlan::Values(Values { values, .. }) => {
1926 let str_values: Vec<_> = values
1927 .iter()
1928 .take(5)
1930 .map(|row| {
1931 let item = row
1932 .iter()
1933 .map(|expr| expr.to_string())
1934 .collect::<Vec<_>>()
1935 .join(", ");
1936 format!("({item})")
1937 })
1938 .collect();
1939
1940 let eclipse = if values.len() > 5 { "..." } else { "" };
1941 write!(f, "Values: {}{}", str_values.join(", "), eclipse)
1942 }
1943
1944 LogicalPlan::TableScan(TableScan {
1945 source,
1946 table_name,
1947 projection,
1948 filters,
1949 fetch,
1950 ..
1951 }) => {
1952 let projected_fields = match projection {
1953 Some(indices) => {
1954 let schema = source.schema();
1955 let names: Vec<&str> = indices
1956 .iter()
1957 .map(|i| schema.field(*i).name().as_str())
1958 .collect();
1959 format!(" projection=[{}]", names.join(", "))
1960 }
1961 _ => "".to_string(),
1962 };
1963
1964 write!(f, "TableScan: {table_name}{projected_fields}")?;
1965
1966 if !filters.is_empty() {
1967 let mut full_filter = vec![];
1968 let mut partial_filter = vec![];
1969 let mut unsupported_filters = vec![];
1970 let filters: Vec<&Expr> = filters.iter().collect();
1971
1972 if let Ok(results) =
1973 source.supports_filters_pushdown(&filters)
1974 {
1975 filters.iter().zip(results.iter()).for_each(
1976 |(x, res)| match res {
1977 TableProviderFilterPushDown::Exact => {
1978 full_filter.push(x)
1979 }
1980 TableProviderFilterPushDown::Inexact => {
1981 partial_filter.push(x)
1982 }
1983 TableProviderFilterPushDown::Unsupported => {
1984 unsupported_filters.push(x)
1985 }
1986 },
1987 );
1988 }
1989
1990 if !full_filter.is_empty() {
1991 write!(
1992 f,
1993 ", full_filters=[{}]",
1994 expr_vec_fmt!(full_filter)
1995 )?;
1996 };
1997 if !partial_filter.is_empty() {
1998 write!(
1999 f,
2000 ", partial_filters=[{}]",
2001 expr_vec_fmt!(partial_filter)
2002 )?;
2003 }
2004 if !unsupported_filters.is_empty() {
2005 write!(
2006 f,
2007 ", unsupported_filters=[{}]",
2008 expr_vec_fmt!(unsupported_filters)
2009 )?;
2010 }
2011 }
2012
2013 if let Some(n) = fetch {
2014 write!(f, ", fetch={n}")?;
2015 }
2016
2017 Ok(())
2018 }
2019 LogicalPlan::Projection(Projection { expr, .. }) => {
2020 write!(f, "Projection:")?;
2021 for (i, expr_item) in expr.iter().enumerate() {
2022 if i > 0 {
2023 write!(f, ",")?;
2024 }
2025 write!(f, " {expr_item}")?;
2026 }
2027 Ok(())
2028 }
2029 LogicalPlan::Dml(DmlStatement { table_name, op, .. }) => {
2030 write!(f, "Dml: op=[{op}] table=[{table_name}]")
2031 }
2032 LogicalPlan::Copy(CopyTo {
2033 input: _,
2034 output_url,
2035 file_type,
2036 options,
2037 ..
2038 }) => {
2039 let op_str = options
2040 .iter()
2041 .map(|(k, v)| format!("{k} {v}"))
2042 .collect::<Vec<String>>()
2043 .join(", ");
2044
2045 write!(
2046 f,
2047 "CopyTo: format={} output_url={output_url} options: ({op_str})",
2048 file_type.get_ext()
2049 )
2050 }
2051 LogicalPlan::Ddl(ddl) => {
2052 write!(f, "{}", ddl.display())
2053 }
2054 LogicalPlan::Filter(Filter {
2055 predicate: expr, ..
2056 }) => write!(f, "Filter: {expr}"),
2057 LogicalPlan::Window(Window { window_expr, .. }) => {
2058 write!(
2059 f,
2060 "WindowAggr: windowExpr=[[{}]]",
2061 expr_vec_fmt!(window_expr)
2062 )
2063 }
2064 LogicalPlan::Aggregate(Aggregate {
2065 group_expr,
2066 aggr_expr,
2067 ..
2068 }) => write!(
2069 f,
2070 "Aggregate: groupBy=[[{}]], aggr=[[{}]]",
2071 expr_vec_fmt!(group_expr),
2072 expr_vec_fmt!(aggr_expr)
2073 ),
2074 LogicalPlan::Sort(Sort { expr, fetch, .. }) => {
2075 write!(f, "Sort: ")?;
2076 for (i, expr_item) in expr.iter().enumerate() {
2077 if i > 0 {
2078 write!(f, ", ")?;
2079 }
2080 write!(f, "{expr_item}")?;
2081 }
2082 if let Some(a) = fetch {
2083 write!(f, ", fetch={a}")?;
2084 }
2085
2086 Ok(())
2087 }
2088 LogicalPlan::Join(Join {
2089 on: keys,
2090 filter,
2091 join_constraint,
2092 join_type,
2093 null_aware,
2094 ..
2095 }) => {
2096 let join_expr: Vec<String> =
2097 keys.iter().map(|(l, r)| format!("{l} = {r}")).collect();
2098 let filter_expr = filter
2099 .as_ref()
2100 .map(|expr| format!(" Filter: {expr}"))
2101 .unwrap_or_else(|| "".to_string());
2102 let null_aware_expr =
2103 if *null_aware { " null_aware" } else { "" };
2104 let join_type = if filter.is_none()
2105 && keys.is_empty()
2106 && *join_type == JoinType::Inner
2107 {
2108 "Cross".to_string()
2109 } else {
2110 join_type.to_string()
2111 };
2112 match join_constraint {
2113 JoinConstraint::On => {
2114 write!(f, "{join_type} Join:",)?;
2115 if !join_expr.is_empty() || !filter_expr.is_empty() {
2116 write!(
2117 f,
2118 " {}{}",
2119 join_expr.join(", "),
2120 filter_expr
2121 )?;
2122 }
2123 write!(f, "{null_aware_expr}")?;
2124 Ok(())
2125 }
2126 JoinConstraint::Using => {
2127 write!(
2128 f,
2129 "{} Join: Using {}{}{}",
2130 join_type,
2131 join_expr.join(", "),
2132 filter_expr,
2133 null_aware_expr,
2134 )
2135 }
2136 }
2137 }
2138 LogicalPlan::Repartition(Repartition {
2139 partitioning_scheme,
2140 ..
2141 }) => match partitioning_scheme {
2142 Partitioning::RoundRobinBatch(n) => {
2143 write!(f, "Repartition: RoundRobinBatch partition_count={n}")
2144 }
2145 Partitioning::Hash(expr, n) => {
2146 let hash_expr: Vec<String> =
2147 expr.iter().map(|e| format!("{e}")).collect();
2148 write!(
2149 f,
2150 "Repartition: Hash({}) partition_count={}",
2151 hash_expr.join(", "),
2152 n
2153 )
2154 }
2155 Partitioning::Range(range) => {
2156 write!(f, "Repartition: {range}")
2157 }
2158 Partitioning::DistributeBy(expr) => {
2159 let dist_by_expr: Vec<String> =
2160 expr.iter().map(|e| format!("{e}")).collect();
2161 write!(
2162 f,
2163 "Repartition: DistributeBy({})",
2164 dist_by_expr.join(", "),
2165 )
2166 }
2167 },
2168 LogicalPlan::Limit(limit) => {
2169 let skip_str = match limit.get_skip_type() {
2171 Ok(SkipType::Literal(n)) => n.to_string(),
2172 _ => limit
2173 .skip
2174 .as_ref()
2175 .map_or_else(|| "None".to_string(), |x| x.to_string()),
2176 };
2177 let fetch_str = match limit.get_fetch_type() {
2178 Ok(FetchType::Literal(Some(n))) => n.to_string(),
2179 Ok(FetchType::Literal(None)) => "None".to_string(),
2180 _ => limit
2181 .fetch
2182 .as_ref()
2183 .map_or_else(|| "None".to_string(), |x| x.to_string()),
2184 };
2185 write!(f, "Limit: skip={skip_str}, fetch={fetch_str}",)
2186 }
2187 LogicalPlan::Subquery(Subquery { .. }) => {
2188 write!(f, "Subquery:")
2189 }
2190 LogicalPlan::SubqueryAlias(SubqueryAlias { alias, .. }) => {
2191 write!(f, "SubqueryAlias: {alias}")
2192 }
2193 LogicalPlan::Statement(statement) => {
2194 write!(f, "{}", statement.display())
2195 }
2196 LogicalPlan::Distinct(distinct) => match distinct {
2197 Distinct::All(_) => write!(f, "Distinct:"),
2198 Distinct::On(DistinctOn {
2199 on_expr,
2200 select_expr,
2201 sort_expr,
2202 ..
2203 }) => write!(
2204 f,
2205 "DistinctOn: on_expr=[[{}]], select_expr=[[{}]], sort_expr=[[{}]]",
2206 expr_vec_fmt!(on_expr),
2207 expr_vec_fmt!(select_expr),
2208 if let Some(sort_expr) = sort_expr {
2209 expr_vec_fmt!(sort_expr)
2210 } else {
2211 "".to_string()
2212 },
2213 ),
2214 },
2215 LogicalPlan::Explain { .. } => write!(f, "Explain"),
2216 LogicalPlan::Analyze { .. } => write!(f, "Analyze"),
2217 LogicalPlan::Union(_) => write!(f, "Union"),
2218 LogicalPlan::Extension(e) => e.node.fmt_for_explain(f),
2219 LogicalPlan::DescribeTable(DescribeTable { .. }) => {
2220 write!(f, "DescribeTable")
2221 }
2222 LogicalPlan::Unnest(Unnest {
2223 input: plan,
2224 list_type_columns: list_col_indices,
2225 struct_type_columns: struct_col_indices,
2226 ..
2227 }) => {
2228 let input_columns = plan.schema().columns();
2229 let list_type_columns = list_col_indices
2230 .iter()
2231 .map(|(i, unnest_info)| {
2232 format!(
2233 "{}|depth={}",
2234 input_columns[*i], unnest_info.depth
2235 )
2236 })
2237 .collect::<Vec<String>>();
2238 let struct_type_columns = struct_col_indices
2239 .iter()
2240 .map(|i| &input_columns[*i])
2241 .collect::<Vec<&Column>>();
2242 write!(
2244 f,
2245 "Unnest: lists[{}] structs[{}]",
2246 expr_vec_fmt!(list_type_columns),
2247 expr_vec_fmt!(struct_type_columns)
2248 )
2249 }
2250 }
2251 }
2252 }
2253 Wrapper(self)
2254 }
2255
2256 pub fn resolve_lambda_variables(self) -> Result<Transformed<LogicalPlan>> {
2260 self.transform_with_subqueries(|plan| {
2261 let schema = merge_schema(&plan.inputs());
2262
2263 plan.map_expressions(|expr| expr.resolve_lambda_variables(&schema))
2264 })
2265 }
2266}
2267
2268impl Display for LogicalPlan {
2269 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
2270 self.display_indent().fmt(f)
2271 }
2272}
2273
2274impl ToStringifiedPlan for LogicalPlan {
2275 fn to_stringified(&self, plan_type: PlanType) -> StringifiedPlan {
2276 StringifiedPlan::new(plan_type, self.display_indent().to_string())
2277 }
2278}
2279
2280#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2284pub struct EmptyRelation {
2285 pub produce_one_row: bool,
2287 pub schema: DFSchemaRef,
2289}
2290
2291impl PartialOrd for EmptyRelation {
2293 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2294 self.produce_one_row
2295 .partial_cmp(&other.produce_one_row)
2296 .filter(|cmp| *cmp != Ordering::Equal || self == other)
2298 }
2299}
2300
2301#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2324pub struct RecursiveQuery {
2325 pub name: String,
2327 pub static_term: Arc<LogicalPlan>,
2329 pub recursive_term: Arc<LogicalPlan>,
2332 pub is_distinct: bool,
2335 pub schema: DFSchemaRef,
2337}
2338
2339impl PartialOrd for RecursiveQuery {
2340 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2341 match self.name.partial_cmp(&other.name) {
2342 Some(Ordering::Equal) => {
2343 match self.static_term.partial_cmp(&other.static_term) {
2344 Some(Ordering::Equal) => {
2345 match self.recursive_term.partial_cmp(&other.recursive_term) {
2346 Some(Ordering::Equal) => {
2347 self.is_distinct.partial_cmp(&other.is_distinct)
2348 }
2349 cmp => cmp,
2350 }
2351 }
2352 cmp => cmp,
2353 }
2354 }
2355 cmp => cmp,
2356 }
2357 .filter(|cmp| *cmp != Ordering::Equal || self == other)
2361 }
2362}
2363
2364impl RecursiveQuery {
2365 pub fn try_new(
2366 name: String,
2367 static_term: Arc<LogicalPlan>,
2368 recursive_term: Arc<LogicalPlan>,
2369 is_distinct: bool,
2370 ) -> Result<Self> {
2371 let schema =
2372 recursive_query_output_schema(static_term.schema(), recursive_term.schema())?;
2373 Ok(Self {
2374 name,
2375 static_term,
2376 recursive_term,
2377 is_distinct,
2378 schema,
2379 })
2380 }
2381}
2382
2383fn recursive_query_output_schema(
2394 static_schema: &DFSchemaRef,
2395 recursive_schema: &DFSchemaRef,
2396) -> Result<DFSchemaRef> {
2397 if static_schema.fields().len() != recursive_schema.fields().len() {
2398 return Err(DataFusionError::Plan(format!(
2399 "Non-recursive term and recursive term must have the same number of columns ({} != {})",
2400 static_schema.fields().len(),
2401 recursive_schema.fields().len()
2402 )));
2403 }
2404
2405 let fields = static_schema
2406 .iter()
2407 .zip(recursive_schema.fields())
2408 .map(|((qualifier, static_field), recursive_field)| {
2409 let nullable = static_field.is_nullable() || recursive_field.is_nullable();
2410 (
2411 qualifier.cloned(),
2412 static_field.as_ref().clone().with_nullable(nullable).into(),
2413 )
2414 })
2415 .collect::<Vec<_>>();
2416
2417 DFSchema::new_with_metadata(fields, static_schema.metadata().clone())
2418 .map(DFSchemaRef::new)
2419}
2420
2421#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2425pub struct Values {
2426 pub schema: DFSchemaRef,
2428 pub values: Vec<Vec<Expr>>,
2430}
2431
2432impl PartialOrd for Values {
2434 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2435 self.values
2436 .partial_cmp(&other.values)
2437 .filter(|cmp| *cmp != Ordering::Equal || self == other)
2439 }
2440}
2441
2442#[derive(Clone, PartialEq, Eq, Hash, Debug)]
2445#[non_exhaustive]
2447pub struct Projection {
2448 pub expr: Vec<Expr>,
2450 pub input: Arc<LogicalPlan>,
2452 pub schema: DFSchemaRef,
2454}
2455
2456impl PartialOrd for Projection {
2458 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2459 match self.expr.partial_cmp(&other.expr) {
2460 Some(Ordering::Equal) => self.input.partial_cmp(&other.input),
2461 cmp => cmp,
2462 }
2463 .filter(|cmp| *cmp != Ordering::Equal || self == other)
2465 }
2466}
2467
2468impl Projection {
2469 pub fn try_new(expr: Vec<Expr>, input: Arc<LogicalPlan>) -> Result<Self> {
2471 let projection_schema = projection_schema(&input, &expr)?;
2472 Self::try_new_with_schema(expr, input, projection_schema)
2473 }
2474
2475 pub fn try_new_with_schema(
2477 expr: Vec<Expr>,
2478 input: Arc<LogicalPlan>,
2479 schema: DFSchemaRef,
2480 ) -> Result<Self> {
2481 #[expect(deprecated)]
2482 if !expr.iter().any(|e| matches!(e, Expr::Wildcard { .. }))
2483 && expr.len() != schema.fields().len()
2484 {
2485 return plan_err!(
2486 "Projection has mismatch between number of expressions ({}) and number of fields in schema ({})",
2487 expr.len(),
2488 schema.fields().len()
2489 );
2490 }
2491 Ok(Self {
2492 expr,
2493 input,
2494 schema,
2495 })
2496 }
2497
2498 pub fn new_from_schema(input: Arc<LogicalPlan>, schema: DFSchemaRef) -> Self {
2500 let expr: Vec<Expr> = schema.columns().into_iter().map(Expr::Column).collect();
2501 Self {
2502 expr,
2503 input,
2504 schema,
2505 }
2506 }
2507}
2508
2509pub fn projection_schema(input: &LogicalPlan, exprs: &[Expr]) -> Result<Arc<DFSchema>> {
2529 let metadata = input.schema().metadata().clone();
2531
2532 let schema =
2534 DFSchema::new_with_metadata(exprlist_to_fields(exprs, input)?, metadata)?
2535 .with_functional_dependencies(calc_func_dependencies_for_project(
2536 exprs, input,
2537 )?)?;
2538
2539 Ok(Arc::new(schema))
2540}
2541
2542#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2544#[non_exhaustive]
2546pub struct SubqueryAlias {
2547 pub input: Arc<LogicalPlan>,
2549 pub alias: TableReference,
2551 pub schema: DFSchemaRef,
2553}
2554
2555impl SubqueryAlias {
2556 pub fn try_new(
2557 plan: Arc<LogicalPlan>,
2558 alias: impl Into<TableReference>,
2559 ) -> Result<Self> {
2560 let alias = alias.into();
2561
2562 let aliases = unique_field_aliases(plan.schema().fields());
2568 let is_projection_needed = aliases.iter().any(Option::is_some);
2569
2570 let plan = if is_projection_needed {
2572 let projection_expressions = aliases
2573 .iter()
2574 .zip(plan.schema().iter())
2575 .map(|(alias, (qualifier, field))| {
2576 let column =
2577 Expr::Column(Column::new(qualifier.cloned(), field.name()));
2578 match alias {
2579 None => column,
2580 Some(alias) => {
2581 Expr::Alias(Alias::new(column, qualifier.cloned(), alias))
2582 }
2583 }
2584 })
2585 .collect();
2586 let projection = Projection::try_new(projection_expressions, plan)?;
2587 Arc::new(LogicalPlan::Projection(projection))
2588 } else {
2589 plan
2590 };
2591
2592 let fields = plan.schema().fields().clone();
2594 let meta_data = plan.schema().metadata().clone();
2595 let func_dependencies = plan.schema().functional_dependencies().clone();
2596
2597 let schema = DFSchema::from_unqualified_fields(fields, meta_data)?;
2598 let schema = schema.as_arrow();
2599
2600 let schema = DFSchemaRef::new(
2601 DFSchema::try_from_qualified_schema(alias.clone(), schema)?
2602 .with_functional_dependencies(func_dependencies)?,
2603 );
2604 Ok(SubqueryAlias {
2605 input: plan,
2606 alias,
2607 schema,
2608 })
2609 }
2610}
2611
2612impl PartialOrd for SubqueryAlias {
2614 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2615 match self.input.partial_cmp(&other.input) {
2616 Some(Ordering::Equal) => self.alias.partial_cmp(&other.alias),
2617 cmp => cmp,
2618 }
2619 .filter(|cmp| *cmp != Ordering::Equal || self == other)
2621 }
2622}
2623
2624#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
2636#[non_exhaustive]
2637pub struct Filter {
2638 pub predicate: Expr,
2640 pub input: Arc<LogicalPlan>,
2642}
2643
2644impl Filter {
2645 #[doc(hidden)]
2654 pub fn new(predicate: Expr, input: Arc<LogicalPlan>) -> Self {
2655 Self { predicate, input }
2656 }
2657
2658 pub fn try_new(predicate: Expr, input: Arc<LogicalPlan>) -> Result<Self> {
2663 Self::try_new_internal(predicate, input)
2664 }
2665
2666 fn is_allowed_filter_type(data_type: &DataType) -> bool {
2667 match data_type {
2668 DataType::Boolean | DataType::Null => true,
2670 DataType::Dictionary(_, value_type) => {
2671 Filter::is_allowed_filter_type(value_type.as_ref())
2672 }
2673 _ => false,
2674 }
2675 }
2676
2677 fn try_new_internal(predicate: Expr, input: Arc<LogicalPlan>) -> Result<Self> {
2678 if let Ok(predicate_type) = predicate.get_type(input.schema())
2683 && !Filter::is_allowed_filter_type(&predicate_type)
2684 {
2685 return plan_err!(
2686 "Cannot create filter with non-boolean predicate '{predicate}' returning {predicate_type}"
2687 );
2688 }
2689
2690 Ok(Self {
2691 predicate: predicate.unalias_nested().data,
2692 input,
2693 })
2694 }
2695
2696 fn is_scalar(&self) -> bool {
2712 let schema = self.input.schema();
2713
2714 let functional_dependencies = self.input.schema().functional_dependencies();
2715 let unique_keys = functional_dependencies.iter().filter(|dep| {
2716 let nullable = dep.nullable
2717 && dep
2718 .source_indices
2719 .iter()
2720 .any(|&source| schema.field(source).is_nullable());
2721 !nullable
2722 && dep.mode == Dependency::Single
2723 && dep.target_indices.len() == schema.fields().len()
2724 });
2725
2726 let exprs = split_conjunction(&self.predicate);
2727 let eq_pred_cols: HashSet<_> = exprs
2728 .iter()
2729 .filter_map(|expr| {
2730 let Expr::BinaryExpr(BinaryExpr {
2731 left,
2732 op: Operator::Eq,
2733 right,
2734 }) = expr
2735 else {
2736 return None;
2737 };
2738 if left == right {
2740 return None;
2741 }
2742
2743 match (left.as_ref(), right.as_ref()) {
2744 (Expr::Column(_), Expr::Column(_)) => None,
2745 (Expr::Column(c), _) | (_, Expr::Column(c)) => {
2746 Some(schema.index_of_column(c).unwrap())
2747 }
2748 _ => None,
2749 }
2750 })
2751 .collect();
2752
2753 for key in unique_keys {
2756 if key.source_indices.iter().all(|c| eq_pred_cols.contains(c)) {
2757 return true;
2758 }
2759 }
2760 false
2761 }
2762}
2763
2764#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2779pub struct Window {
2780 pub input: Arc<LogicalPlan>,
2782 pub window_expr: Vec<Expr>,
2784 pub schema: DFSchemaRef,
2786}
2787
2788impl Window {
2789 pub fn try_new(window_expr: Vec<Expr>, input: Arc<LogicalPlan>) -> Result<Self> {
2791 check_aggregate_and_window_nesting(window_expr.iter())?;
2795
2796 let fields: Vec<(Option<TableReference>, Arc<Field>)> = input
2797 .schema()
2798 .iter()
2799 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
2800 .collect();
2801 let input_len = fields.len();
2802 let mut window_fields = fields;
2803 let expr_fields = exprlist_to_fields(window_expr.as_slice(), &input)?;
2804 window_fields.extend_from_slice(expr_fields.as_slice());
2805 let metadata = input.schema().metadata().clone();
2806
2807 let mut window_func_dependencies =
2809 input.schema().functional_dependencies().clone();
2810 window_func_dependencies.extend_target_indices(window_fields.len());
2811
2812 let mut new_dependencies = window_expr
2816 .iter()
2817 .enumerate()
2818 .filter_map(|(idx, expr)| {
2819 let Expr::WindowFunction(window_fun) = expr else {
2820 return None;
2821 };
2822 let WindowFunction {
2823 fun: WindowFunctionDefinition::WindowUDF(udwf),
2824 params: WindowFunctionParams { partition_by, .. },
2825 } = window_fun.as_ref()
2826 else {
2827 return None;
2828 };
2829 if udwf.name() == "row_number" && partition_by.is_empty() {
2832 Some(idx + input_len)
2833 } else {
2834 None
2835 }
2836 })
2837 .map(|idx| {
2838 FunctionalDependence::new(vec![idx], vec![], false)
2839 .with_mode(Dependency::Single)
2840 })
2841 .collect::<Vec<_>>();
2842
2843 if !new_dependencies.is_empty() {
2844 for dependence in new_dependencies.iter_mut() {
2845 dependence.target_indices = (0..window_fields.len()).collect();
2846 }
2847 let new_deps = FunctionalDependencies::new(new_dependencies);
2849 window_func_dependencies.extend(new_deps);
2850 }
2851
2852 if let Some(e) = window_expr.iter().find(|e| {
2854 matches!(
2855 e,
2856 Expr::WindowFunction(wf)
2857 if !matches!(wf.fun, WindowFunctionDefinition::AggregateUDF(_))
2858 && wf.params.filter.is_some()
2859 )
2860 }) {
2861 return plan_err!(
2862 "FILTER clause can only be used with aggregate window functions. Found in '{e}'"
2863 );
2864 }
2865
2866 Self::try_new_with_schema(
2867 window_expr,
2868 input,
2869 Arc::new(
2870 DFSchema::new_with_metadata(window_fields, metadata)?
2871 .with_functional_dependencies(window_func_dependencies)?,
2872 ),
2873 )
2874 }
2875
2876 pub fn try_new_with_schema(
2882 window_expr: Vec<Expr>,
2883 input: Arc<LogicalPlan>,
2884 schema: DFSchemaRef,
2885 ) -> Result<Self> {
2886 let input_fields_count = input.schema().fields().len();
2887 if schema.fields().len() != input_fields_count + window_expr.len() {
2888 return plan_err!(
2889 "Window schema has wrong number of fields. Expected {} got {}",
2890 input_fields_count + window_expr.len(),
2891 schema.fields().len()
2892 );
2893 }
2894
2895 Ok(Window {
2896 input,
2897 window_expr,
2898 schema,
2899 })
2900 }
2901}
2902
2903impl PartialOrd for Window {
2905 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2906 match self.input.partial_cmp(&other.input)? {
2907 Ordering::Equal => {} not_equal => return Some(not_equal),
2909 }
2910
2911 match self.window_expr.partial_cmp(&other.window_expr)? {
2912 Ordering::Equal => {} not_equal => return Some(not_equal),
2914 }
2915
2916 if self == other {
2919 Some(Ordering::Equal)
2920 } else {
2921 None
2922 }
2923 }
2924}
2925
2926#[derive(Clone)]
2928pub struct TableScan {
2929 pub table_name: TableReference,
2931 pub source: Arc<dyn TableSource>,
2933 pub projection: Option<Vec<usize>>,
2935 pub projected_schema: DFSchemaRef,
2937 pub filters: Vec<Expr>,
2939 pub fetch: Option<usize>,
2941 pub statistics_requests: BTreeSet<StatisticsRequest>,
2947}
2948
2949impl Debug for TableScan {
2950 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
2951 f.debug_struct("TableScan")
2952 .field("table_name", &self.table_name)
2953 .field("source", &"...")
2954 .field("projection", &self.projection)
2955 .field("projected_schema", &self.projected_schema)
2956 .field("filters", &self.filters)
2957 .field("fetch", &self.fetch)
2958 .finish_non_exhaustive()
2959 }
2960}
2961
2962impl PartialEq for TableScan {
2963 fn eq(&self, other: &Self) -> bool {
2964 self.table_name == other.table_name
2965 && self.projection == other.projection
2966 && self.projected_schema == other.projected_schema
2967 && self.filters == other.filters
2968 && self.fetch == other.fetch
2969 }
2970}
2971
2972impl Eq for TableScan {}
2973
2974impl PartialOrd for TableScan {
2977 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2978 #[derive(PartialEq, PartialOrd)]
2979 struct ComparableTableScan<'a> {
2980 pub table_name: &'a TableReference,
2982 pub projection: &'a Option<Vec<usize>>,
2984 pub filters: &'a Vec<Expr>,
2986 pub fetch: &'a Option<usize>,
2988 }
2989 let comparable_self = ComparableTableScan {
2990 table_name: &self.table_name,
2991 projection: &self.projection,
2992 filters: &self.filters,
2993 fetch: &self.fetch,
2994 };
2995 let comparable_other = ComparableTableScan {
2996 table_name: &other.table_name,
2997 projection: &other.projection,
2998 filters: &other.filters,
2999 fetch: &other.fetch,
3000 };
3001 comparable_self
3002 .partial_cmp(&comparable_other)
3003 .filter(|cmp| *cmp != Ordering::Equal || self == other)
3005 }
3006}
3007
3008impl Hash for TableScan {
3009 fn hash<H: Hasher>(&self, state: &mut H) {
3010 self.table_name.hash(state);
3011 self.projection.hash(state);
3012 self.projected_schema.hash(state);
3013 self.filters.hash(state);
3014 self.fetch.hash(state);
3015 }
3016}
3017
3018impl TableScan {
3019 #[deprecated(since = "54.0.0", note = "use `TableScanBuilder` instead")]
3022 pub fn try_new(
3023 table_name: impl Into<TableReference>,
3024 table_source: Arc<dyn TableSource>,
3025 projection: Option<Vec<usize>>,
3026 filters: Vec<Expr>,
3027 fetch: Option<usize>,
3028 ) -> Result<Self> {
3029 TableScanBuilder::new(table_name, table_source)
3030 .with_projection(projection)
3031 .with_filters(filters)
3032 .with_fetch(fetch)
3033 .build()
3034 }
3035}
3036
3037pub struct TableScanBuilder {
3045 table_name: TableReference,
3046 source: Arc<dyn TableSource>,
3047 projection: Option<Vec<usize>>,
3048 filters: Vec<Expr>,
3049 fetch: Option<usize>,
3050 statistics_requests: BTreeSet<StatisticsRequest>,
3051}
3052
3053impl TableScanBuilder {
3054 pub fn new(
3056 table_name: impl Into<TableReference>,
3057 source: Arc<dyn TableSource>,
3058 ) -> Self {
3059 Self {
3060 table_name: table_name.into(),
3061 source,
3062 projection: None,
3063 filters: vec![],
3064 fetch: None,
3065 statistics_requests: BTreeSet::new(),
3066 }
3067 }
3068
3069 pub fn with_projection(mut self, projection: Option<Vec<usize>>) -> Self {
3071 self.projection = projection;
3072 self
3073 }
3074
3075 pub fn with_filters(mut self, filters: Vec<Expr>) -> Self {
3077 self.filters = filters;
3078 self
3079 }
3080
3081 pub fn with_fetch(mut self, fetch: Option<usize>) -> Self {
3083 self.fetch = fetch;
3084 self
3085 }
3086
3087 pub fn with_statistics_requests(
3090 mut self,
3091 statistics_requests: BTreeSet<StatisticsRequest>,
3092 ) -> Self {
3093 self.statistics_requests = statistics_requests;
3094 self
3095 }
3096
3097 pub fn build(self) -> Result<TableScan> {
3100 let TableScanBuilder {
3101 table_name,
3102 source,
3103 projection,
3104 filters,
3105 fetch,
3106 statistics_requests,
3107 } = self;
3108
3109 if table_name.table().is_empty() {
3110 return plan_err!("table_name cannot be empty");
3111 }
3112 let schema = source.schema();
3113 let func_dependencies = FunctionalDependencies::new_from_constraints(
3114 source.constraints(),
3115 schema.fields.len(),
3116 );
3117 let projected_schema = projection
3118 .as_ref()
3119 .map(|p| {
3120 let projected_func_dependencies =
3121 func_dependencies.project_functional_dependencies(p, p.len());
3122
3123 let df_schema = DFSchema::new_with_metadata(
3124 p.iter()
3125 .map(|i| {
3126 (Some(table_name.clone()), Arc::clone(&schema.fields()[*i]))
3127 })
3128 .collect(),
3129 schema.metadata.clone(),
3130 )?;
3131 df_schema.with_functional_dependencies(projected_func_dependencies)
3132 })
3133 .unwrap_or_else(|| {
3134 let df_schema =
3135 DFSchema::try_from_qualified_schema(table_name.clone(), &schema)?;
3136 df_schema.with_functional_dependencies(func_dependencies)
3137 })?;
3138 let projected_schema = Arc::new(projected_schema);
3139
3140 Ok(TableScan {
3141 table_name,
3142 source,
3143 projection,
3144 projected_schema,
3145 filters,
3146 fetch,
3147 statistics_requests,
3148 })
3149 }
3150}
3151
3152impl From<TableScan> for TableScanBuilder {
3153 fn from(scan: TableScan) -> Self {
3154 Self {
3155 table_name: scan.table_name,
3156 source: scan.source,
3157 projection: scan.projection,
3158 filters: scan.filters,
3159 fetch: scan.fetch,
3160 statistics_requests: scan.statistics_requests,
3161 }
3162 }
3163}
3164
3165#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
3167pub struct Repartition {
3168 pub input: Arc<LogicalPlan>,
3170 pub partitioning_scheme: Partitioning,
3172}
3173
3174#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3176pub struct Union {
3177 pub inputs: Vec<Arc<LogicalPlan>>,
3179 pub schema: DFSchemaRef,
3181}
3182
3183impl Union {
3184 pub fn try_new(inputs: Vec<Arc<LogicalPlan>>) -> Result<Self> {
3187 let schema = Self::derive_schema_from_inputs(&inputs, false, false)?;
3188 Ok(Union { inputs, schema })
3189 }
3190
3191 pub fn try_new_with_loose_types(inputs: Vec<Arc<LogicalPlan>>) -> Result<Self> {
3196 let schema = Self::derive_schema_from_inputs(&inputs, true, false)?;
3197 Ok(Union { inputs, schema })
3198 }
3199
3200 pub fn try_new_by_name(inputs: Vec<Arc<LogicalPlan>>) -> Result<Self> {
3204 let schema = Self::derive_schema_from_inputs(&inputs, true, true)?;
3205 let inputs = Self::rewrite_inputs_from_schema(&schema, inputs)?;
3206
3207 Ok(Union { inputs, schema })
3208 }
3209
3210 fn rewrite_inputs_from_schema(
3214 schema: &Arc<DFSchema>,
3215 inputs: Vec<Arc<LogicalPlan>>,
3216 ) -> Result<Vec<Arc<LogicalPlan>>> {
3217 let schema_width = schema.iter().count();
3218 let mut wrapped_inputs = Vec::with_capacity(inputs.len());
3219 for input in inputs {
3220 let mut expr = Vec::with_capacity(schema_width);
3224 for column in schema.columns() {
3225 if input
3226 .schema()
3227 .has_column_with_unqualified_name(column.name())
3228 {
3229 expr.push(Expr::Column(column));
3230 } else {
3231 expr.push(
3232 Expr::Literal(ScalarValue::Null, None).alias(column.name()),
3233 );
3234 }
3235 }
3236 wrapped_inputs.push(Arc::new(LogicalPlan::Projection(
3237 Projection::try_new_with_schema(expr, input, Arc::clone(schema))?,
3238 )));
3239 }
3240
3241 Ok(wrapped_inputs)
3242 }
3243
3244 fn derive_schema_from_inputs(
3253 inputs: &[Arc<LogicalPlan>],
3254 loose_types: bool,
3255 by_name: bool,
3256 ) -> Result<DFSchemaRef> {
3257 if inputs.len() < 2 {
3258 return plan_err!("UNION requires at least two inputs");
3259 }
3260
3261 if by_name {
3262 Self::derive_schema_from_inputs_by_name(inputs, loose_types)
3263 } else {
3264 Self::derive_schema_from_inputs_by_position(inputs, loose_types)
3265 }
3266 }
3267
3268 fn derive_schema_from_inputs_by_name(
3269 inputs: &[Arc<LogicalPlan>],
3270 loose_types: bool,
3271 ) -> Result<DFSchemaRef> {
3272 type FieldData<'a> =
3273 (&'a DataType, bool, Vec<&'a HashMap<String, String>>, usize);
3274 let mut cols: Vec<(&str, FieldData)> = Vec::new();
3275 for input in inputs.iter() {
3276 for field in input.schema().fields() {
3277 if let Some((_, (data_type, is_nullable, metadata, occurrences))) =
3278 cols.iter_mut().find(|(name, _)| name == field.name())
3279 {
3280 if !loose_types && *data_type != field.data_type() {
3281 return plan_err!(
3282 "Found different types for field {}",
3283 field.name()
3284 );
3285 }
3286
3287 metadata.push(field.metadata());
3288 *is_nullable |= field.is_nullable();
3291 *occurrences += 1;
3292 } else {
3293 cols.push((
3294 field.name(),
3295 (
3296 field.data_type(),
3297 field.is_nullable(),
3298 vec![field.metadata()],
3299 1,
3300 ),
3301 ));
3302 }
3303 }
3304 }
3305
3306 let union_fields = cols
3307 .into_iter()
3308 .map(
3309 |(name, (data_type, is_nullable, unmerged_metadata, occurrences))| {
3310 let final_is_nullable = if occurrences == inputs.len() {
3314 is_nullable
3315 } else {
3316 true
3317 };
3318
3319 let mut field =
3320 Field::new(name, data_type.clone(), final_is_nullable);
3321 field.set_metadata(intersect_metadata_for_union(unmerged_metadata));
3322
3323 (None, Arc::new(field))
3324 },
3325 )
3326 .collect::<Vec<(Option<TableReference>, _)>>();
3327
3328 let union_schema_metadata = intersect_metadata_for_union(
3329 inputs.iter().map(|input| input.schema().metadata()),
3330 );
3331
3332 let schema = DFSchema::new_with_metadata(union_fields, union_schema_metadata)?;
3334 let schema = Arc::new(schema);
3335
3336 Ok(schema)
3337 }
3338
3339 fn derive_schema_from_inputs_by_position(
3340 inputs: &[Arc<LogicalPlan>],
3341 loose_types: bool,
3342 ) -> Result<DFSchemaRef> {
3343 let first_schema = inputs[0].schema();
3344 let fields_count = first_schema.fields().len();
3345 for input in inputs.iter().skip(1) {
3346 if fields_count != input.schema().fields().len() {
3347 return plan_err!(
3348 "UNION queries have different number of columns: \
3349 left has {} columns whereas right has {} columns",
3350 fields_count,
3351 input.schema().fields().len()
3352 );
3353 }
3354 }
3355
3356 let mut name_counts: HashMap<String, usize> = HashMap::new();
3357 let union_fields = (0..fields_count)
3358 .map(|i| {
3359 let fields = inputs
3360 .iter()
3361 .map(|input| input.schema().field(i))
3362 .collect::<Vec<_>>();
3363 let first_field = fields[0];
3364 let base_name = first_field.name().to_string();
3365
3366 let data_type = if loose_types {
3367 first_field.data_type()
3371 } else {
3372 fields.iter().skip(1).try_fold(
3373 first_field.data_type(),
3374 |acc, field| {
3375 if acc != field.data_type() {
3376 return plan_err!(
3377 "UNION field {i} have different type in inputs: \
3378 left has {} whereas right has {}",
3379 first_field.data_type(),
3380 field.data_type()
3381 );
3382 }
3383 Ok(acc)
3384 },
3385 )?
3386 };
3387 let nullable = fields.iter().any(|field| field.is_nullable());
3388
3389 let name = if let Some(count) = name_counts.get_mut(&base_name) {
3391 *count += 1;
3392 format!("{base_name}_{count}")
3393 } else {
3394 name_counts.insert(base_name.clone(), 0);
3395 base_name
3396 };
3397
3398 let mut field = Field::new(&name, data_type.clone(), nullable);
3399 let field_metadata = intersect_metadata_for_union(
3400 fields.iter().map(|field| field.metadata()),
3401 );
3402 field.set_metadata(field_metadata);
3403 Ok((None, Arc::new(field)))
3404 })
3405 .collect::<Result<_>>()?;
3406 let union_schema_metadata = intersect_metadata_for_union(
3407 inputs.iter().map(|input| input.schema().metadata()),
3408 );
3409
3410 let schema = DFSchema::new_with_metadata(union_fields, union_schema_metadata)?;
3412 let schema = Arc::new(schema);
3413
3414 Ok(schema)
3415 }
3416}
3417
3418impl PartialOrd for Union {
3420 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
3421 self.inputs
3422 .partial_cmp(&other.inputs)
3423 .filter(|cmp| *cmp != Ordering::Equal || self == other)
3425 }
3426}
3427
3428#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3451pub struct DescribeTable {
3452 pub schema: Arc<Schema>,
3454 pub output_schema: DFSchemaRef,
3456}
3457
3458impl PartialOrd for DescribeTable {
3461 fn partial_cmp(&self, _other: &Self) -> Option<Ordering> {
3462 None
3464 }
3465}
3466
3467#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3469pub struct ExplainOption {
3470 pub verbose: bool,
3472 pub analyze: bool,
3474 pub format: ExplainFormat,
3476 pub show_statistics: Option<bool>,
3479 pub analyze_level: Option<MetricType>,
3482 pub analyze_categories: Option<ExplainAnalyzeCategories>,
3485}
3486
3487impl Default for ExplainOption {
3488 fn default() -> Self {
3489 ExplainOption {
3490 verbose: false,
3491 analyze: false,
3492 format: ExplainFormat::Indent,
3493 show_statistics: None,
3494 analyze_level: None,
3495 analyze_categories: None,
3496 }
3497 }
3498}
3499
3500impl ExplainOption {
3501 pub fn with_verbose(mut self, verbose: bool) -> Self {
3503 self.verbose = verbose;
3504 self
3505 }
3506
3507 pub fn with_analyze(mut self, analyze: bool) -> Self {
3509 self.analyze = analyze;
3510 self
3511 }
3512
3513 pub fn with_format(mut self, format: ExplainFormat) -> Self {
3515 self.format = format;
3516 self
3517 }
3518
3519 pub fn with_show_statistics(mut self, show_statistics: Option<bool>) -> Self {
3522 self.show_statistics = show_statistics;
3523 self
3524 }
3525
3526 pub fn with_analyze_level(mut self, analyze_level: Option<MetricType>) -> Self {
3529 self.analyze_level = analyze_level;
3530 self
3531 }
3532
3533 pub fn with_analyze_categories(
3536 mut self,
3537 analyze_categories: Option<ExplainAnalyzeCategories>,
3538 ) -> Self {
3539 self.analyze_categories = analyze_categories;
3540 self
3541 }
3542}
3543
3544#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3551pub struct Explain {
3552 pub verbose: bool,
3554 pub explain_format: ExplainFormat,
3557 pub plan: Arc<LogicalPlan>,
3559 pub stringified_plans: Vec<StringifiedPlan>,
3561 pub schema: DFSchemaRef,
3563 pub logical_optimization_succeeded: bool,
3565 pub show_statistics: Option<bool>,
3568}
3569
3570impl PartialOrd for Explain {
3572 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
3573 #[derive(PartialEq, PartialOrd)]
3574 struct ComparableExplain<'a> {
3575 pub verbose: &'a bool,
3577 pub plan: &'a Arc<LogicalPlan>,
3579 pub stringified_plans: &'a Vec<StringifiedPlan>,
3581 pub logical_optimization_succeeded: &'a bool,
3583 pub show_statistics: &'a Option<bool>,
3585 }
3586 let comparable_self = ComparableExplain {
3587 verbose: &self.verbose,
3588 plan: &self.plan,
3589 stringified_plans: &self.stringified_plans,
3590 logical_optimization_succeeded: &self.logical_optimization_succeeded,
3591 show_statistics: &self.show_statistics,
3592 };
3593 let comparable_other = ComparableExplain {
3594 verbose: &other.verbose,
3595 plan: &other.plan,
3596 stringified_plans: &other.stringified_plans,
3597 logical_optimization_succeeded: &other.logical_optimization_succeeded,
3598 show_statistics: &other.show_statistics,
3599 };
3600 comparable_self
3601 .partial_cmp(&comparable_other)
3602 .filter(|cmp| *cmp != Ordering::Equal || self == other)
3604 }
3605}
3606
3607#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3610pub struct Analyze {
3611 pub verbose: bool,
3613 pub format: ExplainFormat,
3615 pub input: Arc<LogicalPlan>,
3617 pub schema: DFSchemaRef,
3619 pub analyze_level: Option<MetricType>,
3622 pub analyze_categories: Option<ExplainAnalyzeCategories>,
3625}
3626
3627impl PartialOrd for Analyze {
3632 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
3633 match self.verbose.partial_cmp(&other.verbose) {
3634 Some(Ordering::Equal) => self.input.partial_cmp(&other.input),
3635 cmp => cmp,
3636 }
3637 .filter(|cmp| *cmp != Ordering::Equal || self == other)
3639 }
3640}
3641
3642#[allow(clippy::allow_attributes)]
3647#[allow(clippy::derived_hash_with_manual_eq)]
3648#[derive(Debug, Clone, Eq, Hash)]
3649pub struct Extension {
3650 pub node: Arc<dyn UserDefinedLogicalNode>,
3652}
3653
3654impl PartialEq for Extension {
3658 fn eq(&self, other: &Self) -> bool {
3659 self.node.eq(&other.node)
3660 }
3661}
3662
3663impl PartialOrd for Extension {
3664 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
3665 self.node.partial_cmp(&other.node)
3666 }
3667}
3668
3669#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
3671pub struct Limit {
3672 pub skip: Option<Box<Expr>>,
3674 pub fetch: Option<Box<Expr>>,
3677 pub input: Arc<LogicalPlan>,
3679}
3680
3681pub enum SkipType {
3683 Literal(usize),
3685 UnsupportedExpr,
3687}
3688
3689pub enum FetchType {
3691 Literal(Option<usize>),
3694 UnsupportedExpr,
3696}
3697
3698impl Limit {
3699 pub fn get_skip_type(&self) -> Result<SkipType> {
3701 match self.skip.as_deref() {
3702 Some(expr) => match *expr {
3703 Expr::Literal(ScalarValue::Int64(s), _) => {
3704 let s = s.unwrap_or(0);
3706 if s >= 0 {
3707 Ok(SkipType::Literal(s as usize))
3708 } else {
3709 plan_err!("OFFSET must be >=0, '{}' was provided", s)
3710 }
3711 }
3712 _ => Ok(SkipType::UnsupportedExpr),
3713 },
3714 None => Ok(SkipType::Literal(0)),
3716 }
3717 }
3718
3719 pub fn get_fetch_type(&self) -> Result<FetchType> {
3721 match self.fetch.as_deref() {
3722 Some(expr) => match *expr {
3723 Expr::Literal(ScalarValue::Int64(Some(s)), _) => {
3724 if s >= 0 {
3725 Ok(FetchType::Literal(Some(s as usize)))
3726 } else {
3727 plan_err!("LIMIT must be >= 0, '{}' was provided", s)
3728 }
3729 }
3730 Expr::Literal(ScalarValue::Int64(None), _) => {
3731 Ok(FetchType::Literal(None))
3732 }
3733 _ => Ok(FetchType::UnsupportedExpr),
3734 },
3735 None => Ok(FetchType::Literal(None)),
3736 }
3737 }
3738}
3739
3740#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
3742pub enum Distinct {
3743 All(Arc<LogicalPlan>),
3745 On(DistinctOn),
3747}
3748
3749impl Distinct {
3750 pub fn input(&self) -> &Arc<LogicalPlan> {
3752 match self {
3753 Distinct::All(input) => input,
3754 Distinct::On(DistinctOn { input, .. }) => input,
3755 }
3756 }
3757}
3758
3759#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3761pub struct DistinctOn {
3762 pub on_expr: Vec<Expr>,
3764 pub select_expr: Vec<Expr>,
3766 pub sort_expr: Option<Vec<SortExpr>>,
3770 pub input: Arc<LogicalPlan>,
3772 pub schema: DFSchemaRef,
3774}
3775
3776impl DistinctOn {
3777 pub fn try_new(
3779 on_expr: Vec<Expr>,
3780 select_expr: Vec<Expr>,
3781 sort_expr: Option<Vec<SortExpr>>,
3782 input: Arc<LogicalPlan>,
3783 ) -> Result<Self> {
3784 if on_expr.is_empty() {
3785 return plan_err!("No `ON` expressions provided");
3786 }
3787
3788 let on_expr = normalize_cols(on_expr, input.as_ref())?;
3789 let qualified_fields = exprlist_to_fields(select_expr.as_slice(), &input)?
3790 .into_iter()
3791 .collect();
3792
3793 let dfschema = DFSchema::new_with_metadata(
3794 qualified_fields,
3795 input.schema().metadata().clone(),
3796 )?;
3797
3798 let mut distinct_on = DistinctOn {
3799 on_expr,
3800 select_expr,
3801 sort_expr: None,
3802 input,
3803 schema: Arc::new(dfschema),
3804 };
3805
3806 if let Some(sort_expr) = sort_expr {
3807 distinct_on = distinct_on.with_sort_expr(sort_expr)?;
3808 }
3809
3810 Ok(distinct_on)
3811 }
3812
3813 pub fn with_sort_expr(mut self, sort_expr: Vec<SortExpr>) -> Result<Self> {
3817 let sort_expr = normalize_sorts(sort_expr, self.input.as_ref())?;
3818
3819 let mut matched = true;
3821 for (on, sort) in self.on_expr.iter().zip(sort_expr.iter()) {
3822 if on != &sort.expr {
3823 matched = false;
3824 break;
3825 }
3826 }
3827
3828 if self.on_expr.len() > sort_expr.len() || !matched {
3829 return plan_err!(
3830 "SELECT DISTINCT ON expressions must match initial ORDER BY expressions"
3831 );
3832 }
3833
3834 self.sort_expr = Some(sort_expr);
3835 Ok(self)
3836 }
3837}
3838
3839impl PartialOrd for DistinctOn {
3841 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
3842 #[derive(PartialEq, PartialOrd)]
3843 struct ComparableDistinctOn<'a> {
3844 pub on_expr: &'a Vec<Expr>,
3846 pub select_expr: &'a Vec<Expr>,
3848 pub sort_expr: &'a Option<Vec<SortExpr>>,
3852 pub input: &'a Arc<LogicalPlan>,
3854 }
3855 let comparable_self = ComparableDistinctOn {
3856 on_expr: &self.on_expr,
3857 select_expr: &self.select_expr,
3858 sort_expr: &self.sort_expr,
3859 input: &self.input,
3860 };
3861 let comparable_other = ComparableDistinctOn {
3862 on_expr: &other.on_expr,
3863 select_expr: &other.select_expr,
3864 sort_expr: &other.sort_expr,
3865 input: &other.input,
3866 };
3867 comparable_self
3868 .partial_cmp(&comparable_other)
3869 .filter(|cmp| *cmp != Ordering::Equal || self == other)
3871 }
3872}
3873
3874#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3887#[non_exhaustive]
3889pub struct Aggregate {
3890 pub input: Arc<LogicalPlan>,
3892 pub group_expr: Vec<Expr>,
3894 pub aggr_expr: Vec<Expr>,
3898 pub schema: DFSchemaRef,
3900}
3901
3902impl Aggregate {
3903 pub fn try_new(
3905 input: Arc<LogicalPlan>,
3906 group_expr: Vec<Expr>,
3907 aggr_expr: Vec<Expr>,
3908 ) -> Result<Self> {
3909 check_aggregate_and_window_nesting(group_expr.iter().chain(aggr_expr.iter()))?;
3912
3913 let group_expr = enumerate_grouping_sets(group_expr)?;
3914
3915 let is_grouping_set = matches!(group_expr.as_slice(), [Expr::GroupingSet(_)]);
3916
3917 let grouping_expr: Vec<&Expr> = grouping_set_to_exprlist(group_expr.as_slice())?;
3918
3919 let mut qualified_fields = exprlist_to_fields(grouping_expr, &input)?;
3920
3921 if is_grouping_set {
3923 qualified_fields = qualified_fields
3924 .into_iter()
3925 .map(|(q, f)| (q, f.as_ref().clone().with_nullable(true).into()))
3926 .collect::<Vec<_>>();
3927 let max_ordinal = max_grouping_set_duplicate_ordinal(&group_expr);
3928 qualified_fields.push((
3929 None,
3930 Field::new(
3931 Self::INTERNAL_GROUPING_ID,
3932 Self::grouping_id_type(qualified_fields.len(), max_ordinal),
3933 false,
3934 )
3935 .into(),
3936 ));
3937 }
3938
3939 qualified_fields.extend(exprlist_to_fields(aggr_expr.as_slice(), &input)?);
3940
3941 let schema = DFSchema::new_with_metadata(
3942 qualified_fields,
3943 input.schema().metadata().clone(),
3944 )?;
3945
3946 Self::try_new_with_schema(input, group_expr, aggr_expr, Arc::new(schema))
3947 }
3948
3949 pub fn try_new_with_schema(
3955 input: Arc<LogicalPlan>,
3956 group_expr: Vec<Expr>,
3957 aggr_expr: Vec<Expr>,
3958 schema: DFSchemaRef,
3959 ) -> Result<Self> {
3960 if group_expr.is_empty() && aggr_expr.is_empty() {
3961 return plan_err!(
3962 "Aggregate requires at least one grouping or aggregate expression. \
3963 Aggregate without grouping expressions nor aggregate expressions is \
3964 logically equivalent to, but less efficient than, VALUES producing \
3965 single row. Please use VALUES instead."
3966 );
3967 }
3968 let group_expr_count = grouping_set_expr_count(&group_expr)?;
3969 if schema.fields().len() != group_expr_count + aggr_expr.len() {
3970 return plan_err!(
3971 "Aggregate schema has wrong number of fields. Expected {} got {}",
3972 group_expr_count + aggr_expr.len(),
3973 schema.fields().len()
3974 );
3975 }
3976
3977 let aggregate_func_dependencies =
3978 calc_func_dependencies_for_aggregate(&group_expr, &input, &schema)?;
3979 let new_schema = Arc::unwrap_or_clone(schema);
3980 let schema = Arc::new(
3981 new_schema.with_functional_dependencies(aggregate_func_dependencies)?,
3982 );
3983 Ok(Self {
3984 input,
3985 group_expr,
3986 aggr_expr,
3987 schema,
3988 })
3989 }
3990
3991 fn is_grouping_set(&self) -> bool {
3992 matches!(self.group_expr.as_slice(), [Expr::GroupingSet(_)])
3993 }
3994
3995 fn output_expressions(&self) -> Result<Vec<&Expr>> {
3997 static INTERNAL_ID_EXPR: LazyLock<Expr> = LazyLock::new(|| {
3998 Expr::Column(Column::from_name(Aggregate::INTERNAL_GROUPING_ID))
3999 });
4000 let mut exprs = grouping_set_to_exprlist(self.group_expr.as_slice())?;
4001 if self.is_grouping_set() {
4002 exprs.push(&INTERNAL_ID_EXPR);
4003 }
4004 exprs.extend(self.aggr_expr.iter());
4005 debug_assert!(exprs.len() == self.schema.fields().len());
4006 Ok(exprs)
4007 }
4008
4009 pub fn group_expr_len(&self) -> Result<usize> {
4013 grouping_set_expr_count(&self.group_expr)
4014 }
4015
4016 pub fn grouping_id_type(group_exprs: usize, max_ordinal: usize) -> DataType {
4028 let ordinal_bits = usize::BITS as usize - max_ordinal.leading_zeros() as usize;
4029 let total_bits = group_exprs + ordinal_bits;
4030 if total_bits <= 8 {
4031 DataType::UInt8
4032 } else if total_bits <= 16 {
4033 DataType::UInt16
4034 } else if total_bits <= 32 {
4035 DataType::UInt32
4036 } else {
4037 DataType::UInt64
4038 }
4039 }
4040
4041 pub const INTERNAL_GROUPING_ID: &'static str = "__grouping_id";
4074}
4075
4076impl PartialOrd for Aggregate {
4078 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
4079 match self.input.partial_cmp(&other.input) {
4080 Some(Ordering::Equal) => {
4081 match self.group_expr.partial_cmp(&other.group_expr) {
4082 Some(Ordering::Equal) => self.aggr_expr.partial_cmp(&other.aggr_expr),
4083 cmp => cmp,
4084 }
4085 }
4086 cmp => cmp,
4087 }
4088 .filter(|cmp| *cmp != Ordering::Equal || self == other)
4090 }
4091}
4092
4093#[allow(clippy::allow_attributes, clippy::mutable_key_type)] fn max_grouping_set_duplicate_ordinal(group_expr: &[Expr]) -> usize {
4101 if let Some(Expr::GroupingSet(GroupingSet::GroupingSets(sets))) = group_expr.first() {
4102 let mut counts: HashMap<&[Expr], usize> = HashMap::new();
4103 for set in sets {
4104 *counts.entry(set).or_insert(0) += 1;
4105 }
4106 counts.into_values().max().unwrap_or(0).saturating_sub(1)
4107 } else {
4108 0
4109 }
4110}
4111
4112fn contains_grouping_set(group_expr: &[Expr]) -> bool {
4114 group_expr
4115 .iter()
4116 .any(|expr| matches!(expr, Expr::GroupingSet(_)))
4117}
4118
4119fn calc_func_dependencies_for_aggregate(
4121 group_expr: &[Expr],
4123 input: &LogicalPlan,
4125 aggr_schema: &DFSchema,
4127) -> Result<FunctionalDependencies> {
4128 if !contains_grouping_set(group_expr) {
4134 let group_by_expr_names = group_expr
4135 .iter()
4136 .map(|item| item.schema_name().to_string())
4137 .collect::<IndexSet<_>>()
4138 .into_iter()
4139 .collect::<Vec<_>>();
4140 let aggregate_func_dependencies = aggregate_functional_dependencies(
4141 input.schema(),
4142 &group_by_expr_names,
4143 aggr_schema,
4144 );
4145 Ok(aggregate_func_dependencies)
4146 } else {
4147 Ok(FunctionalDependencies::empty())
4148 }
4149}
4150
4151fn calc_func_dependencies_for_project(
4154 exprs: &[Expr],
4155 input: &LogicalPlan,
4156) -> Result<FunctionalDependencies> {
4157 const COMPUTED_EXPR_INDEX: usize = usize::MAX;
4159
4160 let input_fields = input.schema().field_names();
4161 let proj_indices = exprs
4164 .iter()
4165 .map(|expr| match expr {
4166 #[expect(deprecated)]
4167 Expr::Wildcard { qualifier, options } => {
4168 let wildcard_fields = exprlist_to_fields(
4169 vec![&Expr::Wildcard {
4170 qualifier: qualifier.clone(),
4171 options: options.clone(),
4172 }],
4173 input,
4174 )?;
4175 Ok::<_, DataFusionError>(
4176 wildcard_fields
4177 .into_iter()
4178 .map(|(qualifier, f)| {
4179 let flat_name = qualifier
4180 .map(|t| format!("{}.{}", t, f.name()))
4181 .unwrap_or_else(|| f.name().clone());
4182 input_fields
4183 .iter()
4184 .position(|item| *item == flat_name)
4185 .unwrap_or(COMPUTED_EXPR_INDEX)
4186 })
4187 .collect::<Vec<_>>(),
4188 )
4189 }
4190 Expr::Alias(alias) => {
4191 let name = format!("{}", alias.expr);
4192 let input_index = input_fields
4193 .iter()
4194 .position(|item| *item == name)
4195 .unwrap_or(COMPUTED_EXPR_INDEX);
4196 Ok(vec![input_index])
4197 }
4198 _ => {
4199 let name = format!("{expr}");
4200 let input_index = input_fields
4201 .iter()
4202 .position(|item| *item == name)
4203 .unwrap_or(COMPUTED_EXPR_INDEX);
4204 Ok(vec![input_index])
4205 }
4206 })
4207 .collect::<Result<Vec<_>>>()?
4208 .into_iter()
4209 .flatten()
4210 .collect::<Vec<_>>();
4211
4212 Ok(input
4213 .schema()
4214 .functional_dependencies()
4215 .project_functional_dependencies(&proj_indices, exprs.len()))
4216}
4217
4218#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
4220pub struct Sort {
4221 pub expr: Vec<SortExpr>,
4223 pub input: Arc<LogicalPlan>,
4225 pub fetch: Option<usize>,
4227}
4228
4229#[derive(Debug, Clone, PartialEq, Eq, Hash)]
4231pub struct Join {
4232 pub left: Arc<LogicalPlan>,
4234 pub right: Arc<LogicalPlan>,
4236 pub on: Vec<(Expr, Expr)>,
4238 pub filter: Option<Expr>,
4240 pub join_type: JoinType,
4242 pub join_constraint: JoinConstraint,
4244 pub schema: DFSchemaRef,
4246 pub null_equality: NullEquality,
4248 pub null_aware: bool,
4256}
4257
4258impl Join {
4259 #[expect(clippy::too_many_arguments)]
4279 pub fn try_new(
4280 left: Arc<LogicalPlan>,
4281 right: Arc<LogicalPlan>,
4282 on: Vec<(Expr, Expr)>,
4283 filter: Option<Expr>,
4284 join_type: JoinType,
4285 join_constraint: JoinConstraint,
4286 null_equality: NullEquality,
4287 null_aware: bool,
4288 ) -> Result<Self> {
4289 let join_schema = build_join_schema(left.schema(), right.schema(), &join_type)?;
4290
4291 Ok(Join {
4292 left,
4293 right,
4294 on,
4295 filter,
4296 join_type,
4297 join_constraint,
4298 schema: Arc::new(join_schema),
4299 null_equality,
4300 null_aware,
4301 })
4302 }
4303
4304 pub fn try_new_with_project_input(
4307 original: &LogicalPlan,
4308 left: Arc<LogicalPlan>,
4309 right: Arc<LogicalPlan>,
4310 column_on: (Vec<Column>, Vec<Column>),
4311 ) -> Result<(Self, bool)> {
4312 let original_join = match original {
4313 LogicalPlan::Join(join) => join,
4314 _ => return plan_err!("Could not create join with project input"),
4315 };
4316
4317 let mut left_sch = LogicalPlanBuilder::from(Arc::clone(&left));
4318 let mut right_sch = LogicalPlanBuilder::from(Arc::clone(&right));
4319
4320 let mut requalified = false;
4321
4322 if original_join.join_type == JoinType::Inner
4325 || original_join.join_type == JoinType::Left
4326 || original_join.join_type == JoinType::Right
4327 || original_join.join_type == JoinType::Full
4328 {
4329 (left_sch, right_sch, requalified) =
4330 requalify_sides_if_needed(left_sch.clone(), right_sch.clone())?;
4331 }
4332
4333 let on: Vec<(Expr, Expr)> = column_on
4334 .0
4335 .into_iter()
4336 .zip(column_on.1)
4337 .map(|(l, r)| (Expr::Column(l), Expr::Column(r)))
4338 .collect();
4339
4340 let join_schema = build_join_schema(
4341 left_sch.schema(),
4342 right_sch.schema(),
4343 &original_join.join_type,
4344 )?;
4345
4346 Ok((
4347 Join {
4348 left,
4349 right,
4350 on,
4351 filter: original_join.filter.clone(),
4352 join_type: original_join.join_type,
4353 join_constraint: original_join.join_constraint,
4354 schema: Arc::new(join_schema),
4355 null_equality: original_join.null_equality,
4356 null_aware: original_join.null_aware,
4357 },
4358 requalified,
4359 ))
4360 }
4361}
4362
4363impl PartialOrd for Join {
4365 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
4366 #[derive(PartialEq, PartialOrd)]
4367 struct ComparableJoin<'a> {
4368 pub left: &'a Arc<LogicalPlan>,
4370 pub right: &'a Arc<LogicalPlan>,
4372 pub on: &'a Vec<(Expr, Expr)>,
4374 pub filter: &'a Option<Expr>,
4376 pub join_type: &'a JoinType,
4378 pub join_constraint: &'a JoinConstraint,
4380 pub null_equality: &'a NullEquality,
4382 }
4383 let comparable_self = ComparableJoin {
4384 left: &self.left,
4385 right: &self.right,
4386 on: &self.on,
4387 filter: &self.filter,
4388 join_type: &self.join_type,
4389 join_constraint: &self.join_constraint,
4390 null_equality: &self.null_equality,
4391 };
4392 let comparable_other = ComparableJoin {
4393 left: &other.left,
4394 right: &other.right,
4395 on: &other.on,
4396 filter: &other.filter,
4397 join_type: &other.join_type,
4398 join_constraint: &other.join_constraint,
4399 null_equality: &other.null_equality,
4400 };
4401 comparable_self
4402 .partial_cmp(&comparable_other)
4403 .filter(|cmp| *cmp != Ordering::Equal || self == other)
4405 }
4406}
4407
4408#[derive(Clone, PartialEq, Eq, PartialOrd, Hash)]
4410pub struct Subquery {
4411 pub subquery: Arc<LogicalPlan>,
4413 pub outer_ref_columns: Vec<Expr>,
4415 pub spans: Spans,
4417}
4418
4419impl Normalizeable for Subquery {
4420 fn can_normalize(&self) -> bool {
4421 false
4422 }
4423}
4424
4425impl NormalizeEq for Subquery {
4426 fn normalize_eq(&self, other: &Self) -> bool {
4427 *self.subquery == *other.subquery
4429 && self.outer_ref_columns.len() == other.outer_ref_columns.len()
4430 && self
4431 .outer_ref_columns
4432 .iter()
4433 .zip(other.outer_ref_columns.iter())
4434 .all(|(a, b)| a.normalize_eq(b))
4435 }
4436}
4437
4438impl Subquery {
4439 pub fn try_from_expr(plan: &Expr) -> Result<&Subquery> {
4440 match plan {
4441 Expr::ScalarSubquery(it) => Ok(it),
4442 Expr::Cast(cast) => Subquery::try_from_expr(cast.expr.as_ref()),
4443 _ => plan_err!("Could not coerce into ScalarSubquery!"),
4444 }
4445 }
4446
4447 pub fn with_plan(&self, plan: Arc<LogicalPlan>) -> Subquery {
4448 Subquery {
4449 subquery: plan,
4450 outer_ref_columns: self.outer_ref_columns.clone(),
4451 spans: Spans::new(),
4452 }
4453 }
4454}
4455
4456impl Debug for Subquery {
4457 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
4458 write!(f, "<subquery>")
4459 }
4460}
4461
4462#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
4473pub enum Partitioning {
4474 RoundRobinBatch(usize),
4476 Hash(Vec<Expr>, usize),
4479 Range(RangePartitioning),
4482 DistributeBy(Vec<Expr>),
4484}
4485
4486impl Partitioning {
4487 pub fn partition_count(&self) -> Option<usize> {
4489 match self {
4490 Self::RoundRobinBatch(partition_count) | Self::Hash(_, partition_count) => {
4491 Some(*partition_count)
4492 }
4493 Self::Range(range) => Some(range.partition_count()),
4494 Self::DistributeBy(_) => None,
4495 }
4496 }
4497}
4498
4499#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
4523pub struct RangePartitioning {
4524 ordering: Vec<SortExpr>,
4526 split_points: Vec<SplitPoint>,
4528}
4529
4530impl RangePartitioning {
4531 pub fn try_new(
4534 ordering: Vec<SortExpr>,
4535 split_points: Vec<SplitPoint>,
4536 ) -> Result<Self> {
4537 if ordering.is_empty() {
4538 return plan_err!("Range partitioning requires non-empty ordering");
4539 }
4540
4541 validate_range_split_points(&split_points, &logical_sort_options(&ordering))?;
4542
4543 Ok(Self {
4544 ordering,
4545 split_points,
4546 })
4547 }
4548
4549 pub fn partition_count(&self) -> usize {
4551 self.split_points.len() + 1
4552 }
4553
4554 pub fn ordering(&self) -> &[SortExpr] {
4556 &self.ordering
4557 }
4558
4559 pub fn split_points(&self) -> &[SplitPoint] {
4561 &self.split_points
4562 }
4563}
4564
4565fn logical_sort_options(ordering: &[SortExpr]) -> Vec<SortOptions> {
4566 ordering
4567 .iter()
4568 .map(|sort_expr| SortOptions {
4569 descending: !sort_expr.asc,
4570 nulls_first: sort_expr.nulls_first,
4571 })
4572 .collect()
4573}
4574
4575impl Display for RangePartitioning {
4576 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
4577 let ordering = self.ordering().iter().map(ToString::to_string).join(", ");
4578 let split_points = self
4579 .split_points()
4580 .iter()
4581 .map(ToString::to_string)
4582 .join(", ");
4583 write!(
4584 f,
4585 "Range([{ordering}], [{split_points}], {})",
4586 self.partition_count()
4587 )
4588 }
4589}
4590
4591#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)]
4611pub struct ColumnUnnestList {
4612 pub output_column: Column,
4613 pub depth: usize,
4614}
4615
4616impl Display for ColumnUnnestList {
4617 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
4618 write!(f, "{}|depth={}", self.output_column, self.depth)
4619 }
4620}
4621
4622#[derive(Debug, Clone, PartialEq, Eq, Hash)]
4625pub struct Unnest {
4626 pub input: Arc<LogicalPlan>,
4628 pub exec_columns: Vec<Column>,
4630 pub list_type_columns: Vec<(usize, ColumnUnnestList)>,
4633 pub struct_type_columns: Vec<usize>,
4636 pub dependency_indices: Vec<usize>,
4639 pub schema: DFSchemaRef,
4641 pub options: UnnestOptions,
4643}
4644
4645impl PartialOrd for Unnest {
4647 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
4648 #[derive(PartialEq, PartialOrd)]
4649 struct ComparableUnnest<'a> {
4650 pub input: &'a Arc<LogicalPlan>,
4652 pub exec_columns: &'a Vec<Column>,
4654 pub list_type_columns: &'a Vec<(usize, ColumnUnnestList)>,
4657 pub struct_type_columns: &'a Vec<usize>,
4660 pub dependency_indices: &'a Vec<usize>,
4663 pub options: &'a UnnestOptions,
4665 }
4666 let comparable_self = ComparableUnnest {
4667 input: &self.input,
4668 exec_columns: &self.exec_columns,
4669 list_type_columns: &self.list_type_columns,
4670 struct_type_columns: &self.struct_type_columns,
4671 dependency_indices: &self.dependency_indices,
4672 options: &self.options,
4673 };
4674 let comparable_other = ComparableUnnest {
4675 input: &other.input,
4676 exec_columns: &other.exec_columns,
4677 list_type_columns: &other.list_type_columns,
4678 struct_type_columns: &other.struct_type_columns,
4679 dependency_indices: &other.dependency_indices,
4680 options: &other.options,
4681 };
4682 comparable_self
4683 .partial_cmp(&comparable_other)
4684 .filter(|cmp| *cmp != Ordering::Equal || self == other)
4686 }
4687}
4688
4689impl Unnest {
4690 pub fn try_new(
4691 input: Arc<LogicalPlan>,
4692 exec_columns: Vec<Column>,
4693 options: UnnestOptions,
4694 ) -> Result<Self> {
4695 if exec_columns.is_empty() {
4696 return plan_err!("unnest plan requires at least 1 column to unnest");
4697 }
4698
4699 let mut list_columns: Vec<(usize, ColumnUnnestList)> = vec![];
4700 let mut struct_columns = vec![];
4701 let indices_to_unnest = exec_columns
4702 .iter()
4703 .map(|c| Ok((input.schema().index_of_column(c)?, c)))
4704 .collect::<Result<HashMap<usize, &Column>>>()?;
4705
4706 let input_schema = input.schema();
4707
4708 let mut dependency_indices = vec![];
4709 let fields = input_schema
4725 .iter()
4726 .enumerate()
4727 .map(|(index, (original_qualifier, original_field))| {
4728 match indices_to_unnest.get(&index) {
4729 Some(column_to_unnest) => {
4730 let recursions_on_column = options
4731 .recursions
4732 .iter()
4733 .filter(|p| -> bool { &p.input_column == *column_to_unnest })
4734 .collect::<Vec<_>>();
4735 let mut transformed_columns = recursions_on_column
4736 .iter()
4737 .map(|r| {
4738 list_columns.push((
4739 index,
4740 ColumnUnnestList {
4741 output_column: r.output_column.clone(),
4742 depth: r.depth,
4743 },
4744 ));
4745 Ok(get_unnested_columns(
4746 &r.output_column.name,
4747 original_field.data_type(),
4748 r.depth,
4749 )?
4750 .into_iter()
4751 .next()
4752 .unwrap()) })
4754 .collect::<Result<Vec<(Column, Arc<Field>)>>>()?;
4755 if transformed_columns.is_empty() {
4756 transformed_columns = get_unnested_columns(
4757 &column_to_unnest.name,
4758 original_field.data_type(),
4759 1,
4760 )?;
4761 match original_field.data_type() {
4762 DataType::Struct(_) => {
4763 struct_columns.push(index);
4764 }
4765 DataType::List(_)
4766 | DataType::FixedSizeList(_, _)
4767 | DataType::LargeList(_)
4768 | DataType::ListView(_)
4769 | DataType::LargeListView(_) => {
4770 list_columns.push((
4771 index,
4772 ColumnUnnestList {
4773 output_column: Column::from_name(
4774 &column_to_unnest.name,
4775 ),
4776 depth: 1,
4777 },
4778 ));
4779 }
4780 _ => {}
4781 };
4782 }
4783
4784 dependency_indices.extend(std::iter::repeat_n(
4786 index,
4787 transformed_columns.len(),
4788 ));
4789 Ok(transformed_columns
4790 .iter()
4791 .map(|(col, field)| {
4792 (col.relation.to_owned(), field.to_owned())
4793 })
4794 .collect())
4795 }
4796 None => {
4797 dependency_indices.push(index);
4798 Ok(vec![(
4799 original_qualifier.cloned(),
4800 Arc::clone(original_field),
4801 )])
4802 }
4803 }
4804 })
4805 .collect::<Result<Vec<_>>>()?
4806 .into_iter()
4807 .flatten()
4808 .collect::<Vec<_>>();
4809
4810 let metadata = input_schema.metadata().clone();
4811 let df_schema = DFSchema::new_with_metadata(fields, metadata)?;
4812 let deps = input_schema.functional_dependencies().clone();
4814 let schema = Arc::new(df_schema.with_functional_dependencies(deps)?);
4815
4816 Ok(Unnest {
4817 input,
4818 exec_columns,
4819 list_type_columns: list_columns,
4820 struct_type_columns: struct_columns,
4821 dependency_indices,
4822 schema,
4823 options,
4824 })
4825 }
4826}
4827
4828fn get_unnested_columns(
4837 col_name: &String,
4838 data_type: &DataType,
4839 depth: usize,
4840) -> Result<Vec<(Column, Arc<Field>)>> {
4841 let mut qualified_columns = Vec::with_capacity(1);
4842
4843 match data_type {
4844 DataType::List(_)
4845 | DataType::FixedSizeList(_, _)
4846 | DataType::LargeList(_)
4847 | DataType::ListView(_)
4848 | DataType::LargeListView(_) => {
4849 let data_type = get_unnested_list_datatype_recursive(data_type, depth)?;
4850 let new_field = Arc::new(Field::new(
4851 col_name, data_type,
4852 true,
4855 ));
4856 let column = Column::from_name(col_name);
4857 qualified_columns.push((column, new_field));
4859 }
4860 DataType::Struct(fields) => {
4861 qualified_columns.extend(fields.iter().map(|f| {
4862 let new_name = format!("{}.{}", col_name, f.name());
4863 let column = Column::from_name(&new_name);
4864 let new_field = f.as_ref().clone().with_name(new_name);
4865 (column, Arc::new(new_field))
4867 }))
4868 }
4869 _ => {
4870 return internal_err!("trying to unnest on invalid data type {data_type}");
4871 }
4872 };
4873 Ok(qualified_columns)
4874}
4875
4876fn get_unnested_list_datatype_recursive(
4879 data_type: &DataType,
4880 depth: usize,
4881) -> Result<DataType> {
4882 match data_type {
4883 DataType::List(field)
4884 | DataType::FixedSizeList(field, _)
4885 | DataType::LargeList(field)
4886 | DataType::ListView(field)
4887 | DataType::LargeListView(field) => {
4888 if depth == 1 {
4889 return Ok(field.data_type().clone());
4890 }
4891 return get_unnested_list_datatype_recursive(field.data_type(), depth - 1);
4892 }
4893 _ => {}
4894 };
4895
4896 internal_err!("trying to unnest on invalid data type {data_type}")
4897}
4898
4899#[cfg(test)]
4900mod tests {
4901 use super::*;
4902 use crate::builder::LogicalTableSource;
4903 use crate::logical_plan::table_scan;
4904 use crate::select_expr::SelectExpr;
4905 use crate::test::function_stub::{count, count_udaf};
4906 use crate::{
4907 GroupingSet, binary_expr, col, exists, in_subquery, lit, placeholder,
4908 scalar_subquery,
4909 };
4910 use datafusion_common::metadata::ScalarAndMetadata;
4911 use datafusion_common::tree_node::{
4912 TransformedResult, TreeNodeRewriter, TreeNodeVisitor,
4913 };
4914 use datafusion_common::{Constraint, not_impl_err};
4915 use insta::{assert_debug_snapshot, assert_snapshot};
4916 use std::hash::DefaultHasher;
4917
4918 #[test]
4928 fn test_size_of_logical_plan() {
4929 assert_eq!(size_of::<LogicalPlan>(), 176);
4934 assert!(
4937 size_of::<DdlStatement>() < size_of::<Join>(),
4938 "DdlStatement ({} bytes) should stay smaller than Join ({} bytes); \
4939 box the new large variant rather than letting it dominate `LogicalPlan`.",
4940 size_of::<DdlStatement>(),
4941 size_of::<Join>(),
4942 );
4943 assert_eq!(
4946 size_of::<Box<crate::CreateExternalTable>>(),
4947 8,
4948 "CreateExternalTable should be Box'd inside DdlStatement"
4949 );
4950 assert_eq!(
4951 size_of::<Box<crate::CreateFunction>>(),
4952 8,
4953 "CreateFunction should be Box'd inside DdlStatement"
4954 );
4955 }
4956
4957 fn employee_schema() -> Schema {
4958 Schema::new(vec![
4959 Field::new("id", DataType::Int32, false),
4960 Field::new("first_name", DataType::Utf8, false),
4961 Field::new("last_name", DataType::Utf8, false),
4962 Field::new("state", DataType::Utf8, false),
4963 Field::new("salary", DataType::Int32, false),
4964 ])
4965 }
4966
4967 #[test]
4968 fn projection_with_leading_computed_column_preserves_pk() -> Result<()> {
4969 let constraints =
4970 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
4971 let source = Arc::new(
4972 LogicalTableSource::new(Arc::new(employee_schema()))
4973 .with_constraints(constraints),
4974 );
4975 let plan = LogicalPlanBuilder::scan("employee_csv", source, None)?
4976 .project(vec![
4977 lit(1i32).alias("__common_expr_1"),
4978 col("id"),
4979 col("first_name"),
4980 col("salary"),
4981 ])?
4982 .build()?;
4983
4984 let deps = plan.schema().functional_dependencies();
4985 assert_eq!(deps.len(), 1);
4986 assert_eq!(deps[0].source_indices, vec![1]);
4987
4988 Ok(())
4989 }
4990
4991 #[test]
4992 fn projection_with_leading_computed_column_and_wildcard_preserves_pk() -> Result<()> {
4993 let constraints =
4994 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
4995 let source = Arc::new(
4996 LogicalTableSource::new(Arc::new(employee_schema()))
4997 .with_constraints(constraints),
4998 );
4999 let plan = LogicalPlanBuilder::scan("employee_csv", source, None)?
5000 .project(vec![
5001 SelectExpr::Expression(lit(1i32).alias("__common_expr_1")),
5002 SelectExpr::Wildcard(Default::default()),
5003 ])?
5004 .build()?;
5005
5006 let deps = plan.schema().functional_dependencies();
5007 assert_eq!(plan.schema().fields().len(), 6);
5008 assert_eq!(deps.len(), 1);
5009 assert_eq!(deps[0].source_indices, vec![1]);
5010 assert_eq!(deps[0].target_indices, vec![0, 1, 2, 3, 4, 5]);
5011
5012 Ok(())
5013 }
5014
5015 #[test]
5016 fn projection_with_wildcard_expr_before_pk_preserves_pk() -> Result<()> {
5017 let constraints =
5018 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
5019 let source = Arc::new(
5020 LogicalTableSource::new(Arc::new(employee_schema()))
5021 .with_constraints(constraints),
5022 );
5023 let input = LogicalPlanBuilder::scan("employee_csv", source, None)?.build()?;
5024 #[expect(deprecated)]
5025 let projection = Projection::try_new(
5026 vec![
5027 Expr::Wildcard {
5028 qualifier: None,
5029 options: Box::new(crate::expr::WildcardOptions::default()),
5030 },
5031 col("employee_csv.id"),
5032 ],
5033 Arc::new(input),
5034 )?;
5035
5036 let deps = projection.schema.functional_dependencies();
5037 assert_eq!(deps.len(), 1);
5038 assert_eq!(deps[0].source_indices, vec![1]);
5039
5040 Ok(())
5041 }
5042
5043 fn i32_split_point(value: i32) -> SplitPoint {
5044 SplitPoint::new(vec![ScalarValue::Int32(Some(value))])
5045 }
5046
5047 fn null_i32_split_point() -> SplitPoint {
5048 SplitPoint::new(vec![ScalarValue::Int32(None)])
5049 }
5050
5051 #[test]
5052 fn logical_range_partitioning_validates_shape() {
5053 let range = RangePartitioning::try_new(
5054 vec![col("id").sort(true, true)],
5055 vec![i32_split_point(10), i32_split_point(20)],
5056 )
5057 .unwrap();
5058 assert_eq!(range.partition_count(), 3);
5059
5060 let range = RangePartitioning::try_new(
5061 vec![col("id").sort(false, true)],
5062 vec![i32_split_point(20), i32_split_point(10)],
5063 )
5064 .unwrap();
5065 assert_eq!(range.partition_count(), 3);
5066
5067 let err = RangePartitioning::try_new(vec![], vec![]).unwrap_err();
5068 assert!(err.to_string().contains("non-empty ordering"));
5069
5070 let err = RangePartitioning::try_new(
5071 vec![col("id").sort(true, true), col("salary").sort(true, true)],
5072 vec![i32_split_point(10)],
5073 )
5074 .unwrap_err();
5075 assert!(
5076 err.to_string()
5077 .contains("split point 0 has width 1, but ordering has width 2")
5078 );
5079
5080 let err = RangePartitioning::try_new(
5081 vec![col("id").sort(true, true)],
5082 vec![i32_split_point(20), i32_split_point(10)],
5083 )
5084 .unwrap_err();
5085 assert!(
5086 err.to_string()
5087 .contains("split points must be strictly ordered")
5088 );
5089
5090 let err = RangePartitioning::try_new(
5091 vec![col("id").sort(true, true)],
5092 vec![i32_split_point(10), i32_split_point(10)],
5093 )
5094 .unwrap_err();
5095 assert!(
5096 err.to_string()
5097 .contains("split points must be strictly ordered")
5098 );
5099
5100 let range = RangePartitioning::try_new(
5101 vec![col("id").sort(true, true)],
5102 vec![null_i32_split_point(), i32_split_point(10)],
5103 )
5104 .unwrap();
5105 assert_eq!(range.partition_count(), 3);
5106 }
5107
5108 #[test]
5109 fn logical_partitioning_reports_known_partition_count() -> Result<()> {
5110 let range = RangePartitioning::try_new(
5111 vec![col("id").sort(true, true)],
5112 vec![i32_split_point(10)],
5113 )?;
5114
5115 assert_eq!(Partitioning::RoundRobinBatch(4).partition_count(), Some(4));
5116 assert_eq!(
5117 Partitioning::Hash(vec![col("id")], 8).partition_count(),
5118 Some(8)
5119 );
5120 assert_eq!(Partitioning::Range(range).partition_count(), Some(2));
5121 assert_eq!(
5122 Partitioning::DistributeBy(vec![col("id")]).partition_count(),
5123 None
5124 );
5125
5126 Ok(())
5127 }
5128
5129 #[test]
5130 fn logical_range_partitioning_participates_in_expression_rewrite() -> Result<()> {
5131 let input =
5132 table_scan(Some("employee_csv"), &employee_schema(), None)?.build()?;
5133 let plan = LogicalPlan::Repartition(Repartition {
5134 input: Arc::new(input),
5135 partitioning_scheme: Partitioning::Range(RangePartitioning::try_new(
5136 vec![col("id").sort(true, true)],
5137 vec![i32_split_point(10)],
5138 )?),
5139 });
5140
5141 let mut visited_exprs = vec![];
5142 plan.apply_expressions(|expr| {
5143 visited_exprs.push(expr.to_string());
5144 Ok(TreeNodeRecursion::Continue)
5145 })?;
5146 assert_eq!(visited_exprs, vec!["id"]);
5147
5148 let plan = plan
5149 .map_expressions(|expr| {
5150 if expr == col("id") {
5151 Ok(Transformed::yes(col("salary")))
5152 } else {
5153 Ok(Transformed::no(expr))
5154 }
5155 })?
5156 .data;
5157
5158 let LogicalPlan::Repartition(Repartition {
5159 partitioning_scheme: Partitioning::Range(range),
5160 ..
5161 }) = plan
5162 else {
5163 unreachable!("expected range repartition");
5164 };
5165 assert_eq!(range.ordering()[0].expr, col("salary"));
5166 assert_eq!(range.partition_count(), 2);
5167
5168 Ok(())
5169 }
5170
5171 fn display_plan() -> Result<LogicalPlan> {
5172 let plan1 = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3]))?
5173 .build()?;
5174
5175 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
5176 .filter(in_subquery(col("state"), Arc::new(plan1)))?
5177 .project(vec![col("id")])?
5178 .build()
5179 }
5180
5181 fn recursive_term_scan(name: &str, fields: Vec<Field>) -> Result<Arc<LogicalPlan>> {
5182 Ok(Arc::new(
5183 table_scan(Some(name), &Schema::new(fields), None)?.build()?,
5184 ))
5185 }
5186
5187 #[test]
5188 fn recursive_query_widens_nullability_per_column() -> Result<()> {
5189 let static_term = recursive_term_scan(
5193 "static",
5194 vec![
5195 Field::new("a", DataType::Int32, false),
5196 Field::new("b", DataType::Int32, false),
5197 ],
5198 )?;
5199 let recursive_term = recursive_term_scan(
5200 "rec",
5201 vec![
5202 Field::new("a", DataType::Int32, false),
5203 Field::new("b", DataType::Int32, true),
5204 ],
5205 )?;
5206
5207 let query =
5208 RecursiveQuery::try_new("t".to_string(), static_term, recursive_term, false)?;
5209
5210 assert_eq!(query.schema.field(0).name(), "a");
5212 assert_eq!(query.schema.field(1).name(), "b");
5213 assert_eq!(query.schema.field(0).data_type(), &DataType::Int32);
5214 assert_eq!(query.schema.field(1).data_type(), &DataType::Int32);
5215 assert!(!query.schema.field(0).is_nullable());
5217 assert!(query.schema.field(1).is_nullable());
5218 assert_eq!(
5220 LogicalPlan::RecursiveQuery(query.clone()).schema(),
5221 &query.schema
5222 );
5223 Ok(())
5224 }
5225
5226 #[test]
5227 fn recursive_query_rejects_column_count_mismatch() -> Result<()> {
5228 let static_term =
5229 recursive_term_scan("static", vec![Field::new("a", DataType::Int32, false)])?;
5230 let recursive_term = recursive_term_scan(
5231 "rec",
5232 vec![
5233 Field::new("a", DataType::Int32, false),
5234 Field::new("b", DataType::Int32, false),
5235 ],
5236 )?;
5237
5238 let err =
5239 RecursiveQuery::try_new("t".to_string(), static_term, recursive_term, false)
5240 .unwrap_err();
5241 assert!(
5242 err.strip_backtrace()
5243 .contains("must have the same number of columns"),
5244 "unexpected error: {err}"
5245 );
5246 Ok(())
5247 }
5248
5249 #[test]
5250 fn test_display_indent() -> Result<()> {
5251 let plan = display_plan()?;
5252
5253 assert_snapshot!(plan.display_indent(), @r"
5254 Projection: employee_csv.id
5255 Filter: employee_csv.state IN (<subquery>)
5256 Subquery:
5257 TableScan: employee_csv projection=[state]
5258 TableScan: employee_csv projection=[id, state]
5259 ");
5260 Ok(())
5261 }
5262
5263 #[test]
5264 fn test_display_indent_schema() -> Result<()> {
5265 let plan = display_plan()?;
5266
5267 assert_snapshot!(plan.display_indent_schema(), @r"
5268 Projection: employee_csv.id [id:Int32]
5269 Filter: employee_csv.state IN (<subquery>) [id:Int32, state:Utf8]
5270 Subquery: [state:Utf8]
5271 TableScan: employee_csv projection=[state] [state:Utf8]
5272 TableScan: employee_csv projection=[id, state] [id:Int32, state:Utf8]
5273 ");
5274 Ok(())
5275 }
5276
5277 #[test]
5278 fn test_display_subquery_alias() -> Result<()> {
5279 let plan1 = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3]))?
5280 .build()?;
5281 let plan1 = Arc::new(plan1);
5282
5283 let plan =
5284 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
5285 .project(vec![col("id"), exists(plan1).alias("exists")])?
5286 .build();
5287
5288 assert_snapshot!(plan?.display_indent(), @r"
5289 Projection: employee_csv.id, EXISTS (<subquery>) AS exists
5290 Subquery:
5291 TableScan: employee_csv projection=[state]
5292 TableScan: employee_csv projection=[id, state]
5293 ");
5294 Ok(())
5295 }
5296
5297 #[test]
5298 fn test_display_graphviz() -> Result<()> {
5299 let plan = display_plan()?;
5300
5301 assert_snapshot!(plan.display_graphviz(), @r#"
5304 // Begin DataFusion GraphViz Plan,
5305 // display it online here: https://dreampuf.github.io/GraphvizOnline
5306
5307 digraph {
5308 subgraph cluster_1
5309 {
5310 graph[label="LogicalPlan"]
5311 2[shape=box label="Projection: employee_csv.id"]
5312 3[shape=box label="Filter: employee_csv.state IN (<subquery>)"]
5313 2 -> 3 [arrowhead=none, arrowtail=normal, dir=back]
5314 4[shape=box label="Subquery:"]
5315 3 -> 4 [arrowhead=none, arrowtail=normal, dir=back]
5316 5[shape=box label="TableScan: employee_csv projection=[state]"]
5317 4 -> 5 [arrowhead=none, arrowtail=normal, dir=back]
5318 6[shape=box label="TableScan: employee_csv projection=[id, state]"]
5319 3 -> 6 [arrowhead=none, arrowtail=normal, dir=back]
5320 }
5321 subgraph cluster_7
5322 {
5323 graph[label="Detailed LogicalPlan"]
5324 8[shape=box label="Projection: employee_csv.id\nSchema: [id:Int32]"]
5325 9[shape=box label="Filter: employee_csv.state IN (<subquery>)\nSchema: [id:Int32, state:Utf8]"]
5326 8 -> 9 [arrowhead=none, arrowtail=normal, dir=back]
5327 10[shape=box label="Subquery:\nSchema: [state:Utf8]"]
5328 9 -> 10 [arrowhead=none, arrowtail=normal, dir=back]
5329 11[shape=box label="TableScan: employee_csv projection=[state]\nSchema: [state:Utf8]"]
5330 10 -> 11 [arrowhead=none, arrowtail=normal, dir=back]
5331 12[shape=box label="TableScan: employee_csv projection=[id, state]\nSchema: [id:Int32, state:Utf8]"]
5332 9 -> 12 [arrowhead=none, arrowtail=normal, dir=back]
5333 }
5334 }
5335 // End DataFusion GraphViz Plan
5336 "#);
5337 Ok(())
5338 }
5339
5340 #[test]
5341 fn test_display_pg_json() -> Result<()> {
5342 let plan = display_plan()?;
5343
5344 assert_snapshot!(plan.display_pg_json(), @r#"
5345 [
5346 {
5347 "Plan": {
5348 "Node Type": "Projection",
5349 "Expressions": [
5350 "employee_csv.id"
5351 ],
5352 "Plans": [
5353 {
5354 "Node Type": "Filter",
5355 "Condition": "employee_csv.state IN (<subquery>)",
5356 "Plans": [
5357 {
5358 "Node Type": "Subquery",
5359 "Plans": [
5360 {
5361 "Node Type": "TableScan",
5362 "Relation Name": "employee_csv",
5363 "Plans": [],
5364 "Output": [
5365 "state"
5366 ]
5367 }
5368 ],
5369 "Output": [
5370 "state"
5371 ]
5372 },
5373 {
5374 "Node Type": "TableScan",
5375 "Relation Name": "employee_csv",
5376 "Plans": [],
5377 "Output": [
5378 "id",
5379 "state"
5380 ]
5381 }
5382 ],
5383 "Output": [
5384 "id",
5385 "state"
5386 ]
5387 }
5388 ],
5389 "Output": [
5390 "id"
5391 ]
5392 }
5393 }
5394 ]
5395 "#);
5396 Ok(())
5397 }
5398
5399 #[derive(Debug, Default)]
5401 struct OkVisitor {
5402 strings: Vec<String>,
5403 }
5404
5405 impl<'n> TreeNodeVisitor<'n> for OkVisitor {
5406 type Node = LogicalPlan;
5407
5408 fn f_down(&mut self, plan: &'n LogicalPlan) -> Result<TreeNodeRecursion> {
5409 let s = match plan {
5410 LogicalPlan::Projection { .. } => "pre_visit Projection",
5411 LogicalPlan::Filter { .. } => "pre_visit Filter",
5412 LogicalPlan::TableScan { .. } => "pre_visit TableScan",
5413 _ => {
5414 return not_impl_err!("unknown plan type");
5415 }
5416 };
5417
5418 self.strings.push(s.into());
5419 Ok(TreeNodeRecursion::Continue)
5420 }
5421
5422 fn f_up(&mut self, plan: &'n LogicalPlan) -> Result<TreeNodeRecursion> {
5423 let s = match plan {
5424 LogicalPlan::Projection { .. } => "post_visit Projection",
5425 LogicalPlan::Filter { .. } => "post_visit Filter",
5426 LogicalPlan::TableScan { .. } => "post_visit TableScan",
5427 _ => {
5428 return not_impl_err!("unknown plan type");
5429 }
5430 };
5431
5432 self.strings.push(s.into());
5433 Ok(TreeNodeRecursion::Continue)
5434 }
5435 }
5436
5437 #[test]
5438 fn visit_order() {
5439 let mut visitor = OkVisitor::default();
5440 let plan = test_plan();
5441 let res = plan.visit_with_subqueries(&mut visitor);
5442 assert!(res.is_ok());
5443
5444 assert_debug_snapshot!(visitor.strings, @r#"
5445 [
5446 "pre_visit Projection",
5447 "pre_visit Filter",
5448 "pre_visit TableScan",
5449 "post_visit TableScan",
5450 "post_visit Filter",
5451 "post_visit Projection",
5452 ]
5453 "#);
5454 }
5455
5456 #[derive(Debug, Default)]
5457 struct OptionalCounter {
5459 val: Option<usize>,
5460 }
5461
5462 impl OptionalCounter {
5463 fn new(val: usize) -> Self {
5464 Self { val: Some(val) }
5465 }
5466 fn dec(&mut self) -> bool {
5468 if Some(0) == self.val {
5469 true
5470 } else {
5471 self.val = self.val.take().map(|i| i - 1);
5472 false
5473 }
5474 }
5475 }
5476
5477 #[derive(Debug, Default)]
5478 struct StoppingVisitor {
5480 inner: OkVisitor,
5481 return_false_from_pre_in: OptionalCounter,
5483 return_false_from_post_in: OptionalCounter,
5485 }
5486
5487 impl<'n> TreeNodeVisitor<'n> for StoppingVisitor {
5488 type Node = LogicalPlan;
5489
5490 fn f_down(&mut self, plan: &'n LogicalPlan) -> Result<TreeNodeRecursion> {
5491 if self.return_false_from_pre_in.dec() {
5492 return Ok(TreeNodeRecursion::Stop);
5493 }
5494 self.inner.f_down(plan)?;
5495
5496 Ok(TreeNodeRecursion::Continue)
5497 }
5498
5499 fn f_up(&mut self, plan: &'n LogicalPlan) -> Result<TreeNodeRecursion> {
5500 if self.return_false_from_post_in.dec() {
5501 return Ok(TreeNodeRecursion::Stop);
5502 }
5503
5504 self.inner.f_up(plan)
5505 }
5506 }
5507
5508 #[test]
5510 fn early_stopping_pre_visit() {
5511 let mut visitor = StoppingVisitor {
5512 return_false_from_pre_in: OptionalCounter::new(2),
5513 ..Default::default()
5514 };
5515 let plan = test_plan();
5516 let res = plan.visit_with_subqueries(&mut visitor);
5517 assert!(res.is_ok());
5518
5519 assert_debug_snapshot!(
5520 visitor.inner.strings,
5521 @r#"
5522 [
5523 "pre_visit Projection",
5524 "pre_visit Filter",
5525 ]
5526 "#
5527 );
5528 }
5529
5530 #[test]
5531 fn early_stopping_post_visit() {
5532 let mut visitor = StoppingVisitor {
5533 return_false_from_post_in: OptionalCounter::new(1),
5534 ..Default::default()
5535 };
5536 let plan = test_plan();
5537 let res = plan.visit_with_subqueries(&mut visitor);
5538 assert!(res.is_ok());
5539
5540 assert_debug_snapshot!(
5541 visitor.inner.strings,
5542 @r#"
5543 [
5544 "pre_visit Projection",
5545 "pre_visit Filter",
5546 "pre_visit TableScan",
5547 "post_visit TableScan",
5548 ]
5549 "#
5550 );
5551 }
5552
5553 #[derive(Debug, Default)]
5554 struct ErrorVisitor {
5556 inner: OkVisitor,
5557 return_error_from_pre_in: OptionalCounter,
5559 return_error_from_post_in: OptionalCounter,
5561 }
5562
5563 impl<'n> TreeNodeVisitor<'n> for ErrorVisitor {
5564 type Node = LogicalPlan;
5565
5566 fn f_down(&mut self, plan: &'n LogicalPlan) -> Result<TreeNodeRecursion> {
5567 if self.return_error_from_pre_in.dec() {
5568 return not_impl_err!("Error in pre_visit");
5569 }
5570
5571 self.inner.f_down(plan)
5572 }
5573
5574 fn f_up(&mut self, plan: &'n LogicalPlan) -> Result<TreeNodeRecursion> {
5575 if self.return_error_from_post_in.dec() {
5576 return not_impl_err!("Error in post_visit");
5577 }
5578
5579 self.inner.f_up(plan)
5580 }
5581 }
5582
5583 #[test]
5584 fn error_pre_visit() {
5585 let mut visitor = ErrorVisitor {
5586 return_error_from_pre_in: OptionalCounter::new(2),
5587 ..Default::default()
5588 };
5589 let plan = test_plan();
5590 let res = plan.visit_with_subqueries(&mut visitor).unwrap_err();
5591 assert_snapshot!(
5592 res.strip_backtrace(),
5593 @"This feature is not implemented: Error in pre_visit"
5594 );
5595 assert_debug_snapshot!(
5596 visitor.inner.strings,
5597 @r#"
5598 [
5599 "pre_visit Projection",
5600 "pre_visit Filter",
5601 ]
5602 "#
5603 );
5604 }
5605
5606 #[test]
5607 fn error_post_visit() {
5608 let mut visitor = ErrorVisitor {
5609 return_error_from_post_in: OptionalCounter::new(1),
5610 ..Default::default()
5611 };
5612 let plan = test_plan();
5613 let res = plan.visit_with_subqueries(&mut visitor).unwrap_err();
5614 assert_snapshot!(
5615 res.strip_backtrace(),
5616 @"This feature is not implemented: Error in post_visit"
5617 );
5618 assert_debug_snapshot!(
5619 visitor.inner.strings,
5620 @r#"
5621 [
5622 "pre_visit Projection",
5623 "pre_visit Filter",
5624 "pre_visit TableScan",
5625 "post_visit TableScan",
5626 ]
5627 "#
5628 );
5629 }
5630
5631 #[test]
5632 fn test_partial_eq_hash_and_partial_ord() {
5633 let empty_values = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
5634 produce_one_row: true,
5635 schema: Arc::new(DFSchema::empty()),
5636 }));
5637
5638 let count_window_function = |schema| {
5639 Window::try_new_with_schema(
5640 vec![Expr::WindowFunction(Box::new(WindowFunction::new(
5641 WindowFunctionDefinition::AggregateUDF(count_udaf()),
5642 vec![],
5643 )))],
5644 Arc::clone(&empty_values),
5645 Arc::new(schema),
5646 )
5647 .unwrap()
5648 };
5649
5650 let schema_without_metadata = || {
5651 DFSchema::from_unqualified_fields(
5652 vec![Field::new("count", DataType::Int64, false)].into(),
5653 HashMap::new(),
5654 )
5655 .unwrap()
5656 };
5657
5658 let schema_with_metadata = || {
5659 DFSchema::from_unqualified_fields(
5660 vec![Field::new("count", DataType::Int64, false)].into(),
5661 [("key".to_string(), "value".to_string())].into(),
5662 )
5663 .unwrap()
5664 };
5665
5666 let f = count_window_function(schema_without_metadata());
5668
5669 let f2 = count_window_function(schema_without_metadata());
5671 assert_eq!(f, f2);
5672 assert_eq!(hash(&f), hash(&f2));
5673 assert_eq!(f.partial_cmp(&f2), Some(Ordering::Equal));
5674
5675 let o = count_window_function(schema_with_metadata());
5677 assert_ne!(f, o);
5678 assert_ne!(hash(&f), hash(&o)); assert_eq!(f.partial_cmp(&o), None);
5680 }
5681
5682 fn hash<T: Hash>(value: &T) -> u64 {
5683 let hasher = &mut DefaultHasher::new();
5684 value.hash(hasher);
5685 hasher.finish()
5686 }
5687
5688 #[test]
5689 fn projection_expr_schema_mismatch() -> Result<()> {
5690 let empty_schema = Arc::new(DFSchema::empty());
5691 let p = Projection::try_new_with_schema(
5692 vec![col("a")],
5693 Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
5694 produce_one_row: false,
5695 schema: Arc::clone(&empty_schema),
5696 })),
5697 empty_schema,
5698 );
5699 assert_snapshot!(p.unwrap_err().strip_backtrace(), @"Error during planning: Projection has mismatch between number of expressions (1) and number of fields in schema (0)");
5700 Ok(())
5701 }
5702
5703 fn test_plan() -> LogicalPlan {
5704 let schema = Schema::new(vec![
5705 Field::new("id", DataType::Int32, false),
5706 Field::new("state", DataType::Utf8, false),
5707 ]);
5708
5709 table_scan(TableReference::none(), &schema, Some(vec![0, 1]))
5710 .unwrap()
5711 .filter(col("state").eq(lit("CO")))
5712 .unwrap()
5713 .project(vec![col("id")])
5714 .unwrap()
5715 .build()
5716 .unwrap()
5717 }
5718
5719 #[test]
5720 fn test_replace_invalid_placeholder() {
5721 let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
5723
5724 let plan = table_scan(TableReference::none(), &schema, None)
5725 .unwrap()
5726 .filter(col("id").eq(placeholder("")))
5727 .unwrap()
5728 .build()
5729 .unwrap();
5730
5731 let param_values = vec![ScalarValue::Int32(Some(42))];
5732 plan.replace_params_with_values(¶m_values.clone().into())
5733 .expect_err("unexpectedly succeeded to replace an invalid placeholder");
5734
5735 let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
5737
5738 let plan = table_scan(TableReference::none(), &schema, None)
5739 .unwrap()
5740 .filter(col("id").eq(placeholder("$0")))
5741 .unwrap()
5742 .build()
5743 .unwrap();
5744
5745 plan.replace_params_with_values(¶m_values.clone().into())
5746 .expect_err("unexpectedly succeeded to replace an invalid placeholder");
5747
5748 let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
5750
5751 let plan = table_scan(TableReference::none(), &schema, None)
5752 .unwrap()
5753 .filter(col("id").eq(placeholder("$00")))
5754 .unwrap()
5755 .build()
5756 .unwrap();
5757
5758 plan.replace_params_with_values(¶m_values.into())
5759 .expect_err("unexpectedly succeeded to replace an invalid placeholder");
5760 }
5761
5762 #[test]
5763 fn test_replace_placeholder_mismatched_metadata() {
5764 let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
5765
5766 let plan = table_scan(TableReference::none(), &schema, None)
5768 .unwrap()
5769 .filter(col("id").eq(placeholder("$1")))
5770 .unwrap()
5771 .build()
5772 .unwrap();
5773 let prepared_builder = LogicalPlanBuilder::new(plan)
5774 .prepare(
5775 "".to_string(),
5776 vec![Field::new("", DataType::Int32, true).into()],
5777 )
5778 .unwrap();
5779
5780 let mut scalar_meta = HashMap::new();
5782 scalar_meta.insert("some_key".to_string(), "some_value".to_string());
5783 let param_values = ParamValues::List(vec![ScalarAndMetadata::new(
5784 ScalarValue::Int32(Some(42)),
5785 Some(scalar_meta.into()),
5786 )]);
5787 prepared_builder
5788 .plan()
5789 .clone()
5790 .with_param_values(param_values)
5791 .expect_err("prepared field metadata mismatch unexpectedly succeeded");
5792 }
5793
5794 #[test]
5795 fn test_replace_placeholder_empty_relation_valid_schema() {
5796 let plan = LogicalPlanBuilder::empty(false)
5798 .project(vec![
5799 SelectExpr::from(placeholder("$1")),
5800 SelectExpr::from(placeholder("$2")),
5801 ])
5802 .unwrap()
5803 .build()
5804 .unwrap();
5805
5806 assert_snapshot!(plan.display_indent_schema(), @r"
5808 Projection: $1, $2 [$1:Null;N, $2:Null;N]
5809 EmptyRelation: rows=0 []
5810 ");
5811
5812 let plan = plan
5813 .with_param_values(vec![ScalarValue::from(1i32), ScalarValue::from("s")])
5814 .unwrap();
5815
5816 assert_snapshot!(plan.display_indent_schema(), @r#"
5818 Projection: Int32(1) AS $1, Utf8("s") AS $2 [$1:Int32, $2:Utf8]
5819 EmptyRelation: rows=0 []
5820 "#);
5821 }
5822
5823 #[test]
5824 fn test_nullable_schema_after_grouping_set() {
5825 let schema = Schema::new(vec![
5826 Field::new("foo", DataType::Int32, false),
5827 Field::new("bar", DataType::Int32, false),
5828 ]);
5829
5830 let plan = table_scan(TableReference::none(), &schema, None)
5831 .unwrap()
5832 .aggregate(
5833 vec![Expr::GroupingSet(GroupingSet::GroupingSets(vec![
5834 vec![col("foo")],
5835 vec![col("bar")],
5836 ]))],
5837 vec![count(lit(true))],
5838 )
5839 .unwrap()
5840 .build()
5841 .unwrap();
5842
5843 let output_schema = plan.schema();
5844
5845 assert!(
5846 output_schema
5847 .field_with_name(None, "foo")
5848 .unwrap()
5849 .is_nullable(),
5850 );
5851 assert!(
5852 output_schema
5853 .field_with_name(None, "bar")
5854 .unwrap()
5855 .is_nullable()
5856 );
5857 }
5858
5859 #[test]
5860 fn grouping_id_type_accounts_for_duplicate_ordinal_bits() {
5861 assert_eq!(Aggregate::grouping_id_type(8, 0), DataType::UInt8);
5864 assert_eq!(Aggregate::grouping_id_type(8, 1), DataType::UInt16);
5865 }
5866
5867 #[test]
5868 fn test_filter_is_scalar() {
5869 let schema =
5871 Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
5872
5873 let source = Arc::new(LogicalTableSource::new(schema));
5874 let schema = Arc::new(
5875 DFSchema::try_from_qualified_schema(
5876 TableReference::bare("tab"),
5877 &source.schema(),
5878 )
5879 .unwrap(),
5880 );
5881 let scan = Arc::new(LogicalPlan::TableScan(TableScan {
5882 table_name: TableReference::bare("tab"),
5883 source: Arc::clone(&source) as Arc<dyn TableSource>,
5884 projection: None,
5885 projected_schema: Arc::clone(&schema),
5886 filters: vec![],
5887 fetch: None,
5888 statistics_requests: BTreeSet::new(),
5889 }));
5890 let col = schema.field_names()[0].clone();
5891
5892 let filter = Filter::try_new(
5893 Expr::Column(col.into()).eq(Expr::Literal(ScalarValue::Int32(Some(1)), None)),
5894 scan,
5895 )
5896 .unwrap();
5897 assert!(!filter.is_scalar());
5898 let unique_schema = Arc::new(
5899 schema
5900 .as_ref()
5901 .clone()
5902 .with_functional_dependencies(
5903 FunctionalDependencies::new_from_constraints(
5904 Some(&Constraints::new_unverified(vec![Constraint::Unique(
5905 vec![0],
5906 )])),
5907 1,
5908 ),
5909 )
5910 .unwrap(),
5911 );
5912 let scan = Arc::new(LogicalPlan::TableScan(TableScan {
5913 table_name: TableReference::bare("tab"),
5914 source,
5915 projection: None,
5916 projected_schema: Arc::clone(&unique_schema),
5917 filters: vec![],
5918 fetch: None,
5919 statistics_requests: BTreeSet::new(),
5920 }));
5921 let col = schema.field_names()[0].clone();
5922
5923 let filter =
5924 Filter::try_new(Expr::Column(col.into()).eq(lit(1i32)), scan).unwrap();
5925 assert!(filter.is_scalar());
5926 }
5927
5928 #[test]
5929 fn test_transform_explain() {
5930 let schema = Schema::new(vec![
5931 Field::new("foo", DataType::Int32, false),
5932 Field::new("bar", DataType::Int32, false),
5933 ]);
5934
5935 let plan = table_scan(TableReference::none(), &schema, None)
5936 .unwrap()
5937 .explain(false, false)
5938 .unwrap()
5939 .build()
5940 .unwrap();
5941
5942 let external_filter = col("foo").eq(lit(true));
5943
5944 let plan = plan
5947 .transform(|plan| match plan {
5948 LogicalPlan::TableScan(table) => {
5949 let filter = Filter::try_new(
5950 external_filter.clone(),
5951 Arc::new(LogicalPlan::TableScan(table)),
5952 )
5953 .unwrap();
5954 Ok(Transformed::yes(LogicalPlan::Filter(filter)))
5955 }
5956 x => Ok(Transformed::no(x)),
5957 })
5958 .data()
5959 .unwrap();
5960
5961 let actual = format!("{}", plan.display_indent());
5962 assert_snapshot!(actual, @r"
5963 Explain
5964 Filter: foo = Boolean(true)
5965 TableScan: ?table?
5966 ")
5967 }
5968
5969 #[test]
5970 fn test_plan_partial_ord() {
5971 let empty_relation = LogicalPlan::EmptyRelation(EmptyRelation {
5972 produce_one_row: false,
5973 schema: Arc::new(DFSchema::empty()),
5974 });
5975
5976 let describe_table = LogicalPlan::DescribeTable(DescribeTable {
5977 schema: Arc::new(Schema::new(vec![Field::new(
5978 "foo",
5979 DataType::Int32,
5980 false,
5981 )])),
5982 output_schema: DFSchemaRef::new(DFSchema::empty()),
5983 });
5984
5985 let describe_table_clone = LogicalPlan::DescribeTable(DescribeTable {
5986 schema: Arc::new(Schema::new(vec![Field::new(
5987 "foo",
5988 DataType::Int32,
5989 false,
5990 )])),
5991 output_schema: DFSchemaRef::new(DFSchema::empty()),
5992 });
5993
5994 assert_eq!(
5995 empty_relation.partial_cmp(&describe_table),
5996 Some(Ordering::Less)
5997 );
5998 assert_eq!(
5999 describe_table.partial_cmp(&empty_relation),
6000 Some(Ordering::Greater)
6001 );
6002 assert_eq!(describe_table.partial_cmp(&describe_table_clone), None);
6003 }
6004
6005 #[test]
6006 fn test_limit_with_new_children() {
6007 let input = Arc::new(LogicalPlan::Values(Values {
6008 schema: Arc::new(DFSchema::empty()),
6009 values: vec![vec![]],
6010 }));
6011 let cases = [
6012 LogicalPlan::Limit(Limit {
6013 skip: None,
6014 fetch: None,
6015 input: Arc::clone(&input),
6016 }),
6017 LogicalPlan::Limit(Limit {
6018 skip: None,
6019 fetch: Some(Box::new(Expr::Literal(
6020 ScalarValue::new_ten(&DataType::UInt32).unwrap(),
6021 None,
6022 ))),
6023 input: Arc::clone(&input),
6024 }),
6025 LogicalPlan::Limit(Limit {
6026 skip: Some(Box::new(Expr::Literal(
6027 ScalarValue::new_ten(&DataType::UInt32).unwrap(),
6028 None,
6029 ))),
6030 fetch: None,
6031 input: Arc::clone(&input),
6032 }),
6033 LogicalPlan::Limit(Limit {
6034 skip: Some(Box::new(Expr::Literal(
6035 ScalarValue::new_one(&DataType::UInt32).unwrap(),
6036 None,
6037 ))),
6038 fetch: Some(Box::new(Expr::Literal(
6039 ScalarValue::new_ten(&DataType::UInt32).unwrap(),
6040 None,
6041 ))),
6042 input,
6043 }),
6044 ];
6045
6046 for limit in cases {
6047 let new_limit = limit
6048 .with_new_exprs(
6049 limit.expressions(),
6050 limit.inputs().into_iter().cloned().collect(),
6051 )
6052 .unwrap();
6053 assert_eq!(limit, new_limit);
6054 }
6055 }
6056
6057 #[test]
6058 fn test_with_subqueries_jump() {
6059 let subquery_schema =
6064 Schema::new(vec![Field::new("sub_id", DataType::Int32, false)]);
6065
6066 let subquery_plan =
6067 table_scan(TableReference::none(), &subquery_schema, Some(vec![0]))
6068 .unwrap()
6069 .filter(col("sub_id").eq(lit(0)))
6070 .unwrap()
6071 .build()
6072 .unwrap();
6073
6074 let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
6075
6076 let plan = table_scan(TableReference::none(), &schema, Some(vec![0]))
6077 .unwrap()
6078 .filter(col("id").eq(lit(0)))
6079 .unwrap()
6080 .project(vec![col("id"), scalar_subquery(Arc::new(subquery_plan))])
6081 .unwrap()
6082 .build()
6083 .unwrap();
6084
6085 let mut filter_found = false;
6086 plan.apply_with_subqueries(|plan| {
6087 match plan {
6088 LogicalPlan::Projection(..) => return Ok(TreeNodeRecursion::Jump),
6089 LogicalPlan::Filter(..) => filter_found = true,
6090 _ => {}
6091 }
6092 Ok(TreeNodeRecursion::Continue)
6093 })
6094 .unwrap();
6095 assert!(!filter_found);
6096
6097 struct ProjectJumpVisitor {
6098 filter_found: bool,
6099 }
6100
6101 impl ProjectJumpVisitor {
6102 fn new() -> Self {
6103 Self {
6104 filter_found: false,
6105 }
6106 }
6107 }
6108
6109 impl<'n> TreeNodeVisitor<'n> for ProjectJumpVisitor {
6110 type Node = LogicalPlan;
6111
6112 fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
6113 match node {
6114 LogicalPlan::Projection(..) => return Ok(TreeNodeRecursion::Jump),
6115 LogicalPlan::Filter(..) => self.filter_found = true,
6116 _ => {}
6117 }
6118 Ok(TreeNodeRecursion::Continue)
6119 }
6120 }
6121
6122 let mut visitor = ProjectJumpVisitor::new();
6123 plan.visit_with_subqueries(&mut visitor).unwrap();
6124 assert!(!visitor.filter_found);
6125
6126 let mut filter_found = false;
6127 plan.clone()
6128 .transform_down_with_subqueries(|plan| {
6129 match plan {
6130 LogicalPlan::Projection(..) => {
6131 return Ok(Transformed::new(
6132 plan,
6133 false,
6134 TreeNodeRecursion::Jump,
6135 ));
6136 }
6137 LogicalPlan::Filter(..) => filter_found = true,
6138 _ => {}
6139 }
6140 Ok(Transformed::no(plan))
6141 })
6142 .unwrap();
6143 assert!(!filter_found);
6144
6145 let mut filter_found = false;
6146 plan.clone()
6147 .transform_down_up_with_subqueries(
6148 |plan| {
6149 match plan {
6150 LogicalPlan::Projection(..) => {
6151 return Ok(Transformed::new(
6152 plan,
6153 false,
6154 TreeNodeRecursion::Jump,
6155 ));
6156 }
6157 LogicalPlan::Filter(..) => filter_found = true,
6158 _ => {}
6159 }
6160 Ok(Transformed::no(plan))
6161 },
6162 |plan| Ok(Transformed::no(plan)),
6163 )
6164 .unwrap();
6165 assert!(!filter_found);
6166
6167 struct ProjectJumpRewriter {
6168 filter_found: bool,
6169 }
6170
6171 impl ProjectJumpRewriter {
6172 fn new() -> Self {
6173 Self {
6174 filter_found: false,
6175 }
6176 }
6177 }
6178
6179 impl TreeNodeRewriter for ProjectJumpRewriter {
6180 type Node = LogicalPlan;
6181
6182 fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
6183 match node {
6184 LogicalPlan::Projection(..) => {
6185 return Ok(Transformed::new(
6186 node,
6187 false,
6188 TreeNodeRecursion::Jump,
6189 ));
6190 }
6191 LogicalPlan::Filter(..) => self.filter_found = true,
6192 _ => {}
6193 }
6194 Ok(Transformed::no(node))
6195 }
6196 }
6197
6198 let mut rewriter = ProjectJumpRewriter::new();
6199 plan.rewrite_with_subqueries(&mut rewriter).unwrap();
6200 assert!(!rewriter.filter_found);
6201 }
6202
6203 #[test]
6204 fn test_with_unresolved_placeholders() {
6205 let field_name = "id";
6206 let placeholder_value = "$1";
6207 let schema = Schema::new(vec![Field::new(field_name, DataType::Int32, false)]);
6208
6209 let plan = table_scan(TableReference::none(), &schema, None)
6210 .unwrap()
6211 .filter(col(field_name).eq(placeholder(placeholder_value)))
6212 .unwrap()
6213 .build()
6214 .unwrap();
6215
6216 let params = plan.get_parameter_fields().unwrap();
6218 assert_eq!(params.len(), 1);
6219
6220 let parameter_type = params.clone().get(placeholder_value).unwrap().clone();
6221 assert_eq!(parameter_type, None);
6222 }
6223
6224 #[test]
6225 fn test_join_with_new_exprs() -> Result<()> {
6226 fn create_test_join(
6227 on: Vec<(Expr, Expr)>,
6228 filter: Option<Expr>,
6229 ) -> Result<LogicalPlan> {
6230 let schema = Schema::new(vec![
6231 Field::new("a", DataType::Int32, false),
6232 Field::new("b", DataType::Int32, false),
6233 ]);
6234
6235 let left_schema = DFSchema::try_from_qualified_schema("t1", &schema)?;
6236 let right_schema = DFSchema::try_from_qualified_schema("t2", &schema)?;
6237
6238 Ok(LogicalPlan::Join(Join {
6239 left: Arc::new(
6240 table_scan(Some("t1"), left_schema.as_arrow(), None)?.build()?,
6241 ),
6242 right: Arc::new(
6243 table_scan(Some("t2"), right_schema.as_arrow(), None)?.build()?,
6244 ),
6245 on,
6246 filter,
6247 join_type: JoinType::Inner,
6248 join_constraint: JoinConstraint::On,
6249 schema: Arc::new(left_schema.join(&right_schema)?),
6250 null_equality: NullEquality::NullEqualsNothing,
6251 null_aware: false,
6252 }))
6253 }
6254
6255 {
6256 let join = create_test_join(vec![(col("t1.a"), (col("t2.a")))], None)?;
6257 let LogicalPlan::Join(join) = join.with_new_exprs(
6258 join.expressions(),
6259 join.inputs().into_iter().cloned().collect(),
6260 )?
6261 else {
6262 unreachable!()
6263 };
6264 assert_eq!(join.on, vec![(col("t1.a"), (col("t2.a")))]);
6265 assert_eq!(join.filter, None);
6266 }
6267
6268 {
6269 let join = create_test_join(vec![], Some(col("t1.a").gt(col("t2.a"))))?;
6270 let LogicalPlan::Join(join) = join.with_new_exprs(
6271 join.expressions(),
6272 join.inputs().into_iter().cloned().collect(),
6273 )?
6274 else {
6275 unreachable!()
6276 };
6277 assert_eq!(join.on, vec![]);
6278 assert_eq!(join.filter, Some(col("t1.a").gt(col("t2.a"))));
6279 }
6280
6281 {
6282 let join = create_test_join(
6283 vec![(col("t1.a"), (col("t2.a")))],
6284 Some(col("t1.b").gt(col("t2.b"))),
6285 )?;
6286 let LogicalPlan::Join(join) = join.with_new_exprs(
6287 join.expressions(),
6288 join.inputs().into_iter().cloned().collect(),
6289 )?
6290 else {
6291 unreachable!()
6292 };
6293 assert_eq!(join.on, vec![(col("t1.a"), (col("t2.a")))]);
6294 assert_eq!(join.filter, Some(col("t1.b").gt(col("t2.b"))));
6295 }
6296
6297 {
6298 let join = create_test_join(
6299 vec![(col("t1.a"), (col("t2.a"))), (col("t1.b"), (col("t2.b")))],
6300 None,
6301 )?;
6302 let LogicalPlan::Join(join) = join.with_new_exprs(
6303 vec![
6304 binary_expr(col("t1.a"), Operator::Plus, lit(1)),
6305 binary_expr(col("t2.a"), Operator::Plus, lit(2)),
6306 col("t1.b"),
6307 col("t2.b"),
6308 lit(true),
6309 ],
6310 join.inputs().into_iter().cloned().collect(),
6311 )?
6312 else {
6313 unreachable!()
6314 };
6315 assert_eq!(
6316 join.on,
6317 vec![
6318 (
6319 binary_expr(col("t1.a"), Operator::Plus, lit(1)),
6320 binary_expr(col("t2.a"), Operator::Plus, lit(2))
6321 ),
6322 (col("t1.b"), (col("t2.b")))
6323 ]
6324 );
6325 assert_eq!(join.filter, Some(lit(true)));
6326 }
6327
6328 Ok(())
6329 }
6330
6331 #[test]
6332 fn test_join_try_new() -> Result<()> {
6333 let schema = Schema::new(vec![
6334 Field::new("a", DataType::Int32, false),
6335 Field::new("b", DataType::Int32, false),
6336 ]);
6337
6338 let left_scan = table_scan(Some("t1"), &schema, None)?.build()?;
6339
6340 let right_scan = table_scan(Some("t2"), &schema, None)?.build()?;
6341
6342 let join_types = vec![
6343 JoinType::Inner,
6344 JoinType::Left,
6345 JoinType::Right,
6346 JoinType::Full,
6347 JoinType::LeftSemi,
6348 JoinType::LeftAnti,
6349 JoinType::RightSemi,
6350 JoinType::RightAnti,
6351 JoinType::LeftMark,
6352 ];
6353
6354 for join_type in join_types {
6355 let join = Join::try_new(
6356 Arc::new(left_scan.clone()),
6357 Arc::new(right_scan.clone()),
6358 vec![(col("t1.a"), col("t2.a"))],
6359 Some(col("t1.b").gt(col("t2.b"))),
6360 join_type,
6361 JoinConstraint::On,
6362 NullEquality::NullEqualsNothing,
6363 false,
6364 )?;
6365
6366 match join_type {
6367 JoinType::LeftSemi | JoinType::LeftAnti => {
6368 assert_eq!(join.schema.fields().len(), 2);
6369
6370 let fields = join.schema.fields();
6371 assert_eq!(
6372 fields[0].name(),
6373 "a",
6374 "First field should be 'a' from left table"
6375 );
6376 assert_eq!(
6377 fields[1].name(),
6378 "b",
6379 "Second field should be 'b' from left table"
6380 );
6381 }
6382 JoinType::RightSemi | JoinType::RightAnti => {
6383 assert_eq!(join.schema.fields().len(), 2);
6384
6385 let fields = join.schema.fields();
6386 assert_eq!(
6387 fields[0].name(),
6388 "a",
6389 "First field should be 'a' from right table"
6390 );
6391 assert_eq!(
6392 fields[1].name(),
6393 "b",
6394 "Second field should be 'b' from right table"
6395 );
6396 }
6397 JoinType::LeftMark => {
6398 assert_eq!(join.schema.fields().len(), 3);
6399
6400 let fields = join.schema.fields();
6401 assert_eq!(
6402 fields[0].name(),
6403 "a",
6404 "First field should be 'a' from left table"
6405 );
6406 assert_eq!(
6407 fields[1].name(),
6408 "b",
6409 "Second field should be 'b' from left table"
6410 );
6411 assert_eq!(
6412 fields[2].name(),
6413 "mark",
6414 "Third field should be the mark column"
6415 );
6416
6417 assert!(!fields[0].is_nullable());
6418 assert!(!fields[1].is_nullable());
6419 assert!(!fields[2].is_nullable());
6420 }
6421 _ => {
6422 assert_eq!(join.schema.fields().len(), 4);
6423
6424 let fields = join.schema.fields();
6425 assert_eq!(
6426 fields[0].name(),
6427 "a",
6428 "First field should be 'a' from left table"
6429 );
6430 assert_eq!(
6431 fields[1].name(),
6432 "b",
6433 "Second field should be 'b' from left table"
6434 );
6435 assert_eq!(
6436 fields[2].name(),
6437 "a",
6438 "Third field should be 'a' from right table"
6439 );
6440 assert_eq!(
6441 fields[3].name(),
6442 "b",
6443 "Fourth field should be 'b' from right table"
6444 );
6445
6446 if join_type == JoinType::Left {
6447 assert!(!fields[0].is_nullable());
6449 assert!(!fields[1].is_nullable());
6450 assert!(fields[2].is_nullable());
6452 assert!(fields[3].is_nullable());
6453 } else if join_type == JoinType::Right {
6454 assert!(fields[0].is_nullable());
6456 assert!(fields[1].is_nullable());
6457 assert!(!fields[2].is_nullable());
6459 assert!(!fields[3].is_nullable());
6460 } else if join_type == JoinType::Full {
6461 assert!(fields[0].is_nullable());
6462 assert!(fields[1].is_nullable());
6463 assert!(fields[2].is_nullable());
6464 assert!(fields[3].is_nullable());
6465 }
6466 }
6467 }
6468
6469 assert_eq!(join.on, vec![(col("t1.a"), col("t2.a"))]);
6470 assert_eq!(join.filter, Some(col("t1.b").gt(col("t2.b"))));
6471 assert_eq!(join.join_type, join_type);
6472 assert_eq!(join.join_constraint, JoinConstraint::On);
6473 assert_eq!(join.null_equality, NullEquality::NullEqualsNothing);
6474 }
6475
6476 Ok(())
6477 }
6478
6479 #[test]
6480 fn test_join_try_new_with_using_constraint_and_overlapping_columns() -> Result<()> {
6481 let left_schema = Schema::new(vec![
6482 Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Int32, false), ]);
6486
6487 let right_schema = Schema::new(vec![
6488 Field::new("id", DataType::Int32, false), Field::new("category", DataType::Utf8, false), Field::new("value", DataType::Float64, true), ]);
6492
6493 let left_plan = table_scan(Some("t1"), &left_schema, None)?.build()?;
6494
6495 let right_plan = table_scan(Some("t2"), &right_schema, None)?.build()?;
6496
6497 {
6499 let join = Join::try_new(
6502 Arc::new(left_plan.clone()),
6503 Arc::new(right_plan.clone()),
6504 vec![(col("t1.id"), col("t2.id"))],
6505 None,
6506 JoinType::Inner,
6507 JoinConstraint::Using,
6508 NullEquality::NullEqualsNothing,
6509 false,
6510 )?;
6511
6512 let fields = join.schema.fields();
6513
6514 assert_eq!(fields.len(), 6);
6515
6516 assert_eq!(
6517 fields[0].name(),
6518 "id",
6519 "First field should be 'id' from left table"
6520 );
6521 assert_eq!(
6522 fields[1].name(),
6523 "name",
6524 "Second field should be 'name' from left table"
6525 );
6526 assert_eq!(
6527 fields[2].name(),
6528 "value",
6529 "Third field should be 'value' from left table"
6530 );
6531 assert_eq!(
6532 fields[3].name(),
6533 "id",
6534 "Fourth field should be 'id' from right table"
6535 );
6536 assert_eq!(
6537 fields[4].name(),
6538 "category",
6539 "Fifth field should be 'category' from right table"
6540 );
6541 assert_eq!(
6542 fields[5].name(),
6543 "value",
6544 "Sixth field should be 'value' from right table"
6545 );
6546
6547 assert_eq!(join.join_constraint, JoinConstraint::Using);
6548 }
6549
6550 {
6552 let join = Join::try_new(
6554 Arc::new(left_plan.clone()),
6555 Arc::new(right_plan.clone()),
6556 vec![(col("t1.id"), col("t2.id"))], Some(col("t1.value").lt(col("t2.value"))), JoinType::Inner,
6559 JoinConstraint::On,
6560 NullEquality::NullEqualsNothing,
6561 false,
6562 )?;
6563
6564 let fields = join.schema.fields();
6565 assert_eq!(fields.len(), 6);
6566
6567 assert_eq!(
6568 fields[0].name(),
6569 "id",
6570 "First field should be 'id' from left table"
6571 );
6572 assert_eq!(
6573 fields[1].name(),
6574 "name",
6575 "Second field should be 'name' from left table"
6576 );
6577 assert_eq!(
6578 fields[2].name(),
6579 "value",
6580 "Third field should be 'value' from left table"
6581 );
6582 assert_eq!(
6583 fields[3].name(),
6584 "id",
6585 "Fourth field should be 'id' from right table"
6586 );
6587 assert_eq!(
6588 fields[4].name(),
6589 "category",
6590 "Fifth field should be 'category' from right table"
6591 );
6592 assert_eq!(
6593 fields[5].name(),
6594 "value",
6595 "Sixth field should be 'value' from right table"
6596 );
6597
6598 assert_eq!(join.filter, Some(col("t1.value").lt(col("t2.value"))));
6599 }
6600
6601 {
6603 let join = Join::try_new(
6604 Arc::new(left_plan.clone()),
6605 Arc::new(right_plan.clone()),
6606 vec![(col("t1.id"), col("t2.id"))],
6607 None,
6608 JoinType::Inner,
6609 JoinConstraint::On,
6610 NullEquality::NullEqualsNull,
6611 false,
6612 )?;
6613
6614 assert_eq!(join.null_equality, NullEquality::NullEqualsNull);
6615 }
6616
6617 Ok(())
6618 }
6619
6620 #[test]
6621 fn test_join_try_new_schema_validation() -> Result<()> {
6622 let left_schema = Schema::new(vec![
6623 Field::new("id", DataType::Int32, false),
6624 Field::new("name", DataType::Utf8, false),
6625 Field::new("value", DataType::Float64, true),
6626 ]);
6627
6628 let right_schema = Schema::new(vec![
6629 Field::new("id", DataType::Int32, false),
6630 Field::new("category", DataType::Utf8, true),
6631 Field::new("code", DataType::Int16, false),
6632 ]);
6633
6634 let left_plan = table_scan(Some("t1"), &left_schema, None)?.build()?;
6635
6636 let right_plan = table_scan(Some("t2"), &right_schema, None)?.build()?;
6637
6638 let join_types = vec![
6639 JoinType::Inner,
6640 JoinType::Left,
6641 JoinType::Right,
6642 JoinType::Full,
6643 ];
6644
6645 for join_type in join_types {
6646 let join = Join::try_new(
6647 Arc::new(left_plan.clone()),
6648 Arc::new(right_plan.clone()),
6649 vec![(col("t1.id"), col("t2.id"))],
6650 Some(col("t1.value").gt(lit(5.0))),
6651 join_type,
6652 JoinConstraint::On,
6653 NullEquality::NullEqualsNothing,
6654 false,
6655 )?;
6656
6657 let fields = join.schema.fields();
6658 assert_eq!(fields.len(), 6, "Expected 6 fields for {join_type} join");
6659
6660 for (i, field) in fields.iter().enumerate() {
6661 let expected_nullable = match (i, &join_type) {
6662 (0, JoinType::Right | JoinType::Full) => true, (1, JoinType::Right | JoinType::Full) => true, (2, _) => true, (3, JoinType::Left | JoinType::Full) => true, (4, _) => true, (5, JoinType::Left | JoinType::Full) => true, _ => false,
6673 };
6674
6675 assert_eq!(
6676 field.is_nullable(),
6677 expected_nullable,
6678 "Field {} ({}) nullability incorrect for {:?} join",
6679 i,
6680 field.name(),
6681 join_type
6682 );
6683 }
6684 }
6685
6686 let using_join = Join::try_new(
6687 Arc::new(left_plan.clone()),
6688 Arc::new(right_plan.clone()),
6689 vec![(col("t1.id"), col("t2.id"))],
6690 None,
6691 JoinType::Inner,
6692 JoinConstraint::Using,
6693 NullEquality::NullEqualsNothing,
6694 false,
6695 )?;
6696
6697 assert_eq!(
6698 using_join.schema.fields().len(),
6699 6,
6700 "USING join should have all fields"
6701 );
6702 assert_eq!(using_join.join_constraint, JoinConstraint::Using);
6703
6704 Ok(())
6705 }
6706
6707 #[test]
6708 fn test_unnest_with_new_exprs_accepts_expressions() -> Result<()> {
6709 use crate::LogicalPlanBuilder;
6710 use arrow::datatypes::{DataType, Field, Schema};
6711
6712 let schema = Schema::new(vec![
6713 Field::new("list_col", DataType::new_list(DataType::Int32, true), true),
6714 Field::new("other_col", DataType::Int32, true),
6715 ]);
6716 let plan = table_scan(Some("t"), &schema, None)?.build()?;
6717 let unnest_plan = LogicalPlanBuilder::from(plan)
6718 .unnest_column("list_col")?
6719 .build()?;
6720
6721 let exprs = unnest_plan.expressions();
6722 assert!(!exprs.is_empty(), "Unnest should expose exec_columns");
6723 assert_eq!(exprs.len(), 1);
6724 assert!(matches!(&exprs[0], Expr::Column(c) if c.name == "list_col"));
6725
6726 let inputs: Vec<LogicalPlan> =
6727 unnest_plan.inputs().into_iter().cloned().collect();
6728 let rebuilt = unnest_plan.with_new_exprs(exprs, inputs)?;
6729 assert_eq!(rebuilt.schema(), unnest_plan.schema());
6730
6731 Ok(())
6732 }
6733
6734 #[test]
6735 fn test_unnest_with_new_exprs_empty_preserves_columns() -> Result<()> {
6736 use crate::LogicalPlanBuilder;
6737 use arrow::datatypes::{DataType, Field, Schema};
6738
6739 let schema = Schema::new(vec![
6740 Field::new("list_col", DataType::new_list(DataType::Int32, true), true),
6741 Field::new("other_col", DataType::Int32, true),
6742 ]);
6743 let plan = table_scan(Some("t"), &schema, None)?.build()?;
6744 let unnest_plan = LogicalPlanBuilder::from(plan)
6745 .unnest_column("list_col")?
6746 .build()?;
6747
6748 let inputs: Vec<LogicalPlan> =
6749 unnest_plan.inputs().into_iter().cloned().collect();
6750 let rebuilt = unnest_plan.with_new_exprs(vec![], inputs)?;
6751 assert_eq!(rebuilt.schema(), unnest_plan.schema());
6752
6753 Ok(())
6754 }
6755}