vectors & data

duckdb · Jul 20, 2024 · 720ae5f · 720ae5f
1 parent 75b2ed7
commit 720ae5f
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 2 deletions.
diff --git a/alt/bindings/src/duckdb.d.ts b/alt/bindings/src/duckdb.d.ts
@@ -602,7 +602,7 @@ export function data_chunk_set_size(chunk: DataChunk, size: number): void;
 export function vector_get_column_type(vector: Vector): LogicalType;
 
 // void *duckdb_vector_get_data(duckdb_vector vector)
-export function vector_get_data(vector: Vector): Uint8Array;
+export function vector_get_data(vector: Vector, byteCount: number): Buffer;
 
 // uint64_t *duckdb_vector_get_validity(duckdb_vector vector)
 export function vector_get_validity(vector: Vector): BigUint64Array;

diff --git a/alt/bindings/src/duckdb_node_bindings.cpp b/alt/bindings/src/duckdb_node_bindings.cpp
@@ -1,4 +1,5 @@
 #define NODE_ADDON_API_DISABLE_DEPRECATED
+#define NODE_API_NO_EXTERNAL_BUFFERS_ALLOWED
 #include "napi.h"
 
 #include <optional>
@@ -83,6 +84,18 @@ duckdb_result *GetResultFromExternal(Napi::Env env, Napi::Value value) {
   return GetDataFromExternal<duckdb_result>(env, ResultTypeTag, value, "Invalid result argument");
 }
 
+static const napi_type_tag VectorTypeTag = {
+  0x9FE56DE8E3124D07, 0x9ABF31145EDE1C9E
+};
+
+Napi::External<_duckdb_vector> CreateExternalForVector(Napi::Env env, duckdb_vector vector) {
+  return CreateExternal<_duckdb_vector>(env, VectorTypeTag, vector);
+}
+
+duckdb_vector GetVectorFromExternal(Napi::Env env, Napi::Value value) {
+  return GetDataFromExternal<_duckdb_vector>(env, VectorTypeTag, value, "Invalid vector argument");
+}
+
 class PromiseWorker : public Napi::AsyncWorker {
 
 public:
@@ -410,6 +423,13 @@ class DuckDBNodeAddon : public Napi::Addon<DuckDBNodeAddon> {
       InstanceMethod("vector_size", &DuckDBNodeAddon::vector_size),
 
       InstanceMethod("destroy_data_chunk", &DuckDBNodeAddon::destroy_data_chunk),
+      // TODO: data_chunk_reset
+      InstanceMethod("data_chunk_get_column_count", &DuckDBNodeAddon::data_chunk_get_column_count),
+      InstanceMethod("data_chunk_get_vector", &DuckDBNodeAddon::data_chunk_get_vector),
+      InstanceMethod("data_chunk_get_size", &DuckDBNodeAddon::data_chunk_get_size),
+      // TODO: data_chunk_set_size
+
+      InstanceMethod("vector_get_data", &DuckDBNodeAddon::vector_get_data),
 
       InstanceMethod("fetch_chunk", &DuckDBNodeAddon::fetch_chunk),
     });
