feat: add adbc_insert for bulk ingestion

Author: rustyconover

CLAUDE.md

@@ -21,6 +21,7 @@ The extension provides the following functions:
 ### Query Execution
 - `adbc_scan(handle, query, [params := row(...)])` - Execute a SELECT query and return results as a table. Supports parameterized queries via the optional `params` named parameter.
 - `adbc_execute(handle, query)` - Execute DDL/DML statements (CREATE, INSERT, UPDATE, DELETE). Returns affected row count.
+- `adbc_insert(handle, table_name, <table>, [mode := ...])` - Bulk insert data from a subquery. Modes: 'create', 'append', 'replace', 'create_append'.
 
 ### Catalog Functions
 - `adbc_info(handle)` - Returns driver/database information (vendor name, version, etc.).
@@ -45,6 +46,9 @@ SELECT * FROM adbc_scan(getvariable('conn')::BIGINT, 'SELECT ? AS value', params
 SELECT adbc_execute(getvariable('conn')::BIGINT, 'CREATE TABLE test (id INTEGER, name TEXT)');
 SELECT adbc_execute(getvariable('conn')::BIGINT, 'INSERT INTO test VALUES (1, ''hello'')');
 
+-- Bulk insert from DuckDB query
+SELECT * FROM adbc_insert(getvariable('conn')::BIGINT, 'target', (SELECT * FROM local_table), mode := 'create');
+
 -- Catalog functions
 SELECT * FROM adbc_info(getvariable('conn')::BIGINT);
 SELECT * FROM adbc_tables(getvariable('conn')::BIGINT);

docs/README.md

@@ -163,6 +163,54 @@ SELECT adbc_execute(getvariable('conn')::BIGINT,
     'DELETE FROM products WHERE id = 2');
 ```
 
+### adbc_insert
+
+Bulk insert data from a DuckDB query into a table via ADBC. This is more efficient than executing individual INSERT statements.
+
+```sql
+adbc_insert(connection_id, table_name, <table>, [mode:=]) -> TABLE(rows_inserted BIGINT)
+```
+
+**Parameters:**
+- `connection_id`: Connection handle from `adbc_connect`
+- `table_name`: Target table name in the remote database
+- `<table>`: A subquery providing the data to insert
+- `mode` (optional): Insert mode, one of:
+  - `'create'`: Create the table; error if it exists (default for new tables)
+  - `'append'`: Append to existing table; error if table doesn't exist
+  - `'replace'`: Drop and recreate the table if it exists
+  - `'create_append'`: Create if doesn't exist, append if it does
+
+**Returns:** A table with a single row containing the number of rows inserted.
+
+**Examples:**
+
+```sql
+-- Create a new table and insert data
+SELECT * FROM adbc_insert(getvariable('conn')::BIGINT, 'users',
+    (SELECT id, name, email FROM local_users),
+    mode := 'create');
+
+-- Append data to an existing table
+SELECT * FROM adbc_insert(getvariable('conn')::BIGINT, 'users',
+    (SELECT id, name, email FROM new_users),
+    mode := 'append');
+
+-- Replace an existing table with new data
+SELECT * FROM adbc_insert(getvariable('conn')::BIGINT, 'users',
+    (SELECT * FROM updated_users),
+    mode := 'replace');
+```
+
+Output:
+```
+┌───────────────┐
+│ rows_inserted │
+├───────────────┤
+│          1000 │
+└───────────────┘
+```
+
 ### adbc_info
 
 Returns driver and database information for a connection.

src/adbc_extension.cpp

@@ -21,6 +21,9 @@ static void LoadInternal(ExtensionLoader &loader) {
 	// Register ADBC execute function (adbc_execute for DDL/DML)
 	adbc::RegisterAdbcExecuteFunction(loader.GetDatabaseInstance());
 
+	// Register ADBC insert function (adbc_insert for bulk ingestion)
+	adbc::RegisterAdbcInsertFunction(loader.GetDatabaseInstance());
+
 	QueryFarmSendTelemetry(loader, "adbc", "2025120801");
 }
 

src/adbc_scan.cpp

@@ -10,6 +10,7 @@
 #include "duckdb/common/insertion_order_preserving_map.hpp"
 #include "duckdb/main/client_context.hpp"
 #include <nanoarrow/nanoarrow.h>
+#include <queue>
 
 namespace duckdb {
 namespace adbc {
@@ -578,5 +579,291 @@ void RegisterAdbcExecuteFunction(DatabaseInstance &db) {
     loader.RegisterFunction(adbc_execute_function);
 }
 
+//===--------------------------------------------------------------------===//
+// adbc_insert - Bulk insert data into an ADBC table
+//===--------------------------------------------------------------------===//
+
+struct AdbcInsertBindData : public TableFunctionData {
+    int64_t connection_id;
+    string target_table;
+    string mode;  // "create", "append", "replace", "create_append"
+    shared_ptr<AdbcConnectionWrapper> connection;
+    vector<LogicalType> input_types;
+    vector<string> input_names;
+};
+
+// Custom ArrowArrayStream that we can feed batches into
+// This allows us to use BindStream for proper streaming ingestion
+struct AdbcInsertStream {
+    ArrowArrayStream stream;
+    ArrowSchema schema;
+    bool schema_set = false;
+    queue<ArrowArray> pending_batches;
+    mutex lock;
+    bool finished = false;
+    string last_error;
+
+    AdbcInsertStream() {
+        memset(&stream, 0, sizeof(stream));
+        memset(&schema, 0, sizeof(schema));
+        stream.private_data = this;
+        stream.get_schema = GetSchema;
+        stream.get_next = GetNext;
+        stream.get_last_error = GetLastError;
+        stream.release = Release;
+    }
+
+    ~AdbcInsertStream() {
+        if (schema.release) {
+            schema.release(&schema);
+        }
+        // Release any pending batches
+        while (!pending_batches.empty()) {
+            auto &batch = pending_batches.front();
+            if (batch.release) {
+                batch.release(&batch);
+            }
+            pending_batches.pop();
+        }
+    }
+
+    void SetSchema(ArrowSchema *new_schema) {
+        lock_guard<mutex> l(lock);
+        if (schema.release) {
+            schema.release(&schema);
+        }
+        schema = *new_schema;
+        memset(new_schema, 0, sizeof(*new_schema));  // Transfer ownership
+        schema_set = true;
+    }
+
+    void AddBatch(ArrowArray *batch) {
+        lock_guard<mutex> l(lock);
+        pending_batches.push(*batch);
+        memset(batch, 0, sizeof(*batch));  // Transfer ownership
+    }
+
+    void Finish() {
+        lock_guard<mutex> l(lock);
+        finished = true;
+    }
+
+    static int GetSchema(ArrowArrayStream *stream, ArrowSchema *out) {
+        auto *self = static_cast<AdbcInsertStream *>(stream->private_data);
+        lock_guard<mutex> l(self->lock);
+        if (!self->schema_set) {
+            self->last_error = "Schema not set";
+            return EINVAL;
+        }
+        // Copy the schema (don't transfer ownership)
+        return ArrowSchemaDeepCopy(&self->schema, out);
+    }
+
+    static int GetNext(ArrowArrayStream *stream, ArrowArray *out) {
+        auto *self = static_cast<AdbcInsertStream *>(stream->private_data);
+        lock_guard<mutex> l(self->lock);
+
+        if (self->pending_batches.empty()) {
+            if (self->finished) {
+                // Signal end of stream
+                memset(out, 0, sizeof(*out));
+                return 0;
+            }
+            // No batches available yet - this shouldn't happen in our usage
+            self->last_error = "No batches available";
+            return EAGAIN;
+        }
+
+        *out = self->pending_batches.front();
+        self->pending_batches.pop();
+        return 0;
+    }
+
+    static const char *GetLastError(ArrowArrayStream *stream) {
+        auto *self = static_cast<AdbcInsertStream *>(stream->private_data);
+        return self->last_error.empty() ? nullptr : self->last_error.c_str();
+    }
+
+    static void Release(ArrowArrayStream *stream) {
+        // Don't delete - we manage lifetime externally
+        stream->release = nullptr;
+    }
+};
+
+struct AdbcInsertGlobalState : public GlobalTableFunctionState {
+    mutex lock;
+    shared_ptr<AdbcStatementWrapper> statement;
+    unique_ptr<AdbcInsertStream> insert_stream;
+    int64_t rows_inserted = 0;
+    bool stream_bound = false;
+    bool executed = false;
+    ClientProperties client_properties;
+
+    idx_t MaxThreads() const override {
+        return 1;
+    }
+};
+
+static unique_ptr<FunctionData> AdbcInsertBind(ClientContext &context, TableFunctionBindInput &input,
+                                                vector<LogicalType> &return_types, vector<string> &names) {
+    auto bind_data = make_uniq<AdbcInsertBindData>();
+
+    // First argument is connection handle
+    bind_data->connection_id = input.inputs[0].GetValue<int64_t>();
+    // Second argument is target table name
+    bind_data->target_table = input.inputs[1].GetValue<string>();
+
+    // Check for optional mode parameter (default is "append")
+    auto mode_it = input.named_parameters.find("mode");
+    if (mode_it != input.named_parameters.end() && !mode_it->second.IsNull()) {
+        bind_data->mode = mode_it->second.GetValue<string>();
+        // Validate mode
+        if (bind_data->mode != "create" && bind_data->mode != "append" &&
+            bind_data->mode != "replace" && bind_data->mode != "create_append") {
+            throw InvalidInputException("adbc_insert: Invalid mode '" + bind_data->mode +
+                                         "'. Must be one of: create, append, replace, create_append");
+        }
+    } else {
+        bind_data->mode = "append";  // Default to append
+    }
+
+    // Get connection from registry
+    auto &registry = ConnectionRegistry::Get();
+    bind_data->connection = registry.Get(bind_data->connection_id);
+    if (!bind_data->connection) {
+        throw InvalidInputException("adbc_insert: Invalid connection handle: " + to_string(bind_data->connection_id));
+    }
+
+    if (!bind_data->connection->IsInitialized()) {
+        throw InvalidInputException("adbc_insert: Connection has been closed");
+    }
+
+    // Store input table types and names for Arrow conversion
+    bind_data->input_types = input.input_table_types;
+    bind_data->input_names = input.input_table_names;
+
+    // Return schema: rows_inserted (BIGINT)
+    return_types = {LogicalType::BIGINT};
+    names = {"rows_inserted"};
+
+    return std::move(bind_data);
+}
+
+static unique_ptr<GlobalTableFunctionState> AdbcInsertInitGlobal(ClientContext &context, TableFunctionInitInput &input) {
+    auto &bind_data = input.bind_data->Cast<AdbcInsertBindData>();
+    auto global_state = make_uniq<AdbcInsertGlobalState>();
+
+    // Store client properties for Arrow conversion
+    global_state->client_properties = context.GetClientProperties();
+
+    // Create the statement and set up for bulk ingestion
+    global_state->statement = make_shared_ptr<AdbcStatementWrapper>(bind_data.connection);
+    global_state->statement->Init();
+    global_state->statement->SetOption("adbc.ingest.target_table", bind_data.target_table);
+
+    // Set mode
+    string mode_value;
+    if (bind_data.mode == "create") {
+        mode_value = "adbc.ingest.mode.create";
+    } else if (bind_data.mode == "append") {
+        mode_value = "adbc.ingest.mode.append";
+    } else if (bind_data.mode == "replace") {
+        mode_value = "adbc.ingest.mode.replace";
+    } else if (bind_data.mode == "create_append") {
+        mode_value = "adbc.ingest.mode.create_append";
+    }
+    global_state->statement->SetOption("adbc.ingest.mode", mode_value);
+
+    // Create the insert stream
+    global_state->insert_stream = make_uniq<AdbcInsertStream>();
+
+    // Set up the schema from the input types
+    ArrowSchema schema;
+    ArrowConverter::ToArrowSchema(&schema, bind_data.input_types, bind_data.input_names,
+                                   global_state->client_properties);
+    global_state->insert_stream->SetSchema(&schema);
+
+    // Bind the stream to the statement
+    try {
+        global_state->statement->BindStream(&global_state->insert_stream->stream);
+        global_state->stream_bound = true;
+    } catch (Exception &e) {
+        throw IOException("adbc_insert: Failed to bind stream: " + string(e.what()));
+    }
+
+    return std::move(global_state);
+}
+
+static OperatorResultType AdbcInsertInOut(ExecutionContext &context, TableFunctionInput &data_p,
+                                           DataChunk &input, DataChunk &output) {
+    auto &bind_data = data_p.bind_data->Cast<AdbcInsertBindData>();
+    auto &global_state = data_p.global_state->Cast<AdbcInsertGlobalState>();
+    lock_guard<mutex> l(global_state.lock);
+
+    if (input.size() == 0) {
+        output.SetCardinality(0);
+        return OperatorResultType::NEED_MORE_INPUT;
+    }
+
+    // Convert DuckDB DataChunk to Arrow
+    ArrowAppender appender(bind_data.input_types, input.size(),
+                           global_state.client_properties,
+                           ArrowTypeExtensionData::GetExtensionTypes(context.client, bind_data.input_types));
+    appender.Append(input, 0, input.size(), input.size());
+
+    ArrowArray arr = appender.Finalize();
+
+    // Add the batch to our stream
+    global_state.insert_stream->AddBatch(&arr);
+    global_state.rows_inserted += input.size();
+
+    // Don't output anything during processing - we output the total at the end
+    output.SetCardinality(0);
+    return OperatorResultType::NEED_MORE_INPUT;
+}
+
+static OperatorFinalizeResultType AdbcInsertFinalize(ExecutionContext &context, TableFunctionInput &data_p,
+                                                      DataChunk &output) {
+    auto &global_state = data_p.global_state->Cast<AdbcInsertGlobalState>();
+    lock_guard<mutex> l(global_state.lock);
+
+    // Mark the stream as finished
+    global_state.insert_stream->Finish();
+
+    // Execute the statement to perform the actual insert
+    if (!global_state.executed && global_state.stream_bound) {
+        int64_t rows_affected = -1;
+        try {
+            global_state.statement->ExecuteUpdate(&rows_affected);
+            global_state.executed = true;
+        } catch (Exception &e) {
+            throw IOException("adbc_insert: Failed to execute insert: " + string(e.what()));
+        }
+    }
+
+    // Output the total rows inserted
+    output.SetCardinality(1);
+    output.SetValue(0, 0, Value::BIGINT(global_state.rows_inserted));
+
+    return OperatorFinalizeResultType::FINISHED;
+}
+
+// Register adbc_insert table in-out function
+void RegisterAdbcInsertFunction(DatabaseInstance &db) {
+    ExtensionLoader loader(db, "adbc");
+
+    // adbc_insert(connection_id, table_name, <table>) - Bulk insert data
+    TableFunction adbc_insert_function("adbc_insert",
+                                        {LogicalType::BIGINT, LogicalType::VARCHAR, LogicalType::TABLE},
+                                        nullptr,  // No regular function - use in_out
+                                        AdbcInsertBind,
+                                        AdbcInsertInitGlobal);
+    adbc_insert_function.in_out_function = AdbcInsertInOut;
+    adbc_insert_function.in_out_function_final = AdbcInsertFinalize;
+    adbc_insert_function.named_parameters["mode"] = LogicalType::VARCHAR;
+
+    loader.RegisterFunction(adbc_insert_function);
+}
+
 } // namespace adbc
 } // namespace duckdb