duckdb: don't copy vtab on function bind (#7616)

Store a reference to function's info in bind data to avoid copying vtab.
Derive bind data from base class and not derived one to save a few
bytes.
Replace a few nullptr returns with D_ASSERTs.

Signed-off-by: Mikhail Kot <to@myrrc.dev>

Author: myrrc

vortex-duckdb/cpp/data_chunk.cpp

@@ -1,13 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+#include "duckdb_vx/data_chunk.h"
 #include "duckdb_vx/duckdb_diagnostics.h"
 DUCKDB_INCLUDES_BEGIN
 #include "duckdb/common/types/data_chunk.hpp"
 DUCKDB_INCLUDES_END
 
-#include "duckdb_vx.h"
-
 const char *duckdb_data_chunk_to_string(duckdb_data_chunk chunk, duckdb_vx_error *err) {
     try {
         auto dchunk = reinterpret_cast<duckdb::DataChunk *>(chunk);

vortex-duckdb/cpp/expr.cpp

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-#include "duckdb_vx.h"
-
+#include "duckdb_vx/expr.h"
 #include "duckdb_vx/duckdb_diagnostics.h"
+
 DUCKDB_INCLUDES_BEGIN
 #include "duckdb/planner/expression/bound_between_expression.hpp"
 #include "duckdb/planner/expression/bound_columnref_expression.hpp"

vortex-duckdb/cpp/include/duckdb_vx/copy_function.h

@@ -76,4 +76,4 @@ duckdb_state duckdb_vx_copy_func_register_vtab_one(duckdb_database ffi_db);
 
 #ifdef __cplusplus /* End C ABI */
 };
-#endif
+#endif

vortex-duckdb/cpp/logical_type.cpp

@@ -2,12 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 #include "duckdb_vx/duckdb_diagnostics.h"
-
-DUCKDB_INCLUDES_BEGIN
-#include "duckdb/common/types.hpp"
-DUCKDB_INCLUDES_END
-
-#include "duckdb_vx.h"
 #include "duckdb_vx/logical_type.h"
 
 DUCKDB_INCLUDES_BEGIN
@@ -16,7 +10,7 @@ DUCKDB_INCLUDES_END
 #include <cassert>
 
 duckdb_logical_type duckdb_vx_logical_type_copy(duckdb_logical_type ty) {
-    assert(ty != nullptr);
+    D_ASSERT(ty);
     auto *src = reinterpret_cast<duckdb::LogicalType *>(ty);
     auto copy = duckdb::make_uniq<duckdb::LogicalType>(*src);
     return reinterpret_cast<duckdb_logical_type>(copy.release());

vortex-duckdb/cpp/table_function.cpp

@@ -24,37 +24,45 @@ struct CTableFunctionInfo final : TableFunctionInfo {
     explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) : vtab(vtab) {
     }
 
-    duckdb_vx_tfunc_vtab_t vtab;
+    const duckdb_vx_tfunc_vtab_t vtab;
 };
 
-struct CTableBindData final : TableFunctionData {
-    CTableBindData(unique_ptr<CTableFunctionInfo> info_p,
+struct CTableBindData final : FunctionData {
+    CTableBindData(const CTableFunctionInfo &info,
                    unique_ptr<CData> ffi_data_p,
                    const vector<LogicalType> &types)
-        : info(std::move(info_p)), ffi_data(std::move(ffi_data_p)), types(types) {
+        : info(info), ffi_data(std::move(ffi_data_p)), types(types) {
     }
+    unique_ptr<FunctionData> Copy() const override;
+    bool Equals(const FunctionData &other_base) const override;
 
-    unique_ptr<FunctionData> Copy() const override {
-        D_ASSERT(info->vtab.bind_data_clone);
-
-        duckdb_vx_error error_out = nullptr;
-        const auto copied_ffi_data = info->vtab.bind_data_clone(ffi_data->DataPtr(), &error_out);
-        if (error_out) {
-            throw BinderException(IntoErrString(error_out));
-        }
-
-        auto info_p = make_uniq<CTableFunctionInfo>(info->vtab);
-        auto ffi_data_p = unique_ptr<CData>(reinterpret_cast<CData *>(copied_ffi_data));
-        return make_uniq<CTableBindData>(std::move(info_p), std::move(ffi_data_p), types);
-    }
-
-    unique_ptr<CTableFunctionInfo> info;
+    // Table function info lives for as long as TableFunction is alive as it's
+    // stored inside TableFunction, so it's safe to store a reference.
+    const CTableFunctionInfo &info;
     unique_ptr<CData> ffi_data;
     vector<LogicalType> types;
 };
 
+unique_ptr<FunctionData> CTableBindData::Copy() const {
+    duckdb_vx_error error_out = nullptr;
+    const auto copied_ffi_data = info.vtab.bind_data_clone(ffi_data->DataPtr(), &error_out);
+    if (error_out) {
+        throw BinderException(IntoErrString(error_out));
+    }
+
+    auto ffi_data_p = unique_ptr<CData>(reinterpret_cast<CData *>(copied_ffi_data));
+    return make_uniq<CTableBindData>(info, std::move(ffi_data_p), types);
+}
+
+bool CTableBindData::Equals(const FunctionData &other_base) const {
+    const CTableBindData &other = other_base.Cast<CTableBindData>();
+    // if "types" are different, "ffi_data" would also be different as it
+    // contains types inside, so omit "types" from comparison.
+    return &info == &other.info && ffi_data.get() == other.ffi_data.get();
+}
+
 struct CTableGlobalData final : GlobalTableFunctionState {
-    explicit CTableGlobalData(unique_ptr<CData> ffi_data_p) : ffi_data(std::move(ffi_data_p)) {
+    explicit CTableGlobalData(unique_ptr<CData> ffi_data) : ffi_data(std::move(ffi_data)) {
     }
 
     idx_t MaxThreads() const override {
@@ -65,28 +73,20 @@ struct CTableGlobalData final : GlobalTableFunctionState {
 };
 
 struct CTableLocalData final : LocalTableFunctionState {
-    explicit CTableLocalData(unique_ptr<CData> ffi_data_p) : ffi_data(std::move(ffi_data_p)) {
+    explicit CTableLocalData(unique_ptr<CData> ffi_data) : ffi_data(std::move(ffi_data)) {
     }
 
     unique_ptr<CData> ffi_data;
 };
 
-/**
- * Result of the bind function encapsulates the output schema.
- */
-struct CTableBindResult {
-    vector<LogicalType> &return_types;
-    vector<string> &names;
-};
-
 double c_table_scan_progress(ClientContext &context,
                              const FunctionData *bind_data,
                              const GlobalTableFunctionState *global_state) {
     auto &bind = bind_data->Cast<CTableBindData>();
     duckdb_client_context c_ctx = reinterpret_cast<duckdb_client_context>(&context);
     void *const c_bind_data = bind.ffi_data->DataPtr();
     void *const c_global_state = global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
-    return bind.info->vtab.table_scan_progress(c_ctx, c_bind_data, c_global_state);
+    return bind.info.vtab.table_scan_progress(c_ctx, c_bind_data, c_global_state);
 }
 
 static Value &UnwrapValue(duckdb_value value) {
@@ -148,7 +148,7 @@ c_statistics(ClientContext &context, const FunctionData *bind_data, column_t col
 
     duckdb_client_context c_ctx = reinterpret_cast<duckdb_client_context>(&context);
     duckdb_column_statistics statistics = {};
-    if (!bind.info->vtab.statistics(c_ctx, ffi_bind, column_index, &statistics)) {
+    if (!bind.info.vtab.statistics(c_ctx, ffi_bind, column_index, &statistics)) {
         return {};
     }
 
@@ -186,17 +186,17 @@ c_statistics(ClientContext &context, const FunctionData *bind_data, column_t col
     }
 }
 
+struct CTableBindResult {
+    vector<LogicalType> &return_types;
+    vector<string> &names;
+};
+
 unique_ptr<FunctionData> c_bind(ClientContext &context,
                                 TableFunctionBindInput &input,
                                 vector<LogicalType> &return_types,
                                 vector<string> &names) {
     const auto &info = input.table_function.function_info->Cast<CTableFunctionInfo>();
-
-    // Setup bind info to pass into the callback.
-    CTableBindResult result = {
-        return_types,
-        names,
-    };
+    CTableBindResult result = {return_types, names};
 
     duckdb_vx_error error_out = nullptr;
     auto ctx = reinterpret_cast<duckdb_client_context>(&context);
@@ -208,13 +208,12 @@ unique_ptr<FunctionData> c_bind(ClientContext &context,
         throw BinderException(IntoErrString(error_out));
     }
 
-    return make_uniq<CTableBindData>(make_uniq<CTableFunctionInfo>(info.vtab),
-                                     unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data)),
-                                     return_types);
+    auto cdata = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data));
+    return make_uniq<CTableBindData>(info, std::move(cdata), return_types);
 }
 
 unique_ptr<GlobalTableFunctionState> c_init_global(ClientContext &context, TableFunctionInitInput &input) {
-    const auto &bind = input.bind_data->Cast<CTableBindData>();
+    const CTableBindData &bind = input.bind_data->Cast<CTableBindData>();
 
     duckdb_vx_tfunc_init_input ffi_input = {
         .bind_data = bind.ffi_data->DataPtr(),
@@ -227,7 +226,7 @@ unique_ptr<GlobalTableFunctionState> c_init_global(ClientContext &context, Table
     };
 
     duckdb_vx_error error_out = nullptr;
-    auto ffi_global_data = bind.info->vtab.init_global(&ffi_input, &error_out);
+    duckdb_vx_data ffi_global_data = bind.info.vtab.init_global(&ffi_input, &error_out);
     if (error_out) {
         throw BinderException(IntoErrString(error_out));
     }
@@ -240,7 +239,7 @@ unique_ptr<LocalTableFunctionState> c_init_local(ExecutionContext &context,
                                                  TableFunctionInitInput &input,
                                                  GlobalTableFunctionState *global_state) {
     const auto &bind = input.bind_data->Cast<CTableBindData>();
-    auto global_data = global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
+    void *const ffi_global = global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
 
     duckdb_vx_tfunc_init_input ffi_input = {
         .bind_data = bind.ffi_data->DataPtr(),
@@ -253,25 +252,26 @@ unique_ptr<LocalTableFunctionState> c_init_local(ExecutionContext &context,
     };
 
     duckdb_vx_error error_out = nullptr;
-    auto ffi_local_data = bind.info->vtab.init_local(&ffi_input, global_data, &error_out);
+    duckdb_vx_data ffi_local_data = bind.info.vtab.init_local(&ffi_input, ffi_global, &error_out);
     if (error_out) {
         throw BinderException(IntoErrString(error_out));
     }
 
-    return make_uniq<CTableLocalData>(unique_ptr<CData>(reinterpret_cast<CData *>(ffi_local_data)));
+    auto cdata = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_local_data));
+    return make_uniq<CTableLocalData>(std::move(cdata));
 }
 
 void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) {
     const auto &bind = input.bind_data->Cast<CTableBindData>();
 
-    auto ctx = reinterpret_cast<duckdb_client_context>(&context);
-    const auto bind_data = bind.ffi_data->DataPtr();
-    auto global_data = input.global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
-    auto local_data = input.local_state->Cast<CTableLocalData>().ffi_data->DataPtr();
+    duckdb_client_context ffi_ctx = reinterpret_cast<duckdb_client_context>(&context);
+    void *const ffi_bind = bind.ffi_data->DataPtr();
+    void *const ffi_global = input.global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
+    void *const ffi_local = input.local_state->Cast<CTableLocalData>().ffi_data->DataPtr();
 
     duckdb_data_chunk chunk = reinterpret_cast<duckdb_data_chunk>(&output);
     duckdb_vx_error error_out = nullptr;
-    bind.info->vtab.function(ctx, bind_data, global_data, local_data, chunk, &error_out);
+    bind.info.vtab.function(ffi_ctx, ffi_bind, ffi_global, ffi_local, chunk, &error_out);
     if (error_out) {
         throw InvalidInputException(IntoErrString(error_out));
     }
@@ -283,12 +283,12 @@ void c_pushdown_complex_filter(ClientContext &,
                                vector<unique_ptr<Expression>> &filters) {
     auto &bind = bind_data->Cast<CTableBindData>();
     void *const ffi_bind = bind.ffi_data->DataPtr();
+    duckdb_vx_error error_out = nullptr;
 
     for (auto iter = filters.begin(); iter != filters.end();) {
-        duckdb_vx_error error_out = nullptr;
         duckdb_vx_expr ffi_expr = reinterpret_cast<duckdb_vx_expr>(iter->get());
 
-        const bool pushed = bind.info->vtab.pushdown_complex_filter(ffi_bind, ffi_expr, &error_out);
+        const bool pushed = bind.info.vtab.pushdown_complex_filter(ffi_bind, ffi_expr, &error_out);
         if (error_out) {
             throw BinderException(IntoErrString(error_out));
         }
@@ -302,7 +302,7 @@ unique_ptr<NodeStatistics> c_cardinality(ClientContext &, const FunctionData *bi
     auto &bind = bind_data->Cast<CTableBindData>();
 
     duckdb_vx_node_statistics stats = {};
-    bind.info->vtab.cardinality(bind.ffi_data->DataPtr(), &stats);
+    bind.info.vtab.cardinality(bind.ffi_data->DataPtr(), &stats);
 
     auto out = make_uniq<NodeStatistics>();
     out->has_estimated_cardinality = stats.has_estimated_cardinality;
@@ -315,29 +315,23 @@ unique_ptr<NodeStatistics> c_cardinality(ClientContext &, const FunctionData *bi
 
 extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input,
                                                                  size_t index) {
-    if (!ffi_input) {
-        return nullptr;
-    }
-
+    D_ASSERT(ffi_input);
     const TableFunctionBindInput &input = *reinterpret_cast<TableFunctionBindInput *>(ffi_input);
-    if (index >= input.inputs.size()) {
-        return nullptr;
-    }
     return reinterpret_cast<duckdb_value>(new Value(input.inputs[index]));
 }
 
 extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result,
                                                        const char *name_str,
                                                        size_t name_len,
                                                        duckdb_logical_type ffi_type) {
-    if (!name_str || !ffi_type) {
-        return;
-    }
-    const auto result = reinterpret_cast<CTableBindResult *>(ffi_result);
-    const auto logical_type = reinterpret_cast<LogicalType *>(ffi_type);
-
-    result->names.emplace_back(name_str, name_len);
-    result->return_types.emplace_back(*logical_type);
+    D_ASSERT(ffi_result);
+    D_ASSERT(name_str);
+    D_ASSERT(ffi_type);
+    const CTableBindResult &result = *reinterpret_cast<CTableBindResult *>(ffi_result);
+    const LogicalType logical_type = *reinterpret_cast<LogicalType *>(ffi_type);
+
+    result.names.emplace_back(name_str, name_len);
+    result.return_types.emplace_back(logical_type);
 }
 
 OperatorPartitionData c_get_partition_data(ClientContext &, TableFunctionGetPartitionInput &input) {
@@ -350,7 +344,7 @@ OperatorPartitionData c_get_partition_data(ClientContext &, TableFunctionGetPart
     void *const ffi_local = input.local_state->Cast<CTableLocalData>().ffi_data->DataPtr();
 
     duckdb_vx_error error_out = nullptr;
-    const idx_t batch_index = bind.info->vtab.get_partition_data(ffi_bind, ffi_global, ffi_local, &error_out);
+    const idx_t batch_index = bind.info.vtab.get_partition_data(ffi_bind, ffi_global, ffi_local, &error_out);
     if (error_out) {
         throw InvalidInputException(IntoErrString(error_out));
     }
@@ -368,18 +362,18 @@ InsertionOrderPreservingMap<string> c_to_string(TableFunctionToStringInput &inpu
     InsertionOrderPreservingMap<string> result;
     duckdb_vx_string_map ffi_map = reinterpret_cast<duckdb_vx_string_map>(&result);
     void *const ffi_bind = input.bind_data->Cast<CTableBindData>().ffi_data->DataPtr();
-    static_cast<CTableFunctionInfo &>(*input.table_function.function_info).vtab.to_string(ffi_bind, ffi_map);
+    const auto &info = static_cast<CTableFunctionInfo &>(*input.table_function.function_info);
+    info.vtab.to_string(ffi_bind, ffi_map);
     return result;
 }
 
 extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const duckdb_vx_tfunc_vtab_t *vtab) {
-    if (!ffi_db || !vtab) {
-        return DuckDBError;
-    }
+    D_ASSERT(ffi_db);
+    D_ASSERT(vtab);
 
-    auto wrapper = reinterpret_cast<duckdb::DatabaseWrapper *>(ffi_db);
-    auto db = wrapper->database->instance;
-    auto tf = TableFunction(vtab->name, {}, c_function, c_bind, c_init_global, c_init_local);
+    const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);
+    DatabaseInstance &db = *wrapper.database->instance;
+    TableFunction tf(vtab->name, {}, c_function, c_bind, c_init_global, c_init_local);
 
     tf.projection_pushdown = true;
     tf.filter_pushdown = true;
@@ -400,22 +394,22 @@ extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const d
         return {{COLUMN_IDENTIFIER_EMPTY, TableColumn("", LogicalTypeId::BOOLEAN)}};
     };
 
-    tf.arguments.reserve(vtab->parameter_count);
+    tf.arguments.resize(vtab->parameter_count);
     for (size_t i = 0; i < vtab->parameter_count; i++) {
-        auto logical_type = reinterpret_cast<LogicalType *>(vtab->parameters[i]);
-        tf.arguments.emplace_back(*logical_type);
+        tf.arguments[i] = *reinterpret_cast<LogicalType *>(vtab->parameters[i]);
     }
 
     tf.function_info = make_shared_ptr<CTableFunctionInfo>(*vtab);
 
     try {
-        auto &system_catalog = Catalog::GetSystemCatalog(*db);
-        auto data = CatalogTransaction::GetSystemTransaction(*db);
+        auto &system_catalog = Catalog::GetSystemCatalog(db);
+        auto data = CatalogTransaction::GetSystemTransaction(db);
         CreateTableFunctionInfo tf_info(tf);
-        // Allow registering multiple overloads with the same name but different parameter types.
         tf_info.on_conflict = OnCreateConflict::ALTER_ON_CONFLICT;
         system_catalog.CreateFunction(data, tf_info);
-    } catch (...) {
+    } catch (const std::exception &e) {
+        ErrorData data(e);
+        DUCKDB_LOG_ERROR(db, "Failed to create Vortex table function:\t" + data.Message());
         return DuckDBError;
     }
     return DuckDBSuccess;