1use std::cmp::Ordering;
21use std::collections::{BTreeSet, HashSet};
22use std::sync::Arc;
23
24use crate::expr::{Alias, Sort, WildcardOptions, WindowFunctionParams};
25use crate::expr_rewriter::strip_outer_reference;
26use crate::{
27 BinaryExpr, Expr, ExprSchemable, Filter, GroupingSet, LogicalPlan, Operator, and,
28};
29use datafusion_expr_common::signature::{Signature, TypeSignature};
30
31use arrow::datatypes::{DataType, Field, Schema};
32use datafusion_common::tree_node::{
33 Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
34};
35use datafusion_common::utils::get_at_indices;
36use datafusion_common::{
37 Column, DFSchema, DFSchemaRef, DataFusionError, Diagnostic, HashMap, Result, Span,
38 TableReference, internal_err, plan_datafusion_err, plan_err,
39};
40
41#[cfg(not(feature = "sql"))]
42use crate::sql::{ExceptSelectItem, ExcludeSelectItem, Ident, ObjectName};
43use indexmap::IndexSet;
44#[cfg(feature = "sql")]
45use sqlparser::ast::{ExceptSelectItem, ExcludeSelectItem, Ident, ObjectName};
46
47pub use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity;
48
49pub use datafusion_common::utils::expr::COUNT_STAR_EXPANSION;
52
53pub fn grouping_set_expr_count(group_expr: &[Expr]) -> Result<usize> {
56 if let Some(Expr::GroupingSet(grouping_set)) = group_expr.first() {
57 if group_expr.len() > 1 {
58 return plan_err!(
59 "Invalid group by expressions, GroupingSet must be the only expression"
60 );
61 }
62 Ok(grouping_set.distinct_expr().len() + 1)
64 } else {
65 grouping_set_to_exprlist(group_expr).map(|exprs| exprs.len())
66 }
67}
68
69fn powerset_indices(len: usize) -> impl Iterator<Item = Vec<usize>> {
73 (0..(1 << len)).map(move |mask| {
74 let mut indices = vec![];
75 let mut bitset = mask;
76 while bitset > 0 {
77 let rightmost: u64 = bitset & !(bitset - 1);
78 let idx = rightmost.trailing_zeros() as usize;
79 indices.push(idx);
80 bitset &= bitset - 1;
81 }
82 indices
83 })
84}
85
86pub fn powerset<T>(slice: &[T]) -> Result<Vec<Vec<&T>>> {
104 if slice.len() >= 64 {
105 return plan_err!("The size of the set must be less than 64");
106 }
107
108 Ok(powerset_indices(slice.len())
109 .map(|indices| indices.iter().map(|&idx| &slice[idx]).collect())
110 .collect())
111}
112
113fn check_grouping_set_size_limit(size: usize) -> Result<()> {
115 let max_grouping_set_size = 65535;
116 if size > max_grouping_set_size {
117 return plan_err!(
118 "The number of group_expression in grouping_set exceeds the maximum limit {max_grouping_set_size}, found {size}"
119 );
120 }
121
122 Ok(())
123}
124
125fn check_grouping_sets_size_limit(size: usize) -> Result<()> {
127 let max_grouping_sets_size = 4096;
128 if size > max_grouping_sets_size {
129 return plan_err!(
130 "The number of grouping_set in grouping_sets exceeds the maximum limit {max_grouping_sets_size}, found {size}"
131 );
132 }
133
134 Ok(())
135}
136
137fn merge_grouping_set<T: Clone>(left: &[T], right: &[T]) -> Result<Vec<T>> {
149 check_grouping_set_size_limit(left.len() + right.len())?;
150 Ok(left.iter().chain(right.iter()).cloned().collect())
151}
152
153fn cross_join_grouping_sets<T: Clone>(
166 left: &[Vec<T>],
167 right: &[Vec<T>],
168) -> Result<Vec<Vec<T>>> {
169 let grouping_sets_size = left.len() * right.len();
170
171 check_grouping_sets_size_limit(grouping_sets_size)?;
172
173 let mut result = Vec::with_capacity(grouping_sets_size);
174 for le in left {
175 for re in right {
176 result.push(merge_grouping_set(le, re)?);
177 }
178 }
179 Ok(result)
180}
181
182pub fn enumerate_grouping_sets(group_expr: Vec<Expr>) -> Result<Vec<Expr>> {
203 let has_grouping_set = group_expr
204 .iter()
205 .any(|expr| matches!(expr, Expr::GroupingSet(_)));
206 if !has_grouping_set || group_expr.len() == 1 {
207 return Ok(group_expr);
208 }
209 let partial_sets = group_expr
211 .iter()
212 .map(|expr| {
213 let exprs = match expr {
214 Expr::GroupingSet(GroupingSet::GroupingSets(grouping_sets)) => {
215 check_grouping_sets_size_limit(grouping_sets.len())?;
216 grouping_sets.iter().map(|e| e.iter().collect()).collect()
217 }
218 Expr::GroupingSet(GroupingSet::Cube(group_exprs)) => {
219 let grouping_sets = powerset(group_exprs)?;
220 check_grouping_sets_size_limit(grouping_sets.len())?;
221 grouping_sets
222 }
223 Expr::GroupingSet(GroupingSet::Rollup(group_exprs)) => {
224 let size = group_exprs.len();
225 let slice = group_exprs.as_slice();
226 check_grouping_sets_size_limit(size * (size + 1) / 2 + 1)?;
227 (0..(size + 1))
228 .map(|i| slice[0..i].iter().collect())
229 .collect()
230 }
231 expr => vec![vec![expr]],
232 };
233 Ok(exprs)
234 })
235 .collect::<Result<Vec<_>>>()?;
236
237 let grouping_sets = partial_sets
239 .into_iter()
240 .map(Ok)
241 .reduce(|l, r| cross_join_grouping_sets(&l?, &r?))
242 .transpose()?
243 .map(|e| {
244 e.into_iter()
245 .map(|e| e.into_iter().cloned().collect())
246 .collect()
247 })
248 .unwrap_or_default();
249
250 Ok(vec![Expr::GroupingSet(GroupingSet::GroupingSets(
251 grouping_sets,
252 ))])
253}
254
255pub fn grouping_set_to_exprlist(group_expr: &[Expr]) -> Result<Vec<&Expr>> {
258 if let Some(Expr::GroupingSet(grouping_set)) = group_expr.first() {
259 if group_expr.len() > 1 {
260 return plan_err!(
261 "Invalid group by expressions, GroupingSet must be the only expression"
262 );
263 }
264 Ok(grouping_set.distinct_expr())
265 } else {
266 Ok(group_expr
267 .iter()
268 .collect::<IndexSet<_>>()
269 .into_iter()
270 .collect())
271 }
272}
273
274pub fn expr_to_columns(expr: &Expr, accum: &mut HashSet<Column>) -> Result<()> {
277 expr.apply(|expr| {
278 match expr {
279 Expr::Column(qc) => {
280 accum.insert(qc.clone());
281 }
282 #[expect(deprecated)]
287 Expr::Unnest(_)
288 | Expr::ScalarVariable(_, _)
289 | Expr::Alias(_)
290 | Expr::Literal(_, _)
291 | Expr::BinaryExpr { .. }
292 | Expr::Like { .. }
293 | Expr::SimilarTo { .. }
294 | Expr::Not(_)
295 | Expr::IsNotNull(_)
296 | Expr::IsNull(_)
297 | Expr::IsTrue(_)
298 | Expr::IsFalse(_)
299 | Expr::IsUnknown(_)
300 | Expr::IsNotTrue(_)
301 | Expr::IsNotFalse(_)
302 | Expr::IsNotUnknown(_)
303 | Expr::Negative(_)
304 | Expr::Between { .. }
305 | Expr::Case { .. }
306 | Expr::Cast { .. }
307 | Expr::TryCast { .. }
308 | Expr::ScalarFunction(..)
309 | Expr::WindowFunction { .. }
310 | Expr::AggregateFunction { .. }
311 | Expr::GroupingSet(_)
312 | Expr::InList { .. }
313 | Expr::Exists { .. }
314 | Expr::InSubquery(_)
315 | Expr::SetComparison(_)
316 | Expr::ScalarSubquery(_)
317 | Expr::Wildcard { .. }
318 | Expr::Placeholder(_)
319 | Expr::OuterReferenceColumn { .. }
320 | Expr::HigherOrderFunction(_)
321 | Expr::Lambda(_)
322 | Expr::LambdaVariable(_) => {}
323 }
324 Ok(TreeNodeRecursion::Continue)
325 })
326 .map(|_| ())
327}
328
329fn get_excluded_columns(
332 opt_exclude: Option<&ExcludeSelectItem>,
333 opt_except: Option<&ExceptSelectItem>,
334 schema: &DFSchema,
335 qualifier: Option<&TableReference>,
336) -> Result<Vec<Column>> {
337 let mut idents = vec![];
338 if let Some(excepts) = opt_except {
339 idents.push(&excepts.first_element);
340 idents.extend(&excepts.additional_elements);
341 }
342 let exclude_owned: Vec<Ident>;
345 if let Some(exclude) = opt_exclude {
346 let object_name_to_ident = |name: &ObjectName| -> Result<Ident> {
347 if name.0.len() != 1 {
348 return plan_err!(
349 "EXCLUDE with multi-part identifiers is not supported: {name}"
350 );
351 }
352 let part = &name.0[0];
353 let Some(ident) = part.as_ident() else {
354 return plan_err!(
355 "EXCLUDE with non-identifier name part is not supported: {part}"
356 );
357 };
358 Ok(ident.clone())
359 };
360 exclude_owned = match exclude {
361 ExcludeSelectItem::Single(name) => vec![object_name_to_ident(name)?],
362 ExcludeSelectItem::Multiple(names) => names
363 .iter()
364 .map(object_name_to_ident)
365 .collect::<Result<Vec<_>>>()?,
366 };
367 idents.extend(exclude_owned.iter());
368 }
369 let n_elem = idents.len();
371 let unique_idents = idents.into_iter().collect::<HashSet<_>>();
372 if n_elem != unique_idents.len() {
375 return plan_err!("EXCLUDE or EXCEPT contains duplicate column names");
376 }
377
378 let mut result = vec![];
379 for ident in unique_idents.into_iter() {
380 let col_name = ident.value.as_str();
381 let (qualifier, field) = schema.qualified_field_with_name(qualifier, col_name)?;
382 result.push(Column::from((qualifier, field)));
383 }
384 Ok(result)
385}
386
387fn get_exprs_except_skipped(
389 schema: &DFSchema,
390 columns_to_skip: &HashSet<Column>,
391) -> Vec<Expr> {
392 if columns_to_skip.is_empty() {
393 schema.iter().map(Expr::from).collect::<Vec<Expr>>()
394 } else {
395 schema
396 .columns()
397 .iter()
398 .filter_map(|c| {
399 if !columns_to_skip.contains(c) {
400 Some(Expr::Column(c.clone()))
401 } else {
402 None
403 }
404 })
405 .collect::<Vec<Expr>>()
406 }
407}
408
409fn exclude_using_columns(plan: &LogicalPlan) -> Result<HashSet<Column>> {
414 let output_columns: HashSet<_> = plan.schema().columns().iter().cloned().collect();
415 let mut excluded = HashSet::new();
416 for cols in plan.using_columns()? {
417 let mut cols: Vec<_> = cols
423 .into_iter()
424 .filter(|c| output_columns.contains(c))
425 .collect();
426
427 cols.sort();
430
431 let mut seen_names = HashSet::new();
434 for col in cols {
435 if seen_names.contains(col.name.as_str()) {
436 excluded.insert(col); } else {
438 seen_names.insert(col.name.clone()); }
440 }
441 }
442 Ok(excluded)
443}
444
445pub fn expand_wildcard(
447 schema: &DFSchema,
448 plan: &LogicalPlan,
449 wildcard_options: Option<&WildcardOptions>,
450) -> Result<Vec<Expr>> {
451 let mut columns_to_skip = exclude_using_columns(plan)?;
452 let excluded_columns = if let Some(WildcardOptions {
453 exclude: opt_exclude,
454 except: opt_except,
455 ..
456 }) = wildcard_options
457 {
458 get_excluded_columns(opt_exclude.as_ref(), opt_except.as_ref(), schema, None)?
459 } else {
460 vec![]
461 };
462 columns_to_skip.extend(excluded_columns);
464 Ok(get_exprs_except_skipped(schema, &columns_to_skip))
465}
466
467pub fn expand_qualified_wildcard(
469 qualifier: &TableReference,
470 schema: &DFSchema,
471 wildcard_options: Option<&WildcardOptions>,
472) -> Result<Vec<Expr>> {
473 let qualified_indices = schema.fields_indices_with_qualified(qualifier);
474 let projected_func_dependencies = schema
475 .functional_dependencies()
476 .project_functional_dependencies(&qualified_indices, qualified_indices.len());
477 let fields_with_qualified = get_at_indices(schema.fields(), &qualified_indices)?;
478 if fields_with_qualified.is_empty() {
479 return plan_err!("Invalid qualifier {qualifier}");
480 }
481
482 let qualified_schema = Arc::new(Schema::new_with_metadata(
483 fields_with_qualified,
484 schema.metadata().clone(),
485 ));
486 let qualified_dfschema =
487 DFSchema::try_from_qualified_schema(qualifier.clone(), &qualified_schema)?
488 .with_functional_dependencies(projected_func_dependencies)?;
489 let excluded_columns = if let Some(WildcardOptions {
490 exclude: opt_exclude,
491 except: opt_except,
492 ..
493 }) = wildcard_options
494 {
495 get_excluded_columns(
496 opt_exclude.as_ref(),
497 opt_except.as_ref(),
498 schema,
499 Some(qualifier),
500 )?
501 } else {
502 vec![]
503 };
504 let mut columns_to_skip = HashSet::new();
506 columns_to_skip.extend(excluded_columns);
507 Ok(get_exprs_except_skipped(
508 &qualified_dfschema,
509 &columns_to_skip,
510 ))
511}
512
513type WindowSortKey = Vec<(Sort, bool)>;
516
517pub fn generate_sort_key(
519 partition_by: &[Expr],
520 order_by: &[Sort],
521) -> Result<WindowSortKey> {
522 let normalized_order_by_keys = order_by
523 .iter()
524 .map(|e| {
525 let Sort { expr, .. } = e;
526 Sort::new(expr.clone(), true, false)
527 })
528 .collect::<Vec<_>>();
529
530 let mut final_sort_keys = vec![];
531 let mut is_partition_flag = vec![];
532 partition_by.iter().for_each(|e| {
533 let e = e.clone().sort(true, false);
536 if let Some(pos) = normalized_order_by_keys.iter().position(|key| key.eq(&e)) {
537 let order_by_key = &order_by[pos];
538 if !final_sort_keys.contains(order_by_key) {
539 final_sort_keys.push(order_by_key.clone());
540 is_partition_flag.push(true);
541 }
542 } else if !final_sort_keys.contains(&e) {
543 final_sort_keys.push(e);
544 is_partition_flag.push(true);
545 }
546 });
547
548 order_by.iter().for_each(|e| {
549 if !final_sort_keys.contains(e) {
550 final_sort_keys.push(e.clone());
551 is_partition_flag.push(false);
552 }
553 });
554 let res = final_sort_keys
555 .into_iter()
556 .zip(is_partition_flag)
557 .collect::<Vec<_>>();
558 Ok(res)
559}
560
561pub fn compare_sort_expr(
564 sort_expr_a: &Sort,
565 sort_expr_b: &Sort,
566 schema: &DFSchemaRef,
567) -> Ordering {
568 let Sort {
569 expr: expr_a,
570 asc: asc_a,
571 nulls_first: nulls_first_a,
572 } = sort_expr_a;
573
574 let Sort {
575 expr: expr_b,
576 asc: asc_b,
577 nulls_first: nulls_first_b,
578 } = sort_expr_b;
579
580 let ref_indexes_a = find_column_indexes_referenced_by_expr(expr_a, schema);
581 let ref_indexes_b = find_column_indexes_referenced_by_expr(expr_b, schema);
582 for (idx_a, idx_b) in ref_indexes_a.iter().zip(ref_indexes_b.iter()) {
583 match idx_a.cmp(idx_b) {
584 Ordering::Less => {
585 return Ordering::Less;
586 }
587 Ordering::Greater => {
588 return Ordering::Greater;
589 }
590 Ordering::Equal => {}
591 }
592 }
593 match ref_indexes_a.len().cmp(&ref_indexes_b.len()) {
594 Ordering::Less => return Ordering::Greater,
595 Ordering::Greater => {
596 return Ordering::Less;
597 }
598 Ordering::Equal => {}
599 }
600 match (asc_a, asc_b) {
601 (true, false) => {
602 return Ordering::Greater;
603 }
604 (false, true) => {
605 return Ordering::Less;
606 }
607 _ => {}
608 }
609 match (nulls_first_a, nulls_first_b) {
610 (true, false) => {
611 return Ordering::Less;
612 }
613 (false, true) => {
614 return Ordering::Greater;
615 }
616 _ => {}
617 }
618 Ordering::Equal
619}
620
621pub fn group_window_expr_by_sort_keys(
623 window_expr: impl IntoIterator<Item = Expr>,
624) -> Result<Vec<(WindowSortKey, Vec<Expr>)>> {
625 let mut result = vec![];
626 window_expr.into_iter().try_for_each(|expr| match &expr {
627 Expr::WindowFunction(window_fun) => {
628 let WindowFunctionParams{ partition_by, order_by, ..} = &window_fun.as_ref().params;
629 let sort_key = generate_sort_key(partition_by, order_by)?;
630 if let Some((_, values)) = result.iter_mut().find(
631 |group: &&mut (WindowSortKey, Vec<Expr>)| matches!(group, (key, _) if *key == sort_key),
632 ) {
633 values.push(expr);
634 } else {
635 result.push((sort_key, vec![expr]))
636 }
637 Ok(())
638 }
639 other => internal_err!(
640 "Impossibly got non-window expr {other:?}"
641 ),
642 })?;
643 Ok(result)
644}
645
646pub fn find_aggregate_exprs<'a>(exprs: impl IntoIterator<Item = &'a Expr>) -> Vec<Expr> {
650 find_exprs_in_exprs(exprs, &|nested_expr| {
651 matches!(nested_expr, Expr::AggregateFunction { .. })
652 })
653}
654
655pub(crate) fn check_aggregate_and_window_nesting<'a>(
677 exprs: impl IntoIterator<Item = &'a Expr>,
678) -> Result<()> {
679 for expr in exprs {
680 expr.apply(|outer| {
681 if !matches!(outer, Expr::AggregateFunction(_) | Expr::WindowFunction(_)) {
682 return Ok(TreeNodeRecursion::Continue);
683 }
684
685 let mut err = None;
688 outer.apply_children(|child| {
689 child.apply(|inner| {
690 err = illegal_nesting_err(outer, inner);
691 if err.is_some() {
692 Ok(TreeNodeRecursion::Stop)
693 } else {
694 Ok(TreeNodeRecursion::Continue)
695 }
696 })
697 })?;
698
699 match err {
700 Some(err) => Err(err),
701 None => Ok(TreeNodeRecursion::Continue),
702 }
703 })?;
704 }
705 Ok(())
706}
707
708fn illegal_nesting_err(outer: &Expr, inner: &Expr) -> Option<DataFusionError> {
711 let (message, help) = match (outer, inner) {
713 (Expr::AggregateFunction(_), Expr::AggregateFunction(_)) => (
714 "Aggregate function calls cannot be nested",
715 format!("Compute '{inner}' in an inner query and aggregate its result"),
716 ),
717 (Expr::AggregateFunction(_), Expr::WindowFunction(_)) => (
718 "Aggregate function calls cannot contain window function calls",
719 format!("Compute '{inner}' in an inner query and aggregate its result"),
720 ),
721 (Expr::WindowFunction(_), Expr::WindowFunction(_)) => (
722 "Window function calls cannot be nested",
723 format!("Compute '{inner}' in an inner query and use its result here"),
724 ),
725 _ => return None,
727 };
728
729 Some(
730 plan_datafusion_err!("{message}: '{inner}' is nested inside '{outer}'")
731 .with_diagnostic(
732 Diagnostic::new_error(message, first_span(inner)).with_help(help, None),
733 ),
734 )
735}
736
737fn first_span(expr: &Expr) -> Option<Span> {
742 let mut span = None;
743 expr.apply(|e| {
744 span = e.spans().and_then(|spans| spans.first());
745 if span.is_some() {
746 Ok(TreeNodeRecursion::Stop)
747 } else {
748 Ok(TreeNodeRecursion::Continue)
749 }
750 })
751 .ok()?;
752 span
753}
754
755pub fn find_window_exprs<'a>(exprs: impl IntoIterator<Item = &'a Expr>) -> Vec<Expr> {
758 find_exprs_in_exprs(exprs, &|nested_expr| {
759 matches!(nested_expr, Expr::WindowFunction { .. })
760 })
761}
762
763pub fn find_out_reference_exprs(expr: &Expr) -> Vec<Expr> {
766 find_exprs_in_expr(expr, &|nested_expr| {
767 matches!(nested_expr, Expr::OuterReferenceColumn { .. })
768 })
769}
770
771fn find_exprs_in_exprs<'a, F>(
775 exprs: impl IntoIterator<Item = &'a Expr>,
776 test_fn: &F,
777) -> Vec<Expr>
778where
779 F: Fn(&Expr) -> bool,
780{
781 exprs
782 .into_iter()
783 .flat_map(|expr| find_exprs_in_expr(expr, test_fn))
784 .fold(vec![], |mut acc, expr| {
785 if !acc.contains(&expr) {
786 acc.push(expr)
787 }
788 acc
789 })
790}
791
792fn find_exprs_in_expr<F>(expr: &Expr, test_fn: &F) -> Vec<Expr>
796where
797 F: Fn(&Expr) -> bool,
798{
799 let mut exprs = vec![];
800 expr.apply(|expr| {
801 if test_fn(expr) {
802 if !(exprs.contains(expr)) {
803 exprs.push(expr.clone())
804 }
805 return Ok(TreeNodeRecursion::Jump);
807 }
808
809 Ok(TreeNodeRecursion::Continue)
810 })
811 .expect("no way to return error during recursion");
813 exprs
814}
815
816pub fn inspect_expr_pre<F, E>(expr: &Expr, mut f: F) -> Result<(), E>
818where
819 F: FnMut(&Expr) -> Result<(), E>,
820{
821 let mut err = Ok(());
822 expr.apply(|expr| {
823 if let Err(e) = f(expr) {
824 err = Err(e);
826 Ok(TreeNodeRecursion::Stop)
827 } else {
828 Ok(TreeNodeRecursion::Continue)
830 }
831 })
832 .expect("no way to return error during recursion");
834
835 err
836}
837
838pub fn exprlist_to_fields<'a>(
856 exprs: impl IntoIterator<Item = &'a Expr>,
857 plan: &LogicalPlan,
858) -> Result<Vec<(Option<TableReference>, Arc<Field>)>> {
859 let input_schema = plan.schema();
861 exprs
862 .into_iter()
863 .map(|e| e.to_field(input_schema))
864 .collect()
865}
866
867pub fn columnize_expr(e: Expr, input: &LogicalPlan) -> Result<Expr> {
883 let output_exprs = match input.columnized_output_exprs() {
884 Ok(exprs) if !exprs.is_empty() => exprs,
885 _ => return Ok(e),
886 };
887 let exprs_map: HashMap<&Expr, Column> = output_exprs.into_iter().collect();
888 e.transform_down(|node: Expr| match exprs_map.get(&node) {
889 Some(column) => Ok(Transformed::new(
890 Expr::Column(column.clone()),
891 true,
892 TreeNodeRecursion::Jump,
893 )),
894 None => Ok(Transformed::no(node)),
895 })
896 .data()
897}
898
899pub fn find_column_exprs(exprs: &[Expr]) -> Vec<Expr> {
902 exprs
903 .iter()
904 .flat_map(find_columns_referenced_by_expr)
905 .map(Expr::Column)
906 .collect()
907}
908
909pub(crate) fn find_columns_referenced_by_expr(e: &Expr) -> Vec<Column> {
910 let mut exprs = vec![];
911 e.apply(|expr| {
912 if let Expr::Column(c) = expr {
913 exprs.push(c.clone())
914 }
915 Ok(TreeNodeRecursion::Continue)
916 })
917 .expect("Unexpected error");
919 exprs
920}
921
922pub fn expr_as_column_expr(expr: &Expr, plan: &LogicalPlan) -> Result<Expr> {
924 match expr {
925 Expr::Column(col) => {
926 let (qualifier, field) = plan.schema().qualified_field_from_column(col)?;
927 Ok(Expr::from(Column::from((qualifier, field))))
928 }
929 _ => Ok(Expr::Column(Column::from_name(
930 expr.schema_name().to_string(),
931 ))),
932 }
933}
934
935pub(crate) fn find_column_indexes_referenced_by_expr(
938 e: &Expr,
939 schema: &DFSchemaRef,
940) -> Vec<usize> {
941 let mut indexes = vec![];
942 e.apply(|expr| {
943 match expr {
944 Expr::Column(qc) => {
945 if let Ok(idx) = schema.index_of_column(qc) {
946 indexes.push(idx);
947 }
948 }
949 Expr::Literal(_, _) => {
950 indexes.push(usize::MAX);
951 }
952 _ => {}
953 }
954 Ok(TreeNodeRecursion::Continue)
955 })
956 .unwrap();
957 indexes
958}
959
960pub fn can_hash(data_type: &DataType) -> bool {
964 match data_type {
965 DataType::Null => true,
966 DataType::Boolean => true,
967 DataType::Int8 => true,
968 DataType::Int16 => true,
969 DataType::Int32 => true,
970 DataType::Int64 => true,
971 DataType::UInt8 => true,
972 DataType::UInt16 => true,
973 DataType::UInt32 => true,
974 DataType::UInt64 => true,
975 DataType::Float16 => true,
976 DataType::Float32 => true,
977 DataType::Float64 => true,
978 DataType::Decimal32(_, _) => true,
979 DataType::Decimal64(_, _) => true,
980 DataType::Decimal128(_, _) => true,
981 DataType::Decimal256(_, _) => true,
982 DataType::Timestamp(_, _) => true,
983 DataType::Utf8 => true,
984 DataType::LargeUtf8 => true,
985 DataType::Utf8View => true,
986 DataType::Binary => true,
987 DataType::LargeBinary => true,
988 DataType::BinaryView => true,
989 DataType::Date32 => true,
990 DataType::Date64 => true,
991 DataType::Time32(_) => true,
992 DataType::Time64(_) => true,
993 DataType::Duration(_) => true,
994 DataType::Interval(_) => true,
995 DataType::FixedSizeBinary(_) => true,
996 DataType::Dictionary(key_type, value_type) => {
997 DataType::is_dictionary_key_type(key_type) && can_hash(value_type)
998 }
999 DataType::List(value_type) => can_hash(value_type.data_type()),
1000 DataType::LargeList(value_type) => can_hash(value_type.data_type()),
1001 DataType::FixedSizeList(value_type, _) => can_hash(value_type.data_type()),
1002 DataType::Map(map_struct, true | false) => can_hash(map_struct.data_type()),
1003 DataType::Struct(fields) => fields.iter().all(|f| can_hash(f.data_type())),
1004
1005 DataType::ListView(_)
1006 | DataType::LargeListView(_)
1007 | DataType::Union(_, _)
1008 | DataType::RunEndEncoded(_, _) => false,
1009 }
1010}
1011
1012pub fn check_all_columns_from_schema(
1014 columns: &HashSet<&Column>,
1015 schema: &DFSchema,
1016) -> Result<bool> {
1017 for col in columns.iter() {
1018 let exist = schema.is_column_from_schema(col);
1019 if !exist {
1020 return Ok(false);
1021 }
1022 }
1023
1024 Ok(true)
1025}
1026
1027pub fn find_valid_equijoin_key_pair(
1036 left_key: &Expr,
1037 right_key: &Expr,
1038 left_schema: &DFSchema,
1039 right_schema: &DFSchema,
1040) -> Result<Option<(Expr, Expr)>> {
1041 let left_using_columns = left_key.column_refs();
1042 let right_using_columns = right_key.column_refs();
1043
1044 if left_using_columns.is_empty() || right_using_columns.is_empty() {
1046 return Ok(None);
1047 }
1048
1049 if check_all_columns_from_schema(&left_using_columns, left_schema)?
1050 && check_all_columns_from_schema(&right_using_columns, right_schema)?
1051 {
1052 return Ok(Some((left_key.clone(), right_key.clone())));
1053 } else if check_all_columns_from_schema(&right_using_columns, left_schema)?
1054 && check_all_columns_from_schema(&left_using_columns, right_schema)?
1055 {
1056 return Ok(Some((right_key.clone(), left_key.clone())));
1057 }
1058
1059 Ok(None)
1060}
1061
1062#[expect(clippy::needless_pass_by_value)]
1074#[deprecated(since = "53.0.0", note = "Internal function")]
1075pub fn generate_signature_error_msg(
1076 func_name: &str,
1077 func_signature: Signature,
1078 input_expr_types: &[DataType],
1079) -> String {
1080 let candidate_signatures = func_signature
1081 .type_signature
1082 .to_string_repr_with_names(func_signature.parameter_names.as_deref())
1083 .iter()
1084 .map(|args_str| format!("\t{func_name}({args_str})"))
1085 .collect::<Vec<String>>()
1086 .join("\n");
1087
1088 format!(
1089 "No function matches the given name and argument types '{}({})'. You might need to add explicit type casts.\n\tCandidate functions:\n{}",
1090 func_name,
1091 TypeSignature::join_types(input_expr_types, ", "),
1092 candidate_signatures
1093 )
1094}
1095
1096pub(crate) fn generate_signature_error_message(
1108 func_name: &str,
1109 func_signature: &Signature,
1110 input_expr_types: &[DataType],
1111) -> String {
1112 #[expect(deprecated)]
1113 generate_signature_error_msg(func_name, func_signature.clone(), input_expr_types)
1114}
1115
1116pub fn split_conjunction(expr: &Expr) -> Vec<&Expr> {
1120 split_conjunction_impl(expr, vec![])
1121}
1122
1123fn split_conjunction_impl<'a>(expr: &'a Expr, mut exprs: Vec<&'a Expr>) -> Vec<&'a Expr> {
1124 match expr {
1125 Expr::BinaryExpr(BinaryExpr {
1126 right,
1127 op: Operator::And,
1128 left,
1129 }) => {
1130 let exprs = split_conjunction_impl(left, exprs);
1131 split_conjunction_impl(right, exprs)
1132 }
1133 Expr::Alias(Alias { expr, .. }) => split_conjunction_impl(expr, exprs),
1134 other => {
1135 exprs.push(other);
1136 exprs
1137 }
1138 }
1139}
1140
1141pub fn iter_conjunction(expr: &Expr) -> impl Iterator<Item = &Expr> {
1145 let mut stack = vec![expr];
1146 std::iter::from_fn(move || {
1147 while let Some(expr) = stack.pop() {
1148 match expr {
1149 Expr::BinaryExpr(BinaryExpr {
1150 right,
1151 op: Operator::And,
1152 left,
1153 }) => {
1154 stack.push(right);
1155 stack.push(left);
1156 }
1157 Expr::Alias(Alias { expr, .. }) => stack.push(expr),
1158 other => return Some(other),
1159 }
1160 }
1161 None
1162 })
1163}
1164
1165pub fn iter_conjunction_owned(expr: Expr) -> impl Iterator<Item = Expr> {
1169 let mut stack = vec![expr];
1170 std::iter::from_fn(move || {
1171 while let Some(expr) = stack.pop() {
1172 match expr {
1173 Expr::BinaryExpr(BinaryExpr {
1174 right,
1175 op: Operator::And,
1176 left,
1177 }) => {
1178 stack.push(*right);
1179 stack.push(*left);
1180 }
1181 Expr::Alias(Alias { expr, .. }) => stack.push(*expr),
1182 other => return Some(other),
1183 }
1184 }
1185 None
1186 })
1187}
1188
1189pub fn split_conjunction_owned(expr: Expr) -> Vec<Expr> {
1208 split_binary_owned(expr, Operator::And)
1209}
1210
1211pub fn split_binary_owned(expr: Expr, op: Operator) -> Vec<Expr> {
1231 split_binary_owned_impl(expr, op, vec![])
1232}
1233
1234fn split_binary_owned_impl(
1235 expr: Expr,
1236 operator: Operator,
1237 mut exprs: Vec<Expr>,
1238) -> Vec<Expr> {
1239 match expr {
1240 Expr::BinaryExpr(BinaryExpr { right, op, left }) if op == operator => {
1241 let exprs = split_binary_owned_impl(*left, operator, exprs);
1242 split_binary_owned_impl(*right, operator, exprs)
1243 }
1244 Expr::Alias(Alias { expr, .. }) => {
1245 split_binary_owned_impl(*expr, operator, exprs)
1246 }
1247 other => {
1248 exprs.push(other);
1249 exprs
1250 }
1251 }
1252}
1253
1254pub fn split_binary(expr: &Expr, op: Operator) -> Vec<&Expr> {
1258 split_binary_impl(expr, op, vec![])
1259}
1260
1261fn split_binary_impl<'a>(
1262 expr: &'a Expr,
1263 operator: Operator,
1264 mut exprs: Vec<&'a Expr>,
1265) -> Vec<&'a Expr> {
1266 match expr {
1267 Expr::BinaryExpr(BinaryExpr { right, op, left }) if *op == operator => {
1268 let exprs = split_binary_impl(left, operator, exprs);
1269 split_binary_impl(right, operator, exprs)
1270 }
1271 Expr::Alias(Alias { expr, .. }) => split_binary_impl(expr, operator, exprs),
1272 other => {
1273 exprs.push(other);
1274 exprs
1275 }
1276 }
1277}
1278
1279pub fn conjunction(filters: impl IntoIterator<Item = Expr>) -> Option<Expr> {
1299 filters.into_iter().reduce(Expr::and)
1300}
1301
1302pub fn disjunction(filters: impl IntoIterator<Item = Expr>) -> Option<Expr> {
1322 filters.into_iter().reduce(Expr::or)
1323}
1324
1325pub fn add_filter(plan: LogicalPlan, predicates: &[&Expr]) -> Result<LogicalPlan> {
1340 let predicate = predicates
1342 .iter()
1343 .skip(1)
1344 .fold(predicates[0].clone(), |acc, predicate| {
1345 and(acc, (*predicate).to_owned())
1346 });
1347
1348 Ok(LogicalPlan::Filter(Filter::try_new(
1349 predicate,
1350 Arc::new(plan),
1351 )?))
1352}
1353
1354pub fn find_join_exprs(exprs: Vec<&Expr>) -> Result<(Vec<Expr>, Vec<Expr>)> {
1365 let mut joins = vec![];
1366 let mut others = vec![];
1367 for filter in exprs.into_iter() {
1368 if filter.contains_outer() {
1370 if !matches!(filter, Expr::BinaryExpr(BinaryExpr{ left, op: Operator::Eq, right }) if left.eq(right))
1371 {
1372 joins.push(strip_outer_reference((*filter).clone()));
1373 }
1374 } else {
1375 others.push((*filter).clone());
1376 }
1377 }
1378
1379 Ok((joins, others))
1380}
1381
1382pub fn only_or_err<T>(slice: &[T]) -> Result<&T> {
1392 match slice {
1393 [it] => Ok(it),
1394 [] => plan_err!("No items found!"),
1395 _ => plan_err!("More than one item found!"),
1396 }
1397}
1398
1399pub fn merge_schema(inputs: &[&LogicalPlan]) -> DFSchema {
1404 if inputs.len() == 1 {
1405 inputs[0].schema().as_ref().clone()
1406 } else {
1407 inputs.iter().map(|input| input.schema()).fold(
1408 DFSchema::empty(),
1409 |mut lhs, rhs| {
1410 lhs.merge(rhs);
1411 lhs
1412 },
1413 )
1414 }
1415}
1416
1417pub fn format_state_name(name: &str, state_name: &str) -> String {
1419 format!("{name}[{state_name}]")
1420}
1421
1422pub fn collect_subquery_cols(
1424 exprs: &[Expr],
1425 subquery_schema: &DFSchema,
1426) -> Result<BTreeSet<Column>> {
1427 exprs.iter().try_fold(BTreeSet::new(), |mut cols, expr| {
1428 let mut using_cols: Vec<Column> = vec![];
1429 for col in expr.column_refs().into_iter() {
1430 if subquery_schema.has_column(col) {
1431 using_cols.push(col.clone());
1432 }
1433 }
1434
1435 cols.extend(using_cols);
1436 Result::<_>::Ok(cols)
1437 })
1438}
1439
1440#[cfg(test)]
1441mod tests {
1442 use super::*;
1443 use crate::{
1444 Cast, ExprFunctionExt, WindowFunctionDefinition, col, cube,
1445 expr::WindowFunction,
1446 expr_vec_fmt, grouping_set, lit, rollup,
1447 test::function_stub::{max_udaf, min_udaf, sum_udaf},
1448 };
1449 use arrow::datatypes::{UnionFields, UnionMode};
1450 use datafusion_expr_common::signature::Volatility;
1451
1452 #[test]
1453 fn test_group_window_expr_by_sort_keys_empty_case() -> Result<()> {
1454 let result = group_window_expr_by_sort_keys(vec![])?;
1455 let expected: Vec<(WindowSortKey, Vec<Expr>)> = vec![];
1456 assert_eq!(expected, result);
1457 Ok(())
1458 }
1459
1460 #[test]
1461 fn test_group_window_expr_by_sort_keys_empty_window() -> Result<()> {
1462 let max1 = Expr::from(WindowFunction::new(
1463 WindowFunctionDefinition::AggregateUDF(max_udaf()),
1464 vec![col("name")],
1465 ));
1466 let max2 = Expr::from(WindowFunction::new(
1467 WindowFunctionDefinition::AggregateUDF(max_udaf()),
1468 vec![col("name")],
1469 ));
1470 let min3 = Expr::from(WindowFunction::new(
1471 WindowFunctionDefinition::AggregateUDF(min_udaf()),
1472 vec![col("name")],
1473 ));
1474 let sum4 = Expr::from(WindowFunction::new(
1475 WindowFunctionDefinition::AggregateUDF(sum_udaf()),
1476 vec![col("age")],
1477 ));
1478 let exprs = &[max1.clone(), max2.clone(), min3.clone(), sum4.clone()];
1479 let result = group_window_expr_by_sort_keys(exprs.to_vec())?;
1480 let key = vec![];
1481 let expected: Vec<(WindowSortKey, Vec<Expr>)> =
1482 vec![(key, vec![max1, max2, min3, sum4])];
1483 assert_eq!(expected, result);
1484 Ok(())
1485 }
1486
1487 #[test]
1488 fn test_group_window_expr_by_sort_keys() -> Result<()> {
1489 let age_asc = Sort::new(col("age"), true, true);
1490 let name_desc = Sort::new(col("name"), false, true);
1491 let created_at_desc = Sort::new(col("created_at"), false, true);
1492 let max1 = Expr::from(WindowFunction::new(
1493 WindowFunctionDefinition::AggregateUDF(max_udaf()),
1494 vec![col("name")],
1495 ))
1496 .order_by(vec![age_asc.clone(), name_desc.clone()])
1497 .build()
1498 .unwrap();
1499 let max2 = Expr::from(WindowFunction::new(
1500 WindowFunctionDefinition::AggregateUDF(max_udaf()),
1501 vec![col("name")],
1502 ));
1503 let min3 = Expr::from(WindowFunction::new(
1504 WindowFunctionDefinition::AggregateUDF(min_udaf()),
1505 vec![col("name")],
1506 ))
1507 .order_by(vec![age_asc.clone(), name_desc.clone()])
1508 .build()
1509 .unwrap();
1510 let sum4 = Expr::from(WindowFunction::new(
1511 WindowFunctionDefinition::AggregateUDF(sum_udaf()),
1512 vec![col("age")],
1513 ))
1514 .order_by(vec![
1515 name_desc.clone(),
1516 age_asc.clone(),
1517 created_at_desc.clone(),
1518 ])
1519 .build()
1520 .unwrap();
1521 let exprs = &[max1.clone(), max2.clone(), min3.clone(), sum4.clone()];
1523 let result = group_window_expr_by_sort_keys(exprs.to_vec())?;
1524
1525 let key1 = vec![(age_asc.clone(), false), (name_desc.clone(), false)];
1526 let key2 = vec![];
1527 let key3 = vec![
1528 (name_desc, false),
1529 (age_asc, false),
1530 (created_at_desc, false),
1531 ];
1532
1533 let expected: Vec<(WindowSortKey, Vec<Expr>)> = vec![
1534 (key1, vec![max1, min3]),
1535 (key2, vec![max2]),
1536 (key3, vec![sum4]),
1537 ];
1538 assert_eq!(expected, result);
1539 Ok(())
1540 }
1541
1542 #[test]
1543 fn avoid_generate_duplicate_sort_keys() -> Result<()> {
1544 let asc_or_desc = [true, false];
1545 let nulls_first_or_last = [true, false];
1546 let partition_by = &[col("age"), col("name"), col("created_at")];
1547 for asc_ in asc_or_desc {
1548 for nulls_first_ in nulls_first_or_last {
1549 let order_by = &[
1550 Sort {
1551 expr: col("age"),
1552 asc: asc_,
1553 nulls_first: nulls_first_,
1554 },
1555 Sort {
1556 expr: col("name"),
1557 asc: asc_,
1558 nulls_first: nulls_first_,
1559 },
1560 ];
1561
1562 let expected = vec![
1563 (
1564 Sort {
1565 expr: col("age"),
1566 asc: asc_,
1567 nulls_first: nulls_first_,
1568 },
1569 true,
1570 ),
1571 (
1572 Sort {
1573 expr: col("name"),
1574 asc: asc_,
1575 nulls_first: nulls_first_,
1576 },
1577 true,
1578 ),
1579 (
1580 Sort {
1581 expr: col("created_at"),
1582 asc: true,
1583 nulls_first: false,
1584 },
1585 true,
1586 ),
1587 ];
1588 let result = generate_sort_key(partition_by, order_by)?;
1589 assert_eq!(expected, result);
1590 }
1591 }
1592 Ok(())
1593 }
1594
1595 #[test]
1596 fn test_enumerate_grouping_sets() -> Result<()> {
1597 let multi_cols = vec![col("col1"), col("col2"), col("col3")];
1598 let simple_col = col("simple_col");
1599 let cube = cube(multi_cols.clone());
1600 let rollup = rollup(multi_cols.clone());
1601 let grouping_set = grouping_set(vec![multi_cols]);
1602
1603 let sets = enumerate_grouping_sets(vec![simple_col.clone()])?;
1605 let result = format!("[{}]", expr_vec_fmt!(sets));
1606 assert_eq!("[simple_col]", &result);
1607
1608 let sets = enumerate_grouping_sets(vec![cube.clone()])?;
1610 let result = format!("[{}]", expr_vec_fmt!(sets));
1611 assert_eq!("[CUBE (col1, col2, col3)]", &result);
1612
1613 let sets = enumerate_grouping_sets(vec![rollup.clone()])?;
1615 let result = format!("[{}]", expr_vec_fmt!(sets));
1616 assert_eq!("[ROLLUP (col1, col2, col3)]", &result);
1617
1618 let sets = enumerate_grouping_sets(vec![simple_col.clone(), cube.clone()])?;
1620 let result = format!("[{}]", expr_vec_fmt!(sets));
1621 assert_eq!(
1622 "[GROUPING SETS (\
1623 (simple_col), \
1624 (simple_col, col1), \
1625 (simple_col, col2), \
1626 (simple_col, col1, col2), \
1627 (simple_col, col3), \
1628 (simple_col, col1, col3), \
1629 (simple_col, col2, col3), \
1630 (simple_col, col1, col2, col3))]",
1631 &result
1632 );
1633
1634 let sets = enumerate_grouping_sets(vec![simple_col.clone(), rollup.clone()])?;
1636 let result = format!("[{}]", expr_vec_fmt!(sets));
1637 assert_eq!(
1638 "[GROUPING SETS (\
1639 (simple_col), \
1640 (simple_col, col1), \
1641 (simple_col, col1, col2), \
1642 (simple_col, col1, col2, col3))]",
1643 &result
1644 );
1645
1646 let sets =
1648 enumerate_grouping_sets(vec![simple_col.clone(), grouping_set.clone()])?;
1649 let result = format!("[{}]", expr_vec_fmt!(sets));
1650 assert_eq!(
1651 "[GROUPING SETS (\
1652 (simple_col, col1, col2, col3))]",
1653 &result
1654 );
1655
1656 let sets = enumerate_grouping_sets(vec![
1658 simple_col.clone(),
1659 grouping_set,
1660 rollup.clone(),
1661 ])?;
1662 let result = format!("[{}]", expr_vec_fmt!(sets));
1663 assert_eq!(
1664 "[GROUPING SETS (\
1665 (simple_col, col1, col2, col3), \
1666 (simple_col, col1, col2, col3, col1), \
1667 (simple_col, col1, col2, col3, col1, col2), \
1668 (simple_col, col1, col2, col3, col1, col2, col3))]",
1669 &result
1670 );
1671
1672 let sets = enumerate_grouping_sets(vec![simple_col, cube, rollup])?;
1674 let result = format!("[{}]", expr_vec_fmt!(sets));
1675 assert_eq!(
1676 "[GROUPING SETS (\
1677 (simple_col), \
1678 (simple_col, col1), \
1679 (simple_col, col1, col2), \
1680 (simple_col, col1, col2, col3), \
1681 (simple_col, col1), \
1682 (simple_col, col1, col1), \
1683 (simple_col, col1, col1, col2), \
1684 (simple_col, col1, col1, col2, col3), \
1685 (simple_col, col2), \
1686 (simple_col, col2, col1), \
1687 (simple_col, col2, col1, col2), \
1688 (simple_col, col2, col1, col2, col3), \
1689 (simple_col, col1, col2), \
1690 (simple_col, col1, col2, col1), \
1691 (simple_col, col1, col2, col1, col2), \
1692 (simple_col, col1, col2, col1, col2, col3), \
1693 (simple_col, col3), \
1694 (simple_col, col3, col1), \
1695 (simple_col, col3, col1, col2), \
1696 (simple_col, col3, col1, col2, col3), \
1697 (simple_col, col1, col3), \
1698 (simple_col, col1, col3, col1), \
1699 (simple_col, col1, col3, col1, col2), \
1700 (simple_col, col1, col3, col1, col2, col3), \
1701 (simple_col, col2, col3), \
1702 (simple_col, col2, col3, col1), \
1703 (simple_col, col2, col3, col1, col2), \
1704 (simple_col, col2, col3, col1, col2, col3), \
1705 (simple_col, col1, col2, col3), \
1706 (simple_col, col1, col2, col3, col1), \
1707 (simple_col, col1, col2, col3, col1, col2), \
1708 (simple_col, col1, col2, col3, col1, col2, col3))]",
1709 &result
1710 );
1711
1712 Ok(())
1713 }
1714 #[test]
1715 fn test_split_conjunction() {
1716 let expr = col("a");
1717 let result = split_conjunction(&expr);
1718 assert_eq!(result, vec![&expr]);
1719 }
1720
1721 #[test]
1722 fn test_split_conjunction_two() {
1723 let expr = col("a").eq(lit(5)).and(col("b"));
1724 let expr1 = col("a").eq(lit(5));
1725 let expr2 = col("b");
1726
1727 let result = split_conjunction(&expr);
1728 assert_eq!(result, vec![&expr1, &expr2]);
1729 }
1730
1731 #[test]
1732 fn test_split_conjunction_alias() {
1733 let expr = col("a").eq(lit(5)).and(col("b").alias("the_alias"));
1734 let expr1 = col("a").eq(lit(5));
1735 let expr2 = col("b"); let result = split_conjunction(&expr);
1738 assert_eq!(result, vec![&expr1, &expr2]);
1739 }
1740
1741 #[test]
1742 fn test_split_conjunction_or() {
1743 let expr = col("a").eq(lit(5)).or(col("b"));
1744 let result = split_conjunction(&expr);
1745 assert_eq!(result, vec![&expr]);
1746 }
1747
1748 #[test]
1749 fn test_split_binary_owned() {
1750 let expr = col("a");
1751 assert_eq!(split_binary_owned(expr.clone(), Operator::And), vec![expr]);
1752 }
1753
1754 #[test]
1755 fn test_split_binary_owned_two() {
1756 assert_eq!(
1757 split_binary_owned(col("a").eq(lit(5)).and(col("b")), Operator::And),
1758 vec![col("a").eq(lit(5)), col("b")]
1759 );
1760 }
1761
1762 #[test]
1763 fn test_split_binary_owned_different_op() {
1764 let expr = col("a").eq(lit(5)).or(col("b"));
1765 assert_eq!(
1766 split_binary_owned(expr.clone(), Operator::And),
1768 vec![expr]
1769 );
1770 }
1771
1772 #[test]
1773 fn test_split_conjunction_owned() {
1774 let expr = col("a");
1775 assert_eq!(split_conjunction_owned(expr.clone()), vec![expr]);
1776 }
1777
1778 #[test]
1779 fn test_split_conjunction_owned_two() {
1780 assert_eq!(
1781 split_conjunction_owned(col("a").eq(lit(5)).and(col("b"))),
1782 vec![col("a").eq(lit(5)), col("b")]
1783 );
1784 }
1785
1786 #[test]
1787 fn test_split_conjunction_owned_alias() {
1788 assert_eq!(
1789 split_conjunction_owned(col("a").eq(lit(5)).and(col("b").alias("the_alias"))),
1790 vec![
1791 col("a").eq(lit(5)),
1792 col("b"),
1794 ]
1795 );
1796 }
1797
1798 #[test]
1799 fn test_conjunction_empty() {
1800 assert_eq!(conjunction(vec![]), None);
1801 }
1802
1803 #[test]
1804 fn test_conjunction() {
1805 let expr = conjunction(vec![col("a"), col("b"), col("c")]);
1807
1808 assert_eq!(expr, Some(col("a").and(col("b")).and(col("c"))));
1810
1811 assert_ne!(expr, Some(col("a").and(col("b").and(col("c")))));
1813 }
1814
1815 #[test]
1816 fn test_disjunction_empty() {
1817 assert_eq!(disjunction(vec![]), None);
1818 }
1819
1820 #[test]
1821 fn test_disjunction() {
1822 let expr = disjunction(vec![col("a"), col("b"), col("c")]);
1824
1825 assert_eq!(expr, Some(col("a").or(col("b")).or(col("c"))));
1827
1828 assert_ne!(expr, Some(col("a").or(col("b").or(col("c")))));
1830 }
1831
1832 #[test]
1833 fn test_split_conjunction_owned_or() {
1834 let expr = col("a").eq(lit(5)).or(col("b"));
1835 assert_eq!(split_conjunction_owned(expr.clone()), vec![expr]);
1836 }
1837
1838 #[test]
1839 fn test_collect_expr() -> Result<()> {
1840 let mut accum: HashSet<Column> = HashSet::new();
1841 expr_to_columns(
1842 &Expr::Cast(Cast::new(Box::new(col("a")), DataType::Float64)),
1843 &mut accum,
1844 )?;
1845 expr_to_columns(
1846 &Expr::Cast(Cast::new(Box::new(col("a")), DataType::Float64)),
1847 &mut accum,
1848 )?;
1849 assert_eq!(1, accum.len());
1850 assert!(accum.contains(&Column::from_name("a")));
1851 Ok(())
1852 }
1853
1854 #[test]
1855 fn test_can_hash() {
1856 let union_fields: UnionFields = [
1857 (0, Arc::new(Field::new("A", DataType::Int32, true))),
1858 (1, Arc::new(Field::new("B", DataType::Float64, true))),
1859 ]
1860 .into_iter()
1861 .collect();
1862
1863 let union_type = DataType::Union(union_fields, UnionMode::Sparse);
1864 assert!(!can_hash(&union_type));
1865
1866 let list_union_type =
1867 DataType::List(Arc::new(Field::new("my_union", union_type, true)));
1868 assert!(!can_hash(&list_union_type));
1869 }
1870
1871 #[test]
1872 fn test_generate_signature_error_msg_with_parameter_names() {
1873 let sig = Signature::one_of(
1874 vec![
1875 TypeSignature::Exact(vec![DataType::Utf8, DataType::Int64]),
1876 TypeSignature::Exact(vec![
1877 DataType::Utf8,
1878 DataType::Int64,
1879 DataType::Int64,
1880 ]),
1881 ],
1882 Volatility::Immutable,
1883 )
1884 .with_parameter_names(vec![
1885 "str".to_string(),
1886 "start_pos".to_string(),
1887 "length".to_string(),
1888 ])
1889 .expect("valid parameter names");
1890
1891 let error_msg =
1893 generate_signature_error_message("substr", &sig, &[DataType::Utf8]);
1894
1895 assert!(
1896 error_msg.contains("str: Utf8, start_pos: Int64"),
1897 "Expected 'str: Utf8, start_pos: Int64' in error message, got: {error_msg}"
1898 );
1899 assert!(
1900 error_msg.contains("str: Utf8, start_pos: Int64, length: Int64"),
1901 "Expected 'str: Utf8, start_pos: Int64, length: Int64' in error message, got: {error_msg}"
1902 );
1903 }
1904
1905 #[test]
1906 fn test_generate_signature_error_msg_without_parameter_names() {
1907 let sig = Signature::one_of(
1908 vec![TypeSignature::Any(2), TypeSignature::Any(3)],
1909 Volatility::Immutable,
1910 );
1911
1912 let error_msg =
1913 generate_signature_error_message("my_func", &sig, &[DataType::Int32]);
1914
1915 assert!(
1916 error_msg.contains("Any, Any"),
1917 "Expected 'Any, Any' without parameter names, got: {error_msg}"
1918 );
1919 }
1920
1921 #[test]
1922 fn test_signature_error_msg_exact() {
1923 use insta::assert_snapshot;
1924
1925 let sig = Signature::one_of(
1926 vec![
1927 TypeSignature::Exact(vec![DataType::Float64, DataType::Int64]),
1928 TypeSignature::Exact(vec![DataType::Float32, DataType::Int64]),
1929 TypeSignature::Exact(vec![DataType::Float64]),
1930 TypeSignature::Exact(vec![DataType::Float32]),
1931 ],
1932 Volatility::Immutable,
1933 );
1934 let msg = generate_signature_error_message(
1935 "round",
1936 &sig,
1937 &[DataType::Float64, DataType::Float64],
1938 );
1939 assert_snapshot!(msg, @r"
1940 No function matches the given name and argument types 'round(Float64, Float64)'. You might need to add explicit type casts.
1941 Candidate functions:
1942 round(Float64, Int64)
1943 round(Float32, Int64)
1944 round(Float64)
1945 round(Float32)
1946 ");
1947 }
1948
1949 #[test]
1950 fn test_signature_error_msg_coercible() {
1951 use datafusion_common::types::NativeType;
1952 use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
1953 use insta::assert_snapshot;
1954
1955 let sig = Signature::coercible(
1956 vec![
1957 Coercion::new_implicit(
1958 TypeSignatureClass::Native(
1959 datafusion_common::types::logical_float64(),
1960 ),
1961 vec![TypeSignatureClass::Numeric],
1962 NativeType::Float64,
1963 ),
1964 Coercion::new_implicit(
1965 TypeSignatureClass::Native(datafusion_common::types::logical_int64()),
1966 vec![TypeSignatureClass::Integer],
1967 NativeType::Int64,
1968 ),
1969 ],
1970 Volatility::Immutable,
1971 );
1972 let msg = generate_signature_error_message(
1973 "round",
1974 &sig,
1975 &[DataType::Utf8, DataType::Utf8],
1976 );
1977 assert_snapshot!(msg, @r"
1978 No function matches the given name and argument types 'round(Utf8, Utf8)'. You might need to add explicit type casts.
1979 Candidate functions:
1980 round(Float64, Int64)
1981 ");
1982 }
1983
1984 #[test]
1985 fn test_signature_error_msg_with_names_coercible() {
1986 use datafusion_common::types::NativeType;
1987 use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
1988 use insta::assert_snapshot;
1989
1990 let sig = Signature::coercible(
1991 vec![
1992 Coercion::new_exact(TypeSignatureClass::Native(
1993 datafusion_common::types::logical_string(),
1994 )),
1995 Coercion::new_exact(TypeSignatureClass::Native(
1996 datafusion_common::types::logical_int64(),
1997 )),
1998 Coercion::new_implicit(
1999 TypeSignatureClass::Native(datafusion_common::types::logical_int64()),
2000 vec![TypeSignatureClass::Integer],
2001 NativeType::Int64,
2002 ),
2003 ],
2004 Volatility::Immutable,
2005 )
2006 .with_parameter_names(vec![
2007 "string".to_string(),
2008 "start_pos".to_string(),
2009 "length".to_string(),
2010 ])
2011 .expect("valid parameter names");
2012
2013 let msg = generate_signature_error_message("substr", &sig, &[DataType::Int32]);
2014 assert_snapshot!(msg, @r"
2015 No function matches the given name and argument types 'substr(Int32)'. You might need to add explicit type casts.
2016 Candidate functions:
2017 substr(string: String, start_pos: Int64, length: Int64)
2018 ");
2019 }
2020
2021 fn sum_over(args: Vec<Expr>) -> Expr {
2023 Expr::from(WindowFunction::new(
2024 WindowFunctionDefinition::AggregateUDF(sum_udaf()),
2025 args,
2026 ))
2027 }
2028
2029 #[test]
2030 fn test_check_aggregate_and_window_nesting_ok() -> Result<()> {
2031 use crate::test::function_stub::{count, sum};
2032
2033 let exprs = [
2034 sum(col("a")),
2036 count(col("a")) + lit(1),
2037 sum_over(vec![col("a")]),
2039 sum_over(vec![sum(col("a"))]),
2040 ];
2041
2042 check_aggregate_and_window_nesting(exprs.iter())?;
2043 Ok(())
2044 }
2045
2046 #[test]
2047 fn test_check_aggregate_and_window_nesting_err() {
2048 use crate::test::function_stub::{count, sum};
2049 use insta::assert_snapshot;
2050
2051 let err = check_aggregate_and_window_nesting([&sum(sum(col("a")))]).unwrap_err();
2053 assert_snapshot!(
2054 err.strip_backtrace(),
2055 @"Error during planning: Aggregate function calls cannot be nested: 'sum(a)' is nested inside 'sum(sum(a))'"
2056 );
2057
2058 let err = check_aggregate_and_window_nesting([&sum(col("a") + count(col("b")))])
2060 .unwrap_err();
2061 assert_snapshot!(
2062 err.strip_backtrace(),
2063 @"Error during planning: Aggregate function calls cannot be nested: 'COUNT(b)' is nested inside 'sum(a + COUNT(b))'"
2064 );
2065
2066 let filtered = sum(col("a"))
2068 .filter(sum(col("b")).gt(lit(0)))
2069 .build()
2070 .unwrap();
2071 let err = check_aggregate_and_window_nesting([&filtered]).unwrap_err();
2072 assert_snapshot!(
2073 err.strip_backtrace(),
2074 @"Error during planning: Aggregate function calls cannot be nested: 'sum(b)' is nested inside 'sum(a) FILTER (WHERE sum(b) > Int32(0))'"
2075 );
2076
2077 let ordered = sum(col("a"))
2079 .order_by(vec![Sort::new(sum(col("b")), true, false)])
2080 .build()
2081 .unwrap();
2082 let err = check_aggregate_and_window_nesting([&ordered]).unwrap_err();
2083 assert_snapshot!(
2084 err.strip_backtrace(),
2085 @"Error during planning: Aggregate function calls cannot be nested: 'sum(b)' is nested inside 'sum(a) ORDER BY [sum(b) ASC NULLS LAST]'"
2086 );
2087
2088 let err = check_aggregate_and_window_nesting([&sum(sum_over(vec![col("a")]))])
2090 .unwrap_err();
2091 assert_snapshot!(
2092 err.strip_backtrace(),
2093 @"Error during planning: Aggregate function calls cannot contain window function calls: 'sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)'"
2094 );
2095
2096 let err =
2098 check_aggregate_and_window_nesting([&sum_over(vec![sum_over(vec![col(
2099 "a",
2100 )])])])
2101 .unwrap_err();
2102 assert_snapshot!(
2103 err.strip_backtrace(),
2104 @"Error during planning: Window function calls cannot be nested: 'sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'"
2105 );
2106 }
2107}