rapi_rel_from_altrep_df(enable_materialization)

duckdb · Oct 27, 2024 · b3420e0 · b3420e0
1 parent 26a4926
commit b3420e0
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 12 deletions.
diff --git a/R/cpp11.R b/R/cpp11.R
@@ -181,8 +181,8 @@ rapi_rel_to_altrep <- function(rel, allow_materialization = TRUE) {
   .Call(`_duckdb_rapi_rel_to_altrep`, rel, allow_materialization)
 }
 
-rapi_rel_from_altrep_df <- function(df, strict, allow_materialized) {
-  .Call(`_duckdb_rapi_rel_from_altrep_df`, df, strict, allow_materialized)
+rapi_rel_from_altrep_df <- function(df, strict, allow_materialized, enable_materialization) {
+  .Call(`_duckdb_rapi_rel_from_altrep_df`, df, strict, allow_materialized, enable_materialization)
 }
 
 rapi_release <- function(stmt) {

diff --git a/R/relational.R b/R/relational.R
@@ -389,15 +389,20 @@ rel_to_altrep <- function(rel) {
 
 #' Retrieves the data frame back from a altrep df
 #' @param df the data frame created by rel_to_altrep
+#' @param strict whether to throw an error if the data frame is not an altrep
+#'   or if other criteria are not met
+#' @param allow_materialized whether to succeed if the data frame is already materialized
+#' @param enable_materialization set to `TRUE` for side effect: allow materialization of this a data frame
+#'   if it was not allowed previously
 #' @return the relation object
 #' @noRd
 #' @examples
 #' con <- DBI::dbConnect(duckdb())
 #' rel <- rel_from_df(con, mtcars)
 #' df = rel_to_altrep(rel)
 #' print(rel_from_altrep_df(df))
-rel_from_altrep_df <- function(df, strict = TRUE, allow_materialized = TRUE) {
-  rethrow_rapi_rel_from_altrep_df(df, strict, allow_materialized)
+rel_from_altrep_df <- function(df, strict = TRUE, allow_materialized = TRUE, enable_materialization = FALSE) {
+  rethrow_rapi_rel_from_altrep_df(df, strict, allow_materialized, enable_materialization)
 }
 
 

diff --git a/src/cpp11.cpp b/src/cpp11.cpp
@@ -329,10 +329,10 @@ extern "C" SEXP _duckdb_rapi_rel_to_altrep(SEXP rel, SEXP allow_materialization)
   END_CPP11
 }
 // reltoaltrep.cpp
-SEXP rapi_rel_from_altrep_df(SEXP df, bool strict, bool allow_materialized);
-extern "C" SEXP _duckdb_rapi_rel_from_altrep_df(SEXP df, SEXP strict, SEXP allow_materialized) {
+SEXP rapi_rel_from_altrep_df(SEXP df, bool strict, bool allow_materialized, bool enable_materialization);
+extern "C" SEXP _duckdb_rapi_rel_from_altrep_df(SEXP df, SEXP strict, SEXP allow_materialized, SEXP enable_materialization) {
   BEGIN_CPP11
-    return cpp11::as_sexp(rapi_rel_from_altrep_df(cpp11::as_cpp<cpp11::decay_t<SEXP>>(df), cpp11::as_cpp<cpp11::decay_t<bool>>(strict), cpp11::as_cpp<cpp11::decay_t<bool>>(allow_materialized)));
+    return cpp11::as_sexp(rapi_rel_from_altrep_df(cpp11::as_cpp<cpp11::decay_t<SEXP>>(df), cpp11::as_cpp<cpp11::decay_t<bool>>(strict), cpp11::as_cpp<cpp11::decay_t<bool>>(allow_materialized), cpp11::as_cpp<cpp11::decay_t<bool>>(enable_materialization)));
   END_CPP11
 }
 // statement.cpp
@@ -472,7 +472,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_duckdb_rapi_rel_distinct",            (DL_FUNC) &_duckdb_rapi_rel_distinct,            1},
     {"_duckdb_rapi_rel_explain",             (DL_FUNC) &_duckdb_rapi_rel_explain,             1},
     {"_duckdb_rapi_rel_filter",              (DL_FUNC) &_duckdb_rapi_rel_filter,              2},
-    {"_duckdb_rapi_rel_from_altrep_df",      (DL_FUNC) &_duckdb_rapi_rel_from_altrep_df,      3},
+    {"_duckdb_rapi_rel_from_altrep_df",      (DL_FUNC) &_duckdb_rapi_rel_from_altrep_df,      4},
     {"_duckdb_rapi_rel_from_df",             (DL_FUNC) &_duckdb_rapi_rel_from_df,             3},
     {"_duckdb_rapi_rel_from_sql",            (DL_FUNC) &_duckdb_rapi_rel_from_sql,            2},
     {"_duckdb_rapi_rel_from_table",          (DL_FUNC) &_duckdb_rapi_rel_from_table,          3},

diff --git a/src/reltoaltrep.cpp b/src/reltoaltrep.cpp
@@ -74,7 +74,7 @@ static T *GetFromExternalPtr(SEXP x) {
 	}
 	auto wrapper = (T *)R_ExternalPtrAddr(ptr);
 	if (!wrapper) {
-		cpp11::stop("This looks like it has been freed");
+		cpp11::stop("GetFromExternalPtr: This looks like it has been freed");
 	}
 	return wrapper;
 }
@@ -387,7 +387,7 @@ static R_altrep_class_t LogicalTypeToAltrepType(const LogicalType &type) {
 	return data_frame;
 }
 
-[[cpp11::register]] SEXP rapi_rel_from_altrep_df(SEXP df, bool strict = true, bool allow_materialized = true) {
+[[cpp11::register]] SEXP rapi_rel_from_altrep_df(SEXP df, bool strict, bool allow_materialized, bool enable_materialization) {
 	if (!Rf_inherits(df, "data.frame")) {
 		if (strict) {
 			cpp11::stop("rapi_rel_from_altrep_df: Not a data.frame");
@@ -423,9 +423,8 @@ static R_altrep_class_t LogicalTypeToAltrepType(const LogicalType &type) {
 		}
 	}
 
+	auto wrapper = GetFromExternalPtr<AltrepRownamesWrapper>(row_names);
 	if (!allow_materialized) {
-		auto wrapper = GetFromExternalPtr<AltrepRownamesWrapper>(row_names);
-
 		if (wrapper->rel->res.get()) {
 			if (strict) {
 				cpp11::stop("rapi_rel_from_altrep_df: Data frame already materialized");
@@ -444,6 +443,12 @@ static R_altrep_class_t LogicalTypeToAltrepType(const LogicalType &type) {
 		}
 	}
 
+	// Side effect comes last
+	// FIXME: Add separate rapi_() function for this
+	if (enable_materialization) {
+		wrapper->rel->allow_materialization = true;
+	}
+
 	return res;
 }