diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index e8de34d9371a000034cfa8ccf5dfd6110e851611..7ecd168137e8d9f1c6abc1bd02627dbdc7a3738d 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -249,6 +249,8 @@ exportMethods("%<=>%",
"getField",
"getItem",
"greatest",
+ "grouping_bit",
+ "grouping_id",
"hex",
"histogram",
"hour",
diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index f9687d680e7a251f1a252890f4abed9b86fe8f19..38384a89919a29a260da66b60663e188e427aac6 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -3890,3 +3890,87 @@ setMethod("not",
jc <- callJStatic("org.apache.spark.sql.functions", "not", x@jc)
column(jc)
})
+
+#' grouping_bit
+#'
+#' Indicates whether a specified column in a GROUP BY list is aggregated or not,
+#' returns 1 for aggregated or 0 for not aggregated in the result set.
+#'
+#' Same as \code{GROUPING} in SQL and \code{grouping} function in Scala.
+#'
+#' @param x Column to compute on
+#'
+#' @rdname grouping_bit
+#' @name grouping_bit
+#' @family agg_funcs
+#' @aliases grouping_bit,Column-method
+#' @export
+#' @examples \dontrun{
+#' df <- createDataFrame(mtcars)
+#'
+#' # With cube
+#' agg(
+#' cube(df, "cyl", "gear", "am"),
+#' mean(df$mpg),
+#' grouping_bit(df$cyl), grouping_bit(df$gear), grouping_bit(df$am)
+#' )
+#'
+#' # With rollup
+#' agg(
+#' rollup(df, "cyl", "gear", "am"),
+#' mean(df$mpg),
+#' grouping_bit(df$cyl), grouping_bit(df$gear), grouping_bit(df$am)
+#' )
+#' }
+#' @note grouping_bit since 2.3.0
+setMethod("grouping_bit",
+ signature(x = "Column"),
+ function(x) {
+ jc <- callJStatic("org.apache.spark.sql.functions", "grouping", x@jc)
+ column(jc)
+ })
+
+#' grouping_id
+#'
+#' Returns the level of grouping.
+#'
+#' Equals to \code{
+#' grouping_bit(c1) * 2^(n - 1) + grouping_bit(c2) * 2^(n - 2) + ... + grouping_bit(cn)
+#' }
+#'
+#' @param x Column to compute on
+#' @param ... additional Column(s) (optional).
+#'
+#' @rdname grouping_id
+#' @name grouping_id
+#' @family agg_funcs
+#' @aliases grouping_id,Column-method
+#' @export
+#' @examples \dontrun{
+#' df <- createDataFrame(mtcars)
+#'
+#' # With cube
+#' agg(
+#' cube(df, "cyl", "gear", "am"),
+#' mean(df$mpg),
+#' grouping_id(df$cyl, df$gear, df$am)
+#' )
+#'
+#' # With rollup
+#' agg(
+#' rollup(df, "cyl", "gear", "am"),
+#' mean(df$mpg),
+#' grouping_id(df$cyl, df$gear, df$am)
+#' )
+#' }
+#' @note grouping_id since 2.3.0
+setMethod("grouping_id",
+ signature(x = "Column"),
+ function(x, ...) {
+ jcols <- lapply(list(x, ...), function (x) {
+ stopifnot(class(x) == "Column")
+ x@jc
+ })
+ jc <- callJStatic("org.apache.spark.sql.functions", "grouping_id", jcols)
+ column(jc)
+ })
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index ef36765a7a725ccbf80fb4c0d86ff70f60f23d58..e02d46426a5a6516fd7edb6598fb34c264dbd061 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -1052,6 +1052,14 @@ setGeneric("from_unixtime", function(x, ...) { standardGeneric("from_unixtime")
#' @export
setGeneric("greatest", function(x, ...) { standardGeneric("greatest") })
+#' @rdname grouping_bit
+#' @export
+setGeneric("grouping_bit", function(x) { standardGeneric("grouping_bit") })
+
+#' @rdname grouping_id
+#' @export
+setGeneric("grouping_id", function(x, ...) { standardGeneric("grouping_id") })
+
#' @rdname hex
#' @export
setGeneric("hex", function(x) { standardGeneric("hex") })
diff --git a/R/pkg/inst/tests/testthat/test_sparkSQL.R b/R/pkg/inst/tests/testthat/test_sparkSQL.R
index 08296354ca7ed2d737810b8795937ad9b736e6ac..12867c15d1f95161f96d5eb2640e00f91b791b34 100644
--- a/R/pkg/inst/tests/testthat/test_sparkSQL.R
+++ b/R/pkg/inst/tests/testthat/test_sparkSQL.R
@@ -1848,7 +1848,11 @@ test_that("test multi-dimensional aggregations with cube and rollup", {
orderBy(
agg(
cube(df, "year", "department"),
- expr("sum(salary) AS total_salary"), expr("avg(salary) AS average_salary")
+ expr("sum(salary) AS total_salary"),
+ expr("avg(salary) AS average_salary"),
+ alias(grouping_bit(df$year), "grouping_year"),
+ alias(grouping_bit(df$department), "grouping_department"),
+ alias(grouping_id(df$year, df$department), "grouping_id")
),
"year", "department"
)
@@ -1875,6 +1879,30 @@ test_that("test multi-dimensional aggregations with cube and rollup", {
mean(c(21000, 32000, 22000)), # 2017
22000, 32000, 21000 # 2017 each department
),
+ grouping_year = c(
+ 1, # global
+ 1, 1, 1, # by department
+ 0, # 2016
+ 0, 0, 0, # 2016 by department
+ 0, # 2017
+ 0, 0, 0 # 2017 by department
+ ),
+ grouping_department = c(
+ 1, # global
+ 0, 0, 0, # by department
+ 1, # 2016
+ 0, 0, 0, # 2016 by department
+ 1, # 2017
+ 0, 0, 0 # 2017 by department
+ ),
+ grouping_id = c(
+ 3, # 11
+ 2, 2, 2, # 10
+ 1, # 01
+ 0, 0, 0, # 00
+ 1, # 01
+ 0, 0, 0 # 00
+ ),
stringsAsFactors = FALSE
)
@@ -1896,7 +1924,10 @@ test_that("test multi-dimensional aggregations with cube and rollup", {
orderBy(
agg(
rollup(df, "year", "department"),
- expr("sum(salary) AS total_salary"), expr("avg(salary) AS average_salary")
+ expr("sum(salary) AS total_salary"), expr("avg(salary) AS average_salary"),
+ alias(grouping_bit(df$year), "grouping_year"),
+ alias(grouping_bit(df$department), "grouping_department"),
+ alias(grouping_id(df$year, df$department), "grouping_id")
),
"year", "department"
)
@@ -1920,6 +1951,27 @@ test_that("test multi-dimensional aggregations with cube and rollup", {
mean(c(21000, 32000, 22000)), # 2017
22000, 32000, 21000 # 2017 each department
),
+ grouping_year = c(
+ 1, # global
+ 0, # 2016
+ 0, 0, 0, # 2016 each department
+ 0, # 2017
+ 0, 0, 0 # 2017 each department
+ ),
+ grouping_department = c(
+ 1, # global
+ 1, # 2016
+ 0, 0, 0, # 2016 each department
+ 1, # 2017
+ 0, 0, 0 # 2017 each department
+ ),
+ grouping_id = c(
+ 3, # 11
+ 1, # 01
+ 0, 0, 0, # 00
+ 1, # 01
+ 0, 0, 0 # 00
+ ),
stringsAsFactors = FALSE
)