Implement $sink_ndjson()

NEWS.md

@@ -5,6 +5,7 @@
 ### What's changed
 
 -   New method `$rolling()` for `DataFrame` and `LazyFrame` (#682).
+-   New method `$sink_ndjson()` for LazyFrame (#681).
 
 ## polars 0.12.2
 

R/dataframe__frame.R

@@ -589,6 +589,7 @@ DataFrame_to_series = function(idx = 0) {
 }
 
 #' Sort a DataFrame
+#' @inheritParams DataFrame_unique
 #' @inherit LazyFrame_sort details description params
 #' @return DataFrame
 #' @keywords  DataFrame
@@ -815,6 +816,7 @@ DataFrame_filter = function(...) {
 }
 
 #' Group a DataFrame
+#' @inheritParams DataFrame_unique
 #' @inherit LazyFrame_group_by description params
 #' @keywords DataFrame
 #' @return GroupBy (a DataFrame with special groupby methods like `$agg()`)

R/extendr-wrappers.R

@@ -1033,6 +1033,8 @@ RPolarsLazyFrame$sink_ipc <- function(path, compression_method, maintain_order)
 
 RPolarsLazyFrame$sink_csv <- function(path, include_bom, include_header, separator, line_terminator, quote, batch_size, datetime_format, date_format, time_format, float_precision, null_value, quote_style, maintain_order) .Call(wrap__RPolarsLazyFrame__sink_csv, self, path, include_bom, include_header, separator, line_terminator, quote, batch_size, datetime_format, date_format, time_format, float_precision, null_value, quote_style, maintain_order)
 
+RPolarsLazyFrame$sink_json <- function(path, maintain_order) .Call(wrap__RPolarsLazyFrame__sink_json, self, path, maintain_order)
+
 RPolarsLazyFrame$first <- function() .Call(wrap__RPolarsLazyFrame__first, self)
 
 RPolarsLazyFrame$last <- function() .Call(wrap__RPolarsLazyFrame__last, self)

R/lazyframe__lazy.R

@@ -498,6 +498,7 @@ LazyFrame_collect_in_background = function() {
 #' smaller chunks may reduce memory pressure and improve writing speeds.
 #' @param data_pagesize_limit `NULL` or Integer. If `NULL` (default), the limit
 #' will be ~1MB.
+#' @inheritParams LazyFrame_group_by
 #' @inheritParams DataFrame_unique
 #' @inheritParams LazyFrame_collect
 #'
@@ -575,6 +576,7 @@ LazyFrame_sink_parquet = function(
 #' "lz4" or "zstd". Choose "zstd" for good compression performance. Choose "lz4"
 #' for fast compression/decompression.
 #' @inheritParams LazyFrame_collect
+#' @inheritParams LazyFrame_group_by
 #' @inheritParams DataFrame_unique
 #'
 #' @rdname IO_sink_ipc
@@ -641,6 +643,7 @@ LazyFrame_sink_ipc = function(
 #'
 #' @inheritParams DataFrame_write_csv
 #' @inheritParams LazyFrame_collect
+#' @inheritParams LazyFrame_group_by
 #' @inheritParams DataFrame_unique
 #'
 #' @rdname IO_sink_csv
@@ -721,6 +724,67 @@ LazyFrame_sink_csv = function(
 }
 
 
+#' @title Stream the output of a query to a JSON file
+#' @description
+#' This writes the output of a query directly to a JSON file without collecting
+#' it in the R session first. This is useful if the output of the query is still
+#' larger than RAM as it would crash the R session if it was collected into R.
+#'
+#' @inheritParams DataFrame_write_csv
+#' @inheritParams LazyFrame_collect
+#' @inheritParams LazyFrame_group_by
+#' @inheritParams DataFrame_unique
+#'
+#' @rdname IO_sink_ndjson
+#'
+#' @examples
+#' # sink table 'mtcars' from mem to JSON
+#' tmpf = tempfile(fileext = ".json")
+#' pl$LazyFrame(mtcars)$sink_ndjson(tmpf)
+#'
+#' # load parquet directly into a DataFrame / memory
+#' pl$scan_ndjson(tmpf)$collect()
+LazyFrame_sink_ndjson = function(
+    path,
+    maintain_order = TRUE,
+    type_coercion = TRUE,
+    predicate_pushdown = TRUE,
+    projection_pushdown = TRUE,
+    simplify_expression = TRUE,
+    slice_pushdown = TRUE,
+    no_optimization = FALSE,
+    inherit_optimization = FALSE) {
+  if (isTRUE(no_optimization)) {
+    predicate_pushdown = FALSE
+    projection_pushdown = FALSE
+    slice_pushdown = FALSE
+  }
+
+  lf = self
+
+  if (isFALSE(inherit_optimization)) {
+    lf = self$set_optimization_toggle(
+      type_coercion,
+      predicate_pushdown,
+      projection_pushdown,
+      simplify_expression,
+      slice_pushdown,
+      comm_subplan_elim = FALSE,
+      comm_subexpr_elim = FALSE,
+      streaming = FALSE
+    ) |> unwrap("in $sink_ndjson()")
+  }
+
+  lf |>
+    .pr$LazyFrame$sink_json(
+      path,
+      maintain_order
+    ) |>
+    unwrap("in $sink_ndjson()") |>
+    invisible()
+}
+
+
 #' @title Limit a LazyFrame
 #' @inherit DataFrame_limit description params details
 #' @return A `LazyFrame`
@@ -973,7 +1037,10 @@ LazyFrame_unique = function(subset = NULL, keep = "first", maintain_order = FALS
 #' (`$agg()`, `$filter()`, etc.).
 #' @keywords LazyFrame
 #' @param ... Any Expr(s) or string(s) naming a column.
-#' @inheritParams DataFrame_unique
+#' @param maintain_order Keep the same order as the original `LazyFrame`. Setting
+#'  this to `TRUE` makes it more expensive to compute and blocks the possibility
+#'  to run on the streaming engine. The default value can be changed with
+#' `pl$set_options(maintain_order = TRUE)`.
 #' @return LazyGroupBy (a LazyFrame with special groupby methods like `$agg()`)
 #' @examples
 #' pl$LazyFrame(
@@ -1049,6 +1116,7 @@ LazyFrame_join = function(
 #' either of length 1 or a logical vector of the same length as the number of
 #' Expr(s) specified in `by` and `...`.
 #' @param nulls_last Boolean. Place `NULL`s at the end? Default is `FALSE`.
+#' @inheritParams LazyFrame_group_by
 #' @inheritParams DataFrame_unique
 #' @return LazyFrame
 #' @keywords  LazyFrame

src/rust/src/lazy/dataframe.rs

@@ -175,6 +175,15 @@ impl RPolarsLazyFrame {
             .map_err(polars_to_rpolars_err)
     }
 
+    fn sink_json(&self, path: Robj, maintain_order: Robj) -> RResult<()> {
+        let maintain_order = robj_to!(bool, maintain_order)?;
+        let options = pl::JsonWriterOptions { maintain_order };
+        self.0
+            .clone()
+            .sink_json(robj_to!(String, path)?.into(), options)
+            .map_err(polars_to_rpolars_err)
+    }
+
     fn first(&self) -> Self {
         self.0.clone().first().into()
     }

tests/testthat/_snaps/after-wrappers.md

@@ -151,13 +151,13 @@
       [35] "select"                  "set_optimization_toggle"
       [37] "shift"                   "shift_and_fill"         
       [39] "sink_csv"                "sink_ipc"               
-      [41] "sink_parquet"            "slice"                  
-      [43] "sort"                    "std"                    
-      [45] "sum"                     "tail"                   
-      [47] "unique"                  "unnest"                 
-      [49] "var"                     "width"                  
-      [51] "with_columns"            "with_context"           
-      [53] "with_row_count"         
+      [41] "sink_ndjson"             "sink_parquet"           
+      [43] "slice"                   "sort"                   
+      [45] "std"                     "sum"                    
+      [47] "tail"                    "unique"                 
+      [49] "unnest"                  "var"                    
+      [51] "width"                   "with_columns"           
+      [53] "with_context"            "with_row_count"         
 
 ---
 
@@ -183,13 +183,13 @@
       [33] "select"                  "select_str_as_lit"      
       [35] "set_optimization_toggle" "shift"                  
       [37] "shift_and_fill"          "sink_csv"               
-      [39] "sink_ipc"                "sink_parquet"           
-      [41] "slice"                   "sort_by_exprs"          
-      [43] "std"                     "sum"                    
-      [45] "tail"                    "unique"                 
-      [47] "unnest"                  "var"                    
-      [49] "with_columns"            "with_context"           
-      [51] "with_row_count"         
+      [39] "sink_ipc"                "sink_json"              
+      [41] "sink_parquet"            "slice"                  
+      [43] "sort_by_exprs"           "std"                    
+      [45] "sum"                     "tail"                   
+      [47] "unique"                  "unnest"                 
+      [49] "var"                     "with_columns"           
+      [51] "with_context"            "with_row_count"         
 
 # public and private methods of each class Expr
 

tests/testthat/_snaps/sink_stream.md

@@ -30,3 +30,24 @@
       a,b,c
       "foo",1.0,a
 
+# sink_ndjson works
+
+    Code
+      cat(readLines(path, warn = FALSE), sep = "\n")
+    Output
+      {"drat":3.9,"mpg":21.0}
+      {"drat":3.9,"mpg":21.0}
+      {"drat":3.85,"mpg":22.8}
+      {"drat":3.08,"mpg":21.4}
+      {"drat":3.15,"mpg":18.7}
+      {"drat":2.76,"mpg":18.1}
+      {"drat":3.21,"mpg":14.3}
+      {"drat":3.69,"mpg":24.4}
+      {"drat":3.92,"mpg":22.8}
+      {"drat":3.92,"mpg":19.2}
+      {"drat":3.92,"mpg":17.8}
+      {"drat":3.07,"mpg":16.4}
+      {"drat":3.07,"mpg":17.3}
+      {"drat":3.07,"mpg":15.2}
+      {"drat":2.93,"mpg":10.4}
+

tests/testthat/test-sink_stream.R

@@ -222,3 +222,12 @@ test_that("sink_csv: float_precision works", {
     data.frame(x = c(1.234, 5.600))
   )
 })
+
+
+# sink_ndjson ---------------------------------------------------------
+
+test_that("sink_ndjson works", {
+  temp_out = tempfile(fileext = ".json")
+  pl$LazyFrame(mtcars)$head(15)$select(pl$col("drat", "mpg"))$sink_ndjson(temp_out)
+  expect_snapshot_file(temp_out)
+})