chore(query): fix grouping function in type checker (#18502)

* chore(query): fix grouping function in type checker

* chore(query): fix grouping function in type checker

* chore(query): fix grouping function in type checker

Author: sundy-li

src/query/expression/src/type_check.rs

@@ -316,7 +316,7 @@ pub fn check_function<Index: ColumnIndex>(
 
     // Do not check grouping
     if name == "grouping" {
-        debug_assert!(candidates.len() == 1);
+        debug_assert!(!candidates.is_empty());
         let (id, function) = candidates.into_iter().next().unwrap();
         let return_type = function.signature.return_type.clone();
         return Ok(Expr::FunctionCall(FunctionCall {

src/query/functions/src/scalars/other.rs

@@ -391,7 +391,29 @@ fn register_grouping(registry: &mut FunctionRegistry) {
             },
         }))
     }));
-    registry.register_function_factory("grouping", grouping)
+    registry.register_function_factory("grouping", grouping);
+
+    // dummy grouping
+    // used in type_check before AggregateRewriter
+    let dummy_grouping = FunctionFactory::Closure(Box::new(|_, arg_type: &[DataType]| {
+        Some(Arc::new(Function {
+            signature: FunctionSignature {
+                name: "grouping".to_string(),
+                args_type: vec![DataType::Generic(0); arg_type.len()],
+                return_type: DataType::Number(NumberDataType::UInt32),
+            },
+            eval: FunctionEval::Scalar {
+                calc_domain: Box::new(|_, _| FunctionDomain::Full),
+                eval: Box::new(move |args, _| {
+                    unreachable!(
+                        "grouping function must be rewritten in type_checker, but got: {:?}",
+                        args
+                    )
+                }),
+            },
+        }))
+    }));
+    registry.register_function_factory("grouping", dummy_grouping);
 }
 
 fn register_num_to_char(registry: &mut FunctionRegistry) {

src/query/functions/tests/it/scalars/testdata/function_list.txt

@@ -1998,6 +1998,7 @@ Functions overloads:
 0 great_circle_distance(Float64, Float64, Float64, Float64) :: Float32
 1 great_circle_distance(Float64 NULL, Float64 NULL, Float64 NULL, Float64 NULL) :: Float32 NULL
 0 grouping FACTORY
+1 grouping FACTORY
 0 gt(Variant, Variant) :: Boolean
 1 gt(Variant NULL, Variant NULL) :: Boolean NULL
 2 gt(String, String) :: Boolean

src/query/sql/src/planner/semantic/type_check.rs

@@ -2906,30 +2906,6 @@ impl<'a> TypeChecker<'a> {
             arg_types.push(arg_type);
         }
 
-        // rewrite substr('xx', 0, xx) -> substr('xx', 1, xx)
-        if (func_name == "substr" || func_name == "substring")
-            && self
-                .ctx
-                .get_settings()
-                .get_sql_dialect()
-                .unwrap()
-                .substr_index_zero_literal_as_one()
-        {
-            Self::rewrite_substring(&mut args);
-        }
-
-        if func_name == "grouping" {
-            // `grouping` will be rewritten again after resolving grouping sets.
-            return Ok(Box::new((
-                ScalarExpr::FunctionCall(FunctionCall {
-                    span,
-                    params: vec![],
-                    arguments: args,
-                    func_name: "grouping".to_string(),
-                }),
-                DataType::Number(NumberDataType::UInt32),
-            )));
-        }
         if let Some(rewritten_variant_expr) =
             self.try_rewrite_variant_function(span, func_name, &args, &arg_types)
         {
@@ -2949,11 +2925,22 @@ impl<'a> TypeChecker<'a> {
         span: Span,
         func_name: &str,
         mut params: Vec<Scalar>,
-        args: Vec<ScalarExpr>,
+        mut args: Vec<ScalarExpr>,
     ) -> Result<Box<(ScalarExpr, DataType)>> {
+        // rewrite substr('xx', 0, xx) -> substr('xx', 1, xx)
+        if (func_name == "substr" || func_name == "substring")
+            && self
+                .ctx
+                .get_settings()
+                .get_sql_dialect()
+                .unwrap()
+                .substr_index_zero_literal_as_one()
+        {
+            Self::rewrite_substring(&mut args);
+        }
+
         // Type check
         let mut arguments = args.iter().map(|v| v.as_raw_expr()).collect::<Vec<_>>();
-
         // inject the params
         if ["round", "truncate"].contains(&func_name)
             && !args.is_empty()

tests/sqllogictests/suites/duckdb/sql/aggregate/group/group_by_grouping_sets.test

@@ -17,6 +17,7 @@ statement ok
 select sum(number), number % 3 a, grouping(number % 3)+grouping(number % 4) AS lochierarchy from numbers(10)
 group by rollup(number % 3, number % 4)  order by  grouping(number % 3)+grouping(number % 4) ;
 
+
 query TT
 select number % 2 as a, number % 3 as b from numbers(24) group by grouping sets ((a,b), (a), (b)) order by a,b;
 ----
@@ -91,6 +92,9 @@ NULL B 10 0 1 2 1
 b NULL 11 1 0 1 2
 NULL NULL 18 1 1 3 3
 
+statement ok
+select a, b, sum(c) as sc, grouping(a,b) + 3, grouping(b,a)  > 2  from t group by grouping sets ((a,b),(),(b),(a)) order by sc;
+
 query TTIIIII
 select a, b, sum(c) as sc, grouping(b), grouping(a), grouping(a,b), grouping(b,a) from t group by grouping sets ((1,2),(),(2),(1)) order by sc;
 ----