rebase

R/progress_display.R

@@ -0,0 +1,31 @@
+duckdb_progress_display <- function(x) {
+  if (x <= 100) cat(sprintf("\r%3d%%", x))
+  if (x >= 100) cat("\r    ")
+}
+
+check_progress_display <- function(f) {
+  if (is.null(f)) return()
+  if (is.TRUE(f) || isFALSE(f)) return()
+  if (is.function(f)) {
+    if (length(formals(f)) > 0) return()
+    stop("`progress_display` has no argument, expecting at least one.")
+  }
+  stop("`progress_display` is not function, expecting either a boolean or function.")
+}
+
+set_progress_display <- function(f) {
+  check_progress_display(f)
+  options("duckdb.progress_display" = {
+    if (isTRUE(progress_display)) {
+      duckdb_progress_display
+    } else if (is.function(progress_display)) {
+      progress_display
+    } else {
+      NULL
+    }
+  })
+}
+
+get_progress_display <- function(f) {
+  getOption("duckdb.progress_display", default = duckdb_progress_display)
+}

src/connection.cpp

@@ -1,4 +1,6 @@
 #include "rapi.hpp"
+#include "r_progress_bar_display.hpp"
+#include "duckdb/main/client_context.hpp"
 
 using namespace duckdb;
 
@@ -7,6 +9,50 @@ void duckdb::ConnDeleter(ConnWrapper *conn) {
 	delete conn;
 }
 
+unique_ptr<ProgressBarDisplay> RProgressBarDisplay::Create() {
+  return make_uniq<RProgressBarDisplay>();
+}
+
+void RProgressBarDisplay::Initialize() {
+  auto progress_display = Rf_GetOption(RStrings::get().progress_display_sym, R_BaseEnv);
+  if (Rf_isFunction(progress_display)) {
+    progress_callback = progress_display;
+  }
+  D_ASSERT(progress_callback != R_NilValue);
+}
+
+RProgressBarDisplay::RProgressBarDisplay() : ProgressBarDisplay() {
+  // Empty
+}
+
+void RProgressBarDisplay::Update(double percentage) {
+  if (progress_callback == R_NilValue) {
+    Initialize();
+  }
+  if (progress_callback != R_NilValue) {
+    try {
+      cpp11::sexp call = Rf_lang2(progress_callback, Rf_ScalarReal(percentage));
+      cpp11::safe[Rf_eval](call, R_BaseEnv);
+    } catch (std::exception &e) {
+      // Ignore progress bar error
+    }
+  }
+}
+
+void RProgressBarDisplay::Finish() {
+  Update(100);
+}
+
+static void SetDefaultConfigArguments(ClientContext &context) {
+  auto &config = ClientConfig::GetConfig(context);
+  // Set the function used to create the display for the progress bar
+  config.display_create_func = RProgressBarDisplay::Create;
+  auto progress_display = Rf_GetOption(RStrings::get().progress_display_sym, R_BaseEnv);
+  if (Rf_isFunction(progress_display)) {
+    config.enable_progress_bar = true;
+  }
+}
+
 [[cpp11::register]] duckdb::conn_eptr_t rapi_connect(duckdb::db_eptr_t dual) {
 	if (!dual || !dual.get()) {
 		cpp11::stop("rapi_connect: Invalid database reference");
@@ -20,6 +66,10 @@ void duckdb::ConnDeleter(ConnWrapper *conn) {
 	conn_wrapper->conn = make_uniq<Connection>(*db->db);
 	conn_wrapper->db.swap(db);
 
+	// Set progress display config
+	auto &client_context = *conn_wrapper->conn->context;
+	SetDefaultConfigArguments(client_context);
+
 	// The connection now holds a reference to the database.
 	// This reference is released when the connection is closed.
 	// From the R side, the database pointer will remain valid

src/include/r_progress_bar_display.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "rapi.hpp"
+#include "duckdb/common/progress_bar/progress_bar_display.hpp"
+#include "duckdb/common/helper.hpp"
+
+namespace duckdb {
+
+class RProgressBarDisplay : public ProgressBarDisplay {
+public:
+	RProgressBarDisplay();
+	virtual ~RProgressBarDisplay() {
+	}
+
+	static unique_ptr<ProgressBarDisplay> Create();
+
+public:
+	void Update(double percentage) override;
+	void Finish() override;
+
+private:
+	void Initialize();
+
+private:
+	SEXP progress_callback = R_NilValue;
+};
+
+} // namespace duckdb

src/include/rapi.hpp

@@ -171,6 +171,7 @@ struct RStrings {
 	SEXP ImportRecordBatchReader_sym;
 	SEXP materialize_callback_sym;
 	SEXP materialize_message_sym;
+	SEXP progress_display_sym;
 	SEXP duckdb_row_names_sym;
 	SEXP duckdb_vector_sym;
 

src/utils.cpp

@@ -82,6 +82,7 @@ RStrings::RStrings() {
 	Table__from_record_batches_sym = Rf_install("Table__from_record_batches");
 	materialize_message_sym = Rf_install("duckdb.materialize_message");
 	materialize_callback_sym = Rf_install("duckdb.materialize_callback");
+	progress_display_sym = Rf_install("duckdb.progress_display");
 	duckdb_row_names_sym = Rf_install("duckdb_row_names");
 	duckdb_vector_sym = Rf_install("duckdb_vector");
 }

tests/testthat/test-progress-display.R

@@ -0,0 +1,11 @@
+test_that("progress display", {
+
+  expect_error(check_progress_display(5), "expecting either a boolean or function")
+  expect_error(check_progress_display(function(){}), "has no argument, expecting at least one")
+
+  expect_null(check_progress_display(function(x){}))
+  expect_null(check_progress_display(TRUE))
+  expect_null(check_progress_display(FALSE))
+  expect_null(check_progress_display(NULL))
+
+})