optimization param cluster_with_columns

pola-rs · May 26, 2024 · d73ea04 · d73ea04
1 parent b2aa7d8
commit d73ea04
Show file tree

Hide file tree

Showing 18 changed files with 59 additions and 32 deletions.
diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R
@@ -1244,7 +1244,7 @@ RPolarsLazyFrame$schema <- function() .Call(wrap__RPolarsLazyFrame__schema, self
 
 RPolarsLazyFrame$fetch <- function(n_rows) .Call(wrap__RPolarsLazyFrame__fetch, self, n_rows)
 
-RPolarsLazyFrame$set_optimization_toggle <- function(type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, streaming, eager) .Call(wrap__RPolarsLazyFrame__set_optimization_toggle, self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, streaming, eager)
+RPolarsLazyFrame$set_optimization_toggle <- function(type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, cluster_with_columns, streaming, eager) .Call(wrap__RPolarsLazyFrame__set_optimization_toggle, self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, cluster_with_columns, streaming, eager)
 
 RPolarsLazyFrame$get_optimization_toggle <- function() .Call(wrap__RPolarsLazyFrame__get_optimization_toggle, self)
 

diff --git a/R/lazyframe__lazy.R b/R/lazyframe__lazy.R
@@ -478,6 +478,7 @@ LazyFrame_set_optimization_toggle = function(
       slice_pushdown,
       comm_subplan_elim,
       comm_subexpr_elim,
+      cluster_with_columns,
       streaming,
       eager
     )
@@ -490,7 +491,7 @@ LazyFrame_set_optimization_toggle = function(
 #' @param ... Ignored.
 #' @param no_optimization  Logical. Sets the following parameters to `FALSE`:
 #'  `predicate_pushdown`, `projection_pushdown`, `slice_pushdown`,
-#'  `comm_subplan_elim`, `comm_subexpr_elim`.
+#'  `comm_subplan_elim`, `comm_subexpr_elim`, `cluster_with_columns`.
 #' @param inherit_optimization  Logical. Use existing optimization settings
 #' regardless the settings specified in this function call.
 #' @param collect_in_background Logical. Detach this query from R session.
@@ -533,6 +534,7 @@ LazyFrame_collect = function(
     slice_pushdown = FALSE
     comm_subplan_elim = FALSE
     comm_subexpr_elim = FALSE
+    cluster_with_columns = FALSE
   }
 
   if (isTRUE(streaming)) {
@@ -552,6 +554,7 @@ LazyFrame_collect = function(
       slice_pushdown,
       comm_subplan_elim,
       comm_subexpr_elim,
+      cluster_with_columns,
       streaming
     ) |> unwrap("in $collect():")
   }
@@ -1604,6 +1607,7 @@ LazyFrame_fetch = function(
     slice_pushdown = FALSE
     comm_subplan_elim = FALSE
     comm_subexpr_elim = FALSE
+    cluster_with_columns = FALSE
   }
 
   if (isTRUE(streaming)) {
@@ -1621,6 +1625,7 @@ LazyFrame_fetch = function(
       slice_pushdown,
       comm_subplan_elim,
       comm_subexpr_elim,
+      cluster_with_columns,
       streaming
     ) |> unwrap("in $fetch()")
   }
@@ -1701,6 +1706,7 @@ LazyFrame_profile = function(
     slice_pushdown = FALSE
     comm_subplan_elim = FALSE
     comm_subexpr_elim = FALSE
+    cluster_with_columns = FALSE
   }
 
   if (isTRUE(streaming)) {
@@ -1718,6 +1724,7 @@ LazyFrame_profile = function(
       slice_pushdown,
       comm_subplan_elim,
       comm_subexpr_elim,
+      cluster_with_columns,
       streaming
     ) |> unwrap("in $profile():")
   }

diff --git a/R/series__series.R b/R/series__series.R
@@ -578,8 +578,8 @@ Series_to_list = \(int64_conversion = polars_options()$int64_conversion) {
 #' @return DataFrame
 #' @examples
 #' as_polars_series(iris$Species, name = "flower species")$value_counts()
-Series_value_counts = function(sort = TRUE, parallel = FALSE) {
-  unwrap(.pr$Series$value_counts(self, sort, parallel), "in $value_counts():")
+Series_value_counts = function(..., sort = TRUE, parallel = FALSE, name = "count") {
+  unwrap(.pr$Series$value_counts(self, sort, parallel, name), "in $value_counts():")
 }
 
 #' Apply every value with an R fun

diff --git a/man/IO_sink_csv.Rd b/man/IO_sink_csv.Rd
diff --git a/man/IO_sink_ipc.Rd b/man/IO_sink_ipc.Rd
diff --git a/man/IO_sink_ndjson.Rd b/man/IO_sink_ndjson.Rd
diff --git a/man/IO_sink_parquet.Rd b/man/IO_sink_parquet.Rd
diff --git a/man/LazyFrame_collect.Rd b/man/LazyFrame_collect.Rd
diff --git a/man/LazyFrame_fetch.Rd b/man/LazyFrame_fetch.Rd
diff --git a/man/LazyFrame_join.Rd b/man/LazyFrame_join.Rd
diff --git a/man/LazyFrame_profile.Rd b/man/LazyFrame_profile.Rd
diff --git a/man/LazyFrame_set_optimization_toggle.Rd b/man/LazyFrame_set_optimization_toggle.Rd
diff --git a/man/LazyFrame_to_dot.Rd b/man/LazyFrame_to_dot.Rd
diff --git a/man/S3_as.data.frame.Rd b/man/S3_as.data.frame.Rd
diff --git a/man/as_polars_df.Rd b/man/as_polars_df.Rd
diff --git a/src/rust/src/lazy/dataframe.rs b/src/rust/src/lazy/dataframe.rs
@@ -553,6 +553,7 @@ impl RPolarsLazyFrame {
         slice_pushdown: Robj,
         comm_subplan_elim: Robj,
         comm_subexpr_elim: Robj,
+        cluster_with_columns: Robj,
         streaming: Robj,
         // fast_projection: Robj, // There is no method like with_fast_projection
         eager: Robj,
@@ -568,7 +569,8 @@ impl RPolarsLazyFrame {
             ._with_eager(robj_to!(bool, eager)?)
             .with_projection_pushdown(robj_to!(bool, projection_pushdown)?)
             .with_comm_subplan_elim(robj_to!(bool, comm_subplan_elim)?)
-            .with_comm_subexpr_elim(robj_to!(bool, comm_subexpr_elim)?);
+            .with_comm_subexpr_elim(robj_to!(bool, comm_subexpr_elim)?)
+            .with_cluster_with_columns(robj_to!(bool, cluster_with_columns)?);
 
         Ok(ldf.into())
     }
@@ -577,13 +579,13 @@ impl RPolarsLazyFrame {
         let pl::OptState {
             projection_pushdown,
             predicate_pushdown,
-            cluster_with_columns,
             type_coercion,
             simplify_expr,
             slice_pushdown,
             file_caching: _,
             comm_subplan_elim,
             comm_subexpr_elim,
+            cluster_with_columns,
             streaming,
             fast_projection: _,
             row_estimate: _,
@@ -593,11 +595,11 @@ impl RPolarsLazyFrame {
             type_coercion = type_coercion,
             predicate_pushdown = predicate_pushdown,
             projection_pushdown = projection_pushdown,
-            cluster_with_columns = cluster_with_columns,
             simplify_expression = simplify_expr,
             slice_pushdown = slice_pushdown,
             comm_subplan_elim = comm_subplan_elim,
             comm_subexpr_elim = comm_subexpr_elim,
+            cluster_with_columns = cluster_with_columns,
             streaming = streaming,
             eager = eager,
         )

diff --git a/tests/testthat/_snaps/lazy.md b/tests/testthat/_snaps/lazy.md
@@ -8,7 +8,7 @@
       
       Naive plan:
       FILTER [(col("a")) == (2)] FROM
-      DF ["a", "b"]; PROJECT */2 COLUMNS; SELECTION: None
+        DF ["a", "b"]; PROJECT */2 COLUMNS; SELECTION: None
 
 # LazyFrame serialize/deseialize
 

diff --git a/tests/testthat/test-lazy.R b/tests/testthat/test-lazy.R
@@ -752,7 +752,7 @@ test_that("fetch", {
     lf$select(pl$col("a") * 2L)$fetch(5)$to_list()
   )
 
-  # uszie input can be bit64
+  # usize input can be bit64
   expect_identical(
     lf$select(pl$col("a") * 2L)$fetch(bit64::as.integer64(5))$to_list(),
     lf$select(pl$col("a") * 2L)$fetch(5)$to_list()
@@ -763,17 +763,6 @@ test_that("fetch", {
     result(lf$select(pl$col("a") * 2L)$fetch(-5)$to_list())$err$contexts(),
     list(BadArgument = "n_rows", ValueOutOfScope = "cannot be less than zero", BadValue = "-5")
   )
-
-
-  # bad opt profile arg streaming
-  expect_identical(
-    result(pl$select(pl$lit(2L) * 2L)$lazy()$fetch(-5, streaming = 42)$to_list())$err$contexts(),
-    list(
-      BadArgument = "streaming",
-      TypeMismatch = "bool",
-      BadValue = "Rvalue: 42.0, Rsexp: Doubles, Rclass: [\"numeric\"]"
-    )
-  )
 })
 
 
@@ -821,6 +810,7 @@ test_that("opt_toggles", {
     slice_pushdown = FALSE,
     comm_subplan_elim = FALSE,
     comm_subexpr_elim = FALSE,
+    cluster_with_columns = TRUE,
     streaming = TRUE,
     eager = TRUE
   )