refactor: Define QueryOptFlags S7 class (#1633)

Author: eitsupi

.lintr.R

@@ -1,8 +1,8 @@
 linters <- linters_with_defaults(
   line_length_linter(100L),
   object_name_linter = object_name_linter(
-    styles = c("snake_case", "symbols", "SNAKE_CASE"),
-    regexes = c("^pl__.*", "^_.*")
+    styles = c("snake_case", "symbols", "SNAKE_CASE", "CamelCase"),
+    regexes = c("^.*__.*", "^_.*")
   ),
   commented_code_linter = NULL # TODO: remove commented code
 )

DESCRIPTION

@@ -20,7 +20,8 @@ BugReports: https://github.com/pola-rs/r-polars/issues
 Depends:
     R (>= 4.3)
 Imports:
-    rlang (>= 1.1.0)
+    rlang (>= 1.1.0),
+    S7 (>= 0.2.0)
 Suggests:
     arrow,
     bit64,

NAMESPACE

@@ -196,6 +196,7 @@ S3method(print,polars_series)
 S3method(row.names,polars_data_frame)
 S3method(stats::median,polars_series)
 S3method(sum,polars_series)
+S3method(utils::.DollarNames,polars::QueryOptFlags)
 S3method(utils::.DollarNames,polars_chained_then)
 S3method(utils::.DollarNames,polars_data_frame)
 S3method(utils::.DollarNames,polars_datatype_expr)
@@ -281,6 +282,7 @@ export(polars_envvars)
 export(polars_info)
 export(polars_options)
 export(polars_options_reset)
+import(S7)
 import(rlang)
 importFrom(utils,findMatches)
 importFrom(utils,head)

R/functions-lazy.R

@@ -154,21 +154,14 @@ pl__collect_all <- function(
     check_dots_empty0(...)
     check_list_of_polars_lf(lazy_frames)
     engine <- arg_match0(engine, c("auto", "in-memory", "streaming"))
+    # TODO: add support for argument `optimizations`
+    optflags <- QueryOptFlags()
+    check_is_S7(optflags, QueryOptFlags)
+    validate(optflags)
 
     lfs <- lapply(lazy_frames, \(x) x$`_ldf`)
-    # TODO: add support for argument `optimizations`
-    optflags <- list(
-      comm_subexpr_elim = TRUE,
-      comm_subplan_elim = TRUE,
-      cluster_with_columns = TRUE,
-      predicate_pushdown = TRUE,
-      projection_pushdown = TRUE,
-      simplify_expression = TRUE,
-      slice_pushdown = TRUE,
-      type_coercion = TRUE
-    )
 
-    collect_all(lfs, engine = engine, optflags) |>
+    collect_all(lfs, engine = engine, optflags = optflags) |>
       lapply(\(ptr) .savvy_wrap_PlRDataFrame(ptr) |> wrap())
   })
 }

R/lazyframe-frame.R

