better

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-duckdb/build.rs

@@ -483,6 +483,6 @@ fn main() {
     let duckdb_include_dir = inner_dir.join("src").join("include");
     println!("cargo:rerun-if-changed=cpp/include");
     c2rust(&crate_dir, &duckdb_include_dir);
-    cpp(&duckdb_include_dir);
     rust2c(&crate_dir);
+    cpp(&duckdb_include_dir);
 }

vortex-duckdb/cbindgen.toml

@@ -1,5 +1,8 @@
 language = "C"
 cpp_compat = true
+usize_is_size_t = true
+style = "type"
+braces = "SameLine"
 
 header = """
 // SPDX-License-Identifier: Apache-2.0

vortex-duckdb/cpp/copy_function.cpp

@@ -1,11 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-#include "duckdb_vx.h"
 #include "duckdb_vx/data.hpp"
 #include "duckdb_vx/error.hpp"
-
+#include "duckdb_vx/table_function.h"
+#include "vortex.h"
 #include "duckdb_vx/duckdb_diagnostics.h"
+
 DUCKDB_INCLUDES_BEGIN
 #include "duckdb/function/copy_function.hpp"
 #include "duckdb/main/capi/capi_internal.hpp"
@@ -15,161 +15,122 @@ DUCKDB_INCLUDES_BEGIN
 DUCKDB_INCLUDES_END
 
 using namespace duckdb;
+using vortex::CData;
+using vortex::IntoErrString;
 
-namespace vortex {
-
-struct CCopyBindData final : TableFunctionData {
-    CCopyBindData(const duckdb_vx_copy_func_vtab_t vtab_p, unique_ptr<CData> ffi_data_p)
-        : vtab(vtab_p), ffi_data(std::move(ffi_data_p)) {
-    }
-
-    const duckdb_vx_copy_func_vtab_t vtab;
-    unique_ptr<CData> ffi_data;
-};
-
-struct CCopyGlobalData final : GlobalFunctionData {
-    explicit CCopyGlobalData(unique_ptr<CData> ffi_data_p) : ffi_data(std::move(ffi_data_p)) {
+struct CopyBindData final : TableFunctionData {
+    CopyBindData(unique_ptr<CData> ffi_data) : ffi_data(std::move(ffi_data)) {
     }
-
     unique_ptr<CData> ffi_data;
 };
 
-struct CCopyLocalData final : LocalFunctionData {
-    explicit CCopyLocalData(unique_ptr<CData> ffi_data_p) : ffi_data(std::move(ffi_data_p)) {
+struct CopyGlobalData final : GlobalFunctionData {
+    CopyGlobalData(unique_ptr<CData> ffi_data) : ffi_data(std::move(ffi_data)) {
     }
 
     unique_ptr<CData> ffi_data;
 };
 
-static duckdb_vx_copy_func_vtab_t copy_vtab_one;
-
-unique_ptr<FunctionData> c_bind_one(ClientContext & /*context*/,
-                                    CopyFunctionBindInput &info,
-                                    const vector<string> &column_names,
-                                    const vector<LogicalType> &column_types) {
-
-    auto c_column_names = vector<char *>();
-    c_column_names.reserve(column_names.size());
-    for (const auto &col_id : column_names) {
-        c_column_names.push_back(const_cast<char *>(col_id.c_str()));
+unique_ptr<FunctionData> copy_to_bind(ClientContext &,
+                                      CopyFunctionBindInput &,
+                                      const vector<string> &column_names,
+                                      const vector<LogicalType> &column_types) {
+    vector<const char *> ffi_column_names(column_names.size());
+    for (size_t i = 0; i < column_names.size(); ++i) {
+        ffi_column_names[i] = column_names[i].c_str();
     }
 
-    auto c_column_types = vector<duckdb_logical_type>();
-    c_column_types.reserve(c_column_types.size());
-    for (auto &col_type : column_types) {
-        c_column_types.push_back(reinterpret_cast<duckdb_logical_type>(const_cast<LogicalType *>(&col_type)));
+    vector<duckdb_logical_type> ffi_column_types(column_types.size());
+    for (size_t i = 0; i < column_types.size(); ++i) {
+        // duckdb C api doesn't allow passing const LogicalTypes. We never
+        // modify input in copy function.
+        ffi_column_types[i] =
+            reinterpret_cast<duckdb_logical_type>(const_cast<LogicalType *>(&column_types[i]));
     }
 
     duckdb_vx_error error_out = nullptr;
-    // TODO(myrrc): do we pass ownership of c_column_names in bind?
-    // If yes, it's a UB as we'd double delete on function return
-    auto ffi_bind_data = copy_vtab_one.bind(reinterpret_cast<duckdb_vx_copy_func_bind_input>(&info),
-                                            c_column_names.data(),
-                                            c_column_names.size(),
-                                            c_column_types.data(),
-                                            c_column_types.size(),
-                                            &error_out);
+    const duckdb_vx_data ffi_bind_data = duckdb_copy_function_copy_to_bind(ffi_column_names.data(),
+                                                                           ffi_column_names.size(),
+                                                                           ffi_column_types.data(),
+                                                                           ffi_column_types.size(),
+                                                                           &error_out);
     if (error_out) {
         throw BinderException(IntoErrString(error_out));
     }
-
-    return make_uniq<CCopyBindData>(
-        // This should only be filled out once
-        copy_vtab_one,
-        unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data)));
+    auto cdata = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data));
+    return make_uniq<CopyBindData>(std::move(cdata));
 }
 
 unique_ptr<GlobalFunctionData>
-c_init_global(ClientContext &context, FunctionData &bind_data, const string &file_path) {
-    auto &bind = bind_data.Cast<CCopyBindData>();
-    duckdb_vx_error error_out = nullptr;
-    auto global_data = bind.vtab.init_global(reinterpret_cast<duckdb_client_context>(&context),
-                                             bind.ffi_data->DataPtr(),
-                                             file_path.c_str(),
-                                             &error_out);
-    if (error_out) {
-        throw ExecutorException(IntoErrString(error_out));
-    }
+copy_to_initialize_global(ClientContext &context, FunctionData &bind_data, const string &file_path) {
+    void *const ffi_bind = bind_data.Cast<CopyBindData>().ffi_data->DataPtr();
+    const auto ffi_ctx = reinterpret_cast<duckdb_client_context>(&context);
 
-    return make_uniq<CCopyGlobalData>(unique_ptr<CData>(reinterpret_cast<CData *>(global_data)));
-}
-
-unique_ptr<LocalFunctionData> c_init_local(ExecutionContext & /*context*/, FunctionData &bind_data) {
-    auto &bind = bind_data.Cast<CCopyBindData>();
     duckdb_vx_error error_out = nullptr;
-    auto data = bind.vtab.init_local(bind.ffi_data->DataPtr(), &error_out);
+    const duckdb_vx_data ffi_global =
+        duckdb_copy_function_copy_to_initialize_global(ffi_ctx, ffi_bind, file_path.c_str(), &error_out);
     if (error_out) {
         throw ExecutorException(IntoErrString(error_out));
     }
 
-    return make_uniq<CCopyLocalData>(unique_ptr<CData>(reinterpret_cast<CData *>(data)));
+    auto cdata = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_global));
+    return make_uniq<CopyGlobalData>(std::move(cdata));
 }
 
-void c_copy_to_sink(ExecutionContext & /*context*/,
-                    FunctionData &bind_data,
-                    GlobalFunctionData &gstate,
-                    LocalFunctionData &lstate,
-                    DataChunk &input) {
-    auto &bind = bind_data.Cast<CCopyBindData>();
-    auto &global = gstate.Cast<CCopyGlobalData>();
-    auto &local = lstate.Cast<CCopyLocalData>();
+void copy_to_sink(ExecutionContext &,
+                  FunctionData &bind_data,
+                  GlobalFunctionData &gstate,
+                  LocalFunctionData &,
+                  DataChunk &input) {
+    void *const ffi_bind = bind_data.Cast<CopyBindData>().ffi_data->DataPtr();
+    void *const ffi_global = gstate.Cast<CopyGlobalData>().ffi_data->DataPtr();
+    auto ffi_chunk = reinterpret_cast<duckdb_data_chunk>(&input);
     duckdb_vx_error error_out = nullptr;
-    bind.vtab.copy_to_sink(bind.ffi_data->DataPtr(),
-                           global.ffi_data->DataPtr(),
-                           local.ffi_data->DataPtr(),
-                           reinterpret_cast<duckdb_data_chunk>(&input),
-                           &error_out);
+    duckdb_copy_function_copy_to_sink(ffi_bind, ffi_global, ffi_chunk, &error_out);
     if (error_out) {
         throw ExecutorException(IntoErrString(error_out));
     }
 }
 
-void copy_to_finalize(ClientContext & /*context*/, FunctionData &bind_data, GlobalFunctionData &gstate) {
-    auto &bind = bind_data.Cast<CCopyBindData>();
-    auto &global = gstate.Cast<CCopyGlobalData>();
+void copy_to_finalize(ClientContext &, FunctionData &, GlobalFunctionData &gstate) {
+    void *const ffi_global = gstate.Cast<CopyGlobalData>().ffi_data->DataPtr();
     duckdb_vx_error error_out = nullptr;
-    bind.vtab.copy_to_finalize(bind.ffi_data->DataPtr(), global.ffi_data->DataPtr(), &error_out);
+    duckdb_copy_function_copy_to_finalize(ffi_global, &error_out);
     if (error_out) {
         throw ExecutorException(IntoErrString(error_out));
     }
 }
 
-extern "C" duckdb_vx_copy_func_vtab_t *get_vtab_one() {
-    return &copy_vtab_one;
-}
+extern "C" duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db) {
+    D_ASSERT(ffi_db);
+    const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);
+    DatabaseInstance &db = *wrapper.database->instance;
 
-extern "C" duckdb_state duckdb_vx_copy_func_register_vtab_one(duckdb_database ffi_db) {
-    if (!ffi_db) {
-        return DuckDBError;
-    }
-
-    auto wrapper = reinterpret_cast<duckdb::DatabaseWrapper *>(ffi_db);
-    auto db = wrapper->database->instance;
-    auto copy_function = CopyFunction(copy_vtab_one.name);
-
-    copy_function.copy_to_bind = c_bind_one;
-    copy_function.copy_to_initialize_global = c_init_global;
-    copy_function.copy_to_initialize_local = c_init_local;
-
-    copy_function.copy_to_sink = c_copy_to_sink;
-    copy_function.copy_to_finalize = copy_to_finalize;
-    copy_function.extension = copy_vtab_one.extension;
+    CopyFunction fn("vortex");
+    fn.copy_to_bind = copy_to_bind;
+    fn.copy_to_initialize_global = copy_to_initialize_global;
+    fn.copy_to_initialize_local = [](auto &, auto &) {
+        return make_uniq<LocalFunctionData>();
+    };
+    fn.copy_to_sink = copy_to_sink;
+    fn.copy_to_finalize = copy_to_finalize;
+    fn.extension = "vortex";
 
     // TODO(joe): expose this via c our api
-    copy_function.execution_mode = [](bool /*preserve_insertion_order*/, bool /*supports_batch_index*/) {
+    fn.execution_mode = [](bool, bool) {
         return CopyFunctionExecutionMode::REGULAR_COPY_TO_FILE;
     };
     // TODO(joe): handle parameters as in table_function
 
     try {
-        auto &system_catalog = Catalog::GetSystemCatalog(*db);
-        auto data = CatalogTransaction::GetSystemTransaction(*db);
-        CreateCopyFunctionInfo copy_info(std::move(copy_function));
+        Catalog &system_catalog = Catalog::GetSystemCatalog(db);
+        CatalogTransaction data = CatalogTransaction::GetSystemTransaction(db);
+        CreateCopyFunctionInfo copy_info(std::move(fn));
         system_catalog.CreateCopyFunction(data, copy_info);
-    } catch (...) {
+    } catch (const std::exception &e) {
+        ErrorData data(e);
+        DUCKDB_LOG_ERROR(db, "Failed to create Vortex copy function:\t" + data.Message());
         return DuckDBError;
     }
     return DuckDBSuccess;
 }
-
-} // namespace vortex

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

@@ -1,25 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-/**
- * We redefine a C API for DuckDB Copy Functions used to copy/write values from duckdb to a vortex file.
- * See table_filter.h for more info
- */
 #pragma once
-
 #include "data.h"
 #include "error.h"
-#include "logical_type.h"
 #include "duckdb_vx/data.h"
 
-#ifdef __cplusplus /* If compiled as C++, use C ABI */
+#ifdef __cplusplus
 extern "C" {
 #endif
 
-// Info passed into the bind callback. The callback should set error or else add result columns.
 typedef struct duckdb_vx_copy_func_bind_input_ *duckdb_vx_copy_func_bind_input;
 
-// Input data passed into the init_global and init_local callbacks.
 typedef struct {
     const void *bind_data;
     const void *local_state;
@@ -32,6 +23,10 @@ typedef struct {
 // PARALLEL_COPY_TO_FILE,
 // BATCH_COPY_TO_FILE
 // } copy_function_execution_mode;
+//
+// TODO(joe): expose via c api
+// copy_function_execution_mode (*execution_mode)(bool preserve_insertion_order, bool
+// supports_batch_index);
 
 // A transparent DuckDB copy function vtable, which can be used to configure a copy function.
 typedef struct {
@@ -63,17 +58,10 @@ typedef struct {
 
     void (*copy_to_finalize)(const void *bind_data, void *global_data, duckdb_vx_error *error_out);
 
-    // TODO(joe): expose via c api
-    // copy_function_execution_mode (*execution_mode)(bool preserve_insertion_order, bool
-    // supports_batch_index);
 } duckdb_vx_copy_func_vtab_t;
 
-// Due to a limitation in the copy function duckdb api we have to have global (copy) vtabs.
-duckdb_vx_copy_func_vtab_t *get_vtab_one();
-
-// A single function for configuring the DuckDB table function vtable.
-duckdb_state duckdb_vx_copy_func_register_vtab_one(duckdb_database ffi_db);
+duckdb_state duckdb_vx_register_copy_function(duckdb_database ffi_db);
 
-#ifdef __cplusplus /* End C ABI */
+#ifdef __cplusplus
 };
 #endif

vortex-duckdb/cpp/table_function.cpp

@@ -37,12 +37,7 @@ struct CTableBindData final : FunctionData {
 };
 
 unique_ptr<FunctionData> CTableBindData::Copy() const {
-    duckdb_vx_error error_out = nullptr;
-    const auto copied_ffi_data = duckdb_table_function_bind_data_clone(ffi_data->DataPtr(), &error_out);
-    if (error_out) {
-        throw BinderException(IntoErrString(error_out));
-    }
-
+    const auto copied_ffi_data = duckdb_table_function_bind_data_clone(ffi_data->DataPtr());
     auto ffi_data_p = unique_ptr<CData>(reinterpret_cast<CData *>(copied_ffi_data));
     return make_uniq<CTableBindData>(std::move(ffi_data_p), types);
 }

vortex-duckdb/include/vortex.h

@@ -46,7 +46,7 @@ extern void duckdb_table_function_to_string(void *bind_data, duckdb_vx_string_ma
 
 extern
 bool duckdb_table_function_statistics(const void *bind_data,
-                                      uintptr_t column_index,
+                                      size_t column_index,
                                       duckdb_column_statistics *stats_out);
 
 extern double duckdb_table_function_scan_progress(void *global_state);
@@ -85,9 +85,28 @@ duckdb_vx_data duckdb_table_function_bind(duckdb_client_context ctx,
                                           duckdb_vx_tfunc_bind_result bind_result,
                                           duckdb_vx_error *error_out);
 
+extern duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind_data);
+
+extern
+duckdb_vx_data duckdb_copy_function_copy_to_bind(const char *const *column_names,
+                                                 size_t column_name_count,
+                                                 const duckdb_logical_type *column_types,
+                                                 size_t column_type_count,
+                                                 duckdb_vx_error *error_out);
+
 extern
-duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind_data,
-                                                     duckdb_vx_error *error_out);
+duckdb_vx_data duckdb_copy_function_copy_to_initialize_global(duckdb_client_context client_context,
+                                                              const void *bind_data,
+                                                              const char *file_path,
+                                                              duckdb_vx_error *error_out);
+
+extern
+void duckdb_copy_function_copy_to_sink(const void *bind_data,
+                                       void *global_data,
+                                       duckdb_data_chunk data_chunk,
+                                       duckdb_vx_error *error_out);
+
+extern void duckdb_copy_function_copy_to_finalize(void *global_data, duckdb_vx_error *error_out);
 
 #ifdef __cplusplus
 }  // extern "C"

vortex-duckdb/src/convert/dtype.rs

@@ -33,7 +33,6 @@ use std::sync::Arc;
 
 use vortex::dtype::DType;
 use vortex::dtype::DecimalDType;
-use vortex::dtype::FieldName;
 use vortex::dtype::Nullability;
 use vortex::dtype::PType;
 use vortex::dtype::PType::F32;
@@ -177,20 +176,6 @@ impl FromLogicalType for DType {
     }
 }
 
-pub fn from_duckdb_table<'a, I, S>(iter: I) -> VortexResult<StructFields>
-where
-    I: Iterator<Item = (S, &'a LogicalTypeRef, Nullability)>,
-    S: AsRef<str>,
-{
-    iter.map(|(name, type_, nullability)| {
-        Ok((
-            FieldName::from(name.as_ref()),
-            DType::from_logical_type(type_, nullability)?,
-        ))
-    })
-    .collect::<VortexResult<StructFields>>()
-}
-
 impl TryFrom<DType> for LogicalType {
     type Error = VortexError;