Create a separate class RollingGroupBy

DESCRIPTION

@@ -91,6 +91,7 @@ Collate:
     'functions__lazy.R'
     'functions__whenthen.R'
     'group_by.R'
+    'group_by_rolling.R'
     'info.R'
     'ipc.R'
     'is_polars.R'

NAMESPACE

@@ -26,6 +26,7 @@ S3method("$",RPolarsProtoExprArray)
 S3method("$",RPolarsRField)
 S3method("$",RPolarsRNullValues)
 S3method("$",RPolarsRThreadHandle)
+S3method("$",RPolarsRollingGroupBy)
 S3method("$",RPolarsSQLContext)
 S3method("$",RPolarsSeries)
 S3method("$",RPolarsStringCacheHolder)
@@ -77,6 +78,7 @@ S3method("[[",RPolarsProtoExprArray)
 S3method("[[",RPolarsRField)
 S3method("[[",RPolarsRNullValues)
 S3method("[[",RPolarsRThreadHandle)
+S3method("[[",RPolarsRollingGroupBy)
 S3method("[[",RPolarsSQLContext)
 S3method("[[",RPolarsSeries)
 S3method("[[",RPolarsStringCacheHolder)
@@ -94,6 +96,7 @@ S3method(.DollarNames,RPolarsGroupBy)
 S3method(.DollarNames,RPolarsLazyFrame)
 S3method(.DollarNames,RPolarsRField)
 S3method(.DollarNames,RPolarsRThreadHandle)
+S3method(.DollarNames,RPolarsRollingGroupBy)
 S3method(.DollarNames,RPolarsSQLContext)
 S3method(.DollarNames,RPolarsSeries)
 S3method(.DollarNames,RPolarsThen)
@@ -115,6 +118,7 @@ S3method(as_polars_df,RPolarsDataFrame)
 S3method(as_polars_df,RPolarsGroupBy)
 S3method(as_polars_df,RPolarsLazyFrame)
 S3method(as_polars_df,RPolarsLazyGroupBy)
+S3method(as_polars_df,RPolarsRollingGroupBy)
 S3method(as_polars_df,RPolarsSeries)
 S3method(as_polars_df,data.frame)
 S3method(as_polars_df,default)
@@ -165,6 +169,7 @@ S3method(print,RPolarsLazyFrame)
 S3method(print,RPolarsLazyGroupBy)
 S3method(print,RPolarsRField)
 S3method(print,RPolarsRThreadHandle)
+S3method(print,RPolarsRollingGroupBy)
 S3method(print,RPolarsSQLContext)
 S3method(print,RPolarsSeries)
 S3method(print,RPolarsThen)

NEWS.md

@@ -4,7 +4,8 @@
 
 ### What's changed
 
