Merge pull request #17 from samansmink/switch-to-extension-util

Switch to extension util for loading extension

Author: samansmink

Makefile

@@ -31,6 +31,7 @@ EXTENSION_FLAGS=\
 -DDUCKDB_EXTENSION_NAMES="arrow" \
 -DDUCKDB_EXTENSION_ARROW_PATH="$(PROJ_DIR)" \
 -DDUCKDB_EXTENSION_ARROW_LOAD_TESTS=1 \
+-DDUCKDB_EXTENSION_ARROW_SHOULD_LINK=0 \
 -DDUCKDB_EXTENSION_ARROW_TEST_PATH=$(PROJ_DIR)test \
 -DDUCKDB_EXTENSION_ARROW_INCLUDE_PATH="$(PROJ_DIR)src/include"
 
@@ -46,18 +47,18 @@ clean:
 # Main build
 debug:
 	mkdir -p  build/debug && \
-	cmake $(GENERATOR) $(FORCE_COLOR) $(EXTENSION_FLAGS) -DDUCKDB_EXTENSION_ARROW_SHOULD_LINK=0 ${CLIENT_FLAGS} ${CMAKE_VARS} -DEXTENSION_STATIC_BUILD=1 -DCMAKE_BUILD_TYPE=Debug ${BUILD_FLAGS} -S ./duckdb/ -B build/debug && \
+	cmake $(GENERATOR) $(FORCE_COLOR) $(EXTENSION_FLAGS) ${CLIENT_FLAGS} ${CMAKE_VARS} -DEXTENSION_STATIC_BUILD=1 -DCMAKE_BUILD_TYPE=Debug ${BUILD_FLAGS} -S ./duckdb/ -B build/debug && \
 	cmake --build build/debug --config Debug
 
 release:
 	mkdir -p build/release && \
-	cmake $(GENERATOR) $(FORCE_COLOR) $(EXTENSION_FLAGS) -DDUCKDB_EXTENSION_ARROW_SHOULD_LINK=0 ${CLIENT_FLAGS} ${CMAKE_VARS} -DEXTENSION_STATIC_BUILD=1 -DCMAKE_BUILD_TYPE=Release ${BUILD_FLAGS} -S ./duckdb/ -B build/release && \
+	cmake $(GENERATOR) $(FORCE_COLOR) $(EXTENSION_FLAGS) ${CLIENT_FLAGS} ${CMAKE_VARS} -DEXTENSION_STATIC_BUILD=1 -DCMAKE_BUILD_TYPE=Release ${BUILD_FLAGS} -S ./duckdb/ -B build/release && \
 	cmake --build build/release --config Release
 
 # Client build
-debug_js: CLIENT_FLAGS=-DBUILD_NODE=1 -DBUILD_EXTENSIONS=json -DDUCKDB_EXTENSION_ARROW_SHOULD_LINK=1
+debug_js: CLIENT_FLAGS=-DBUILD_NODE=1
 debug_js: debug
-release_js: CLIENT_FLAGS=-DBUILD_NODE=1 -DBUILD_EXTENSIONS=json -DDUCKDB_EXTENSION_ARROW_SHOULD_LINK=1
+release_js: CLIENT_FLAGS=-DBUILD_NODE=1
 release_js: release
 
 # Main tests
@@ -70,12 +71,13 @@ test_debug: debug
 	./build/release/test/unittest "$(PROJ_DIR)test/*"
 
 # Client tests
+DEBUG_EXT_PATH='$(PROJ_DIR)build/debug/extension/arrow/arrow.duckdb_extension'
+RELEASE_EXT_PATH='$(PROJ_DIR)build/debug/extension/arrow/arrow.duckdb_extension'
 test_js: test_debug_js
 test_debug_js: debug_js
-	cd duckdb/tools/nodejs && npm run test-path -- "../../../test/nodejs/**/*.js"
-
+	cd duckdb/tools/nodejs && ARROW_EXTENSION_BINARY_PATH=$(DEBUG_EXT_PATH) npm run test-path -- "../../../test/nodejs/**/*.js"
 test_release_js: release_js
-	cd duckdb/tools/nodejs && npm run test-path -- "../../../test/nodejs/**/*.js"
+	cd duckdb/tools/nodejs && ARROW_EXTENSION_BINARY_PATH=$(DEBUG_EXT_PATH) npm run test-path -- "../../../test/nodejs/**/*.js"
 
 format:
 	find src/ -iname *.hpp -o -iname *.cpp | xargs clang-format --sort-includes=0 -style=file -i

duckdb

