From b9eb0cdbfbd5353bf65ec3e36bad8e1674301007 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 13 Mar 2026 15:37:21 +0000 Subject: [PATCH 01/21] initial --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-duckdb/Cargo.toml | 1 + vortex-duckdb/cpp/CMakeLists.txt | 2 +- .../include/duckdb_vx/duckdb_diagnostics.h | 3 +- .../cpp/include/duckdb_vx/table_function.h | 7 ++- vortex-duckdb/cpp/table_function.cpp | 57 ++++++++++++----- .../src/duckdb/table_function/mod.rs | 3 - .../table_function/table_scan_progress.rs | 19 ------ vortex-ffi/build.rs | 24 ++++++- vortex-ffi/cinclude/vortex.h | 58 +++++++++++++++++ vortex-ffi/src/lib.rs | 1 + vortex-ffi/src/scan.rs | 63 +++++++++++++++++++ 13 files changed, 195 insertions(+), 45 deletions(-) delete mode 100644 vortex-duckdb/src/duckdb/table_function/table_scan_progress.rs create mode 100644 vortex-ffi/src/scan.rs diff --git a/Cargo.lock b/Cargo.lock index 0c1d847db67..55c7986b199 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10024,6 +10024,7 @@ dependencies = [ "url", "vortex", "vortex-array", + "vortex-ffi", "vortex-runend", "vortex-sequence", "vortex-utils", diff --git a/Cargo.toml b/Cargo.toml index 260e5300ab4..6444bf60bba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -263,6 +263,7 @@ vortex-decimal-byte-parts = { version = "0.1.0", path = "encodings/decimal-byte- vortex-error = { version = "0.1.0", path = "./vortex-error", default-features = false } vortex-fastlanes = { version = "0.1.0", path = "./encodings/fastlanes", default-features = false } vortex-file = { version = "0.1.0", path = "./vortex-file", default-features = false } +vortex-ffi = { version = "0.1.0", path = "./vortex-ffi", default-features = false } vortex-flatbuffers = { version = "0.1.0", path = "./vortex-flatbuffers", default-features = false } vortex-fsst = { version = "0.1.0", path = "./encodings/fsst", default-features = false } vortex-io = { version = "0.1.0", path = "./vortex-io", default-features = false } diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index 2d095c1c3e5..d93f213e59b 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -37,6 +37,7 @@ parking_lot = { workspace = true } paste = { workspace = true } tracing = { workspace = true } url = { workspace = true } +vortex-ffi = { workspace = true } vortex = { workspace = true, features = ["files", "tokio", "object_store"] } vortex-utils = { workspace = true, features = ["dashmap"] } diff --git a/vortex-duckdb/cpp/CMakeLists.txt b/vortex-duckdb/cpp/CMakeLists.txt index 9671d93dd6d..e7d9bbbfc11 100644 --- a/vortex-duckdb/cpp/CMakeLists.txt +++ b/vortex-duckdb/cpp/CMakeLists.txt @@ -23,7 +23,7 @@ if (NOT CMAKE_BUILD_TYPE) endif() # Enable compiler warnings (matching build.rs flags). -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wpedantic") +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wpedantic -Wno-unused-parameter") # Find DuckDB include directory via the symlink created by build.rs. # The symlink points to target/duckdb-source-vX.Y.Z which contains duckdb-X.Y.Z/ diff --git a/vortex-duckdb/cpp/include/duckdb_vx/duckdb_diagnostics.h b/vortex-duckdb/cpp/include/duckdb_vx/duckdb_diagnostics.h index e294af29bad..65d4e7fff9d 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/duckdb_diagnostics.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/duckdb_diagnostics.h @@ -18,9 +18,10 @@ _Pragma("GCC diagnostic ignored \"-Wall\"") \ _Pragma("GCC diagnostic ignored \"-Wextra\"") \ _Pragma("GCC diagnostic ignored \"-Wpedantic\"") + _Pragma("GCC diagnostic ignored \"-Wunused-parameter\"") #define DUCKDB_INCLUDES_END _Pragma("GCC diagnostic pop") #else #define DUCKDB_INCLUDES_BEGIN #define DUCKDB_INCLUDES_END #endif -// clang-format on \ No newline at end of file +// clang-format on diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h index e8f483514a5..4985f0292a5 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h @@ -150,8 +150,6 @@ typedef struct { duckdb_vx_string_map (*to_string)(void *bind_data); // void *dynamic_to_string; - double (*table_scan_progress)(duckdb_client_context ctx, void *bind_data, void *global_state); - idx_t (*get_partition_data)(const void *bind_data, void *init_global_data, void *init_local_data, @@ -174,7 +172,10 @@ typedef struct { } duckdb_vx_tfunc_vtab_t; // A single function for configuring the DuckDB table function vtable. -duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const duckdb_vx_tfunc_vtab_t *vtab); +duckdb_state duckdb_vx_tfunc_register( + duckdb_database ffi_db, + const duckdb_vx_tfunc_vtab_t *vtab +); #ifdef __cplusplus /* End C ABI */ } diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index fc4ac1b66da..6a6ecd0a0f1 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -17,16 +17,18 @@ DUCKDB_INCLUDES_END #include "duckdb_vx/data.hpp" #include "duckdb_vx/error.hpp" +// TODO +#include "../../vortex-ffi/cinclude/vortex.h" + using namespace duckdb; namespace vortex { struct CTableFunctionInfo final : TableFunctionInfo { explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) - : vtab(vtab), max_threads(vtab.max_threads) { - } + : vtab(vtab), max_threads(vtab.max_threads) { } duckdb_vx_tfunc_vtab_t vtab; - idx_t max_threads; + idx_t max_threads; // derived from vtab.max_threads }; struct CTableBindData final : TableFunctionData { @@ -50,11 +52,32 @@ struct CTableBindData final : TableFunctionData { unique_ptr ffi_data; }; +std::string move_vx_err(vx_error* error) { + const vx_string * error_vx_str = vx_error_get_message(error); + string str{vx_string_ptr(error_vx_str), vx_string_len(error_vx_str)}; + vx_error_free(error); + return str; +} + struct CTableGlobalData final : GlobalTableFunctionState { - explicit CTableGlobalData(unique_ptr ffi_data_p, idx_t max_threads_p) - : ffi_data(std::move(ffi_data_p)), max_threads(max_threads_p) { + explicit CTableGlobalData( + unique_ptr ffi_data_p, + const vx_scan_options& options) + : ffi_data(std::move(ffi_data_p)) + , max_threads(options.max_threads) + { + vx_error* error; + if (scan = vx_scan_init(&options, &error); !scan) + throw BinderException(move_vx_err(error)); } + ~CTableGlobalData() override { + if (scan) vx_scan_free(scan); + } + + // TODO Ideally should be C++ api for lifetime checks + const vx_scan * scan; + unique_ptr ffi_data; idx_t max_threads; @@ -78,14 +101,10 @@ struct CTableBindResult { vector &names; }; -double c_table_scan_progress(ClientContext &context, - const FunctionData *bind_data, +double c_table_scan_progress(ClientContext &, + const FunctionData *, const GlobalTableFunctionState *global_state) { - auto &bind = bind_data->Cast(); - duckdb_client_context c_ctx = reinterpret_cast(&context); - void *const c_bind_data = bind.ffi_data->DataPtr(); - void *const c_global_state = global_state->Cast().ffi_data->DataPtr(); - return bind.info->vtab.table_scan_progress(c_ctx, c_bind_data, c_global_state); + return vx_scan_progress(global_state->Cast().scan); } unique_ptr c_bind(ClientContext &context, @@ -133,9 +152,12 @@ unique_ptr c_init_global(ClientContext &context, Table throw BinderException(IntoErrString(error_out)); } - return make_uniq( - unique_ptr(reinterpret_cast(ffi_global_data)), - bind.info->max_threads); + vx_scan_options scan_options = { + .max_threads = bind.info->max_threads, + }; + + auto cdata = unique_ptr(reinterpret_cast(ffi_global_data)); + return make_uniq(std::move(cdata), scan_options); } unique_ptr c_init_local(ExecutionContext &context, @@ -342,7 +364,10 @@ InsertionOrderPreservingMap c_to_string(TableFunctionToStringInput &inpu return result; } -extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const duckdb_vx_tfunc_vtab_t *vtab) { +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; } diff --git a/vortex-duckdb/src/duckdb/table_function/mod.rs b/vortex-duckdb/src/duckdb/table_function/mod.rs index f20e844d381..c9fa08177b3 100644 --- a/vortex-duckdb/src/duckdb/table_function/mod.rs +++ b/vortex-duckdb/src/duckdb/table_function/mod.rs @@ -14,7 +14,6 @@ mod cardinality; mod init; mod partition; mod pushdown_complex_filter; -mod table_scan_progress; mod virtual_columns; pub use bind::*; @@ -34,7 +33,6 @@ use crate::duckdb::expr::ExpressionRef; use crate::duckdb::table_function::cardinality::cardinality_callback; use crate::duckdb::table_function::partition::get_partition_data_callback; use crate::duckdb::table_function::pushdown_complex_filter::pushdown_complex_filter_callback; -use crate::duckdb::table_function::table_scan_progress::table_scan_progress_callback; use crate::duckdb::table_function::virtual_columns::get_virtual_columns_callback; use crate::duckdb_try; @@ -194,7 +192,6 @@ impl DatabaseRef { pushdown_expression: ptr::null_mut::(), get_virtual_columns: Some(get_virtual_columns_callback::), to_string: Some(to_string_callback::), - table_scan_progress: Some(table_scan_progress_callback::), get_partition_data: Some(get_partition_data_callback::), projection_pushdown: T::PROJECTION_PUSHDOWN, filter_pushdown: T::FILTER_PUSHDOWN, diff --git a/vortex-duckdb/src/duckdb/table_function/table_scan_progress.rs b/vortex-duckdb/src/duckdb/table_function/table_scan_progress.rs deleted file mode 100644 index cfe3dd43b45..00000000000 --- a/vortex-duckdb/src/duckdb/table_function/table_scan_progress.rs +++ /dev/null @@ -1,19 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex::error::VortexExpect; - -use crate::duckdb::TableFunction; - -pub(crate) unsafe extern "C-unwind" fn table_scan_progress_callback( - ctx: crate::cpp::duckdb_client_context, - bind_data: *mut ::std::os::raw::c_void, - global_state: *mut ::std::os::raw::c_void, -) -> f64 { - let ctx = unsafe { crate::duckdb::ClientContext::borrow(ctx) }; - let bind_data = - unsafe { bind_data.cast::().as_ref() }.vortex_expect("bind_data null pointer"); - let global_state = unsafe { global_state.cast::().as_ref() } - .vortex_expect("global_init_data null pointer"); - T::table_scan_progress(ctx, bind_data, global_state) -} diff --git a/vortex-ffi/build.rs b/vortex-ffi/build.rs index 8c14001790f..223a3ab9073 100644 --- a/vortex-ffi/build.rs +++ b/vortex-ffi/build.rs @@ -7,9 +7,28 @@ use std::env; use std::path::PathBuf; use std::process::Command; +const SOURCE_FILES: [&str; 15] = [ + "array.rs", + "array_iterator.rs", + "binary.rs", + "dtype.rs", + "error.rs", + "file.rs", + "lib.rs", + "log.rs", + "macros.rs", + "ptype.rs", + "scan.rs", + "session.rs", + "sink.rs", + "string.rs", + "struct_fields.rs", +]; + fn main() { - // Set up dependency tracking - println!("cargo:rerun-if-changed=src/"); + for f in SOURCE_FILES { + println!("cargo:rerun-if-changed=src/{f}"); + } println!("cargo:rerun-if-changed=cbindgen.toml"); println!("cargo:rerun-if-changed=Cargo.toml"); println!("cargo:rerun-if-changed=build.rs"); @@ -43,6 +62,7 @@ fn main() { let result = cbindgen::Builder::new() .with_crate(&crate_dir) .with_config(config) + .with_documentation(true) .generate(); match result { diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index fc28a952fdc..5261a47b584 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -302,6 +302,13 @@ typedef struct vx_error vx_error; */ typedef struct vx_file vx_file; +/** + * A scan request. May point to a local file or a remote file. + * Multiple files are supported via globs. + * Functions operating on it are thread-safe. + */ +typedef struct vx_scan vx_scan; + /** * A handle to a Vortex session. */ @@ -384,6 +391,40 @@ typedef struct { unsigned long row_offset; } vx_file_scan_options; +typedef struct { + const char *files; + /** + * Column names to project out in the scan. These must be null-terminated C strings. + */ + const char *projection_expression; + /** + * Number of columns in `projection`. + */ + unsigned int projection_expr_len; + /** + * Serialized expressions for pushdown + */ + const char *filter_expression; + /** + * The len in bytes of the filter expression + */ + unsigned int filter_expression_len; + /** + * Splits files into chunks of this size, if zero then we use the write layout. + */ + int split_by_row_count; + /** + * First row of a range to scan. + */ + unsigned long row_range_start; + /** + * Last row of a range to scan. + */ + unsigned long row_range_end; + unsigned long limit; + unsigned long max_threads; +} vx_scan_options; + #ifdef __cplusplus extern "C" { #endif // __cplusplus @@ -735,6 +776,23 @@ vx_array_iterator *vx_file_scan(const vx_session *session, */ void vx_set_log_level(vx_log_level level); +/** + * Clone a borrowed [`vx_scan`], returning an owned [`vx_scan`]. + * + * + * Must be released with [`vx_scan_free`]. + */ +const vx_scan *vx_scan_clone(const vx_scan *ptr); + +/** + * Free an owned [`vx_scan`] object. + */ +void vx_scan_free(const vx_scan *ptr); + +const vx_scan *vx_scan_init(const vx_scan_options *_options, vx_error **_error_out); + +double vx_scan_progress(const vx_scan *); + /** * Free an owned [`vx_session`] object. */ diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 1d97f29a01b..c7e32c346d9 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -13,6 +13,7 @@ mod dtype; mod error; mod file; mod log; +mod scan; mod macros; mod ptype; mod session; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs new file mode 100644 index 00000000000..fd76c31258b --- /dev/null +++ b/vortex-ffi/src/scan.rs @@ -0,0 +1,63 @@ +use std::{ffi::{c_char, c_int, c_uint, c_ulong}, sync::Arc}; + +use crate::error::vx_error; + +pub struct Scan { + +} + +crate::arc_wrapper!( + /// A scan request. May point to a local file or a remote file. + /// Multiple files are supported via globs. + /// Functions operating on it are thread-safe. + Scan, + vx_scan); + +// TODO(myrrc): currently copy of vx_file_scan_options +#[repr(C)] +#[allow(non_camel_case_types)] +pub struct vx_scan_options { + pub files: *const c_char, + + /// Column names to project out in the scan. These must be null-terminated C strings. + pub projection_expression: *const c_char, + + /// Number of columns in `projection`. + pub projection_expr_len: c_uint, + + /// Serialized expressions for pushdown + pub filter_expression: *const c_char, + + /// The len in bytes of the filter expression + pub filter_expression_len: c_uint, + + /// Splits files into chunks of this size, if zero then we use the write layout. + pub split_by_row_count: c_int, + + /// First row of a range to scan. + pub row_range_start: c_ulong, + + /// Last row of a range to scan. + pub row_range_end: c_ulong, + + pub limit: c_ulong, + + pub max_threads: c_ulong, +} + +// TODO mut pointer? +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_scan_init( + _options: *const vx_scan_options, + _error_out: *mut *mut vx_error, +) -> *const vx_scan { + let scan = Arc::new(Scan { + + }); + vx_scan::new(scan) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_scan_progress(_: *const vx_scan) -> f64 { + return 0.0; +} From fa3b275a4b22b82e53097f97a24033a9dec29fdc Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 13 Mar 2026 17:09:44 +0000 Subject: [PATCH 02/21] fix --- vortex-duckdb/cpp/table_function.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 6a6ecd0a0f1..213ffdcf167 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -49,7 +49,13 @@ struct CTableBindData final : TableFunctionData { } unique_ptr info; + unique_ptr ffi_data; + // ^ In Rust it contains + // Data source ref (already populated) + // filter expressions + // column names + // column types }; std::string move_vx_err(vx_error* error) { From 15f9e5f0bf2319196ba0f19b278f48302f3adaad Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 09:57:32 +0000 Subject: [PATCH 03/21] better --- .../cpp/include/duckdb_vx/table_function.h | 12 +++++++ vortex-duckdb/cpp/table_function.cpp | 33 +++++++++---------- vortex-ffi/cinclude/vortex.h | 15 +++++++++ vortex-ffi/src/scan.rs | 16 ++++++--- 4 files changed, 54 insertions(+), 22 deletions(-) diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h index 4985f0292a5..09bdc921fd3 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h @@ -14,6 +14,9 @@ #include "duckdb_vx/data.h" #include "duckdb_vx/client_context.h" +// TODO +#include "../../../vortex-ffi/cinclude/vortex.h" + #ifdef __cplusplus /* If compiled as C++, use C ABI */ extern "C" { #endif @@ -169,6 +172,15 @@ typedef struct { bool sampling_pushdown; bool late_materialization; idx_t max_threads; + + // PoC hack: retain Rust exporter code + void (*export_array)( + const vx_array* arr, + duckdb_data_chunk chunk, + duckdb_vx_error* err, + // needed to initialize rust-side exporter + void* local_state + ); } duckdb_vx_tfunc_vtab_t; // A single function for configuring the DuckDB table function vtable. diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 213ffdcf167..2a5e5dd6e1d 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -17,9 +17,6 @@ DUCKDB_INCLUDES_END #include "duckdb_vx/data.hpp" #include "duckdb_vx/error.hpp" -// TODO -#include "../../vortex-ffi/cinclude/vortex.h" - using namespace duckdb; namespace vortex { @@ -193,22 +190,22 @@ unique_ptr c_init_local(ExecutionContext &context, } void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { - const auto &bind = input.bind_data->Cast(); - - auto ctx = reinterpret_cast(&context); - const auto bind_data = bind.ffi_data->DataPtr(); - auto global_data = input.global_state->Cast().ffi_data->DataPtr(); - auto local_data = input.local_state->Cast().ffi_data->DataPtr(); + const vx_scan* const scan = input.global_state->Cast().scan; + vx_error * vx_err = nullptr; + const vx_array* arr = vx_scan_next_arr(scan, &vx_err); + if (vx_err) { + throw InvalidInputException(move_vx_err(vx_err)); + } + if (!arr) { // everything exported + return; + } - duckdb_vx_error error_out = nullptr; - bind.info->vtab.function(ctx, - bind_data, - global_data, - local_data, - reinterpret_cast(&output), - &error_out); - if (error_out) { - throw InvalidInputException(IntoErrString(error_out)); + duckdb_vx_error err = nullptr; + void* local_data = input.local_state->Cast().ffi_data->DataPtr(); + input.bind_data->Cast().info->vtab.export_array( + arr, reinterpret_cast(&output), &err, local_data); + if (err) { + throw InvalidInputException(IntoErrString(err)); } } diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 5261a47b584..81c6fcc15de 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -319,6 +319,8 @@ typedef struct vx_session vx_session; */ typedef struct vx_string vx_string; +typedef struct vx_struct_array vx_struct_array; + /** * Represents a Vortex struct data type, without top-level nullability. */ @@ -776,6 +778,19 @@ vx_array_iterator *vx_file_scan(const vx_session *session, */ void vx_set_log_level(vx_log_level level); +/** + * Clone a borrowed [`vx_struct_array`], returning an owned [`vx_struct_array`]. + * + * + * Must be released with [`vx_struct_array_free`]. + */ +const vx_struct_array *vx_struct_array_clone(const vx_struct_array *ptr); + +/** + * Free an owned [`vx_struct_array`] object. + */ +void vx_struct_array_free(const vx_struct_array *ptr); + /** * Clone a borrowed [`vx_scan`], returning an owned [`vx_scan`]. * diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index fd76c31258b..b024d2d6d69 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -1,6 +1,11 @@ use std::{ffi::{c_char, c_int, c_uint, c_ulong}, sync::Arc}; +use vortex::{array::arrays::StructArray, scan::api::DataSource}; +use crate::{array::vx_array, error::vx_error}; -use crate::error::vx_error; +crate::arc_dyn_wrapper!( + dyn DataSource, + vx_data_source +); pub struct Scan { @@ -10,6 +15,7 @@ crate::arc_wrapper!( /// A scan request. May point to a local file or a remote file. /// Multiple files are supported via globs. /// Functions operating on it are thread-safe. + /// A Scan can be consumed only once, it can't be restarted. Scan, vx_scan); @@ -57,7 +63,9 @@ pub unsafe extern "C-unwind" fn vx_scan_init( vx_scan::new(scan) } -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_scan_progress(_: *const vx_scan) -> f64 { - return 0.0; +pub unsafe extern "C-unwind" fn vx_scan_next( + scan: *const vx_scan, + err: *mut *mut vx_error, +) -> *const vx_array { + } From 15991e52defb8402654da4ed18aa06a181037819 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 11:13:08 +0000 Subject: [PATCH 04/21] better --- vortex-ffi/src/dtype.rs | 7 ++ vortex-ffi/src/error.rs | 3 + vortex-ffi/src/scan.rs | 191 ++++++++++++++++++++++++++++++++-------- 3 files changed, 165 insertions(+), 36 deletions(-) diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index a05a53459ca..7a6e97fa4dd 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -324,6 +324,13 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *cons } } + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_type_to_arrow_schema( + dtype: *const DType, out: *mut ArrowSchema, err: vx_error_out) { + todo!(); +} + #[cfg(test)] #[allow(clippy::cast_possible_truncation)] mod tests { diff --git a/vortex-ffi/src/error.rs b/vortex-ffi/src/error.rs index 53c91bca7eb..888967e25dc 100644 --- a/vortex-ffi/src/error.rs +++ b/vortex-ffi/src/error.rs @@ -19,6 +19,9 @@ box_wrapper!( vx_error ); +#[allow(non_camel_case_types)] +pub type vx_error_out = *mut *mut vx_error; + #[inline] pub fn try_or_default( error_out: *mut *mut vx_error, diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index b024d2d6d69..7ed83b43883 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -1,61 +1,156 @@ -use std::{ffi::{c_char, c_int, c_uint, c_ulong}, sync::Arc}; -use vortex::{array::arrays::StructArray, scan::api::DataSource}; -use crate::{array::vx_array, error::vx_error}; +#![allow(non_camel_case_types)] + +use std::{ffi::{c_char, c_int, c_uint, c_ulong, c_void}, sync::Arc}; +use vortex::{expr::Expression, scan::api::{DataSource, Partition}}; +use crate::{array::vx_array, dtype::vx_dtype, error::{vx_error, vx_error_out}}; +use std::option::Option; crate::arc_dyn_wrapper!( dyn DataSource, vx_data_source ); -pub struct Scan { +struct FFIFileHandle { } -crate::arc_wrapper!( - /// A scan request. May point to a local file or a remote file. - /// Multiple files are supported via globs. - /// Functions operating on it are thread-safe. - /// A Scan can be consumed only once, it can't be restarted. - Scan, - vx_scan); +crate::box_wrapper!( + FFIFileHandle, + vx_file_handle +); + +pub type vx_list_callback = Option; +pub type vx_glob_callback = Option; + +pub type FsUseVortex = Option; +pub type FsSetUserdata = Option; + +pub type FsOpen = Option; +pub type FsCreate = Option; + +pub type FsList = Option; + +pub type FsClose = Option; +pub type FsSize = Option u64>; + +pub type FsRead = Option u64>; + +pub type FsWrite = Option u64>; + +pub type FsSync = Option; + +pub type Glob = Option; + +pub type CacheInit = Option; +pub type CacheFree = Option; +pub type CacheGet = Option; +pub type CachePut = Option; +pub type CacheDelete = Option; -// TODO(myrrc): currently copy of vx_file_scan_options #[repr(C)] -#[allow(non_camel_case_types)] -pub struct vx_scan_options { - pub files: *const c_char, +pub struct vx_datasource_options { + fs_use_vortex: FsUseVortex, + fs_set_userdata: FsSetUserdata, + fs_open: FsOpen, + fs_create: FsCreate, + fs_list: FsList, + fs_close: FsClose, + fs_size: FsSize, + fs_read: FsRead, + fs_write: FsWrite, + fs_sync: FsSync, + glob: Glob, + cache_init: CacheInit, + cache_get: CacheGet, + cache_put: CachePut, + cache_delete: CacheDelete, +} - /// Column names to project out in the scan. These must be null-terminated C strings. - pub projection_expression: *const c_char, +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_datasource_dtype(ds: *const vx_data_source) -> *const vx_dtype { - /// Number of columns in `projection`. - pub projection_expr_len: c_uint, +} - /// Serialized expressions for pushdown - pub filter_expression: *const c_char, +#[repr(C)] +enum VxCardinality { + VX_CARD_UNKNOWN = 0, + VX_CARD_ESTIMATE = 1, + VX_CARD_MAXIMUM = 2, +} - /// The len in bytes of the filter expression - pub filter_expression_len: c_uint, +#[repr(C)] +struct vx_datasource_row_count { + cardinality: VxCardinality, + rows: u64 +} - /// Splits files into chunks of this size, if zero then we use the write layout. - pub split_by_row_count: c_int, +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_datasource_get_row_count(ds: *const vx_data_source, rc: *mut vx_datasource_row_count) { - /// First row of a range to scan. - pub row_range_start: c_ulong, +} - /// Last row of a range to scan. - pub row_range_end: c_ulong, +pub struct VxArrayIt; +crate::box_wrapper!( + VxArrayIt, + vx_array_it +); - pub limit: c_ulong, +#[repr(C)] +enum ScanSelection { + VX_S_INCLUDE_ALL = 0, + VX_S_INCLUDE_RANGE = 1, + VX_S_EXCLUDE_RANGE = 2 +} - pub max_threads: c_ulong, +#[repr(C)] +struct vx_scan_selection { + selection: ScanSelection, + idx: *mut usize, + idx_len: usize } -// TODO mut pointer? +crate::arc_dyn_wrapper!( + Expression, + vx_expression +); + +// TODO why make it distinct from upstream scan options? +#[repr(C)] +struct vx_scan_options { + projection: vx_expression, + filter: vx_expression, + row_range_begin: u64, + row_range_end: u64, + selection: vx_scan_selection, + ordered: c_int, + limit: u64 +} + +enum VxPartitionEstimate { + VX_P_ESTIMATE_UNKNOWN = 0, + VX_P_ESTIMATE_EXACT = 1, + VX_P_ESTIMATE_INEXACT = 2 +} + +#[repr(C)] +struct vx_partition_estimate { + boundary: VxPartitionEstimate, + estimate: u64, +} + +pub struct Scan; +crate::arc_wrapper!( + Scan, + vx_scan); + #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_scan_init( - _options: *const vx_scan_options, - _error_out: *mut *mut vx_error, +pub unsafe extern "C-unwind" fn vx_datasource_scan( + ds: vx_data_source, + opts: *const vx_scan_options, + est: *mut vx_partition_estimate, + err: vx_error_out ) -> *const vx_scan { let scan = Arc::new(Scan { @@ -63,9 +158,33 @@ pub unsafe extern "C-unwind" fn vx_scan_init( vx_scan::new(scan) } +crate::arc_dyn_wrapper!( + dyn Partition, + vx_partition +); + pub unsafe extern "C-unwind" fn vx_scan_next( scan: *const vx_scan, - err: *mut *mut vx_error, + err: vx_error_out +) -> vx_partition { + +} + +pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( + part: vx_partition, + out: *mut ArrowArrayStream, + err: vx_error_out +) { + +} + +pub unsafe extern "C-unwind" fn vx_partition_next( + partition: vx_partition, + err: vx_error_out ) -> *const vx_array { } + +pub unsafe extern "C-unwind" fn vx_scan_cancel(scan: vx_scan) { + +} From a886b025e7f7e217da3a5a20935f4bb185609e81 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 12:03:26 +0000 Subject: [PATCH 05/21] better --- vortex-ffi/src/scan.rs | 163 ++++++++++++++++++++++++++++------------- 1 file changed, 112 insertions(+), 51 deletions(-) diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 7ed83b43883..fa7945effe9 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -1,8 +1,9 @@ #![allow(non_camel_case_types)] -use std::{ffi::{c_char, c_int, c_uint, c_ulong, c_void}, sync::Arc}; -use vortex::{expr::Expression, scan::api::{DataSource, Partition}}; -use crate::{array::vx_array, dtype::vx_dtype, error::{vx_error, vx_error_out}}; +use core::slice; +use std::{ffi::{c_char, c_int, c_void}, ops::Range, sync::Arc}; +use vortex::{buffer::Buffer, expr::{stats::Precision::{Exact, Inexact}, Expression}, scan::{api::{DataSource, DataSourceScan, Partition, ScanRequest}, Selection}}; +use crate::{array::vx_array, dtype::vx_dtype, error::{try_or_default, vx_error_out}}; use std::option::Option; crate::arc_dyn_wrapper!( @@ -10,13 +11,31 @@ crate::arc_dyn_wrapper!( vx_data_source ); -struct FFIFileHandle { +pub struct FileHandle; +crate::box_wrapper!( + FileHandle, + vx_file_handle +); -} +crate::arc_dyn_wrapper!( + dyn Partition, + vx_partition +); crate::box_wrapper!( - FFIFileHandle, - vx_file_handle + Expression, + vx_expression +); + +// TODO should be akin to DataSourceScanRef +crate::arc_dyn_wrapper!( + dyn DataSourceScan, + vx_scan); + +pub struct VxArrayIt; +crate::box_wrapper!( + VxArrayIt, + vx_array_it ); pub type vx_list_callback = Option; @@ -50,7 +69,7 @@ pub type CachePut = Option; #[repr(C)] -pub struct vx_datasource_options { +pub struct vx_data_source_options { fs_use_vortex: FsUseVortex, fs_set_userdata: FsSetUserdata, fs_open: FsOpen, @@ -69,8 +88,8 @@ pub struct vx_datasource_options { } #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_datasource_dtype(ds: *const vx_data_source) -> *const vx_dtype { - +pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { + vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) } #[repr(C)] @@ -81,22 +100,28 @@ enum VxCardinality { } #[repr(C)] -struct vx_datasource_row_count { +struct vx_data_source_row_count { cardinality: VxCardinality, rows: u64 } #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_datasource_get_row_count(ds: *const vx_data_source, rc: *mut vx_datasource_row_count) { - +pub unsafe extern "C-unwind" fn vx_data_source_get_row_count(ds: *const vx_data_source, rc: *mut vx_data_source_row_count) { + match vx_data_source::as_ref(ds).row_count() { + Some(Exact(rows)) => unsafe { + (*rc).cardinality = VxCardinality::VX_CARD_MAXIMUM; + (*rc).rows = rows; + }, + Some(Inexact(rows)) => unsafe { + (*rc).cardinality = VxCardinality::VX_CARD_ESTIMATE; + (*rc).rows = rows; + }, + None => unsafe { + (*rc).cardinality = VxCardinality::VX_CARD_UNKNOWN; + } + } } -pub struct VxArrayIt; -crate::box_wrapper!( - VxArrayIt, - vx_array_it -); - #[repr(C)] enum ScanSelection { VX_S_INCLUDE_ALL = 0, @@ -107,20 +132,15 @@ enum ScanSelection { #[repr(C)] struct vx_scan_selection { selection: ScanSelection, - idx: *mut usize, + idx: *mut u64, idx_len: usize } -crate::arc_dyn_wrapper!( - Expression, - vx_expression -); - -// TODO why make it distinct from upstream scan options? +// TODO why make it distinct from ScanRequest? Easier option handling #[repr(C)] struct vx_scan_options { - projection: vx_expression, - filter: vx_expression, + projection: *const vx_expression, + filter: *const vx_expression, row_range_begin: u64, row_range_end: u64, selection: vx_scan_selection, @@ -140,38 +160,81 @@ struct vx_partition_estimate { estimate: u64, } -pub struct Scan; -crate::arc_wrapper!( - Scan, - vx_scan); - #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_datasource_scan( - ds: vx_data_source, +pub unsafe extern "C-unwind" fn vx_data_source_scan( + ds: *const vx_data_source, opts: *const vx_scan_options, est: *mut vx_partition_estimate, err: vx_error_out ) -> *const vx_scan { - let scan = Arc::new(Scan { - - }); - vx_scan::new(scan) + let request = if opts.is_null() { + ScanRequest::default() + } else { + let opts = &unsafe { *opts }; + + let projection = if opts.projection.is_null() { + panic!() // TODO + } else { + *vx_expression::as_ref(opts.projection) + }; + + let filter = if opts.filter.is_null() { + None + } else { + Some(*vx_expression::as_ref(opts.filter)) + }; + + let selection = &opts.selection; + let selection: Selection = match selection.selection { + ScanSelection::VX_S_INCLUDE_ALL => { + Selection::All + }, + ScanSelection::VX_S_INCLUDE_RANGE => { + let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; + let buf = Buffer::copy_from(buf); + Selection::IncludeByIndex(buf) + } + ScanSelection::VX_S_EXCLUDE_RANGE => { + let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; + let buf = Buffer::copy_from(buf); + Selection::ExcludeByIndex(buf) + } + }; + + let ordered = opts.ordered == 1; + + let row_range: Option> = if opts.row_range_begin == 0 && opts.row_range_end == 0 { + None + } else { + Some(Range{ start: opts.row_range_begin, end: opts.row_range_end} ) + }; + + let limit = if opts.limit == 0 { None } else { Some(opts.limit) }; + + ScanRequest { + projection, + filter, + row_range, + selection, + ordered, + limit + } + }; + try_or_default(err, || { + let scan = vx_data_source::as_ref(ds).scan(request); + Ok(vx_scan::new(scan)) + }) } -crate::arc_dyn_wrapper!( - dyn Partition, - vx_partition -); - pub unsafe extern "C-unwind" fn vx_scan_next( - scan: *const vx_scan, + scan: vx_scan, err: vx_error_out -) -> vx_partition { +) -> *const vx_partition { } pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( - part: vx_partition, + part: *const vx_partition, out: *mut ArrowArrayStream, err: vx_error_out ) { @@ -179,12 +242,10 @@ pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( } pub unsafe extern "C-unwind" fn vx_partition_next( - partition: vx_partition, + partition: *const vx_partition, err: vx_error_out ) -> *const vx_array { } -pub unsafe extern "C-unwind" fn vx_scan_cancel(scan: vx_scan) { - -} +pub unsafe extern "C-unwind" fn vx_scan_cancel(_: vx_scan) { } From 1bba8bdddd4e078b93e30668764a759b4c1dcd23 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 12:21:52 +0000 Subject: [PATCH 06/21] better --- Cargo.lock | 1 + vortex-cxx/cpp/include/vortex/scan.hpp | 2 +- vortex-ffi/Cargo.toml | 5 +- vortex-ffi/src/dtype.rs | 6 +- vortex-ffi/src/lib.rs | 2 +- vortex-ffi/src/scan.rs | 117 +++++++++++++------------ 6 files changed, 73 insertions(+), 60 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 55c7986b199..776a631edf7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10071,6 +10071,7 @@ dependencies = [ name = "vortex-ffi" version = "0.1.0" dependencies = [ + "arrow-array", "async-fs", "cbindgen", "futures", diff --git a/vortex-cxx/cpp/include/vortex/scan.hpp b/vortex-cxx/cpp/include/vortex/scan.hpp index 7debb9ef6aa..b5cd075f9f7 100644 --- a/vortex-cxx/cpp/include/vortex/scan.hpp +++ b/vortex-cxx/cpp/include/vortex/scan.hpp @@ -105,4 +105,4 @@ class ScanBuilder { rust::Box impl_; }; -} // namespace vortex \ No newline at end of file +} // namespace vortex diff --git a/vortex-ffi/Cargo.toml b/vortex-ffi/Cargo.toml index 377e1125dc1..0a777fd2fa4 100644 --- a/vortex-ffi/Cargo.toml +++ b/vortex-ffi/Cargo.toml @@ -2,8 +2,8 @@ name = "vortex-ffi" description = "Native C FFI bindings for Vortex" readme = "README.md" -# This crate is not meant to be consumed by other Rust code but rather produces a static binary -# that can be linked to by other languages. +# This crate is not meant to be consumed by Rust code outside of Vortex but rather +# producee a static binary that can be linked to by other languages. publish = false version = { workspace = true } homepage = { workspace = true } @@ -31,6 +31,7 @@ tracing = { workspace = true, features = ["std", "log"] } tracing-subscriber = { workspace = true, features = ["env-filter"] } url = { workspace = true, features = [] } vortex = { workspace = true, features = ["object_store"] } +arrow-array = { workspace = true, features = ["ffi"] } [dev-dependencies] tempfile = { workspace = true } diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index 7a6e97fa4dd..1cd41b1d898 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -324,10 +324,12 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *cons } } - #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_type_to_arrow_schema( - dtype: *const DType, out: *mut ArrowSchema, err: vx_error_out) { + dtype: *const DType, + out: *mut ArrowSchema, + err: vx_error_out, +) { todo!(); } diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index c7e32c346d9..30764572c6a 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -13,9 +13,9 @@ mod dtype; mod error; mod file; mod log; -mod scan; mod macros; mod ptype; +mod scan; mod session; mod sink; mod string; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index fa7945effe9..1e75fa6f1fd 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -1,42 +1,43 @@ #![allow(non_camel_case_types)] use core::slice; -use std::{ffi::{c_char, c_int, c_void}, ops::Range, sync::Arc}; -use vortex::{buffer::Buffer, expr::{stats::Precision::{Exact, Inexact}, Expression}, scan::{api::{DataSource, DataSourceScan, Partition, ScanRequest}, Selection}}; -use crate::{array::vx_array, dtype::vx_dtype, error::{try_or_default, vx_error_out}}; +use std::ffi::c_char; +use std::ffi::c_int; +use std::ffi::c_void; +use std::ops::Range; use std::option::Option; - -crate::arc_dyn_wrapper!( - dyn DataSource, - vx_data_source -); +use std::sync::Arc; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use vortex::buffer::Buffer; +use vortex::expr::Expression; +use vortex::expr::stats::Precision::Exact; +use vortex::expr::stats::Precision::Inexact; +use vortex::scan::Selection; +use vortex::scan::api::DataSource; +use vortex::scan::api::DataSourceScan; +use vortex::scan::api::Partition; +use vortex::scan::api::ScanRequest; + +use crate::array::vx_array; +use crate::dtype::vx_dtype; +use crate::error::try_or_default; +use crate::error::vx_error_out; + +crate::arc_dyn_wrapper!(dyn DataSource, vx_data_source); pub struct FileHandle; -crate::box_wrapper!( - FileHandle, - vx_file_handle -); +crate::box_wrapper!(FileHandle, vx_file_handle); -crate::arc_dyn_wrapper!( - dyn Partition, - vx_partition -); +crate::arc_dyn_wrapper!(dyn Partition, vx_partition); -crate::box_wrapper!( - Expression, - vx_expression -); +crate::box_wrapper!(Expression, vx_expression); // TODO should be akin to DataSourceScanRef -crate::arc_dyn_wrapper!( - dyn DataSourceScan, - vx_scan); +crate::arc_dyn_wrapper!(dyn DataSourceScan, vx_scan); pub struct VxArrayIt; -crate::box_wrapper!( - VxArrayIt, - vx_array_it -); +crate::box_wrapper!(VxArrayIt, vx_array_it); pub type vx_list_callback = Option; pub type vx_glob_callback = Option; @@ -52,11 +53,13 @@ pub type FsList = Option; pub type FsSize = Option u64>; -pub type FsRead = Option u64>; +pub type FsRead = Option< + unsafe extern "C" fn(vx_file_handle, usize, usize, *mut u8, *mut usize, vx_error_out) -> u64, +>; -pub type FsWrite = Option u64>; +pub type FsWrite = Option< + unsafe extern "C" fn(vx_file_handle, usize, usize, *mut u8, *mut usize, vx_error_out) -> u64, +>; pub type FsSync = Option; @@ -102,11 +105,14 @@ enum VxCardinality { #[repr(C)] struct vx_data_source_row_count { cardinality: VxCardinality, - rows: u64 + rows: u64, } #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_data_source_get_row_count(ds: *const vx_data_source, rc: *mut vx_data_source_row_count) { +pub unsafe extern "C-unwind" fn vx_data_source_get_row_count( + ds: *const vx_data_source, + rc: *mut vx_data_source_row_count, +) { match vx_data_source::as_ref(ds).row_count() { Some(Exact(rows)) => unsafe { (*rc).cardinality = VxCardinality::VX_CARD_MAXIMUM; @@ -118,7 +124,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_get_row_count(ds: *const vx_data_ }, None => unsafe { (*rc).cardinality = VxCardinality::VX_CARD_UNKNOWN; - } + }, } } @@ -126,14 +132,14 @@ pub unsafe extern "C-unwind" fn vx_data_source_get_row_count(ds: *const vx_data_ enum ScanSelection { VX_S_INCLUDE_ALL = 0, VX_S_INCLUDE_RANGE = 1, - VX_S_EXCLUDE_RANGE = 2 + VX_S_EXCLUDE_RANGE = 2, } #[repr(C)] struct vx_scan_selection { selection: ScanSelection, idx: *mut u64, - idx_len: usize + idx_len: usize, } // TODO why make it distinct from ScanRequest? Easier option handling @@ -145,13 +151,13 @@ struct vx_scan_options { row_range_end: u64, selection: vx_scan_selection, ordered: c_int, - limit: u64 + limit: u64, } enum VxPartitionEstimate { VX_P_ESTIMATE_UNKNOWN = 0, VX_P_ESTIMATE_EXACT = 1, - VX_P_ESTIMATE_INEXACT = 2 + VX_P_ESTIMATE_INEXACT = 2, } #[repr(C)] @@ -165,7 +171,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( ds: *const vx_data_source, opts: *const vx_scan_options, est: *mut vx_partition_estimate, - err: vx_error_out + err: vx_error_out, ) -> *const vx_scan { let request = if opts.is_null() { ScanRequest::default() @@ -186,9 +192,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( let selection = &opts.selection; let selection: Selection = match selection.selection { - ScanSelection::VX_S_INCLUDE_ALL => { - Selection::All - }, + ScanSelection::VX_S_INCLUDE_ALL => Selection::All, ScanSelection::VX_S_INCLUDE_RANGE => { let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; let buf = Buffer::copy_from(buf); @@ -203,13 +207,21 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( let ordered = opts.ordered == 1; - let row_range: Option> = if opts.row_range_begin == 0 && opts.row_range_end == 0 { + let row_range: Option> = if opts.row_range_begin == 0 && opts.row_range_end == 0 + { None } else { - Some(Range{ start: opts.row_range_begin, end: opts.row_range_end} ) + Some(Range { + start: opts.row_range_begin, + end: opts.row_range_end, + }) }; - let limit = if opts.limit == 0 { None } else { Some(opts.limit) }; + let limit = if opts.limit == 0 { + None + } else { + Some(opts.limit) + }; ScanRequest { projection, @@ -217,7 +229,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( row_range, selection, ordered, - limit + limit, } }; try_or_default(err, || { @@ -228,24 +240,21 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( pub unsafe extern "C-unwind" fn vx_scan_next( scan: vx_scan, - err: vx_error_out + err: vx_error_out, ) -> *const vx_partition { - } pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( part: *const vx_partition, - out: *mut ArrowArrayStream, - err: vx_error_out + out: *mut FFI_ArrowArrayStream, + err: vx_error_out, ) { - } pub unsafe extern "C-unwind" fn vx_partition_next( partition: *const vx_partition, - err: vx_error_out + err: vx_error_out, ) -> *const vx_array { - } -pub unsafe extern "C-unwind" fn vx_scan_cancel(_: vx_scan) { } +pub unsafe extern "C-unwind" fn vx_scan_cancel(_: vx_scan) {} From 8274a1703047ac4321762f8d7f54f9dde815a80f Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 14:35:45 +0000 Subject: [PATCH 07/21] better --- vortex-duckdb/build.rs | 1 + vortex-duckdb/cpp/CMakeLists.txt | 2 +- .../cpp/include/duckdb_vx/table_function.h | 8 +++- vortex-duckdb/cpp/table_function.cpp | 43 +++++++++++-------- 4 files changed, 33 insertions(+), 21 deletions(-) diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index be20dde4b8a..0613f163338 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -330,6 +330,7 @@ fn cpp(duckdb_include_dir: &Path) { .cpp(true) .include(duckdb_include_dir) .include("cpp/include") + .include("../vortex-ffi/cinclude") .files(SOURCE_FILES) .compile("vortex-duckdb-extras"); // bindgen generates rerun-if-changed for .h/.hpp files diff --git a/vortex-duckdb/cpp/CMakeLists.txt b/vortex-duckdb/cpp/CMakeLists.txt index e7d9bbbfc11..fdc96e472a5 100644 --- a/vortex-duckdb/cpp/CMakeLists.txt +++ b/vortex-duckdb/cpp/CMakeLists.txt @@ -39,7 +39,7 @@ else() ) endif() -include_directories(include ${DUCKDB_INCLUDE}) +include_directories(include ${DUCKDB_INCLUDE} ../../vortex-ffi/cinclude) # Auto-discover C++ source files file(GLOB CPP_SOURCES "*.cpp") diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h index 09bdc921fd3..258a0c47fae 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h @@ -13,9 +13,13 @@ #include "table_filter.h" #include "duckdb_vx/data.h" #include "duckdb_vx/client_context.h" +#include "duckdb_vx/duckdb_diagnostics.h" -// TODO -#include "../../../vortex-ffi/cinclude/vortex.h" +DUCKDB_INCLUDES_BEGIN +#include "duckdb/common/arrow/arrow.hpp" +DUCKDB_INCLUDES_END + +#include "vortex.h" #ifdef __cplusplus /* If compiled as C++, use C ABI */ extern "C" { diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 2a5e5dd6e1d..b70ac10de7f 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -49,10 +49,11 @@ struct CTableBindData final : TableFunctionData { unique_ptr ffi_data; // ^ In Rust it contains - // Data source ref (already populated) // filter expressions // column names // column types + + const vx_data_source* data_source; }; std::string move_vx_err(vx_error* error) { @@ -234,24 +235,30 @@ void c_pushdown_complex_filter(ClientContext & /*context*/, } } -unique_ptr c_cardinality(ClientContext & /*context*/, const FunctionData *bind_data) { - auto &bind = bind_data->Cast(); - - duckdb_vx_node_statistics node_stats_out = { - .estimated_cardinality = 0, - .max_cardinality = 0, - .has_estimated_cardinality = false, - .has_max_cardinality = false, - }; - bind.info->vtab.cardinality(bind_data->Cast().ffi_data->DataPtr(), &node_stats_out); - +unique_ptr c_cardinality(ClientContext &, const FunctionData *bind_data) { auto stats = make_uniq(); - stats->has_estimated_cardinality = node_stats_out.has_estimated_cardinality; - stats->estimated_cardinality = node_stats_out.estimated_cardinality; - stats->has_max_cardinality = node_stats_out.has_max_cardinality; - stats->max_cardinality = node_stats_out.max_cardinality; - - return stats; + vx_data_source_row_count rc = {}; + vx_data_source_get_row_count(bind_data->Cast().data_source, &rc); + switch (rc.cardinality) { + case VX_CARD_ESTIMATE: { + stats->has_estimated_cardinality = true; + stats->has_max_cardinality = false; + stats->estimated_cardinality = rc.rows; + return stats; + } + case VX_CARD_MAXIMUM: { + stats->has_estimated_cardinality = true; + stats->has_max_cardinality = true; + stats->estimated_cardinality = rc.rows; + stats->max_cardinality = rc.rows; + return stats; + } + case VX_CARD_UNKNOWN: { + stats->has_estimated_cardinality = false; + stats->has_max_cardinality = false; + return stats; + } + } } extern "C" size_t duckdb_vx_tfunc_bind_input_get_parameter_count(duckdb_vx_tfunc_bind_input ffi_input) { From c089d9bbe7fd16397911b52dd4bc5a7501562087 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 17:10:34 +0000 Subject: [PATCH 08/21] better --- vortex-duckdb/build.rs | 4 +- vortex-duckdb/cpp/table_function.cpp | 102 +++++++------ vortex-duckdb/src/datasource.rs | 8 -- vortex-ffi/cbindgen.toml | 3 + vortex-ffi/cinclude/vortex.h | 206 +++++++++++++++++++++------ vortex-ffi/src/dtype.rs | 4 +- vortex-ffi/src/scan.rs | 132 ++++++++++++----- 7 files changed, 322 insertions(+), 137 deletions(-) diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index 0613f163338..90f264a119a 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -41,6 +41,8 @@ const SOURCE_FILES: [&str; 18] = [ "cpp/vector_buffer.cpp", ]; +const FFI_INCLUDE: &str = "../vortex-ffi/cinclude"; + const DOWNLOAD_MAX_RETRIES: i32 = 3; const DOWNLOAD_TIMEOUT: u64 = 90; @@ -330,7 +332,7 @@ fn cpp(duckdb_include_dir: &Path) { .cpp(true) .include(duckdb_include_dir) .include("cpp/include") - .include("../vortex-ffi/cinclude") + .include(FFI_INCLUDE) .files(SOURCE_FILES) .compile("vortex-duckdb-extras"); // bindgen generates rerun-if-changed for .h/.hpp files diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index b70ac10de7f..459cce3c992 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -19,6 +19,10 @@ DUCKDB_INCLUDES_END using namespace duckdb; +std::string to_string(const vx_string* str) { + return { vx_string_ptr(str), size_t vx_string_len(str) }; +} + namespace vortex { struct CTableFunctionInfo final : TableFunctionInfo { explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) @@ -29,9 +33,13 @@ struct CTableFunctionInfo final : TableFunctionInfo { }; struct CTableBindData final : TableFunctionData { - CTableBindData(unique_ptr info_p, unique_ptr ffi_data_p) - : info(std::move(info_p)), ffi_data(std::move(ffi_data_p)) { - } + CTableBindData( + unique_ptr info_p, + unique_ptr ffi_data_p, + const vx_data_source* data_source) + : info(std::move(info_p)) + , ffi_data(std::move(ffi_data_p)), + , data_source(data_source) { } unique_ptr Copy() const override { assert(info->vtab.bind_data_clone != nullptr); @@ -45,13 +53,12 @@ struct CTableBindData final : TableFunctionData { unique_ptr(reinterpret_cast(copied_ffi_data))); } + // TODO remove with unique_ptr info; unique_ptr ffi_data; // ^ In Rust it contains // filter expressions - // column names - // column types const vx_data_source* data_source; }; @@ -66,9 +73,10 @@ std::string move_vx_err(vx_error* error) { struct CTableGlobalData final : GlobalTableFunctionState { explicit CTableGlobalData( unique_ptr ffi_data_p, + idx_t max_threads, const vx_scan_options& options) : ffi_data(std::move(ffi_data_p)) - , max_threads(options.max_threads) + , max_threads(max_threads) { vx_error* error; if (scan = vx_scan_init(&options, &error); !scan) @@ -97,44 +105,40 @@ struct CTableLocalData final : LocalTableFunctionState { unique_ptr ffi_data; }; -/** - * Result of the bind function encapsulates the output schema. - */ -struct CTableBindResult { - vector &return_types; - vector &names; -}; - -double c_table_scan_progress(ClientContext &, - const FunctionData *, - const GlobalTableFunctionState *global_state) { - return vx_scan_progress(global_state->Cast().scan); -} - -unique_ptr c_bind(ClientContext &context, - TableFunctionBindInput &input, - vector &return_types, - vector &names) { - const auto &info = input.table_function.function_info->Cast(); - - // Setup bind info to pass into the callback. - CTableBindResult result = { - return_types, - names, - }; - - duckdb_vx_error error_out = nullptr; - auto ctx = reinterpret_cast(&context); - auto ffi_bind_data = info.vtab.bind(ctx, - reinterpret_cast(&input), - reinterpret_cast(&result), - &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); +LogicalType from_dtype(const vx_dtype* dtype) { } + +unique_ptr c_bind( + ClientContext &context, + TableFunctionBindInput &info, + vector &types, + vector &names) +{ + if (info->inputs.empty()) + throw BinderException("missing file glob parameter"); + + const vx_error* err = nullptr; + vx_data_source_options opts = {0}; + opts.files = input->inputs[0]; + const vx_data_source* data_source = vx_data_source_new(opts, &err); + + const vx_dtype* dtype = vx_data_source_dtype(data_source); + const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); + + for (uint64_t i = 0; i < vx_struct_fields_nfields(struct_dtype); ++i) { + const vx_string* field_name = vx_struct_fields_field_name(struct_dtype, i); + const vx_dtype* field_dtype = vx_struct_fields_field_dtype(struct_dtype, i); + if (!field_dtype) + throw BinderException(StringUtil::Format( + "Field dtype %s at index %d can't be parsed", + to_string(field_name), i)); + names.emplace_back(to_string(field_name)); + types.emplace_back(from_dtype(field_dtype)); } - return make_uniq(make_uniq(info.vtab), - unique_ptr(reinterpret_cast(ffi_bind_data))); + return make_uniq( + make_uniq(info.vtab), + unique_ptr(reinterpret_cast(ffi_bind_data)), + data_source); } unique_ptr c_init_global(ClientContext &context, TableFunctionInitInput &input) { @@ -192,7 +196,7 @@ unique_ptr c_init_local(ExecutionContext &context, void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { const vx_scan* const scan = input.global_state->Cast().scan; - vx_error * vx_err = nullptr; + vx_error* vx_err = nullptr; const vx_array* arr = vx_scan_next_arr(scan, &vx_err); if (vx_err) { throw InvalidInputException(move_vx_err(vx_err)); @@ -395,9 +399,17 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( tf.late_materialization = vtab->late_materialization; tf.cardinality = c_cardinality; tf.get_partition_data = c_get_partition_data; - tf.get_virtual_columns = c_get_virtual_columns; tf.to_string = c_to_string; - tf.table_scan_progress = c_table_scan_progress; + + tf.get_virtual_columns = [](auto&, auto) { + auto res = virtual_column_map_t(); + res->emplace(COLUMN_IDENTIFIER_EMPTY, "", LogicalType::bool()); + return; + } + + tf.table_scan_progress = [](auto&, auto*, const GlobalTableFunctionState& state) { + return vx_scan_progress(state->Cast().scan); + }; // Set up the parameters tf.arguments.reserve(vtab->parameter_count); diff --git a/vortex-duckdb/src/datasource.rs b/vortex-duckdb/src/datasource.rs index 9e4d39c677d..ea7520af591 100644 --- a/vortex-duckdb/src/datasource.rs +++ b/vortex-duckdb/src/datasource.rs @@ -406,14 +406,6 @@ impl TableFunction for T { Ok(false) } - fn cardinality(bind_data: &Self::BindData) -> Cardinality { - match bind_data.data_source.row_count() { - Some(Precision::Exact(v)) => Cardinality::Maximum(v), - Some(Precision::Inexact(v)) => Cardinality::Estimate(v), - None => Cardinality::Unknown, - } - } - fn partition_data( _bind_data: &Self::BindData, _global_init_data: &Self::GlobalState, diff --git a/vortex-ffi/cbindgen.toml b/vortex-ffi/cbindgen.toml index 12e484f093e..65e175ba995 100644 --- a/vortex-ffi/cbindgen.toml +++ b/vortex-ffi/cbindgen.toml @@ -12,6 +12,9 @@ header = """ // // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY // + +typedef ArrowSchema FFI_ArrowSchema; +typedef ArrowArrayStream FFI_ArrowArrayStream; """ [export.rename] diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 81c6fcc15de..8266f6bed89 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -5,6 +5,9 @@ // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY // +typedef ArrowSchema FFI_ArrowSchema; +typedef ArrowArrayStream FFI_ArrowArrayStream; + #pragma once #include @@ -144,6 +147,24 @@ typedef enum { LOG_LEVEL_TRACE = 5, } vx_log_level; +typedef enum { + VX_CARD_UNKNOWN = 0, + VX_CARD_ESTIMATE = 1, + VX_CARD_MAXIMUM = 2, +} VxCardinality; + +typedef enum { + VX_S_INCLUDE_ALL = 0, + VX_S_INCLUDE_RANGE = 1, + VX_S_EXCLUDE_RANGE = 2, +} ScanSelection; + +typedef enum { + VX_P_ESTIMATE_UNKNOWN = 0, + VX_P_ESTIMATE_EXACT = 1, + VX_P_ESTIMATE_INEXACT = 2, +} vx_partition_estimate_boundary; + /** * Physical type enum, represents the in-memory physical layout but might represent a different logical type. */ @@ -228,6 +249,8 @@ typedef uint8_t PType; */ typedef struct DType DType; +typedef struct FileHandle FileHandle; + /** * Whether an instance of a DType can be `null or not */ @@ -248,6 +271,8 @@ typedef struct Primitive Primitive; */ typedef struct vx_array vx_array; +typedef struct vx_array_it vx_array_it; + /** * A Vortex array iterator. * @@ -284,6 +309,8 @@ typedef struct vx_array_sink vx_array_sink; */ typedef struct vx_binary vx_binary; +typedef struct vx_data_source vx_data_source; + /** * A Vortex data type. * @@ -297,16 +324,15 @@ typedef struct vx_dtype vx_dtype; */ typedef struct vx_error vx_error; +typedef struct vx_expression vx_expression; + /** * A handle to a Vortex file encapsulating the footer and logic for instantiating a reader. */ typedef struct vx_file vx_file; -/** - * A scan request. May point to a local file or a remote file. - * Multiple files are supported via globs. - * Functions operating on it are thread-safe. - */ +typedef struct vx_partition vx_partition; + typedef struct vx_scan vx_scan; /** @@ -319,8 +345,6 @@ typedef struct vx_session vx_session; */ typedef struct vx_string vx_string; -typedef struct vx_struct_array vx_struct_array; - /** * Represents a Vortex struct data type, without top-level nullability. */ @@ -331,6 +355,8 @@ typedef struct vx_struct_fields vx_struct_fields; */ typedef struct vx_struct_fields_builder vx_struct_fields_builder; +typedef vx_error **vx_error_out; + /** * Options supplied for opening a file. */ @@ -393,40 +419,87 @@ typedef struct { unsigned long row_offset; } vx_file_scan_options; +typedef void (*FsUseVortex)(char, char); + +typedef void (*FsSetUserdata)(void); + +typedef void (*FsOpen)(void, char, vx_error_out); + +typedef void (*FsCreate)(void, char, vx_error_out); + +typedef void (*vx_list_callback)(void, char, int); + +typedef void (*FsList)(void, char, vx_list_callback, vx_error_out); + +typedef const FileHandle *vx_file_handle; + +typedef void (*FsClose)(vx_file_handle); + +typedef uint64_t (*FsSize)(vx_file_handle, vx_error_out); + +typedef uint64_t (*FsRead)(vx_file_handle, size_t, size_t, uint8_t *, size_t *, vx_error_out); + +typedef uint64_t (*FsWrite)(vx_file_handle, size_t, size_t, uint8_t *, size_t *, vx_error_out); + +typedef void (*FsSync)(vx_file_handle, vx_error_out); + +typedef void (*vx_glob_callback)(void, char); + +typedef void (*Glob)(char, vx_glob_callback, vx_error_out); + +typedef void (*CacheInit)(vx_error_out); + +typedef void (*CacheGet)(void, char, void *, vx_error_out); + +typedef void (*CachePut)(void, char, void, vx_error_out); + +typedef void (*CacheDelete)(void, char, vx_error_out); + typedef struct { const char *files; - /** - * Column names to project out in the scan. These must be null-terminated C strings. - */ - const char *projection_expression; - /** - * Number of columns in `projection`. - */ - unsigned int projection_expr_len; - /** - * Serialized expressions for pushdown - */ - const char *filter_expression; - /** - * The len in bytes of the filter expression - */ - unsigned int filter_expression_len; - /** - * Splits files into chunks of this size, if zero then we use the write layout. - */ - int split_by_row_count; - /** - * First row of a range to scan. - */ - unsigned long row_range_start; - /** - * Last row of a range to scan. - */ - unsigned long row_range_end; - unsigned long limit; - unsigned long max_threads; + FsUseVortex fs_use_vortex; + FsSetUserdata fs_set_userdata; + FsOpen fs_open; + FsCreate fs_create; + FsList fs_list; + FsClose fs_close; + FsSize fs_size; + FsRead fs_read; + FsWrite fs_write; + FsSync fs_sync; + Glob glob; + CacheInit cache_init; + CacheGet cache_get; + CachePut cache_put; + CacheDelete cache_delete; +} vx_data_source_options; + +typedef struct { + VxCardinality cardinality; + uint64_t rows; +} vx_data_source_row_count; + +typedef struct { + ScanSelection selection; + uint64_t *idx; + size_t idx_len; +} vx_scan_selection; + +typedef struct { + const vx_expression *projection; + const vx_expression *filter; + uint64_t row_range_begin; + uint64_t row_range_end; + vx_scan_selection selection; + int ordered; + uint64_t limit; } vx_scan_options; +typedef struct { + vx_partition_estimate_boundary boundary; + uint64_t estimate; +} vx_partition_estimate; + #ifdef __cplusplus extern "C" { #endif // __cplusplus @@ -707,6 +780,8 @@ uint8_t vx_dtype_time_unit(const DType *dtype); */ const vx_string *vx_dtype_time_zone(const DType *dtype); +void vx_type_to_arrow_schema(const DType *dtype, FFI_ArrowSchema *out, vx_error_out err); + /** * Free an owned [`vx_error`] object. */ @@ -779,17 +854,35 @@ vx_array_iterator *vx_file_scan(const vx_session *session, void vx_set_log_level(vx_log_level level); /** - * Clone a borrowed [`vx_struct_array`], returning an owned [`vx_struct_array`]. + * Clone a borrowed [`vx_data_source`], returning an owned [`vx_data_source`]. * * - * Must be released with [`vx_struct_array_free`]. + * Must be released with [`vx_data_source_free`]. */ -const vx_struct_array *vx_struct_array_clone(const vx_struct_array *ptr); +const vx_data_source *vx_data_source_clone(const vx_data_source *ptr); /** - * Free an owned [`vx_struct_array`] object. + * Free an owned [`vx_data_source`] object. */ -void vx_struct_array_free(const vx_struct_array *ptr); +void vx_data_source_free(const vx_data_source *ptr); + +/** + * Clone a borrowed [`vx_partition`], returning an owned [`vx_partition`]. + * + * + * Must be released with [`vx_partition_free`]. + */ +const vx_partition *vx_partition_clone(const vx_partition *ptr); + +/** + * Free an owned [`vx_partition`] object. + */ +void vx_partition_free(const vx_partition *ptr); + +/** + * Free an owned [`vx_expression`] object. + */ +void vx_expression_free(vx_expression *ptr); /** * Clone a borrowed [`vx_scan`], returning an owned [`vx_scan`]. @@ -804,9 +897,34 @@ const vx_scan *vx_scan_clone(const vx_scan *ptr); */ void vx_scan_free(const vx_scan *ptr); -const vx_scan *vx_scan_init(const vx_scan_options *_options, vx_error **_error_out); +/** + * Free an owned [`vx_array_it`] object. + */ +void vx_array_it_free(vx_array_it *ptr); + +const vx_data_source *vx_data_source_new(const vx_data_source_options *opts, vx_error_out err); + +const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); + +void vx_data_source_get_row_count(const vx_data_source *ds, vx_data_source_row_count *rc); + +const vx_scan *vx_data_source_scan(const vx_data_source *ds, + const vx_scan_options *opts, + vx_partition_estimate *est, + vx_error_out err); + +const vx_partition *vx_scan_next(const vx_scan *scan, vx_error_out err); + +void vx_partition_scan_arrow(const vx_partition *part, FFI_ArrowArrayStream *out, vx_error_out err); + +const vx_array *vx_partition_next(const vx_partition *partition, vx_error_out err); + +/** + * Scan progress between 0.0 and 1.0 + */ +double vx_scan_progress(const vx_scan *_scan); -double vx_scan_progress(const vx_scan *); +void vx_scan_cancel(const vx_scan *); /** * Free an owned [`vx_session`] object. diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index 1cd41b1d898..18fc722d295 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -4,6 +4,7 @@ use std::ptr; use std::sync::Arc; +use arrow_array::ffi::FFI_ArrowSchema; use vortex::dtype::DType; use vortex::dtype::DecimalDType; use vortex::error::VortexExpect; @@ -14,6 +15,7 @@ use vortex::extension::datetime::Time; use vortex::extension::datetime::Timestamp; use crate::arc_wrapper; +use crate::error::vx_error_out; use crate::ptype::vx_ptype; use crate::string::vx_string; use crate::struct_fields::vx_struct_fields; @@ -327,7 +329,7 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *cons #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_type_to_arrow_schema( dtype: *const DType, - out: *mut ArrowSchema, + out: *mut FFI_ArrowSchema, err: vx_error_out, ) { todo!(); diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 1e75fa6f1fd..8a94c216f64 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -6,6 +6,7 @@ use std::ffi::c_int; use std::ffi::c_void; use std::ops::Range; use std::option::Option; +use std::ptr; use std::sync::Arc; use arrow_array::ffi_stream::FFI_ArrowArrayStream; @@ -21,13 +22,19 @@ use vortex::scan::api::ScanRequest; use crate::array::vx_array; use crate::dtype::vx_dtype; -use crate::error::try_or_default; +//use crate::error::try_or_default; use crate::error::vx_error_out; -crate::arc_dyn_wrapper!(dyn DataSource, vx_data_source); +crate::arc_dyn_wrapper!( + /// A data source is a reference to multiple possibly remote files. When + /// created, it opens first file to determine the schema from DType, all + /// other operations are deferred till a scan is requested. You can request + /// multiple file scans from a data source + dyn DataSource, + vx_data_source); pub struct FileHandle; -crate::box_wrapper!(FileHandle, vx_file_handle); +pub type vx_file_handle = *const FileHandle; crate::arc_dyn_wrapper!(dyn Partition, vx_partition); @@ -39,25 +46,35 @@ crate::arc_dyn_wrapper!(dyn DataSourceScan, vx_scan); pub struct VxArrayIt; crate::box_wrapper!(VxArrayIt, vx_array_it); -pub type vx_list_callback = Option; -pub type vx_glob_callback = Option; +pub type vx_list_callback = Option; +pub type vx_glob_callback = Option; -pub type FsUseVortex = Option; -pub type FsSetUserdata = Option; +pub type vx_fs_use_vortex = Option c_int>; +pub type vx_fs_set_userdata = Option; -pub type FsOpen = Option; -pub type FsCreate = Option; +pub type vx_fs_open = Option; +pub type vx_fs_create = Option; -pub type FsList = Option; +pub type vx_fs_list = Option; -pub type FsClose = Option; -pub type FsSize = Option u64>; +pub type vx_fs_close = Option; +pub type vx_fs_size = Option u64>; -pub type FsRead = Option< - unsafe extern "C" fn(vx_file_handle, usize, usize, *mut u8, *mut usize, vx_error_out) -> u64, +pub type vx_fs_read = Option u64 >; -pub type FsWrite = Option< +pub type vx_fs_write = Option< unsafe extern "C" fn(vx_file_handle, usize, usize, *mut u8, *mut usize, vx_error_out) -> u64, >; @@ -73,15 +90,23 @@ pub type CacheDelete = Option *const vx_data_source { + ptr::null_mut() +} + + #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) @@ -103,7 +137,7 @@ enum VxCardinality { } #[repr(C)] -struct vx_data_source_row_count { +pub struct vx_data_source_row_count { cardinality: VxCardinality, rows: u64, } @@ -129,7 +163,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_get_row_count( } #[repr(C)] -enum ScanSelection { +pub enum ScanSelection { VX_S_INCLUDE_ALL = 0, VX_S_INCLUDE_RANGE = 1, VX_S_EXCLUDE_RANGE = 2, @@ -144,7 +178,7 @@ struct vx_scan_selection { // TODO why make it distinct from ScanRequest? Easier option handling #[repr(C)] -struct vx_scan_options { +pub struct vx_scan_options { projection: *const vx_expression, filter: *const vx_expression, row_range_begin: u64, @@ -154,15 +188,16 @@ struct vx_scan_options { limit: u64, } -enum VxPartitionEstimate { +#[repr(C)] +pub enum vx_partition_estimate_boundary { VX_P_ESTIMATE_UNKNOWN = 0, VX_P_ESTIMATE_EXACT = 1, VX_P_ESTIMATE_INEXACT = 2, } #[repr(C)] -struct vx_partition_estimate { - boundary: VxPartitionEstimate, +pub struct vx_partition_estimate { + boundary: vx_partition_estimate_boundary, estimate: u64, } @@ -176,18 +211,18 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( let request = if opts.is_null() { ScanRequest::default() } else { - let opts = &unsafe { *opts }; + let opts = unsafe { &*opts }; let projection = if opts.projection.is_null() { panic!() // TODO } else { - *vx_expression::as_ref(opts.projection) + vx_expression::as_ref(opts.projection).clone() }; let filter = if opts.filter.is_null() { None } else { - Some(*vx_expression::as_ref(opts.filter)) + Some(vx_expression::as_ref(opts.filter).clone()) }; let selection = &opts.selection; @@ -232,18 +267,27 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( limit, } }; - try_or_default(err, || { + //try_or_default(err, || { let scan = vx_data_source::as_ref(ds).scan(request); - Ok(vx_scan::new(scan)) - }) + //Ok(vx_scan::new(scan)) + //}) + ptr::null_mut() } +#[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scan_next( - scan: vx_scan, + scan: *const vx_scan, err: vx_error_out, ) -> *const vx_partition { + ptr::null_mut() + //try_or_default(err, || { + // Ok(vx_partition::new(Arc::new())) + //}) } +// TODO how to export nanoarrow headers? + +#[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( part: *const vx_partition, out: *mut FFI_ArrowArrayStream, @@ -251,10 +295,22 @@ pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( ) { } +#[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_partition_next( partition: *const vx_partition, err: vx_error_out, ) -> *const vx_array { + ptr::null_mut() + //try_or_default(err, || { + // Ok(vx_array::new(Arc::new())) + //}) } -pub unsafe extern "C-unwind" fn vx_scan_cancel(_: vx_scan) {} +#[unsafe(no_mangle)] +/// Scan progress between 0.0 and 1.0 +pub unsafe extern "C-unwind" fn vx_scan_progress(_scan: *const vx_scan) -> f64 { + 0.0 +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_scan_cancel(_: *const vx_scan) {} From 41bb5efe70aff6b34ea358e511e099e6dcef8653 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Wed, 18 Mar 2026 18:22:58 +0000 Subject: [PATCH 09/21] better --- vortex-duckdb/cpp/table_function.cpp | 21 +++- vortex-ffi/cinclude/vortex.h | 96 ++++++++++------ vortex-ffi/src/error.rs | 14 ++- vortex-ffi/src/scan.rs | 165 +++++++++++++++++++-------- 4 files changed, 203 insertions(+), 93 deletions(-) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 459cce3c992..7d4e4c205ee 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -36,10 +36,17 @@ struct CTableBindData final : TableFunctionData { CTableBindData( unique_ptr info_p, unique_ptr ffi_data_p, + const vx_session* session, const vx_data_source* data_source) : info(std::move(info_p)) , ffi_data(std::move(ffi_data_p)), - , data_source(data_source) { } + , session(session) + , data_source(data_source) {} + + ~CTableBindData() override { + vx_data_source_free(data_source); + vx_session_free(session); + } unique_ptr Copy() const override { assert(info->vtab.bind_data_clone != nullptr); @@ -53,13 +60,12 @@ struct CTableBindData final : TableFunctionData { unique_ptr(reinterpret_cast(copied_ffi_data))); } - // TODO remove with + // TODO remove unique_ptr info; - unique_ptr ffi_data; - // ^ In Rust it contains - // filter expressions + // ^ In Rust it contains filter expressions + const vx_session* session; const vx_data_source* data_source; }; @@ -119,7 +125,9 @@ unique_ptr c_bind( const vx_error* err = nullptr; vx_data_source_options opts = {0}; opts.files = input->inputs[0]; - const vx_data_source* data_source = vx_data_source_new(opts, &err); + + const vx_session* session = vx_session_new(); + const vx_data_source* data_source = vx_data_source_new(session, opts, &err); const vx_dtype* dtype = vx_data_source_dtype(data_source); const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); @@ -138,6 +146,7 @@ unique_ptr c_bind( return make_uniq( make_uniq(info.vtab), unique_ptr(reinterpret_cast(ffi_bind_data)), + session, data_source); } diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 8266f6bed89..fdc2f59b2aa 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -309,6 +309,12 @@ typedef struct vx_array_sink vx_array_sink; */ typedef struct vx_binary vx_binary; +/** + * A data source is a reference to multiple possibly remote files. When + * created, it opens first file to determine the schema from DType, all + * other operations are deferred till a scan is requested. You can request + * multiple file scans from a data source + */ typedef struct vx_data_source vx_data_source; /** @@ -419,59 +425,77 @@ typedef struct { unsigned long row_offset; } vx_file_scan_options; -typedef void (*FsUseVortex)(char, char); +typedef int (*vx_fs_use_vortex)(const char *schema, const char *path); -typedef void (*FsSetUserdata)(void); +typedef void (*vx_fs_set_userdata)(void *userdata); -typedef void (*FsOpen)(void, char, vx_error_out); +typedef void (*vx_fs_open)(void *userdata, const char *path, vx_error_out err); -typedef void (*FsCreate)(void, char, vx_error_out); +typedef void (*vx_fs_create)(void *userdata, const char *path, vx_error_out err); -typedef void (*vx_list_callback)(void, char, int); +typedef void (*vx_list_callback)(void *userdata, const char *name, int is_dir); -typedef void (*FsList)(void, char, vx_list_callback, vx_error_out); +typedef void (*vx_fs_list)(const void *userdata, + const char *path, + vx_list_callback callback, + vx_error_out error); typedef const FileHandle *vx_file_handle; -typedef void (*FsClose)(vx_file_handle); +typedef void (*vx_fs_close)(vx_file_handle handle); + +typedef uint64_t (*vx_fs_size)(vx_file_handle handle, vx_error_out err); + +typedef uint64_t ( + *vx_fs_read)(vx_file_handle handle, uint64_t offset, size_t len, uint8_t *buffer, vx_error_out err); + +typedef uint64_t ( + *vx_fs_write)(vx_file_handle handle, uint64_t offset, size_t len, uint8_t *buffer, vx_error_out err); -typedef uint64_t (*FsSize)(vx_file_handle, vx_error_out); +typedef void (*vx_fs_sync)(vx_file_handle handle, vx_error_out err); -typedef uint64_t (*FsRead)(vx_file_handle, size_t, size_t, uint8_t *, size_t *, vx_error_out); +typedef void (*vx_glob_callback)(void *userdata, const char *file); -typedef uint64_t (*FsWrite)(vx_file_handle, size_t, size_t, uint8_t *, size_t *, vx_error_out); +typedef void (*vx_glob)(const char *glob, vx_glob_callback callback, vx_error_out err); -typedef void (*FsSync)(vx_file_handle, vx_error_out); +typedef void *vx_cache; -typedef void (*vx_glob_callback)(void, char); +typedef vx_cache (*vx_cache_init)(vx_error_out err); -typedef void (*Glob)(char, vx_glob_callback, vx_error_out); +typedef void (*vx_cache_free)(vx_cache cache, vx_error_out err); -typedef void (*CacheInit)(vx_error_out); +typedef const char *vx_cache_key; -typedef void (*CacheGet)(void, char, void *, vx_error_out); +typedef void (*vx_cache_get)(vx_cache cache, vx_cache_key key, void **value, vx_error_out err); -typedef void (*CachePut)(void, char, void, vx_error_out); +typedef void (*vx_cache_put)(vx_cache cache, vx_cache_key key, void *value, vx_error_out err); -typedef void (*CacheDelete)(void, char, vx_error_out); +typedef void (*vx_cache_delete)(vx_cache cache, vx_cache_key key, vx_error_out err); typedef struct { const char *files; - FsUseVortex fs_use_vortex; - FsSetUserdata fs_set_userdata; - FsOpen fs_open; - FsCreate fs_create; - FsList fs_list; - FsClose fs_close; - FsSize fs_size; - FsRead fs_read; - FsWrite fs_write; - FsSync fs_sync; - Glob glob; - CacheInit cache_init; - CacheGet cache_get; - CachePut cache_put; - CacheDelete cache_delete; + /** + * Whether to use Vortex filesystem or host's filesystem. + * Return 1 to use Vortex for a given schema ("file", "s3") and path. + * Return 0 to use host's filesystem. If 0 is returned, host must also + * Implement all fs_* callbacks. + */ + vx_fs_use_vortex fs_use_vortex; + vx_fs_set_userdata fs_set_userdata; + vx_fs_open fs_open; + vx_fs_create fs_create; + vx_fs_list fs_list; + vx_fs_close fs_close; + vx_fs_size fs_size; + vx_fs_read fs_read; + vx_fs_write fs_write; + vx_fs_sync fs_sync; + vx_glob glob; + vx_cache_init cache_init; + vx_cache_free cache_free; + vx_cache_get cache_get; + vx_cache_put cache_put; + vx_cache_delete cache_delete; } vx_data_source_options; typedef struct { @@ -902,7 +926,11 @@ void vx_scan_free(const vx_scan *ptr); */ void vx_array_it_free(vx_array_it *ptr); -const vx_data_source *vx_data_source_new(const vx_data_source_options *opts, vx_error_out err); +/** + * Create a new owned datasource which must be freed by the caller + */ +const vx_data_source * +vx_data_source_new(const vx_session *session, const vx_data_source_options *opts, vx_error_out err); const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); @@ -924,7 +952,7 @@ const vx_array *vx_partition_next(const vx_partition *partition, vx_error_out er */ double vx_scan_progress(const vx_scan *_scan); -void vx_scan_cancel(const vx_scan *); +void vx_scan_cancel(const vx_scan *_scan); /** * Free an owned [`vx_session`] object. diff --git a/vortex-ffi/src/error.rs b/vortex-ffi/src/error.rs index 888967e25dc..372a4a07775 100644 --- a/vortex-ffi/src/error.rs +++ b/vortex-ffi/src/error.rs @@ -22,9 +22,16 @@ box_wrapper!( #[allow(non_camel_case_types)] pub type vx_error_out = *mut *mut vx_error; +pub(crate) fn write_error(error_out: vx_error_out, error_string: &str) { + let err = vx_error::new(Box::new(VortexError { + message: error_string.into(), + })); + unsafe { error_out.write(err) }; +} + #[inline] pub fn try_or_default( - error_out: *mut *mut vx_error, + error_out: vx_error_out, function: impl FnOnce() -> VortexResult, ) -> T { match function() { @@ -33,10 +40,7 @@ pub fn try_or_default( value } Err(err) => { - let err = vx_error::new(Box::new(VortexError { - message: err.to_string().into(), - })); - unsafe { error_out.write(err) }; + write_error(error_out, &err.to_string()); T::default() } } diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 8a94c216f64..96ead24d7fd 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -11,19 +11,27 @@ use std::sync::Arc; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use vortex::buffer::Buffer; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; use vortex::expr::Expression; use vortex::expr::stats::Precision::Exact; use vortex::expr::stats::Precision::Inexact; +use vortex::file::multi::MultiFileDataSource; +use vortex::io::runtime::BlockingRuntime; use vortex::scan::Selection; use vortex::scan::api::DataSource; +use vortex::scan::api::DataSourceRef; use vortex::scan::api::DataSourceScan; use vortex::scan::api::Partition; use vortex::scan::api::ScanRequest; +use crate::RUNTIME; use crate::array::vx_array; use crate::dtype::vx_dtype; -//use crate::error::try_or_default; +use crate::error::try_or_default; use crate::error::vx_error_out; +use crate::session::vx_session; +use crate::to_string; crate::arc_dyn_wrapper!( /// A data source is a reference to multiple possibly remote files. When @@ -46,47 +54,77 @@ crate::arc_dyn_wrapper!(dyn DataSourceScan, vx_scan); pub struct VxArrayIt; crate::box_wrapper!(VxArrayIt, vx_array_it); -pub type vx_list_callback = Option; -pub type vx_glob_callback = Option; - -pub type vx_fs_use_vortex = Option c_int>; -pub type vx_fs_set_userdata = Option; - -pub type vx_fs_open = Option; -pub type vx_fs_create = Option; - -pub type vx_fs_list = Option; +pub type vx_list_callback = + Option; +pub type vx_glob_callback = + Option; + +pub type vx_fs_use_vortex = + Option c_int>; +pub type vx_fs_set_userdata = Option; + +pub type vx_fs_open = + Option; +pub type vx_fs_create = + Option; + +pub type vx_fs_list = Option< + unsafe extern "C" fn( + userdata: *const c_void, + path: *const c_char, + callback: vx_list_callback, + error: vx_error_out, + ), +>; pub type vx_fs_close = Option; -pub type vx_fs_size = Option u64>; - -pub type vx_fs_read = Option u64 +pub type vx_fs_size = + Option u64>; + +pub type vx_fs_read = Option< + unsafe extern "C" fn( + handle: vx_file_handle, + offset: u64, + len: usize, + buffer: *mut u8, + err: vx_error_out, + ) -> u64, >; pub type vx_fs_write = Option< - unsafe extern "C" fn(vx_file_handle, usize, usize, *mut u8, *mut usize, vx_error_out) -> u64, + unsafe extern "C" fn( + handle: vx_file_handle, + offset: u64, + len: usize, + buffer: *mut u8, + err: vx_error_out, + ) -> u64, >; -pub type FsSync = Option; +pub type vx_fs_sync = Option; -pub type Glob = Option; +pub type vx_glob = Option< + unsafe extern "C" fn(glob: *const c_char, callback: vx_glob_callback, err: vx_error_out), +>; -pub type CacheInit = Option; -pub type CacheFree = Option; -pub type CacheGet = Option; -pub type CachePut = Option; -pub type CacheDelete = Option; +pub type vx_cache = *mut c_void; +pub type vx_cache_key = *const c_char; + +pub type vx_cache_init = Option vx_cache>; +pub type vx_cache_free = Option; +pub type vx_cache_get = Option< + unsafe extern "C" fn( + cache: vx_cache, + key: vx_cache_key, + value: *mut *mut c_void, + err: vx_error_out, + ), +>; +pub type vx_cache_put = Option< + unsafe extern "C" fn(cache: vx_cache, key: vx_cache_key, value: *mut c_void, err: vx_error_out), +>; +pub type vx_cache_delete = + Option; #[repr(C)] pub struct vx_data_source_options { @@ -98,7 +136,6 @@ pub struct vx_data_source_options { /// Return 0 to use host's filesystem. If 0 is returned, host must also /// Implement all fs_* callbacks. fs_use_vortex: vx_fs_use_vortex, - fs_set_userdata: vx_fs_set_userdata, fs_open: vx_fs_open, fs_create: vx_fs_create, @@ -107,23 +144,55 @@ pub struct vx_data_source_options { fs_size: vx_fs_size, fs_read: vx_fs_read, fs_write: vx_fs_write, - fs_sync: FsSync, - glob: Glob, - cache_init: CacheInit, - cache_get: CacheGet, - cache_put: CachePut, - cache_delete: CacheDelete, + fs_sync: vx_fs_sync, + + glob: vx_glob, + + cache_init: vx_cache_init, + cache_free: vx_cache_free, + cache_get: vx_cache_get, + cache_put: vx_cache_put, + cache_delete: vx_cache_delete, } +unsafe fn data_source_new( + session: *const vx_session, + opts: *const vx_data_source_options, +) -> VortexResult<*const vx_data_source> { + if session.is_null() { + vortex_bail!("empty session"); + } + let session = vx_session::as_ref(session).clone(); + + if opts.is_null() { + vortex_bail!("empty opts"); + } + let opts = unsafe { &*opts }; + + if opts.files.is_null() { + vortex_bail!("empty opts.files"); + } + let glob = unsafe { to_string(opts.files) }; + + RUNTIME.block_on(async { + let builder = MultiFileDataSource::new(session) + .with_filesystem(fs) + .with_glob(glob); + let data_source = builder.build().await?; + Ok(vx_data_source::new(Arc::new(data_source) as DataSourceRef)) + }) +} + +/// Create a new owned datasource which must be freed by the caller #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_new( + session: *const vx_session, opts: *const vx_data_source_options, - err: vx_error_out -)-> *const vx_data_source { - ptr::null_mut() + err: vx_error_out, +) -> *const vx_data_source { + try_or_default(err, || unsafe { data_source_new(session, opts) }) } - #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) @@ -176,7 +245,7 @@ struct vx_scan_selection { idx_len: usize, } -// TODO why make it distinct from ScanRequest? Easier option handling +// Distinct from ScanRequest for easier option handling from C #[repr(C)] pub struct vx_scan_options { projection: *const vx_expression, @@ -268,7 +337,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( } }; //try_or_default(err, || { - let scan = vx_data_source::as_ref(ds).scan(request); + let scan = vx_data_source::as_ref(ds).scan(request); //Ok(vx_scan::new(scan)) //}) ptr::null_mut() @@ -285,7 +354,7 @@ pub unsafe extern "C-unwind" fn vx_scan_next( //}) } -// TODO how to export nanoarrow headers? +// TODO export nanoarrow headers? #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( @@ -313,4 +382,4 @@ pub unsafe extern "C-unwind" fn vx_scan_progress(_scan: *const vx_scan) -> f64 { } #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_scan_cancel(_: *const vx_scan) {} +pub unsafe extern "C-unwind" fn vx_scan_cancel(_scan: *const vx_scan) {} From 2853a77015027e4af93bbc3be7abcce2bdf2fd4a Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 11:30:08 +0000 Subject: [PATCH 10/21] better --- vortex-array/src/dtype/field_names.rs | 6 + vortex-duckdb/cpp/table_function.cpp | 333 +++++++++--------------- vortex-ffi/cinclude/vortex.h | 217 ++++++++-------- vortex-ffi/src/data_source.rs | 212 ++++++++++++++++ vortex-ffi/src/dtype.rs | 6 +- vortex-ffi/src/expression.rs | 27 ++ vortex-ffi/src/lib.rs | 2 + vortex-ffi/src/scan.rs | 349 +++++--------------------- 8 files changed, 553 insertions(+), 599 deletions(-) create mode 100644 vortex-ffi/src/data_source.rs create mode 100644 vortex-ffi/src/expression.rs diff --git a/vortex-array/src/dtype/field_names.rs b/vortex-array/src/dtype/field_names.rs index 737ad0cfd52..f148954b30b 100644 --- a/vortex-array/src/dtype/field_names.rs +++ b/vortex-array/src/dtype/field_names.rs @@ -341,6 +341,12 @@ impl From> for FieldNames { } } +impl From> for FieldNames { + fn from(value: Vec) -> Self { + Self(value.into_iter().map(FieldName::from).collect()) + } +} + impl From<&[FieldName]> for FieldNames { fn from(value: &[FieldName]) -> Self { Self(Arc::from(value)) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 7d4e4c205ee..f86954446dd 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -1,6 +1,9 @@ // 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/duckdb_diagnostics.h" DUCKDB_INCLUDES_BEGIN @@ -13,16 +16,19 @@ DUCKDB_INCLUDES_BEGIN #include "duckdb/parser/parsed_data/create_table_function_info.hpp" DUCKDB_INCLUDES_END -#include "duckdb_vx.h" -#include "duckdb_vx/data.hpp" -#include "duckdb_vx/error.hpp" - using namespace duckdb; std::string to_string(const vx_string* str) { return { vx_string_ptr(str), size_t vx_string_len(str) }; } +std::string move_vx_err(vx_error* error) { + const vx_string* vx_str = vx_error_get_message(error); + string str{vx_string_ptr(vx_str), vx_string_len(vx_str)}; + vx_error_free(error); + return str; +} + namespace vortex { struct CTableFunctionInfo final : TableFunctionInfo { explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) @@ -36,10 +42,12 @@ struct CTableBindData final : TableFunctionData { CTableBindData( unique_ptr info_p, unique_ptr ffi_data_p, + const vector& column_names, const vx_session* session, const vx_data_source* data_source) : info(std::move(info_p)) , ffi_data(std::move(ffi_data_p)), + , column_names(column_names) , session(session) , data_source(data_source) {} @@ -65,53 +73,35 @@ struct CTableBindData final : TableFunctionData { unique_ptr ffi_data; // ^ In Rust it contains filter expressions + // pre-computed from data_source using vx_struct_fields_field_name, see c_bind() + vector column_names; + const vx_session* session; const vx_data_source* data_source; }; -std::string move_vx_err(vx_error* error) { - const vx_string * error_vx_str = vx_error_get_message(error); - string str{vx_string_ptr(error_vx_str), vx_string_len(error_vx_str)}; - vx_error_free(error); - return str; -} - struct CTableGlobalData final : GlobalTableFunctionState { - explicit CTableGlobalData( - unique_ptr ffi_data_p, - idx_t max_threads, - const vx_scan_options& options) - : ffi_data(std::move(ffi_data_p)) - , max_threads(max_threads) - { - vx_error* error; - if (scan = vx_scan_init(&options, &error); !scan) - throw BinderException(move_vx_err(error)); - } + explicit CTableGlobalData(const vx_scan* scan, idx_t max_threads) + : scan(scan), max_threads(max_threads) {} ~CTableGlobalData() override { - if (scan) vx_scan_free(scan); + vx_scan_free(scan); } - // TODO Ideally should be C++ api for lifetime checks - const vx_scan * scan; + idx_t MaxThreads() const override { return max_threads; } - unique_ptr ffi_data; + const vx_scan* scan; idx_t max_threads; - - idx_t MaxThreads() const override { - return max_threads; - } }; struct CTableLocalData final : LocalTableFunctionState { - explicit CTableLocalData(unique_ptr ffi_data_p) : ffi_data(std::move(ffi_data_p)) { - } - - unique_ptr ffi_data; + explicit CTableLocalData() {} + std::optional batch_id; }; -LogicalType from_dtype(const vx_dtype* dtype) { } +LogicalType from_dtype(const vx_dtype* dtype) { + return LogicalType::NULL; +} unique_ptr c_bind( ClientContext &context, @@ -146,35 +136,62 @@ unique_ptr c_bind( return make_uniq( make_uniq(info.vtab), unique_ptr(reinterpret_cast(ffi_bind_data)), + names, session, data_source); } -unique_ptr c_init_global(ClientContext &context, TableFunctionInitInput &input) { - const auto &bind = input.bind_data->Cast(); +const vx_expression* projection( + const vector& column_ids, + const vector& projection_ids, + const vector& column_names +) { + vector projected_column_names; + projected_column_names.reserve(column_names); + + if (projection_ids.empty()) { + for (column_t id : column_ids) { + if (column_names.size() < id) + throw InvalidInputException(); + string& name = column_names[id]; + projected_column_names.emplace_back(name.data()); + } + } else { + for (idx_t projection_id : projection_ids) { + column_it id = column_ids[projection_id]; + if (column_names.size() < id) + throw InvalidInputException(); + string& name = column_names[id]; + projected_column_names.emplace_back(name.data()); + } + } - duckdb_vx_tfunc_init_input ffi_input = { - .bind_data = bind.ffi_data->DataPtr(), - .column_ids = input.column_ids.data(), - .column_ids_count = input.column_ids.size(), - .projection_ids = input.projection_ids.data(), - .projection_ids_count = input.projection_ids.size(), - .filters = reinterpret_cast(input.filters.get()), - .client_context = reinterpret_cast(&context), - }; + const vx_expression* root = vx_expression_root(); + // TODO vx_expression may take a string vx_array + const vx_expression* expr = vx_expression_select( + projected_column_names.data(), + projected_column_names.size(), + root); + vx_expression_free(root); + return expr; +} - duckdb_vx_error error_out = nullptr; - auto ffi_global_data = bind.info->vtab.init_global(&ffi_input, &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } +unique_ptr c_init_global(ClientContext &, TableFunctionInitInput &input) { + const auto &bind = input.bind_data->Cast(); - vx_scan_options scan_options = { - .max_threads = bind.info->max_threads, - }; + vx_scan_options options = {0}; + options.projection = projection(input.column_ids, input.projection_ids, bind.names); + option.filter = filter(); + + vx_error* err = nullptr; + const vx_scan* scan = vx_data_source_scan( + bind.data_source, + options, + /* partition estimate */ nullptr, + &err); + if (err) throw InvalidInputException(move_vx_err(err)); - auto cdata = unique_ptr(reinterpret_cast(ffi_global_data)); - return make_uniq(std::move(cdata), scan_options); + return make_uniq(scan, bind.info->max_threads); } unique_ptr c_init_local(ExecutionContext &context, @@ -223,30 +240,30 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou } } -void c_pushdown_complex_filter(ClientContext & /*context*/, - LogicalGet & /*get*/, - FunctionData *bind_data, - vector> &filters) { - if (filters.empty()) { - return; - } - - auto &bind = bind_data->Cast(); - - for (auto iter = filters.begin(); iter != filters.end();) { - duckdb_vx_error error_out = nullptr; - auto pushed = - bind.info->vtab.pushdown_complex_filter(bind_data->Cast().ffi_data->DataPtr(), - reinterpret_cast(iter->get()), - &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } - - // If the pushdown complex filter returns true, we can remove the filter from the list. - iter = pushed ? filters.erase(iter) : std::next(iter); - } -} +//void c_pushdown_complex_filter(ClientContext & /*context*/, +// LogicalGet & /*get*/, +// FunctionData *bind_data, +// vector> &filters) { +// if (filters.empty()) { +// return; +// } +// +// auto &bind = bind_data->Cast(); +// +// for (auto iter = filters.begin(); iter != filters.end();) { +// duckdb_vx_error error_out = nullptr; +// auto pushed = +// bind.info->vtab.pushdown_complex_filter(bind_data->Cast().ffi_data->DataPtr(), +// reinterpret_cast(iter->get()), +// &error_out); +// if (error_out) { +// throw BinderException(IntoErrString(error_out)); +// } +// +// // If the pushdown complex filter returns true, we can remove the filter from the list. +// iter = pushed ? filters.erase(iter) : std::next(iter); +// } +//} unique_ptr c_cardinality(ClientContext &, const FunctionData *bind_data) { auto stats = make_uniq(); @@ -274,23 +291,6 @@ unique_ptr c_cardinality(ClientContext &, const FunctionData *bi } } -extern "C" size_t duckdb_vx_tfunc_bind_input_get_parameter_count(duckdb_vx_tfunc_bind_input ffi_input) { - if (!ffi_input) { - return 0; - } - const auto input = reinterpret_cast(ffi_input); - return input->inputs.size(); -} - -extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, - size_t index) { - if (!ffi_input || index >= duckdb_vx_tfunc_bind_input_get_parameter_count(ffi_input)) { - return nullptr; - } - const auto info = reinterpret_cast(ffi_input); - return reinterpret_cast(new Value(info->inputs[index])); -} - extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_named_parameter(duckdb_vx_tfunc_bind_input ffi_input, const char *name) { if (!ffi_input || !name) { @@ -305,86 +305,13 @@ extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_named_parameter(duckdb_vx return reinterpret_cast(value.release()); } -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(ffi_result); - const auto logical_type = reinterpret_cast(ffi_type); - - result->names.emplace_back(name_str, name_len); - result->return_types.emplace_back(*logical_type); -} - -virtual_column_map_t c_get_virtual_columns(ClientContext & /*context*/, - optional_ptr bind_data) { - auto &bind = bind_data->Cast(); - - auto result = virtual_column_map_t(); - bind.info->vtab.get_virtual_columns(bind_data->Cast().ffi_data->DataPtr(), - reinterpret_cast(&result)); - return result; -} - -extern "C" void duckdb_vx_tfunc_virtual_columns_push(duckdb_vx_tfunc_virtual_cols_result ffi_result, - idx_t column_idx, - const char *name_str, - size_t name_len, - duckdb_logical_type ffi_type) { - if (!ffi_result || !name_str || !ffi_type) { - return; - } - - auto result = reinterpret_cast(ffi_result); - const auto logical_type = reinterpret_cast(ffi_type); - const auto name = string(name_str, name_len); - - auto table_col = TableColumn(std::move(name), *logical_type); - result->emplace(column_idx, std::move(table_col)); -} - -OperatorPartitionData c_get_partition_data(ClientContext & /*context*/, +OperatorPartitionData c_get_partition_data(ClientContext&, TableFunctionGetPartitionInput &input) { - if (input.partition_info.RequiresPartitionColumns()) { + if (input.partition_info.RequiresPartitionColumns()) throw InternalException("TableScan::GetPartitionData: partition columns not supported"); - } - auto &bind = input.bind_data->Cast(); - auto &global = input.global_state->Cast(); - auto &local = input.local_state->Cast(); - - duckdb_vx_error error_out = nullptr; - auto index = bind.info->vtab.get_partition_data(bind.ffi_data->DataPtr(), - global.ffi_data->DataPtr(), - local.ffi_data->DataPtr(), - &error_out); - if (error_out) { - throw InvalidInputException(IntoErrString(error_out)); - } - return OperatorPartitionData(index); -} - -InsertionOrderPreservingMap c_to_string(TableFunctionToStringInput &input) { - InsertionOrderPreservingMap result; - auto &bind = input.bind_data->Cast(); - - // Call the Rust side to get custom string representation if available - if (bind.info->vtab.to_string) { - auto map = bind.info->vtab.to_string(bind.ffi_data->DataPtr()); - if (map) { - // Copy the map contents to the result - auto *cpp_map = reinterpret_cast *>(map); - for (const auto &[key, value] : *cpp_map) { - result[key] = value; - } - // Free the map allocated by Rust - duckdb_vx_string_map_free(map); - } - } - - return result; + if (auto &batch_id = input.local_state->Cast().batch_id; batch_id) + return OperatorPartitionData(*batch_id); + throw InvalidInputException("Batch id missing, no batches exported"); } extern "C" duckdb_state duckdb_vx_tfunc_register( @@ -399,40 +326,49 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( auto db = wrapper->database->instance; auto tf = TableFunction(vtab->name, {}, c_function, c_bind, c_init_global, c_init_local); - tf.pushdown_complex_filter = c_pushdown_complex_filter; + //tf.pushdown_complex_filter = c_pushdown_complex_filter; + + //tf.projection_pushdown = vtab->projection_pushdown; + //tf.filter_pushdown = vtab->filter_pushdown; + //tf.filter_prune = vtab->filter_prune; + //tf.sampling_pushdown = vtab->sampling_pushdown; + //tf.late_materialization = vtab->late_materialization; + + tf.projection_pushdown = true; + tf.filter_pushdown = false; + tf.filter_prune = false + tf.sampling_pushdown = false; + tf.late_materialization = false; - tf.projection_pushdown = vtab->projection_pushdown; - tf.filter_pushdown = vtab->filter_pushdown; - tf.filter_prune = vtab->filter_prune; - tf.sampling_pushdown = vtab->sampling_pushdown; - tf.late_materialization = vtab->late_materialization; tf.cardinality = c_cardinality; tf.get_partition_data = c_get_partition_data; - tf.to_string = c_to_string; + + tf.to_string = [](TableFunctionToStringInput&) { + InsertionOrderPreservingMap result; + result.emplace("Function", "Vortex Scan"); + // TODO filters + return result; + } tf.get_virtual_columns = [](auto&, auto) { - auto res = virtual_column_map_t(); - res->emplace(COLUMN_IDENTIFIER_EMPTY, "", LogicalType::bool()); - return; + return {{COLUMN_IDENTIFIER_EMPTY, TableColumn{"", LogicalType::BOOL}}}; } tf.table_scan_progress = [](auto&, auto*, const GlobalTableFunctionState& state) { return vx_scan_progress(state->Cast().scan); }; - // Set up the parameters tf.arguments.reserve(vtab->parameter_count); for (size_t i = 0; i < vtab->parameter_count; i++) { auto logical_type = reinterpret_cast(vtab->parameters[i]); tf.arguments.emplace_back(*logical_type); } - // And the named parameters + for (size_t i = 0; i < vtab->named_parameter_count; i++) { auto logical_type = reinterpret_cast(vtab->named_parameter_types[i]); tf.named_parameters.emplace(vtab->named_parameter_names[i], *logical_type); } - // Assign the VTable to the function info so we can access it later to invoke the callbacks. tf.function_info = make_shared_ptr(*vtab); try { @@ -445,25 +381,4 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( } return DuckDBSuccess; } - -extern "C" duckdb_vx_string_map duckdb_vx_string_map_create() { - auto map = new InsertionOrderPreservingMap(); - return reinterpret_cast(map); -} - -extern "C" void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, const char *value) { - if (!map || !key || !value) { - return; - } - auto *cpp_map = reinterpret_cast *>(map); - (*cpp_map)[string(key)] = string(value); -} - -extern "C" void duckdb_vx_string_map_free(duckdb_vx_string_map map) { - if (!map) { - return; - } - auto *cpp_map = reinterpret_cast *>(map); - delete cpp_map; -} } // namespace vortex diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index fdc2f59b2aa..f3b7a34609c 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -21,6 +21,12 @@ typedef ArrowArrayStream FFI_ArrowArrayStream; */ #define BinaryView_MAX_INLINED_SIZE 12 +typedef enum { + VX_CARD_UNKNOWN = 0, + VX_CARD_ESTIMATE = 1, + VX_CARD_MAXIMUM = 2, +} vx_cardinality; + /** * Variant enum for Vortex primitive types. */ @@ -147,17 +153,11 @@ typedef enum { LOG_LEVEL_TRACE = 5, } vx_log_level; -typedef enum { - VX_CARD_UNKNOWN = 0, - VX_CARD_ESTIMATE = 1, - VX_CARD_MAXIMUM = 2, -} VxCardinality; - typedef enum { VX_S_INCLUDE_ALL = 0, VX_S_INCLUDE_RANGE = 1, VX_S_EXCLUDE_RANGE = 2, -} ScanSelection; +} vx_scan_selection_include; typedef enum { VX_P_ESTIMATE_UNKNOWN = 0, @@ -361,74 +361,12 @@ typedef struct vx_struct_fields vx_struct_fields; */ typedef struct vx_struct_fields_builder vx_struct_fields_builder; -typedef vx_error **vx_error_out; - -/** - * Options supplied for opening a file. - */ -typedef struct { - /** - * URI for opening the file. - * This must be a valid URI, even for files (file:///path/to/file) - */ - const char *uri; - /** - * Additional configuration for the file source (e.g. "s3.accessKey"). - * This may be null, in which case it is treated as empty. - */ - const char *const *property_keys; - /** - * Additional configuration values for the file source (e.g. S3 credentials). - */ - const char *const *property_vals; - /** - * Number of properties in `property_keys` and `property_vals`. - */ - int property_len; -} vx_file_open_options; - -/** - * Scan options provided by an FFI client calling the `vx_file_scan` function. - */ -typedef struct { - /** - * Column names to project out in the scan. These must be null-terminated C strings. - */ - const char *projection_expression; - /** - * Number of columns in `projection`. - */ - unsigned int projection_expr_len; - /** - * Serialized expressions for pushdown - */ - const char *filter_expression; - /** - * The len in bytes of the filter expression - */ - unsigned int filter_expression_len; - /** - * Splits the file into chunks of this size, if zero then we use the write layout. - */ - int split_by_row_count; - /** - * First row of a range to scan. - */ - unsigned long row_range_start; - /** - * Last row of a range to scan. - */ - unsigned long row_range_end; - /** - * The row offset of the file in a multi-file scan. - */ - unsigned long row_offset; -} vx_file_scan_options; - typedef int (*vx_fs_use_vortex)(const char *schema, const char *path); typedef void (*vx_fs_set_userdata)(void *userdata); +typedef vx_error **vx_error_out; + typedef void (*vx_fs_open)(void *userdata, const char *path, vx_error_out err); typedef void (*vx_fs_create)(void *userdata, const char *path, vx_error_out err); @@ -472,13 +410,15 @@ typedef void (*vx_cache_put)(vx_cache cache, vx_cache_key key, void *value, vx_e typedef void (*vx_cache_delete)(vx_cache cache, vx_cache_key key, vx_error_out err); +/** + * Host must either implement all or none of fs_* callbacks. + */ typedef struct { const char *files; /** * Whether to use Vortex filesystem or host's filesystem. * Return 1 to use Vortex for a given schema ("file", "s3") and path. - * Return 0 to use host's filesystem. If 0 is returned, host must also - * Implement all fs_* callbacks. + * Return 0 to use host's filesystem. */ vx_fs_use_vortex fs_use_vortex; vx_fs_set_userdata fs_set_userdata; @@ -499,12 +439,74 @@ typedef struct { } vx_data_source_options; typedef struct { - VxCardinality cardinality; + vx_cardinality cardinality; uint64_t rows; } vx_data_source_row_count; +/** + * Options supplied for opening a file. + */ typedef struct { - ScanSelection selection; + /** + * URI for opening the file. + * This must be a valid URI, even for files (file:///path/to/file) + */ + const char *uri; + /** + * Additional configuration for the file source (e.g. "s3.accessKey"). + * This may be null, in which case it is treated as empty. + */ + const char *const *property_keys; + /** + * Additional configuration values for the file source (e.g. S3 credentials). + */ + const char *const *property_vals; + /** + * Number of properties in `property_keys` and `property_vals`. + */ + int property_len; +} vx_file_open_options; + +/** + * Scan options provided by an FFI client calling the `vx_file_scan` function. + */ +typedef struct { + /** + * Column names to project out in the scan. These must be null-terminated C strings. + */ + const char *projection_expression; + /** + * Number of columns in `projection`. + */ + unsigned int projection_expr_len; + /** + * Serialized expressions for pushdown + */ + const char *filter_expression; + /** + * The len in bytes of the filter expression + */ + unsigned int filter_expression_len; + /** + * Splits the file into chunks of this size, if zero then we use the write layout. + */ + int split_by_row_count; + /** + * First row of a range to scan. + */ + unsigned long row_range_start; + /** + * Last row of a range to scan. + */ + unsigned long row_range_end; + /** + * The row offset of the file in a multi-file scan. + */ + unsigned long row_offset; +} vx_file_scan_options; + +typedef struct { + vx_scan_selection_include include; uint64_t *idx; size_t idx_len; } vx_scan_selection; @@ -618,6 +620,29 @@ const vx_string *vx_array_get_utf8(const vx_array *array, uint32_t index); */ const vx_binary *vx_array_get_binary(const vx_array *array, uint32_t index); +/** + * Clone a borrowed [`vx_data_source`], returning an owned [`vx_data_source`]. + * + * + * Must be released with [`vx_data_source_free`]. + */ +const vx_data_source *vx_data_source_clone(const vx_data_source *ptr); + +/** + * Free an owned [`vx_data_source`] object. + */ +void vx_data_source_free(const vx_data_source *ptr); + +/** + * Create a new owned datasource which must be freed by the caller + */ +const vx_data_source * +vx_data_source_new(const vx_session *session, const vx_data_source_options *opts, vx_error_out err); + +const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); + +void vx_data_source_get_row_count(const vx_data_source *ds, vx_data_source_row_count *rc); + /** * Free an owned [`vx_array_iterator`] object. */ @@ -804,7 +829,7 @@ uint8_t vx_dtype_time_unit(const DType *dtype); */ const vx_string *vx_dtype_time_zone(const DType *dtype); -void vx_type_to_arrow_schema(const DType *dtype, FFI_ArrowSchema *out, vx_error_out err); +void vx_type_to_arrow_schema(const vx_dtype *_dtype, FFI_ArrowSchema *_schema, vx_error_out _err); /** * Free an owned [`vx_error`] object. @@ -878,17 +903,14 @@ vx_array_iterator *vx_file_scan(const vx_session *session, void vx_set_log_level(vx_log_level level); /** - * Clone a borrowed [`vx_data_source`], returning an owned [`vx_data_source`]. - * - * - * Must be released with [`vx_data_source_free`]. + * Free an owned [`vx_expression`] object. */ -const vx_data_source *vx_data_source_clone(const vx_data_source *ptr); +void vx_expression_free(vx_expression *ptr); -/** - * Free an owned [`vx_data_source`] object. - */ -void vx_data_source_free(const vx_data_source *ptr); +const vx_expression *vx_expression_root(void); + +const vx_expression * +vx_expression_select(const char *const *names, size_t names_len, const vx_expression *child); /** * Clone a borrowed [`vx_partition`], returning an owned [`vx_partition`]. @@ -903,11 +925,6 @@ const vx_partition *vx_partition_clone(const vx_partition *ptr); */ void vx_partition_free(const vx_partition *ptr); -/** - * Free an owned [`vx_expression`] object. - */ -void vx_expression_free(vx_expression *ptr); - /** * Clone a borrowed [`vx_scan`], returning an owned [`vx_scan`]. * @@ -926,20 +943,10 @@ void vx_scan_free(const vx_scan *ptr); */ void vx_array_it_free(vx_array_it *ptr); -/** - * Create a new owned datasource which must be freed by the caller - */ -const vx_data_source * -vx_data_source_new(const vx_session *session, const vx_data_source_options *opts, vx_error_out err); - -const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); - -void vx_data_source_get_row_count(const vx_data_source *ds, vx_data_source_row_count *rc); - -const vx_scan *vx_data_source_scan(const vx_data_source *ds, - const vx_scan_options *opts, - vx_partition_estimate *est, - vx_error_out err); +const vx_scan *vx_data_source_scan(const vx_data_source *_data_source, + const vx_scan_options *_options, + vx_partition_estimate *_partition_estimate, + vx_error_out _err); const vx_partition *vx_scan_next(const vx_scan *scan, vx_error_out err); @@ -952,8 +959,6 @@ const vx_array *vx_partition_next(const vx_partition *partition, vx_error_out er */ double vx_scan_progress(const vx_scan *_scan); -void vx_scan_cancel(const vx_scan *_scan); - /** * Free an owned [`vx_session`] object. */ diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs new file mode 100644 index 00000000000..560895a5175 --- /dev/null +++ b/vortex-ffi/src/data_source.rs @@ -0,0 +1,212 @@ +#![allow(non_camel_case_types)] + +use std::ffi::c_char; +use std::ffi::c_int; +use std::ffi::c_void; +use std::sync::Arc; + +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::expr::stats::Precision::Exact; +use vortex::expr::stats::Precision::Inexact; +use vortex::file::multi::MultiFileDataSource; +use vortex::io::runtime::BlockingRuntime; +use vortex::scan::api::DataSource; +use vortex::scan::api::DataSourceRef; + +use crate::RUNTIME; +use crate::dtype::vx_dtype; +use crate::error::try_or_default; +use crate::error::vx_error_out; +use crate::session::vx_session; +use crate::to_string; + +crate::arc_dyn_wrapper!( + /// A data source is a reference to multiple possibly remote files. When + /// created, it opens first file to determine the schema from DType, all + /// other operations are deferred till a scan is requested. You can request + /// multiple file scans from a data source + dyn DataSource, + vx_data_source); + +pub struct FileHandle; +pub type vx_file_handle = *const FileHandle; + +pub type vx_list_callback = + Option; +pub type vx_glob_callback = + Option; + +pub type vx_fs_use_vortex = + Option c_int>; +pub type vx_fs_set_userdata = Option; + +pub type vx_fs_open = + Option; +pub type vx_fs_create = + Option; + +pub type vx_fs_list = Option< + unsafe extern "C" fn( + userdata: *const c_void, + path: *const c_char, + callback: vx_list_callback, + error: vx_error_out, + ), +>; + +pub type vx_fs_close = Option; +pub type vx_fs_size = + Option u64>; + +pub type vx_fs_read = Option< + unsafe extern "C" fn( + handle: vx_file_handle, + offset: u64, + len: usize, + buffer: *mut u8, + err: vx_error_out, + ) -> u64, +>; + +pub type vx_fs_write = Option< + unsafe extern "C" fn( + handle: vx_file_handle, + offset: u64, + len: usize, + buffer: *mut u8, + err: vx_error_out, + ) -> u64, +>; + +pub type vx_fs_sync = Option; + +pub type vx_glob = Option< + unsafe extern "C" fn(glob: *const c_char, callback: vx_glob_callback, err: vx_error_out), +>; + +pub type vx_cache = *mut c_void; +pub type vx_cache_key = *const c_char; + +pub type vx_cache_init = Option vx_cache>; +pub type vx_cache_free = Option; +pub type vx_cache_get = Option< + unsafe extern "C" fn( + cache: vx_cache, + key: vx_cache_key, + value: *mut *mut c_void, + err: vx_error_out, + ), +>; +pub type vx_cache_put = Option< + unsafe extern "C" fn(cache: vx_cache, key: vx_cache_key, value: *mut c_void, err: vx_error_out), +>; +pub type vx_cache_delete = + Option; + +#[repr(C)] +/// Host must either implement all or none of fs_* callbacks. +pub struct vx_data_source_options { + // TODO what if the program wants to read a Vortex file from an existing buffer? + files: *const c_char, + + /// Whether to use Vortex filesystem or host's filesystem. + /// Return 1 to use Vortex for a given schema ("file", "s3") and path. + /// Return 0 to use host's filesystem. + fs_use_vortex: vx_fs_use_vortex, + fs_set_userdata: vx_fs_set_userdata, + fs_open: vx_fs_open, + fs_create: vx_fs_create, + fs_list: vx_fs_list, + fs_close: vx_fs_close, + fs_size: vx_fs_size, + fs_read: vx_fs_read, + fs_write: vx_fs_write, + fs_sync: vx_fs_sync, + + glob: vx_glob, + + cache_init: vx_cache_init, + cache_free: vx_cache_free, + cache_get: vx_cache_get, + cache_put: vx_cache_put, + cache_delete: vx_cache_delete, +} + +unsafe fn data_source_new( + session: *const vx_session, + opts: *const vx_data_source_options, +) -> VortexResult<*const vx_data_source> { + if session.is_null() { + vortex_bail!("empty session"); + } + let session = vx_session::as_ref(session).clone(); + + if opts.is_null() { + vortex_bail!("empty opts"); + } + let opts = unsafe { &*opts }; + + if opts.files.is_null() { + vortex_bail!("empty opts.files"); + } + let glob = unsafe { to_string(opts.files) }; + + RUNTIME.block_on(async { + let data_source = MultiFileDataSource::new(session) + //.with_filesystem(fs) + .with_glob(glob) + .build() + .await?; + Ok(vx_data_source::new(Arc::new(data_source) as DataSourceRef)) + }) +} + +/// Create a new owned datasource which must be freed by the caller +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_new( + session: *const vx_session, + opts: *const vx_data_source_options, + err: vx_error_out, +) -> *const vx_data_source { + try_or_default(err, || unsafe { data_source_new(session, opts) }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { + vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) +} + +#[repr(C)] +enum vx_cardinality { + VX_CARD_UNKNOWN = 0, + VX_CARD_ESTIMATE = 1, + VX_CARD_MAXIMUM = 2, +} + +#[repr(C)] +pub struct vx_data_source_row_count { + cardinality: vx_cardinality, + rows: u64, +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_get_row_count( + ds: *const vx_data_source, + rc: *mut vx_data_source_row_count, +) { + let rc = unsafe { &mut *rc }; + match vx_data_source::as_ref(ds).row_count() { + Some(Exact(rows)) => { + rc.cardinality = vx_cardinality::VX_CARD_MAXIMUM; + rc.rows = rows; + } + Some(Inexact(rows)) => { + rc.cardinality = vx_cardinality::VX_CARD_ESTIMATE; + rc.rows = rows; + } + None => { + rc.cardinality = vx_cardinality::VX_CARD_UNKNOWN; + } + } +} diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index 18fc722d295..98a52f59444 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -328,9 +328,9 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *cons #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_type_to_arrow_schema( - dtype: *const DType, - out: *mut FFI_ArrowSchema, - err: vx_error_out, + _dtype: *const vx_dtype, + _schema: *mut FFI_ArrowSchema, + _err: vx_error_out, ) { todo!(); } diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs new file mode 100644 index 00000000000..142ca752ae7 --- /dev/null +++ b/vortex-ffi/src/expression.rs @@ -0,0 +1,27 @@ +use std::ffi::c_char; +use std::ffi::c_int; + +use vortex::expr::Expression; +use vortex::expr::root; +use vortex::expr::select; + +use crate::to_string_vec; + +crate::box_wrapper!(Expression, vx_expression); + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_expression_root() -> *const vx_expression { + vx_expression::new(Box::new(root())) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_expression_select( + names: *const *const c_char, + names_len: usize, + child: *const vx_expression, +) -> *const vx_expression { + // TODO don't allocate, convert to [&str] + let names = unsafe { to_string_vec(names, names_len as c_int) }; + let expr = select(names, vx_expression::as_ref(child).clone()); + vx_expression::new(Box::new(expr)) +} diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 30764572c6a..ef6ccf93cb3 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -9,8 +9,10 @@ mod array; mod array_iterator; mod binary; +mod data_source; mod dtype; mod error; +mod expression; mod file; mod log; mod macros; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 96ead24d7fd..8b0a5a4fa36 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -1,260 +1,57 @@ #![allow(non_camel_case_types)] use core::slice; -use std::ffi::c_char; use std::ffi::c_int; -use std::ffi::c_void; use std::ops::Range; -use std::option::Option; use std::ptr; -use std::sync::Arc; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use vortex::buffer::Buffer; use vortex::error::VortexResult; use vortex::error::vortex_bail; -use vortex::expr::Expression; -use vortex::expr::stats::Precision::Exact; -use vortex::expr::stats::Precision::Inexact; -use vortex::file::multi::MultiFileDataSource; -use vortex::io::runtime::BlockingRuntime; use vortex::scan::Selection; -use vortex::scan::api::DataSource; -use vortex::scan::api::DataSourceRef; use vortex::scan::api::DataSourceScan; use vortex::scan::api::Partition; use vortex::scan::api::ScanRequest; -use crate::RUNTIME; use crate::array::vx_array; -use crate::dtype::vx_dtype; +use crate::data_source::vx_data_source; use crate::error::try_or_default; use crate::error::vx_error_out; -use crate::session::vx_session; -use crate::to_string; - -crate::arc_dyn_wrapper!( - /// A data source is a reference to multiple possibly remote files. When - /// created, it opens first file to determine the schema from DType, all - /// other operations are deferred till a scan is requested. You can request - /// multiple file scans from a data source - dyn DataSource, - vx_data_source); - -pub struct FileHandle; -pub type vx_file_handle = *const FileHandle; +use crate::expression::vx_expression; crate::arc_dyn_wrapper!(dyn Partition, vx_partition); -crate::box_wrapper!(Expression, vx_expression); - // TODO should be akin to DataSourceScanRef crate::arc_dyn_wrapper!(dyn DataSourceScan, vx_scan); pub struct VxArrayIt; crate::box_wrapper!(VxArrayIt, vx_array_it); -pub type vx_list_callback = - Option; -pub type vx_glob_callback = - Option; - -pub type vx_fs_use_vortex = - Option c_int>; -pub type vx_fs_set_userdata = Option; - -pub type vx_fs_open = - Option; -pub type vx_fs_create = - Option; - -pub type vx_fs_list = Option< - unsafe extern "C" fn( - userdata: *const c_void, - path: *const c_char, - callback: vx_list_callback, - error: vx_error_out, - ), ->; - -pub type vx_fs_close = Option; -pub type vx_fs_size = - Option u64>; - -pub type vx_fs_read = Option< - unsafe extern "C" fn( - handle: vx_file_handle, - offset: u64, - len: usize, - buffer: *mut u8, - err: vx_error_out, - ) -> u64, ->; - -pub type vx_fs_write = Option< - unsafe extern "C" fn( - handle: vx_file_handle, - offset: u64, - len: usize, - buffer: *mut u8, - err: vx_error_out, - ) -> u64, ->; - -pub type vx_fs_sync = Option; - -pub type vx_glob = Option< - unsafe extern "C" fn(glob: *const c_char, callback: vx_glob_callback, err: vx_error_out), ->; - -pub type vx_cache = *mut c_void; -pub type vx_cache_key = *const c_char; - -pub type vx_cache_init = Option vx_cache>; -pub type vx_cache_free = Option; -pub type vx_cache_get = Option< - unsafe extern "C" fn( - cache: vx_cache, - key: vx_cache_key, - value: *mut *mut c_void, - err: vx_error_out, - ), ->; -pub type vx_cache_put = Option< - unsafe extern "C" fn(cache: vx_cache, key: vx_cache_key, value: *mut c_void, err: vx_error_out), ->; -pub type vx_cache_delete = - Option; - #[repr(C)] -pub struct vx_data_source_options { - // TODO what if the program wants to read a Vortex file from an existing buffer? - files: *const c_char, - - /// Whether to use Vortex filesystem or host's filesystem. - /// Return 1 to use Vortex for a given schema ("file", "s3") and path. - /// Return 0 to use host's filesystem. If 0 is returned, host must also - /// Implement all fs_* callbacks. - fs_use_vortex: vx_fs_use_vortex, - fs_set_userdata: vx_fs_set_userdata, - fs_open: vx_fs_open, - fs_create: vx_fs_create, - fs_list: vx_fs_list, - fs_close: vx_fs_close, - fs_size: vx_fs_size, - fs_read: vx_fs_read, - fs_write: vx_fs_write, - fs_sync: vx_fs_sync, - - glob: vx_glob, - - cache_init: vx_cache_init, - cache_free: vx_cache_free, - cache_get: vx_cache_get, - cache_put: vx_cache_put, - cache_delete: vx_cache_delete, -} - -unsafe fn data_source_new( - session: *const vx_session, - opts: *const vx_data_source_options, -) -> VortexResult<*const vx_data_source> { - if session.is_null() { - vortex_bail!("empty session"); - } - let session = vx_session::as_ref(session).clone(); - - if opts.is_null() { - vortex_bail!("empty opts"); - } - let opts = unsafe { &*opts }; - - if opts.files.is_null() { - vortex_bail!("empty opts.files"); - } - let glob = unsafe { to_string(opts.files) }; - - RUNTIME.block_on(async { - let builder = MultiFileDataSource::new(session) - .with_filesystem(fs) - .with_glob(glob); - let data_source = builder.build().await?; - Ok(vx_data_source::new(Arc::new(data_source) as DataSourceRef)) - }) -} - -/// Create a new owned datasource which must be freed by the caller -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_data_source_new( - session: *const vx_session, - opts: *const vx_data_source_options, - err: vx_error_out, -) -> *const vx_data_source { - try_or_default(err, || unsafe { data_source_new(session, opts) }) -} - -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { - vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) -} - -#[repr(C)] -enum VxCardinality { - VX_CARD_UNKNOWN = 0, - VX_CARD_ESTIMATE = 1, - VX_CARD_MAXIMUM = 2, -} - -#[repr(C)] -pub struct vx_data_source_row_count { - cardinality: VxCardinality, - rows: u64, -} - -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_data_source_get_row_count( - ds: *const vx_data_source, - rc: *mut vx_data_source_row_count, -) { - match vx_data_source::as_ref(ds).row_count() { - Some(Exact(rows)) => unsafe { - (*rc).cardinality = VxCardinality::VX_CARD_MAXIMUM; - (*rc).rows = rows; - }, - Some(Inexact(rows)) => unsafe { - (*rc).cardinality = VxCardinality::VX_CARD_ESTIMATE; - (*rc).rows = rows; - }, - None => unsafe { - (*rc).cardinality = VxCardinality::VX_CARD_UNKNOWN; - }, - } -} - -#[repr(C)] -pub enum ScanSelection { +pub enum vx_scan_selection_include { VX_S_INCLUDE_ALL = 0, VX_S_INCLUDE_RANGE = 1, VX_S_EXCLUDE_RANGE = 2, } #[repr(C)] -struct vx_scan_selection { - selection: ScanSelection, - idx: *mut u64, - idx_len: usize, +pub struct vx_scan_selection { + pub include: vx_scan_selection_include, + pub idx: *mut u64, + pub idx_len: usize, } // Distinct from ScanRequest for easier option handling from C #[repr(C)] pub struct vx_scan_options { - projection: *const vx_expression, - filter: *const vx_expression, - row_range_begin: u64, - row_range_end: u64, - selection: vx_scan_selection, - ordered: c_int, - limit: u64, + pub projection: *const vx_expression, + pub filter: *const vx_expression, + pub row_range_begin: u64, + pub row_range_end: u64, + pub selection: vx_scan_selection, + pub ordered: c_int, + pub limit: u64, } #[repr(C)] @@ -270,77 +67,70 @@ pub struct vx_partition_estimate { estimate: u64, } -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_data_source_scan( - ds: *const vx_data_source, - opts: *const vx_scan_options, - est: *mut vx_partition_estimate, - err: vx_error_out, -) -> *const vx_scan { - let request = if opts.is_null() { - ScanRequest::default() +fn scan_request(opts: *const vx_scan_options) -> VortexResult { + if opts.is_null() { + return Ok(ScanRequest::default()); + } + let opts = unsafe { &*opts }; + + let projection = if opts.projection.is_null() { + vortex_bail!("empty opts.projection"); } else { - let opts = unsafe { &*opts }; + vx_expression::as_ref(opts.projection).clone() + }; - let projection = if opts.projection.is_null() { - panic!() // TODO - } else { - vx_expression::as_ref(opts.projection).clone() - }; + let filter = if opts.filter.is_null() { + None + } else { + Some(vx_expression::as_ref(opts.filter).clone()) + }; - let filter = if opts.filter.is_null() { - None - } else { - Some(vx_expression::as_ref(opts.filter).clone()) - }; + let selection = &opts.selection; + let selection = match selection.include { + vx_scan_selection_include::VX_S_INCLUDE_ALL => Selection::All, + vx_scan_selection_include::VX_S_INCLUDE_RANGE => { + let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; + let buf = Buffer::copy_from(buf); + Selection::IncludeByIndex(buf) + } + vx_scan_selection_include::VX_S_EXCLUDE_RANGE => { + let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; + let buf = Buffer::copy_from(buf); + Selection::ExcludeByIndex(buf) + } + }; - let selection = &opts.selection; - let selection: Selection = match selection.selection { - ScanSelection::VX_S_INCLUDE_ALL => Selection::All, - ScanSelection::VX_S_INCLUDE_RANGE => { - let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; - let buf = Buffer::copy_from(buf); - Selection::IncludeByIndex(buf) - } - ScanSelection::VX_S_EXCLUDE_RANGE => { - let buf = unsafe { slice::from_raw_parts(selection.idx, selection.idx_len) }; - let buf = Buffer::copy_from(buf); - Selection::ExcludeByIndex(buf) - } - }; + let ordered = opts.ordered == 1; - let ordered = opts.ordered == 1; + let start = opts.row_range_begin; + let end = opts.row_range_end; + let row_range = (start > 0 || end > 0).then_some(Range { start, end }); - let row_range: Option> = if opts.row_range_begin == 0 && opts.row_range_end == 0 - { - None - } else { - Some(Range { - start: opts.row_range_begin, - end: opts.row_range_end, - }) - }; + let limit = (opts.limit == 0).then_some(opts.limit); - let limit = if opts.limit == 0 { - None - } else { - Some(opts.limit) - }; + Ok(ScanRequest { + projection, + filter, + row_range, + selection, + ordered, + limit, + }) +} - ScanRequest { - projection, - filter, - row_range, - selection, - ordered, - limit, - } - }; +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_scan( + _data_source: *const vx_data_source, + _options: *const vx_scan_options, + _partition_estimate: *mut vx_partition_estimate, + _err: vx_error_out, +) -> *const vx_scan { //try_or_default(err, || { - let scan = vx_data_source::as_ref(ds).scan(request); - //Ok(vx_scan::new(scan)) + // let request = scan_request(opts)?; + // let scan = vx_data_source::as_ref(ds).scan(request); + // Ok(vx_scan::new(Arc::new(scan))) //}) - ptr::null_mut() + ptr::null() } #[unsafe(no_mangle)] @@ -380,6 +170,3 @@ pub unsafe extern "C-unwind" fn vx_partition_next( pub unsafe extern "C-unwind" fn vx_scan_progress(_scan: *const vx_scan) -> f64 { 0.0 } - -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_scan_cancel(_scan: *const vx_scan) {} From 08e408b5009aead798072ade47e15b411f10d3d5 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 13:11:16 +0000 Subject: [PATCH 11/21] better --- .../cpp/include/duckdb_vx/table_function.h | 24 +- vortex-duckdb/cpp/table_function.cpp | 231 +++++++++--------- 2 files changed, 131 insertions(+), 124 deletions(-) diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h index 258a0c47fae..9a3d566c045 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h @@ -15,9 +15,17 @@ #include "duckdb_vx/client_context.h" #include "duckdb_vx/duckdb_diagnostics.h" -DUCKDB_INCLUDES_BEGIN -#include "duckdb/common/arrow/arrow.hpp" -DUCKDB_INCLUDES_END +#ifdef __cplusplus + DUCKDB_INCLUDES_BEGIN + #include "duckdb/common/arrow/arrow.hpp" + DUCKDB_INCLUDES_END + + using FFI_ArrowSchema = ArrowSchema; + using FFI_ArrowArrayStream = ArrowArrayStream; +#else + typedef void FFI_ArrowSchema; + typedef void FFI_ArrowArrayStream; +#endif #include "vortex.h" @@ -178,13 +186,9 @@ typedef struct { idx_t max_threads; // PoC hack: retain Rust exporter code - void (*export_array)( - const vx_array* arr, - duckdb_data_chunk chunk, - duckdb_vx_error* err, - // needed to initialize rust-side exporter - void* local_state - ); + // return local batch id + uint64_t (*export_array)(const vx_array* arr, duckdb_data_chunk chunk); + } duckdb_vx_tfunc_vtab_t; // A single function for configuring the DuckDB table function vtable. diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index f86954446dd..510966c8fd7 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -2,9 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "duckdb_vx.h" -#include "duckdb_vx/data.hpp" #include "duckdb_vx/error.hpp" -#include "duckdb_vx/duckdb_diagnostics.h" +#include "duckdb_vx/table_function.h" DUCKDB_INCLUDES_BEGIN #include "duckdb.h" @@ -19,7 +18,7 @@ DUCKDB_INCLUDES_END using namespace duckdb; std::string to_string(const vx_string* str) { - return { vx_string_ptr(str), size_t vx_string_len(str) }; + return { vx_string_ptr(str), vx_string_len(str) }; } std::string move_vx_err(vx_error* error) { @@ -29,24 +28,25 @@ std::string move_vx_err(vx_error* error) { return str; } +LogicalType from_dtype(const vx_dtype* dtype) { + return LogicalType::UNKNOWN; +} + namespace vortex { struct CTableFunctionInfo final : TableFunctionInfo { - explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) - : vtab(vtab), max_threads(vtab.max_threads) { } + explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) : vtab(vtab) {} duckdb_vx_tfunc_vtab_t vtab; - idx_t max_threads; // derived from vtab.max_threads + // vtab.max_threads }; struct CTableBindData final : TableFunctionData { CTableBindData( unique_ptr info_p, - unique_ptr ffi_data_p, const vector& column_names, - const vx_session* session, + vx_session* session, const vx_data_source* data_source) : info(std::move(info_p)) - , ffi_data(std::move(ffi_data_p)), , column_names(column_names) , session(session) , data_source(data_source) {} @@ -57,26 +57,21 @@ struct CTableBindData final : TableFunctionData { } unique_ptr Copy() const override { - assert(info->vtab.bind_data_clone != nullptr); - - 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)); - } - return make_uniq(make_uniq(info->vtab), - unique_ptr(reinterpret_cast(copied_ffi_data))); + return make_uniq( + make_uniq(info->vtab), + column_names, + // TODO what's with this ownership? + session, + // TODO arcd inside, so fine, don't need to clone it? + data_source); } - // TODO remove - unique_ptr info; - unique_ptr ffi_data; - // ^ In Rust it contains filter expressions + unique_ptr info; // TODO remove - // pre-computed from data_source using vx_struct_fields_field_name, see c_bind() + // precomputed from data_source using vx_struct_fields_field_name, see c_bind() vector column_names; - const vx_session* session; + vx_session* session; const vx_data_source* data_source; }; @@ -99,25 +94,41 @@ struct CTableLocalData final : LocalTableFunctionState { std::optional batch_id; }; -LogicalType from_dtype(const vx_dtype* dtype) { - return LogicalType::NULL; -} - unique_ptr c_bind( ClientContext &context, TableFunctionBindInput &info, vector &types, vector &names) { - if (info->inputs.empty()) + if (info.inputs.empty()) throw BinderException("missing file glob parameter"); - const vx_error* err = nullptr; - vx_data_source_options opts = {0}; - opts.files = input->inputs[0]; + vx_error* err = nullptr; + // TODO check if it is string + std::string glob = info.inputs[0].GetValue(); + + vx_data_source_options opts = { + .files = glob.data(), + .fs_use_vortex = nullptr, + .fs_set_userdata = nullptr, + .fs_open = nullptr, + .fs_create = nullptr, + .fs_list = nullptr, + .fs_close = nullptr, + .fs_size = nullptr, + .fs_read = nullptr, + .fs_write = nullptr, + .fs_sync = nullptr, + .glob = nullptr, + .cache_init = nullptr, + .cache_free = nullptr, + .cache_get = nullptr, + .cache_put = nullptr, + .cache_delete = nullptr + }; - const vx_session* session = vx_session_new(); - const vx_data_source* data_source = vx_data_source_new(session, opts, &err); + vx_session* session = vx_session_new(); + const vx_data_source* data_source = vx_data_source_new(session, &opts, &err); const vx_dtype* dtype = vx_data_source_dtype(data_source); const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); @@ -133,9 +144,9 @@ unique_ptr c_bind( types.emplace_back(from_dtype(field_dtype)); } + auto& vtab = info.table_function.function_info->Cast().vtab; return make_uniq( - make_uniq(info.vtab), - unique_ptr(reinterpret_cast(ffi_bind_data)), + make_uniq(vtab), names, session, data_source); @@ -146,31 +157,39 @@ const vx_expression* projection( const vector& projection_ids, const vector& column_names ) { - vector projected_column_names; - projected_column_names.reserve(column_names); + vector projected_names; + projected_names.reserve(column_names.size()); if (projection_ids.empty()) { for (column_t id : column_ids) { if (column_names.size() < id) - throw InvalidInputException(); - string& name = column_names[id]; - projected_column_names.emplace_back(name.data()); + throw InvalidInputException(StringUtil::Format( + "Expected column id %d but there are %d columns", + id, column_names.size())); + // column_names[id] lives till end of projection(). Initialized + // vx_expression copies the buffer, so it's safe to use data() + projected_names.emplace_back(column_names[id].data()); } } else { for (idx_t projection_id : projection_ids) { - column_it id = column_ids[projection_id]; + if (column_ids.size() < projection_id) + throw InvalidInputException(StringUtil::Format( + "Expected projection id %d but there are %d columns", + projection_id, column_ids.size())); + column_t id = column_ids[projection_id]; if (column_names.size() < id) - throw InvalidInputException(); - string& name = column_names[id]; - projected_column_names.emplace_back(name.data()); + throw InvalidInputException(StringUtil::Format( + "Expected column id %d but there are %d columns", + id, column_names.size())); + projected_names.emplace_back(column_names[id].data()); } } - const vx_expression* root = vx_expression_root(); + vx_expression* root = vx_expression_root(); // TODO vx_expression may take a string vx_array const vx_expression* expr = vx_expression_select( - projected_column_names.data(), - projected_column_names.size(), + projected_names.data(), + projected_names.size(), root); vx_expression_free(root); return expr; @@ -179,65 +198,60 @@ const vx_expression* projection( unique_ptr c_init_global(ClientContext &, TableFunctionInitInput &input) { const auto &bind = input.bind_data->Cast(); - vx_scan_options options = {0}; - options.projection = projection(input.column_ids, input.projection_ids, bind.names); - option.filter = filter(); + vx_scan_selection selection = { + .include = VX_S_INCLUDE_ALL, + .idx = nullptr, + .idx_len = 0 + }; + vx_scan_options options = { + .projection = projection(input.column_ids, input.projection_ids, bind.column_names), + .filter = nullptr, + .row_range_begin = 0, + .row_range_end = 0, + .selection = selection, + .ordered = 0, + .limit = 0 + }; vx_error* err = nullptr; const vx_scan* scan = vx_data_source_scan( bind.data_source, - options, + &options, /* partition estimate */ nullptr, &err); - if (err) throw InvalidInputException(move_vx_err(err)); + if (err) + throw InvalidInputException(move_vx_err(err)); - return make_uniq(scan, bind.info->max_threads); + return make_uniq(scan, bind.info->vtab.max_threads); } -unique_ptr c_init_local(ExecutionContext &context, - TableFunctionInitInput &input, - GlobalTableFunctionState *global_state) { - const auto &bind = input.bind_data->Cast(); - auto global_data = global_state->Cast().ffi_data->DataPtr(); - - duckdb_vx_tfunc_init_input ffi_input = { - .bind_data = bind.ffi_data->DataPtr(), - .column_ids = input.column_ids.data(), - .column_ids_count = input.column_ids.size(), - .projection_ids = input.projection_ids.data(), - .projection_ids_count = input.projection_ids.size(), - .filters = reinterpret_cast(input.filters.get()), - .client_context = reinterpret_cast(&context), - }; - - duckdb_vx_error error_out = nullptr; - auto ffi_local_data = bind.info->vtab.init_local(&ffi_input, global_data, &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } - - return make_uniq( - unique_ptr(reinterpret_cast(ffi_local_data))); +unique_ptr c_init_local(ExecutionContext &, + TableFunctionInitInput &, + GlobalTableFunctionState*) { + return make_uniq(); } void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { const vx_scan* const scan = input.global_state->Cast().scan; - vx_error* vx_err = nullptr; - const vx_array* arr = vx_scan_next_arr(scan, &vx_err); - if (vx_err) { - throw InvalidInputException(move_vx_err(vx_err)); - } - if (!arr) { // everything exported - return; - } - - duckdb_vx_error err = nullptr; - void* local_data = input.local_state->Cast().ffi_data->DataPtr(); - input.bind_data->Cast().info->vtab.export_array( - arr, reinterpret_cast(&output), &err, local_data); - if (err) { - throw InvalidInputException(IntoErrString(err)); + vx_error* err = nullptr; + const vx_partition* partition = nullptr; + const vx_array* array = nullptr; + + auto export_array = input.bind_data->Cast().info->vtab.export_array; + auto& batch_id = input.local_state->Cast().batch_id; + + // TODO this is essentially single-threaded processing + while (err == nullptr && (partition = vx_scan_next(scan, &err))) { + while ((array = vx_partition_next(partition, &err))) { + *batch_id = export_array( + array, + reinterpret_cast(&output)); + vx_array_free(array); + } + vx_partition_free(partition); } + if (err) + throw InvalidInputException(move_vx_err(err)); } //void c_pushdown_complex_filter(ClientContext & /*context*/, @@ -283,7 +297,7 @@ unique_ptr c_cardinality(ClientContext &, const FunctionData *bi stats->max_cardinality = rc.rows; return stats; } - case VX_CARD_UNKNOWN: { + default: { stats->has_estimated_cardinality = false; stats->has_max_cardinality = false; return stats; @@ -291,20 +305,6 @@ unique_ptr c_cardinality(ClientContext &, const FunctionData *bi } } -extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_named_parameter(duckdb_vx_tfunc_bind_input ffi_input, - const char *name) { - if (!ffi_input || !name) { - return nullptr; - } - const auto info = reinterpret_cast(ffi_input); - - if (!info->named_parameters.contains(name)) { - return nullptr; - } - auto value = duckdb::make_uniq(info->named_parameters.at(name)); - return reinterpret_cast(value.release()); -} - OperatorPartitionData c_get_partition_data(ClientContext&, TableFunctionGetPartitionInput &input) { if (input.partition_info.RequiresPartitionColumns()) @@ -336,7 +336,7 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( tf.projection_pushdown = true; tf.filter_pushdown = false; - tf.filter_prune = false + tf.filter_prune = false; tf.sampling_pushdown = false; tf.late_materialization = false; @@ -345,16 +345,19 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( tf.to_string = [](TableFunctionToStringInput&) { InsertionOrderPreservingMap result; - result.emplace("Function", "Vortex Scan"); + result.insert("Function", "Vortex Scan"); // TODO filters return result; - } + }; tf.get_virtual_columns = [](auto&, auto) { - return {{COLUMN_IDENTIFIER_EMPTY, TableColumn{"", LogicalType::BOOL}}}; - } + virtual_column_map_t map = { + {COLUMN_IDENTIFIER_EMPTY, TableColumn{"", LogicalType::BOOLEAN}} + }; + return map; + }; - tf.table_scan_progress = [](auto&, auto*, const GlobalTableFunctionState& state) { + tf.table_scan_progress = [](auto&, auto*, const GlobalTableFunctionState* state) { return vx_scan_progress(state->Cast().scan); }; From e97c1190c2f1ee5682ef7288dcec8e489bb89598 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 13:50:41 +0000 Subject: [PATCH 12/21] better --- vortex-duckdb/build.rs | 1 + .../cpp/include/duckdb_vx/table_function.h | 1 + vortex-duckdb/cpp/table_function.cpp | 6 +- vortex-duckdb/include/vortex.h | 2 + vortex-duckdb/src/datasource.rs | 4 + .../src/duckdb/table_function/export_array.rs | 82 +++++++++++++++++++ .../src/duckdb/table_function/mod.rs | 3 + vortex-ffi/cbindgen.toml | 3 - vortex-ffi/cinclude/vortex.h | 82 +++++++++---------- vortex-ffi/src/data_source.rs | 4 +- vortex-ffi/src/expression.rs | 4 +- vortex-ffi/src/scan.rs | 22 +++-- 12 files changed, 156 insertions(+), 58 deletions(-) create mode 100644 vortex-duckdb/src/duckdb/table_function/export_array.rs diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index 90f264a119a..660e4b822e3 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -304,6 +304,7 @@ fn c2rust(crate_dir: &Path, duckdb_include_dir: &Path) { .rustified_non_exhaustive_enum("DUCKDB_TYPE") .size_t_is_usize(true) .clang_arg(format!("-I{}", duckdb_include_dir.display())) + .clang_arg(format!("-I{}", crate_dir.join(FFI_INCLUDE).display())) .clang_arg(format!("-I{}", crate_dir.join("cpp/include").display())) .generate_comments(true) // Tell cargo to invalidate the built crate whenever any of the diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h index 9a3d566c045..1fa7759f5df 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h @@ -23,6 +23,7 @@ using FFI_ArrowSchema = ArrowSchema; using FFI_ArrowArrayStream = ArrowArrayStream; #else + // TODO nanoarrow typedef void FFI_ArrowSchema; typedef void FFI_ArrowArrayStream; #endif diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 510966c8fd7..8676b71b062 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -243,10 +243,14 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou // TODO this is essentially single-threaded processing while (err == nullptr && (partition = vx_scan_next(scan, &err))) { while ((array = vx_partition_next(partition, &err))) { - *batch_id = export_array( + uint64_t export_res = export_array( array, reinterpret_cast(&output)); vx_array_free(array); + if (export_res == std::numeric_limits::max()) + batch_id = std::nullopt; + else + batch_id = export_res; } vx_partition_free(partition); } diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index 7148c4d9799..72341cc523d 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -38,6 +38,8 @@ const char *vortex_version_rust(void); */ const char *vortex_extension_version_rust(void); +uint64_t export_array_callback(const vx_array *array, duckdb_data_chunk chunk); + #ifdef __cplusplus } #endif diff --git a/vortex-duckdb/src/datasource.rs b/vortex-duckdb/src/datasource.rs index ea7520af591..dcdc43e11a8 100644 --- a/vortex-duckdb/src/datasource.rs +++ b/vortex-duckdb/src/datasource.rs @@ -320,6 +320,10 @@ impl TableFunction for T { global_state: &Self::GlobalState, chunk: &mut DataChunkRef, ) -> VortexResult<()> { + // TODO create exporter for every chunk, use global + // variable to keep track of exporter till it fully + // consumes generated array + loop { if local_state.exporter.is_none() { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-duckdb/src/duckdb/table_function/export_array.rs b/vortex-duckdb/src/duckdb/table_function/export_array.rs new file mode 100644 index 00000000000..3cc448ba2b5 --- /dev/null +++ b/vortex-duckdb/src/duckdb/table_function/export_array.rs @@ -0,0 +1,82 @@ +use vortex::{error::VortexExpect, scalar_fn::fns::pack::Pack}; +use vortex::array::Canonical; +use vortex::array::arrays::ScalarFnVTable; +use vortex::array::arrays::Struct; +use vortex::array::arrays::StructArray; +use vortex::array::optimizer::ArrayOptimizer; +use vortex::array::VortexSessionExecute; +use crate::{cpp::{duckdb_data_chunk, vx_array}, duckdb::TableFunction, exporter::{ArrayExporter, ConversionCache}, SESSION}; +use crate::duckdb::DataChunkRef; + +thread_local! { + +static EXPORTER: Option = None; +static CONVERSION_CACHE: ConversionCache = ConversionCache::default(); +} + +// TODO create exporter for every chunk, use global +// variable to keep track of exporter till it fully +// consumes generated array. Use global, not per-partition, +// conversion cache +#[unsafe(no_mangle)] +pub(crate) unsafe extern "C-unwind" fn export_array_callback( + array: *const vx_array, + chunk: duckdb_data_chunk +) -> u64 { + let chunk: &mut DataChunkRef = unsafe { duckdb_data_chunk::as_mut(chunk) }; + let mut batch_id = u64::MAX; + if array.is_null() { + return batch_id; + } + let array_result = vx_array::as_ref(array); + + if EXPORTER.is_none() { + let mut ctx = SESSION.create_execution_ctx(); + let array_result = array_result.optimize_recursive() + .vortex_expect("failed to optimize array"); + + let array_result = if let Some(array) = array_result.as_opt::() { + array.clone() + } else if let Some(array) = array_result.as_opt::() + && let Some(pack_options) = array.scalar_fn().as_opt::() + { + StructArray::new( + pack_options.names.clone(), + array.children(), + array.len(), + pack_options.nullability.into(), + ) + } else { + array_result.execute::(&mut ctx) + .vortex_expect("failed to canonicalize array") + .into_struct() + }; + + EXPORTER = Some(ArrayExporter::try_new( + &array_result, + &CONVERSION_CACHE, + ctx, + ).vortex_expect("failed to initialize array exporter")); + // Relaxed since there is no intra-instruction ordering required. + //batch_id = Some(global_state.batch_id.fetch_add(1, Ordering::Relaxed)); + batch_id += 1; + } + + let exporter = EXPORTER + .as_mut() + .vortex_expect("error: exporter missing"); + + let has_more_data = exporter.export(chunk)?; + //global_state + // .bytes_read + // .fetch_add(chunk.len(), Ordering::Relaxed); + + if !has_more_data { + // This exporter is fully consumed. + EXPORTER = None; + batch_id = u64::MAX; + } + + assert!(!chunk.is_empty()); + batch_id +} diff --git a/vortex-duckdb/src/duckdb/table_function/mod.rs b/vortex-duckdb/src/duckdb/table_function/mod.rs index c9fa08177b3..0accb557adc 100644 --- a/vortex-duckdb/src/duckdb/table_function/mod.rs +++ b/vortex-duckdb/src/duckdb/table_function/mod.rs @@ -13,6 +13,7 @@ mod bind; mod cardinality; mod init; mod partition; +mod export_array; mod pushdown_complex_filter; mod virtual_columns; @@ -30,6 +31,7 @@ use crate::duckdb::LogicalType; use crate::duckdb::client_context::ClientContextRef; use crate::duckdb::data_chunk::DataChunkRef; use crate::duckdb::expr::ExpressionRef; +use crate::duckdb::table_function::export_array::export_array_callback; use crate::duckdb::table_function::cardinality::cardinality_callback; use crate::duckdb::table_function::partition::get_partition_data_callback; use crate::duckdb::table_function::pushdown_complex_filter::pushdown_complex_filter_callback; @@ -199,6 +201,7 @@ impl DatabaseRef { sampling_pushdown: false, late_materialization: false, max_threads: T::MAX_THREADS, + export_array: Some(export_array_callback::), }; duckdb_try!( diff --git a/vortex-ffi/cbindgen.toml b/vortex-ffi/cbindgen.toml index 65e175ba995..12e484f093e 100644 --- a/vortex-ffi/cbindgen.toml +++ b/vortex-ffi/cbindgen.toml @@ -12,9 +12,6 @@ header = """ // // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY // - -typedef ArrowSchema FFI_ArrowSchema; -typedef ArrowArrayStream FFI_ArrowArrayStream; """ [export.rename] diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index f3b7a34609c..29327e7b8b1 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -5,9 +5,6 @@ // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY // -typedef ArrowSchema FFI_ArrowSchema; -typedef ArrowArrayStream FFI_ArrowArrayStream; - #pragma once #include @@ -249,8 +246,6 @@ typedef uint8_t PType; */ typedef struct DType DType; -typedef struct FileHandle FileHandle; - /** * Whether an instance of a DType can be `null or not */ @@ -258,6 +253,8 @@ typedef struct Nullability Nullability; typedef struct Primitive Primitive; +typedef struct VxFileHandle VxFileHandle; + /** * Base type for all Vortex arrays. * @@ -378,7 +375,7 @@ typedef void (*vx_fs_list)(const void *userdata, vx_list_callback callback, vx_error_out error); -typedef const FileHandle *vx_file_handle; +typedef const VxFileHandle *vx_file_handle; typedef void (*vx_fs_close)(vx_file_handle handle); @@ -620,29 +617,6 @@ const vx_string *vx_array_get_utf8(const vx_array *array, uint32_t index); */ const vx_binary *vx_array_get_binary(const vx_array *array, uint32_t index); -/** - * Clone a borrowed [`vx_data_source`], returning an owned [`vx_data_source`]. - * - * - * Must be released with [`vx_data_source_free`]. - */ -const vx_data_source *vx_data_source_clone(const vx_data_source *ptr); - -/** - * Free an owned [`vx_data_source`] object. - */ -void vx_data_source_free(const vx_data_source *ptr); - -/** - * Create a new owned datasource which must be freed by the caller - */ -const vx_data_source * -vx_data_source_new(const vx_session *session, const vx_data_source_options *opts, vx_error_out err); - -const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); - -void vx_data_source_get_row_count(const vx_data_source *ds, vx_data_source_row_count *rc); - /** * Free an owned [`vx_array_iterator`] object. */ @@ -686,6 +660,29 @@ size_t vx_binary_len(const vx_binary *ptr); */ const char *vx_binary_ptr(const vx_binary *ptr); +/** + * Clone a borrowed [`vx_data_source`], returning an owned [`vx_data_source`]. + * + * + * Must be released with [`vx_data_source_free`]. + */ +const vx_data_source *vx_data_source_clone(const vx_data_source *ptr); + +/** + * Free an owned [`vx_data_source`] object. + */ +void vx_data_source_free(const vx_data_source *ptr); + +/** + * Create a new owned datasource which must be freed by the caller + */ +const vx_data_source * +vx_data_source_new(const vx_session *session, const vx_data_source_options *opts, vx_error_out err); + +const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); + +void vx_data_source_get_row_count(const vx_data_source *ds, vx_data_source_row_count *rc); + /** * Clone a borrowed [`vx_dtype`], returning an owned [`vx_dtype`]. * @@ -844,6 +841,15 @@ void vx_error_free(vx_error *ptr); */ const vx_string *vx_error_get_message(const vx_error *error); +/** + * Free an owned [`vx_expression`] object. + */ +void vx_expression_free(vx_expression *ptr); + +vx_expression *vx_expression_root(void); + +vx_expression *vx_expression_select(const char *const *names, size_t names_len, const vx_expression *child); + /** * Clone a borrowed [`vx_file`], returning an owned [`vx_file`]. * @@ -902,16 +908,6 @@ vx_array_iterator *vx_file_scan(const vx_session *session, */ void vx_set_log_level(vx_log_level level); -/** - * Free an owned [`vx_expression`] object. - */ -void vx_expression_free(vx_expression *ptr); - -const vx_expression *vx_expression_root(void); - -const vx_expression * -vx_expression_select(const char *const *names, size_t names_len, const vx_expression *child); - /** * Clone a borrowed [`vx_partition`], returning an owned [`vx_partition`]. * @@ -946,13 +942,13 @@ void vx_array_it_free(vx_array_it *ptr); const vx_scan *vx_data_source_scan(const vx_data_source *_data_source, const vx_scan_options *_options, vx_partition_estimate *_partition_estimate, - vx_error_out _err); + vx_error_out err); -const vx_partition *vx_scan_next(const vx_scan *scan, vx_error_out err); +const vx_partition *vx_scan_next(const vx_scan *_scan, vx_error_out err); -void vx_partition_scan_arrow(const vx_partition *part, FFI_ArrowArrayStream *out, vx_error_out err); +void vx_partition_scan_arrow(const vx_partition *_partition, FFI_ArrowArrayStream *_stream, vx_error_out err); -const vx_array *vx_partition_next(const vx_partition *partition, vx_error_out err); +const vx_array *vx_partition_next(const vx_partition *_partition, vx_error_out err); /** * Scan progress between 0.0 and 1.0 diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index 560895a5175..f2f55b40fcd 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -29,8 +29,8 @@ crate::arc_dyn_wrapper!( dyn DataSource, vx_data_source); -pub struct FileHandle; -pub type vx_file_handle = *const FileHandle; +pub struct VxFileHandle; +pub type vx_file_handle = *const VxFileHandle; pub type vx_list_callback = Option; diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 142ca752ae7..2aa93db21e8 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -10,7 +10,7 @@ use crate::to_string_vec; crate::box_wrapper!(Expression, vx_expression); #[unsafe(no_mangle)] -pub unsafe extern "C" fn vx_expression_root() -> *const vx_expression { +pub unsafe extern "C" fn vx_expression_root() -> *mut vx_expression { vx_expression::new(Box::new(root())) } @@ -19,7 +19,7 @@ pub unsafe extern "C" fn vx_expression_select( names: *const *const c_char, names_len: usize, child: *const vx_expression, -) -> *const vx_expression { +) -> *mut vx_expression { // TODO don't allocate, convert to [&str] let names = unsafe { to_string_vec(names, names_len as c_int) }; let expr = select(names, vx_expression::as_ref(child).clone()); diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 8b0a5a4fa36..b11c009fd4c 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -18,6 +18,7 @@ use crate::array::vx_array; use crate::data_source::vx_data_source; use crate::error::try_or_default; use crate::error::vx_error_out; +use crate::error::write_error; use crate::expression::vx_expression; crate::arc_dyn_wrapper!(dyn Partition, vx_partition); @@ -123,46 +124,53 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( _data_source: *const vx_data_source, _options: *const vx_scan_options, _partition_estimate: *mut vx_partition_estimate, - _err: vx_error_out, + err: vx_error_out, ) -> *const vx_scan { //try_or_default(err, || { // let request = scan_request(opts)?; // let scan = vx_data_source::as_ref(ds).scan(request); // Ok(vx_scan::new(Arc::new(scan))) //}) + + write_error(err, "failed to create scan"); ptr::null() } #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scan_next( - scan: *const vx_scan, + _scan: *const vx_scan, err: vx_error_out, ) -> *const vx_partition { - ptr::null_mut() //try_or_default(err, || { // Ok(vx_partition::new(Arc::new())) //}) + + write_error(err, "failed to get next partition"); + ptr::null() } // TODO export nanoarrow headers? #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( - part: *const vx_partition, - out: *mut FFI_ArrowArrayStream, + _partition: *const vx_partition, + _stream: *mut FFI_ArrowArrayStream, err: vx_error_out, ) { + write_error(err, "failed to scan partition to Arrow"); } #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_partition_next( - partition: *const vx_partition, + _partition: *const vx_partition, err: vx_error_out, ) -> *const vx_array { - ptr::null_mut() //try_or_default(err, || { // Ok(vx_array::new(Arc::new())) //}) + + write_error(err, "failed to get next array out of partition"); + ptr::null() } #[unsafe(no_mangle)] From fc58552c343f18c36f57310add36269d2ad10c38 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 14:28:40 +0000 Subject: [PATCH 13/21] better --- vortex-duckdb/include/vortex.h | 2 - .../src/duckdb/table_function/export_array.rs | 104 +++++++++--------- .../src/duckdb/table_function/mod.rs | 4 +- vortex-ffi/Cargo.toml | 3 +- vortex-ffi/src/lib.rs | 3 + vortex-ffi/src/macros.rs | 3 +- 6 files changed, 62 insertions(+), 57 deletions(-) diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index 72341cc523d..7148c4d9799 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -38,8 +38,6 @@ const char *vortex_version_rust(void); */ const char *vortex_extension_version_rust(void); -uint64_t export_array_callback(const vx_array *array, duckdb_data_chunk chunk); - #ifdef __cplusplus } #endif diff --git a/vortex-duckdb/src/duckdb/table_function/export_array.rs b/vortex-duckdb/src/duckdb/table_function/export_array.rs index 3cc448ba2b5..d1f5a9db73a 100644 --- a/vortex-duckdb/src/duckdb/table_function/export_array.rs +++ b/vortex-duckdb/src/duckdb/table_function/export_array.rs @@ -1,79 +1,81 @@ -use vortex::{error::VortexExpect, scalar_fn::fns::pack::Pack}; +use std::sync::Arc; + use vortex::array::Canonical; +use vortex::array::VortexSessionExecute; use vortex::array::arrays::ScalarFnVTable; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; use vortex::array::optimizer::ArrayOptimizer; -use vortex::array::VortexSessionExecute; -use crate::{cpp::{duckdb_data_chunk, vx_array}, duckdb::TableFunction, exporter::{ArrayExporter, ConversionCache}, SESSION}; -use crate::duckdb::DataChunkRef; +use vortex::error::VortexExpect; +use vortex::scalar_fn::fns::pack::Pack; +use vortex::array::DynArray; -thread_local! { +use crate::duckdb::DataChunk; +use crate::SESSION; +use crate::cpp::duckdb_data_chunk; +use crate::duckdb::TableFunction; +use crate::exporter::ArrayExporter; +use crate::exporter::ConversionCache; -static EXPORTER: Option = None; -static CONVERSION_CACHE: ConversionCache = ConversionCache::default(); -} +// TODO In the original implementation, exporter is initialized in the +// local state, and conversion cache is scoped per partition. -// TODO create exporter for every chunk, use global -// variable to keep track of exporter till it fully -// consumes generated array. Use global, not per-partition, -// conversion cache -#[unsafe(no_mangle)] pub(crate) unsafe extern "C-unwind" fn export_array_callback( - array: *const vx_array, - chunk: duckdb_data_chunk + array: *const crate::cpp::vx_array, + chunk: duckdb_data_chunk, ) -> u64 { - let chunk: &mut DataChunkRef = unsafe { duckdb_data_chunk::as_mut(chunk) }; + let chunk = unsafe { DataChunk::borrow_mut(chunk) }; let mut batch_id = u64::MAX; if array.is_null() { return batch_id; } - let array_result = vx_array::as_ref(array); + let array_result: Arc = + vortex_ffi::vx_array::as_ref(array as *const vortex_ffi::vx_array).clone(); - if EXPORTER.is_none() { - let mut ctx = SESSION.create_execution_ctx(); - let array_result = array_result.optimize_recursive() - .vortex_expect("failed to optimize array"); + // TODO this will produce incorrect results as exporter may export + // multiple data chunks. This exporter exports only one data chunk - let array_result = if let Some(array) = array_result.as_opt::() { - array.clone() - } else if let Some(array) = array_result.as_opt::() - && let Some(pack_options) = array.scalar_fn().as_opt::() - { - StructArray::new( - pack_options.names.clone(), - array.children(), - array.len(), - pack_options.nullability.into(), - ) - } else { - array_result.execute::(&mut ctx) - .vortex_expect("failed to canonicalize array") - .into_struct() - }; + let conversion_cache = ConversionCache::default(); - EXPORTER = Some(ArrayExporter::try_new( - &array_result, - &CONVERSION_CACHE, - ctx, - ).vortex_expect("failed to initialize array exporter")); - // Relaxed since there is no intra-instruction ordering required. - //batch_id = Some(global_state.batch_id.fetch_add(1, Ordering::Relaxed)); - batch_id += 1; - } + let mut ctx = SESSION.create_execution_ctx(); + let array_result = array_result + .optimize_recursive() + .vortex_expect("failed to optimize array"); + + let array_result = if let Some(array) = array_result.as_opt::() { + array.clone() + } else if let Some(array) = array_result.as_opt::() + && let Some(pack_options) = array.scalar_fn().as_opt::() + { + StructArray::new( + pack_options.names.clone(), + array.children(), + array.len(), + pack_options.nullability.into(), + ) + } else { + array_result + .execute::(&mut ctx) + .vortex_expect("failed to canonicalize array") + .into_struct() + }; + + let mut exporter = ArrayExporter::try_new(&array_result, &conversion_cache, ctx) + .vortex_expect("failed to initialize array exporter"); - let exporter = EXPORTER - .as_mut() - .vortex_expect("error: exporter missing"); + // Relaxed since there is no intra-instruction ordering required. + //batch_id = Some(global_state.batch_id.fetch_add(1, Ordering::Relaxed)); + batch_id += 1; - let has_more_data = exporter.export(chunk)?; + let has_more_data = exporter.export(chunk) + .vortex_expect("failed to export chunk"); //global_state // .bytes_read // .fetch_add(chunk.len(), Ordering::Relaxed); if !has_more_data { // This exporter is fully consumed. - EXPORTER = None; + //EXPORTER = None; batch_id = u64::MAX; } diff --git a/vortex-duckdb/src/duckdb/table_function/mod.rs b/vortex-duckdb/src/duckdb/table_function/mod.rs index 0accb557adc..ec16cfc9b97 100644 --- a/vortex-duckdb/src/duckdb/table_function/mod.rs +++ b/vortex-duckdb/src/duckdb/table_function/mod.rs @@ -11,9 +11,9 @@ use vortex::error::VortexExpect; use vortex::error::VortexResult; mod bind; mod cardinality; +mod export_array; mod init; mod partition; -mod export_array; mod pushdown_complex_filter; mod virtual_columns; @@ -31,8 +31,8 @@ use crate::duckdb::LogicalType; use crate::duckdb::client_context::ClientContextRef; use crate::duckdb::data_chunk::DataChunkRef; use crate::duckdb::expr::ExpressionRef; -use crate::duckdb::table_function::export_array::export_array_callback; use crate::duckdb::table_function::cardinality::cardinality_callback; +use crate::duckdb::table_function::export_array::export_array_callback; use crate::duckdb::table_function::partition::get_partition_data_callback; use crate::duckdb::table_function::pushdown_complex_filter::pushdown_complex_filter_callback; use crate::duckdb::table_function::virtual_columns::get_virtual_columns_callback; diff --git a/vortex-ffi/Cargo.toml b/vortex-ffi/Cargo.toml index 0a777fd2fa4..949cbb0a1d5 100644 --- a/vortex-ffi/Cargo.toml +++ b/vortex-ffi/Cargo.toml @@ -45,7 +45,8 @@ mimalloc = ["dep:mimalloc"] [lib] name = "vortex_ffi" -crate-type = ["staticlib", "cdylib"] +# TODO lib for vortex-duckdb exporter hack +crate-type = ["lib", "staticlib", "cdylib"] [lints] workspace = true diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index ef6ccf93cb3..343473b7d5b 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -28,6 +28,9 @@ use std::ffi::c_char; use std::ffi::c_int; use std::sync::LazyLock; +// TODO hack for duckdb exporter +pub use array::vx_array; + pub use log::vx_log_level; use vortex::io::runtime::current::CurrentThreadRuntime; diff --git a/vortex-ffi/src/macros.rs b/vortex-ffi/src/macros.rs index a0a0f9b4711..52760e32caa 100644 --- a/vortex-ffi/src/macros.rs +++ b/vortex-ffi/src/macros.rs @@ -66,7 +66,8 @@ macro_rules! arc_dyn_wrapper { } /// Extract a borrowed reference from a const pointer. - pub(crate) fn as_ref<'a>(ptr: *const $ffi_ident) -> &'a std::sync::Arc<$T> { + /// TODO hack for duckdb exporter + pub fn as_ref<'a>(ptr: *const $ffi_ident) -> &'a std::sync::Arc<$T> { use vortex::error::VortexExpect; // TODO(joe): propagate this error up instead of expecting &unsafe { ptr.as_ref() } From 3de17ad6d22d17aea5f187f0f0e467497486be5e Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 14:41:34 +0000 Subject: [PATCH 14/21] re-add stuff needed for rust --- vortex-duckdb/cpp/table_function.cpp | 82 ++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 8676b71b062..ad13e186a8a 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -29,6 +29,7 @@ std::string move_vx_err(vx_error* error) { } LogicalType from_dtype(const vx_dtype* dtype) { + throw BinderException("cannot convert dtype to logical type"); return LogicalType::UNKNOWN; } @@ -388,4 +389,85 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( } return DuckDBSuccess; } + +// +// TODO this stuff should be removed +// + +extern "C" duckdb_vx_string_map duckdb_vx_string_map_create() { + auto map = new InsertionOrderPreservingMap(); + return reinterpret_cast(map); +} + +extern "C" void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, const char *value) { + if (!map || !key || !value) { + return; + } + auto *cpp_map = reinterpret_cast *>(map); + (*cpp_map)[string(key)] = string(value); +} + +extern "C" void duckdb_vx_string_map_free(duckdb_vx_string_map map) { + if (!map) { + return; + } + auto *cpp_map = reinterpret_cast *>(map); + delete cpp_map; +} + +extern "C" void duckdb_vx_tfunc_virtual_columns_push(duckdb_vx_tfunc_virtual_cols_result ffi_result, + idx_t column_idx, + const char *name_str, + size_t name_len, + duckdb_logical_type ffi_type) { + if (!ffi_result || !name_str || !ffi_type) { + return; + } + + auto result = reinterpret_cast(ffi_result); + const auto logical_type = reinterpret_cast(ffi_type); + const auto name = string(name_str, name_len); + + auto table_col = TableColumn(std::move(name), *logical_type); + result->emplace(column_idx, std::move(table_col)); +} + +extern "C" size_t duckdb_vx_tfunc_bind_input_get_parameter_count(duckdb_vx_tfunc_bind_input ffi_input) { + if (!ffi_input) { + return 0; + } + const auto input = reinterpret_cast(ffi_input); + return input->inputs.size(); +} + +extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, + size_t index) { + if (!ffi_input || index >= duckdb_vx_tfunc_bind_input_get_parameter_count(ffi_input)) { + return nullptr; + } + const auto info = reinterpret_cast(ffi_input); + return reinterpret_cast(new Value(info->inputs[index])); +} + +extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_named_parameter(duckdb_vx_tfunc_bind_input ffi_input, + const char *name) { + if (!ffi_input || !name) { + return nullptr; + } + const auto info = reinterpret_cast(ffi_input); + + if (!info->named_parameters.contains(name)) { + return nullptr; + } + auto value = duckdb::make_uniq(info->named_parameters.at(name)); + return reinterpret_cast(value.release()); +} + +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) { + return; +} + } // namespace vortex From 9d74371386dc39e5bcd528186c3f374463369175 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 15:35:43 +0000 Subject: [PATCH 15/21] better --- vortex-duckdb/cpp/table_function.cpp | 77 ++++++++++++++++++++++++---- vortex-ffi/src/data_source.rs | 5 +- 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index ad13e186a8a..e27524e66c0 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -28,17 +28,70 @@ std::string move_vx_err(vx_error* error) { return str; } +LogicalTypeId from_ptype(vx_ptype ptype) { + using enum LogicalTypeId; + switch (ptype) { + case PTYPE_U8: return UTINYINT; + case PTYPE_U16: return USMALLINT; + case PTYPE_U32: return UINTEGER; + case PTYPE_U64: return UBIGINT; + case PTYPE_I8: return TINYINT; + case PTYPE_I16: return SMALLINT; + case PTYPE_I32: return INTEGER; + case PTYPE_I64: return BIGINT; + case PTYPE_F16: throw BinderException("F16 type not supported in Duckdb"); + case PTYPE_F32: return FLOAT; + case PTYPE_F64: return DOUBLE; + } +} + +LogicalType from_dtype(const vx_dtype* dtype); +LogicalType from_struct(const vx_dtype* dtype) { + const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); + uint64_t struct_size = vx_struct_fields_nfields(struct_dtype); + child_list_t children(struct_size); + for (uint64_t i = 0; i < struct_size; ++i) { + const vx_string* field_name = vx_struct_fields_field_name(struct_dtype, i); + const vx_dtype* field_dtype = vx_struct_fields_field_dtype(struct_dtype, i); + children[i] = { to_string(field_name), from_dtype(field_dtype) }; + } + + return LogicalType::STRUCT(children); +} + LogicalType from_dtype(const vx_dtype* dtype) { - throw BinderException("cannot convert dtype to logical type"); - return LogicalType::UNKNOWN; + using enum LogicalTypeId; + switch (vx_dtype_get_variant(dtype)) { + case DTYPE_NULL: return SQLNULL; + case DTYPE_BOOL: return BOOLEAN; + case DTYPE_PRIMITIVE: return from_ptype(vx_dtype_primitive_ptype(dtype)); + case DTYPE_UTF8: return VARCHAR; + case DTYPE_BINARY: return BLOB; + case DTYPE_STRUCT: return from_struct(dtype); + case DTYPE_DECIMAL: { + uint8_t width = vx_dtype_decimal_precision(dtype); + uint8_t scale = vx_dtype_decimal_scale(dtype); + return LogicalType::DECIMAL(width, scale); + }; + case DTYPE_LIST: { + LogicalType child_type = from_dtype(vx_dtype_list_element(dtype)); + return LogicalType::LIST(child_type); + } + case DTYPE_FIXED_SIZE_LIST: { + LogicalType child_type = from_dtype(vx_dtype_fixed_size_list_element(dtype)); + idx_t idx = vx_dtype_fixed_size_list_size(dtype); + return LogicalType::ARRAY(child_type, idx); + }; + case DTYPE_EXTENSION: { // TODO Temporal + throw BinderException("DTYPE_EXTENSION not supported"); + }; + }; } namespace vortex { struct CTableFunctionInfo final : TableFunctionInfo { explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) : vtab(vtab) {} - duckdb_vx_tfunc_vtab_t vtab; - // vtab.max_threads }; struct CTableBindData final : TableFunctionData { @@ -101,15 +154,15 @@ unique_ptr c_bind( vector &types, vector &names) { - if (info.inputs.empty()) - throw BinderException("missing file glob parameter"); + if (info.inputs.size() != 1) + throw BinderException("expected single file glob parameter"); + std::string files_glob = StringValue::Get(info.inputs[0]); - vx_error* err = nullptr; - // TODO check if it is string - std::string glob = info.inputs[0].GetValue(); + vx_session* session = vx_session_new(); vx_data_source_options opts = { - .files = glob.data(), + // files_glob lives till end of c_bind, vx_data_source_new copies the argument + .files = files_glob.data(), .fs_use_vortex = nullptr, .fs_set_userdata = nullptr, .fs_open = nullptr, @@ -128,8 +181,10 @@ unique_ptr c_bind( .cache_delete = nullptr }; - vx_session* session = vx_session_new(); + vx_error* err = nullptr; const vx_data_source* data_source = vx_data_source_new(session, &opts, &err); + if (err) + throw BinderException(move_vx_err(err)); const vx_dtype* dtype = vx_data_source_dtype(data_source); const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index f2f55b40fcd..2cbe897eb9d 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -173,8 +173,11 @@ pub unsafe extern "C-unwind" fn vx_data_source_new( } #[unsafe(no_mangle)] +// Create a non-owned dtype referencing dataframe. +// This dtype's lifetime is bound to underlying data source. +// Caller should not free this dtype manually pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { - vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) + vx_dtype::new_ref(vx_data_source::as_ref(ds).dtype()) } #[repr(C)] From 702c675de90a9e89eab8d86cf3cd81573e071705 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 19 Mar 2026 16:08:03 +0000 Subject: [PATCH 16/21] better --- vortex-duckdb/cpp/table_function.cpp | 12 +++-- .../src/duckdb/table_function/export_array.rs | 9 ++-- vortex-ffi/cinclude/vortex.h | 18 ++----- vortex-ffi/src/lib.rs | 1 - vortex-ffi/src/scan.rs | 51 ++++++++++++++----- 5 files changed, 56 insertions(+), 35 deletions(-) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index e27524e66c0..e72bf782fea 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -43,6 +43,8 @@ LogicalTypeId from_ptype(vx_ptype ptype) { case PTYPE_F32: return FLOAT; case PTYPE_F64: return DOUBLE; } + throw BinderException( + StringUtil::Format("value %d out of range for vx_ptype", ptype)); } LogicalType from_dtype(const vx_dtype* dtype); @@ -86,6 +88,7 @@ LogicalType from_dtype(const vx_dtype* dtype) { throw BinderException("DTYPE_EXTENSION not supported"); }; }; + throw BinderException(StringUtil::Format("value %d out of range for vx_dtype", dtype)); } namespace vortex { @@ -130,7 +133,7 @@ struct CTableBindData final : TableFunctionData { }; struct CTableGlobalData final : GlobalTableFunctionState { - explicit CTableGlobalData(const vx_scan* scan, idx_t max_threads) + explicit CTableGlobalData(vx_scan* scan, idx_t max_threads) : scan(scan), max_threads(max_threads) {} ~CTableGlobalData() override { @@ -139,7 +142,7 @@ struct CTableGlobalData final : GlobalTableFunctionState { idx_t MaxThreads() const override { return max_threads; } - const vx_scan* scan; + vx_scan* scan; idx_t max_threads; }; @@ -270,10 +273,9 @@ unique_ptr c_init_global(ClientContext &, TableFunctio }; vx_error* err = nullptr; - const vx_scan* scan = vx_data_source_scan( + vx_scan* scan = vx_data_source_scan( bind.data_source, &options, - /* partition estimate */ nullptr, &err); if (err) throw InvalidInputException(move_vx_err(err)); @@ -288,7 +290,7 @@ unique_ptr c_init_local(ExecutionContext &, } void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { - const vx_scan* const scan = input.global_state->Cast().scan; + vx_scan* const scan = input.global_state->Cast().scan; vx_error* err = nullptr; const vx_partition* partition = nullptr; const vx_array* array = nullptr; diff --git a/vortex-duckdb/src/duckdb/table_function/export_array.rs b/vortex-duckdb/src/duckdb/table_function/export_array.rs index d1f5a9db73a..69ecef6156f 100644 --- a/vortex-duckdb/src/duckdb/table_function/export_array.rs +++ b/vortex-duckdb/src/duckdb/table_function/export_array.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use vortex::array::Canonical; +use vortex::array::DynArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ScalarFnVTable; use vortex::array::arrays::Struct; @@ -8,11 +9,10 @@ use vortex::array::arrays::StructArray; use vortex::array::optimizer::ArrayOptimizer; use vortex::error::VortexExpect; use vortex::scalar_fn::fns::pack::Pack; -use vortex::array::DynArray; -use crate::duckdb::DataChunk; use crate::SESSION; use crate::cpp::duckdb_data_chunk; +use crate::duckdb::DataChunk; use crate::duckdb::TableFunction; use crate::exporter::ArrayExporter; use crate::exporter::ConversionCache; @@ -61,13 +61,14 @@ pub(crate) unsafe extern "C-unwind" fn export_array_callback( }; let mut exporter = ArrayExporter::try_new(&array_result, &conversion_cache, ctx) - .vortex_expect("failed to initialize array exporter"); + .vortex_expect("failed to initialize array exporter"); // Relaxed since there is no intra-instruction ordering required. //batch_id = Some(global_state.batch_id.fetch_add(1, Ordering::Relaxed)); batch_id += 1; - let has_more_data = exporter.export(chunk) + let has_more_data = exporter + .export(chunk) .vortex_expect("failed to export chunk"); //global_state // .bytes_read diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 29327e7b8b1..87029260998 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -921,28 +921,20 @@ const vx_partition *vx_partition_clone(const vx_partition *ptr); */ void vx_partition_free(const vx_partition *ptr); -/** - * Clone a borrowed [`vx_scan`], returning an owned [`vx_scan`]. - * - * - * Must be released with [`vx_scan_free`]. - */ -const vx_scan *vx_scan_clone(const vx_scan *ptr); - /** * Free an owned [`vx_scan`] object. */ -void vx_scan_free(const vx_scan *ptr); +void vx_scan_free(vx_scan *ptr); /** * Free an owned [`vx_array_it`] object. */ void vx_array_it_free(vx_array_it *ptr); -const vx_scan *vx_data_source_scan(const vx_data_source *_data_source, - const vx_scan_options *_options, - vx_partition_estimate *_partition_estimate, - vx_error_out err); +vx_scan * +vx_data_source_scan(const vx_data_source *data_source, const vx_scan_options *options, vx_error_out err); + +void vx_scan_partition_count(const vx_scan *scan, vx_partition_estimate *count); const vx_partition *vx_scan_next(const vx_scan *_scan, vx_error_out err); diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 343473b7d5b..655b3af2ccc 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -30,7 +30,6 @@ use std::sync::LazyLock; // TODO hack for duckdb exporter pub use array::vx_array; - pub use log::vx_log_level; use vortex::io::runtime::current::CurrentThreadRuntime; diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index b11c009fd4c..f691d6e45f9 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -9,10 +9,12 @@ use arrow_array::ffi_stream::FFI_ArrowArrayStream; use vortex::buffer::Buffer; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::io::runtime::BlockingRuntime; use vortex::scan::Selection; use vortex::scan::api::DataSourceScan; use vortex::scan::api::Partition; use vortex::scan::api::ScanRequest; +use vortex::array::expr::stats::Precision; use crate::array::vx_array; use crate::data_source::vx_data_source; @@ -20,11 +22,12 @@ use crate::error::try_or_default; use crate::error::vx_error_out; use crate::error::write_error; use crate::expression::vx_expression; +use crate::RUNTIME; crate::arc_dyn_wrapper!(dyn Partition, vx_partition); // TODO should be akin to DataSourceScanRef -crate::arc_dyn_wrapper!(dyn DataSourceScan, vx_scan); +crate::box_dyn_wrapper!(dyn DataSourceScan, vx_scan); pub struct VxArrayIt; crate::box_wrapper!(VxArrayIt, vx_array_it); @@ -120,20 +123,44 @@ fn scan_request(opts: *const vx_scan_options) -> VortexResult { } #[unsafe(no_mangle)] +// Create a new owned data source scan which must be freed by the caller. +// Scan can be consumed only once. +// Returns NULL and sets err on error. +// options may not be NULL. pub unsafe extern "C-unwind" fn vx_data_source_scan( - _data_source: *const vx_data_source, - _options: *const vx_scan_options, - _partition_estimate: *mut vx_partition_estimate, + data_source: *const vx_data_source, + options: *const vx_scan_options, err: vx_error_out, -) -> *const vx_scan { - //try_or_default(err, || { - // let request = scan_request(opts)?; - // let scan = vx_data_source::as_ref(ds).scan(request); - // Ok(vx_scan::new(Arc::new(scan))) - //}) +) -> *mut vx_scan { + try_or_default(err, || { + let request = scan_request(options)?; + RUNTIME.block_on(async { + let scan = vx_data_source::as_ref(data_source).scan(request).await?; + Ok(vx_scan::new(scan)) + }) + }) +} - write_error(err, "failed to create scan"); - ptr::null() +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_scan_partition_count( + scan: *const vx_scan, + count: *mut vx_partition_estimate, +) { + let scan = vx_scan::as_ref(scan); + let count = unsafe { &mut *count }; + match scan.partition_count() { + Some(Precision::Exact(value)) => { + count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_EXACT; + count.estimate = value as u64; + } + Some(Precision::Inexact(value)) => { + count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_INEXACT; + count.estimate = value as u64; + }, + None => { + count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_UNKNOWN; + } + } } #[unsafe(no_mangle)] From 7a30556ae3ef990a7027c2f850ca22f0b394c281 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 20 Mar 2026 10:44:18 +0000 Subject: [PATCH 17/21] better --- .../cpp/include/duckdb_vx/table_function.h | 1 - vortex-duckdb/cpp/table_function.cpp | 183 +++++++++++++----- vortex-ffi/cinclude/vortex.h | 39 ++-- vortex-ffi/src/scan.rs | 85 ++++++-- vortex-ffi/src/session.rs | 9 + 5 files changed, 233 insertions(+), 84 deletions(-) diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h index 1fa7759f5df..8bd54315da9 100644 --- a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h +++ b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h @@ -12,7 +12,6 @@ #include "error.h" #include "table_filter.h" #include "duckdb_vx/data.h" -#include "duckdb_vx/client_context.h" #include "duckdb_vx/duckdb_diagnostics.h" #ifdef __cplusplus diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index e72bf782fea..97b157f3a21 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#include "duckdb_vx.h" -#include "duckdb_vx/error.hpp" #include "duckdb_vx/table_function.h" DUCKDB_INCLUDES_BEGIN @@ -91,6 +89,107 @@ LogicalType from_dtype(const vx_dtype* dtype) { throw BinderException(StringUtil::Format("value %d out of range for vx_dtype", dtype)); } +// TODO This belongs in C++ part + +struct Session { + Session(): session(vx_session_new()) {} + + Session(const Session& other): session(vx_session_clone(other.session)) { } + + Session& operator=(const Session& other) { + session = vx_session_clone(other.session); + return *this; + } + + Session(Session&& other) noexcept : session(other.session) { + other.session = nullptr; + } + + Session& operator=(Session&& other) noexcept { + session = std::exchange(other.session, nullptr); + return *this; + } + + ~Session() { + if (session) vx_session_free(session); + } + + vx_session* session; +}; + +struct DataSource { + DataSource(const Session& session, vx_data_source_options& opts) { + vx_error* err; + data_source = vx_data_source_new(session.session, &opts, &err); + if (err) { + throw BinderException(move_vx_err(err)); + } + } + + DataSource(const DataSource& other) + : data_source(vx_data_source_clone(other.data_source)) { } + + DataSource& operator=(const DataSource& other) { + data_source = vx_data_source_clone(other.data_source); + return *this; + } + + DataSource(DataSource&& other) noexcept + : data_source(other.data_source) + { + other.data_source = nullptr; + } + + DataSource& operator=(DataSource&& other) noexcept { + data_source = std::exchange(other.data_source, nullptr); + return *this; + } + + ~DataSource() { + if (data_source) vx_data_source_free(data_source); + } + + vx_data_source_row_count row_count() const noexcept { + vx_data_source_row_count rc; + vx_data_source_get_row_count(data_source, &rc); + return rc; + } + + const vx_data_source* data_source; +}; + +struct Scan { + Scan(const DataSource& data_source, vx_scan_options& options) { + vx_error* err = nullptr; + scan = vx_data_source_scan(data_source.data_source, &options, &err); + if (err) { + throw BinderException(move_vx_err(err)); + } + } + + Scan(const Scan&) = delete; + Scan& operator=(const Scan&) = delete; + + Scan(Scan&& other) noexcept: scan(other.scan) { + other.scan = nullptr; + } + + Scan& operator= (Scan&& other) noexcept { + scan = std::exchange(other.scan, nullptr); + return *this; + } + + ~Scan() { + if (scan) vx_scan_free(scan); + } + + double progress() const noexcept { + return vx_scan_progress(scan); + } + + vx_scan* scan; +}; + namespace vortex { struct CTableFunctionInfo final : TableFunctionInfo { explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) : vtab(vtab) {} @@ -101,26 +200,21 @@ struct CTableBindData final : TableFunctionData { CTableBindData( unique_ptr info_p, const vector& column_names, - vx_session* session, - const vx_data_source* data_source) + Session&& session, + DataSource&& data_source) : info(std::move(info_p)) , column_names(column_names) - , session(session) - , data_source(data_source) {} + , session(std::move(session)) + , data_source(std::move(data_source)) {} - ~CTableBindData() override { - vx_data_source_free(data_source); - vx_session_free(session); - } + ~CTableBindData() override = default; unique_ptr Copy() const override { return make_uniq( make_uniq(info->vtab), column_names, - // TODO what's with this ownership? - session, - // TODO arcd inside, so fine, don't need to clone it? - data_source); + Session{session}, + DataSource{data_source}); } unique_ptr info; // TODO remove @@ -128,21 +222,19 @@ struct CTableBindData final : TableFunctionData { // precomputed from data_source using vx_struct_fields_field_name, see c_bind() vector column_names; - vx_session* session; - const vx_data_source* data_source; + Session session; + DataSource data_source; }; struct CTableGlobalData final : GlobalTableFunctionState { - explicit CTableGlobalData(vx_scan* scan, idx_t max_threads) - : scan(scan), max_threads(max_threads) {} + explicit CTableGlobalData(Scan&& scan, idx_t max_threads) + : scan(std::move(scan)), max_threads(max_threads) {} - ~CTableGlobalData() override { - vx_scan_free(scan); - } + ~CTableGlobalData() override = default; idx_t MaxThreads() const override { return max_threads; } - vx_scan* scan; + Scan scan; idx_t max_threads; }; @@ -161,7 +253,7 @@ unique_ptr c_bind( throw BinderException("expected single file glob parameter"); std::string files_glob = StringValue::Get(info.inputs[0]); - vx_session* session = vx_session_new(); + Session session; vx_data_source_options opts = { // files_glob lives till end of c_bind, vx_data_source_new copies the argument @@ -184,12 +276,9 @@ unique_ptr c_bind( .cache_delete = nullptr }; - vx_error* err = nullptr; - const vx_data_source* data_source = vx_data_source_new(session, &opts, &err); - if (err) - throw BinderException(move_vx_err(err)); + DataSource data_source{session, opts}; - const vx_dtype* dtype = vx_data_source_dtype(data_source); + const vx_dtype* dtype = vx_data_source_dtype(data_source.data_source); const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); for (uint64_t i = 0; i < vx_struct_fields_nfields(struct_dtype); ++i) { @@ -207,8 +296,8 @@ unique_ptr c_bind( return make_uniq( make_uniq(vtab), names, - session, - data_source); + std::move(session), + std::move(data_source)); } const vx_expression* projection( @@ -258,9 +347,9 @@ unique_ptr c_init_global(ClientContext &, TableFunctio const auto &bind = input.bind_data->Cast(); vx_scan_selection selection = { - .include = VX_S_INCLUDE_ALL, .idx = nullptr, - .idx_len = 0 + .idx_len = 0, + .include = VX_S_INCLUDE_ALL, }; vx_scan_options options = { .projection = projection(input.column_ids, input.projection_ids, bind.column_names), @@ -268,19 +357,12 @@ unique_ptr c_init_global(ClientContext &, TableFunctio .row_range_begin = 0, .row_range_end = 0, .selection = selection, + .limit = 0, .ordered = 0, - .limit = 0 }; - vx_error* err = nullptr; - vx_scan* scan = vx_data_source_scan( - bind.data_source, - &options, - &err); - if (err) - throw InvalidInputException(move_vx_err(err)); - - return make_uniq(scan, bind.info->vtab.max_threads); + Scan scan{bind.data_source, options}; + return make_uniq(std::move(scan), bind.info->vtab.max_threads); } unique_ptr c_init_local(ExecutionContext &, @@ -290,17 +372,19 @@ unique_ptr c_init_local(ExecutionContext &, } void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { - vx_scan* const scan = input.global_state->Cast().scan; + const Scan& scan = input.global_state->Cast().scan; vx_error* err = nullptr; - const vx_partition* partition = nullptr; + vx_partition* partition = nullptr; const vx_array* array = nullptr; auto export_array = input.bind_data->Cast().info->vtab.export_array; auto& batch_id = input.local_state->Cast().batch_id; // TODO this is essentially single-threaded processing - while (err == nullptr && (partition = vx_scan_next(scan, &err))) { - while ((array = vx_partition_next(partition, &err))) { + while ((partition = vx_scan_next(scan.scan, &err)) && !err) { + printf("next partition: %p, err: %p\n", (void*) partition, (void *) err); + while ((array = vx_partition_next(partition, &err)) && !err) { + printf("next array: %p, err: %p\n", (void*) array, (void *) err); uint64_t export_res = export_array( array, reinterpret_cast(&output)); @@ -310,6 +394,8 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou else batch_id = export_res; } + if (err) + throw InvalidInputException(move_vx_err(err)); vx_partition_free(partition); } if (err) @@ -343,8 +429,7 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou unique_ptr c_cardinality(ClientContext &, const FunctionData *bind_data) { auto stats = make_uniq(); - vx_data_source_row_count rc = {}; - vx_data_source_get_row_count(bind_data->Cast().data_source, &rc); + vx_data_source_row_count rc = bind_data->Cast().data_source.row_count(); switch (rc.cardinality) { case VX_CARD_ESTIMATE: { stats->has_estimated_cardinality = true; @@ -420,7 +505,7 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( }; tf.table_scan_progress = [](auto&, auto*, const GlobalTableFunctionState* state) { - return vx_scan_progress(state->Cast().scan); + return state->Cast().scan.progress(); }; tf.arguments.reserve(vtab->parameter_count); diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 87029260998..0d936e0396a 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -334,6 +334,12 @@ typedef struct vx_expression vx_expression; */ typedef struct vx_file vx_file; +/** + * A partition is a contiguous chunk of memory from which you can + * interatively get vx_arrays. + * TODO We're going away from exposing partitions to user, revise + * design + */ typedef struct vx_partition vx_partition; typedef struct vx_scan vx_scan; @@ -503,9 +509,9 @@ typedef struct { } vx_file_scan_options; typedef struct { - vx_scan_selection_include include; uint64_t *idx; size_t idx_len; + vx_scan_selection_include include; } vx_scan_selection; typedef struct { @@ -514,13 +520,13 @@ typedef struct { uint64_t row_range_begin; uint64_t row_range_end; vx_scan_selection selection; - int ordered; uint64_t limit; + int ordered; } vx_scan_options; typedef struct { - vx_partition_estimate_boundary boundary; uint64_t estimate; + vx_partition_estimate_boundary boundary; } vx_partition_estimate; #ifdef __cplusplus @@ -908,18 +914,10 @@ vx_array_iterator *vx_file_scan(const vx_session *session, */ void vx_set_log_level(vx_log_level level); -/** - * Clone a borrowed [`vx_partition`], returning an owned [`vx_partition`]. - * - * - * Must be released with [`vx_partition_free`]. - */ -const vx_partition *vx_partition_clone(const vx_partition *ptr); - /** * Free an owned [`vx_partition`] object. */ -void vx_partition_free(const vx_partition *ptr); +void vx_partition_free(vx_partition *ptr); /** * Free an owned [`vx_scan`] object. @@ -934,9 +932,15 @@ void vx_array_it_free(vx_array_it *ptr); vx_scan * vx_data_source_scan(const vx_data_source *data_source, const vx_scan_options *options, vx_error_out err); -void vx_scan_partition_count(const vx_scan *scan, vx_partition_estimate *count); +void vx_scan_partition_count(const vx_scan *scan, vx_partition_estimate *count, vx_error_out err); -const vx_partition *vx_scan_next(const vx_scan *_scan, vx_error_out err); +/** + * Get next owned partition out of a scan request which caller must free. + * Thread-safe TODO actually it's not thread-safe now + * Returns NULL and doesn't set err on exhaustion. + * Returns NULL and sets err on error. + */ +vx_partition *vx_scan_next(vx_scan *scan, vx_error_out err); void vx_partition_scan_arrow(const vx_partition *_partition, FFI_ArrowArrayStream *_stream, vx_error_out err); @@ -959,6 +963,13 @@ void vx_session_free(vx_session *ptr); */ vx_session *vx_session_new(void); +/** + * Clone a Vortex session, returning an owned copy. + * + * The caller is responsible for freeing the session with [`vx_session_free`]. + */ +vx_session *vx_session_clone(vx_session *session); + /** * Opens a writable array stream, where sink is used to push values into the stream. * To close the stream close the sink with `vx_array_sink_close`. diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index f691d6e45f9..fb289f5eab9 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -6,10 +6,12 @@ use std::ops::Range; use std::ptr; use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use futures::StreamExt; use vortex::buffer::Buffer; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::io::runtime::BlockingRuntime; +use vortex::scan::api::PartitionStream; use vortex::scan::Selection; use vortex::scan::api::DataSourceScan; use vortex::scan::api::Partition; @@ -24,10 +26,23 @@ use crate::error::write_error; use crate::expression::vx_expression; use crate::RUNTIME; -crate::arc_dyn_wrapper!(dyn Partition, vx_partition); - -// TODO should be akin to DataSourceScanRef -crate::box_dyn_wrapper!(dyn DataSourceScan, vx_scan); +crate::box_dyn_wrapper!( + /// A partition is a contiguous chunk of memory from which you can + /// interatively get vx_arrays. + /// TODO We're going away from exposing partitions to user, revise + /// design + dyn Partition, + vx_partition); + +// TODO To be able to do _next() functions in C we need to store the output. +// of partitions(). To be able to estimate number of partitions after scan is +// started, we save the value in enum's second option +pub enum VxDataSourceScan { + Pending(Box), + Started(PartitionStream), + Finished, +} +crate::box_wrapper!(VxDataSourceScan, vx_scan); pub struct VxArrayIt; crate::box_wrapper!(VxArrayIt, vx_array_it); @@ -41,9 +56,9 @@ pub enum vx_scan_selection_include { #[repr(C)] pub struct vx_scan_selection { - pub include: vx_scan_selection_include, pub idx: *mut u64, pub idx_len: usize, + pub include: vx_scan_selection_include, } // Distinct from ScanRequest for easier option handling from C @@ -54,8 +69,8 @@ pub struct vx_scan_options { pub row_range_begin: u64, pub row_range_end: u64, pub selection: vx_scan_selection, - pub ordered: c_int, pub limit: u64, + pub ordered: c_int, } #[repr(C)] @@ -67,8 +82,8 @@ pub enum vx_partition_estimate_boundary { #[repr(C)] pub struct vx_partition_estimate { - boundary: vx_partition_estimate_boundary, estimate: u64, + boundary: vx_partition_estimate_boundary, } fn scan_request(opts: *const vx_scan_options) -> VortexResult { @@ -136,7 +151,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( let request = scan_request(options)?; RUNTIME.block_on(async { let scan = vx_data_source::as_ref(data_source).scan(request).await?; - Ok(vx_scan::new(scan)) + Ok(vx_scan::new(Box::new(VxDataSourceScan::Pending(scan)))) }) }) } @@ -145,9 +160,13 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( pub unsafe extern "C-unwind" fn vx_scan_partition_count( scan: *const vx_scan, count: *mut vx_partition_estimate, + err: vx_error_out ) { - let scan = vx_scan::as_ref(scan); let count = unsafe { &mut *count }; + let VxDataSourceScan::Pending(scan) = vx_scan::as_ref(scan) else { + write_error(err, "can't get partition count of a scan that's already started"); + return; + }; match scan.partition_count() { Some(Precision::Exact(value)) => { count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_EXACT; @@ -164,16 +183,46 @@ pub unsafe extern "C-unwind" fn vx_scan_partition_count( } #[unsafe(no_mangle)] +/// Get next owned partition out of a scan request which caller must free. +/// Thread-safe TODO actually it's not thread-safe now +/// Returns NULL and doesn't set err on exhaustion. +/// Returns NULL and sets err on error. pub unsafe extern "C-unwind" fn vx_scan_next( - _scan: *const vx_scan, + scan: *mut vx_scan, err: vx_error_out, -) -> *const vx_partition { - //try_or_default(err, || { - // Ok(vx_partition::new(Arc::new())) - //}) +) -> *mut vx_partition { + let scan = vx_scan::as_mut(scan); + if let VxDataSourceScan::Finished = scan { + return ptr::null_mut(); + } - write_error(err, "failed to get next partition"); - ptr::null() + // No other function can observe intermediate state + // TODO actually it can because vx_scan_next should be thread_safe + unsafe { + let ptr = scan as *mut VxDataSourceScan; + let owned = ptr::read(ptr); + if let VxDataSourceScan::Pending(inner) = owned { + let stream = inner.partitions(); + ptr::write(ptr, VxDataSourceScan::Started(stream)); + } else { + ptr::write(ptr, owned); + } + } + let VxDataSourceScan::Started(stream) = scan else { + write_error(err, "internal error on vx_scan_next"); + return ptr::null_mut(); + }; + + try_or_default(err, || { + Ok(match RUNTIME.block_on(stream.next()) { + Some(partition) => vx_partition::new(partition?), + None => { + // TODO + //std::mem::replace(scan, VxDataSourceScan::Finished); + ptr::null_mut() + } + }) + }) } // TODO export nanoarrow headers? @@ -192,10 +241,6 @@ pub unsafe extern "C-unwind" fn vx_partition_next( _partition: *const vx_partition, err: vx_error_out, ) -> *const vx_array { - //try_or_default(err, || { - // Ok(vx_array::new(Arc::new())) - //}) - write_error(err, "failed to get next array out of partition"); ptr::null() } diff --git a/vortex-ffi/src/session.rs b/vortex-ffi/src/session.rs index 8950abb90d4..721b5a74322 100644 --- a/vortex-ffi/src/session.rs +++ b/vortex-ffi/src/session.rs @@ -24,3 +24,12 @@ pub unsafe extern "C-unwind" fn vx_session_new() -> *mut vx_session { VortexSession::default().with_handle(RUNTIME.handle()), )) } + +#[unsafe(no_mangle)] +/// Clone a Vortex session, returning an owned copy. +/// +/// The caller is responsible for freeing the session with [`vx_session_free`]. +pub unsafe extern "C-unwind" fn vx_session_clone(session: *mut vx_session) -> *mut vx_session { + let session = vx_session::as_mut(session); + vx_session::new(Box::new(session.clone())) +} From a30fab03b5d6f75a290e110200e975e7b2ebb6f8 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 20 Mar 2026 12:09:03 +0000 Subject: [PATCH 18/21] better --- vortex-duckdb/cpp/table_function.cpp | 156 +++++++++++++++++---------- vortex-ffi/src/scan.rs | 133 ++++++++++++++--------- 2 files changed, 179 insertions(+), 110 deletions(-) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 97b157f3a21..f8145132c77 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -91,77 +91,112 @@ LogicalType from_dtype(const vx_dtype* dtype) { // TODO This belongs in C++ part -struct Session { +class Session { +public: Session(): session(vx_session_new()) {} - Session(const Session& other): session(vx_session_clone(other.session)) { } - - Session& operator=(const Session& other) { - session = vx_session_clone(other.session); - return *this; + Session(const Session& other): session(vx_session_clone(other.session.get())) { } + Session(Session&& other) noexcept { + std::swap(session, other.session); } - Session(Session&& other) noexcept : session(other.session) { - other.session = nullptr; + Session& operator=(const Session& other) + { + session.reset(vx_session_clone(other.session.get())); + return *this; } Session& operator=(Session&& other) noexcept { - session = std::exchange(other.session, nullptr); + std::swap(session, other.session); return *this; } - ~Session() { - if (session) vx_session_free(session); - } + vx_session* handle() const noexcept { return session.get(); } - vx_session* session; +private: + struct deleter { + void operator()(vx_session* session) { + vx_session_free(session); + } + }; + std::unique_ptr session; }; -struct DataSource { +class DataSource { +public: DataSource(const Session& session, vx_data_source_options& opts) { vx_error* err; - data_source = vx_data_source_new(session.session, &opts, &err); + data_source.reset(vx_data_source_new(session.handle(), &opts, &err)); if (err) { throw BinderException(move_vx_err(err)); } } DataSource(const DataSource& other) - : data_source(vx_data_source_clone(other.data_source)) { } + : data_source(vx_data_source_clone(other.data_source.get())) { } - DataSource& operator=(const DataSource& other) { - data_source = vx_data_source_clone(other.data_source); - return *this; + DataSource(DataSource&& other) noexcept { + std::swap(data_source, other.data_source); } - DataSource(DataSource&& other) noexcept - : data_source(other.data_source) - { - other.data_source = nullptr; + DataSource& operator=(const DataSource& other) { + data_source.reset(vx_data_source_clone(other.data_source.get())); + return *this; } DataSource& operator=(DataSource&& other) noexcept { - data_source = std::exchange(other.data_source, nullptr); + std::swap(data_source, other.data_source); return *this; } - ~DataSource() { - if (data_source) vx_data_source_free(data_source); - } - vx_data_source_row_count row_count() const noexcept { vx_data_source_row_count rc; - vx_data_source_get_row_count(data_source, &rc); + vx_data_source_get_row_count(handle(), &rc); return rc; } - const vx_data_source* data_source; + const vx_data_source* handle() const noexcept { return data_source.get(); } + +private: + struct deleter { + void operator()(const vx_data_source* data_source) { + vx_data_source_free(data_source); + } + }; + std::unique_ptr data_source; }; -struct Scan { +struct Array { + ~Array() { + vx_array_free(array); + } + const vx_array* array; +}; + +struct Partition { + ~Partition() { + vx_partition_free(partition); + } + + std::optional next_array() { + vx_error* err = nullptr; + const vx_array* array = vx_partition_next(partition, &err); + if (err) { + throw BinderException(move_vx_err(err)); + } + if (!array) return std::nullopt; + return Array{array}; + } + + vx_partition* partition; +}; + + +class Scan { +public: Scan(const DataSource& data_source, vx_scan_options& options) { vx_error* err = nullptr; - scan = vx_data_source_scan(data_source.data_source, &options, &err); + scan.reset(vx_data_source_scan(data_source.handle(), &options, &err)); if (err) { throw BinderException(move_vx_err(err)); } @@ -170,24 +205,38 @@ struct Scan { Scan(const Scan&) = delete; Scan& operator=(const Scan&) = delete; - Scan(Scan&& other) noexcept: scan(other.scan) { - other.scan = nullptr; + Scan(Scan&& other) noexcept { + std::swap(scan, other.scan); } Scan& operator= (Scan&& other) noexcept { - scan = std::exchange(other.scan, nullptr); + std::swap(scan, other.scan); return *this; } - ~Scan() { - if (scan) vx_scan_free(scan); + double progress() const noexcept { + return vx_scan_progress(handle()); } - double progress() const noexcept { - return vx_scan_progress(scan); + std::optional next_partition() { + vx_error* err = nullptr; + vx_partition* partition = vx_scan_next(handle(), &err); + if (err) { + throw BinderException(move_vx_err(err)); + } + if (!partition) return std::nullopt; + return Partition{partition}; } - vx_scan* scan; + vx_scan* handle() const noexcept { return scan.get(); } + +private: + struct deleter { + void operator()(vx_scan* scan) { + vx_scan_free(scan); + } + }; + std::unique_ptr scan; }; namespace vortex { @@ -278,7 +327,7 @@ unique_ptr c_bind( DataSource data_source{session, opts}; - const vx_dtype* dtype = vx_data_source_dtype(data_source.data_source); + const vx_dtype* dtype = vx_data_source_dtype(data_source.handle()); const vx_struct_fields* struct_dtype = vx_dtype_struct_dtype(dtype); for (uint64_t i = 0; i < vx_struct_fields_nfields(struct_dtype); ++i) { @@ -372,34 +421,27 @@ unique_ptr c_init_local(ExecutionContext &, } void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { - const Scan& scan = input.global_state->Cast().scan; - vx_error* err = nullptr; - vx_partition* partition = nullptr; - const vx_array* array = nullptr; + Scan& scan = input.global_state->Cast().scan; + std::optional partition; + std::optional array; auto export_array = input.bind_data->Cast().info->vtab.export_array; auto& batch_id = input.local_state->Cast().batch_id; - // TODO this is essentially single-threaded processing - while ((partition = vx_scan_next(scan.scan, &err)) && !err) { - printf("next partition: %p, err: %p\n", (void*) partition, (void *) err); - while ((array = vx_partition_next(partition, &err)) && !err) { - printf("next array: %p, err: %p\n", (void*) array, (void *) err); + while ((partition = scan.next_partition())) { + printf("next partition: %p\n", (void*)partition->partition); + while ((array = partition->next_array())) { + printf("next array: %p\n", (void*) array->array); + uint64_t export_res = export_array( - array, + array->array, reinterpret_cast(&output)); - vx_array_free(array); if (export_res == std::numeric_limits::max()) batch_id = std::nullopt; else batch_id = export_res; } - if (err) - throw InvalidInputException(move_vx_err(err)); - vx_partition_free(partition); } - if (err) - throw InvalidInputException(move_vx_err(err)); } //void c_pushdown_complex_filter(ClientContext & /*context*/, diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index fb289f5eab9..9e77161d4bf 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -4,48 +4,49 @@ use core::slice; use std::ffi::c_int; use std::ops::Range; use std::ptr; +use std::sync::Arc; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use futures::StreamExt; +use vortex::array::expr::stats::Precision; +use vortex::array::stream::SendableArrayStream; use vortex::buffer::Buffer; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::io::runtime::BlockingRuntime; -use vortex::scan::api::PartitionStream; use vortex::scan::Selection; use vortex::scan::api::DataSourceScan; use vortex::scan::api::Partition; +use vortex::scan::api::PartitionStream; use vortex::scan::api::ScanRequest; -use vortex::array::expr::stats::Precision; +use crate::RUNTIME; use crate::array::vx_array; use crate::data_source::vx_data_source; use crate::error::try_or_default; use crate::error::vx_error_out; use crate::error::write_error; use crate::expression::vx_expression; -use crate::RUNTIME; -crate::box_dyn_wrapper!( - /// A partition is a contiguous chunk of memory from which you can - /// interatively get vx_arrays. - /// TODO We're going away from exposing partitions to user, revise - /// design - dyn Partition, - vx_partition); - -// TODO To be able to do _next() functions in C we need to store the output. -// of partitions(). To be able to estimate number of partitions after scan is -// started, we save the value in enum's second option -pub enum VxDataSourceScan { +pub enum VxScan { Pending(Box), Started(PartitionStream), Finished, } -crate::box_wrapper!(VxDataSourceScan, vx_scan); +crate::box_wrapper!(VxScan, vx_scan); -pub struct VxArrayIt; -crate::box_wrapper!(VxArrayIt, vx_array_it); +pub enum VxPartitionScan { + Pending(Box), + Started(SendableArrayStream), + Finished, +} +crate::box_wrapper!( + /// A partition is a contiguous chunk of memory from which you can + /// interatively get vx_arrays. + /// TODO We're going away from exposing partitions to user, revise + /// design + VxPartitionScan, + vx_partition); #[repr(C)] pub enum vx_scan_selection_include { @@ -151,7 +152,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( let request = scan_request(options)?; RUNTIME.block_on(async { let scan = vx_data_source::as_ref(data_source).scan(request).await?; - Ok(vx_scan::new(Box::new(VxDataSourceScan::Pending(scan)))) + Ok(vx_scan::new(Box::new(VxScan::Pending(scan)))) }) }) } @@ -160,11 +161,14 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( pub unsafe extern "C-unwind" fn vx_scan_partition_count( scan: *const vx_scan, count: *mut vx_partition_estimate, - err: vx_error_out + err: vx_error_out, ) { let count = unsafe { &mut *count }; - let VxDataSourceScan::Pending(scan) = vx_scan::as_ref(scan) else { - write_error(err, "can't get partition count of a scan that's already started"); + let VxScan::Pending(scan) = vx_scan::as_ref(scan) else { + write_error( + err, + "can't get partition count of a scan that's already started", + ); return; }; match scan.partition_count() { @@ -175,7 +179,7 @@ pub unsafe extern "C-unwind" fn vx_scan_partition_count( Some(Precision::Inexact(value)) => { count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_INEXACT; count.estimate = value as u64; - }, + } None => { count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_UNKNOWN; } @@ -192,37 +196,33 @@ pub unsafe extern "C-unwind" fn vx_scan_next( err: vx_error_out, ) -> *mut vx_partition { let scan = vx_scan::as_mut(scan); - if let VxDataSourceScan::Finished = scan { - return ptr::null_mut(); - } - - // No other function can observe intermediate state - // TODO actually it can because vx_scan_next should be thread_safe unsafe { - let ptr = scan as *mut VxDataSourceScan; - let owned = ptr::read(ptr); - if let VxDataSourceScan::Pending(inner) = owned { - let stream = inner.partitions(); - ptr::write(ptr, VxDataSourceScan::Started(stream)); - } else { - ptr::write(ptr, owned); - } - } - let VxDataSourceScan::Started(stream) = scan else { - write_error(err, "internal error on vx_scan_next"); - return ptr::null_mut(); - }; + let ptr = scan as *mut VxScan; - try_or_default(err, || { - Ok(match RUNTIME.block_on(stream.next()) { - Some(partition) => vx_partition::new(partition?), - None => { - // TODO - //std::mem::replace(scan, VxDataSourceScan::Finished); - ptr::null_mut() + let on_finish = || -> VortexResult<*mut vx_partition> { + ptr::write(ptr, VxScan::Finished); + Ok(ptr::null_mut()) + }; + + let on_stream = |mut stream: PartitionStream| -> VortexResult<*mut vx_partition> { + match RUNTIME.block_on(stream.next()) { + Some(partition) => { + let partition = VxPartitionScan::Pending(partition?); + let partition = vx_partition::new(Box::new(partition)); + ptr::write(ptr, VxScan::Started(stream)); + Ok(partition) + } + None => on_finish(), } + }; + + let owned = ptr::read(ptr); + try_or_default(err, || match owned { + VxScan::Pending(scan) => on_stream(scan.partitions()), + VxScan::Started(stream) => on_stream(stream), + VxScan::Finished => on_finish(), }) - }) + } } // TODO export nanoarrow headers? @@ -237,12 +237,39 @@ pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( } #[unsafe(no_mangle)] +/// Get next vx_array out of this partition. +/// Thread-unsafe. pub unsafe extern "C-unwind" fn vx_partition_next( - _partition: *const vx_partition, + partition: *mut vx_partition, err: vx_error_out, ) -> *const vx_array { - write_error(err, "failed to get next array out of partition"); - ptr::null() + let partition = vx_partition::as_mut(partition); + unsafe { + let ptr = partition as *mut VxPartitionScan; + + let on_finish = || -> VortexResult<*const vx_array> { + ptr::write(ptr, VxPartitionScan::Finished); + Ok(ptr::null_mut()) + }; + + let on_stream = |mut stream: SendableArrayStream| -> VortexResult<*const vx_array> { + match RUNTIME.block_on(stream.next()) { + Some(array) => { + let array = vx_array::new(Arc::new(array?)); + ptr::write(ptr, VxPartitionScan::Started(stream)); + Ok(array) + } + None => on_finish(), + } + }; + + let owned = ptr::read(ptr); + try_or_default(err, || match owned { + VxPartitionScan::Pending(partition) => on_stream(partition.execute()?), + VxPartitionScan::Started(stream) => on_stream(stream), + VxPartitionScan::Finished => on_finish(), + }) + } } #[unsafe(no_mangle)] From 92fede127f896b12539bc5834913a397882e144a Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 20 Mar 2026 15:28:36 +0000 Subject: [PATCH 19/21] better --- vortex-duckdb/Cargo.toml | 2 +- vortex-duckdb/build.rs | 4 ++ vortex-duckdb/cpp/table_function.cpp | 80 ++++++++++++++++++---------- vortex-ffi/src/scan.rs | 32 +++++++---- 4 files changed, 79 insertions(+), 39 deletions(-) diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index d93f213e59b..6c821c2ac4b 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -56,6 +56,6 @@ workspace = true [build-dependencies] bindgen = { workspace = true } cbindgen = { workspace = true } -cc = { workspace = true } +cc = { workspace = true , features = ["parallel"] } reqwest = { workspace = true } zip = { workspace = true } diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index 660e4b822e3..a15e7854cdd 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -326,11 +326,15 @@ fn c2rust(crate_dir: &Path, duckdb_include_dir: &Path) { } fn cpp(duckdb_include_dir: &Path) { + //println!("cargo:rustc-link-arg=-fsanitize=address"); cc::Build::new() .std("c++20") // Duckdb sources fail -Wno-unused-parameter .flags(["-Wall", "-Wextra", "-Wpedantic", "-Wno-unused-parameter"]) + // TODO + //.flag("-fsanitize=address") .cpp(true) + .debug(true) .include(duckdb_include_dir) .include("cpp/include") .include(FFI_INCLUDE) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index f8145132c77..907fb87f882 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -166,21 +166,33 @@ class DataSource { std::unique_ptr data_source; }; -struct Array { - ~Array() { - vx_array_free(array); - } - const vx_array* array; +class Array { +public: + const vx_array* handle() const noexcept { return array.get(); } +private: + friend class Partition; + + Array(const vx_array* array): array(array) {} + + struct deleter { + void operator()(const vx_array* array) { + vx_array_free(array); + } + }; + std::unique_ptr array; }; -struct Partition { - ~Partition() { - vx_partition_free(partition); +class Partition { +public: + Partition(const Partition&) = delete; + Partition& operator=(const Partition&) = delete; + Partition(Partition&& other) noexcept { + std::swap(partition, other.partition); } std::optional next_array() { vx_error* err = nullptr; - const vx_array* array = vx_partition_next(partition, &err); + const vx_array* array = vx_partition_next(handle(), &err); if (err) { throw BinderException(move_vx_err(err)); } @@ -188,7 +200,18 @@ struct Partition { return Array{array}; } - vx_partition* partition; + vx_partition* handle() const noexcept { return partition.get(); } + +private: + friend class Scan; + Partition(vx_partition* partition): partition(partition) {} + + struct deleter { + void operator()(vx_partition* partition) { + vx_partition_free(partition); + } + }; + std::unique_ptr partition; }; @@ -422,25 +445,28 @@ unique_ptr c_init_local(ExecutionContext &, void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { Scan& scan = input.global_state->Cast().scan; - std::optional partition; - std::optional array; - - auto export_array = input.bind_data->Cast().info->vtab.export_array; auto& batch_id = input.local_state->Cast().batch_id; - while ((partition = scan.next_partition())) { - printf("next partition: %p\n", (void*)partition->partition); - while ((array = partition->next_array())) { - printf("next array: %p\n", (void*) array->array); - - uint64_t export_res = export_array( - array->array, - reinterpret_cast(&output)); - if (export_res == std::numeric_limits::max()) - batch_id = std::nullopt; - else - batch_id = export_res; - } + printf("thread id: %lu\n", std::hash{}(std::this_thread::get_id())); + + auto next_partition = scan.next_partition(); + if (!next_partition) return; + Partition partition = std::move(*next_partition); + + printf("partition: %p\n", (void*)partition.handle()); + + std::optional array; + auto export_array = input.bind_data->Cast().info->vtab.export_array; + while ((array = partition.next_array())) { + printf("next array: %p\n", (void*) array->handle()); + uint64_t export_res = export_array( + array->handle(), + reinterpret_cast(&output)); + if (export_res == std::numeric_limits::max()) + batch_id = std::nullopt; + else + batch_id = export_res; + printf("batch id: %lu\n", batch_id.value_or(-1)); } } diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 9e77161d4bf..c639334e016 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -5,6 +5,7 @@ use std::ffi::c_int; use std::ops::Range; use std::ptr; use std::sync::Arc; +use std::sync::Mutex; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use futures::StreamExt; @@ -28,11 +29,12 @@ use crate::error::vx_error_out; use crate::error::write_error; use crate::expression::vx_expression; -pub enum VxScan { +pub enum VxScanState { Pending(Box), Started(PartitionStream), Finished, } +pub type VxScan = Mutex; crate::box_wrapper!(VxScan, vx_scan); pub enum VxPartitionScan { @@ -152,7 +154,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( let request = scan_request(options)?; RUNTIME.block_on(async { let scan = vx_data_source::as_ref(data_source).scan(request).await?; - Ok(vx_scan::new(Box::new(VxScan::Pending(scan)))) + Ok(vx_scan::new(Box::new(Mutex::new(VxScanState::Pending(scan))))) }) }) } @@ -164,7 +166,10 @@ pub unsafe extern "C-unwind" fn vx_scan_partition_count( err: vx_error_out, ) { let count = unsafe { &mut *count }; - let VxScan::Pending(scan) = vx_scan::as_ref(scan) else { + let scan = vx_scan::as_ref(scan); + let mut scan = scan.lock().expect("failed to lock mutex"); + let scan = &mut *scan; + let VxScanState::Pending(scan) = scan else { write_error( err, "can't get partition count of a scan that's already started", @@ -187,8 +192,11 @@ pub unsafe extern "C-unwind" fn vx_scan_partition_count( } #[unsafe(no_mangle)] -/// Get next owned partition out of a scan request which caller must free. -/// Thread-safe TODO actually it's not thread-safe now +/// Get next owned partition out of a scan request. +/// Caller must free this partition using vx_partition_free. +/// This method is thread-safe. +/// If using in a sync multi-thread runtime, users are encouraged to create a +/// worker thread per partition. /// Returns NULL and doesn't set err on exhaustion. /// Returns NULL and sets err on error. pub unsafe extern "C-unwind" fn vx_scan_next( @@ -196,11 +204,13 @@ pub unsafe extern "C-unwind" fn vx_scan_next( err: vx_error_out, ) -> *mut vx_partition { let scan = vx_scan::as_mut(scan); + let mut scan = scan.lock().expect("failed to lock mutex"); + let scan = &mut *scan; unsafe { - let ptr = scan as *mut VxScan; + let ptr = scan as *mut VxScanState; let on_finish = || -> VortexResult<*mut vx_partition> { - ptr::write(ptr, VxScan::Finished); + ptr::write(ptr, VxScanState::Finished); Ok(ptr::null_mut()) }; @@ -209,7 +219,7 @@ pub unsafe extern "C-unwind" fn vx_scan_next( Some(partition) => { let partition = VxPartitionScan::Pending(partition?); let partition = vx_partition::new(Box::new(partition)); - ptr::write(ptr, VxScan::Started(stream)); + ptr::write(ptr, VxScanState::Started(stream)); Ok(partition) } None => on_finish(), @@ -218,9 +228,9 @@ pub unsafe extern "C-unwind" fn vx_scan_next( let owned = ptr::read(ptr); try_or_default(err, || match owned { - VxScan::Pending(scan) => on_stream(scan.partitions()), - VxScan::Started(stream) => on_stream(stream), - VxScan::Finished => on_finish(), + VxScanState::Pending(scan) => on_stream(scan.partitions()), + VxScanState::Started(stream) => on_stream(stream), + VxScanState::Finished => on_finish(), }) } } From 9af9fd6ad06e4f6cf946c2c0d496a2a69ec9d7ba Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 20 Mar 2026 16:11:58 +0000 Subject: [PATCH 20/21] better --- vortex-duckdb/cpp/table_function.cpp | 8 +--- vortex-ffi/cinclude/vortex.h | 40 +++++++++-------- vortex-ffi/src/scan.rs | 66 +++++++++++++++++++--------- 3 files changed, 68 insertions(+), 46 deletions(-) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 907fb87f882..3135e05eba9 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "duckdb_vx/table_function.h" +#include "vortex.h" DUCKDB_INCLUDES_BEGIN #include "duckdb.h" @@ -382,6 +383,7 @@ const vx_expression* projection( if (projection_ids.empty()) { for (column_t id : column_ids) { + if (id == COLUMN_IDENTIFIER_EMPTY) continue; if (column_names.size() < id) throw InvalidInputException(StringUtil::Format( "Expected column id %d but there are %d columns", @@ -447,18 +449,13 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou Scan& scan = input.global_state->Cast().scan; auto& batch_id = input.local_state->Cast().batch_id; - printf("thread id: %lu\n", std::hash{}(std::this_thread::get_id())); - auto next_partition = scan.next_partition(); if (!next_partition) return; Partition partition = std::move(*next_partition); - printf("partition: %p\n", (void*)partition.handle()); - std::optional array; auto export_array = input.bind_data->Cast().info->vtab.export_array; while ((array = partition.next_array())) { - printf("next array: %p\n", (void*) array->handle()); uint64_t export_res = export_array( array->handle(), reinterpret_cast(&output)); @@ -466,7 +463,6 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou batch_id = std::nullopt; else batch_id = export_res; - printf("batch id: %lu\n", batch_id.value_or(-1)); } } diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 0d936e0396a..529b83ad360 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -157,10 +157,10 @@ typedef enum { } vx_scan_selection_include; typedef enum { - VX_P_ESTIMATE_UNKNOWN = 0, - VX_P_ESTIMATE_EXACT = 1, - VX_P_ESTIMATE_INEXACT = 2, -} vx_partition_estimate_boundary; + VX_ESTIMATE_UNKNOWN = 0, + VX_ESTIMATE_EXACT = 1, + VX_ESTIMATE_INEXACT = 2, +} vx_estimate_boundary; /** * Physical type enum, represents the in-memory physical layout but might represent a different logical type. @@ -268,8 +268,6 @@ typedef struct VxFileHandle VxFileHandle; */ typedef struct vx_array vx_array; -typedef struct vx_array_it vx_array_it; - /** * A Vortex array iterator. * @@ -526,8 +524,8 @@ typedef struct { typedef struct { uint64_t estimate; - vx_partition_estimate_boundary boundary; -} vx_partition_estimate; + vx_estimate_boundary boundary; +} vx_estimate; #ifdef __cplusplus extern "C" { @@ -914,37 +912,41 @@ vx_array_iterator *vx_file_scan(const vx_session *session, */ void vx_set_log_level(vx_log_level level); -/** - * Free an owned [`vx_partition`] object. - */ -void vx_partition_free(vx_partition *ptr); - /** * Free an owned [`vx_scan`] object. */ void vx_scan_free(vx_scan *ptr); /** - * Free an owned [`vx_array_it`] object. + * Free an owned [`vx_partition`] object. */ -void vx_array_it_free(vx_array_it *ptr); +void vx_partition_free(vx_partition *ptr); vx_scan * vx_data_source_scan(const vx_data_source *data_source, const vx_scan_options *options, vx_error_out err); -void vx_scan_partition_count(const vx_scan *scan, vx_partition_estimate *count, vx_error_out err); +void vx_scan_partition_count(const vx_scan *scan, vx_estimate *count, vx_error_out err); /** - * Get next owned partition out of a scan request which caller must free. - * Thread-safe TODO actually it's not thread-safe now + * Get next owned partition out of a scan request. + * Caller must free this partition using vx_partition_free. + * This method is thread-safe. + * If using in a sync multi-thread runtime, users are encouraged to create a + * worker thread per partition. * Returns NULL and doesn't set err on exhaustion. * Returns NULL and sets err on error. */ vx_partition *vx_scan_next(vx_scan *scan, vx_error_out err); +void vx_partition_row_count(const vx_partition *partition, vx_estimate *count, vx_error_out err); + void vx_partition_scan_arrow(const vx_partition *_partition, FFI_ArrowArrayStream *_stream, vx_error_out err); -const vx_array *vx_partition_next(const vx_partition *_partition, vx_error_out err); +/** + * Get next vx_array out of this partition. + * Thread-unsafe. + */ +const vx_array *vx_partition_next(vx_partition *partition, vx_error_out err); /** * Scan progress between 0.0 and 1.0 diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index c639334e016..bab4b888947 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -77,16 +77,16 @@ pub struct vx_scan_options { } #[repr(C)] -pub enum vx_partition_estimate_boundary { - VX_P_ESTIMATE_UNKNOWN = 0, - VX_P_ESTIMATE_EXACT = 1, - VX_P_ESTIMATE_INEXACT = 2, +pub enum vx_estimate_boundary { + VX_ESTIMATE_UNKNOWN = 0, + VX_ESTIMATE_EXACT = 1, + VX_ESTIMATE_INEXACT = 2, } #[repr(C)] -pub struct vx_partition_estimate { +pub struct vx_estimate { estimate: u64, - boundary: vx_partition_estimate_boundary, + boundary: vx_estimate_boundary, } fn scan_request(opts: *const vx_scan_options) -> VortexResult { @@ -128,7 +128,7 @@ fn scan_request(opts: *const vx_scan_options) -> VortexResult { let end = opts.row_range_end; let row_range = (start > 0 || end > 0).then_some(Range { start, end }); - let limit = (opts.limit == 0).then_some(opts.limit); + let limit = (opts.limit != 0).then_some(opts.limit); Ok(ScanRequest { projection, @@ -159,10 +159,26 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( }) } +fn estimate>(estimate: Option>, out: &mut vx_estimate) { + match estimate { + Some(Precision::Exact(value)) => { + out.boundary = vx_estimate_boundary::VX_ESTIMATE_EXACT; + out.estimate = value.into(); + } + Some(Precision::Inexact(value)) => { + out.boundary = vx_estimate_boundary::VX_ESTIMATE_INEXACT; + out.estimate = value.into(); + } + None => { + out.boundary = vx_estimate_boundary::VX_ESTIMATE_UNKNOWN; + } + } +} + #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scan_partition_count( scan: *const vx_scan, - count: *mut vx_partition_estimate, + count: *mut vx_estimate, err: vx_error_out, ) { let count = unsafe { &mut *count }; @@ -176,19 +192,10 @@ pub unsafe extern "C-unwind" fn vx_scan_partition_count( ); return; }; - match scan.partition_count() { - Some(Precision::Exact(value)) => { - count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_EXACT; - count.estimate = value as u64; - } - Some(Precision::Inexact(value)) => { - count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_INEXACT; - count.estimate = value as u64; - } - None => { - count.boundary = vx_partition_estimate_boundary::VX_P_ESTIMATE_UNKNOWN; - } - } + estimate(scan.partition_count().map(|x| match x { + Precision::Exact(v) => Precision::Exact(v as u64), + Precision::Inexact(v) => Precision::Inexact(v as u64), + }), count) } #[unsafe(no_mangle)] @@ -235,6 +242,23 @@ pub unsafe extern "C-unwind" fn vx_scan_next( } } +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_partition_row_count( + partition: *const vx_partition, + count: *mut vx_estimate, + err: vx_error_out +) { + let partition = vx_partition::as_ref(partition); + let VxPartitionScan::Pending(partition) = partition else { + write_error( + err, + "can't get row count of a partition that's already started", + ); + return; + }; + estimate(partition.row_count(), unsafe { &mut *count} ) +} + // TODO export nanoarrow headers? #[unsafe(no_mangle)] From 2fcd327a6dab141f67d0aa54666ec64966cd1196 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 20 Mar 2026 17:45:58 +0000 Subject: [PATCH 21/21] better --- vortex-duckdb/cpp/table_function.cpp | 86 +++++++++++++++------------- 1 file changed, 45 insertions(+), 41 deletions(-) diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 3135e05eba9..8405941ea07 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -272,31 +272,31 @@ struct CTableFunctionInfo final : TableFunctionInfo { struct CTableBindData final : TableFunctionData { CTableBindData( unique_ptr info_p, - const vector& column_names, Session&& session, - DataSource&& data_source) + DataSource&& data_source, + const vector& column_names) : info(std::move(info_p)) - , column_names(column_names) , session(std::move(session)) - , data_source(std::move(data_source)) {} + , data_source(std::move(data_source)) + , column_names(column_names) {} ~CTableBindData() override = default; unique_ptr Copy() const override { return make_uniq( make_uniq(info->vtab), - column_names, Session{session}, - DataSource{data_source}); + DataSource{data_source}, + column_names); } unique_ptr info; // TODO remove - // precomputed from data_source using vx_struct_fields_field_name, see c_bind() - vector column_names; - Session session; DataSource data_source; + + vector column_names; + vector filters; }; struct CTableGlobalData final : GlobalTableFunctionState { @@ -368,12 +368,12 @@ unique_ptr c_bind( auto& vtab = info.table_function.function_info->Cast().vtab; return make_uniq( make_uniq(vtab), - names, std::move(session), - std::move(data_source)); + std::move(data_source), + names); } -const vx_expression* projection( +const vx_expression* make_projection( const vector& column_ids, const vector& projection_ids, const vector& column_names @@ -417,6 +417,16 @@ const vx_expression* projection( return expr; } +const vx_expression* make_filter( + optional_ptr table_filters, + const vector& column_ids, + const vector& column_names, + const vector& additional_filters, + const vx_dtype* dtype +) { + return nullptr; +} + unique_ptr c_init_global(ClientContext &, TableFunctionInitInput &input) { const auto &bind = input.bind_data->Cast(); @@ -425,9 +435,16 @@ unique_ptr c_init_global(ClientContext &, TableFunctio .idx_len = 0, .include = VX_S_INCLUDE_ALL, }; + + const vx_dtype* dtype = vx_data_source_dtype(bind.data_source.handle()); + const vx_expression* projection = make_projection( + input.column_ids, input.projection_ids, bind.column_names); + const vx_expression* filter = make_filter(input.filters, + input.column_ids, bind.column_names, bind.filters, dtype); + vx_scan_options options = { - .projection = projection(input.column_ids, input.projection_ids, bind.column_names), - .filter = nullptr, + .projection = projection, + .filter = filter, .row_range_begin = 0, .row_range_end = 0, .selection = selection, @@ -446,10 +463,10 @@ unique_ptr c_init_local(ExecutionContext &, } void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &output) { - Scan& scan = input.global_state->Cast().scan; + auto& global_state = input.global_state->Cast(); auto& batch_id = input.local_state->Cast().batch_id; - auto next_partition = scan.next_partition(); + auto next_partition = global_state.scan.next_partition(); if (!next_partition) return; Partition partition = std::move(*next_partition); @@ -466,30 +483,17 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou } } -//void c_pushdown_complex_filter(ClientContext & /*context*/, -// LogicalGet & /*get*/, -// FunctionData *bind_data, -// vector> &filters) { -// if (filters.empty()) { -// return; -// } -// -// auto &bind = bind_data->Cast(); -// -// for (auto iter = filters.begin(); iter != filters.end();) { -// duckdb_vx_error error_out = nullptr; -// auto pushed = -// bind.info->vtab.pushdown_complex_filter(bind_data->Cast().ffi_data->DataPtr(), -// reinterpret_cast(iter->get()), -// &error_out); -// if (error_out) { -// throw BinderException(IntoErrString(error_out)); -// } -// -// // If the pushdown complex filter returns true, we can remove the filter from the list. -// iter = pushed ? filters.erase(iter) : std::next(iter); -// } -//} +const vx_expression* from_filter(const Expression& filter) { + return nullptr; +} + +void c_pushdown_complex_filter(ClientContext &, LogicalGet &, FunctionData *bind_data, vector> &filters) { + auto &bind = bind_data->Cast(); + // We don't erase filters, see Nick's comment in datasource.rs + for (auto iter = filters.begin(); iter != filters.end(); ++iter) { + bind.filters.emplace_back(from_filter(**iter)); + } +} unique_ptr c_cardinality(ClientContext &, const FunctionData *bind_data) { auto stats = make_uniq(); @@ -537,7 +541,7 @@ extern "C" duckdb_state duckdb_vx_tfunc_register( auto db = wrapper->database->instance; auto tf = TableFunction(vtab->name, {}, c_function, c_bind, c_init_global, c_init_local); - //tf.pushdown_complex_filter = c_pushdown_complex_filter; + tf.pushdown_complex_filter = c_pushdown_complex_filter; //tf.projection_pushdown = vtab->projection_pushdown; //tf.filter_pushdown = vtab->filter_pushdown;