feat: New duckdb.materialize_callback option, supersedes get_last_rel() (#589)

Author: krlmlr

R/cpp11.R

@@ -172,10 +172,6 @@ rapi_rel_from_table_function <- function(con, function_name, positional_paramete
   .Call(`_duckdb_rapi_rel_from_table_function`, con, function_name, positional_parameters_sexps, named_parameters_sexps)
 }
 
-rapi_get_last_rel <- function() {
-  .Call(`_duckdb_rapi_get_last_rel`)
-}
-
 # allow_materialization = TRUE: compatibility with duckplyr <= 0.4.1
 rapi_rel_to_altrep <- function(rel, allow_materialization = TRUE) {
   .Call(`_duckdb_rapi_rel_to_altrep`, rel, allow_materialization)

R/relational.R

@@ -86,7 +86,7 @@ rel_from_df <- function(con, df, experimental=FALSE) {
 
 #' @export
 print.duckdb_relation <- function(x, ...) {
-  message("DuckDB Relation: \n", rapi_rel_tostring(x))
+  message("DuckDB Relation: \n", rel_tostring(x))
 }
 
 #' @export
@@ -503,7 +503,3 @@ rel_names <- function(rel) {
 load_rfuns <- function() {
   rethrow_rapi_load_rfuns()
 }
-
-get_last_rel <- function() {
-  rethrow_rapi_get_last_rel()
-}

R/rethrow-gen.R

@@ -387,15 +387,6 @@ rethrow_rapi_rel_from_table_function <- function(con, function_name, positional_
   )
 }
 
-rethrow_rapi_get_last_rel <- function(call = parent.frame(2)) {
-  rlang::try_fetch(
-    rapi_get_last_rel(),
-    error = function(e) {
-      rethrow_error_from_rapi(e, call)
-    }
-  )
-}
-
 rethrow_rapi_rel_to_altrep <- function(rel, allow_materialization = TRUE, call = parent.frame(2)) {
   rlang::try_fetch(
     # duckplyr compat
@@ -585,7 +576,6 @@ rethrow_restore <- function() {
   rethrow_rapi_rel_from_sql <<- rapi_rel_from_sql
   rethrow_rapi_rel_from_table <<- rapi_rel_from_table
   rethrow_rapi_rel_from_table_function <<- rapi_rel_from_table_function
-  rethrow_rapi_get_last_rel <<- rapi_get_last_rel
   rethrow_rapi_rel_to_altrep <<- rapi_rel_to_altrep
   rethrow_rapi_rel_from_altrep_df <<- rapi_rel_from_altrep_df
   rethrow_rapi_release <<- rapi_release

src/cpp11.cpp

@@ -315,13 +315,6 @@ extern "C" SEXP _duckdb_rapi_rel_from_table_function(SEXP con, SEXP function_nam
   END_CPP11
 }
 // reltoaltrep.cpp
-SEXP rapi_get_last_rel();
-extern "C" SEXP _duckdb_rapi_get_last_rel() {
-  BEGIN_CPP11
-    return cpp11::as_sexp(rapi_get_last_rel());
-  END_CPP11
-}
-// reltoaltrep.cpp
 SEXP rapi_rel_to_altrep(duckdb::rel_extptr_t rel, bool allow_materialization);
 extern "C" SEXP _duckdb_rapi_rel_to_altrep(SEXP rel, SEXP allow_materialization) {
   BEGIN_CPP11
@@ -452,7 +445,6 @@ static const R_CallMethodDef CallEntries[] = {
     {"_duckdb_rapi_expr_set_alias",          (DL_FUNC) &_duckdb_rapi_expr_set_alias,          2},
     {"_duckdb_rapi_expr_tostring",           (DL_FUNC) &_duckdb_rapi_expr_tostring,           1},
     {"_duckdb_rapi_expr_window",             (DL_FUNC) &_duckdb_rapi_expr_window,             9},
-    {"_duckdb_rapi_get_last_rel",            (DL_FUNC) &_duckdb_rapi_get_last_rel,            0},
     {"_duckdb_rapi_get_null_SEXP_ptr",       (DL_FUNC) &_duckdb_rapi_get_null_SEXP_ptr,       0},
     {"_duckdb_rapi_get_substrait",           (DL_FUNC) &_duckdb_rapi_get_substrait,           3},
     {"_duckdb_rapi_get_substrait_json",      (DL_FUNC) &_duckdb_rapi_get_substrait_json,      3},

src/include/rapi.hpp

@@ -165,7 +165,8 @@ struct RStrings {
 	SEXP ImportSchema_sym;
 	SEXP ImportRecordBatch_sym;
 	SEXP ImportRecordBatchReader_sym;
-	SEXP materialize_sym;
+	SEXP materialize_callback_sym;
+	SEXP materialize_message_sym;
 	SEXP duckdb_row_names_sym;
 	SEXP duckdb_vector_sym;
 

src/reltoaltrep.cpp

@@ -1,12 +1,12 @@
 #define __STDC_FORMAT_MACROS
 
-#include "httplib.hpp"
 #include "rapi.hpp"
 #include "typesr.hpp"
 #include "reltoaltrep.hpp"
 #include "signal.hpp"
 #include "cpp11/declarations.hpp"
 
+#include "httplib.hpp"
 #include <cinttypes>
 
 using namespace duckdb;
@@ -85,8 +85,8 @@ struct AltrepRelationWrapper {
 		return GetFromExternalPtr<AltrepRelationWrapper>(x);
 	}
 
-	AltrepRelationWrapper(duckdb::shared_ptr<Relation> rel_p, bool allow_materialization_)
-	    : allow_materialization(allow_materialization_), rel(rel_p) {
+	AltrepRelationWrapper(rel_extptr_t rel_, bool allow_materialization_)
+	    : allow_materialization(allow_materialization_), rel_eptr(rel_), rel(rel_->rel) {
 	}
 
 	bool HasQueryResult() const {
@@ -99,12 +99,15 @@ struct AltrepRelationWrapper {
 				cpp11::stop("Materialization is disabled, use collect() or as_tibble() to materialize");
 			}
 
-			auto option = Rf_GetOption(RStrings::get().materialize_sym, R_BaseEnv);
-			if (option != R_NilValue && !Rf_isNull(option) && LOGICAL_ELT(option, 0) == true) {
-				Rprintf("duckplyr: materializing, review details with duckplyr::last_rel()\n");
+			auto materialize_callback = Rf_GetOption(RStrings::get().materialize_callback_sym, R_BaseEnv);
+			if (Rf_isFunction(materialize_callback)) {
+				sexp call = Rf_lang2(materialize_callback, rel_eptr);
+				Rf_eval(call, R_BaseEnv);
+			}
+			else if (Rf_isLogical(materialize_callback) && Rf_length(materialize_callback) == 1 && LOGICAL_ELT(materialize_callback, 0) == true) {
+				// Legacy
+				Rprintf("duckplyr: materializing\n");
 			}
-
-			last_rel = rel;
 
 			ScopedInterruptHandler signal_handler(rel->context.GetContext());
 
@@ -142,14 +145,11 @@ struct AltrepRelationWrapper {
 
 	bool allow_materialization;
 
+	rel_extptr_t rel_eptr;
 	duckdb::shared_ptr<Relation> rel;
 	duckdb::unique_ptr<QueryResult> res;
-
-	static duckdb::shared_ptr<Relation> last_rel;
 };
 
-duckdb::shared_ptr<Relation> AltrepRelationWrapper::last_rel;
-
 struct AltrepRownamesWrapper {
 
 	AltrepRownamesWrapper(duckdb::shared_ptr<AltrepRelationWrapper> rel_p) : rel(rel_p) {
@@ -338,17 +338,12 @@ static R_altrep_class_t LogicalTypeToAltrepType(const LogicalType &type) {
 	}
 }
 
-[[cpp11::register]] SEXP rapi_get_last_rel() {
-	auto last_rel = AltrepRelationWrapper::last_rel;
-	return sexp(make_external_prot<RelationWrapper>("duckdb_relation", R_NilValue, std::move(last_rel)));
-}
-
 [[cpp11::register]] SEXP rapi_rel_to_altrep(duckdb::rel_extptr_t rel, bool allow_materialization) {
 	D_ASSERT(rel && rel->rel);
 	auto drel = rel->rel;
 	auto ncols = drel->Columns().size();
 
-	auto relation_wrapper = make_shared_ptr<AltrepRelationWrapper>(drel, allow_materialization);
+	auto relation_wrapper = make_shared_ptr<AltrepRelationWrapper>(rel, allow_materialization);
 
 	cpp11::writable::list data_frame;
 	data_frame.reserve(ncols);

src/utils.cpp

@@ -80,7 +80,8 @@ RStrings::RStrings() {
 	ImportRecordBatch_sym = Rf_install("ImportRecordBatch");
 	ImportRecordBatchReader_sym = Rf_install("ImportRecordBatchReader");
 	Table__from_record_batches_sym = Rf_install("Table__from_record_batches");
-	materialize_sym = Rf_install("duckdb.materialize_message");
+	materialize_message_sym = Rf_install("duckdb.materialize_message");
+	materialize_callback_sym = Rf_install("duckdb.materialize_callback");
 	duckdb_row_names_sym = Rf_install("duckdb_row_names");
 	duckdb_vector_sym = Rf_install("duckdb_vector");
 }

tests/testthat/_snaps/relational.md

@@ -64,6 +64,27 @@
       Error:
       ! expr_comparison: Invalid comparison operator
 
+# the altrep-conversion for relations works
+
+    Code
+      last_rel
+    Message
+      DuckDB Relation: 
+      ---------------------
+      --- Relation Tree ---
+      ---------------------
+      r_dataframe_scan(0x...)
+      
+      ---------------------
+      -- Result Columns  --
+      ---------------------
+      - Sepal.Length (DOUBLE)
+      - Sepal.Width (DOUBLE)
+      - Petal.Length (DOUBLE)
+      - Petal.Width (DOUBLE)
+      - Species (VARCHAR)
+      
+
 # rel_tostring()
 
     Code

tests/testthat/test-relational.R

@@ -162,6 +162,15 @@ test_that("we can cast R strings to DuckDB strings", {
 })
 
 test_that("the altrep-conversion for relations works", {
+  local_edition(3)
+
+  n_callback <- 0
+  last_rel <- NULL
+  rlang::local_options(duckdb.materialize_callback = function(rel) {
+    n_callback <<- n_callback + 1
+    last_rel <<- rel
+  })
+
   iris$Species <- as.character(iris$Species)
   rel <- rel_from_df(con, iris)
   df <- rel_to_altrep(rel)
@@ -170,8 +179,15 @@ test_that("the altrep-conversion for relations works", {
   expect_true(any(grepl("DUCKDB_ALTREP_REL_VECTOR", inspect_output, fixed = TRUE)))
   expect_true(any(grepl("DUCKDB_ALTREP_REL_ROWNAMES", inspect_output, fixed = TRUE)))
   expect_false(df_is_materialized(df))
+  expect_equal(n_callback, 0)
   dim(df)
   expect_true(df_is_materialized(df))
+
+  expect_snapshot(transform = function(x) gsub("0x[0-9a-f]+", "0x...", x), {
+    last_rel
+  })
+
+  expect_equal(n_callback, 1)
   expect_equal(iris, df)
 })