@@ -10,26 +10,16 @@
 #include "duckdb/common/arrow/result_arrow_wrapper.hpp"
 #include "duckdb/common/arrow/arrow_appender.hpp"
 #include "duckdb/common/arrow/arrow_converter.hpp"
+#include "duckdb/main/extension_util.hpp"
 #include "duckdb/parser/parsed_data/create_table_function_info.hpp"
 #include "duckdb/function/table/arrow.hpp"
 #endif
 
 namespace duckdb {
 
 static void LoadInternal(DatabaseInstance &instance) {
-	Connection con(instance);
-	con.BeginTransaction();
-	auto &catalog = Catalog::GetSystemCatalog(*con.context);
-
-	auto to_arrow_fun = ToArrowIPCFunction::GetFunction();
-	CreateTableFunctionInfo to_arrow_ipc_info(to_arrow_fun);
-	catalog.CreateTableFunction(*con.context, &to_arrow_ipc_info);
-
-	auto scan_arrow_fun = ArrowIPCTableFunction::GetFunction();
-	CreateTableFunctionInfo scan_arrow_ipc_info(scan_arrow_fun);
-	catalog.CreateTableFunction(*con.context, &scan_arrow_ipc_info);
-
-	con.Commit();
+        ExtensionUtil::RegisterFunction(instance, ToArrowIPCFunction::GetFunction());
+        ExtensionUtil::RegisterFunction(instance, ArrowIPCTableFunction::GetFunction());
 }
 
 void ArrowExtension::Load(DuckDB &db) {

src/arrow_extension.cpp

@@ -60,13 +60,15 @@ unique_ptr <FunctionData> ArrowIPCTableFunction::ArrowScanBind(ClientContext &co
         if (!schema.release) {
             throw InvalidInputException("arrow_scan: released schema passed");
         }
+        auto arrow_type = GetArrowLogicalType(schema);
         if (schema.dictionary) {
-            res->arrow_convert_data[col_idx] =
-                    make_uniq<ArrowConvertData>(GetArrowLogicalType(schema, res->arrow_convert_data, col_idx));
-            return_types.emplace_back(GetArrowLogicalType(*schema.dictionary, res->arrow_convert_data, col_idx));
+            auto dictionary_type = GetArrowLogicalType(*schema.dictionary);
+            return_types.emplace_back(dictionary_type->GetDuckType());
+            arrow_type->SetDictionary(std::move(dictionary_type));
         } else {
-            return_types.emplace_back(GetArrowLogicalType(schema, res->arrow_convert_data, col_idx));
+            return_types.emplace_back(arrow_type->GetDuckType());
         }
+        res->arrow_table.AddColumn(col_idx, std::move(arrow_type));
         auto format = string(schema.format);
         auto name = string(schema.name);
         if (name.empty()) {
@@ -97,7 +99,7 @@ void ArrowIPCTableFunction::ArrowScanFunction(ClientContext &context, TableFunct
             MinValue<int64_t>(STANDARD_VECTOR_SIZE, state.chunk->arrow_array.length - state.chunk_offset);
     data.lines_read += output_size;
     output.SetCardinality(output_size);
-    ArrowTableFunction::ArrowToDuckDB(state, data.arrow_convert_data, output, data.lines_read - output_size, false);
+    ArrowTableFunction::ArrowToDuckDB(state, data.arrow_table.GetColumns(), output, data.lines_read - output_size, false);
     output.Verify();
     state.chunk_offset += output.size();
 }

src/arrow_scan_ipc.cpp

@@ -22,7 +22,7 @@
 #include "duckdb/common/arrow/arrow_converter.hpp"
 #include "duckdb/parser/parsed_data/create_table_function_info.hpp"
 #include "duckdb/function/table/arrow.hpp"
-#include "duckdb/common/arrow/arrow_options.hpp"
+#include "duckdb/main/client_properties.hpp"
 #endif
 
 namespace duckdb {
@@ -72,10 +72,7 @@ unique_ptr<FunctionData> ToArrowIPCFunction::Bind(ClientContext &context, TableF
 
     // Create the Arrow schema
     ArrowSchema schema;
-    ArrowOptions options;
-    options.time_zone = context.GetClientProperties().time_zone;
-    options.offset_size = ArrowOffsetSize::REGULAR;
-    ArrowConverter::ToArrowSchema(&schema, input.input_table_types, input.input_table_names, options);
+    ArrowConverter::ToArrowSchema(&schema, input.input_table_types, input.input_table_names, context.GetClientProperties());
     result->schema = arrow::ImportSchema(&schema).ValueOrDie();
 
     return std::move(result);
@@ -91,9 +88,6 @@ OperatorResultType ToArrowIPCFunction::Function(ExecutionContext &context, Table
     bool sending_schema = false;
 
     bool caching_disabled = context.pipeline && !context.pipeline->GetSink();
-    ArrowOptions options;
-    options.offset_size = ArrowOffsetSize::REGULAR;
-//    bool caching_disabled = true;
 
     if (!local_state.checked_schema) {
         if (!global_state.sent_schema) {
@@ -113,7 +107,7 @@ OperatorResultType ToArrowIPCFunction::Function(ExecutionContext &context, Table
         output.data[1].SetValue(0, Value::BOOLEAN(1));
     } else {
         if (!local_state.appender) {
-            local_state.appender = make_uniq<ArrowAppender>(input.GetTypes(), data.chunk_size, options);
+            local_state.appender = make_uniq<ArrowAppender>(input.GetTypes(), data.chunk_size, context.client.GetClientProperties());
         }
 
         // Append input chunk

src/arrow_to_ipc.cpp

@@ -28,13 +28,25 @@ const to_ipc_functions = {
     'materialized': arrow_ipc_materialized,
 }
 
+function getDatabase() {
+    return new duckdb.Database(':memory:', {"allow_unsigned_extensions":"true"});
+}
+
+function getConnection(db, done) {
+    let conn = new duckdb.Connection(db);
+    conn.exec(`LOAD '${process.env.ARROW_EXTENSION_BINARY_PATH}';`, function (err) {
+        if (err) throw err;
+        done();
+    });
+    return conn
+}
+
 describe(`Arrow IPC`, () => {
     let db;
     let conn;
     before((done) => {
-        db = new duckdb.Database(':memory:', {"allow_unsigned_extensions": "true"});
-        conn = new duckdb.Connection(db);
-        done()
+        db = getDatabase();
+        conn = getConnection(db, () => done())
     });
 
     it(`Basic examples`, async () => {
@@ -110,9 +122,8 @@ for (const [name, fun] of Object.entries(to_ipc_functions)) {
         let db;
         let conn;
         before((done) => {
-            db = new duckdb.Database(':memory:', {"allow_unsigned_extensions": "true"});
-            conn = new duckdb.Connection(db);
-            done()
+            db = getDatabase();
+            conn = getConnection(db, () => done())
         });
 
         it(`Buffers are not garbage collected`, async () => {
@@ -211,12 +222,13 @@ describe('[Benchmark] Arrow IPC Single Int Column (50M tuples)',() => {
     let conn;
 
     before((done) => {
-        db = new duckdb.Database(':memory:',  {"allow_unsigned_extensions":"true"});
-        conn = new duckdb.Connection(db);
-        conn.run("CREATE TABLE test AS select * FROM range(0,?) tbl(i);", column_size, (err) => {
-            assert(!err);
-            done()
-        });
+        db = getDatabase();
+        conn = getConnection(db, () => {
+            conn.run("CREATE TABLE test AS select * FROM range(0,?) tbl(i);", column_size, (err) => {
+                if (err) throw err;
+                done()
+            });
+        })
     });
 
     it('DuckDB table -> DuckDB table', (done) => {
@@ -253,6 +265,13 @@ describe('Buffer registration',() => {
         done();
     });
 
+    before((done) => {
+        db = getDatabase();
+        conn1 = getConnection(db, () => {
+            conn2 = getConnection(db, () => done());
+        })
+    });
+
     it('Buffers can only be overwritten with force flag',  async () => {
         const arrow_buffer = await arrow_ipc_materialized(conn1, "SELECT 1337 as a");
 
@@ -368,9 +387,8 @@ describe('[Benchmark] Arrow IPC TPC-H lineitem.parquet', () => {
     let conn;
 
     before((done) => {
-        db = new duckdb.Database(':memory:',  {"allow_unsigned_extensions":"true"});
-        conn = new duckdb.Connection(db);
-        done();
+        db = getDatabase();
+        conn = getConnection(db, () => done())
     });
 
     it('Parquet -> DuckDB Streaming-> Arrow IPC -> DuckDB Query', async () => {
@@ -455,9 +473,8 @@ for (const [name, fun] of Object.entries(to_ipc_functions)) {
         let db;
         let conn;
         before((done) => {
-            db = new duckdb.Database(':memory:', {"allow_unsigned_extensions": "true"});
-            conn = new duckdb.Connection(db);
-            done();
+            db = getDatabase();
+            conn = getConnection(db, () => done())
         });
 
         for (const query of queries) {