feat: Reuse intermediate materialization results

Author: krlmlr

R/relational-duckdb.R

@@ -99,7 +99,13 @@ duckdb_rel_from_df <- function(df, call = caller_env()) {
   # FIXME: make generic
   stopifnot(is.data.frame(df))
 
-  rel <- duckdb$rel_from_altrep_df(df, strict = FALSE, allow_materialized = FALSE)
+  rel <- duckdb$rel_from_altrep_df(
+    df,
+    strict = FALSE,
+    allow_materialized = FALSE,
+    wrap = TRUE
+  )
+
   if (!is.null(rel)) {
     # Once we're here, we know it's an ALTREP data frame
     # We don't get here if it's already materialized

tests/testthat/_snaps/compute.md

@@ -7,7 +7,8 @@
       ---------------------
       --- Relation Tree ---
       ---------------------
-      Scan Table [duckplyr_4hYuvhNS26]
+      AltrepDataFrame [0xdeadbeef]
+        Scan Table [duckplyr_4hYuvhNS26]
       
       ---------------------
       -- Result Columns  --

tests/testthat/_snaps/relational-duckdb.md

@@ -42,11 +42,12 @@
       ---------------------
       --- Relation Tree ---
       ---------------------
-      Projection [a as a]
-        Order [___row_number ASC]
-          Filter [(a = 1.0)]
-            Projection [a as a, row_number() OVER () as ___row_number]
-              r_dataframe_scan(0xdeadbeef)
+      AltrepDataFrame [0xdeadbeef]
+        Projection [a as a]
+          Order [___row_number ASC]
+            Filter [(a = 1.0)]
+              Projection [a as a, row_number() OVER () as ___row_number]
+                r_dataframe_scan(0xdeadbeef)
       
       ---------------------
       -- Result Columns  --

tests/testthat/test-compute.R

@@ -1,9 +1,14 @@
 test_that("compute()", {
   set.seed(20241230)
 
+  transform <- function(x) {
+    x <- gsub("0x[0-9a-f]+", "0xdeadbeef", x)
+    x
+  }
+
   df <- duckdb_tibble(x = c(1, 2))
   out <- compute(df)
-  expect_snapshot({
+  expect_snapshot(transform = transform, {
     duckdb_rel_from_df(out)
   })
 

tests/testthat/test-relational-duckdb.R

@@ -138,3 +138,28 @@ test_that("duckdb_rel_from_df() uses materialized results", {
 
   expect_equal(n_calls, 1)
 })
+
+test_that("duckdb_rel_from_df() uses materialized intermediate results", {
+  skip_if(identical(Sys.getenv("R_COVR"), "true"))
+
+  withr::local_envvar(DUCKPLYR_OUTPUT_ORDER = FALSE)
+
+  df1 <- duckdb_tibble(a = 1)
+  df2 <- df1 |> arrange(a)
+  df3 <- df2 |> mutate(b = 2)
+
+  rel2 <- duckdb:::rel_from_altrep_df(df2, wrap = TRUE)
+  expect_length(strsplit(duckdb:::rel_tostring(rel2, "tree"), "\n")[[1]], 4)
+
+  rel3 <- duckdb:::rel_from_altrep_df(df3, wrap = TRUE)
+  expect_length(strsplit(duckdb:::rel_tostring(rel3, "tree"), "\n")[[1]], 6)
+
+  # Side effect: trigger intermediate materialization
+  nrow(df2)
+
+  # The depth of the rel2 tree is shorter thanks to `wrap = TRUE`
+  expect_length(strsplit(duckdb:::rel_tostring(rel2, "tree"), "\n")[[1]], 2)
+
+  # The depth of the rel3 tree is shorter now too
+  expect_length(strsplit(duckdb:::rel_tostring(rel3, "tree"), "\n")[[1]], 4)
+})