feat: <SQLContext>$register_globals() (#1064)

NEWS.md

@@ -2,6 +2,10 @@
 
 ## Polars R Package (development version)
 
+### New features
+
+- New method `<SQLContext>$register_globals()` (#1064).
+
 ## Polars R Package 0.16.2
 
 ### New features

R/sql.R

@@ -29,10 +29,10 @@ print.RPolarsSQLContext = function(x, ...) {
 
 
 #' Initialise a new SQLContext
-#' @name pl_SQLContext
-#' @description Create a new SQLContext and register the given LazyFrames.
+#'
+#' Create a new SQLContext and register the given LazyFrames.
 #' @param ... Name-value pairs of [LazyFrame][LazyFrame_class] like objects to register.
-#' @return RPolarsSQLContext
+#' @return An [SQLContext][SQLContext_class]
 #' @examplesIf polars_info()$features$sql
 #' ctx = pl$SQLContext(mtcars = mtcars)
 #' ctx
@@ -57,8 +57,9 @@ pl_SQLContext = function(...) {
 }
 
 
-#' @title Execute SQL query against the registered data
-#' @description Parse the given SQL query and execute it against the registered frame data.
+#' Execute SQL query against the registered data
+#'
+#' Parse the given SQL query and execute it against the registered frame data.
 #' @param query A valid string SQL query.
 #' @return A [LazyFrame][LazyFrame_class]
 #' @examplesIf polars_info()$features$sql
@@ -71,11 +72,14 @@ SQLContext_execute = function(query) {
 }
 
 
-#' @title Register a single data as a table
-#' @description Register a single frame as a table, using the given name.
+#' Register a single data as a table
+#'
+#' Register a single frame as a table, using the given name.
+#'
+#' If a table with the same name is already registered, it will be overwritten.
 #' @param name A string name to register the frame as.
 #' @param frame A [LazyFrame][LazyFrame_class] like object to register.
-#' @return Returns the [SQLContext_class] object invisibly.
+#' @return Returns the [SQLContext][SQLContext_class] object invisibly.
 #' @examplesIf polars_info()$features$sql
 #' ctx = pl$SQLContext()
 #' ctx$register("mtcars", mtcars)
@@ -88,10 +92,11 @@ SQLContext_register = function(name, frame) {
 }
 
 
-#' @title Register multiple data as tables
-#' @description Register multiple frames as tables.
+#' Register multiple data as tables
+#'
+#' Register multiple frames as tables.
+#' @inherit SQLContext_register details return
 #' @param ... Name-value pairs of [LazyFrame][LazyFrame_class] like objects to register.
-#' @return Returns the [SQLContext_class] object invisibly.
 #' @examplesIf polars_info()$features$sql
 #' ctx = pl$SQLContext()
 #' r_df = mtcars
@@ -125,10 +130,11 @@ SQLContext_register_many = function(...) {
 }
 
 
-#' @title Unregister tables by name
-#' @description Unregister tables by name.
+#' Unregister tables by name
+#'
+#' Unregister tables by name.
+#' @inherit SQLContext_register return
 #' @param names A character vector of table names to unregister.
-#' @return Returns the [SQLContext_class] object invisibly.
 #' @examplesIf polars_info()$features$sql
 #' # Initialise a new SQLContext and register the given tables.
 #' ctx = pl$SQLContext(x = mtcars, y = mtcars, z = mtcars)
@@ -146,8 +152,9 @@ SQLContext_unregister = function(names) {
 }
 
 
-#' @title List registered tables
-#' @description Return a character vector of the registered table names.
+#' List registered tables
+#'
+#' Return a character vector of the registered table names.
 #' @return A character vector of the registered table names.
 #' @examplesIf polars_info()$features$sql
 #' ctx = pl$SQLContext()
@@ -160,3 +167,38 @@ SQLContext_tables = function() {
   .pr$SQLContext$get_tables(self) |>
     unwrap("in $tables()")
 }
+
+
+#' Register all polars DataFrames/LazyFrames found in the environment
+#'
+#' Automatically maps variable names to table names.
+#' @inherit SQLContext_register details return
+#' @param ... Ignored.
+#' @param envir The environment to search for polars DataFrames/LazyFrames.
+#' @seealso
+#' - [`<SQLContext>$register()`][SQLContext_register]
+#' - [`<SQLContext>$register_many()`][SQLContext_register_many]
+#' - [`<SQLContext>$unregister()`][SQLContext_unregister]
+#' @examplesIf polars_info()$features$sql
+#' df1 = pl$DataFrame(a = 1:3, b = c("x", NA, "z"))
+#' df2 = pl$LazyFrame(a = 2:4, c = c("t", "w", "v"))
+#'
+#' # Register frames directly from variables found in the current environment.
+#' ctx = pl$SQLContext()$register_globals()
+#' ctx$tables()
+#'
+#' ctx$execute(
+#'   "SELECT a, b, c FROM df1 LEFT JOIN df2 USING (a) ORDER BY a DESC"
+#' )$collect()
+SQLContext_register_globals = function(..., envir = parent.frame()) {
+  Filter(
+    \(x) inherits(get(x, envir = envir), c("RPolarsDataFrame", "RPolarsLazyFrame")),
+    ls(envir = envir)
+  ) |>
+    sapply(\(x) get(x, envir = envir), simplify = FALSE, USE.NAMES = TRUE) |>
+    do.call(self$register_many, args = _) |>
+    result() |>
+    unwrap("in $register_globals()")
+
+  invisible(self)
+}

tests/testthat/_snaps/after-wrappers.md

@@ -745,8 +745,8 @@
     Code
       ls(.pr$env[[class_name]])
     Output
-      [1] "execute"       "register"      "register_many" "tables"       
-      [5] "unregister"   
+      [1] "execute"          "register"         "register_globals" "register_many"   
+      [5] "tables"           "unregister"      
 
 ---
 

tests/testthat/test-sql.R

@@ -46,3 +46,32 @@ test_that("SQLContext_register, register_many, unregister", {
   ctx$unregister(c("mtcars", "mtcars2"))
   expect_equal(sort(ctx$tables()), sort(c("mtcars3", "mtcars4")))
 })
+
+
+test_that("SQLContext_register_globals", {
+  skip_if_not(polars_info()$features$sql)
+
+  ctx = pl$SQLContext()
+  f1 = pl$DataFrame(x = 1)
+  f2 = pl$LazyFrame(x = 2)
+
+  ctx$register_globals()
+  expect_equal(ctx$tables(), c("f1", "f2"))
+
+  func1 = function(ctx) {
+    f3 = pl$DataFrame(x = 3)
+    ctx$register_globals()
+  }
+  func2 = function(ctx) {
+    f4 = pl$LazyFrame(x = 4)
+    ctx$register_globals(envir = parent.frame())
+  }
+
+  ctx2 = pl$SQLContext()
+  func1(ctx2)
+  expect_equal(ctx2$tables(), c("f3"))
+
+  ctx3 = pl$SQLContext()
+  func2(ctx3)
+  expect_equal(ctx3$tables(), c("f1", "f2"))
+})