--   New method `$rolling()` for `DataFrame` and `LazyFrame` (#682).
+-   New method `$rolling()` for `DataFrame` and `LazyFrame`. When this is
+    applied, it creates an object of class `RPolarsRollingGroupBy` (#682, #694).
 -   New method `$sink_ndjson()` for LazyFrame (#681).
 -   New function `pl$duration()` to create a duration by components (week, day,
     hour, etc.), and use them with date(time) variables (#692).

R/as_polars.R

@@ -102,6 +102,11 @@ as_polars_df.RPolarsGroupBy = function(x, ...) {
   x$ungroup()
 }
 
+#' @rdname as_polars_df
+#' @export
+as_polars_df.RPolarsRollingGroupBy = function(x, ...) {
+  x$ungroup()
+}
 
 #' @rdname as_polars_df
 #' @export

R/dataframe__frame.R

@@ -1828,8 +1828,8 @@ DataFrame_write_ndjson = function(file) {
 #'   pl$max("a")$alias("max_a")
 #' )
 DataFrame_rolling = function(index_column, period, offset = NULL, closed = "right", by = NULL, check_sorted = TRUE) {
-  out = self$lazy()$rolling(index_column, period, offset, closed, by, check_sorted)
-  attr(out, "is_rolling_group_by") = TRUE
-  class(out) = "RPolarsGroupBy"
-  out
+  if (is.null(offset)) {
+    offset = paste0("-", period)
+  }
+  construct_rolling_group_by(self, index_column, period, offset, closed, by, check_sorted)
 }

R/group_by.R

@@ -12,7 +12,6 @@
 NULL
 
 
-
 RPolarsGroupBy = new.env(parent = emptyenv())
 
 #' @export
@@ -25,28 +24,30 @@ RPolarsGroupBy = new.env(parent = emptyenv())
 #' @export
 `[[.RPolarsGroupBy` = `$.RPolarsGroupBy`
 
-#' @title auto complete $-access into a polars object
-#' @description called by the interactive R session internally
-#' @param x GroupBy
-#' @param pattern code-stump as string to auto-complete
-#' @return char vec
 #' @export
-#' @inherit .DollarNames.RPolarsDataFrame return
 #' @noRd
 .DollarNames.RPolarsGroupBy = function(x, pattern = "") {
   paste0(ls(RPolarsGroupBy, pattern = pattern), "()")
 }
 
-
 #' The internal GroupBy constructor
 #' @return The input as grouped DataFrame
 #' @noRd
 construct_group_by = function(df, groupby_input, maintain_order) {
-  if (!inherits(df, "RPolarsDataFrame")) stop("internal error: construct_group called not on DataFrame")
-  df = df$clone()
-  attr(df, "private") = list(groupby_input = unlist(groupby_input), maintain_order = maintain_order)
-  class(df) = "RPolarsGroupBy"
-  df
+  if (!inherits(df, "RPolarsDataFrame")) {
+    stop("internal error: construct_group called not on DataFrame")
+  }
+  # Make an empty object. Store everything (including data) in attributes, so
+  # that we can keep the RPolarsDataFrame class on the data but still return
+  # a RPolarsGroupBy object here.
+  out = c(" ")
+  attr(out, "private") = list(
+    dat = df$clone(),
+    groupby_input = unlist(groupby_input),
+    maintain_order = maintain_order
+  )
+  class(out) = "RPolarsGroupBy"
+  out
 }
 
 
@@ -58,13 +59,13 @@ construct_group_by = function(df, groupby_input, maintain_order) {
 #' @return self
 #' @export
 #'
-#' @examples pl$DataFrame(iris)$group_by("Species")
+#' @examples
+#' pl$DataFrame(iris)$group_by("Species")
 print.RPolarsGroupBy = function(x, ...) {
-  .pr$DataFrame$print(x)
-  cat("groups: ")
   prv = attr(x, "private")
-  cat(toString(prv$groupby_input))
-  cat("\nmaintain order: ", prv$maintain_order)
+  .pr$DataFrame$print(prv$dat)
+  cat("groups:", toString(prv$groupby_input))
+  cat("\nmaintain order:", prv$maintain_order)
   invisible(x)
 }
 
@@ -86,18 +87,13 @@ print.RPolarsGroupBy = function(x, ...) {
 #'   pl$col("bar")$mean()$alias("bar_tail_sum")
 #' )
 GroupBy_agg = function(...) {
-  if (isTRUE(attributes(self)[["is_rolling_group_by"]])) {
-    class(self) = "RPolarsLazyGroupBy"
-    self$agg(unpack_list(..., .context = "in $agg():"))$collect(no_optimization = TRUE)
-  } else {
-    class(self) = "RPolarsDataFrame"
-    self$lazy()$group_by(
-      attr(self, "private")$groupby_input,
-      maintain_order = attr(self, "private")$maintain_order
-    )$
-      agg(...)$
-      collect(no_optimization = TRUE)
-  }
+  prv = attr(self, "private")
+  prv$dat$lazy()$group_by(
+    prv$groupby_input,
+    maintain_order = prv$maintain_order
+  )$
+    agg(...)$
+    collect(no_optimization = TRUE)
 }
 
 
@@ -300,13 +296,6 @@ GroupBy_null_count = function() {
 #'
 #' gb$ungroup()
 GroupBy_ungroup = function() {
-  if (isTRUE(attributes(self)[["is_rolling_group_by"]])) {
-    class(self) = "RPolarsLazyGroupBy"
-    self = self$ungroup()$collect(no_optimization = TRUE)
-  } else {
-    self = .pr$DataFrame$clone_in_rust(self)
-    class(self) = "RPolarsDataFrame"
-    attr(self, "private") = NULL
-  }
-  self
+  prv = attr(self, "private")
+  prv$dat
 }

R/group_by_rolling.R

@@ -0,0 +1,133 @@
+#' Operations on Polars DataFrame grouped by rolling windows
+#'
+#' @return not applicable
+#' @name RollingGroupBy_class
+NULL
+
+RPolarsRollingGroupBy = new.env(parent = emptyenv())
+
+#' @export
+`$.RPolarsRollingGroupBy` = function(self, name) {
+  func = RPolarsRollingGroupBy[[name]]
+  environment(func) = environment()
+  func
+}
+
+#' @export
+`[[.RPolarsRollingGroupBy` = `$.RPolarsRollingGroupBy`
+
+#' @export
+#' @noRd
+.DollarNames.RPolarsRollingGroupBy = function(x, pattern = "") {
+  paste0(ls(RPolarsRollingGroupBy, pattern = pattern), "()")
+}
+
+#' The internal RollingGroupBy constructor
+#' @return The input as grouped DataFrame
+#' @noRd
+construct_rolling_group_by = function(df, index_column, period, offset, closed, by, check_sorted) {
+  if (!inherits(df, "RPolarsDataFrame")) {
+    stop("internal error: construct_group called not on DataFrame")
+  }
+  # Make an empty object. Store everything (including data) in attributes, so
+  # that we can keep the RPolarsDataFrame class on the data but still return
+  # a RPolarsRollingGroupBy object here.
+  out = c(" ")
+  attr(out, "private") = list(
+    dat = df$clone(),
+    index_column = index_column,
+    period = period,
+    offset = offset,
+    closed = closed,
+    by = by,
+    check_sorted = check_sorted
+  )
+  class(out) = "RPolarsRollingGroupBy"
+  out
+}
+
+#' print RollingGroupBy
+#'
+#' @param x DataFrame
+#' @param ... not used
+#' @noRd
+#' @return self
+#' @export
+#'
+#' @examples
+#' df = pl$DataFrame(
+#'   dt = c("2020-01-01", "2020-01-01", "2020-01-01", "2020-01-02", "2020-01-03", "2020-01-08"),
+#'   a = c(3, 7, 5, 9, 2, 1)
+#' )$with_columns(
+#'   pl$col("dt")$str$strptime(pl$Date, format = NULL)$set_sorted()
+#' )
+#'
+#' df$rolling(index_column = "dt", period = "2d")
+print.RPolarsRollingGroupBy = function(x, ...) {
+  prv = attr(x, "private")
+  .pr$DataFrame$print(prv$dat)
+  cat(paste("index column:", prv$index))
+  cat(paste("\nother groups:", toString(prv$by)))
+  cat(paste("\nperiod:", prv$period))
+  cat(paste("\noffset:", prv$offset))
+  cat(paste("\nclosed:", prv$closed))
+}
+
+
+#' Aggregate over a RollingGroupBy
+#'
+#' Aggregate a DataFrame over a rolling window created with `$rolling()`.
+#'
+#' @param ... Exprs to aggregate over. Those can also be passed wrapped in a
+#' list, e.g `$agg(list(e1,e2,e3))`.
+#'
+#' @return An aggregated [DataFrame][DataFrame_class]
+#' @examples
+#' df = pl$DataFrame(
+#'   dt = c("2020-01-01", "2020-01-01", "2020-01-01", "2020-01-02", "2020-01-03", "2020-01-08"),
+#'   a = c(3, 7, 5, 9, 2, 1)
+#' )$with_columns(
+#'   pl$col("dt")$str$strptime(pl$Date, format = NULL)$set_sorted()
+#' )
+#'
+#' df$rolling(index_column = "dt", period = "2d")$agg(
+#'   pl$col("a"),
+#'   pl$sum("a")$alias("sum_a"),
+#'   pl$min("a")$alias("min_a"),
+#'   pl$max("a")$alias("max_a")
+#' )
+RollingGroupBy_agg = function(...) {
+  prv = attr(self, "private")
+  prv$dat$
+    lazy()$
+    rolling(
+    index_column = prv$index,
+    period = prv$period,
+    offset = prv$offset,
+    closed = prv$closed,
+    by = prv$by,
+    check_sorted = prv$check_sorted
+  )$
+    agg(unpack_list(..., .context = "in $agg():"))$
+    collect(no_optimization = TRUE)
+}
+
+#' Ungroup a RollingGroupBy object
+#'
+#' Revert the `$rolling()` operation. Doing `<DataFrame>$rolling(...)$ungroup()`
+#' returns the original `DataFrame`.
+#'
+#' @return [DataFrame][DataFrame_class]
+#' @examples
+#' df = pl$DataFrame(
+#'   dt = c("2020-01-01", "2020-01-01", "2020-01-01", "2020-01-02", "2020-01-03", "2020-01-08"),
+#'   a = c(3, 7, 5, 9, 2, 1)
+#' )$with_columns(
+#'   pl$col("dt")$str$strptime(pl$Date, format = NULL)$set_sorted()
+#' )
+#'
+#' df$rolling(index_column = "dt", period = "2d")$ungroup()
+RollingGroupBy_ungroup = function() {
+  prv = attr(self, "private")
+  prv$dat
+}

R/zzz.R

@@ -22,6 +22,9 @@ replace_private_with_pub_methods(RPolarsLazyFrame, "^LazyFrame_")
 # LazyGroupBy
 replace_private_with_pub_methods(RPolarsLazyGroupBy, "^LazyGroupBy_")
 
+# LazyGroupBy
+replace_private_with_pub_methods(RPolarsRollingGroupBy, "^RollingGroupBy_")
+
 # Expr
 replace_private_with_pub_methods(RPolarsExpr, "^Expr_")
 

altdoc/altdoc_preprocessing.R

@@ -36,7 +36,7 @@ out = list()
 # order determines order in sidebar
 classes = c(
   "pl", "Series", "DataFrame", "LazyFrame", "GroupBy",
-  "LazyGroupBy", "ExprList", "ExprBin", "ExprCat", "ExprDT",
+  "LazyGroupBy", "RollingGroupBy", "ExprList", "ExprBin", "ExprCat", "ExprDT",
   "ExprMeta", "ExprName", "ExprStr", "ExprStruct",
   "Expr", "IO", "RField", "RThreadHandle", "SQLContext", "S3"
 )