duckdb 0.5.2-dev1745.0 → 0.5.2-dev1806.0

package/lib/duckdb.d.ts

@@ -20,13 +20,14 @@ export class Connection {
   constructor(db: Database, callback?: Callback<any>);
 
   all(sql: string, ...args: [...any, Callback<TableData>] | []): void;
+  arrowIPCAll(sql: string, ...args: [...any, Callback<TableData>] | []): void;
   each(sql: string, ...args: [...any, Callback<RowData>] | []): void;
   exec(sql: string, ...args: [...any, Callback<void>] | []): void;
 
   prepare(sql: string, ...args: [...any, Callback<Statement>] | []): Statement;
   run(sql: string, ...args: [...any, Callback<void>] | []): Statement;
 
-  register(
+  register_udf(
     name: string,
     return_type: string,
     fun: (...args: any[]) => any
@@ -37,9 +38,10 @@ export class Connection {
     return_type: string,
     fun: (...args: any[]) => any
   ): void;
-  unregister(name: string, callback: Callback<any>): void;
+  unregister_udf(name: string, callback: Callback<any>): void;
 
   stream(sql: any, ...args: any[]): QueryResult;
+  arrowIPCStream(sql: any, ...args: any[]): QueryResult;
 }
 
 export class QueryResult {
@@ -54,24 +56,31 @@ export class Database {
   connect(): Connection;
 
   all(sql: string, ...args: [...any, Callback<TableData>] | []): void;
+  arrowIPCAll(sql: string, ...args: [...any, Callback<TableData>] | []): void;
   each(sql: string, ...args: [...any, Callback<RowData>] | []): void;
   exec(sql: string, ...args: [...any, Callback<void>] | []): void;
 
   prepare(sql: string, ...args: [...any, Callback<Statement>] | []): Statement;
   run(sql: string, ...args: [...any, Callback<void>] | []): Statement;
 
-  register(
+  register_udf(
     name: string,
     return_type: string,
     fun: (...args: any[]) => any
   ): void;
-  unregister(name: string, callback: Callback<any>): void;
+  unregister_udf(name: string, callback: Callback<any>): void;
+
+  stream(sql: any, ...args: any[]): QueryResult;
+  arrowIPCStream(sql: any, ...args: any[]): QueryResult;
 }
 
 export class Statement {
   constructor();
 
   all(...args: [...any, Callback<TableData>] | []): void;
+
+  arrowIPCAll(...args: [...any, Callback<TableData>] | []): void;
+
   each(...args: [...any, Callback<RowData>] | []): void;
 
   finalize(callback?: Callback<void>): void;

package/lib/duckdb.js

@@ -67,6 +67,15 @@ var QueryResult = duckdb.QueryResult;
  */
 QueryResult.prototype.nextChunk;
 
+/**
+ * Function to fetch the next result blob of an Arrow IPC Stream in a zero-copy way.
+ * (requires arrow extension to be loaded)
+ *
+ * @method
+ * @return data chunk
+ */
+QueryResult.prototype.nextIpcBuffer;
+
 /**
  * @name asyncIterator
  * @memberof module:duckdb~QueryResult
@@ -115,6 +124,75 @@ Connection.prototype.all = function (sql) {
     return statement.all.apply(statement, arguments);
 }
 
+// Utility class for streaming Apache Arrow IPC
+class IpcResultStreamIterator {
+    constructor(stream_result_p) {
+        this._depleted = false;
+        this.stream_result = stream_result_p;
+    }
+
+    async next() {
+        if (this._depleted) {
+            return { done: true, value: null };
+        }
+
+        const ipc_raw = await this.stream_result.nextIpcBuffer();
+        const res = new Uint8Array(ipc_raw);
+
+        this._depleted = res.length == 0;
+        return {
+            done: this._depleted,
+            value: res,
+        };
+    }
+
+    [Symbol.asyncIterator]() {
+        return this;
+    }
+
+    // Materialize the IPC stream into a list of Uint8Arrays
+    async toArray () {
+        const retval = []
+
+        for await (const ipc_buf of this) {
+            retval.push(ipc_buf);
+        }
+
+        // Push EOS message containing 4 bytes of 0
+        retval.push(new Uint8Array([0,0,0,0]));
+
+        return retval;
+    }
+}
+
+/**
+ * Run a SQL query and serialize the result into the Apache Arrow IPC format (requires arrow extension to be loaded)
+ * @arg sql
+ * @param {...*} params
+ * @param callback
+ * @return {void}
+ */
+Connection.prototype.arrowIPCAll = function (sql) {
+    const query = "SELECT * FROM to_arrow_ipc((" + sql + "));";
+    var statement = new Statement(this, query);
+    return statement.arrowIPCAll.apply(statement, arguments);
+}
+
+/**
+ * Run a SQL query, returns a IpcResultStreamIterator that allows streaming the result into the Apache Arrow IPC format
+ * (requires arrow extension to be loaded)
+ *
+ * @arg sql
+ * @param {...*} params
+ * @param callback
+ * @return IpcResultStreamIterator
+ */
+Connection.prototype.arrowIPCStream = async function (sql) {
+    const query = "SELECT * FROM to_arrow_ipc((" + sql + "));";
+    const statement = new Statement(this, query);
+    return new IpcResultStreamIterator(await statement.stream.apply(statement, arguments));
+}
+
 /**
  * Runs a SQL query and triggers the callback for each result row
  * @arg sql
@@ -149,9 +227,9 @@ Connection.prototype.stream = async function* (sql) {
  * @return {void}
  * @note this follows the wasm udfs somewhat but is simpler because we can pass data much more cleanly
  */
-Connection.prototype.register = function (name, return_type, fun) {
+Connection.prototype.register_udf = function (name, return_type, fun) {
     // TODO what if this throws an error somewhere? do we need a try/catch?
-    return this.register_bulk(name, return_type, function (desc) {
+    return this.register_udf_bulk(name, return_type, function (desc) {
         try {
             // Build an argument resolver
             const buildResolver = (arg) => {
@@ -289,7 +367,7 @@ Connection.prototype.exec;
  * @param callback
  * @return {void}
  */
-Connection.prototype.register_bulk;
+Connection.prototype.register_udf_bulk;
 /**
  * Unregister a User Defined Function
  *
@@ -299,7 +377,7 @@ Connection.prototype.register_bulk;
  * @param callback
  * @return {void}
  */
-Connection.prototype.unregister;
+Connection.prototype.unregister_udf;
 
 var default_connection = function (o) {
     if (o.default_connection == undefined) {
@@ -308,6 +386,29 @@ var default_connection = function (o) {
     return o.default_connection;
 }
 
+/**
+ * Register a Buffer to be scanned using the Apache Arrow IPC scanner
+ * (requires arrow extension to be loaded)
+ *
+ * @method
+ * @arg name
+ * @arg array
+ * @arg force
+ * @param callback
+ * @return {void}
+ */
+Connection.prototype.register_buffer;
+
+/**
+ * Unregister the Buffer
+ *
+ * @method
+ * @arg name
+ * @param callback
+ * @return {void}
+ */
+Connection.prototype.unregister_buffer;
+
 
 /**
  * Closes database instance
@@ -390,7 +491,18 @@ Database.prototype.run = function () {
 }
 
 /**
- * Convenience method for Connection#each using a built-in default connection
+ * Convenience method for Connection#scanArrowIpc using a built-in default connection
+ * @arg sql
+ * @param {...*} params
+ * @param callback
+ * @return {void}
+ */
+Database.prototype.scanArrowIpc = function () {
+    default_connection(this).scanArrowIpc.apply(this.default_connection, arguments);
+    return this;
+}
+
+/**
  * @arg sql
  * @param {...*} params
  * @param callback
@@ -414,7 +526,29 @@ Database.prototype.all = function () {
 }
 
 /**
- * Convenience method for Connection#exec using a built-in default connection
+ * Convenience method for Connection#arrowIPCAll using a built-in default connection
+ * @arg sql
+ * @param {...*} params
+ * @param callback
+ * @return {void}
+ */
+Database.prototype.arrowIPCAll = function () {
+    default_connection(this).arrowIPCAll.apply(this.default_connection, arguments);
+    return this;
+}
+
+/**
+ * Convenience method for Connection#arrowIPCStream using a built-in default connection
+ * @arg sql
+ * @param {...*} params
+ * @param callback
+ * @return {void}
+ */
+Database.prototype.arrowIPCStream = function () {
+    return default_connection(this).arrowIPCStream.apply(this.default_connection, arguments);
+}
+
+/**
  * @arg sql
  * @param {...*} params
  * @param callback
@@ -426,24 +560,52 @@ Database.prototype.exec = function () {
 }
 
 /**
- * Convenience method for Connection#register using a built-in default connection
+ * Register a User Defined Function
+ *
+ * Convenience method for Connection#register_udf
  * @arg name
  * @arg return_type
  * @arg fun
  * @return {this}
  */
-Database.prototype.register = function () {
-    default_connection(this).register.apply(this.default_connection, arguments);
+Database.prototype.register_udf = function () {
+    default_connection(this).register_udf.apply(this.default_connection, arguments);
+    return this;
+}
+
+/**
+ * Register a buffer containing serialized data to be scanned from DuckDB.
+ *
+ * Convenience method for Connection#unregister_buffer
+ * @arg name
+ * @return {this}
+ */
+Database.prototype.register_buffer = function () {
+    default_connection(this).register_buffer.apply(this.default_connection, arguments);
+    return this;
+}
+
+/**
+ * Unregister a Buffer
+ *
+ * Convenience method for Connection#unregister_buffer
+ * @arg name
+ * @return {this}
+ */
+Database.prototype.unregister_buffer = function () {
+    default_connection(this).unregister_buffer.apply(this.default_connection, arguments);
     return this;
 }
 
 /**
- * Convenience method for Connection#unregister using a built-in default connection
+ * Unregister a UDF
+ *
+ * Convenience method for Connection#unregister_udf
  * @arg name
  * @return {this}
  */
-Database.prototype.unregister = function () {
-    default_connection(this).unregister.apply(this.default_connection, arguments);
+Database.prototype.unregister_udf = function () {
+    default_connection(this).unregister_udf.apply(this.default_connection, arguments);
     return this;
 }
 
@@ -477,6 +639,14 @@ Statement.prototype.run;
  * @return {void}
  */
 Statement.prototype.all;
+/**
+ * @method
+ * @arg sql
+ * @param {...*} params
+ * @param callback
+ * @return {void}
+ */
+Statement.prototype.arrowIPCAll;
 /**
  * @method
  * @arg sql

package/package.json

@@ -2,7 +2,7 @@
   "name": "duckdb",
   "main": "./lib/duckdb.js",
   "types": "./lib/duckdb.d.ts",
-  "version": "0.5.2-dev1745.0",
+  "version": "0.5.2-dev1806.0",
   "description": "DuckDB node.js API",
   "gypfile": true,
   "dependencies": {
@@ -18,7 +18,7 @@
   "scripts": {
     "install": "node-pre-gyp install --fallback-to-build",
     "pretest": "node test/support/createdb.js",
-    "test": "mocha -R spec --timeout 480000",
+    "test": "mocha -R spec --timeout 480000 --expose-gc",
     "pack": "node-pre-gyp package"
   },
   "directories": {
@@ -28,6 +28,7 @@
   "devDependencies": {
     "@types/mocha": "^10.0.0",
     "@types/node": "^18.11.0",
+    "apache-arrow": "^9.0.0",
     "aws-sdk": "^2.790.0",
     "chai": "^4.3.6",
     "jsdoc3-parser": "^2.0.0",

package/src/connection.cpp

@@ -15,8 +15,10 @@ Napi::Object Connection::Init(Napi::Env env, Napi::Object exports) {
 	Napi::Function t =
 	    DefineClass(env, "Connection",
 	                {InstanceMethod("prepare", &Connection::Prepare), InstanceMethod("exec", &Connection::Exec),
-	                 InstanceMethod("register_bulk", &Connection::Register),
-	                 InstanceMethod("unregister", &Connection::Unregister)});
+	                 InstanceMethod("register_udf_bulk", &Connection::RegisterUdf),
+	                 InstanceMethod("register_buffer", &Connection::RegisterBuffer),
+	                 InstanceMethod("unregister_udf", &Connection::UnregisterUdf),
+	                 InstanceMethod("unregister_buffer", &Connection::UnRegisterBuffer)});
 
 	constructor = Napi::Persistent(t);
 	constructor.SuppressDestruct();
@@ -56,6 +58,46 @@ struct ConnectTask : public Task {
 	bool success = false;
 };
 
+struct NodeReplacementScanData : duckdb::ReplacementScanData {
+	NodeReplacementScanData(Connection *con_p) : connection_ref(con_p) {};
+	Connection *connection_ref;
+};
+
+static duckdb::unique_ptr<duckdb::TableFunctionRef>
+ScanReplacement(duckdb::ClientContext &context, const std::string &table_name, duckdb::ReplacementScanData *data) {
+	auto &buffers = ((NodeReplacementScanData *)data)->connection_ref->buffers;
+	// Lookup buffer
+	auto lookup = buffers.find(table_name);
+	if (lookup == buffers.end()) {
+		return nullptr;
+	}
+
+	// Create table scan on ipc buffers
+	auto name = lookup->first;
+	auto ipc_buffer_array = lookup->second;
+
+	auto table_function = duckdb::make_unique<duckdb::TableFunctionRef>();
+	std::vector<duckdb::unique_ptr<duckdb::ParsedExpression>> table_fun_children;
+
+	duckdb::vector<duckdb::Value> list_children;
+
+	for (uint64_t ipc_idx = 0; ipc_idx < ipc_buffer_array.size(); ipc_idx++) {
+		auto &v = ipc_buffer_array[ipc_idx];
+		duckdb::child_list_t<duckdb::Value> struct_children;
+		struct_children.push_back(make_pair("ptr", duckdb::Value::UBIGINT(v.first)));
+		struct_children.push_back(make_pair("size", duckdb::Value::UBIGINT(v.second)));
+
+		// Push struct into table fun
+		list_children.push_back(duckdb::Value::STRUCT(move(struct_children)));
+	}
+
+	table_fun_children.push_back(
+	    duckdb::make_unique<duckdb::ConstantExpression>(duckdb::Value::LIST(move(list_children))));
+	table_function->function =
+	    duckdb::make_unique<duckdb::FunctionExpression>("scan_arrow_ipc", move(table_fun_children));
+	return table_function;
+}
+
 Connection::Connection(const Napi::CallbackInfo &info) : Napi::ObjectWrap<Connection>(info) {
 	Napi::Env env = info.Env();
 	int length = info.Length();
@@ -68,6 +110,15 @@ Connection::Connection(const Napi::CallbackInfo &info) : Napi::ObjectWrap<Connec
 	database_ref = Napi::ObjectWrap<Database>::Unwrap(info[0].As<Napi::Object>());
 	database_ref->Ref();
 
+	if (!database_ref->database) {
+		Napi::Error::New(env, "Connection created on database that was not yet initialized")
+		    .ThrowAsJavaScriptException();
+		return;
+	}
+	// Register replacement scan
+	database_ref->database->instance->config.replacement_scans.emplace_back(
+	    ScanReplacement, duckdb::make_unique<NodeReplacementScanData>(this));
+
 	Napi::Function callback;
 	if (info.Length() > 0 && info[1].IsFunction()) {
 		callback = info[1].As<Napi::Function>();
@@ -217,8 +268,8 @@ void DuckDBNodeUDFLauncher(Napi::Env env, Napi::Function jsudf, std::nullptr_t *
 	jsargs->done = true;
 }
 
-struct RegisterTask : public Task {
-	RegisterTask(Connection &connection, std::string name, std::string return_type_name, Napi::Function callback)
+struct RegisterUdfTask : public Task {
+	RegisterUdfTask(Connection &connection, std::string name, std::string return_type_name, Napi::Function callback)
 	    : Task(connection, callback), name(std::move(name)), return_type_name(std::move(return_type_name)) {
 	}
 
@@ -258,7 +309,7 @@ struct RegisterTask : public Task {
 	std::string return_type_name;
 };
 
-Napi::Value Connection::Register(const Napi::CallbackInfo &info) {
+Napi::Value Connection::RegisterUdf(const Napi::CallbackInfo &info) {
 	auto env = info.Env();
 	if (info.Length() < 3 || !info[0].IsString() || !info[1].IsString() || !info[2].IsFunction()) {
 		Napi::TypeError::New(env, "Holding it wrong").ThrowAsJavaScriptException();
@@ -287,13 +338,69 @@ Napi::Value Connection::Register(const Napi::CallbackInfo &info) {
 	udfs[name] = udf;
 
 	database_ref->Schedule(info.Env(),
-	                       duckdb::make_unique<RegisterTask>(*this, name, return_type_name, completion_callback));
+	                       duckdb::make_unique<RegisterUdfTask>(*this, name, return_type_name, completion_callback));
+
+	return Value();
+}
+
+// Register Arrow IPC buffers for scanning from DuckDB
+Napi::Value Connection::RegisterBuffer(const Napi::CallbackInfo &info) {
+	auto env = info.Env();
+	if (info.Length() < 2 || !info[0].IsString() || !info[1].IsObject()) {
+		Napi::TypeError::New(env, "Incorrect params").ThrowAsJavaScriptException();
+		return env.Null();
+	}
 
+	std::string name = info[0].As<Napi::String>();
+	Napi::Array array = info[1].As<Napi::Array>();
+	bool force_register = false;
+
+	if (info.Length() > 2) {
+		if (!info[2].IsBoolean()) {
+			Napi::TypeError::New(env, "Incorrect params").ThrowAsJavaScriptException();
+			return env.Null();
+		}
+		force_register = info[2].As<Napi::Boolean>().Value();
+	}
+
+	array_references[name] = Napi::Persistent(array);
+
+	if (!force_register && buffers.find(name) != buffers.end()) {
+		Napi::TypeError::New(env, "Buffer with this name already exists").ThrowAsJavaScriptException();
+		return env.Null();
+	}
+
+	buffers[name] = std::vector<std::pair<uint64_t, uint64_t>>();
+
+	for (uint64_t ipc_idx = 0; ipc_idx < array.Length(); ipc_idx++) {
+		Napi::Value v = array[ipc_idx];
+		if (!v.IsObject()) {
+			Napi::TypeError::New(env, "Incorrect params").ThrowAsJavaScriptException();
+			return env.Null();
+		}
+		Napi::Uint8Array arr = v.As<Napi::Uint8Array>();
+		auto raw_ptr = reinterpret_cast<uint64_t>(arr.ArrayBuffer().Data());
+		auto length = (uint64_t)arr.ElementLength();
+
+		buffers[name].push_back(std::pair<uint64_t, uint64_t>({raw_ptr, length}));
+	}
+
+	return Value();
+}
+
+Napi::Value Connection::UnRegisterBuffer(const Napi::CallbackInfo &info) {
+	auto env = info.Env();
+	if (info.Length() != 1 || !info[0].IsString()) {
+		Napi::TypeError::New(env, "Holding it wrong").ThrowAsJavaScriptException();
+		return env.Null();
+	}
+	std::string name = info[0].As<Napi::String>();
+	buffers.erase(name);
 	return Value();
 }
 
-struct UnregisterTask : public Task {
-	UnregisterTask(Connection &connection, std::string name, Napi::Function callback)
+struct UnregisterUdfTask : public Task {
+	UnregisterUdfTask(Connection &connection, std::string name, Napi::Function callback)
 	    : Task(connection, callback), name(std::move(name)) {
 	}
 
@@ -318,7 +425,7 @@ struct UnregisterTask : public Task {
 	std::string name;
 };
 
-Napi::Value Connection::Unregister(const Napi::CallbackInfo &info) {
+Napi::Value Connection::UnregisterUdf(const Napi::CallbackInfo &info) {
 	auto env = info.Env();
 	if (info.Length() < 1 || !info[0].IsString()) {
 		Napi::TypeError::New(env, "Holding it wrong").ThrowAsJavaScriptException();
@@ -331,7 +438,7 @@ Napi::Value Connection::Unregister(const Napi::CallbackInfo &info) {
 		callback = info[1].As<Napi::Function>();
 	}
 
-	database_ref->Schedule(info.Env(), duckdb::make_unique<UnregisterTask>(*this, name, callback));
+	database_ref->Schedule(info.Env(), duckdb::make_unique<UnregisterUdfTask>(*this, name, callback));
 	return Value();
 }