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

package/src/statement.cpp

@@ -4,6 +4,9 @@
 
 #include <algorithm>
 #include <cassert>
+#include <iostream>
+#include <string>
+#include <regex>
 
 namespace node_duckdb {
 
@@ -15,8 +18,8 @@ Napi::Object Statement::Init(Napi::Env env, Napi::Object exports) {
 	Napi::Function t =
 	    DefineClass(env, "Statement",
 	                {InstanceMethod("run", &Statement::Run), InstanceMethod("all", &Statement::All),
-	                 InstanceMethod("each", &Statement::Each), InstanceMethod("finalize", &Statement::Finish),
-	                 InstanceMethod("stream", &Statement::Stream)});
+	                 InstanceMethod("arrowIPCAll", &Statement::ArrowIPCAll), InstanceMethod("each", &Statement::Each),
+	                 InstanceMethod("finalize", &Statement::Finish), InstanceMethod("stream", &Statement::Stream)});
 
 	constructor = Napi::Persistent(t);
 	constructor.SuppressDestruct();
@@ -254,7 +257,7 @@ static Napi::Value convert_chunk(Napi::Env &env, std::vector<std::string> names,
 	return scope.Escape(result);
 }
 
-enum RunType { RUN, EACH, ALL };
+enum RunType { RUN, EACH, ALL, ARROW_ALL };
 
 struct StatementParam {
 	std::vector<duckdb::Value> params;
@@ -274,7 +277,8 @@ struct RunPreparedTask : public Task {
 			return;
 		}
 
-		result = statement.statement->Execute(params->params, run_type != RunType::ALL);
+		result =
+		    statement.statement->Execute(params->params, run_type != RunType::ALL && run_type != RunType::ARROW_ALL);
 	}
 
 	void Callback() override {
@@ -349,6 +353,62 @@ struct RunPreparedTask : public Task {
 
 			cb.MakeCallback(statement.Value(), {env.Null(), result_arr});
 		} break;
+		case RunType::ARROW_ALL: {
+			auto materialized_result = (duckdb::MaterializedQueryResult *)result.get();
+			// +1 is for null bytes at end of stream
+			Napi::Array result_arr(Napi::Array::New(env, materialized_result->RowCount() + 1));
+
+			auto deleter = [](Napi::Env, void *finalizeData, void *hint) {
+				delete static_cast<std::shared_ptr<duckdb::QueryResult> *>(hint);
+			};
+
+			std::shared_ptr<duckdb::QueryResult> result_ptr = move(result);
+
+			duckdb::idx_t out_idx = 1;
+			while (true) {
+				auto chunk = result_ptr->Fetch();
+
+				if (!chunk || chunk->size() == 0) {
+					break;
+				}
+
+				D_ASSERT(chunk->ColumnCount() == 2);
+				D_ASSERT(chunk->data[0].GetType() == duckdb::LogicalType::BLOB);
+				D_ASSERT(chunk->data[1].GetType() == duckdb::LogicalType::BOOLEAN);
+
+				for (duckdb::idx_t row_idx = 0; row_idx < chunk->size(); row_idx++) {
+					duckdb::string_t blob = ((duckdb::string_t *)(chunk->data[0].GetData()))[row_idx];
+					bool is_header = chunk->data[1].GetData()[row_idx];
+
+					// Create shared pointer to give (shared) ownership to ArrayBuffer, not that for these materialized
+					// query results, the string data is owned by the QueryResult
+					auto result_ref_ptr = new std::shared_ptr<duckdb::QueryResult>(result_ptr);
+
+					auto array_buffer = Napi::ArrayBuffer::New(env, (void *)blob.GetDataUnsafe(), blob.GetSize(),
+					                                           deleter, result_ref_ptr);
+
+					auto typed_array = Napi::Uint8Array::New(env, blob.GetSize(), array_buffer, 0);
+
+					// TODO we should handle this in duckdb probably
+					if (is_header) {
+						result_arr.Set((uint32_t)0, typed_array);
+					} else {
+						D_ASSERT(out_idx < materialized_result->RowCount());
+						result_arr.Set(out_idx++, typed_array);
+					}
+				}
+			}
+
+			// TODO we should handle this in duckdb probably
+			auto null_arr = Napi::Uint8Array::New(env, 4);
+			memset(null_arr.Data(), '\0', 4);
+			result_arr.Set(out_idx++, null_arr);
+
+			// Confirm all rows are set
+			D_ASSERT(out_idx == materialized_result->RowCount() + 1);
+
+			cb.MakeCallback(statement.Value(), {env.Null(), result_arr});
+		} break;
 		}
 	}
 	std::unique_ptr<duckdb::QueryResult> result;
