fix(query): group by item allow set returning functions (#16986)

* fix(query): group by item allow set returning functions * fix
databendlabs · Dec 3, 2024 · 5581b05 · 5581b05
1 parent 9bb9441
commit 5581b05
Show file tree

Hide file tree

Showing 15 changed files with 174 additions and 256 deletions.
diff --git a/src/query/functions/src/lib.rs b/src/query/functions/src/lib.rs
@@ -71,6 +71,9 @@ pub const GENERAL_WINDOW_FUNCTIONS: [&str; 13] = [
     "cume_dist",
 ];
 
+pub const RANK_WINDOW_FUNCTIONS: [&str; 5] =
+    ["first_value", "first", "last_value", "last", "nth_value"];
+
 pub const GENERAL_LAMBDA_FUNCTIONS: [&str; 16] = [
     "array_transform",
     "array_apply",

diff --git a/src/query/sql/src/planner/binder/aggregate.rs b/src/query/sql/src/planner/binder/aggregate.rs
@@ -33,7 +33,7 @@ use itertools::Itertools;
 use super::prune_by_children;
 use super::ExprContext;
 use super::Finder;
-use crate::binder::project_set::SetReturningRewriter;
+use crate::binder::project_set::SetReturningAnalyzer;
 use crate::binder::scalar::ScalarBinder;
 use crate::binder::select::SelectList;
 use crate::binder::Binder;
@@ -349,13 +349,6 @@ impl Binder {
         bind_context: &mut BindContext,
         select_list: &mut SelectList,
     ) -> Result<()> {
-        if !bind_context.srf_info.srfs.is_empty() {
-            // Rewrite the Set-returning functions in Aggregate function as columns.
-            let mut srf_rewriter = SetReturningRewriter::new(bind_context, true);
-            for item in select_list.items.iter_mut() {
-                srf_rewriter.visit(&mut item.scalar)?;
-            }
-        }
         let mut rewriter = AggregateRewriter::new(bind_context, self.metadata.clone());
         for item in select_list.items.iter_mut() {
             rewriter.visit(&mut item.scalar)?;
@@ -714,6 +707,9 @@ impl Binder {
                 .bind(expr)
                 .or_else(|e| self.resolve_alias_item(bind_context, expr, available_aliases, e))?;
 
+            let mut analyzer = SetReturningAnalyzer::new(bind_context, self.metadata.clone());
+            analyzer.visit(&mut scalar_expr)?;
+
             if collect_grouping_sets && !grouping_sets.last().unwrap().contains(&scalar_expr) {
                 grouping_sets.last_mut().unwrap().push(scalar_expr.clone());
             }
@@ -727,11 +723,6 @@ impl Binder {
                 continue;
             }
 
-            if !bind_context.srf_info.srfs.is_empty() {
-                let mut srf_rewriter = SetReturningRewriter::new(bind_context, false);
-                srf_rewriter.visit(&mut scalar_expr)?;
-            }
-
             let group_item_name = format!("{:#}", expr);
             let index = if let ScalarExpr::BoundColumnRef(BoundColumnRef {
                 column: ColumnBinding { index, .. },
@@ -877,12 +868,7 @@ impl Binder {
                     .set_span(expr.span()),
             )
         } else {
-            let (alias, mut scalar) = available_aliases[result[0]].clone();
-
-            if !bind_context.srf_info.srfs.is_empty() {
-                let mut srf_rewriter = SetReturningRewriter::new(bind_context, false);
-                srf_rewriter.visit(&mut scalar)?;
-            }
+            let (alias, scalar) = available_aliases[result[0]].clone();
 
             // check scalar first, avoid duplicate create column.
             let mut scalar_column_index = None;

diff --git a/src/query/sql/src/planner/binder/bind_query/bind_select.rs b/src/query/sql/src/planner/binder/bind_query/bind_select.rs
@@ -141,10 +141,10 @@ impl Binder {
             .map(|item| (item.alias.clone(), item.scalar.clone()))
             .collect::<Vec<_>>();
 
-        // Check Set-returning functions, if the argument contains aggregation function or group item,
+        // Rewrite Set-returning functions, if the argument contains aggregation function or group item,
         // set as lazy Set-returning functions.
         if !from_context.srf_info.srfs.is_empty() {
-            self.check_project_set_select(&mut from_context)?;
+            self.rewrite_project_set_select(&mut from_context)?;
         }
 
         // Bind Set-returning functions before filter plan and aggregate plan.

diff --git a/src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs b/src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs
@@ -38,7 +38,6 @@ use databend_common_storages_result_cache::ResultCacheMetaManager;
 use databend_common_storages_result_cache::ResultScan;
 use databend_common_users::UserApiProvider;
 
-use crate::binder::project_set::SetReturningRewriter;
 use crate::binder::scalar::ScalarBinder;
 use crate::binder::table_args::bind_table_args;
 use crate::binder::Binder;
@@ -50,7 +49,6 @@ use crate::plans::EvalScalar;
 use crate::plans::FunctionCall;
 use crate::plans::RelOperator;
 use crate::plans::ScalarItem;
-use crate::plans::VisitorMut;
 use crate::BindContext;
 use crate::ScalarExpr;
 
@@ -356,11 +354,6 @@ impl Binder {
                         self.normalize_select_list(&mut bind_context, &select_list)?;
                     // analyze Set-returning functions.
                     self.analyze_project_set_select(&mut bind_context, &mut select_list)?;
-                    // rewrite Set-returning functions as columns.
-                    let mut srf_rewriter = SetReturningRewriter::new(&mut bind_context, false);
-                    for item in select_list.items.iter_mut() {
-                        srf_rewriter.visit(&mut item.scalar)?;
-                    }
                     // bind Set-returning functions.
                     let srf_expr = self.bind_project_set(&mut bind_context, child, false)?;
                     // clear Set-returning functions, avoid duplicate bind.

diff --git a/src/query/sql/src/planner/binder/distinct.rs b/src/query/sql/src/planner/binder/distinct.rs
@@ -18,7 +18,6 @@ use std::sync::Arc;
 use databend_common_ast::Span;
 use databend_common_exception::Result;
 
-use crate::binder::project_set::SetReturningRewriter;
 use crate::binder::Binder;
 use crate::binder::ColumnBinding;
 use crate::optimizer::SExpr;
@@ -43,14 +42,6 @@ impl Binder {
         scalar_items: &mut HashMap<IndexType, ScalarItem>,
         child: SExpr,
     ) -> Result<SExpr> {
-        if !bind_context.srf_info.srfs.is_empty() {
-            // Rewrite the Set-returning functions as columns.
-            let mut srf_rewriter = SetReturningRewriter::new(bind_context, false);
-            for (_, item) in scalar_items.iter_mut() {
-                srf_rewriter.visit(&mut item.scalar)?;
-            }
-        }
-
         let scalar_items: Vec<ScalarItem> = scalar_items
             .drain()
             .map(|(_, item)| {