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, TableScan, Union, Unnest, Values,
37 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 "type mismatch and can't cast to got {} and {}",
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 for (i, row) in values.iter().enumerate() {
307 let value = &row[j];
308 let metadata = value.metadata(&schema)?;
309 if let Some(ref cm) = common_metadata {
310 if &metadata != cm {
311 return plan_err!(
312 "Inconsistent metadata across values list at row {i} column {j}. Was {:?} but found {:?}",
313 cm,
314 metadata
315 );
316 }
317 } else {
318 common_metadata = Some(metadata.clone());
319 }
320 let data_type = value.get_type(&schema)?;
321 if data_type == DataType::Null {
322 continue;
323 }
324
325 if let Some(prev_type) = common_type {
326 let data_types = vec![prev_type.clone(), data_type.clone()];
328 let Some(new_type) = type_union_resolution(&data_types) else {
329 return plan_err!(
330 "Inconsistent data type across values list at row {i} column {j}. Was {prev_type} but found {data_type}"
331 );
332 };
333 common_type = Some(new_type);
334 } else {
335 common_type = Some(data_type);
336 }
337 }
338 fields.push_with_metadata(
341 common_type.unwrap_or(DataType::Null),
342 true,
343 common_metadata,
344 );
345 }
346
347 Self::infer_inner(values, fields, &schema)
348 }
349
350 fn infer_inner(
351 mut values: Vec<Vec<Expr>>,
352 fields: ValuesFields,
353 schema: &DFSchema,
354 ) -> Result<Self> {
355 let fields = fields.into_fields();
356 for row in &mut values {
358 for (j, field_type) in fields.iter().map(|f| f.data_type()).enumerate() {
359 if let Expr::Literal(ScalarValue::Null, metadata) = &row[j] {
360 row[j] = Expr::Literal(
361 ScalarValue::try_from(field_type)?,
362 metadata.clone(),
363 );
364 } else {
365 row[j] = std::mem::take(&mut row[j]).cast_to(field_type, schema)?;
366 }
367 }
368 }
369
370 let dfschema = DFSchema::from_unqualified_fields(fields, HashMap::new())?;
371 let schema = DFSchemaRef::new(dfschema);
372
373 Ok(Self::new(LogicalPlan::Values(Values { schema, values })))
374 }
375
376 pub fn scan(
409 table_name: impl Into<TableReference>,
410 table_source: Arc<dyn TableSource>,
411 projection: Option<Vec<usize>>,
412 ) -> Result<Self> {
413 Self::scan_with_filters(table_name, table_source, projection, vec![])
414 }
415
416 pub fn copy_to(
418 input: LogicalPlan,
419 output_url: String,
420 file_type: Arc<dyn FileType>,
421 options: HashMap<String, String>,
422 partition_by: Vec<String>,
423 ) -> Result<Self> {
424 Ok(Self::new(LogicalPlan::Copy(CopyTo::new(
425 Arc::new(input),
426 output_url,
427 partition_by,
428 file_type,
429 options,
430 ))))
431 }
432
433 pub fn insert_into(
467 input: LogicalPlan,
468 table_name: impl Into<TableReference>,
469 target: Arc<dyn TableSource>,
470 insert_op: InsertOp,
471 ) -> Result<Self> {
472 Ok(Self::new(LogicalPlan::Dml(DmlStatement::new(
473 table_name.into(),
474 target,
475 WriteOp::Insert(insert_op),
476 Arc::new(input),
477 ))))
478 }
479
480 pub fn scan_with_filters(
482 table_name: impl Into<TableReference>,
483 table_source: Arc<dyn TableSource>,
484 projection: Option<Vec<usize>>,
485 filters: Vec<Expr>,
486 ) -> Result<Self> {
487 Self::scan_with_filters_inner(table_name, table_source, projection, filters, None)
488 }
489
490 pub fn scan_with_filters_fetch(
492 table_name: impl Into<TableReference>,
493 table_source: Arc<dyn TableSource>,
494 projection: Option<Vec<usize>>,
495 filters: Vec<Expr>,
496 fetch: Option<usize>,
497 ) -> Result<Self> {
498 Self::scan_with_filters_inner(
499 table_name,
500 table_source,
501 projection,
502 filters,
503 fetch,
504 )
505 }
506
507 fn scan_with_filters_inner(
508 table_name: impl Into<TableReference>,
509 table_source: Arc<dyn TableSource>,
510 projection: Option<Vec<usize>>,
511 filters: Vec<Expr>,
512 fetch: Option<usize>,
513 ) -> Result<Self> {
514 let table_scan =
515 TableScan::try_new(table_name, table_source, projection, filters, fetch)?;
516
517 if table_scan.filters.is_empty()
519 && let Some(p) = table_scan.source.get_logical_plan()
520 {
521 let sub_plan = p.into_owned();
522
523 if let Some(proj) = table_scan.projection {
524 let projection_exprs = proj
525 .into_iter()
526 .map(|i| {
527 Expr::Column(Column::from(sub_plan.schema().qualified_field(i)))
528 })
529 .collect::<Vec<_>>();
530 return Self::new(sub_plan)
531 .project(projection_exprs)?
532 .alias(table_scan.table_name);
533 }
534
535 return Self::new(sub_plan).alias(table_scan.table_name);
539 }
540
541 Ok(Self::new(LogicalPlan::TableScan(table_scan)))
542 }
543
544 pub fn window_plan(
546 input: LogicalPlan,
547 window_exprs: impl IntoIterator<Item = Expr>,
548 ) -> Result<LogicalPlan> {
549 let mut plan = input;
550 let mut groups = group_window_expr_by_sort_keys(window_exprs)?;
551 groups.sort_by(|(key_a, _), (key_b, _)| {
557 for ((first, _), (second, _)) in key_a.iter().zip(key_b.iter()) {
558 let key_ordering = compare_sort_expr(first, second, plan.schema());
559 match key_ordering {
560 Ordering::Less => {
561 return Ordering::Less;
562 }
563 Ordering::Greater => {
564 return Ordering::Greater;
565 }
566 Ordering::Equal => {}
567 }
568 }
569 key_b.len().cmp(&key_a.len())
570 });
571 for (_, exprs) in groups {
572 let window_exprs = exprs.into_iter().collect::<Vec<_>>();
573 plan = LogicalPlanBuilder::from(plan)
576 .window(window_exprs)?
577 .build()?;
578 }
579 Ok(plan)
580 }
581
582 pub fn project(
584 self,
585 expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
586 ) -> Result<Self> {
587 project(Arc::unwrap_or_clone(self.plan), expr).map(Self::new)
588 }
589
590 pub fn project_with_validation(
593 self,
594 expr: Vec<(impl Into<SelectExpr>, bool)>,
595 ) -> Result<Self> {
596 project_with_validation(Arc::unwrap_or_clone(self.plan), expr, None)
597 .map(Self::new)
598 }
599
600 pub fn project_with_validation_and_schema(
603 self,
604 expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
605 schema: &DFSchemaRef,
606 ) -> Result<Self> {
607 project_with_validation(
608 Arc::unwrap_or_clone(self.plan),
609 expr.into_iter().map(|e| (e, true)),
610 Some(schema),
611 )
612 .map(Self::new)
613 }
614
615 pub fn select(self, indices: impl IntoIterator<Item = usize>) -> Result<Self> {
617 let exprs: Vec<_> = indices
618 .into_iter()
619 .map(|x| Expr::Column(Column::from(self.plan.schema().qualified_field(x))))
620 .collect();
621 self.project(exprs)
622 }
623
624 pub fn filter(self, expr: impl Into<Expr>) -> Result<Self> {
626 let expr = normalize_col(expr.into(), &self.plan)?;
627 Filter::try_new(expr, self.plan)
628 .map(LogicalPlan::Filter)
629 .map(Self::new)
630 }
631
632 pub fn having(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::from)
638 }
639
640 pub fn prepare(self, name: String, fields: Vec<FieldRef>) -> Result<Self> {
642 Ok(Self::new(LogicalPlan::Statement(Statement::Prepare(
643 Prepare {
644 name,
645 fields,
646 input: self.plan,
647 },
648 ))))
649 }
650
651 pub fn limit(self, skip: usize, fetch: Option<usize>) -> Result<Self> {
658 let skip_expr = if skip == 0 {
659 None
660 } else {
661 Some(lit(skip as i64))
662 };
663 let fetch_expr = fetch.map(|f| lit(f as i64));
664 self.limit_by_expr(skip_expr, fetch_expr)
665 }
666
667 pub fn limit_by_expr(self, skip: Option<Expr>, fetch: Option<Expr>) -> Result<Self> {
671 Ok(Self::new(LogicalPlan::Limit(Limit {
672 skip: skip.map(Box::new),
673 fetch: fetch.map(Box::new),
674 input: self.plan,
675 })))
676 }
677
678 pub fn alias(self, alias: impl Into<TableReference>) -> Result<Self> {
680 subquery_alias(Arc::unwrap_or_clone(self.plan), alias).map(Self::new)
681 }
682
683 fn add_missing_columns(
712 curr_plan: LogicalPlan,
713 missing_cols: &IndexSet<Column>,
714 is_distinct: bool,
715 ) -> Result<LogicalPlan> {
716 match curr_plan {
717 LogicalPlan::Projection(Projection {
718 input,
719 mut expr,
720 schema: _,
721 }) if missing_cols.iter().all(|c| input.schema().has_column(c)) => {
722 let mut missing_exprs = missing_cols
723 .iter()
724 .map(|c| normalize_col(Expr::Column(c.clone()), &input))
725 .collect::<Result<Vec<_>>>()?;
726
727 missing_exprs.retain(|e| !expr.contains(e));
731 if is_distinct {
732 Self::ambiguous_distinct_check(&missing_exprs, missing_cols, &expr)?;
733 }
734 expr.extend(missing_exprs);
735 project(Arc::unwrap_or_clone(input), expr)
736 }
737 _ => {
738 let is_distinct =
739 is_distinct || matches!(curr_plan, LogicalPlan::Distinct(_));
740 let new_inputs = curr_plan
741 .inputs()
742 .into_iter()
743 .map(|input_plan| {
744 Self::add_missing_columns(
745 (*input_plan).clone(),
746 missing_cols,
747 is_distinct,
748 )
749 })
750 .collect::<Result<Vec<_>>>()?;
751 curr_plan.with_new_exprs(curr_plan.expressions(), new_inputs)
752 }
753 }
754 }
755
756 fn ambiguous_distinct_check(
757 missing_exprs: &[Expr],
758 missing_cols: &IndexSet<Column>,
759 projection_exprs: &[Expr],
760 ) -> Result<()> {
761 if missing_exprs.is_empty() {
762 return Ok(());
763 }
764
765 let all_aliases = missing_exprs.iter().all(|e| {
773 projection_exprs.iter().any(|proj_expr| {
774 if let Expr::Alias(Alias { expr, .. }) = proj_expr {
775 e == expr.as_ref()
776 } else {
777 false
778 }
779 })
780 });
781 if all_aliases {
782 return Ok(());
783 }
784
785 let missing_col_names = missing_cols
786 .iter()
787 .map(|col| col.flat_name())
788 .collect::<String>();
789
790 plan_err!(
791 "For SELECT DISTINCT, ORDER BY expressions {missing_col_names} must appear in select list"
792 )
793 }
794
795 pub fn sort_by(
797 self,
798 expr: impl IntoIterator<Item = impl Into<Expr>> + Clone,
799 ) -> Result<Self> {
800 self.sort(
801 expr.into_iter()
802 .map(|e| e.into().sort(true, false))
803 .collect::<Vec<SortExpr>>(),
804 )
805 }
806
807 pub fn sort(
808 self,
809 sorts: impl IntoIterator<Item = impl Into<SortExpr>> + Clone,
810 ) -> Result<Self> {
811 self.sort_with_limit(sorts, None)
812 }
813
814 pub fn sort_with_limit(
816 self,
817 sorts: impl IntoIterator<Item = impl Into<SortExpr>> + Clone,
818 fetch: Option<usize>,
819 ) -> Result<Self> {
820 let sorts = rewrite_sort_cols_by_aggs(sorts, &self.plan)?;
821
822 let schema = self.plan.schema();
823
824 let mut missing_cols: IndexSet<Column> = IndexSet::new();
826 sorts.iter().try_for_each::<_, Result<()>>(|sort| {
827 let columns = sort.expr.column_refs();
828
829 missing_cols.extend(
830 columns
831 .into_iter()
832 .filter(|c| !schema.has_column(c))
833 .cloned(),
834 );
835
836 Ok(())
837 })?;
838
839 if missing_cols.is_empty() {
840 return Ok(Self::new(LogicalPlan::Sort(Sort {
841 expr: normalize_sorts(sorts, &self.plan)?,
842 input: self.plan,
843 fetch,
844 })));
845 }
846
847 let new_expr = schema.columns().into_iter().map(Expr::Column).collect();
849
850 let is_distinct = false;
851 let plan = Self::add_missing_columns(
852 Arc::unwrap_or_clone(self.plan),
853 &missing_cols,
854 is_distinct,
855 )?;
856
857 let sort_plan = LogicalPlan::Sort(Sort {
858 expr: normalize_sorts(sorts, &plan)?,
859 input: Arc::new(plan),
860 fetch,
861 });
862
863 Projection::try_new(new_expr, Arc::new(sort_plan))
864 .map(LogicalPlan::Projection)
865 .map(Self::new)
866 }
867
868 pub fn union(self, plan: LogicalPlan) -> Result<Self> {
870 union(Arc::unwrap_or_clone(self.plan), plan).map(Self::new)
871 }
872
873 pub fn union_by_name(self, plan: LogicalPlan) -> Result<Self> {
875 union_by_name(Arc::unwrap_or_clone(self.plan), plan).map(Self::new)
876 }
877
878 pub fn union_by_name_distinct(self, plan: LogicalPlan) -> Result<Self> {
880 let left_plan: LogicalPlan = Arc::unwrap_or_clone(self.plan);
881 let right_plan: LogicalPlan = plan;
882
883 Ok(Self::new(LogicalPlan::Distinct(Distinct::All(Arc::new(
884 union_by_name(left_plan, right_plan)?,
885 )))))
886 }
887
888 pub fn union_distinct(self, plan: LogicalPlan) -> Result<Self> {
890 let left_plan: LogicalPlan = Arc::unwrap_or_clone(self.plan);
891 let right_plan: LogicalPlan = plan;
892
893 Ok(Self::new(LogicalPlan::Distinct(Distinct::All(Arc::new(
894 union(left_plan, right_plan)?,
895 )))))
896 }
897
898 pub fn distinct(self) -> Result<Self> {
900 Ok(Self::new(LogicalPlan::Distinct(Distinct::All(self.plan))))
901 }
902
903 pub fn distinct_on(
906 self,
907 on_expr: Vec<Expr>,
908 select_expr: Vec<Expr>,
909 sort_expr: Option<Vec<SortExpr>>,
910 ) -> Result<Self> {
911 Ok(Self::new(LogicalPlan::Distinct(Distinct::On(
912 DistinctOn::try_new(on_expr, select_expr, sort_expr, self.plan)?,
913 ))))
914 }
915
916 pub fn join(
930 self,
931 right: LogicalPlan,
932 join_type: JoinType,
933 join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
934 filter: Option<Expr>,
935 ) -> Result<Self> {
936 self.join_detailed(
937 right,
938 join_type,
939 join_keys,
940 filter,
941 NullEquality::NullEqualsNothing,
942 )
943 }
944
945 pub fn join_on(
986 self,
987 right: LogicalPlan,
988 join_type: JoinType,
989 on_exprs: impl IntoIterator<Item = Expr>,
990 ) -> Result<Self> {
991 let filter = on_exprs.into_iter().reduce(Expr::and);
992
993 self.join_detailed(
994 right,
995 join_type,
996 (Vec::<Column>::new(), Vec::<Column>::new()),
997 filter,
998 NullEquality::NullEqualsNothing,
999 )
1000 }
1001
1002 pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result<Column> {
1003 if column.relation.is_some() {
1004 return Ok(column);
1006 }
1007
1008 let schema = plan.schema();
1009 let fallback_schemas = plan.fallback_normalize_schemas();
1010 let using_columns = plan.using_columns()?;
1011 column.normalize_with_schemas_and_ambiguity_check(
1012 &[&[schema], &fallback_schemas],
1013 &using_columns,
1014 )
1015 }
1016
1017 pub fn join_detailed(
1024 self,
1025 right: LogicalPlan,
1026 join_type: JoinType,
1027 join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
1028 filter: Option<Expr>,
1029 null_equality: NullEquality,
1030 ) -> Result<Self> {
1031 self.join_detailed_with_options(
1032 right,
1033 join_type,
1034 join_keys,
1035 filter,
1036 null_equality,
1037 false,
1038 )
1039 }
1040
1041 pub fn join_detailed_with_options(
1042 self,
1043 right: LogicalPlan,
1044 join_type: JoinType,
1045 join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
1046 filter: Option<Expr>,
1047 null_equality: NullEquality,
1048 null_aware: bool,
1049 ) -> Result<Self> {
1050 if join_keys.0.len() != join_keys.1.len() {
1051 return plan_err!("left_keys and right_keys were not the same length");
1052 }
1053
1054 let filter = if let Some(expr) = filter {
1055 let filter = normalize_col_with_schemas_and_ambiguity_check(
1056 expr,
1057 &[&[self.schema(), right.schema()]],
1058 &[],
1059 )?;
1060 Some(filter)
1061 } else {
1062 None
1063 };
1064
1065 let (left_keys, right_keys): (Vec<Result<Column>>, Vec<Result<Column>>) =
1066 join_keys
1067 .0
1068 .into_iter()
1069 .zip(join_keys.1)
1070 .map(|(l, r)| {
1071 let l = l.into();
1072 let r = r.into();
1073
1074 match (&l.relation, &r.relation) {
1075 (Some(lr), Some(rr)) => {
1076 let l_is_left =
1077 self.plan.schema().field_with_qualified_name(lr, &l.name);
1078 let l_is_right =
1079 right.schema().field_with_qualified_name(lr, &l.name);
1080 let r_is_left =
1081 self.plan.schema().field_with_qualified_name(rr, &r.name);
1082 let r_is_right =
1083 right.schema().field_with_qualified_name(rr, &r.name);
1084
1085 match (l_is_left, l_is_right, r_is_left, r_is_right) {
1086 (_, Ok(_), Ok(_), _) => (Ok(r), Ok(l)),
1087 (Ok(_), _, _, Ok(_)) => (Ok(l), Ok(r)),
1088 _ => (
1089 Self::normalize(&self.plan, l),
1090 Self::normalize(&right, r),
1091 ),
1092 }
1093 }
1094 (Some(lr), None) => {
1095 let l_is_left =
1096 self.plan.schema().field_with_qualified_name(lr, &l.name);
1097 let l_is_right =
1098 right.schema().field_with_qualified_name(lr, &l.name);
1099
1100 match (l_is_left, l_is_right) {
1101 (Ok(_), _) => (Ok(l), Self::normalize(&right, r)),
1102 (_, Ok(_)) => (Self::normalize(&self.plan, r), Ok(l)),
1103 _ => (
1104 Self::normalize(&self.plan, l),
1105 Self::normalize(&right, r),
1106 ),
1107 }
1108 }
1109 (None, Some(rr)) => {
1110 let r_is_left =
1111 self.plan.schema().field_with_qualified_name(rr, &r.name);
1112 let r_is_right =
1113 right.schema().field_with_qualified_name(rr, &r.name);
1114
1115 match (r_is_left, r_is_right) {
1116 (Ok(_), _) => (Ok(r), Self::normalize(&right, l)),
1117 (_, Ok(_)) => (Self::normalize(&self.plan, l), Ok(r)),
1118 _ => (
1119 Self::normalize(&self.plan, l),
1120 Self::normalize(&right, r),
1121 ),
1122 }
1123 }
1124 (None, None) => {
1125 let mut swap = false;
1126 let left_key = Self::normalize(&self.plan, l.clone())
1127 .or_else(|_| {
1128 swap = true;
1129 Self::normalize(&right, l)
1130 });
1131 if swap {
1132 (Self::normalize(&self.plan, r), left_key)
1133 } else {
1134 (left_key, Self::normalize(&right, r))
1135 }
1136 }
1137 }
1138 })
1139 .unzip();
1140
1141 let left_keys = left_keys.into_iter().collect::<Result<Vec<Column>>>()?;
1142 let right_keys = right_keys.into_iter().collect::<Result<Vec<Column>>>()?;
1143
1144 let on: Vec<_> = left_keys
1145 .into_iter()
1146 .zip(right_keys)
1147 .map(|(l, r)| (Expr::Column(l), Expr::Column(r)))
1148 .collect();
1149 let join_schema =
1150 build_join_schema(self.plan.schema(), right.schema(), &join_type)?;
1151
1152 if join_type != JoinType::Inner && on.is_empty() && filter.is_none() {
1154 return plan_err!("join condition should not be empty");
1155 }
1156
1157 Ok(Self::new(LogicalPlan::Join(Join {
1158 left: self.plan,
1159 right: Arc::new(right),
1160 on,
1161 filter,
1162 join_type,
1163 join_constraint: JoinConstraint::On,
1164 schema: DFSchemaRef::new(join_schema),
1165 null_equality,
1166 null_aware,
1167 })))
1168 }
1169
1170 pub fn join_using(
1172 self,
1173 right: LogicalPlan,
1174 join_type: JoinType,
1175 using_keys: Vec<Column>,
1176 ) -> Result<Self> {
1177 let left_keys: Vec<Column> = using_keys
1178 .clone()
1179 .into_iter()
1180 .map(|c| Self::normalize(&self.plan, c))
1181 .collect::<Result<_>>()?;
1182 let right_keys: Vec<Column> = using_keys
1183 .into_iter()
1184 .map(|c| Self::normalize(&right, c))
1185 .collect::<Result<_>>()?;
1186
1187 let on: Vec<(_, _)> = left_keys.into_iter().zip(right_keys).collect();
1188 let mut join_on: Vec<(Expr, Expr)> = vec![];
1189 let mut filters: Option<Expr> = None;
1190 for (l, r) in &on {
1191 if self.plan.schema().has_column(l)
1192 && right.schema().has_column(r)
1193 && can_hash(
1194 datafusion_common::ExprSchema::field_from_column(
1195 self.plan.schema(),
1196 l,
1197 )?
1198 .data_type(),
1199 )
1200 {
1201 join_on.push((Expr::Column(l.clone()), Expr::Column(r.clone())));
1202 } else if self.plan.schema().has_column(l)
1203 && right.schema().has_column(r)
1204 && can_hash(
1205 datafusion_common::ExprSchema::field_from_column(
1206 self.plan.schema(),
1207 r,
1208 )?
1209 .data_type(),
1210 )
1211 {
1212 join_on.push((Expr::Column(r.clone()), Expr::Column(l.clone())));
1213 } else {
1214 let expr = binary_expr(
1215 Expr::Column(l.clone()),
1216 Operator::Eq,
1217 Expr::Column(r.clone()),
1218 );
1219 match filters {
1220 None => filters = Some(expr),
1221 Some(filter_expr) => filters = Some(and(expr, filter_expr)),
1222 }
1223 }
1224 }
1225
1226 if join_on.is_empty() {
1227 let join = Self::from(self.plan).cross_join(right)?;
1228 join.filter(filters.ok_or_else(|| {
1229 internal_datafusion_err!("filters should not be None here")
1230 })?)
1231 } else {
1232 let join = Join::try_new(
1233 self.plan,
1234 Arc::new(right),
1235 join_on,
1236 filters,
1237 join_type,
1238 JoinConstraint::Using,
1239 NullEquality::NullEqualsNothing,
1240 false, )?;
1242
1243 Ok(Self::new(LogicalPlan::Join(join)))
1244 }
1245 }
1246
1247 pub fn cross_join(self, right: LogicalPlan) -> Result<Self> {
1249 let join = Join::try_new(
1250 self.plan,
1251 Arc::new(right),
1252 vec![],
1253 None,
1254 JoinType::Inner,
1255 JoinConstraint::On,
1256 NullEquality::NullEqualsNothing,
1257 false, )?;
1259
1260 Ok(Self::new(LogicalPlan::Join(join)))
1261 }
1262
1263 pub fn repartition(self, partitioning_scheme: Partitioning) -> Result<Self> {
1265 Ok(Self::new(LogicalPlan::Repartition(Repartition {
1266 input: self.plan,
1267 partitioning_scheme,
1268 })))
1269 }
1270
1271 pub fn window(
1273 self,
1274 window_expr: impl IntoIterator<Item = impl Into<Expr>>,
1275 ) -> Result<Self> {
1276 let window_expr = normalize_cols(window_expr, &self.plan)?;
1277 validate_unique_names("Windows", &window_expr)?;
1278 Ok(Self::new(LogicalPlan::Window(Window::try_new(
1279 window_expr,
1280 self.plan,
1281 )?)))
1282 }
1283
1284 pub fn aggregate(
1288 self,
1289 group_expr: impl IntoIterator<Item = impl Into<Expr>>,
1290 aggr_expr: impl IntoIterator<Item = impl Into<Expr>>,
1291 ) -> Result<Self> {
1292 let group_expr = normalize_cols(group_expr, &self.plan)?;
1293 let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
1294
1295 let group_expr = if self.options.add_implicit_group_by_exprs {
1296 add_group_by_exprs_from_dependencies(group_expr, self.plan.schema())?
1297 } else {
1298 group_expr
1299 };
1300
1301 Aggregate::try_new(self.plan, group_expr, aggr_expr)
1302 .map(LogicalPlan::Aggregate)
1303 .map(Self::new)
1304 }
1305
1306 pub fn explain(self, verbose: bool, analyze: bool) -> Result<Self> {
1313 self.explain_option_format(
1315 ExplainOption::default()
1316 .with_verbose(verbose)
1317 .with_analyze(analyze),
1318 )
1319 }
1320
1321 pub fn explain_option_format(self, explain_option: ExplainOption) -> Result<Self> {
1325 let schema = LogicalPlan::explain_schema();
1326 let schema = schema.to_dfschema_ref()?;
1327
1328 if explain_option.analyze {
1329 Ok(Self::new(LogicalPlan::Analyze(Analyze {
1330 verbose: explain_option.verbose,
1331 input: self.plan,
1332 schema,
1333 })))
1334 } else {
1335 let stringified_plans =
1336 vec![self.plan.to_stringified(PlanType::InitialLogicalPlan)];
1337
1338 Ok(Self::new(LogicalPlan::Explain(Explain {
1339 verbose: explain_option.verbose,
1340 plan: self.plan,
1341 explain_format: explain_option.format,
1342 stringified_plans,
1343 schema,
1344 logical_optimization_succeeded: false,
1345 })))
1346 }
1347 }
1348
1349 pub fn intersect(
1351 left_plan: LogicalPlan,
1352 right_plan: LogicalPlan,
1353 is_all: bool,
1354 ) -> Result<LogicalPlan> {
1355 LogicalPlanBuilder::intersect_or_except(
1356 left_plan,
1357 right_plan,
1358 JoinType::LeftSemi,
1359 is_all,
1360 )
1361 }
1362
1363 pub fn except(
1365 left_plan: LogicalPlan,
1366 right_plan: LogicalPlan,
1367 is_all: bool,
1368 ) -> Result<LogicalPlan> {
1369 LogicalPlanBuilder::intersect_or_except(
1370 left_plan,
1371 right_plan,
1372 JoinType::LeftAnti,
1373 is_all,
1374 )
1375 }
1376
1377 fn intersect_or_except(
1379 left_plan: LogicalPlan,
1380 right_plan: LogicalPlan,
1381 join_type: JoinType,
1382 is_all: bool,
1383 ) -> Result<LogicalPlan> {
1384 let left_len = left_plan.schema().fields().len();
1385 let right_len = right_plan.schema().fields().len();
1386
1387 if left_len != right_len {
1388 return plan_err!(
1389 "INTERSECT/EXCEPT query must have the same number of columns. Left is {left_len} and right is {right_len}."
1390 );
1391 }
1392
1393 let left_builder = LogicalPlanBuilder::from(left_plan);
1396 let right_builder = LogicalPlanBuilder::from(right_plan);
1397 let (left_builder, right_builder, _requalified) =
1398 requalify_sides_if_needed(left_builder, right_builder)?;
1399 let left_plan = left_builder.build()?;
1400 let right_plan = right_builder.build()?;
1401
1402 let join_keys = left_plan
1403 .schema()
1404 .fields()
1405 .iter()
1406 .zip(right_plan.schema().fields().iter())
1407 .map(|(left_field, right_field)| {
1408 (
1409 (Column::from_name(left_field.name())),
1410 (Column::from_name(right_field.name())),
1411 )
1412 })
1413 .unzip();
1414 if is_all {
1415 LogicalPlanBuilder::from(left_plan)
1416 .join_detailed(
1417 right_plan,
1418 join_type,
1419 join_keys,
1420 None,
1421 NullEquality::NullEqualsNull,
1422 )?
1423 .build()
1424 } else {
1425 LogicalPlanBuilder::from(left_plan)
1426 .distinct()?
1427 .join_detailed(
1428 right_plan,
1429 join_type,
1430 join_keys,
1431 None,
1432 NullEquality::NullEqualsNull,
1433 )?
1434 .build()
1435 }
1436 }
1437
1438 pub fn build(self) -> Result<LogicalPlan> {
1440 Ok(Arc::unwrap_or_clone(self.plan))
1441 }
1442
1443 pub fn join_with_expr_keys(
1458 self,
1459 right: LogicalPlan,
1460 join_type: JoinType,
1461 equi_exprs: (Vec<impl Into<Expr>>, Vec<impl Into<Expr>>),
1462 filter: Option<Expr>,
1463 ) -> Result<Self> {
1464 if equi_exprs.0.len() != equi_exprs.1.len() {
1465 return plan_err!("left_keys and right_keys were not the same length");
1466 }
1467
1468 let join_key_pairs = equi_exprs
1469 .0
1470 .into_iter()
1471 .zip(equi_exprs.1)
1472 .map(|(l, r)| {
1473 let left_key = l.into();
1474 let right_key = r.into();
1475 let mut left_using_columns = HashSet::new();
1476 expr_to_columns(&left_key, &mut left_using_columns)?;
1477 let normalized_left_key = normalize_col_with_schemas_and_ambiguity_check(
1478 left_key,
1479 &[&[self.plan.schema()]],
1480 &[],
1481 )?;
1482
1483 let mut right_using_columns = HashSet::new();
1484 expr_to_columns(&right_key, &mut right_using_columns)?;
1485 let normalized_right_key = normalize_col_with_schemas_and_ambiguity_check(
1486 right_key,
1487 &[&[right.schema()]],
1488 &[],
1489 )?;
1490
1491 find_valid_equijoin_key_pair(
1493 &normalized_left_key,
1494 &normalized_right_key,
1495 self.plan.schema(),
1496 right.schema(),
1497 )?.ok_or_else(||
1498 plan_datafusion_err!(
1499 "can't create join plan, join key should belong to one input, error key: ({normalized_left_key},{normalized_right_key})"
1500 ))
1501 })
1502 .collect::<Result<Vec<_>>>()?;
1503
1504 let join = Join::try_new(
1505 self.plan,
1506 Arc::new(right),
1507 join_key_pairs,
1508 filter,
1509 join_type,
1510 JoinConstraint::On,
1511 NullEquality::NullEqualsNothing,
1512 false, )?;
1514
1515 Ok(Self::new(LogicalPlan::Join(join)))
1516 }
1517
1518 pub fn unnest_column(self, column: impl Into<Column>) -> Result<Self> {
1520 unnest(Arc::unwrap_or_clone(self.plan), vec![column.into()]).map(Self::new)
1521 }
1522
1523 pub fn unnest_column_with_options(
1525 self,
1526 column: impl Into<Column>,
1527 options: UnnestOptions,
1528 ) -> Result<Self> {
1529 unnest_with_options(
1530 Arc::unwrap_or_clone(self.plan),
1531 vec![column.into()],
1532 options,
1533 )
1534 .map(Self::new)
1535 }
1536
1537 pub fn unnest_columns_with_options(
1539 self,
1540 columns: Vec<Column>,
1541 options: UnnestOptions,
1542 ) -> Result<Self> {
1543 unnest_with_options(Arc::unwrap_or_clone(self.plan), columns, options)
1544 .map(Self::new)
1545 }
1546}
1547
1548impl From<LogicalPlan> for LogicalPlanBuilder {
1549 fn from(plan: LogicalPlan) -> Self {
1550 LogicalPlanBuilder::new(plan)
1551 }
1552}
1553
1554impl From<Arc<LogicalPlan>> for LogicalPlanBuilder {
1555 fn from(plan: Arc<LogicalPlan>) -> Self {
1556 LogicalPlanBuilder::new_from_arc(plan)
1557 }
1558}
1559
1560#[derive(Default)]
1562struct ValuesFields {
1563 inner: Vec<Field>,
1564}
1565
1566impl ValuesFields {
1567 pub fn new() -> Self {
1568 Self::default()
1569 }
1570
1571 pub fn push(&mut self, data_type: DataType, nullable: bool) {
1572 self.push_with_metadata(data_type, nullable, None);
1573 }
1574
1575 pub fn push_with_metadata(
1576 &mut self,
1577 data_type: DataType,
1578 nullable: bool,
1579 metadata: Option<FieldMetadata>,
1580 ) {
1581 let name = format!("column{}", self.inner.len() + 1);
1584 let mut field = Field::new(name, data_type, nullable);
1585 if let Some(metadata) = metadata {
1586 field.set_metadata(metadata.to_hashmap());
1587 }
1588 self.inner.push(field);
1589 }
1590
1591 pub fn into_fields(self) -> Fields {
1592 self.inner.into()
1593 }
1594}
1595
1596pub fn unique_field_aliases(fields: &Fields) -> Vec<Option<String>> {
1608 let mut name_map = HashMap::<&str, usize>::new();
1616 let mut seen = HashSet::<Cow<String>>::new();
1618
1619 fields
1620 .iter()
1621 .map(|field| {
1622 let original_name = field.name();
1623 let mut name = Cow::Borrowed(original_name);
1624
1625 let count = name_map.entry(original_name).or_insert(0);
1626
1627 while seen.contains(&name) {
1629 *count += 1;
1630 name = Cow::Owned(format!("{original_name}:{count}"));
1631 }
1632
1633 seen.insert(name.clone());
1634
1635 match name {
1636 Cow::Borrowed(_) => None,
1637 Cow::Owned(alias) => Some(alias),
1638 }
1639 })
1640 .collect()
1641}
1642
1643fn mark_field(schema: &DFSchema) -> (Option<TableReference>, Arc<Field>) {
1644 let mut table_references = schema
1645 .iter()
1646 .filter_map(|(qualifier, _)| qualifier)
1647 .collect::<Vec<_>>();
1648 table_references.dedup();
1649 let table_reference = if table_references.len() == 1 {
1650 table_references.pop().cloned()
1651 } else {
1652 None
1653 };
1654
1655 (
1656 table_reference,
1657 Arc::new(Field::new("mark", DataType::Boolean, false)),
1658 )
1659}
1660
1661pub fn build_join_schema(
1664 left: &DFSchema,
1665 right: &DFSchema,
1666 join_type: &JoinType,
1667) -> Result<DFSchema> {
1668 fn nullify_fields<'a>(
1669 fields: impl Iterator<Item = (Option<&'a TableReference>, &'a Arc<Field>)>,
1670 ) -> Vec<(Option<TableReference>, Arc<Field>)> {
1671 fields
1672 .map(|(q, f)| {
1673 let field = f.as_ref().clone().with_nullable(true);
1675 (q.cloned(), Arc::new(field))
1676 })
1677 .collect()
1678 }
1679
1680 let right_fields = right.iter();
1681 let left_fields = left.iter();
1682
1683 let qualified_fields: Vec<(Option<TableReference>, Arc<Field>)> = match join_type {
1684 JoinType::Inner => {
1685 let left_fields = left_fields
1687 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1688 .collect::<Vec<_>>();
1689 let right_fields = right_fields
1690 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1691 .collect::<Vec<_>>();
1692 left_fields.into_iter().chain(right_fields).collect()
1693 }
1694 JoinType::Left => {
1695 let left_fields = left_fields
1697 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1698 .collect::<Vec<_>>();
1699 left_fields
1700 .into_iter()
1701 .chain(nullify_fields(right_fields))
1702 .collect()
1703 }
1704 JoinType::Right => {
1705 let right_fields = right_fields
1707 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1708 .collect::<Vec<_>>();
1709 nullify_fields(left_fields)
1710 .into_iter()
1711 .chain(right_fields)
1712 .collect()
1713 }
1714 JoinType::Full => {
1715 nullify_fields(left_fields)
1717 .into_iter()
1718 .chain(nullify_fields(right_fields))
1719 .collect()
1720 }
1721 JoinType::LeftSemi | JoinType::LeftAnti => {
1722 left_fields
1724 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1725 .collect()
1726 }
1727 JoinType::LeftMark => left_fields
1728 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1729 .chain(once(mark_field(right)))
1730 .collect(),
1731 JoinType::RightSemi | JoinType::RightAnti => {
1732 right_fields
1734 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1735 .collect()
1736 }
1737 JoinType::RightMark => right_fields
1738 .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1739 .chain(once(mark_field(left)))
1740 .collect(),
1741 };
1742 let func_dependencies = left.functional_dependencies().join(
1743 right.functional_dependencies(),
1744 join_type,
1745 left.fields().len(),
1746 );
1747
1748 let (schema1, schema2) = match join_type {
1749 JoinType::Right
1750 | JoinType::RightSemi
1751 | JoinType::RightAnti
1752 | JoinType::RightMark => (left, right),
1753 _ => (right, left),
1754 };
1755
1756 let metadata = schema1
1757 .metadata()
1758 .clone()
1759 .into_iter()
1760 .chain(schema2.metadata().clone())
1761 .collect();
1762
1763 let dfschema = DFSchema::new_with_metadata(qualified_fields, metadata)?;
1764 dfschema.with_functional_dependencies(func_dependencies)
1765}
1766
1767pub fn requalify_sides_if_needed(
1777 left: LogicalPlanBuilder,
1778 right: LogicalPlanBuilder,
1779) -> Result<(LogicalPlanBuilder, LogicalPlanBuilder, bool)> {
1780 let left_cols = left.schema().columns();
1781 let right_cols = right.schema().columns();
1782
1783 for l in &left_cols {
1797 for r in &right_cols {
1798 if l.name != r.name {
1799 continue;
1800 }
1801
1802 match (&l.relation, &r.relation) {
1804 (Some(l_rel), Some(r_rel)) if l_rel == r_rel => {
1806 return Ok((
1807 left.alias(TableReference::bare("left"))?,
1808 right.alias(TableReference::bare("right"))?,
1809 true,
1810 ));
1811 }
1812 (None, None) => {
1814 return Ok((
1815 left.alias(TableReference::bare("left"))?,
1816 right.alias(TableReference::bare("right"))?,
1817 true,
1818 ));
1819 }
1820 (Some(_), None) | (None, Some(_)) => {
1822 return Ok((
1823 left.alias(TableReference::bare("left"))?,
1824 right.alias(TableReference::bare("right"))?,
1825 true,
1826 ));
1827 }
1828 _ => {}
1830 }
1831 }
1832 }
1833
1834 Ok((left, right, false))
1836}
1837pub fn add_group_by_exprs_from_dependencies(
1847 mut group_expr: Vec<Expr>,
1848 schema: &DFSchemaRef,
1849) -> Result<Vec<Expr>> {
1850 let mut group_by_field_names = group_expr
1853 .iter()
1854 .map(|e| e.schema_name().to_string())
1855 .collect::<Vec<_>>();
1856
1857 if let Some(target_indices) =
1858 get_target_functional_dependencies(schema, &group_by_field_names)
1859 {
1860 for idx in target_indices {
1861 let expr = Expr::Column(Column::from(schema.qualified_field(idx)));
1862 let expr_name = expr.schema_name().to_string();
1863 if !group_by_field_names.contains(&expr_name) {
1864 group_by_field_names.push(expr_name);
1865 group_expr.push(expr);
1866 }
1867 }
1868 }
1869 Ok(group_expr)
1870}
1871
1872pub fn validate_unique_names<'a>(
1874 node_name: &str,
1875 expressions: impl IntoIterator<Item = &'a Expr>,
1876) -> Result<()> {
1877 let mut unique_names = HashMap::new();
1878
1879 expressions.into_iter().enumerate().try_for_each(|(position, expr)| {
1880 let name = expr.schema_name().to_string();
1881 match unique_names.get(&name) {
1882 None => {
1883 unique_names.insert(name, (position, expr));
1884 Ok(())
1885 },
1886 Some((existing_position, existing_expr)) => {
1887 plan_err!("{node_name} require unique expression names \
1888 but the expression \"{existing_expr}\" at position {existing_position} and \"{expr}\" \
1889 at position {position} have the same name. Consider aliasing (\"AS\") one of them."
1890 )
1891 }
1892 }
1893 })
1894}
1895
1896pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalPlan> {
1908 Ok(LogicalPlan::Union(Union::try_new_with_loose_types(vec![
1909 Arc::new(left_plan),
1910 Arc::new(right_plan),
1911 ])?))
1912}
1913
1914pub fn union_by_name(
1917 left_plan: LogicalPlan,
1918 right_plan: LogicalPlan,
1919) -> Result<LogicalPlan> {
1920 Ok(LogicalPlan::Union(Union::try_new_by_name(vec![
1921 Arc::new(left_plan),
1922 Arc::new(right_plan),
1923 ])?))
1924}
1925
1926pub fn project(
1932 plan: LogicalPlan,
1933 expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
1934) -> Result<LogicalPlan> {
1935 project_with_validation(plan, expr.into_iter().map(|e| (e, true)), None)
1936}
1937
1938fn project_with_validation(
1946 plan: LogicalPlan,
1947 expr: impl IntoIterator<Item = (impl Into<SelectExpr>, bool)>,
1948 schema: Option<&DFSchemaRef>,
1949) -> Result<LogicalPlan> {
1950 let mut projected_expr = vec![];
1951 let mut has_wildcard = false;
1952 for (e, validate) in expr {
1953 let e = e.into();
1954 match e {
1955 SelectExpr::Wildcard(opt) => {
1956 has_wildcard = true;
1957 let expanded = expand_wildcard(plan.schema(), &plan, Some(&opt))?;
1958
1959 let expanded = if let Some(replace) = opt.replace {
1962 replace_columns(expanded, &replace)?
1963 } else {
1964 expanded
1965 };
1966
1967 for e in expanded {
1968 if validate {
1969 projected_expr
1970 .push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
1971 } else {
1972 projected_expr.push(e)
1973 }
1974 }
1975 }
1976 SelectExpr::QualifiedWildcard(table_ref, opt) => {
1977 has_wildcard = true;
1978 let expanded =
1979 expand_qualified_wildcard(&table_ref, plan.schema(), Some(&opt))?;
1980
1981 let expanded = if let Some(replace) = opt.replace {
1984 replace_columns(expanded, &replace)?
1985 } else {
1986 expanded
1987 };
1988
1989 for e in expanded {
1990 if validate {
1991 projected_expr
1992 .push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
1993 } else {
1994 projected_expr.push(e)
1995 }
1996 }
1997 }
1998 SelectExpr::Expression(e) => {
1999 if validate {
2000 projected_expr.push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
2001 } else {
2002 projected_expr.push(e)
2003 }
2004 }
2005 }
2006 }
2007
2008 if has_wildcard && projected_expr.is_empty() && !plan.schema().fields().is_empty() {
2009 return plan_err!(
2010 "SELECT list is empty after resolving * expressions, \
2011 the wildcard expanded to zero columns"
2012 );
2013 }
2014
2015 if let Some(schema) = &schema {
2018 for (expr, field) in projected_expr.iter_mut().zip(schema.fields()) {
2019 if !matches!(expr, Expr::Column(_) | Expr::Alias(_)) {
2020 *expr = std::mem::take(expr).alias(field.name());
2021 }
2022 }
2023 }
2024
2025 validate_unique_names("Projections", projected_expr.iter())?;
2026
2027 Projection::try_new(projected_expr, Arc::new(plan)).map(LogicalPlan::Projection)
2028}
2029
2030fn replace_columns(
2035 mut exprs: Vec<Expr>,
2036 replace: &PlannedReplaceSelectItem,
2037) -> Result<Vec<Expr>> {
2038 for expr in exprs.iter_mut() {
2039 if let Expr::Column(Column { name, .. }) = expr
2040 && let Some((_, new_expr)) = replace
2041 .items()
2042 .iter()
2043 .zip(replace.expressions().iter())
2044 .find(|(item, _)| item.column_name.value == *name)
2045 {
2046 *expr = new_expr.clone().alias(name.clone())
2047 }
2048 }
2049 Ok(exprs)
2050}
2051
2052pub fn subquery_alias(
2054 plan: LogicalPlan,
2055 alias: impl Into<TableReference>,
2056) -> Result<LogicalPlan> {
2057 SubqueryAlias::try_new(Arc::new(plan), alias).map(LogicalPlan::SubqueryAlias)
2058}
2059
2060pub fn table_scan(
2063 name: Option<impl Into<TableReference>>,
2064 table_schema: &Schema,
2065 projection: Option<Vec<usize>>,
2066) -> Result<LogicalPlanBuilder> {
2067 table_scan_with_filters(name, table_schema, projection, vec![])
2068}
2069
2070pub fn table_scan_with_filters(
2074 name: Option<impl Into<TableReference>>,
2075 table_schema: &Schema,
2076 projection: Option<Vec<usize>>,
2077 filters: Vec<Expr>,
2078) -> Result<LogicalPlanBuilder> {
2079 let table_source = table_source(table_schema);
2080 let name = name
2081 .map(|n| n.into())
2082 .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
2083 LogicalPlanBuilder::scan_with_filters(name, table_source, projection, filters)
2084}
2085
2086pub fn table_scan_with_filter_and_fetch(
2090 name: Option<impl Into<TableReference>>,
2091 table_schema: &Schema,
2092 projection: Option<Vec<usize>>,
2093 filters: Vec<Expr>,
2094 fetch: Option<usize>,
2095) -> Result<LogicalPlanBuilder> {
2096 let table_source = table_source(table_schema);
2097 let name = name
2098 .map(|n| n.into())
2099 .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
2100 LogicalPlanBuilder::scan_with_filters_fetch(
2101 name,
2102 table_source,
2103 projection,
2104 filters,
2105 fetch,
2106 )
2107}
2108
2109pub fn table_source(table_schema: &Schema) -> Arc<dyn TableSource> {
2110 let table_schema = Arc::new(table_schema.clone());
2112 Arc::new(LogicalTableSource {
2113 table_schema,
2114 constraints: Default::default(),
2115 })
2116}
2117
2118pub fn table_source_with_constraints(
2119 table_schema: &Schema,
2120 constraints: Constraints,
2121) -> Arc<dyn TableSource> {
2122 let table_schema = Arc::new(table_schema.clone());
2124 Arc::new(LogicalTableSource {
2125 table_schema,
2126 constraints,
2127 })
2128}
2129
2130pub fn wrap_projection_for_join_if_necessary(
2132 join_keys: &[Expr],
2133 input: LogicalPlan,
2134) -> Result<(LogicalPlan, Vec<Column>, bool)> {
2135 let input_schema = input.schema();
2136 let alias_join_keys: Vec<Expr> = join_keys
2137 .iter()
2138 .map(|key| {
2139 if matches!(key, Expr::Cast(_)) || matches!(key, Expr::TryCast(_)) {
2148 let alias = format!("{key}");
2149 key.clone().alias(alias)
2150 } else {
2151 key.clone()
2152 }
2153 })
2154 .collect::<Vec<_>>();
2155
2156 let need_project = join_keys.iter().any(|key| !matches!(key, Expr::Column(_)));
2157 let plan = if need_project {
2158 let mut projection = input_schema
2160 .columns()
2161 .into_iter()
2162 .map(Expr::Column)
2163 .collect::<Vec<_>>();
2164 #[allow(clippy::allow_attributes, clippy::mutable_key_type)]
2165 let join_key_items = alias_join_keys
2167 .iter()
2168 .flat_map(|expr| expr.try_as_col().is_none().then_some(expr))
2169 .cloned()
2170 .collect::<HashSet<Expr>>();
2171 projection.extend(join_key_items);
2172
2173 LogicalPlanBuilder::from(input)
2174 .project(projection.into_iter().map(SelectExpr::from))?
2175 .build()?
2176 } else {
2177 input
2178 };
2179
2180 let join_on = alias_join_keys
2181 .into_iter()
2182 .map(|key| {
2183 if let Some(col) = key.try_as_col() {
2184 Ok(col.clone())
2185 } else {
2186 let name = key.schema_name().to_string();
2187 Ok(Column::from_name(name))
2188 }
2189 })
2190 .collect::<Result<Vec<_>>>()?;
2191
2192 Ok((plan, join_on, need_project))
2193}
2194
2195pub struct LogicalTableSource {
2199 table_schema: SchemaRef,
2200 constraints: Constraints,
2201}
2202
2203impl LogicalTableSource {
2204 pub fn new(table_schema: SchemaRef) -> Self {
2206 Self {
2207 table_schema,
2208 constraints: Constraints::default(),
2209 }
2210 }
2211
2212 pub fn with_constraints(mut self, constraints: Constraints) -> Self {
2213 self.constraints = constraints;
2214 self
2215 }
2216}
2217
2218impl TableSource for LogicalTableSource {
2219 fn schema(&self) -> SchemaRef {
2220 Arc::clone(&self.table_schema)
2221 }
2222
2223 fn constraints(&self) -> Option<&Constraints> {
2224 Some(&self.constraints)
2225 }
2226
2227 fn supports_filters_pushdown(
2228 &self,
2229 filters: &[&Expr],
2230 ) -> Result<Vec<TableProviderFilterPushDown>> {
2231 Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
2232 }
2233}
2234
2235pub fn unnest(input: LogicalPlan, columns: Vec<Column>) -> Result<LogicalPlan> {
2237 unnest_with_options(input, columns, UnnestOptions::default())
2238}
2239
2240pub fn get_struct_unnested_columns(
2241 col_name: &String,
2242 inner_fields: &Fields,
2243) -> Vec<Column> {
2244 inner_fields
2245 .iter()
2246 .map(|f| Column::from_name(format!("{}.{}", col_name, f.name())))
2247 .collect()
2248}
2249
2250pub fn unnest_with_options(
2280 input: LogicalPlan,
2281 columns_to_unnest: Vec<Column>,
2282 options: UnnestOptions,
2283) -> Result<LogicalPlan> {
2284 Ok(LogicalPlan::Unnest(Unnest::try_new(
2285 Arc::new(input),
2286 columns_to_unnest,
2287 options,
2288 )?))
2289}
2290
2291#[cfg(test)]
2292mod tests {
2293 use std::vec;
2294
2295 use super::*;
2296 use crate::lit_with_metadata;
2297 use crate::logical_plan::StringifiedPlan;
2298 use crate::{col, expr, expr_fn::exists, in_subquery, scalar_subquery};
2299
2300 use crate::test::function_stub::sum;
2301 use datafusion_common::{
2302 Constraint, DataFusionError, RecursionUnnestOption, SchemaError,
2303 };
2304 use insta::assert_snapshot;
2305
2306 #[test]
2307 fn plan_builder_simple() -> Result<()> {
2308 let plan =
2309 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
2310 .filter(col("state").eq(lit("CO")))?
2311 .project(vec![col("id")])?
2312 .build()?;
2313
2314 assert_snapshot!(plan, @r#"
2315 Projection: employee_csv.id
2316 Filter: employee_csv.state = Utf8("CO")
2317 TableScan: employee_csv projection=[id, state]
2318 "#);
2319
2320 Ok(())
2321 }
2322
2323 #[test]
2324 fn plan_builder_schema() {
2325 let schema = employee_schema();
2326 let projection = None;
2327 let plan =
2328 LogicalPlanBuilder::scan("employee_csv", table_source(&schema), projection)
2329 .unwrap();
2330 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:{}");
2331
2332 let projection = None;
2335 let plan =
2336 LogicalPlanBuilder::scan("EMPLOYEE_CSV", table_source(&schema), projection)
2337 .unwrap();
2338 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:{}");
2339 }
2340
2341 #[test]
2342 fn plan_builder_empty_name() {
2343 let schema = employee_schema();
2344 let projection = None;
2345 let err =
2346 LogicalPlanBuilder::scan("", table_source(&schema), projection).unwrap_err();
2347 assert_snapshot!(
2348 err.strip_backtrace(),
2349 @"Error during planning: table_name cannot be empty"
2350 );
2351 }
2352
2353 #[test]
2354 fn plan_builder_sort() -> Result<()> {
2355 let plan =
2356 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?
2357 .sort(vec![
2358 expr::Sort::new(col("state"), true, true),
2359 expr::Sort::new(col("salary"), false, false),
2360 ])?
2361 .build()?;
2362
2363 assert_snapshot!(plan, @r"
2364 Sort: employee_csv.state ASC NULLS FIRST, employee_csv.salary DESC NULLS LAST
2365 TableScan: employee_csv projection=[state, salary]
2366 ");
2367
2368 Ok(())
2369 }
2370
2371 #[test]
2372 fn plan_builder_union() -> Result<()> {
2373 let plan =
2374 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?;
2375
2376 let plan = plan
2377 .clone()
2378 .union(plan.clone().build()?)?
2379 .union(plan.clone().build()?)?
2380 .union(plan.build()?)?
2381 .build()?;
2382
2383 assert_snapshot!(plan, @r"
2384 Union
2385 Union
2386 Union
2387 TableScan: employee_csv projection=[state, salary]
2388 TableScan: employee_csv projection=[state, salary]
2389 TableScan: employee_csv projection=[state, salary]
2390 TableScan: employee_csv projection=[state, salary]
2391 ");
2392
2393 Ok(())
2394 }
2395
2396 #[test]
2397 fn plan_builder_union_distinct() -> Result<()> {
2398 let plan =
2399 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?;
2400
2401 let plan = plan
2402 .clone()
2403 .union_distinct(plan.clone().build()?)?
2404 .union_distinct(plan.clone().build()?)?
2405 .union_distinct(plan.build()?)?
2406 .build()?;
2407
2408 assert_snapshot!(plan, @r"
2409 Distinct:
2410 Union
2411 Distinct:
2412 Union
2413 Distinct:
2414 Union
2415 TableScan: employee_csv projection=[state, salary]
2416 TableScan: employee_csv projection=[state, salary]
2417 TableScan: employee_csv projection=[state, salary]
2418 TableScan: employee_csv projection=[state, salary]
2419 ");
2420
2421 Ok(())
2422 }
2423
2424 #[test]
2425 fn plan_builder_simple_distinct() -> Result<()> {
2426 let plan =
2427 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
2428 .filter(col("state").eq(lit("CO")))?
2429 .project(vec![col("id")])?
2430 .distinct()?
2431 .build()?;
2432
2433 assert_snapshot!(plan, @r#"
2434 Distinct:
2435 Projection: employee_csv.id
2436 Filter: employee_csv.state = Utf8("CO")
2437 TableScan: employee_csv projection=[id, state]
2438 "#);
2439
2440 Ok(())
2441 }
2442
2443 #[test]
2444 fn exists_subquery() -> Result<()> {
2445 let foo = test_table_scan_with_name("foo")?;
2446 let bar = test_table_scan_with_name("bar")?;
2447
2448 let subquery = LogicalPlanBuilder::from(foo)
2449 .project(vec![col("a")])?
2450 .filter(col("a").eq(col("bar.a")))?
2451 .build()?;
2452
2453 let outer_query = LogicalPlanBuilder::from(bar)
2454 .project(vec![col("a")])?
2455 .filter(exists(Arc::new(subquery)))?
2456 .build()?;
2457
2458 assert_snapshot!(outer_query, @r"
2459 Filter: EXISTS (<subquery>)
2460 Subquery:
2461 Filter: foo.a = bar.a
2462 Projection: foo.a
2463 TableScan: foo
2464 Projection: bar.a
2465 TableScan: bar
2466 ");
2467
2468 Ok(())
2469 }
2470
2471 #[test]
2472 fn filter_in_subquery() -> Result<()> {
2473 let foo = test_table_scan_with_name("foo")?;
2474 let bar = test_table_scan_with_name("bar")?;
2475
2476 let subquery = LogicalPlanBuilder::from(foo)
2477 .project(vec![col("a")])?
2478 .filter(col("a").eq(col("bar.a")))?
2479 .build()?;
2480
2481 let outer_query = LogicalPlanBuilder::from(bar)
2483 .project(vec![col("a")])?
2484 .filter(in_subquery(col("a"), Arc::new(subquery)))?
2485 .build()?;
2486
2487 assert_snapshot!(outer_query, @r"
2488 Filter: bar.a IN (<subquery>)
2489 Subquery:
2490 Filter: foo.a = bar.a
2491 Projection: foo.a
2492 TableScan: foo
2493 Projection: bar.a
2494 TableScan: bar
2495 ");
2496
2497 Ok(())
2498 }
2499
2500 #[test]
2501 fn select_scalar_subquery() -> Result<()> {
2502 let foo = test_table_scan_with_name("foo")?;
2503 let bar = test_table_scan_with_name("bar")?;
2504
2505 let subquery = LogicalPlanBuilder::from(foo)
2506 .project(vec![col("b")])?
2507 .filter(col("a").eq(col("bar.a")))?
2508 .build()?;
2509
2510 let outer_query = LogicalPlanBuilder::from(bar)
2512 .project(vec![scalar_subquery(Arc::new(subquery))])?
2513 .build()?;
2514
2515 assert_snapshot!(outer_query, @r"
2516 Projection: (<subquery>)
2517 Subquery:
2518 Filter: foo.a = bar.a
2519 Projection: foo.b
2520 TableScan: foo
2521 TableScan: bar
2522 ");
2523
2524 Ok(())
2525 }
2526
2527 #[test]
2528 fn projection_non_unique_names() -> Result<()> {
2529 let plan = table_scan(
2530 Some("employee_csv"),
2531 &employee_schema(),
2532 Some(vec![0, 1]),
2534 )?
2535 .project(vec![col("id"), col("first_name").alias("id")]);
2537
2538 match plan {
2539 Err(DataFusionError::SchemaError(err, _)) => {
2540 if let SchemaError::AmbiguousReference { field } = *err {
2541 let Column {
2542 relation,
2543 name,
2544 spans: _,
2545 } = *field;
2546 let Some(TableReference::Bare { table }) = relation else {
2547 return plan_err!(
2548 "wrong relation: {relation:?}, expected table name"
2549 );
2550 };
2551 assert_eq!(*"employee_csv", *table);
2552 assert_eq!("id", &name);
2553 Ok(())
2554 } else {
2555 plan_err!("Plan should have returned an DataFusionError::SchemaError")
2556 }
2557 }
2558 _ => plan_err!("Plan should have returned an DataFusionError::SchemaError"),
2559 }
2560 }
2561
2562 fn employee_schema() -> Schema {
2563 Schema::new(vec![
2564 Field::new("id", DataType::Int32, false),
2565 Field::new("first_name", DataType::Utf8, false),
2566 Field::new("last_name", DataType::Utf8, false),
2567 Field::new("state", DataType::Utf8, false),
2568 Field::new("salary", DataType::Int32, false),
2569 ])
2570 }
2571
2572 #[test]
2573 fn stringified_plan() {
2574 let stringified_plan =
2575 StringifiedPlan::new(PlanType::InitialLogicalPlan, "...the plan...");
2576 assert!(stringified_plan.should_display(true));
2577 assert!(!stringified_plan.should_display(false)); let stringified_plan =
2580 StringifiedPlan::new(PlanType::FinalLogicalPlan, "...the plan...");
2581 assert!(stringified_plan.should_display(true));
2582 assert!(stringified_plan.should_display(false)); let stringified_plan =
2585 StringifiedPlan::new(PlanType::InitialPhysicalPlan, "...the plan...");
2586 assert!(stringified_plan.should_display(true));
2587 assert!(!stringified_plan.should_display(false)); let stringified_plan =
2590 StringifiedPlan::new(PlanType::FinalPhysicalPlan, "...the plan...");
2591 assert!(stringified_plan.should_display(true));
2592 assert!(stringified_plan.should_display(false)); let stringified_plan = StringifiedPlan::new(
2595 PlanType::OptimizedLogicalPlan {
2596 optimizer_name: "random opt pass".into(),
2597 },
2598 "...the plan...",
2599 );
2600 assert!(stringified_plan.should_display(true));
2601 assert!(!stringified_plan.should_display(false));
2602 }
2603
2604 fn test_table_scan_with_name(name: &str) -> Result<LogicalPlan> {
2605 let schema = Schema::new(vec![
2606 Field::new("a", DataType::UInt32, false),
2607 Field::new("b", DataType::UInt32, false),
2608 Field::new("c", DataType::UInt32, false),
2609 ]);
2610 table_scan(Some(name), &schema, None)?.build()
2611 }
2612
2613 #[test]
2614 fn plan_builder_intersect_different_num_columns_error() -> Result<()> {
2615 let plan1 =
2616 table_scan(TableReference::none(), &employee_schema(), Some(vec![3]))?;
2617 let plan2 =
2618 table_scan(TableReference::none(), &employee_schema(), Some(vec![3, 4]))?;
2619
2620 let err_msg1 =
2621 LogicalPlanBuilder::intersect(plan1.build()?, plan2.build()?, true)
2622 .unwrap_err();
2623
2624 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.");
2625
2626 Ok(())
2627 }
2628
2629 #[test]
2630 fn plan_builder_unnest() -> Result<()> {
2631 let err = nested_table_scan("test_table")?
2633 .unnest_column("scalar")
2634 .unwrap_err();
2635
2636 let DataFusionError::Internal(desc) = err else {
2637 return plan_err!("Plan should have returned an DataFusionError::Internal");
2638 };
2639
2640 let desc = (*desc
2641 .split(DataFusionError::BACK_TRACE_SEP)
2642 .collect::<Vec<&str>>()
2643 .first()
2644 .unwrap_or(&""))
2645 .to_string();
2646
2647 assert_snapshot!(desc, @"trying to unnest on invalid data type UInt32");
2648
2649 let plan = nested_table_scan("test_table")?
2651 .unnest_column("strings")?
2652 .build()?;
2653
2654 assert_snapshot!(plan, @r"
2655 Unnest: lists[test_table.strings|depth=1] structs[]
2656 TableScan: test_table
2657 ");
2658
2659 let field = plan.schema().field_with_name(None, "strings").unwrap();
2661 assert_eq!(&DataType::Utf8, field.data_type());
2662
2663 let plan = nested_table_scan("test_table")?
2665 .unnest_column("struct_singular")?
2666 .build()?;
2667
2668 assert_snapshot!(plan, @r"
2669 Unnest: lists[] structs[test_table.struct_singular]
2670 TableScan: test_table
2671 ");
2672
2673 for field_name in &["a", "b"] {
2674 let field = plan
2676 .schema()
2677 .field_with_name(None, &format!("struct_singular.{field_name}"))
2678 .unwrap();
2679 assert_eq!(&DataType::UInt32, field.data_type());
2680 }
2681
2682 let plan = nested_table_scan("test_table")?
2684 .unnest_column("strings")?
2685 .unnest_column("structs")?
2686 .unnest_column("struct_singular")?
2687 .build()?;
2688
2689 assert_snapshot!(plan, @r"
2690 Unnest: lists[] structs[test_table.struct_singular]
2691 Unnest: lists[test_table.structs|depth=1] structs[]
2692 Unnest: lists[test_table.strings|depth=1] structs[]
2693 TableScan: test_table
2694 ");
2695
2696 let field = plan.schema().field_with_name(None, "structs").unwrap();
2698 assert!(matches!(field.data_type(), DataType::Struct(_)));
2699
2700 let cols = vec!["strings", "structs", "struct_singular"]
2702 .into_iter()
2703 .map(|c| c.into())
2704 .collect();
2705
2706 let plan = nested_table_scan("test_table")?
2707 .unnest_columns_with_options(cols, UnnestOptions::default())?
2708 .build()?;
2709
2710 assert_snapshot!(plan, @r"
2711 Unnest: lists[test_table.strings|depth=1, test_table.structs|depth=1] structs[test_table.struct_singular]
2712 TableScan: test_table
2713 ");
2714
2715 let plan = nested_table_scan("test_table")?.unnest_column("missing");
2717 assert!(plan.is_err());
2718
2719 let plan = nested_table_scan("test_table")?
2721 .unnest_columns_with_options(
2722 vec!["stringss".into(), "struct_singular".into()],
2723 UnnestOptions::default()
2724 .with_recursions(RecursionUnnestOption {
2725 input_column: "stringss".into(),
2726 output_column: "stringss_depth_1".into(),
2727 depth: 1,
2728 })
2729 .with_recursions(RecursionUnnestOption {
2730 input_column: "stringss".into(),
2731 output_column: "stringss_depth_2".into(),
2732 depth: 2,
2733 }),
2734 )?
2735 .build()?;
2736
2737 assert_snapshot!(plan, @r"
2738 Unnest: lists[test_table.stringss|depth=1, test_table.stringss|depth=2] structs[test_table.struct_singular]
2739 TableScan: test_table
2740 ");
2741
2742 let field = plan
2744 .schema()
2745 .field_with_name(None, "stringss_depth_1")
2746 .unwrap();
2747 assert_eq!(
2748 &DataType::new_list(DataType::Utf8, false),
2749 field.data_type()
2750 );
2751 let field = plan
2752 .schema()
2753 .field_with_name(None, "stringss_depth_2")
2754 .unwrap();
2755 assert_eq!(&DataType::Utf8, field.data_type());
2756 for field_name in &["a", "b"] {
2758 let field = plan
2759 .schema()
2760 .field_with_name(None, &format!("struct_singular.{field_name}"))
2761 .unwrap();
2762 assert_eq!(&DataType::UInt32, field.data_type());
2763 }
2764
2765 Ok(())
2766 }
2767
2768 fn nested_table_scan(table_name: &str) -> Result<LogicalPlanBuilder> {
2769 let struct_field_in_list = Field::new_struct(
2772 "item",
2773 vec![
2774 Field::new("a", DataType::UInt32, false),
2775 Field::new("b", DataType::UInt32, false),
2776 ],
2777 false,
2778 );
2779 let string_field = Field::new_list_field(DataType::Utf8, false);
2780 let strings_field = Field::new_list("item", string_field.clone(), false);
2781 let schema = Schema::new(vec![
2782 Field::new("scalar", DataType::UInt32, false),
2783 Field::new_list("strings", string_field, false),
2784 Field::new_list("structs", struct_field_in_list, false),
2785 Field::new(
2786 "struct_singular",
2787 DataType::Struct(Fields::from(vec![
2788 Field::new("a", DataType::UInt32, false),
2789 Field::new("b", DataType::UInt32, false),
2790 ])),
2791 false,
2792 ),
2793 Field::new_list("stringss", strings_field, false),
2794 ]);
2795
2796 table_scan(Some(table_name), &schema, None)
2797 }
2798
2799 #[test]
2800 fn test_union_after_join() -> Result<()> {
2801 let values = vec![vec![lit(1)]];
2802
2803 let left = LogicalPlanBuilder::values(values.clone())?
2804 .alias("left")?
2805 .build()?;
2806 let right = LogicalPlanBuilder::values(values)?
2807 .alias("right")?
2808 .build()?;
2809
2810 let join = LogicalPlanBuilder::from(left).cross_join(right)?.build()?;
2811
2812 let plan = LogicalPlanBuilder::from(join.clone())
2813 .union(join)?
2814 .build()?;
2815
2816 assert_snapshot!(plan, @r"
2817 Union
2818 Cross Join:
2819 SubqueryAlias: left
2820 Values: (Int32(1))
2821 SubqueryAlias: right
2822 Values: (Int32(1))
2823 Cross Join:
2824 SubqueryAlias: left
2825 Values: (Int32(1))
2826 SubqueryAlias: right
2827 Values: (Int32(1))
2828 ");
2829
2830 Ok(())
2831 }
2832
2833 #[test]
2834 fn plan_builder_from_logical_plan() -> Result<()> {
2835 let plan =
2836 table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?
2837 .sort(vec![
2838 expr::Sort::new(col("state"), true, true),
2839 expr::Sort::new(col("salary"), false, false),
2840 ])?
2841 .build()?;
2842
2843 let plan_expected = format!("{plan}");
2844 let plan_builder: LogicalPlanBuilder = Arc::new(plan).into();
2845 assert_eq!(plan_expected, format!("{}", plan_builder.plan));
2846
2847 Ok(())
2848 }
2849
2850 #[test]
2851 fn plan_builder_aggregate_without_implicit_group_by_exprs() -> Result<()> {
2852 let constraints =
2853 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
2854 let table_source = table_source_with_constraints(&employee_schema(), constraints);
2855
2856 let plan =
2857 LogicalPlanBuilder::scan("employee_csv", table_source, Some(vec![0, 3, 4]))?
2858 .aggregate(vec![col("id")], vec![sum(col("salary"))])?
2859 .build()?;
2860
2861 assert_snapshot!(plan, @r"
2862 Aggregate: groupBy=[[employee_csv.id]], aggr=[[sum(employee_csv.salary)]]
2863 TableScan: employee_csv projection=[id, state, salary]
2864 ");
2865
2866 Ok(())
2867 }
2868
2869 #[test]
2870 fn plan_builder_aggregate_with_implicit_group_by_exprs() -> Result<()> {
2871 let constraints =
2872 Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
2873 let table_source = table_source_with_constraints(&employee_schema(), constraints);
2874
2875 let options =
2876 LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
2877 let plan =
2878 LogicalPlanBuilder::scan("employee_csv", table_source, Some(vec![0, 3, 4]))?
2879 .with_options(options)
2880 .aggregate(vec![col("id")], vec![sum(col("salary"))])?
2881 .build()?;
2882
2883 assert_snapshot!(plan, @r"
2884 Aggregate: groupBy=[[employee_csv.id, employee_csv.state, employee_csv.salary]], aggr=[[sum(employee_csv.salary)]]
2885 TableScan: employee_csv projection=[id, state, salary]
2886 ");
2887
2888 Ok(())
2889 }
2890
2891 #[test]
2892 fn test_join_metadata() -> Result<()> {
2893 let left_schema = DFSchema::new_with_metadata(
2894 vec![(None, Arc::new(Field::new("a", DataType::Int32, false)))],
2895 HashMap::from([("key".to_string(), "left".to_string())]),
2896 )?;
2897 let right_schema = DFSchema::new_with_metadata(
2898 vec![(None, Arc::new(Field::new("b", DataType::Int32, false)))],
2899 HashMap::from([("key".to_string(), "right".to_string())]),
2900 )?;
2901
2902 let join_schema =
2903 build_join_schema(&left_schema, &right_schema, &JoinType::Left)?;
2904 assert_eq!(
2905 join_schema.metadata(),
2906 &HashMap::from([("key".to_string(), "left".to_string())])
2907 );
2908 let join_schema =
2909 build_join_schema(&left_schema, &right_schema, &JoinType::Right)?;
2910 assert_eq!(
2911 join_schema.metadata(),
2912 &HashMap::from([("key".to_string(), "right".to_string())])
2913 );
2914
2915 Ok(())
2916 }
2917
2918 #[test]
2919 fn test_values_metadata() -> Result<()> {
2920 let metadata: HashMap<String, String> =
2921 [("ARROW:extension:metadata".to_string(), "test".to_string())]
2922 .into_iter()
2923 .collect();
2924 let metadata = FieldMetadata::from(metadata);
2925 let values = LogicalPlanBuilder::values(vec![
2926 vec![lit_with_metadata(1, Some(metadata.clone()))],
2927 vec![lit_with_metadata(2, Some(metadata.clone()))],
2928 ])?
2929 .build()?;
2930 assert_eq!(*values.schema().field(0).metadata(), metadata.to_hashmap());
2931
2932 let metadata2: HashMap<String, String> =
2934 [("ARROW:extension:metadata".to_string(), "test2".to_string())]
2935 .into_iter()
2936 .collect();
2937 let metadata2 = FieldMetadata::from(metadata2);
2938 assert!(
2939 LogicalPlanBuilder::values(vec![
2940 vec![lit_with_metadata(1, Some(metadata.clone()))],
2941 vec![lit_with_metadata(2, Some(metadata2.clone()))],
2942 ])
2943 .is_err()
2944 );
2945
2946 Ok(())
2947 }
2948
2949 #[test]
2950 fn test_unique_field_aliases() {
2951 let t1_field_1 = Field::new("a", DataType::Int32, false);
2952 let t2_field_1 = Field::new("a", DataType::Int32, false);
2953 let t2_field_3 = Field::new("a", DataType::Int32, false);
2954 let t2_field_4 = Field::new("a:1", DataType::Int32, false);
2955 let t1_field_2 = Field::new("b", DataType::Int32, false);
2956 let t2_field_2 = Field::new("b", DataType::Int32, false);
2957
2958 let fields = vec![
2959 t1_field_1, t2_field_1, t1_field_2, t2_field_2, t2_field_3, t2_field_4,
2960 ];
2961 let fields = Fields::from(fields);
2962
2963 let remove_redundant = unique_field_aliases(&fields);
2964
2965 assert_eq!(
2972 remove_redundant,
2973 vec![
2974 None,
2975 Some("a:1".to_string()),
2976 None,
2977 Some("b:1".to_string()),
2978 Some("a:2".to_string()),
2979 Some("a:1:1".to_string()),
2980 ]
2981 );
2982 }
2983}