@@ -423,6 +483,12 @@ Napi::Value Statement::All(const Napi::CallbackInfo &info) {
 	return info.This();
 }
 
+Napi::Value Statement::ArrowIPCAll(const Napi::CallbackInfo &info) {
+	connection_ref->database_ref->Schedule(
+	    info.Env(), duckdb::make_unique<RunPreparedTask>(*this, HandleArgs(info), RunType::ARROW_ALL));
+	return info.This();
+}
+
 Napi::Value Statement::Run(const Napi::CallbackInfo &info) {
 	connection_ref->database_ref->Schedule(info.Env(),
 	                                       duckdb::make_unique<RunPreparedTask>(*this, HandleArgs(info), RunType::RUN));
@@ -471,7 +537,9 @@ Napi::FunctionReference QueryResult::constructor;
 Napi::Object QueryResult::Init(Napi::Env env, Napi::Object exports) {
 	Napi::HandleScope scope(env);
 
-	Napi::Function t = DefineClass(env, "QueryResult", {InstanceMethod("nextChunk", &QueryResult::NextChunk)});
+	Napi::Function t = DefineClass(env, "QueryResult",
+	                               {InstanceMethod("nextChunk", &QueryResult::NextChunk),
+	                                InstanceMethod("nextIpcBuffer", &QueryResult::NextIpcBuffer)});
 
 	constructor = Napi::Persistent(t);
 	constructor.SuppressDestruct();
@@ -521,6 +589,48 @@ struct GetChunkTask : public Task {
 	std::unique_ptr<duckdb::DataChunk> chunk;
 };
 
+struct GetNextArrowIpcTask : public Task {
+	GetNextArrowIpcTask(QueryResult &query_result, Napi::Promise::Deferred deferred)
+	    : Task(query_result), deferred(deferred) {
+	}
+
+	void DoWork() override {
+		auto &query_result = Get<QueryResult>();
+		chunk = query_result.result->Fetch();
+	}
+
+	void DoCallback() override {
+		auto &query_result = Get<QueryResult>();
+		Napi::Env env = query_result.Env();
+		Napi::HandleScope scope(env);
+
+		if (chunk == nullptr || chunk->size() == 0) {
+			deferred.Resolve(env.Null());
+			return;
+		}
+
+		// Arrow IPC streams should be a single column of a single blob
+		D_ASSERT(chunk->size() == 1 && chunk->ColumnCount() == 2);
+		D_ASSERT(chunk->data[0].GetType() == duckdb::LogicalType::BLOB);
+
+		duckdb::string_t blob = *(duckdb::string_t *)(chunk->data[0].GetData());
+
+		// Transfer ownership and Construct ArrayBuffer
+		auto data_chunk_ptr = new std::unique_ptr<duckdb::DataChunk>();
+		*data_chunk_ptr = std::move(chunk);
+		auto deleter = [](Napi::Env, void *finalizeData, void *hint) {
+			delete static_cast<std::unique_ptr<duckdb::DataChunk> *>(hint);
+		};
+		auto array_buffer =
+		    Napi::ArrayBuffer::New(env, (void *)blob.GetDataUnsafe(), blob.GetSize(), deleter, data_chunk_ptr);
+
+		deferred.Resolve(array_buffer);
+	}
+
+	Napi::Promise::Deferred deferred;
+	std::unique_ptr<duckdb::DataChunk> chunk;
+};
+
 Napi::Value QueryResult::NextChunk(const Napi::CallbackInfo &info) {
 	auto env = info.Env();
 	auto deferred = Napi::Promise::Deferred::New(env);
@@ -529,4 +639,12 @@ Napi::Value QueryResult::NextChunk(const Napi::CallbackInfo &info) {
 	return deferred.Promise();
 }
 
+// Should only be called on an arrow ipc query
+Napi::Value QueryResult::NextIpcBuffer(const Napi::CallbackInfo &info) {
+	auto env = info.Env();
+	auto deferred = Napi::Promise::Deferred::New(env);
+	database_ref->Schedule(env, duckdb::make_unique<GetNextArrowIpcTask>(*this, deferred));
+	return deferred.Promise();
+}
+
 } // namespace node_duckdb

package/test/open_close.test.js

@@ -150,13 +150,13 @@ describe('open/close', function() {
         });
     });
 
