1use std::borrow::Cow;
21use std::cmp::Ordering;
22use std::collections::{HashMap, HashSet};
23use std::iter::once;
24use std::sync::Arc;
25
26use crate::dml::CopyTo;
27use crate::expr::{Alias, PlannedReplaceSelectItem, Sort as SortExpr};
28use crate::expr_rewriter::{
29 coerce_plan_expr_for_schema, normalize_col,
30 normalize_col_with_schemas_and_ambiguity_check, normalize_cols, normalize_sorts,
31 rewrite_sort_cols_by_aggs,
32};
33use crate::logical_plan::{
34 Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
35 JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
36 Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest,
37 Values, Window,
38};
39use crate::select_expr::SelectExpr;
40use crate::utils::{
41 can_hash, columnize_expr, compare_sort_expr, expand_qualified_wildcard,
42 expand_wildcard, expr_to_columns, find_valid_equijoin_key_pair,
43 group_window_expr_by_sort_keys,
44};
45use crate::{
46 DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, RecursiveQuery,
47 Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, binary_expr, lit,
48};
49
50use super::dml::InsertOp;
51use arrow::compute::can_cast_types;
52use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema, SchemaRef};
53use datafusion_common::display::ToStringifiedPlan;
54use datafusion_common::file_options::file_type::FileType;
55use datafusion_common::metadata::FieldMetadata;
56use datafusion_common::{
57 Column, Constraints, DFSchema, DFSchemaRef, NullEquality, Result, ScalarValue,
58 TableReference, ToDFSchema, UnnestOptions, exec_err,
59 get_target_functional_dependencies, internal_datafusion_err, plan_datafusion_err,
60 plan_err,
61};
62use datafusion_expr_common::type_coercion::binary::type_union_resolution;
63
64use indexmap::IndexSet;
65
66pub const UNNAMED_TABLE: &str = "?table?";
68
69#[derive(Default, Debug, Clone)]
71pub struct LogicalPlanBuilderOptions {
72 add_implicit_group_by_exprs: bool,
75}
76
77impl LogicalPlanBuilderOptions {
78 pub fn new() -> Self {
79 Default::default()
80 }
81
82 pub fn with_add_implicit_group_by_exprs(mut self, add: bool) -> Self {
84 self.add_implicit_group_by_exprs = add;
85 self
86 }
87}
88
89#[derive(Debug, Clone)]
127pub struct LogicalPlanBuilder {
128 plan: Arc<LogicalPlan>,
129 options: LogicalPlanBuilderOptions,
130}
131
132impl LogicalPlanBuilder {
133 pub fn new(plan: LogicalPlan) -> Self {
135 Self {
136 plan: Arc::new(plan),
137 options: LogicalPlanBuilderOptions::default(),
138 }
139 }
140
141 pub fn new_from_arc(plan: Arc<LogicalPlan>) -> Self {
143 Self {
144 plan,
145 options: LogicalPlanBuilderOptions::default(),
146 }
147 }
148
149 pub fn with_options(mut self, options: LogicalPlanBuilderOptions) -> Self {
150 self.options = options;
151 self
152 }
153
154 pub fn schema(&self) -> &DFSchemaRef {
156 self.plan.schema()
157 }
158
159 pub fn plan(&self) -> &LogicalPlan {
161 &self.plan
162 }
163
164 pub fn empty(produce_one_row: bool) -> Self {
168 Self::new(LogicalPlan::EmptyRelation(EmptyRelation {
169 produce_one_row,
170 schema: DFSchemaRef::new(DFSchema::empty()),
171 }))
172 }
173
174 pub fn to_recursive_query(
177 self,
178 name: String,
179 recursive_term: LogicalPlan,
180 is_distinct: bool,
181 ) -> Result<Self> {
182 let static_fields_len = self.plan.schema().fields().len();
184 let recursive_fields_len = recursive_term.schema().fields().len();
185 if static_fields_len != recursive_fields_len {
186 return plan_err!(
187 "Non-recursive term and recursive term must have the same number of columns ({} != {})",
188 static_fields_len,
189 recursive_fields_len
190 );
191 }
192 let coerced_recursive_term =
194 coerce_plan_expr_for_schema(recursive_term, self.plan.schema())?;
195 let recursive_query = RecursiveQuery::try_new(
196 name,
197 self.plan,
198 Arc::new(coerced_recursive_term),
199 is_distinct,
200 )?;
201 Ok(Self::from(LogicalPlan::RecursiveQuery(recursive_query)))
202 }
203
204 pub fn values(values: Vec<Vec<Expr>>) -> Result<Self> {
212 if values.is_empty() {
213 return plan_err!("Values list cannot be empty");
214 }
215 let n_cols = values[0].len();
216 if n_cols == 0 {
217 return plan_err!("Values list cannot be zero length");
218 }
219 for (i, row) in values.iter().enumerate() {
220 if row.len() != n_cols {
221 return plan_err!(
222 "Inconsistent data length across values list: got {} values in row {} but expected {}",
223 row.len(),
224 i,
225 n_cols
226 );
227 }
228 }
229
230 Self::infer_data(values)
232 }
233
234 pub fn values_with_schema(
244 values: Vec<Vec<Expr>>,
245 schema: &DFSchemaRef,
246 ) -> Result<Self> {
247 if values.is_empty() {
248 return plan_err!("Values list cannot be empty");
249 }
250 let n_cols = schema.fields().len();
251 if n_cols == 0 {
252 return plan_err!("Values list cannot be zero length");
253 }
254 for (i, row) in values.iter().enumerate() {
255 if row.len() != n_cols {
256 return plan_err!(
257 "Inconsistent data length across values list: got {} values in row {} but expected {}",
258 row.len(),
259 i,
260 n_cols
261 );
262 }
263 }
264
265 Self::infer_values_from_schema(values, schema)
267 }
268
269 fn infer_values_from_schema(
270 values: Vec<Vec<Expr>>,
271 schema: &DFSchema,
272 ) -> Result<Self> {
273 let n_cols = values[0].len();
274 let mut fields = ValuesFields::new();
275 for j in 0..n_cols {
276 let field_type = schema.field(j).data_type();
277 let field_nullable = schema.field(j).is_nullable();
278 for row in values.iter() {
279 let value = &row[j];
280 let data_type = value.get_type(schema)?;
281
282 if !data_type.equals_datatype(field_type)
283 && !can_cast_types(&data_type, field_type)
284 {
285 return exec_err!(
286 "Types don't match and no valid cast exists, received data of type {} for field of type {}",
287 data_type,
288 field_type
289 );
290 }
291 }
292 fields.push(field_type.to_owned(), field_nullable);
293 }
294
295 Self::infer_inner(values, fields, schema)
296 }
297
298 fn infer_data(values: Vec<Vec<Expr>>) -> Result<Self> {
299 let n_cols = values[0].len();
300 let schema = DFSchema::empty();
301 let mut fields = ValuesFields::new();
302
303 for j in 0..n_cols {
304 let mut common_type: Option<DataType> = None;
305 let mut common_metadata: Option<FieldMetadata> = None;
306 let mut nullable = false;
307 for (i, row) in values.iter().enumerate() {
308 let value = &row[j];
309 let metadata = value.metadata(&schema)?;
310 if let Some(ref cm) = common_metadata {
311 if &metadata != cm {
312 return plan_err!(
313 "Inconsistent metadata across values list at row {i} column {j}. Was {:?} but found {:?}",
314 cm,
315 metadata
316 );
317 }
318 } else {
319 common_metadata = Some(metadata.clone());
320 }
321 if !nullable && value.nullable(&schema)? {
322 nullable = true;
323 }
324 let data_type = value.get_type(&schema)?;
325 if data_type == DataType::Null {
326 continue;
327 }
328
329 if let Some(prev_type) = common_type {
330 let data_types = vec![prev_type.clone(), data_type.clone()];
333 let Some(new_type) = type_union_resolution(&data_types) else {
334 return plan_err!(
335 "Inconsistent data type across values list at row {i} column {j}. Was {prev_type} but found {data_type}"
336 );
337 };
338 common_type = Some(new_type);
339 } else {
340 common_type = Some(data_type);
341 }
342 }
343 let (data_type, nullable) = match common_type {
346 Some(t) => (t, nullable),
347 None => (DataType::Null, true),
348 };
349 fields.push_with_metadata(data_type, nullable, common_metadata);
350 }
351
352 Self::infer_inner(values, fields, &schema)
353 }
354
355 fn infer_inner(
356 mut values: Vec<Vec<Expr>>,
357 fields: ValuesFields,
358 schema: &DFSchema,
359 ) -> Result<Self> {
360 let fields = fields.into_fields();
361 for row in &mut values {
363 for (j, field_type) in fields.iter().map(|f| f.data_type()).enumerate() {
364 if let Expr::Literal(ScalarValue::Null, metadata) = &row[j] {
365 row[j] = Expr::Literal(
366 ScalarValue::try_from(field_type)?,
367 metadata.clone(),
368 );
369 } else {
370 row[j] = std::mem::take(&mut row[j]).cast_to(field_type, schema)?;
371 }
372 }
373 }
374
375 let dfschema = DFSchema::from_unqualified_fields(fields, HashMap::new())?;
376 let schema = DFSchemaRef::new(dfschema);
377
378 Ok(Self::new(LogicalPlan::Values(Values { schema, values })))
379 }
380
381 pub fn scan(
414 table_name: impl Into<TableReference>,
415 table_source: Arc<dyn TableSource>,
416 projection: Option<Vec<usize>>,
417 ) -> Result<Self> {
418 Self::scan_with_filters(table_name, table_source, projection, vec![])
419 }
420
421 pub fn copy_to(
423 input: LogicalPlan,
424 output_url: String,
425 file_type: Arc<dyn FileType>,
426 options: HashMap<String, String>,
427 partition_by: Vec<String>,
428 ) -> Result<Self> {
429 Ok(Self::new(LogicalPlan::Copy(CopyTo::new(
430 Arc::new(input),
431 output_url,
432 partition_by,
433 file_type,
434 options,
435 ))))
436 }
437
438 pub fn insert_into(
472 input: LogicalPlan,
473 table_name: impl Into<TableReference>,
474 target: Arc<dyn TableSource>,
475 insert_op: InsertOp,
476 ) -> Result<Self> {
477 Ok(Self::new(LogicalPlan::Dml(DmlStatement::new(
478 table_name.into(),
479 target,
480 WriteOp::Insert(insert_op),
481 Arc::new(input),
482 ))))
483 }
484
485 pub fn scan_with_filters(
487 table_name: impl Into<TableReference>,
488 table_source: Arc<dyn TableSource>,
489 projection: Option<Vec<usize>>,
490 filters: Vec<Expr>,
491 ) -> Result<Self> {
492 Self::scan_with_filters_inner(table_name, table_source, projection, filters, None)
493 }
494
495 pub fn scan_with_filters_fetch(
497 table_name: impl Into<TableReference>,
498 table_source: Arc<dyn TableSource>,
499 projection: Option<Vec<usize>>,
500 filters: Vec<Expr>,
501 fetch: Option<usize>,
502 ) -> Result<Self> {
503 Self::scan_with_filters_inner(
504 table_name,
505 table_source,
506 projection,
507 filters,
508 fetch,
509 )
510 }
511
512 fn scan_with_filters_inner(
513 table_name: impl Into<TableReference>,
514 table_source: Arc<dyn TableSource>,
515 projection: Option<Vec<usize>>,
516 filters: Vec<Expr>,
517 fetch: Option<usize>,
518 ) -> Result<Self> {
519 let table_scan = TableScanBuilder::new(table_name, table_source)
520 .with_projection(projection)
521 .with_filters(filters)
522 .with_fetch(fetch)
523 .build()?;
524
525 if table_scan.filters.is_empty()
527 && let Some(p) = table_scan.source.get_logical_plan()
528 {
529 let sub_plan = p.into_owned();
530
531 if let Some(proj) = table_scan.projection {
532 let projection_exprs = proj
533 .into_iter()
534 .map(|i| {
535 Expr::Column(Column::from(sub_plan.schema().qualified_field(i)))
536 })
537 .collect::<Vec<_>>();
538 return Self::new(sub_plan)
539 .project(projection_exprs)?
540 .alias(table_scan.table_name);
541 }
542
543 return Self::new(sub_plan).alias(table_scan.table_name);
547 }
548
549 Ok(Self::new(LogicalPlan::TableScan(table_scan)))
550 }
551
552 pub fn window_plan(
554 input: LogicalPlan,
555 window_exprs: impl IntoIterator<Item = Expr>,
556 ) -> Result<LogicalPlan> {
557 let mut plan = input;
558 let mut groups = group_window_expr_by_sort_keys(window_exprs)?;
559 groups.sort_by(|(key_a, _), (key_b, _)| {
565 for ((first, _), (second, _)) in key_a.iter().zip(key_b.iter()) {
566 let key_ordering = compare_sort_expr(first, second, plan.schema());
567 match key_ordering {
568 Ordering::Less => {
569 return Ordering::Less;
570 }
571 Ordering::Greater => {
572 return Ordering::Greater;
573 }
574 Ordering::Equal => {}
575 }
576 }
577 key_b.len().cmp(&key_a.len())
578 });
579 for (_, exprs) in groups {
580 let window_exprs = exprs.into_iter().collect::<Vec<_>>();
581 plan = LogicalPlanBuilder::from(plan)
584 .window(window_exprs)?
585 .build()?;
586 }
587 Ok(plan)
588 }
589
590 pub fn project(
592 self,
593 expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
594 ) -> Result<Self> {
595 project(Arc::unwrap_or_clone(self.plan), expr).map(Self::new)
596 }
597
598 pub fn project_with_validation(
601 self,
602 expr: Vec<(impl Into<SelectExpr>, bool)>,
603 ) -> Result<Self> {
604 project_with_validation(Arc::unwrap_or_clone(self.plan), expr, None)
605 .map(Self::new)
606 }
607
608 pub fn project_with_validation_and_schema(
611 self,
612 expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
613 schema: &DFSchemaRef,
614 ) -> Result<Self> {
615 project_with_validation(
616 Arc::unwrap_or_clone(self.plan),
617 expr.into_iter().map(|e| (e, true)),
618 Some(schema),
619 )
620 .map(Self::new)
621 }
622
623 pub fn select(self, indices: impl IntoIterator<Item = usize>) -> Result<Self> {
625 let exprs: Vec<_> = indices
626 .into_iter()
627 .map(|x| Expr::Column(Column::from(self.plan.schema().qualified_field(x))))
628 .collect();
629 self.project(exprs)
630 }
631
632 pub fn filter(self, expr: impl Into<Expr>) -> Result<Self> {
634 let expr = normalize_col(expr.into(), &self.plan)?;
635 Filter::try_new(expr, self.plan)
636 .map(LogicalPlan::Filter)
637 .map(Self::new)
638 }
639
640 pub fn having(self, expr: impl Into<Expr>) -> Result<Self> {
642 let expr = normalize_col(expr.into(), &self.plan)?;
643 Filter::try_new(expr, self.plan)
644 .map(LogicalPlan::Filter)
645 .map(Self::from)
646 }
647
648 pub fn prepare(self, name: String, fields: Vec<FieldRef>) -> Result<Self> {
650 Ok(Self::new(LogicalPlan::Statement(Statement::Prepare(
651 Prepare {
652 name,
653 fields,
654 input: self.plan,
655 },
656 ))))
657 }
658
659 pub fn limit(self, skip: usize, fetch: Option<usize>) -> Result<Self> {
666 let skip_expr = if skip == 0 {
667 None
668 } else {
669 Some(lit(skip as i64))
670 };
671 let fetch_expr = fetch.map(|f| lit(f as i64));
672 self.limit_by_expr(skip_expr, fetch_expr)
673 }
674
675 pub fn limit_by_expr(self, skip: Option<Expr>, fetch: Option<Expr>) -> Result<Self> {
679 Ok(Self::new(LogicalPlan::Limit(Limit {
680 skip: skip.map(Box::new),
681 fetch: fetch.map(Box::new),
682 input: self.plan,
683 })))
684 }
685
686 pub fn alias(self, alias: impl Into<TableReference>) -> Result<Self> {
688 subquery_alias(Arc::unwrap_or_clone(self.plan), alias).map(Self::new)
689 }
690
691 fn add_missing_columns(
720 curr_plan: LogicalPlan,
721 missing_cols: &IndexSet<Column>,
722 is_distinct: bool,
723 ) -> Result<LogicalPlan> {
724 match curr_plan {
725 LogicalPlan::Projection(Projection {
726 input,
727 mut expr,
728 schema: _,
729 }) if missing_cols.iter().all(|c| input.schema().has_column(c)) => {
730 let mut missing_exprs = missing_cols
731 .iter()
732 .map(|c| normalize_col(Expr::Column(c.clone()), &input))
733 .collect::<Result<Vec<_>>>()?;
734
735 missing_exprs.retain(|e| !expr.contains(e));
739 if is_distinct {
740 Self::ambiguous_distinct_check(&missing_exprs, missing_cols, &expr)?;
741 }
742 expr.extend(missing_exprs);
743 project(Arc::unwrap_or_clone(input), expr)
744 }
745 _ => {
746 let is_distinct =
747 is_distinct || matches!(curr_plan, LogicalPlan::Distinct(_));
748 let new_inputs = curr_plan
749 .inputs()
750 .into_iter()
751 .map(|input_plan| {
752 Self::add_missing_columns(
753 (*input_plan).clone(),
754 missing_cols,
755 is_distinct,
756 )
757 })
758 .collect::<Result<Vec<_>>>()?;
759 curr_plan.with_new_exprs(curr_plan.expressions(), new_inputs)
760 }
761 }
762 }
763
764 fn ambiguous_distinct_check(
765 missing_exprs: &[Expr],
766 missing_cols: &IndexSet<Column>,
767 projection_exprs: &[Expr],
768 ) -> Result<()> {
769 if missing_exprs.is_empty() {
770 return Ok(());
771 }
772
773 let all_aliases = missing_exprs.iter().all(|e| {
781 projection_exprs.iter().any(|proj_expr| {
782 if let Expr::Alias(Alias { expr, .. }) = proj_expr {
783 e == expr.as_ref()
784 } else {
785 false
786 }
787 })
788 });
789 if all_aliases {
790 return Ok(());
791 }
792
793 let missing_col_names = missing_cols
794 .iter()
795 .map(|col| col.flat_name())
796 .collect::<String>();
797
798 plan_err!(
799 "For SELECT DISTINCT, ORDER BY expressions {missing_col_names} must appear in select list"
800 )
801 }
802
803 pub fn sort_by(
805 self,
806 expr: impl IntoIterator<Item = impl Into<Expr>> + Clone,
807 ) -> Result<Self> {
808 self.sort(
809 expr.into_iter()
810 .map(|e| e.into().sort(true, false))
811 .collect::<Vec<SortExpr>>(),
812 )
813 }
814
815 pub fn sort(
816 self,
817 sorts: impl IntoIterator<Item = impl Into<SortExpr>> + Clone,
818 ) -> Result<Self> {
819 self.sort_with_limit(sorts, None)
820 }
821
822 pub fn sort_with_limit(
824 self,
825 sorts: impl IntoIterator<Item = impl Into<SortExpr>> + Clone,
826 fetch: Option<usize>,
827 ) -> Result<Self> {
828 let sorts = rewrite_sort_cols_by_aggs(sorts, &self.plan)?;
829
830 let schema = self.plan.schema();
831
832 let mut missing_cols: IndexSet<Column> = IndexSet::new();
834 sorts.iter().try_for_each::<_, Result<()>>(|sort| {
835 let columns = sort.expr.column_refs();
836
837 missing_cols.extend(
838 columns
839 .into_iter()
840 .filter(|c| !schema.has_column(c))
841 .cloned(),
842 );
843
844 Ok(())
845 })?;
846
847 if missing_cols.is_empty() {
848 return Ok(Self::new(LogicalPlan::Sort(Sort {
849 expr: normalize_sorts(sorts, &self.plan)?,
850 input: self.plan,
851 fetch,
852 })));
853 }
854
855 let new_expr = schema.columns().into_iter().map(Expr::Column).collect();
857
858 let is_distinct = false;
859 let plan = Self::add_missing_columns(
860 Arc::unwrap_or_clone(self.plan),
861 &missing_cols,
862 is_distinct,
863 )?;
864
865 let sort_plan = LogicalPlan::Sort(Sort {
866 expr: normalize_sorts(sorts, &plan)?,
867 input: Arc::new(plan),
868 fetch,
869 });
870
871 Projection::try_new(new_expr, Arc::new(sort_plan))
872 .map(LogicalPlan::Projection)
873 .map(Self::new)
874 }
875
876 pub fn union(self, plan: LogicalPlan) -> Result<Self> {
878 union(Arc::unwrap_or_clone(self.plan), plan).map(Self::new)
879 }
880
881 pub fn union_by_name(self, plan: LogicalPlan) -> Result<Self> {
883 union_by_name(Arc::unwrap_or_clone(self.plan), plan).map(Self::new)
884 }
885
886 pub fn union_by_name_distinct(self, plan: LogicalPlan) -> Result<Self> {
888 let left_plan: LogicalPlan = Arc::unwrap_or_clone(self.plan);
889 let right_plan: LogicalPlan = plan;
890
891 Ok(Self::new(LogicalPlan::Distinct(Distinct::All(Arc::new(
892 union_by_name(left_plan, right_plan)?,
893 )))))
894 }
895
896 pub fn union_distinct(self, plan: LogicalPlan) -> Result<Self> {
898 let left_plan: LogicalPlan = Arc::unwrap_or_clone(self.plan);
899 let right_plan: LogicalPlan = plan;
900
901 Ok(Self::new(LogicalPlan::Distinct(Distinct::All(Arc::new(
902 union(left_plan, right_plan)?,
903 )))))
904 }
905
906 pub fn distinct(self) -> Result<Self> {
908 Ok(Self::new(LogicalPlan::Distinct(Distinct::All(self.plan))))
909 }
910
911 pub fn distinct_on(
914 self,
915 on_expr: Vec<Expr>,
916 select_expr: Vec<Expr>,
917 sort_expr: Option<Vec<SortExpr>>,
918 ) -> Result<Self> {
919 Ok(Self::new(LogicalPlan::Distinct(Distinct::On(
920 DistinctOn::try_new(on_expr, select_expr, sort_expr, self.plan)?,
921 ))))
922 }
923
924 pub fn join(
938 self,
939 right: LogicalPlan,
940 join_type: JoinType,
941 join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
942 filter: Option<Expr>,
943 ) -> Result<Self> {
944 self.join_detailed(
945 right,
946 join_type,
947 join_keys,
948 filter,
949 NullEquality::NullEqualsNothing,
950 )
951 }
952
953 pub fn join_on(
994 self,
995 right: LogicalPlan,
996 join_type: JoinType,
997 on_exprs: impl IntoIterator<Item = Expr>,
998 ) -> Result<Self> {
999 let filter = on_exprs.into_iter().reduce(Expr::and);
1000
1001 self.join_detailed(
1002 right,
1003 join_type,
1004 (Vec::<Column>::new(), Vec::<Column>::new()),
1005 filter,
1006 NullEquality::NullEqualsNothing,
1007 )
1008 }
1009
1010 pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result<Column> {
1011 if column.relation.is_some() {
1012 return Ok(column);
1014 }
1015
1016 let schema = plan.schema();
1017 let fallback_schemas = plan.fallback_normalize_schemas();
1018 let using_columns = plan.using_columns()?;
1019 column.normalize_with_schemas_and_ambiguity_check(
1020 &[&[schema], &fallback_schemas],
1021 &using_columns,
1022 )
1023 }
1024
1025 pub fn join_detailed(
1032 self,
1033 right: LogicalPlan,
1034 join_type: JoinType,
1035 join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
1036 filter: Option<Expr>,
1037 null_equality: NullEquality,
1038 ) -> Result<Self> {
1039 self.join_detailed_with_options(
1040 right,
1041 join_type,
1042 join_keys,
1043 filter,
1044 null_equality,
1045 false,
1046 )
1047 }
1048
1049 pub fn join_detailed_with_options(
1050 self,
1051 right: LogicalPlan,
1052 join_type: JoinType,
1053 join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
1054 filter: Option<Expr>,
1055 null_equality: NullEquality,
1056 null_aware: bool,
1057 ) -> Result<Self> {
1058 if join_keys.0.len() != join_keys.1.len() {
1059 return plan_err!("left_keys and right_keys were not the same length");
1060 }
1061
1062 let filter = if let Some(expr) = filter {
1063 let filter = normalize_col_with_schemas_and_ambiguity_check(
1064 expr,
1065 &[&[self.schema(), right.schema()]],
1066 &[],
1067 )?;
1068 Some(filter)
1069 } else {
1070 None
1071 };
1072
1073 let (left_keys, right_keys): (Vec<Result<Column>>, Vec<Result<Column>>) =
1074 join_keys
1075 .0
1076 .into_iter()
1077 .zip(join_keys.1)
1078 .map(|(l, r)| {
1079 let l = l.into();
1080 let r = r.into();
1081
1082 match (&l.relation, &r.relation) {
1083 (Some(lr), Some(rr)) => {
1084 let l_is_left =
1085 self.plan.schema().field_with_qualified_name(lr, &l.name);
1086 let l_is_right =
1087 right.schema().field_with_qualified_name(lr, &l.name);
1088 let r_is_left =
1089 self.plan.schema().field_with_qualified_name(rr, &r.name);
1090 let r_is_right =
1091 right.schema().field_with_qualified_name(rr, &r.name);
1092
1093 match (l_is_left, l_is_right, r_is_left, r_is_right) {
1094 (_, Ok(_), Ok(_), _) => (Ok(r), Ok(l)),
1095 (Ok(_), _, _, Ok(_)) => (Ok(l), Ok(r)),
1096 _ => (
1097 Self::normalize(&self.plan, l),
1098 Self::normalize(&right, r),
1099 ),
1100 }
1101 }
1102 (Some(lr), None) => {
1103 let l_is_left =
1104 self.plan.schema().field_with_qualified_name(lr, &l.name);
1105 let l_is_right =
1106 right.schema().field_with_qualified_name(lr, &l.name);
1107
1108 match (l_is_left, l_is_right) {
1109 (Ok(_), _) => (Ok(l), Self::normalize(&right, r)),
1110 (_, Ok(_)) => (Self::normalize(&self.plan, r), Ok(l)),
1111 _ => (
1112 Self::normalize(&self.plan, l),
1113 Self::normalize(&right, r),
1114 ),
1115 }
1116 }
1117 (None, Some(rr)) => {
1118 let r_is_left =
1119 self.plan.schema().field_with_qualified_name(rr, &r.name);
1120 let r_is_right =
1121 right.schema().field_with_qualified_name(rr, &r.name);
1122
1123 match (r_is_left, r_is_right) {
1124 (Ok(_), _) => (Ok(r), Self::normalize(&right, l)),
1125 (_, Ok(_)) => (Self::normalize(&self.plan, l), Ok(r)),
1126 _ => (
1127 Self::normalize(&self.plan, l),
1128 Self::normalize(&right, r),
1129 ),
1130 }
1131 }
1132 (None, None) => {
1133 let mut swap = false;
1134 let left_key = Self::normalize(&self.plan, l.clone())
1135 .or_else(|_| {
1136 swap = true;
1137 Self::normalize(&right, l)
1138 });
1139 if swap {
1140 (Self::normalize(&self.plan, r), left_key)
1141 } else {
1142 (left_key, Self::normalize(&right, r))
1143 }
1144 }
1145 }
1146 })
1147 .unzip();
1148
1149 let left_keys = left_keys.into_iter().collect::<Result<Vec<Column>>>()?;
1150 let right_keys = right_keys.into_iter().collect::<Result<Vec<Column>>>()?;
1151
1152 let on: Vec<_> = left_keys
1153 .into_iter()
1154 .zip(right_keys)
1155 .map(|(l, r)| (Expr::Column(l), Expr::Column(r)))
1156 .collect();
1157 let join_schema =
1158 build_join_schema(self.plan.schema(), right.schema(), &join_type)?;
1159
1160 if join_type != JoinType::Inner && on.is_empty() && filter.is_none() {
1162 return plan_err!("join condition should not be empty");
1163 }
1164
1165 Ok(Self::new(LogicalPlan::Join(Join {
1166 left: self.plan,
1167 right: Arc::new(right),
1168 on,
1169 filter,
1170 join_type,
1171 join_constraint: JoinConstraint::On,
1172 schema: DFSchemaRef::new(join_schema),
1173 null_equality,
1174 null_aware,
1175 })))
1176 }
1177
1178 pub fn join_using(
1180 self,
1181 right: LogicalPlan,
1182 join_type: JoinType,
1183 using_keys: Vec<Column>,
1184 ) -> Result<Self> {
1185 let left_keys: Vec<Column> = using_keys
1186 .clone()
1187 .into_iter()
1188 .map(|c| Self::normalize(&self.plan, c))
1189 .collect::<Result<_>>()?;
1190 let right_keys: Vec<Column> = using_keys
1191 .into_iter()
1192 .map(|c| Self::normalize(&right, c))
1193 .collect::<Result<_>>()?;
1194
1195 let on: Vec<(_, _)> = left_keys.into_iter().zip(right_keys).collect();
1196 let mut join_on: Vec<(Expr, Expr)> = vec![];
1197 let mut filters: Option<Expr> = None;
1198 for (l, r) in &on {
1199 if self.plan.schema().has_column(l)
1200 && right.schema().has_column(r)
1201 && can_hash(
1202 datafusion_common::ExprSchema::field_from_column(
1203 self.plan.schema(),
1204 l,
1205 )?
1206 .data_type(),
1207 )
1208 {
1209 join_on.push((Expr::Column(l.clone()), Expr::Column(r.clone())));
1210 } else if self.plan.schema().has_column(l)
1211 && right.schema().has_column(r)
1212 && can_hash(
1213 datafusion_common::ExprSchema::field_from_column(
1214 self.plan.schema(),
1215 r,
1216 )?
1217 .data_type(),
1218 )
1219 {
1220 join_on.push((Expr::Column(r.clone()), Expr::Column(l.clone())));
1221 } else {
1222 let expr = binary_expr(
1223 Expr::Column(l.clone()),
1224 Operator::Eq,
1225 Expr::Column(r.clone()),
1226 );
1227 match filters {
1228 None => filters = Some(expr),
1229 Some(filter_expr) => filters = Some(and(expr, filter_expr)),
1230 }
1231 }
1232 }
1233
1234 if join_on.is_empty() {
1235 let join = Self::from(self.plan).cross_join(right)?;
1236 join.filter(filters.ok_or_else(|| {
1237 internal_datafusion_err!("filters should not be None here")
1238 })?)
1239 } else {
1240 let join = Join::try_new(
1241 self.plan,
1242 Arc::new(right),
1243 join_on,
1244 filters,
1245 join_type,
1246 JoinConstraint::Using,
1247 NullEquality::NullEqualsNothing,
1248 false, )?;
1250
1251 Ok(Self::new(LogicalPlan::Join(join)))
1252 }
1253 }
1254
1255 pub fn cross_join(self, right: LogicalPlan) -> Result<Self> {
1257 let join = Join::try_new(
1258 self.plan,
1259 Arc::new(right),
1260 vec![],
1261 None,
1262 JoinType::Inner,
1263 JoinConstraint::On,
1264 NullEquality::NullEqualsNothing,
1265 false, )?;
1267
1268 Ok(Self::new(LogicalPlan::Join(join)))
1269 }
1270
1271 pub fn repartition(self, partitioning_scheme: Partitioning) -> Result<Self> {
1273 Ok(Self::new(LogicalPlan::Repartition(Repartition {
1274 input: self.plan,
1275 partitioning_scheme,
1276 })))
1277 }
1278
1279 pub fn window(
1281 self,
1282 window_expr: impl IntoIterator<Item = impl Into<Expr>>,
1283 ) -> Result<Self> {
1284 let window_expr = normalize_cols(window_expr, &self.plan)?;
1285 validate_unique_names("Windows", &window_expr)?;
1286 Ok(Self::new(LogicalPlan::Window(Window::try_new(
1287 window_expr,
1288 self.plan,
1289 )?)))
1290 }
1291
1292 pub fn aggregate(
1296 self,
1297 group_expr: impl IntoIterator<Item = impl Into<Expr>>,
1298 aggr_expr: impl IntoIterator<Item = impl Into<Expr>>,
1299 ) -> Result<Self> {
1300 let group_expr = normalize_cols(group_expr, &self.plan)?;
1301 let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
1302
1303 let group_expr = if self.options.add_implicit_group_by_exprs {
1304 add_group_by_exprs_from_dependencies(group_expr, self.plan.schema())?
1305 } else {
1306 group_expr
1307 };
1308
1309 Aggregate::try_new(self.plan, group_expr, aggr_expr)
1310 .map(LogicalPlan::Aggregate)
1311 .map(Self::new)
1312 }
1313
1314 pub fn explain(self, verbose: bool, analyze: bool) -> Result<Self> {
1321 self.explain_option_format(
1323 ExplainOption::default()
1324 .with_verbose(verbose)
1325 .with_analyze(analyze),
1326 )
1327 }
1328
1329 pub fn explain_option_format(self, explain_option: ExplainOption) -> Result<Self> {
1333 let schema = LogicalPlan::explain_schema();
1334 let schema = schema.to_dfschema_ref()?;
1335
1336 if explain_option.analyze {
1337 Ok(Self::new(LogicalPlan::Analyze(Analyze {
1338 verbose: explain_option.verbose,
1339 format: explain_option.format,
1340 input: self.plan,
1341 schema,
1342 analyze_level: explain_option.analyze_level,
1343 analyze_categories: explain_option.analyze_categories,
1344 })))
1345 } else {
1346 let stringified_plans =
1347 vec![self.plan.to_stringified(PlanType::InitialLogicalPlan)];
1348
1349 Ok(Self::new(LogicalPlan::Explain(Explain {
1350 verbose: explain_option.verbose,
1351 plan: self.plan,
1352 explain_format: explain_option.format,
1353 stringified_plans,
1354 schema,
1355 logical_optimization_succeeded: false,
1356 show_statistics: explain_option.show_statistics,
1357 })))
1358 }
1359 }
1360
1361 pub fn intersect(
1363 left_plan: LogicalPlan,
1364 right_plan: LogicalPlan,
1365 is_all: bool,
1366 ) -> Result<LogicalPlan> {
1367 LogicalPlanBuilder::intersect_or_except(
1368 left_plan,
1369 right_plan,
1370 JoinType::LeftSemi,
1371 is_all,
1372 )
1373 }
1374
1375 pub fn except(
1377 left_plan: LogicalPlan,
1378 right_plan: LogicalPlan,
1379 is_all: bool,
1380 ) -> Result<LogicalPlan> {
1381 LogicalPlanBuilder::intersect_or_except(
1382 left_plan,
1383 right_plan,
1384 JoinType::LeftAnti,
1385 is_all,
1386 )
1387 }
1388
1389 fn intersect_or_except(
1391 left_plan: LogicalPlan,
1392 right_plan: LogicalPlan,
1393 join_type: JoinType,
1394 is_all: bool,
1395 ) -> Result<LogicalPlan> {
1396 let left_len = left_plan.schema().fields().len();
1397 let right_len = right_plan.schema().fields().len();
1398
1399 if left_len != right_len {
1400 return plan_err!(
1401 "INTERSECT/EXCEPT query must have the same number of columns. Left is {left_len} and right is {right_len}."
1402 );
1403 }
1404
1405 let left_builder = LogicalPlanBuilder::from(left_plan);
1408 let right_builder = LogicalPlanBuilder::from(right_plan);
1409 let (left_builder, right_builder, _requalified) =
1410 requalify_sides_if_needed(left_builder, right_builder)?;
1411 let left_plan = left_builder.build()?;
1412 let right_plan = right_builder.build()?;
1413
1414 let join_keys = left_plan
1415 .schema()
1416 .fields()
1417 .iter()
1418 .zip(right_plan.schema().fields().iter())
1419 .map(|(left_field, right_field)| {
1420 (
1421 (Column::from_name(left_field.name())),
1422 (Column::from_name(right_field.name())),
1423 )
1424 })
1425 .unzip();
1426 if is_all {
1427 LogicalPlanBuilder::from(left_plan)
1428 .join_detailed(
1429 right_plan,
1430 join_type,
1431 join_keys,
1432 None,
1433 NullEquality::NullEqualsNull,
1434 )?
1435 .build()
1436 } else {
1437 LogicalPlanBuilder::from(left_plan)
1438 .distinct()?
1439 .join_detailed(
1440 right_plan,
1441 join_type,
1442 join_keys,
1443 None,
1444 NullEquality::NullEqualsNull,
1445 )?
1446 .build()
1447 }
1448 }
1449
1450 pub fn build(self) -> Result<LogicalPlan> {
1452 Ok(Arc::unwrap_or_clone(self.plan))
1453 }
1454
1455 pub fn join_with_expr_keys(
1470 self,
1471 right: LogicalPlan,
1472 join_type: JoinType,
1473 equi_exprs: (Vec<impl Into<Expr>>, Vec<impl Into<Expr>>),
1474 filter: Option<Expr>,
1475 ) -> Result<Self> {
1476 if equi_exprs.0.len() != equi_exprs.1.len() {
1477 return plan_err!("left_keys and right_keys were not the same length");
1478 }
1479
1480 let join_key_pairs = equi_exprs
1481 .0
1482 .into_iter()
1483 .zip(equi_exprs.1)
1484 .map(|(l, r)| {
1485 let left_key = l.into();
1486 let right_key = r.into();
1487 let mut left_using_columns = HashSet::new();
1488 expr_to_columns(&left_key, &mut left_using_columns)?;
1489 let normalized_left_key = normalize_col_with_schemas_and_ambiguity_check(
1490 left_key,
1491 &[&[self.plan.schema()]],
1492 &[],
1493 )?;
1494
1495 let mut right_using_columns = HashSet::new();
1496 expr_to_columns(&right_key, &mut right_using_columns)?;
1497 let normalized_right_key = normalize_col_with_schemas_and_ambiguity_check(
1498 right_key,
1499 &[&[right.schema()]],
1500 &[],
1501 )?;
1502
1503 find_valid_equijoin_key_pair(
1505 &normalized_left_key,
1506 &normalized_right_key,
1507 self.plan.schema(),
1508 right.schema(),
1509 )?.ok_or_else(||
1510 plan_datafusion_err!(
1511 "can't create join plan, join key should belong to one input, error key: ({normalized_left_key},{normalized_right_key})"
1512 ))
1513 })
1514 .collect::<Result<Vec<_>>>()?;
1515
1516 let join = Join::try_new(
1517 self.plan,
1518 Arc::new(right),
1519 join_key_pairs,
1520 filter,
1521 join_type,
1522 JoinConstraint::On,
1523 NullEquality::NullEqualsNothing,
1524 false, )?;
1526
1527 Ok(Self::new(LogicalPlan::Join(join)))
1528 }
1529
1530 pub fn unnest_column(self, column: impl Into<Column>) -> Result<Self> {
1532 unnest(Arc::unwrap_or_clone(self.plan), vec![column.into()]).map(Self::new)
1533 }
1534
1535 pub fn unnest_column_with_options(
1537 self,
1538 column: impl Into<Column>,
1539 options: UnnestOptions,
1540 ) -> Result<Self> {
1541 unnest_with_options(
1542 Arc::unwrap_or_clone(self.plan),
1543 vec![column.into()],
1544 options,
1545 )
1546 .map(Self::new)
1547 }
1548
1549 pub fn unnest_columns_with_options(
1551 self,
1552 columns: Vec<Column>,
1553 options: UnnestOptions,
1554 ) -> Result<Self> {
1555 unnest_with_options(Arc::unwrap_or_clone(self.plan), columns, options)
1556 .map(Self::new)
1557 }
1558}
1559
1560impl From<LogicalPlan> for LogicalPlanBuilder {
1561 fn from(plan: LogicalPlan) -> Self {
1562 LogicalPlanBuilder::new(plan)
1563 }
1564}
1565
1566impl From<Arc<LogicalPlan>> for LogicalPlanBuilder {
1567 fn from(plan: Arc<LogicalPlan>) -> Self {
1568 LogicalPlanBuilder::new_from_arc(plan)
1569 }
1570}
1571
1572#[derive(Default)]
1574struct ValuesFields {
1575 inner: Vec<Field>,
1576}
1577
1578impl ValuesFields {
1579 pub fn new() -> Self {
1580 Self::default()
1581 }
1582
1583 pub fn push(&mut self, data_type: DataType, nullable: bool) {
1584 self.push_with_metadata(data_type, nullable, None);
1585 }
1586
1587 pub fn push_with_metadata(
1588 &mut self,
1589 data_type: DataType,
1590 nullable: bool,
1591 metadata: Option<FieldMetadata>,
1592 ) {
1593 let name = format!("column{}", self.inner.len() + 1);
1596 let mut field = Field::new(name, data_type, nullable);
1597 if let Some(metadata) = metadata {
1598 field.set_metadata(metadata.to_hashmap());
1599 }
1600 self.inner.push(field);
1601 }
1602
1603 pub fn into_fields(self) -> Fields {
1604 self.inner.into()
1605 }
1606}
1607
1608pub fn unique_field_aliases(fields: &Fields) -> Vec<Option<String>> {
1620 let mut name_map = HashMap::<&str, usize>::new();
1628 let mut seen = HashSet::<Cow<String>>::new();
1630
1631 fields
1632 .iter()
1633 .map(|field| {
1634 let original_name = field.name();
1635 let mut name = Cow::Borrowed(original_name);
1636
1637 let count = name_map.entry(original_name).or_insert(0);
1638
1639 while seen.contains(&name) {
1641 *count += 1;
1642 name = Cow::Owned(format!("{original_name}:{count}"));
1643 }
1644
1645 seen.insert(name.clone());
1646
1647 match name {
1648 Cow::Borrowed(_) => None,
1649 Cow::Owned(alias) => Some(alias),
1650 }
1651 })
1652 .collect()
1653}
1654
1655fn mark_field(schema: &DFSchema) -> (Option<TableReference>, Arc<Field>) {
1656 let mut table_references = schema
1657 .iter()
1658 .filter_map(|(qualifier, _)| qualifier)
1659 .collect::<Vec<_>>();
1660 table_references.dedup();
1661 let table_reference = if table_references.len() == 1 {
1662 table_references.pop().cloned()
1663 } else {
1664 None
1665 };
1666
1667 (
1668 table_reference,
1669 Arc::new(Field::new("mark", DataType::Boolean, false)),
1670 )
1671}
1672
1673pub fn build_join_schema(
1676 left: &DFSchema,
1677 right: &DFSchema,
1678 join_type: &JoinType,
1679) -> Result<DFSchema> {
1680 fn nullify_fields<'a>(
1681 fields: impl Iterator<Item = (Option<&'a TableReference>, &'a Arc<Field>)>,
1682 ) -> Vec<(Option<TableReference>, Arc<Field>)> {
1683 fields
1684 .map(|(q, f)| {
1685 let field = f.as_ref().clone().with_nullable(true);
1687 (q.cloned(), Arc::new(field))
1688 })
1689 .collect()
1690 }
1691
1692 let right_fields = right.iter();
1693 let left_fields = left.iter();
1694
1695 let qualified_fields: Vec<(Option<TableReference>, Arc<Field>)> = match join_type {
1696 JoinType::Inner => {
1697 let left_fields = left_fields
1699 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1700 .collect::<Vec<_>>();
1701 let right_fields = right_fields
1702 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1703 .collect::<Vec<_>>();
1704 left_fields.into_iter().chain(right_fields).collect()
1705 }
1706 JoinType::Left => {
1707 let left_fields = left_fields
1709 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1710 .collect::<Vec<_>>();
1711 left_fields
1712 .into_iter()
1713 .chain(nullify_fields(right_fields))
1714 .collect()
1715 }
1716 JoinType::Right => {
1717 let right_fields = right_fields
1719 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1720 .collect::<Vec<_>>();
1721 nullify_fields(left_fields)
1722 .into_iter()
1723 .chain(right_fields)
1724 .collect()
1725 }
1726 JoinType::Full => {
1727 nullify_fields(left_fields)
1729 .into_iter()
1730 .chain(nullify_fields(right_fields))
1731 .collect()
1732 }
1733 JoinType::LeftSemi | JoinType::LeftAnti => {
1734 left_fields
1736 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1737 .collect()
1738 }
1739 JoinType::LeftMark => left_fields
1740 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1741 .chain(once(mark_field(right)))
1742 .collect(),
1743 JoinType::RightSemi | JoinType::RightAnti => {
1744 right_fields
1746 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1747 .collect()
1748 }
1749 JoinType::RightMark => right_fields
1750 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1751 .chain(once(mark_field(left)))
1752 .collect(),
1753 };
1754 let func_dependencies = left.functional_dependencies().join(
1755 right.functional_dependencies(),
1756 join_type,
1757 left.fields().len(),
1758 );
1759
1760 let (schema1, schema2) = match join_type {
1761 JoinType::Right
1762 | JoinType::RightSemi
1763 | JoinType::RightAnti
1764 | JoinType::RightMark => (left, right),
1765 _ => (right, left),
1766 };
1767
1768 let metadata = schema1
1769 .metadata()
1770 .clone()
1771 .into_iter()
1772 .chain(schema2.metadata().clone())
1773 .collect();
1774
1775 let dfschema = DFSchema::new_with_metadata(qualified_fields, metadata)?;
1776 dfschema.with_functional_dependencies(func_dependencies)
1777}
1778
1779pub fn requalify_sides_if_needed(
1789 left: LogicalPlanBuilder,
1790 right: LogicalPlanBuilder,
1791) -> Result<(LogicalPlanBuilder, LogicalPlanBuilder, bool)> {
1792 let left_cols = left.schema().columns();
1793 let right_cols = right.schema().columns();
1794
1795 for l in &left_cols {
1809 for r in &right_cols {
1810 if l.name != r.name {
1811 continue;
1812 }
1813
1814 match (&l.relation, &r.relation) {
1816 (Some(l_rel), Some(r_rel)) if l_rel == r_rel => {
1818 return Ok((
1819 left.alias(TableReference::bare("left"))?,
1820 right.alias(TableReference::bare("right"))?,
1821 true,
1822 ));
1823 }
1824 (None, None) => {
1826 return Ok((
1827 left.alias(TableReference::bare("left"))?,
1828 right.alias(TableReference::bare("right"))?,
1829 true,
1830 ));
1831 }
1832 (Some(_), None) | (None, Some(_)) => {
1834 return Ok((
1835 left.alias(TableReference::bare("left"))?,
1836 right.alias(TableReference::bare("right"))?,
1837 true,
1838 ));
1839 }
1840 _ => {}
1842 }
1843 }
1844 }
1845
1846 Ok((left, right, false))
1848}
1849pub fn add_group_by_exprs_from_dependencies(
1859 mut group_expr: Vec<Expr>,
1860 schema: &DFSchemaRef,
1861) -> Result<Vec<Expr>> {
1862 let mut group_by_field_names = group_expr
1865 .iter()
1866 .map(|e| e.schema_name().to_string())
1867 .collect::<Vec<_>>();
1868
1869 if let Some(target_indices) =
1870 get_target_functional_dependencies(schema, &group_by_field_names)
1871 {
1872 for idx in target_indices {
1873 let expr = Expr::Column(Column::from(schema.qualified_field(idx)));
1874 let expr_name = expr.schema_name().to_string();
1875 if !group_by_field_names.contains(&expr_name) {
1876 group_by_field_names.push(expr_name);
1877 group_expr.push(expr);
1878 }
1879 }
1880 }
1881 Ok(group_expr)
1882}
1883
1884pub fn validate_unique_names<'a>(
1886 node_name: &str,
1887 expressions: impl IntoIterator<Item = &'a Expr>,
1888) -> Result<()> {
1889 let mut unique_names = HashMap::new();
1890
1891 expressions.into_iter().enumerate().try_for_each(|(position, expr)| {
1892 let name = expr.schema_name().to_string();
1893 match unique_names.get(&name) {
1894 None => {
1895 unique_names.insert(name, (position, expr));
1896 Ok(())
1897 },
1898 Some((existing_position, existing_expr)) => {
1899 plan_err!("{node_name} require unique expression names \
1900 but the expression \"{existing_expr}\" at position {existing_position} and \"{expr}\" \
1901 at position {position} have the same name. Consider aliasing (\"AS\") one of them."
1902 )
1903 }
1904 }
1905 })
1906}
1907
1908pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalPlan> {
1920 Ok(LogicalPlan::Union(Union::try_new_with_loose_types(vec![
1921 Arc::new(left_plan),
1922 Arc::new(right_plan),
1923 ])?))
1924}
1925
1926pub fn union_by_name(
1929 left_plan: LogicalPlan,
1930 right_plan: LogicalPlan,
1931) -> Result<LogicalPlan> {
1932 Ok(LogicalPlan::Union(Union::try_new_by_name(vec![
1933 Arc::new(left_plan),
1934 Arc::new(right_plan),
1935 ])?))
1936}
1937
1938pub fn project(
1944 plan: LogicalPlan,
1945 expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
1946) -> Result<LogicalPlan> {
1947 project_with_validation(plan, expr.into_iter().map(|e| (e, true)), None)
1948}
1949
1950fn project_with_validation(
1958 plan: LogicalPlan,
1959 expr: impl IntoIterator<Item = (impl Into<SelectExpr>, bool)>,
1960 schema: Option<&DFSchemaRef>,
1961) -> Result<LogicalPlan> {
1962 let mut projected_expr = vec![];
1963 let mut has_wildcard = false;
1964 for (e, validate) in expr {
1965 let e = e.into();
1966 match e {
1967 SelectExpr::Wildcard(opt) => {
1968 has_wildcard = true;
1969 let expanded = expand_wildcard(plan.schema(), &plan, Some(&opt))?;
1970
1971 let expanded = if let Some(replace) = opt.replace {
1974 replace_columns(expanded, &replace)?
1975 } else {
1976 expanded
1977 };
1978
1979 for e in expanded {
1980 if validate {
1981 projected_expr
1982 .push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
1983 } else {
1984 projected_expr.push(e)
1985 }
1986 }
1987 }
1988 SelectExpr::QualifiedWildcard(table_ref, opt) => {
1989 has_wildcard = true;
1990 let expanded =
1991 expand_qualified_wildcard(&table_ref, plan.schema(), Some(&opt))?;
1992
1993 let expanded = if let Some(replace) = opt.replace {
1996 replace_columns(expanded, &replace)?
1997 } else {
1998 expanded
1999 };
2000
2001 for e in expanded {
2002 if validate {
2003 projected_expr
2004 .push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
2005 } else {
2006 projected_expr.push(e)
2007 }
2008 }
2009 }
2010 SelectExpr::Expression(e) => {
2011 if validate {
2012 projected_expr.push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
2013 } else {
2014 projected_expr.push(e)
2015 }
2016 }
2017 }
2018 }
2019
2020 if has_wildcard && projected_expr.is_empty() && !plan.schema().fields().is_empty() {
2021 return plan_err!(
2022 "SELECT list is empty after resolving * expressions, \
2023 the wildcard expanded to zero columns"
2024 );
2025 }
2026
2027 if let Some(schema) = &schema {
2030 for (expr, field) in projected_expr.iter_mut().zip(schema.fields()) {
2031 if !matches!(expr, Expr::Column(_) | Expr::Alias(_)) {
2032 *expr = std::mem::take(expr).alias(field.name());
2033 }
2034 }
2035 }
2036
2037 validate_unique_names("Projections", projected_expr.iter())?;
2038
2039 Projection::try_new(projected_expr, Arc::new(plan)).map(LogicalPlan::Projection)
2040}
2041
2042fn replace_columns(
2047 mut exprs: Vec<Expr>,
2048 replace: &PlannedReplaceSelectItem,
2049) -> Result<Vec<Expr>> {
2050 for expr in exprs.iter_mut() {
2051 if let Expr::Column(Column { name, .. }) = expr
2052 && let Some((_, new_expr)) = replace
2053 .items()
2054 .iter()
2055 .zip(replace.expressions().iter())
2056 .find(|(item, _)| item.column_name.value == *name)
2057 {
2058 *expr = new_expr.clone().alias(name.clone())
2059 }
2060 }
2061 Ok(exprs)
2062}
2063
2064pub fn subquery_alias(
2066 plan: LogicalPlan,
2067 alias: impl Into<TableReference>,
2068) -> Result<LogicalPlan> {
2069 SubqueryAlias::try_new(Arc::new(plan), alias).map(LogicalPlan::SubqueryAlias)
2070}
2071
2072pub fn table_scan(
2075 name: Option<impl Into<TableReference>>,
2076 table_schema: &Schema,
2077 projection: Option<Vec<usize>>,
2078) -> Result<LogicalPlanBuilder> {
2079 table_scan_with_filters(name, table_schema, projection, vec![])
2080}
2081
2082pub fn table_scan_with_filters(
2086 name: Option<impl Into<TableReference>>,
2087 table_schema: &Schema,
2088 projection: Option<Vec<usize>>,
2089 filters: Vec<Expr>,
2090) -> Result<LogicalPlanBuilder> {
2091 let table_source = table_source(table_schema);
2092 let name = name
2093 .map(|n| n.into())
2094 .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
2095 LogicalPlanBuilder::scan_with_filters(name, table_source, projection, filters)
2096}
2097
2098pub fn table_scan_with_filter_and_fetch(
2102 name: Option<impl Into<TableReference>>,
2103 table_schema: &Schema,
2104 projection: Option<Vec<usize>>,
2105 filters: Vec<Expr>,
2106 fetch: Option<usize>,
2107) -> Result<LogicalPlanBuilder> {
2108 let table_source = table_source(table_schema);
2109 let name = name
2110 .map(|n| n.into())
2111 .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
2112 LogicalPlanBuilder::scan_with_filters_fetch(
2113 name,
2114 table_source,
2115 projection,
2116 filters,
2117 fetch,
2118 )
2119}
2120
2121pub fn table_source(table_schema: &Schema) -> Arc<dyn TableSource> {
2122 let table_schema = Arc::new(table_schema.clone());
2124 Arc::new(LogicalTableSource {
2125 table_schema,
2126 constraints: Default::default(),
2127 })
2128}
2129
2130pub fn table_source_with_constraints(
2131 table_schema: &Schema,
2132 constraints: Constraints,
2133) -> Arc<dyn TableSource> {
2134 let table_schema = Arc::new(table_schema.clone());
2136 Arc::new(LogicalTableSource {
2137 table_schema,
2138 constraints,
2139 })
2140}
2141
2142pub fn wrap_projection_for_join_if_necessary(
2144 join_keys: &[Expr],
2145 input: LogicalPlan,
2146) -> Result<(LogicalPlan, Vec<Column>, bool)> {
2147 let input_schema = input.schema();
2148 let alias_join_keys: Vec<Expr> = join_keys
2149 .iter()
2150 .map(|key| {
2151 if matches!(key, Expr::Cast(_)) || matches!(key, Expr::TryCast(_)) {
2160 let alias = format!("{key}");
2161 key.clone().alias(alias)
2162 } else {
2163 key.clone()
2164 }
2165 })
2166 .collect::<Vec<_>>();
2167
2168 let need_project = join_keys.iter().any(|key| !matches!(key, Expr::Column(_)));
2169 let plan = if need_project {
2170 let mut projection = input_schema
2172 .columns()
2173 .into_iter()
2174 .map(Expr::Column)
2175 .collect::<Vec<_>>();
2176 #[allow(clippy::allow_attributes, clippy::mutable_key_type)]
2177 let join_key_items = alias_join_keys
2179 .iter()
2180 .flat_map(|expr| expr.try_as_col().is_none().then_some(expr))
2181 .cloned()
2182 .collect::<HashSet<Expr>>();
2183 projection.extend(join_key_items);
2184
2185 LogicalPlanBuilder::from(input)
2186 .project(projection.into_iter().map(SelectExpr::from))?
2187 .build()?
2188 } else {
2189 input
2190 };
2191
2192 let join_on = alias_join_keys
2193 .into_iter()
2194 .map(|key| {
2195 if let Some(col) = key.try_as_col() {
2196 Ok(col.clone())
2197 } else {
2198 let name = key.schema_name().to_string();
2199 Ok(Column::from_name(name))
2200 }
2201 })
2202 .collect::<Result<Vec<_>>>()?;
2203
2204 Ok((plan, join_on, need_project))
2205}
2206
2207pub struct LogicalTableSource {
2211 table_schema: SchemaRef,
2212 constraints: Constraints,
2213}
2214
2215impl LogicalTableSource {
2216 pub fn new(table_schema: SchemaRef) -> Self {
2218 Self {
2219 table_schema,
2220 constraints: Constraints::default(),
2221 }
2222 }
2223
2224 pub fn with_constraints(mut self, constraints: Constraints) -> Self {
2225 self.constraints = constraints;
2226 self
2227 }
2228}
2229
2230impl TableSource for LogicalTableSource {
2231 fn schema(&self) -> SchemaRef {
2232 Arc::clone(&self.table_schema)
2233 }
2234
2235 fn constraints(&self) -> Option<&Constraints> {
2236 Some(&self.constraints)
2237 }
2238
2239 fn supports_filters_pushdown(
2240 &self,
2241 filters: &[&Expr],
2242 ) -> Result<Vec<TableProviderFilterPushDown>> {
2243 Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
2244 }
2245}
2246
2247pub fn unnest(input: LogicalPlan, columns: Vec<Column>) -> Result<LogicalPlan> {
2249 unnest_with_options(input, columns, UnnestOptions::default())
2250}
2251
2252pub fn get_struct_unnested_columns(
2253 col_name: &String,
2254 inner_fields: &Fields,
2255) -> Vec<Column> {
2256 inner_fields
2257 .iter()
2258 .map(|f| Column::from_name(format!("{}.{}", col_name, f.name())))
2259 .collect()
2260}
2261
2262pub fn unnest_with_options(
2292 input: LogicalPlan,
2293 columns_to_unnest: Vec<Column>,
2294 options: UnnestOptions,
2295) -> Result<LogicalPlan> {
2296 Ok(LogicalPlan::Unnest(Unnest::try_new(
2297 Arc::new(input),
2298 columns_to_unnest,
2299 options,
2300 )?))
2301}
2302
2303#[cfg(test)]
2304mod tests {
2305 use std::vec;
2306
2307 use super::*;
2308 use crate::lit_with_metadata;
2309 use crate::logical_plan::StringifiedPlan;
2310 use crate::{col, expr, expr_fn::exists, in_subquery, scalar_subquery};
2311
2312 use crate::test::function_stub::sum;
2313 use datafusion_common::{
2314 Constraint, DataFusionError, RecursionUnnestOption, SchemaError,
2315 };
2316 use insta::assert_snapshot;
2317
2318 #[test]
2319 fn plan_builder_simple() -> Result<()> {
2320 let plan =
2321 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
2322 .filter(col("state").eq(lit("CO")))?
2323 .project(vec![col("id")])?
2324 .build()?;
2325
2326 assert_snapshot!(plan, @r#"
2327 Projection: employee_csv.id
2328 Filter: employee_csv.state = Utf8("CO")
2329 TableScan: employee_csv projection=[id, state]
2330 "#);
2331
2332 Ok(())
2333 }
2334
2335 #[test]
2336 fn plan_builder_schema() {
2337 let schema = employee_schema();
2338 let projection = None;
2339 let plan =
2340 LogicalPlanBuilder::scan("employee_csv", table_source(&schema), projection)
2341 .unwrap();
2342 assert_snapshot!(plan.schema().as_ref(), @"fields:[employee_csv.id, employee_csv.first_name, employee_csv.last_name, employee_csv.state, employee_csv.salary], metadata:{}");
2343
2344 let projection = None;
2347 let plan =
2348 LogicalPlanBuilder::scan("EMPLOYEE_CSV", table_source(&schema), projection)
2349 .unwrap();
2350 assert_snapshot!(plan.schema().as_ref(), @"fields:[employee_csv.id, employee_csv.first_name, employee_csv.last_name, employee_csv.state, employee_csv.salary], metadata:{}");
2351 }
2352
2353 #[test]
2354 fn plan_builder_empty_name() {
2355 let schema = employee_schema();
2356 let projection = None;
2357 let err =
2358 LogicalPlanBuilder::scan("", table_source(&schema), projection).unwrap_err();
2359 assert_snapshot!(
2360 err.strip_backtrace(),
2361 @"Error during planning: table_name cannot be empty"
2362 );
2363 }
2364
2365 #[test]
2366 fn plan_builder_sort() -> Result<()> {
2367 let plan =
2368 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?
2369 .sort(vec![
2370 expr::Sort::new(col("state"), true, true),
2371 expr::Sort::new(col("salary"), false, false),
2372 ])?
2373 .build()?;
2374
2375 assert_snapshot!(plan, @r"
2376 Sort: employee_csv.state ASC NULLS FIRST, employee_csv.salary DESC NULLS LAST
2377 TableScan: employee_csv projection=[state, salary]
2378 ");
2379
2380 Ok(())
2381 }
2382
2383 #[test]
2384 fn plan_builder_union() -> Result<()> {
2385 let plan =
2386 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?;
2387
2388 let plan = plan
2389 .clone()
2390 .union(plan.clone().build()?)?
2391 .union(plan.clone().build()?)?
2392 .union(plan.build()?)?
2393 .build()?;
2394
2395 assert_snapshot!(plan, @r"
2396 Union
2397 Union
2398 Union
2399 TableScan: employee_csv projection=[state, salary]
2400 TableScan: employee_csv projection=[state, salary]
2401 TableScan: employee_csv projection=[state, salary]
2402 TableScan: employee_csv projection=[state, salary]
2403 ");
2404
2405 Ok(())
2406 }
2407
2408 #[test]
2409 fn plan_builder_union_distinct() -> Result<()> {
2410 let plan =
2411 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?;
2412
2413 let plan = plan
2414 .clone()
2415 .union_distinct(plan.clone().build()?)?
2416 .union_distinct(plan.clone().build()?)?
2417 .union_distinct(plan.build()?)?
2418 .build()?;
2419
2420 assert_snapshot!(plan, @r"
2421 Distinct:
2422 Union
2423 Distinct:
2424 Union
2425 Distinct:
2426 Union
2427 TableScan: employee_csv projection=[state, salary]
2428 TableScan: employee_csv projection=[state, salary]
2429 TableScan: employee_csv projection=[state, salary]
2430 TableScan: employee_csv projection=[state, salary]
2431 ");
2432
2433 Ok(())
2434 }
2435
2436 #[test]
2437 fn plan_builder_simple_distinct() -> Result<()> {
2438 let plan =
2439 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
2440 .filter(col("state").eq(lit("CO")))?
2441 .project(vec![col("id")])?
2442 .distinct()?
2443 .build()?;
2444
2445 assert_snapshot!(plan, @r#"
2446 Distinct:
2447 Projection: employee_csv.id
2448 Filter: employee_csv.state = Utf8("CO")
2449 TableScan: employee_csv projection=[id, state]
2450 "#);
2451
2452 Ok(())
2453 }
2454
2455 #[test]
2456 fn exists_subquery() -> Result<()> {
2457 let foo = test_table_scan_with_name("foo")?;
2458 let bar = test_table_scan_with_name("bar")?;
2459
2460 let subquery = LogicalPlanBuilder::from(foo)
2461 .project(vec![col("a")])?
2462 .filter(col("a").eq(col("bar.a")))?
2463 .build()?;
2464
2465 let outer_query = LogicalPlanBuilder::from(bar)
2466 .project(vec![col("a")])?
2467 .filter(exists(Arc::new(subquery)))?
2468 .build()?;
2469
2470 assert_snapshot!(outer_query, @r"
2471 Filter: EXISTS (<subquery>)
2472 Subquery:
2473 Filter: foo.a = bar.a
2474 Projection: foo.a
2475 TableScan: foo
2476 Projection: bar.a
2477 TableScan: bar
2478 ");
2479
2480 Ok(())
2481 }
2482
2483 #[test]
2484 fn filter_in_subquery() -> Result<()> {
2485 let foo = test_table_scan_with_name("foo")?;
2486 let bar = test_table_scan_with_name("bar")?;
2487
2488 let subquery = LogicalPlanBuilder::from(foo)
2489 .project(vec![col("a")])?
2490 .filter(col("a").eq(col("bar.a")))?
2491 .build()?;
2492
2493 let outer_query = LogicalPlanBuilder::from(bar)
2495 .project(vec![col("a")])?
2496 .filter(in_subquery(col("a"), Arc::new(subquery)))?
2497 .build()?;
2498
2499 assert_snapshot!(outer_query, @r"
2500 Filter: bar.a IN (<subquery>)
2501 Subquery:
2502 Filter: foo.a = bar.a
2503 Projection: foo.a
2504 TableScan: foo
2505 Projection: bar.a
2506 TableScan: bar
2507 ");
2508
2509 Ok(())
2510 }
2511
2512 #[test]
2513 fn select_scalar_subquery() -> Result<()> {
2514 let foo = test_table_scan_with_name("foo")?;
2515 let bar = test_table_scan_with_name("bar")?;
2516
2517 let subquery = LogicalPlanBuilder::from(foo)
2518 .project(vec![col("b")])?
2519 .filter(col("a").eq(col("bar.a")))?
2520 .build()?;
2521
2522 let outer_query = LogicalPlanBuilder::from(bar)
2524 .project(vec![scalar_subquery(Arc::new(subquery))])?
2525 .build()?;
2526
2527 assert_snapshot!(outer_query, @r"
2528 Projection: (<subquery>)
2529 Subquery:
2530 Filter: foo.a = bar.a
2531 Projection: foo.b
2532 TableScan: foo
2533 TableScan: bar
2534 ");
2535
2536 Ok(())
2537 }
2538
2539 #[test]
2540 fn projection_non_unique_names() -> Result<()> {
2541 let plan = table_scan(
2542 Some("employee_csv"),
2543 &employee_schema(),
2544 Some(vec![0, 1]),
2546 )?
2547 .project(vec![col("id"), col("first_name").alias("id")]);
2549
2550 match plan {
2551 Err(DataFusionError::SchemaError(err, _)) => {
2552 if let SchemaError::AmbiguousReference { field } = *err {
2553 let Column {
2554 relation,
2555 name,
2556 spans: _,
2557 } = *field;
2558 let Some(TableReference::Bare { table }) = relation else {
2559 return plan_err!(
2560 "wrong relation: {relation:?}, expected table name"
2561 );
2562 };
2563 assert_eq!(*"employee_csv", *table);
2564 assert_eq!("id", &name);
2565 Ok(())
2566 } else {
2567 plan_err!("Plan should have returned an DataFusionError::SchemaError")
2568 }
2569 }
2570 _ => plan_err!("Plan should have returned an DataFusionError::SchemaError"),
2571 }
2572 }
2573
2574 fn employee_schema() -> Schema {
2575 Schema::new(vec![
2576 Field::new("id", DataType::Int32, false),
2577 Field::new("first_name", DataType::Utf8, false),
2578 Field::new("last_name", DataType::Utf8, false),
2579 Field::new("state", DataType::Utf8, false),
2580 Field::new("salary", DataType::Int32, false),
2581 ])
2582 }
2583
2584 #[test]
2585 fn stringified_plan() {
2586 let stringified_plan =
2587 StringifiedPlan::new(PlanType::InitialLogicalPlan, "...the plan...");
2588 assert!(stringified_plan.should_display(true));
2589 assert!(!stringified_plan.should_display(false)); let stringified_plan =
2592 StringifiedPlan::new(PlanType::FinalLogicalPlan, "...the plan...");
2593 assert!(stringified_plan.should_display(true));
2594 assert!(stringified_plan.should_display(false)); let stringified_plan =
2597 StringifiedPlan::new(PlanType::InitialPhysicalPlan, "...the plan...");
2598 assert!(stringified_plan.should_display(true));
2599 assert!(!stringified_plan.should_display(false)); let stringified_plan =
2602 StringifiedPlan::new(PlanType::FinalPhysicalPlan, "...the plan...");
2603 assert!(stringified_plan.should_display(true));
2604 assert!(stringified_plan.should_display(false)); let stringified_plan = StringifiedPlan::new(
2607 PlanType::OptimizedLogicalPlan {
2608 optimizer_name: "random opt pass".into(),
2609 },
2610 "...the plan...",
2611 );
2612 assert!(stringified_plan.should_display(true));
2613 assert!(!stringified_plan.should_display(false));
2614 }
2615
2616 fn test_table_scan_with_name(name: &str) -> Result<LogicalPlan> {
2617 let schema = Schema::new(vec![
2618 Field::new("a", DataType::UInt32, false),
2619 Field::new("b", DataType::UInt32, false),
2620 Field::new("c", DataType::UInt32, false),
2621 ]);
2622 table_scan(Some(name), &schema, None)?.build()
2623 }
2624
2625 #[test]
2626 fn plan_builder_intersect_different_num_columns_error() -> Result<()> {
2627 let plan1 =
2628 table_scan(TableReference::none(), &employee_schema(), Some(vec![3]))?;
2629 let plan2 =
2630 table_scan(TableReference::none(), &employee_schema(), Some(vec![3, 4]))?;
2631
2632 let err_msg1 =
2633 LogicalPlanBuilder::intersect(plan1.build()?, plan2.build()?, true)
2634 .unwrap_err();
2635
2636 assert_snapshot!(err_msg1.strip_backtrace(), @"Error during planning: INTERSECT/EXCEPT query must have the same number of columns. Left is 1 and right is 2.");
2637
2638 Ok(())
2639 }
2640
2641 #[test]
2642 fn plan_builder_unnest() -> Result<()> {
2643 let err = nested_table_scan("test_table")?
2645 .unnest_column("scalar")
2646 .unwrap_err();
2647
2648 let DataFusionError::Internal(desc) = err else {
2649 return plan_err!("Plan should have returned an DataFusionError::Internal");
2650 };
2651
2652 let desc = (*desc
2653 .split(DataFusionError::BACK_TRACE_SEP)
2654 .collect::<Vec<&str>>()
2655 .first()
2656 .unwrap_or(&""))
2657 .to_string();
2658
2659 assert_snapshot!(desc, @"trying to unnest on invalid data type UInt32");
2660
2661 let plan = nested_table_scan("test_table")?
2663 .unnest_column("strings")?
2664 .build()?;
2665
2666 assert_snapshot!(plan, @r"
2667 Unnest: lists[test_table.strings|depth=1] structs[]
2668 TableScan: test_table
2669 ");
2670
2671 let field = plan.schema().field_with_name(None, "strings").unwrap();
2673 assert_eq!(&DataType::Utf8, field.data_type());
2674
2675 let plan = nested_table_scan("test_table")?
2677 .unnest_column("struct_singular")?
2678 .build()?;
2679
2680 assert_snapshot!(plan, @r"
2681 Unnest: lists[] structs[test_table.struct_singular]
2682 TableScan: test_table
2683 ");
2684
2685 for field_name in &["a", "b"] {
2686 let field = plan
2688 .schema()
2689 .field_with_name(None, &format!("struct_singular.{field_name}"))
2690 .unwrap();
2691 assert_eq!(&DataType::UInt32, field.data_type());
2692 }
2693
2694 let plan = nested_table_scan("test_table")?
2696 .unnest_column("strings")?
2697 .unnest_column("structs")?
2698 .unnest_column("struct_singular")?
2699 .build()?;
2700
2701 assert_snapshot!(plan, @r"
2702 Unnest: lists[] structs[test_table.struct_singular]
2703 Unnest: lists[test_table.structs|depth=1] structs[]
2704 Unnest: lists[test_table.strings|depth=1] structs[]
2705 TableScan: test_table
2706 ");
2707
2708 let field = plan.schema().field_with_name(None, "structs").unwrap();
2710 assert!(matches!(field.data_type(), DataType::Struct(_)));
2711
2712 let cols = vec!["strings", "structs", "struct_singular"]
2714 .into_iter()
2715 .map(|c| c.into())
2716 .collect();
2717
2718 let plan = nested_table_scan("test_table")?
2719 .unnest_columns_with_options(cols, UnnestOptions::default())?
2720 .build()?;
2721
2722 assert_snapshot!(plan, @r"
2723 Unnest: lists[test_table.strings|depth=1, test_table.structs|depth=1] structs[test_table.struct_singular]
2724 TableScan: test_table
2725 ");
2726
2727 let plan = nested_table_scan("test_table")?.unnest_column("missing");
2729 assert!(plan.is_err());
2730
2731 let plan = nested_table_scan("test_table")?
2733 .unnest_columns_with_options(
2734 vec!["stringss".into(), "struct_singular".into()],
2735 UnnestOptions::default()
2736 .with_recursions(RecursionUnnestOption {
2737 input_column: "stringss".into(),
2738 output_column: "stringss_depth_1".into(),
2739 depth: 1,
2740 })
2741 .with_recursions(RecursionUnnestOption {
2742 input_column: "stringss".into(),
2743 output_column: "stringss_depth_2".into(),
2744 depth: 2,
2745 }),
2746 )?
2747 .build()?;
2748
2749 assert_snapshot!(plan, @r"
2750 Unnest: lists[test_table.stringss|depth=1, test_table.stringss|depth=2] structs[test_table.struct_singular]
2751 TableScan: test_table
2752 ");
2753
2754 let field = plan
2756 .schema()
2757 .field_with_name(None, "stringss_depth_1")
2758 .unwrap();
2759 assert_eq!(
2760 &DataType::new_list(DataType::Utf8, false),
2761 field.data_type()
2762 );
2763 let field = plan
2764 .schema()
2765 .field_with_name(None, "stringss_depth_2")
2766 .unwrap();
2767 assert_eq!(&DataType::Utf8, field.data_type());
2768 for field_name in &["a", "b"] {
2770 let field = plan
2771 .schema()
2772 .field_with_name(None, &format!("struct_singular.{field_name}"))
2773 .unwrap();
2774 assert_eq!(&DataType::UInt32, field.data_type());
2775 }
2776
2777 Ok(())
2778 }
2779
2780 fn nested_table_scan(table_name: &str) -> Result<LogicalPlanBuilder> {
2781 let struct_field_in_list = Field::new_struct(
2784 "item",
2785 vec![
2786 Field::new("a", DataType::UInt32, false),
2787 Field::new("b", DataType::UInt32, false),
2788 ],
2789 false,
2790 );
2791 let string_field = Field::new_list_field(DataType::Utf8, false);
2792 let strings_field = Field::new_list("item", string_field.clone(), false);
2793 let schema = Schema::new(vec![
2794 Field::new("scalar", DataType::UInt32, false),
2795 Field::new_list("strings", string_field, false),
2796 Field::new_list("structs", struct_field_in_list, false),
2797 Field::new(
2798 "struct_singular",
2799 DataType::Struct(Fields::from(vec![
2800 Field::new("a", DataType::UInt32, false),
2801 Field::new("b", DataType::UInt32, false),
2802 ])),
2803 false,
2804 ),
2805 Field::new_list("stringss", strings_field, false),
2806 ]);
2807
2808 table_scan(Some(table_name), &schema, None)
2809 }
2810
2811 #[test]
2812 fn test_union_after_join() -> Result<()> {
2813 let values = vec![vec![lit(1)]];
2814
2815 let left = LogicalPlanBuilder::values(values.clone())?
2816 .alias("left")?
2817 .build()?;
2818 let right = LogicalPlanBuilder::values(values)?
2819 .alias("right")?
2820 .build()?;
2821
2822 let join = LogicalPlanBuilder::from(left).cross_join(right)?.build()?;
2823
2824 let plan = LogicalPlanBuilder::from(join.clone())
2825 .union(join)?
2826 .build()?;
2827
2828 assert_snapshot!(plan, @r"
2829 Union
2830 Cross Join:
2831 SubqueryAlias: left
2832 Values: (Int32(1))
2833 SubqueryAlias: right
2834 Values: (Int32(1))
2835 Cross Join:
2836 SubqueryAlias: left
2837 Values: (Int32(1))
2838 SubqueryAlias: right
2839 Values: (Int32(1))
2840 ");
2841
2842 Ok(())
2843 }
2844
2845 #[test]
2846 fn plan_builder_from_logical_plan() -> Result<()> {
2847 let plan =
2848 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?
2849 .sort(vec![
2850 expr::Sort::new(col("state"), true, true),
2851 expr::Sort::new(col("salary"), false, false),
2852 ])?
2853 .build()?;
2854
2855 let plan_expected = format!("{plan}");
2856 let plan_builder: LogicalPlanBuilder = Arc::new(plan).into();
2857 assert_eq!(plan_expected, format!("{}", plan_builder.plan));
2858
2859 Ok(())
2860 }
2861
2862 #[test]
2863 fn plan_builder_aggregate_without_implicit_group_by_exprs() -> Result<()> {
2864 let constraints =
2865 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
2866 let table_source = table_source_with_constraints(&employee_schema(), constraints);
2867
2868 let plan =
2869 LogicalPlanBuilder::scan("employee_csv", table_source, Some(vec![0, 3, 4]))?
2870 .aggregate(vec![col("id")], vec![sum(col("salary"))])?
2871 .build()?;
2872
2873 assert_snapshot!(plan, @r"
2874 Aggregate: groupBy=[[employee_csv.id]], aggr=[[sum(employee_csv.salary)]]
2875 TableScan: employee_csv projection=[id, state, salary]
2876 ");
2877
2878 Ok(())
2879 }
2880
2881 #[test]
2882 fn plan_builder_aggregate_with_implicit_group_by_exprs() -> Result<()> {
2883 let constraints =
2884 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
2885 let table_source = table_source_with_constraints(&employee_schema(), constraints);
2886
2887 let options =
2888 LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
2889 let plan =
2890 LogicalPlanBuilder::scan("employee_csv", table_source, Some(vec![0, 3, 4]))?
2891 .with_options(options)
2892 .aggregate(vec![col("id")], vec![sum(col("salary"))])?
2893 .build()?;
2894
2895 assert_snapshot!(plan, @r"
2896 Aggregate: groupBy=[[employee_csv.id, employee_csv.state, employee_csv.salary]], aggr=[[sum(employee_csv.salary)]]
2897 TableScan: employee_csv projection=[id, state, salary]
2898 ");
2899
2900 Ok(())
2901 }
2902
2903 #[test]
2904 fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> {
2905 let err = table_scan(
2907 Some("employee_csv"),
2908 &employee_schema(),
2909 Some(vec![0, 3, 4]),
2910 )?
2911 .aggregate(vec![col("id")], vec![sum(sum(col("salary")))])
2912 .expect_err("nested aggregates should be rejected");
2913
2914 assert_snapshot!(
2915 err.strip_backtrace(),
2916 @"Error during planning: Aggregate function calls cannot be nested: 'sum(employee_csv.salary)' is nested inside 'sum(sum(employee_csv.salary))'"
2917 );
2918
2919 Ok(())
2920 }
2921
2922 #[test]
2923 fn plan_builder_window_rejects_nested_window_functions() -> Result<()> {
2924 let sum_over = |arg| {
2926 Expr::from(expr::WindowFunction::new(
2927 crate::WindowFunctionDefinition::AggregateUDF(
2928 crate::test::function_stub::sum_udaf(),
2929 ),
2930 vec![arg],
2931 ))
2932 };
2933 let err = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![4]))?
2934 .window(vec![sum_over(sum_over(col("salary")))])
2935 .expect_err("nested window functions should be rejected");
2936
2937 assert_snapshot!(
2938 err.strip_backtrace(),
2939 @"Error during planning: Window function calls cannot be nested: 'sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'"
2940 );
2941
2942 Ok(())
2943 }
2944
2945 #[test]
2946 fn test_join_metadata() -> Result<()> {
2947 let left_schema = DFSchema::new_with_metadata(
2948 vec![(None, Arc::new(Field::new("a", DataType::Int32, false)))],
2949 HashMap::from([("key".to_string(), "left".to_string())]),
2950 )?;
2951 let right_schema = DFSchema::new_with_metadata(
2952 vec![(None, Arc::new(Field::new("b", DataType::Int32, false)))],
2953 HashMap::from([("key".to_string(), "right".to_string())]),
2954 )?;
2955
2956 let join_schema =
2957 build_join_schema(&left_schema, &right_schema, &JoinType::Left)?;
2958 assert_eq!(
2959 join_schema.metadata(),
2960 &HashMap::from([("key".to_string(), "left".to_string())])
2961 );
2962 let join_schema =
2963 build_join_schema(&left_schema, &right_schema, &JoinType::Right)?;
2964 assert_eq!(
2965 join_schema.metadata(),
2966 &HashMap::from([("key".to_string(), "right".to_string())])
2967 );
2968
2969 Ok(())
2970 }
2971
2972 #[test]
2973 fn test_values_metadata() -> Result<()> {
2974 let metadata: HashMap<String, String> =
2975 [("ARROW:extension:metadata".to_string(), "test".to_string())]
2976 .into_iter()
2977 .collect();
2978 let metadata = FieldMetadata::from(metadata);
2979 let values = LogicalPlanBuilder::values(vec![
2980 vec![lit_with_metadata(1, Some(metadata.clone()))],
2981 vec![lit_with_metadata(2, Some(metadata.clone()))],
2982 ])?
2983 .build()?;
2984 assert_eq!(*values.schema().field(0).metadata(), metadata.to_hashmap());
2985
2986 let metadata2: HashMap<String, String> =
2988 [("ARROW:extension:metadata".to_string(), "test2".to_string())]
2989 .into_iter()
2990 .collect();
2991 let metadata2 = FieldMetadata::from(metadata2);
2992 assert!(
2993 LogicalPlanBuilder::values(vec![
2994 vec![lit_with_metadata(1, Some(metadata.clone()))],
2995 vec![lit_with_metadata(2, Some(metadata2.clone()))],
2996 ])
2997 .is_err()
2998 );
2999
3000 Ok(())
3001 }
3002
3003 #[test]
3004 fn test_unique_field_aliases() {
3005 let t1_field_1 = Field::new("a", DataType::Int32, false);
3006 let t2_field_1 = Field::new("a", DataType::Int32, false);
3007 let t2_field_3 = Field::new("a", DataType::Int32, false);
3008 let t2_field_4 = Field::new("a:1", DataType::Int32, false);
3009 let t1_field_2 = Field::new("b", DataType::Int32, false);
3010 let t2_field_2 = Field::new("b", DataType::Int32, false);
3011
3012 let fields = vec![
3013 t1_field_1, t2_field_1, t1_field_2, t2_field_2, t2_field_3, t2_field_4,
3014 ];
3015 let fields = Fields::from(fields);
3016
3017 let remove_redundant = unique_field_aliases(&fields);
3018
3019 assert_eq!(
3026 remove_redundant,
3027 vec![
3028 None,
3029 Some("a:1".to_string()),
3030 None,
3031 Some("b:1".to_string()),
3032 Some("a:2".to_string()),
3033 Some("a:1:1".to_string()),
3034 ]
3035 );
3036 }
3037
3038 #[test]
3039 fn test_values_with_schema_type_mismatch_error_message() {
3040 let schema = Arc::new(
3042 DFSchema::from_unqualified_fields(
3043 vec![Field::new("a", DataType::Date32, false)].into(),
3044 HashMap::new(),
3045 )
3046 .unwrap(),
3047 );
3048
3049 let err = LogicalPlanBuilder::values_with_schema(vec![vec![lit(true)]], &schema)
3050 .unwrap_err();
3051
3052 assert_eq!(
3053 err.strip_backtrace(),
3054 "Execution error: Types don't match and no valid cast exists, \
3055 received data of type Boolean for field of type Date32"
3056 );
3057 }
3058}