@@ -234,16 +234,8 @@ lazyframe__group_by <- function(..., .maintain_order = FALSE) {
 #' Individual optimizations may be disabled by setting the corresponding parameter to `FALSE`.
 #' @inherit pl__DataFrame return
 #' @inheritParams rlang::args_dots_empty
+#' @inheritParams QueryOptFlags
 #' @param type_coercion A logical, indicates type coercion optimization.
-#' @param predicate_pushdown A logical, indicates predicate pushdown optimization.
-#' @param projection_pushdown A logical, indicates projection pushdown optimization.
-#' @param simplify_expression A logical, indicates simplify expression optimization.
-#' @param slice_pushdown A logical, indicates slice pushdown optimization.
-#' @param comm_subplan_elim A logical, indicates trying to cache branching subplans that occur
-#' on self-joins or unions.
-#' @param comm_subexpr_elim A logical, indicates trying to cache common subexpressions.
-#' @param cluster_with_columns A logical, indicates to combine sequential independent calls
-#' to with_columns.
 #' @param collapse_joins `r lifecycle::badge("deprecated")`
 #'   Use `predicate_pushdown` instead.
 #' @param no_optimization A logical. If `TRUE`, turn off (certain) optimizations.

R/lazyframe-opt_flags.R

@@ -0,0 +1,167 @@
+# nolint start: object_length_linter
+
+# The env for storing QueryOptFlags methods
+polars_query_opt_flags__methods <- new.env(parent = emptyenv())
+
+#' The set of the optimizations considered during query optimization
+#'
+#' `r lifecycle::badge("experimental")`
+#'
+#' @inheritParams rlang::args_dots_empty
+#' @param predicate_pushdown A logical, indicates predicate pushdown optimization.
+#' @param projection_pushdown A logical, indicates projection pushdown optimization.
+#' @param simplify_expression A logical, indicates simplify expression optimization.
+#' @param slice_pushdown A logical, indicates slice pushdown optimization.
+#' @param comm_subplan_elim A logical, indicates trying to cache branching subplans that occur
+#'   on self-joins or unions.
+#' @param comm_subexpr_elim A logical, indicates trying to cache common subexpressions.
+#' @param cluster_with_columns A logical, indicates to combine sequential independent calls
+#'   to with_columns.
+#' @param check_order_observe A logical, indicates not to maintain order
+#'   if the order would not be observed.
+#' @param fast_projection A logical, indicates to replace simple projections
+#'   with a faster inlined projection that skips the expression engine.
+#' @return A `QueryOptFlags` object.
+#' @name QueryOptFlags
+#' @examples
+#' opt_flags <- pl$QueryOptFlags()
+#' opt_flags
+#'
+#' S7::check_is_S7(opt_flags, pl$QueryOptFlags)
+#' @keywords internal
+NULL
+
+QueryOptFlags <- new_class(
+  "QueryOptFlags",
+  properties = list(
+    # TODO: switch to scalar properties when supported
+    # https://github.com/RConsortium/S7/pull/433
+    type_coercion = class_logical,
+    type_check = class_logical,
+    predicate_pushdown = class_logical,
+    projection_pushdown = class_logical,
+    simplify_expression = class_logical,
+    slice_pushdown = class_logical,
+    comm_subplan_elim = class_logical,
+    comm_subexpr_elim = class_logical,
+    cluster_with_columns = class_logical,
+    check_order_observe = class_logical,
+    fast_projection = class_logical,
+    eager = class_logical,
+    streaming = class_logical
+  ),
+  constructor = function(
+    ...,
+    predicate_pushdown = TRUE,
+    projection_pushdown = TRUE,
+    simplify_expression = TRUE,
+    slice_pushdown = TRUE,
+    comm_subplan_elim = TRUE,
+    comm_subexpr_elim = TRUE,
+    cluster_with_columns = TRUE,
+    check_order_observe = TRUE,
+    fast_projection = TRUE
+  ) {
+    check_dots_empty0(...)
+
+    # Should be synced with impl Default for OptFlags in Rust
+    # https://github.com/pola-rs/polars/blob/d0be4c06447f2e530c6f68d3cd73b43f2e33e2f3/crates/polars-plan/src/frame/opt_state.rs#L77-L81 # nolint: line_length_linter
+    new_object(
+      S7_object(),
+      type_coercion = TRUE,
+      type_check = TRUE,
+      predicate_pushdown = predicate_pushdown,
+      projection_pushdown = projection_pushdown,
+      simplify_expression = simplify_expression,
+      slice_pushdown = slice_pushdown,
+      comm_subplan_elim = comm_subplan_elim,
+      comm_subexpr_elim = comm_subexpr_elim,
+      cluster_with_columns = cluster_with_columns,
+      check_order_observe = check_order_observe,
+      fast_projection = fast_projection,
+      eager = FALSE,
+      streaming = FALSE
+    )
+  },
+  validator = function(self) {
+    call <- caller_env(4L)
+
+    check_bool(self@type_coercion, arg = "type_coercion", call = call)
+    check_bool(self@type_check, arg = "type_check", call = call)
+
+    check_bool(self@predicate_pushdown, arg = "predicate_pushdown", call = call)
+    check_bool(self@projection_pushdown, arg = "projection_pushdown", call = call)
+    check_bool(self@simplify_expression, arg = "simplify_expression", call = call)
+    check_bool(self@slice_pushdown, arg = "slice_pushdown", call = call)
+    check_bool(self@comm_subplan_elim, arg = "comm_subplan_elim", call = call)
+    check_bool(self@comm_subexpr_elim, arg = "comm_subexpr_elim", call = call)
+    check_bool(self@cluster_with_columns, arg = "cluster_with_columns", call = call)
+    check_bool(self@check_order_observe, arg = "check_order_observe", call = call)
+    check_bool(self@fast_projection, arg = "fast_projection", call = call)
+
+    check_bool(self@eager, arg = "eager", call = call)
+    check_bool(self@streaming, arg = "streaming", call = call)
+  }
+)
+
+#' @rdname QueryOptFlags
+pl__QueryOptFlags <- QueryOptFlags
+
+# TODO: moved to generated file
+# `local({})` is needed for `$` methods.
+# https://github.com/RConsortium/S7/issues/390
+local({
+  method(`$`, QueryOptFlags) <- function(self, name) {
+    method_names <- names(polars_query_opt_flags__methods)
+
+    if (name %in% method_names) {
+      fn <- polars_query_opt_flags__methods[[name]]
+      environment(fn) <- environment()
+      fn
+    } else {
+      # TODO: use abstract class and better handling
+      prop(self, name) |>
+        wrap()
+    }
+  }
+})
+
+#' @exportS3Method utils::.DollarNames
+`.DollarNames.polars::QueryOptFlags` <- function(x, pattern = "") {
+  property_names <- prop_names(x)
+  method_names <- names(polars_query_opt_flags__methods)
+
+  all_names <- c(property_names, method_names)
+  filtered_names <- findMatches(pattern, all_names)
+
+  filtered_names[!startsWith(filtered_names, "_")]
+}
+
+# TODO: This should be a constructor
+# https://github.com/RConsortium/S7/issues/522
+eager_opt_flags <- function() {
+  QueryOptFlags()$no_optimizations() |>
+    set_props_uncheck(
+      eager = TRUE,
+      simplify_expression = TRUE
+    )
+}
+
+# Because of immutability, unlike Python Polars, there are no side effects,
+# and a new modified OptFlags is returned
+QueryOptFlags__no_optimizations <- function() {
+  set_props_uncheck(
+    self,
+    predicate_pushdown = FALSE,
+    projection_pushdown = FALSE,
+    simplify_expression = FALSE,
+    slice_pushdown = FALSE,
+    comm_subplan_elim = FALSE,
+    comm_subexpr_elim = FALSE,
+    cluster_with_columns = FALSE,
+    check_order_observe = FALSE,
+    fast_projection = FALSE
+  )
+}
+
+# nolint end

R/polars-package.R

@@ -3,6 +3,7 @@
 
 ## usethis namespace: start
 #' @import rlang
+#' @import S7
 #' @importFrom utils findMatches head tail
 ## usethis namespace: end
 NULL

R/utils-s7.R

@@ -0,0 +1,17 @@
+# Because validation step is very expensive,
+# this is a copy of `props<-`, the only difference is `check = FALSE`
+# https://github.com/RConsortium/S7/issues/574
+`props_uncheck<-` <- function(object, value) {
+  for (name in names(value)) {
+    prop(object, name, check = FALSE) <- value[[name]]
+  }
+
+  object
+}
+
+# A copy of `set_props()`, using `props_uncheck<-`
+# instead of `props<-` to avoid validation
+set_props_uncheck <- function(object, ...) {
+  props_uncheck(object) <- list(...)
+  object
+}

R/zzz.R

@@ -71,7 +71,8 @@ POLARS_STORE_ENVS <- list(
   "groupby__" = polars_groupby__methods,
   "rolling_groupby__" = polars_rolling_groupby__methods,
   "group_by_dynamic__" = polars_group_by_dynamic__methods,
-  "sql_context__" = polars_sql_context__methods
+  "sql_context__" = polars_sql_context__methods,
+  QueryOptFlags__ = polars_query_opt_flags__methods
 )
 
 lapply(names(POLARS_STORE_ENVS), function(name) {
@@ -92,6 +93,7 @@ utils::globalVariables("self")
 on_load(local_use_cli())
 
 .onLoad <- function(libname, pkgname) {
+  methods_register()
   run_on_load()
 
   # Register S3 methods for optional packages