@@ -739,6 +759,7 @@ class DuckDBNodeAddon : public Napi::Addon<DuckDBNodeAddon> {
   // void duckdb_destroy_logical_type(duckdb_logical_type *type)
 
   // duckdb_data_chunk duckdb_create_data_chunk(duckdb_logical_type *types, idx_t column_count)
+  // TODO
 
   // void duckdb_destroy_data_chunk(duckdb_data_chunk *chunk)
   // function destroy_data_chunk(chunk: DataChunk): void
@@ -750,14 +771,54 @@ class DuckDBNodeAddon : public Napi::Addon<DuckDBNodeAddon> {
   }
 
   // void duckdb_data_chunk_reset(duckdb_data_chunk chunk)
+  // TODO
+
   // idx_t duckdb_data_chunk_get_column_count(duckdb_data_chunk chunk)
+  // function data_chunk_get_column_count(chunk: DataChunk): number
+  Napi::Value data_chunk_get_column_count(const Napi::CallbackInfo& info) {
+    auto env = info.Env();
+    auto chunk = GetDataChunkFromExternal(env, info[0]);
+    auto column_count = duckdb_data_chunk_get_column_count(chunk);
+    return Napi::Number::New(env, column_count);
+  }
+
   // duckdb_vector duckdb_data_chunk_get_vector(duckdb_data_chunk chunk, idx_t col_idx)
+  // function data_chunk_get_vector(chunk: DataChunk, column_index: number): Vector
+  Napi::Value data_chunk_get_vector(const Napi::CallbackInfo& info) {
+    auto env = info.Env();
+    auto chunk = GetDataChunkFromExternal(env, info[0]);
+    auto column_index = info[1].As<Napi::Number>().Uint32Value();
+    auto vector = duckdb_data_chunk_get_vector(chunk, column_index);
+    return CreateExternalForVector(env, vector);
+  }
+
   // idx_t duckdb_data_chunk_get_size(duckdb_data_chunk chunk)
+  // function data_chunk_get_size(chunk: DataChunk): number
+  Napi::Value data_chunk_get_size(const Napi::CallbackInfo& info) {
+    auto env = info.Env();
+    auto chunk = GetDataChunkFromExternal(env, info[0]);
+    auto size = duckdb_data_chunk_get_size(chunk);
+    return Napi::Number::New(env, size);
+  }
+
   // void duckdb_data_chunk_set_size(duckdb_data_chunk chunk, idx_t size)
+  // TODO
 
   // duckdb_logical_type duckdb_vector_get_column_type(duckdb_vector vector)
+  // TODO
+
   // void *duckdb_vector_get_data(duckdb_vector vector)
+  // function vector_get_data(vector: Vector, length: number): Buffer
+  Napi::Value vector_get_data(const Napi::CallbackInfo& info) {
+    auto env = info.Env();
+    auto vector = GetVectorFromExternal(env, info[0]);
+    auto byteCount = info[1].As<Napi::Number>().Uint32Value();
+    auto data = duckdb_vector_get_data(vector);
+    return Napi::Buffer<uint8_t>::NewOrCopy(env, reinterpret_cast<uint8_t*>(data), byteCount);
+  }
+
   // uint64_t *duckdb_vector_get_validity(duckdb_vector vector)
+
   // void duckdb_vector_ensure_validity_writable(duckdb_vector vector)
   // void duckdb_vector_assign_string_element(duckdb_vector vector, idx_t index, const char *str)
   // void duckdb_vector_assign_string_element_len(duckdb_vector vector, idx_t index, const char *str, idx_t str_len)

diff --git a/alt/bindings/test/query.test.ts b/alt/bindings/test/query.test.ts
@@ -17,7 +17,13 @@ suite('query', () => {
           expect(duckdb.column_type(res, 0)).toBe(duckdb.Type.INTEGER);
           const chunk = await duckdb.fetch_chunk(res);
           try {
-            expect(chunk).toBeTruthy();
+            expect(duckdb.data_chunk_get_column_count(chunk)).toBe(1);
+            expect(duckdb.data_chunk_get_size(chunk)).toBe(1);
+            const vector = duckdb.data_chunk_get_vector(chunk, 0);
+            const data = duckdb.vector_get_data(vector, 4);
+            const dv = new DataView(data.buffer);
+            const value = dv.getInt32(0, true);
+            expect(value).toBe(17);
           } finally {
             duckdb.destroy_data_chunk(chunk);
           }
@@ -58,6 +64,20 @@ suite('query', () => {
           expect(duckdb.column_type(res, 0)).toBe(duckdb.Type.BOOLEAN);
           expect(duckdb.column_name(res, 52)).toBe('list_of_fixed_int_array');
           expect(duckdb.column_type(res, 52)).toBe(duckdb.Type.LIST);
+          const chunk = await duckdb.fetch_chunk(res);
+          try {
+            expect(duckdb.data_chunk_get_column_count(chunk)).toBe(53);
+            expect(duckdb.data_chunk_get_size(chunk)).toBe(3);
+            const vector = duckdb.data_chunk_get_vector(chunk, 0);
+            const data = duckdb.vector_get_data(vector, 3);
+            const dv = new DataView(data.buffer);
+            const value0 = dv.getUint8(0) !== 0;
+            expect(value0).toBe(false);
+            const value1 = dv.getUint8(1) !== 0;
+            expect(value1).toBe(true);
+          } finally {
+            duckdb.destroy_data_chunk(chunk);
+          }
         } finally {
           duckdb.destroy_result(res);
         }
@@ -81,6 +101,13 @@ suite('query', () => {
           expect(duckdb.column_count(res)).toBe(1);
           expect(duckdb.column_name(res, 0)).toBe('Count');
           expect(duckdb.column_type(res, 0)).toBe(duckdb.Type.BIGINT);
+          const chunk = await duckdb.fetch_chunk(res);
+          try {
+            expect(duckdb.data_chunk_get_column_count(chunk)).toBe(0);
+            expect(duckdb.data_chunk_get_size(chunk)).toBe(0);
+          } finally {
+            duckdb.destroy_data_chunk(chunk);
+          }
         } finally {
           duckdb.destroy_result(res);
         }
@@ -92,6 +119,18 @@ suite('query', () => {
           expect(duckdb.column_count(res2)).toBe(1);
           expect(duckdb.column_name(res2, 0)).toBe('Count');
           expect(duckdb.column_type(res2, 0)).toBe(duckdb.Type.BIGINT);
+          const chunk = await duckdb.fetch_chunk(res2);
+          try {
+            expect(duckdb.data_chunk_get_column_count(chunk)).toBe(1);
+            expect(duckdb.data_chunk_get_size(chunk)).toBe(1);
+            const vector = duckdb.data_chunk_get_vector(chunk, 0);
+            const data = duckdb.vector_get_data(vector, 8);
+            const dv = new DataView(data.buffer);
+            const value = dv.getBigInt64(0, true);
+            expect(value).toBe(17n);
+          } finally {
+            duckdb.destroy_data_chunk(chunk);
+          }
         } finally {
           duckdb.destroy_result(res2);
         }