-    describe('closing with unfinalized statements', function(done) {
+    describe('closing with unfinalized statements', function() {
         var completed = false;
         var completedSecond = false;
         var closed = false;
 
         var db;
-        before(function() {
+        before(function(done) {
             db = new sqlite3.Database(':memory:', done);
         });
 

package/test/support/createdb.js

@@ -24,17 +24,19 @@ function createdb(callback) {
         if (callback) callback();
     } else {
         console.log("Creating test database... This may take several minutes.");
-        var db = new sqlite3.Database(db_path);
-        db.serialize(function() {
-            db.run("CREATE TABLE foo (id INT, txt TEXT)");
-            db.run("BEGIN TRANSACTION");
-            var stmt = db.prepare("INSERT INTO foo VALUES(?, ?)");
-            for (var i = 0; i < count; i++) {
-                stmt.run(i, randomString());
-            }
-            stmt.finalize();
-            db.run("COMMIT TRANSACTION", function () {
-                db.close(callback);
+        var db = new sqlite3.Database(db_path, () => {
+            db.serialize(function() {
+                db.run("CREATE TABLE foo (id INT, txt TEXT)");
+                db.run("BEGIN TRANSACTION");
+                var stmt = db.prepare("INSERT INTO foo VALUES(?, ?)");
+                for (var i = 0; i < count; i++) {
+                    stmt.run(i, randomString());
+                }
+
+                stmt.finalize();
+                db.run("COMMIT TRANSACTION", function () {
+                    db.close(callback);
+                });
             });
         });
     }

package/test/typescript_decls.test.ts

@@ -95,7 +95,7 @@ describe("TypeScript declarataions", function () {
   });
 
   it("typescript: ternary int udf", function (done) {
-    db.register(
+    db.register_udf(
       "udf",
       "integer",
       (x: number, y: number, z: number) => x + y + z
@@ -107,7 +107,7 @@ describe("TypeScript declarataions", function () {
         assert.equal(rows[0].v, 42);
       }
     );
-    db.unregister("udf", done);
+    db.unregister_udf("udf", done);
   });
   it("typescript: retrieve 100,000 rows with Statement#each", function (done) {
     var total = 100000;

package/test/udf.test.js

@@ -9,142 +9,142 @@ describe('UDFs', function() {
         });
 
         it('0ary int', function(done) {
-            db.register("udf", "integer", () => 42);
+            db.register_udf("udf", "integer", () => 42);
             db.all("select udf() v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 42);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('0ary double', function(done) {
-            db.register("udf", "double", () => 4.2);
+            db.register_udf("udf", "double", () => 4.2);
             db.all("select udf() v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 4.2);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('0ary string', function(done) {
-            db.register("udf", "varchar", () => 'hello');
+            db.register_udf("udf", "varchar", () => 'hello');
             db.all("select udf() v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 'hello');
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('0ary int null', function(done) {
-            db.register("udf", "integer", () => undefined);
+            db.register_udf("udf", "integer", () => undefined);
             db.all("select udf() v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, undefined);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
 
         it('0ary string null', function(done) {
-            db.register("udf", "varchar", () => undefined);
+            db.register_udf("udf", "varchar", () => undefined);
             db.all("select udf() v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, undefined);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
 
         it('unary int', function(done) {
-            db.register("udf", "integer", (x) => x+1);
+            db.register_udf("udf", "integer", (x) => x+1);
             db.all("select udf(42) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 43);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('unary double', function(done) {
-            db.register("udf", "double", (x) => x);
+            db.register_udf("udf", "double", (x) => x);
             db.all("select udf(4.2::double) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 4.2);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('unary int null', function(done) {
-            db.register("udf", "integer", (x) => undefined);
+            db.register_udf("udf", "integer", (x) => undefined);
             db.all("select udf(42) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, undefined);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
 
         it('unary double null', function(done) {
-            db.register("udf", "double", (x) => undefined);
+            db.register_udf("udf", "double", (x) => undefined);
             db.all("select udf(4.2::double) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, undefined);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
 
         it('unary string', function(done) {
-            db.register("udf", "varchar", (x) => 'hello ' + x);
+            db.register_udf("udf", "varchar", (x) => 'hello ' + x);
             db.all("select udf('world') v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 'hello world');
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('unary string null', function(done) {
-            db.register("udf", "varchar", (x) => undefined);
+            db.register_udf("udf", "varchar", (x) => undefined);
             db.all("select udf('world') v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, undefined);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('binary int', function(done) {
-            db.register("udf", "integer", (x, y) => x + y);
+            db.register_udf("udf", "integer", (x, y) => x + y);
             db.all("select udf(40, 2) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 42);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('binary string', function(done) {
-            db.register("udf", "varchar", (x, y) => x + ' ' + y);
+            db.register_udf("udf", "varchar", (x, y) => x + ' ' + y);
             db.all("select udf('hello', 'world') v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 'hello world');
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('ternary int', function(done) {
-            db.register("udf", "integer", (x, y, z) => x + y + z);
+            db.register_udf("udf", "integer", (x, y, z) => x + y + z);
             db.all("select udf(21, 20, 1) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 42);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('unary larger series', function(done) {
-            db.register("udf", "integer", (x) => 1);
+            db.register_udf("udf", "integer", (x) => 1);
             db.all("select sum(udf(range::double)) v from range(10000)", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 10000);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
     });
 
@@ -155,60 +155,60 @@ describe('UDFs', function() {
         });
 
         it('tinyint', function(done) {
-            db.register("udf", "integer", (x) => x+1);
+            db.register_udf("udf", "integer", (x) => x+1);
             db.all("select udf(42::tinyint) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 43);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('smallint', function(done) {
-            db.register("udf", "integer", (x) => x+1);
+            db.register_udf("udf", "integer", (x) => x+1);
             db.all("select udf(42::smallint) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 43);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('int', function(done) {
-            db.register("udf", "integer", (x) => x+1);
+            db.register_udf("udf", "integer", (x) => x+1);
             db.all("select udf(42::integer) v", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].v, 43);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('timestamp', function(done) {
-            db.register("udf", "timestamp", (x) => x);
+            db.register_udf("udf", "timestamp", (x) => x);
             db.all("select udf(timestamp '1992-09-20 11:30:00') v", function(err, rows) {
                 if (err) throw err;
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('struct', function(done) {
-            db.register("udf", "integer", a => {
+            db.register_udf("udf", "integer", a => {
                 return (a.x == null ? -100 : a.x);
             });
             db.all("SELECT min(udf({'x': (case when v % 2 = 0 then v else null end)::INTEGER, 'y': 42}))::INTEGER as foo FROM generate_series(1, 10000) as t(v)", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].foo, -100);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
 
         it('structnestednull', function(done) {
-            db.register("udf", "integer", a => {
+            db.register_udf("udf", "integer", a => {
                 return (a.x == null ? -100 : a.x.y);
             });
             db.all("SELECT min(udf({'x': (case when v % 2 = 0 then {'y': v::INTEGER } else null end), 'z': 42}))::INTEGER as foo FROM generate_series(1, 10000) as t(v)", function(err, rows) {
                 if (err) throw err;
                 assert.equal(rows[0].foo, -100);
             });
-            db.unregister("udf", done);
+            db.unregister_udf("udf", done);
         });
     });
 });