MultiFileReader interface for vortex-duckdb

vortex-duckdb/build.rs

@@ -30,7 +30,7 @@ const BUILD_MARKER: &str = ".vx-build-complete";
 const DUCKDB_CACHE_DIR: &str = "vortex-duckdb-cache";
 const EXTRACT_MARKER: &str = ".vx-extract-complete";
 
-const SOURCE_FILES: [&str; 11] = [
+const SOURCE_FILES: [&str; 12] = [
     "cpp/vortex_duckdb.cpp",
     "cpp/copy_function.cpp",
     "cpp/expr.cpp",
@@ -40,6 +40,7 @@ const SOURCE_FILES: [&str; 11] = [
     "cpp/cast_pushdown.cpp",
     "cpp/aggregate_fn_pushdown.cpp",
     "cpp/table_filter.cpp",
+    "cpp/multi_file_reader.cpp",
     "cpp/table_function.cpp",
     "cpp/vector.cpp",
 ];
@@ -352,6 +353,48 @@ fn extract(archive: &Path, dest: &Path) {
     zip::ZipArchive::new(file).unwrap().extract(dest).unwrap();
 }
 
+fn git_apply(repo_dir: &Path, patch: &Path, args: &[&str]) -> bool {
+    let output = Command::new("git")
+        .current_dir(repo_dir)
+        .args(["apply", "-p1"])
+        .args(args)
+        .arg(patch)
+        .output();
+    match output {
+        Ok(out) => out.status.success(),
+        Err(e) => {
+            println!("cargo:error=git is required to patch DuckDB sources: {e}");
+            exit(1);
+        }
+    }
+}
+
+fn apply_source_patches(crate_dir: &Path, repo_dir: &Path) {
+    let mut patches: Vec<PathBuf> = fs::read_dir(crate_dir.join("patches"))
+        .unwrap()
+        .map(|entry| entry.unwrap().path())
+        .filter(|path| path.extension().is_some_and(|ext| ext == "diff"))
+        .collect();
+    patches.sort();
+
+    for patch in patches {
+        // A successful reverse dry-run means the patch is already applied.
+        if git_apply(repo_dir, &patch, &["--check", "--reverse"]) {
+            continue;
+        }
+        if !git_apply(repo_dir, &patch, &[]) {
+            println!(
+                "cargo:error=Failed to apply {} to {}; delete that directory to re-extract \
+                DuckDB sources",
+                patch.display(),
+                repo_dir.display()
+            );
+            exit(1);
+        }
+        println!("cargo:info=Applied {}", patch.display());
+    }
+}
+
 /// Download DuckDB library archive from R2 and extract it.
 /// Return false if archive is not available or download failed
 fn download_prebuilt(version: &DuckDBVersion, library_dir: &Path, target: &str) -> bool {
@@ -576,6 +619,7 @@ fn cbindgen_rust2c(crate_dir: &Path) {
 
 fn main() {
     println!("cargo:rerun-if-changed=cpp/include");
+    println!("cargo:rerun-if-changed=patches");
     println!("cargo:rerun-if-env-changed=VX_DUCKDB_DEBUG");
     println!("cargo:rerun-if-env-changed=VX_DUCKDB_SAN");
     println!("cargo:rerun-if-env-changed=CARGO_HTTP_TIMEOUT");
@@ -656,6 +700,8 @@ fn main() {
         fs::write(&extract_marker, version.to_string()).unwrap();
     }
 
+    apply_source_patches(&crate_dir, &inner_dir);
+
     drop(fs::remove_file(&duckdb_dir));
     drop(fs::remove_dir_all(&duckdb_dir));
     symlink(&source_dir, &duckdb_dir).unwrap();

vortex-duckdb/cpp/aggregate_fn_pushdown.cpp

@@ -136,5 +136,5 @@ LogicalGet *GetChildGet(const LogicalAggregate &agg) {
         return nullptr;
     }
     LogicalGet &get = op->Cast<LogicalGet>();
-    return get.function.bind == duckdb_vx_table_function_bind ? &get : nullptr;
+    return is_vortex_scan(get.function) ? &get : nullptr;
 }

vortex-duckdb/cpp/cast_pushdown.cpp

@@ -25,7 +25,7 @@ static bool ReachesPushdownGet(const LogicalOperator &op) {
         cur = cur->children[0].get();
         switch (cur->type) {
         case LogicalOperatorType::LOGICAL_GET:
-            return cur->Cast<LogicalGet>().function.bind == duckdb_vx_table_function_bind;
+            return is_vortex_scan(cur->Cast<LogicalGet>().function);
         case LogicalOperatorType::LOGICAL_PROJECTION:
             continue;
         default:

vortex-duckdb/cpp/include/multi_file_reader.hpp

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+#pragma once
+
+#include "data.hpp"
+#include "duckdb/common/multi_file/multi_file_function.hpp"
+
+using namespace duckdb;
+
+struct VortexBindData final : TableFunctionData {
+    VortexBindData() = default;
+    unique_ptr<FunctionData> Copy() const override;
+    bool Equals(const FunctionData &other) const override;
+
+    unique_ptr<CData> ffi_bind_data;
+};
+
+struct VortexBindResult {
+    vector<LogicalType> &return_types;
+    vector<string> &names;
+};
+
+struct VortexGlobalState final : GlobalTableFunctionState {
+    VortexGlobalState() = default;
+    ~VortexGlobalState() override = default;
+
+    void *ffi_bind_data = nullptr; // needed for local state partial accumulation
+    unique_ptr<CData> ffi_global_state;
+};
+
+struct VortexLocalState final : LocalTableFunctionState {
+    VortexLocalState() = default;
+    unique_ptr<CData> ffi_local_state;
+};
+
+struct VortexMultiFileReader final : MultiFileReader {
+    inline unique_ptr<MultiFileReader> Copy() const override {
+        return make_uniq<VortexMultiFileReader>();
+    }
+
+    // Prune reader if file statistics prove false for "table_filters"
+    ReaderInitializeType InitializeReader(MultiFileReaderData &reader_data,
+                                          const MultiFileBindData &bind_data,
+                                          const vector<MultiFileColumnDefinition> &global_columns,
+                                          const vector<ColumnIndex> &global_column_ids,
+                                          optional_ptr<TableFilterSet> table_filters,
+                                          ClientContext &context,
+                                          MultiFileGlobalState &gstate) override;
+};
+
+struct VortexReaderInterface final : MultiFileReaderInterface {
+    static unique_ptr<MultiFileReaderInterface> CreateInterface(ClientContext &) {
+        return make_uniq<VortexReaderInterface>();
+    }
+
+    inline unique_ptr<BaseFileReaderOptions> InitializeOptions(ClientContext &,
+                                                               optional_ptr<TableFunctionInfo>) override {
+        return make_uniq<BaseFileReaderOptions>();
+    }
+
+    inline bool ParseCopyOption(ClientContext &,
+                                const string &,
+                                const vector<Value> &,
+                                BaseFileReaderOptions &,
+                                vector<string> &,
+                                vector<LogicalType> &) override {
+        return false;
+    };
+
+    inline bool ParseOption(ClientContext &,
+                            const string &,
+                            const Value &,
+                            MultiFileOptions &,
+                            BaseFileReaderOptions &) override {
+        return false;
+    }
+
+    inline unique_ptr<TableFunctionData> InitializeBindData(MultiFileBindData &,
+                                                            unique_ptr<BaseFileReaderOptions>) override {
+        return make_uniq<VortexBindData>();
+    }
+
+    void BindReader(ClientContext &,
+                    vector<LogicalType> &return_types,
+                    vector<string> &names,
+                    MultiFileBindData &bind_data) override;
+
+    unique_ptr<GlobalTableFunctionState> InitializeGlobalState(ClientContext &context,
+                                                               MultiFileBindData &bind_data,
+                                                               MultiFileGlobalState &global_state) override;
+
+    unique_ptr<LocalTableFunctionState> InitializeLocalState(ExecutionContext &context,
+                                                             GlobalTableFunctionState &global_state) override;
+
+    inline shared_ptr<BaseFileReader> CreateReader(ClientContext &,
+                                                   GlobalTableFunctionState &,
+                                                   BaseUnionData &,
+                                                   const MultiFileBindData &) override {
+        throw BinderException("UNION BY NAME for Vortex files is not supported");
+    }
+
+    shared_ptr<BaseFileReader> CreateReader(ClientContext &context,
+                                            GlobalTableFunctionState &gstate,
+                                            const OpenFileInfo &file,
+                                            idx_t file_idx,
+                                            const MultiFileBindData &bind_data) override;
+
+    shared_ptr<BaseFileReader> CreateReader(ClientContext &context,
+                                            const OpenFileInfo &file,
+                                            BaseFileReaderOptions &options,
+                                            const MultiFileOptions &file_options) override;
+
+    unique_ptr<NodeStatistics> GetCardinality(const MultiFileBindData &bind_data, idx_t file_count) override;
+
+    inline FileGlobInput GetGlobInput() override {
+        return {FileGlobOptions::FALLBACK_GLOB, "vortex"};
+    }
+
+    inline unique_ptr<MultiFileReaderInterface> Copy() override {
+        return make_uniq<VortexReaderInterface>();
+    }
+
+    void GetVirtualColumns(ClientContext &, MultiFileBindData &, virtual_column_map_t &result) override;
+
+    bool FinalizeScan(ClientContext &, GlobalTableFunctionState &gstate, DataChunk &output) override;
+};
+
+struct VortexBaseReader final : BaseFileReader {
+    VortexBaseReader(OpenFileInfo file, unique_ptr<CData> ffi_file)
+        : BaseFileReader(file), ffi_file(std::move(ffi_file)) {
+    }
+
+    unique_ptr<CData> ffi_file;
+    unique_ptr<CData> ffi_file_scan;
+    vector<column_t> virtual_ids;
+
+    inline void AddVirtualColumn(column_t virtual_column_id) override {
+        virtual_ids.push_back(virtual_column_id);
+    }
+
+    void StartScan(GlobalTableFunctionState &gstate);
+
+    // Returns false when file is exhausted
+    bool TryInitializeScan(ClientContext &context,
+                           GlobalTableFunctionState &gstate,
+                           LocalTableFunctionState &lstate) override;
+
+    AsyncResult Scan(ClientContext &context,
+                     GlobalTableFunctionState &global_state,
+                     LocalTableFunctionState &local_state,
+                     DataChunk &chunk) override;
+
+    inline void FinishFile(ClientContext &, GlobalTableFunctionState &) override {
+    }
+
+    double GetProgressInFile(ClientContext &context) override;
+
+    unique_ptr<BaseStatistics> GetStatistics(ClientContext &context, const string &name) override;
+
+    inline string GetReaderType() const override {
+        return "Vortex";
+    }
+};

vortex-duckdb/cpp/include/table_function.h

@@ -10,14 +10,8 @@
 extern "C" {
 #endif
 
-// Info passed into the bind callback. The callback should set error or else add result columns.
-typedef struct duckdb_vx_tfunc_bind_input_ *duckdb_vx_tfunc_bind_input;
 typedef struct duckdb_vx_tfunc_bind_result_ *duckdb_vx_tfunc_bind_result;
 
-// Fetch a parameter from the bind info.
-// The caller is responsible for freeing the value using duckdb_value_free.
-duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, size_t index);
-
 // Add a result column to the bind info.
 void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result,
                                             const char *name_str,
@@ -45,14 +39,11 @@ typedef struct {
      * after filter pushdown and filter pruning. May be empty, in which case
      * column_ids should be used.
      * Indices in this list reference values from column_ids. I.e. if
-     * column_ids=[1,5,6], projection_ids=[1], output column should be
      * column_ids[1] = 5
      *
      * Example usage:
      * https://github.com/duckdb/duckdb/blob/dc11eadd8f0a7c600f0034810706605ebe10d5b9/src/include/duckdb/function/table_function.hpp#L147
      */
-    const idx_t *projection_ids;
-    size_t projection_ids_count;
 
     duckdb_vx_table_filter_set filters;
     duckdb_client_context client_context;
@@ -76,16 +67,6 @@ typedef struct {
     bool has_null;
 } duckdb_column_statistics;
 
-const idx_t INVALID_IDX = UINT64_MAX;
-
-typedef struct {
-    idx_t partition_index;
-    // Either INVALID_IDX or position of column in output for file_index column
-    size_t file_index_column_pos;
-    // File index for the exported partition.
-    size_t file_index;
-} duckdb_vx_partition_data;
-
 duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db);
 
 typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input;

vortex-duckdb/cpp/include/table_function.hpp

@@ -3,7 +3,6 @@
 
 #pragma once
 
-#include "data.hpp"
 #include "duckdb.h"
 #include "duckdb/function/function.hpp"
 #include "duckdb/function/table_function.hpp"
@@ -12,11 +11,7 @@ using namespace duckdb;
 
 static_assert(sizeof(idx_t) == 8);
 
-// We need this exposed to compare function addresses in optimizer.cpp
-unique_ptr<FunctionData> duckdb_vx_table_function_bind(ClientContext &context,
-                                                       TableFunctionBindInput &input,
-                                                       vector<LogicalType> &return_types,
-                                                       vector<string> &names);
+bool is_vortex_scan(const TableFunction &function);
 
 struct TableFunctionProjectionExpressionInput {
     const LogicalGet &get;
@@ -35,36 +30,3 @@ struct TableFunctionUngroupedAggregateInput {
 };
 
 bool aggregate_pushdown(ClientContext &context, const TableFunctionUngroupedAggregateInput &input);
-
-struct VortexBindData final : FunctionData {
-    VortexBindData(unique_ptr<CData> ffi_data, const vector<LogicalType> &types)
-        : ffi_data(std::move(ffi_data)), types(types) {
-    }
-    unique_ptr<FunctionData> Copy() const override;
-    bool Equals(const FunctionData &other) const override;
-
-    unique_ptr<CData> ffi_data;
-    vector<LogicalType> types;
-};
-
-struct VortexGlobalData final : GlobalTableFunctionState {
-    explicit VortexGlobalData(unique_ptr<CData> ffi_data) : ffi_data(std::move(ffi_data)) {
-    }
-
-    idx_t MaxThreads() const override {
-        return GlobalTableFunctionState::MAX_THREADS;
-    }
-
-    unique_ptr<CData> ffi_data;
-};
-
-struct VortexLocalData final : LocalTableFunctionState {
-    explicit VortexLocalData(unique_ptr<CData> ffi_data) : ffi_data(std::move(ffi_data)) {
-    }
-    unique_ptr<CData> ffi_data;
-};
-
-struct VortexBindResults {
-    vector<LogicalType> &return_types;
-    vector<string> &names;
-};