move grouping udf

Author: chenkovsky

Cargo.lock

@@ -27,7 +27,6 @@ pub mod expr_ext;
 pub mod getfield;
 pub mod greatest;
 mod greatest_least_utils;
-pub mod grouping;
 pub mod least;
 pub mod named_struct;
 pub mod nullif;
@@ -56,7 +55,6 @@ make_udf_function!(least::LeastFunc, least);
 make_udf_function!(union_extract::UnionExtractFun, union_extract);
 make_udf_function!(union_tag::UnionTagFunc, union_tag);
 make_udf_function!(version::VersionFunc, version);
-make_udf_function!(grouping::GroupingFunc, grouping);
 
 pub mod expr_fn {
     use datafusion_expr::{Expr, Literal};

datafusion/functions/src/core/mod.rs

@@ -45,7 +45,6 @@ arrow = { workspace = true }
 chrono = { workspace = true }
 datafusion-common = { workspace = true, default-features = true }
 datafusion-expr = { workspace = true }
-datafusion-functions = { workspace = true }
 datafusion-physical-expr = { workspace = true }
 indexmap = { workspace = true }
 itertools = { workspace = true }

datafusion/optimizer/Cargo.toml

@@ -26,11 +26,14 @@ use datafusion_expr::{
     ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
     Volatility,
 };
-use datafusion_macros::user_doc;
 use std::any::Any;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
-use crate::utils::make_scalar_function;
+pub(crate) fn grouping_func() -> Arc<datafusion_expr::ScalarUDF> {
+    static INSTANCE: LazyLock<Arc<datafusion_expr::ScalarUDF>> =
+        LazyLock::new(|| Arc::new(datafusion_expr::ScalarUDF::from(GroupingFunc::new())));
+    Arc::clone(&INSTANCE)
+}
 
 macro_rules! grouping_id {
     ($grouping_id:expr, $indices:expr, $type:ty, $array_type:ty) => {{
@@ -69,28 +72,6 @@ macro_rules! grouping_id {
     }};
 }
 
-#[user_doc(
-    doc_section(label = "Other Functions"),
-    description = "Developer API: Returns the level of grouping, equals to (((grouping_id >> array[0]) & 1) << (n-1)) + (((grouping_id >> array[1]) & 1) << (n-2)) + ... + (((grouping_id >> array[n-1]) & 1) << 0). Returns grouping_id if indices is not provided.",
-    syntax_example = "grouping(grouping_id[, indices])",
-    sql_example = r#"```sql
-> SELECT grouping(__grouping_id, make_array(0)) FROM table GROUP BY GROUPING SETS ((a), (b));
-+----------------+
-| grouping       |
-+----------------+
-| 1              |
-| 0              |
-+----------------+
-```"#,
-    argument(
-        name = "grouping_id",
-        description = "The internal grouping ID column (UInt8/16/32/64)"
-    ),
-    argument(
-        name = "indices",
-        description = "The indices of the column in the grouping set (Int32)"
-    )
-)]
 #[derive(Debug)]
 pub struct GroupingFunc {
     signature: Signature,
@@ -128,7 +109,22 @@ impl ScalarUDFImpl for GroupingFunc {
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
-        make_scalar_function(grouping_inner, vec![])(&args.args)
+        let mut array_args = Vec::new();
+        for arg in &args.args {
+            match arg {
+                ColumnarValue::Array(array) => {
+                    array_args.push(Arc::clone(array));
+                }
+                ColumnarValue::Scalar(scalar) => {
+                    // Convert scalar to array by repeating the value
+                    let array = scalar.to_array_of_size(args.number_rows)?;
+                    array_args.push(array);
+                }
+            }
+        }
+
+        let result = grouping_inner(&array_args)?;
+        Ok(ColumnarValue::Array(result))
     }
 
     fn short_circuits(&self) -> bool {
@@ -175,7 +171,7 @@ impl ScalarUDFImpl for GroupingFunc {
     }
 
     fn documentation(&self) -> Option<&Documentation> {
-        self.doc()
+        None
     }
 }
 

datafusion/functions/src/core/grouping.rs renamed to datafusion/optimizer/src/analyzer/grouping.rs

@@ -38,6 +38,8 @@ pub mod function_rewrite;
 pub mod resolve_grouping_function;
 pub mod type_coercion;
 
+mod grouping;
+
 pub mod subquery {
     #[deprecated(
         since = "44.0.0",

datafusion/optimizer/src/analyzer/mod.rs

@@ -21,18 +21,18 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use crate::analyzer::grouping::grouping_func;
 use crate::analyzer::AnalyzerRule;
 
 use arrow::array::ListArray;
 use arrow::datatypes::Int32Type;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{plan_err, Column, DFSchemaRef, Result, ScalarValue};
-use datafusion_expr::expr::{AggregateFunction, Alias};
+use datafusion_expr::expr::{AggregateFunction, Alias, ScalarFunction};
 use datafusion_expr::logical_plan::LogicalPlan;
 use datafusion_expr::utils::grouping_set_to_exprlist;
 use datafusion_expr::{Aggregate, Expr, Projection};
-use datafusion_functions::core::grouping;
 use itertools::Itertools;
 
 /// Replaces grouping aggregation function with value derived from internal grouping id
@@ -198,7 +198,10 @@ fn grouping_function_on_id(
             .enumerate()
             .all(|(idx, expr)| group_by_expr.get(expr) == Some(&idx))
     {
-        return Ok(grouping().call(vec![grouping_id_column]));
+        return Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
+            grouping_func(),
+            vec![grouping_id_column],
+        )));
     }
 
     let args = args
@@ -216,5 +219,8 @@ fn grouping_function_on_id(
         ))),
         None,
     );
-    Ok(grouping().call(vec![grouping_id_column, indices]))
+    Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
+        grouping_func(),
+        vec![grouping_id_column, indices],
+    )))
 }

datafusion/optimizer/src/analyzer/resolve_grouping_function.rs

@@ -21,6 +21,7 @@ use datafusion_common::{
     assert_contains, Column, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue,
     TableReference,
 };
+use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::test::function_stub::{
     count_udaf, max_udaf, min_udaf, sum, sum_udaf,
 };
@@ -29,7 +30,6 @@ use datafusion_expr::{
     LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode,
     UserDefinedLogicalNodeCore,
 };
-use datafusion_functions::core::grouping;
 use datafusion_functions::unicode;
 use datafusion_functions_aggregate::grouping::grouping_udaf;
 use datafusion_functions_nested::make_array::make_array_udf;
@@ -2567,6 +2567,47 @@ fn test_not_ilike_filter_with_escape() {
 
 #[test]
 fn test_grouping() -> Result<()> {
+    #[derive(Debug)]
+    struct MockGroupingFunc {
+        signature: Signature,
+    }
+
+    use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
+    use std::any::Any;
+
+    impl MockGroupingFunc {
+        fn new() -> Self {
+            Self {
+                signature: Signature::user_defined(Volatility::Immutable),
+            }
+        }
+    }
+
+    impl ScalarUDFImpl for MockGroupingFunc {
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn name(&self) -> &str {
+            "grouping"
+        }
+
+        fn signature(&self) -> &Signature {
+            &self.signature
+        }
+
+        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+            Ok(DataType::Int32)
+        }
+
+        fn invoke_with_args(
+            &self,
+            _args: datafusion_expr::ScalarFunctionArgs,
+        ) -> Result<datafusion_expr::ColumnarValue> {
+            todo!()
+        }
+    }
+
     let schema = Schema::new(vec![
         Field::new("c1", DataType::Int32, false),
         Field::new("c2", DataType::Int32, false),
@@ -2587,6 +2628,8 @@ fn test_grouping() -> Result<()> {
         )?
         .build()?;
 
+    let udf = Arc::new(datafusion_expr::ScalarUDF::from(MockGroupingFunc::new()));
+
     let group1 =
         ScalarValue::List(Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(
             vec![Some(vec![Some(0)])],
@@ -2601,15 +2644,21 @@ fn test_grouping() -> Result<()> {
         )));
     let project = LogicalPlanBuilder::from(plan)
         .project(vec![
-            grouping()
-                .call(vec![col(Aggregate::INTERNAL_GROUPING_ID), lit(group1)])
-                .alias("grouping(test.c1)"),
-            grouping()
-                .call(vec![col(Aggregate::INTERNAL_GROUPING_ID), lit(group2)])
-                .alias("grouping(test.c2)"),
-            grouping()
-                .call(vec![col(Aggregate::INTERNAL_GROUPING_ID), lit(group3)])
-                .alias("grouping(test.c1,test.c2)"),
+            Expr::ScalarFunction(ScalarFunction::new_udf(
+                Arc::clone(&udf),
+                vec![col(Aggregate::INTERNAL_GROUPING_ID), lit(group1)],
+            ))
+            .alias("grouping(test.c1)"),
+            Expr::ScalarFunction(ScalarFunction::new_udf(
+                Arc::clone(&udf),
+                vec![col(Aggregate::INTERNAL_GROUPING_ID), lit(group2)],
+            ))
+            .alias("grouping(test.c2)"),
+            Expr::ScalarFunction(ScalarFunction::new_udf(
+                Arc::clone(&udf),
+                vec![col(Aggregate::INTERNAL_GROUPING_ID), lit(group3)],
+            ))
+            .alias("grouping(test.c1,test.c2)"),
         ])?
         .build()?;
     let unparser = Unparser::new(&UnparserPostgreSqlDialect {});