diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index b76a3b45fff..50bb093f95c 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -30,7 +30,7 @@ const BUILD_MARKER: &str = ".vx-build-complete"; const DUCKDB_CACHE_DIR: &str = "vortex-duckdb-cache"; const EXTRACT_MARKER: &str = ".vx-extract-complete"; -const SOURCE_FILES: [&str; 11] = [ +const SOURCE_FILES: [&str; 12] = [ "cpp/vortex_duckdb.cpp", "cpp/copy_function.cpp", "cpp/expr.cpp", @@ -40,6 +40,7 @@ const SOURCE_FILES: [&str; 11] = [ "cpp/cast_pushdown.cpp", "cpp/aggregate_fn_pushdown.cpp", "cpp/table_filter.cpp", + "cpp/multi_file_reader.cpp", "cpp/table_function.cpp", "cpp/vector.cpp", ]; @@ -352,6 +353,48 @@ fn extract(archive: &Path, dest: &Path) { zip::ZipArchive::new(file).unwrap().extract(dest).unwrap(); } +fn git_apply(repo_dir: &Path, patch: &Path, args: &[&str]) -> bool { + let output = Command::new("git") + .current_dir(repo_dir) + .args(["apply", "-p1"]) + .args(args) + .arg(patch) + .output(); + match output { + Ok(out) => out.status.success(), + Err(e) => { + println!("cargo:error=git is required to patch DuckDB sources: {e}"); + exit(1); + } + } +} + +fn apply_source_patches(crate_dir: &Path, repo_dir: &Path) { + let mut patches: Vec = fs::read_dir(crate_dir.join("patches")) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "diff")) + .collect(); + patches.sort(); + + for patch in patches { + // A successful reverse dry-run means the patch is already applied. + if git_apply(repo_dir, &patch, &["--check", "--reverse"]) { + continue; + } + if !git_apply(repo_dir, &patch, &[]) { + println!( + "cargo:error=Failed to apply {} to {}; delete that directory to re-extract \ + DuckDB sources", + patch.display(), + repo_dir.display() + ); + exit(1); + } + println!("cargo:info=Applied {}", patch.display()); + } +} + /// Download DuckDB library archive from R2 and extract it. /// Return false if archive is not available or download failed fn download_prebuilt(version: &DuckDBVersion, library_dir: &Path, target: &str) -> bool { @@ -576,6 +619,7 @@ fn cbindgen_rust2c(crate_dir: &Path) { fn main() { println!("cargo:rerun-if-changed=cpp/include"); + println!("cargo:rerun-if-changed=patches"); println!("cargo:rerun-if-env-changed=VX_DUCKDB_DEBUG"); println!("cargo:rerun-if-env-changed=VX_DUCKDB_SAN"); println!("cargo:rerun-if-env-changed=CARGO_HTTP_TIMEOUT"); @@ -656,6 +700,8 @@ fn main() { fs::write(&extract_marker, version.to_string()).unwrap(); } + apply_source_patches(&crate_dir, &inner_dir); + drop(fs::remove_file(&duckdb_dir)); drop(fs::remove_dir_all(&duckdb_dir)); symlink(&source_dir, &duckdb_dir).unwrap(); diff --git a/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp b/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp index 2caa3463d53..c08d7bedab6 100644 --- a/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp @@ -136,5 +136,5 @@ LogicalGet *GetChildGet(const LogicalAggregate &agg) { return nullptr; } LogicalGet &get = op->Cast(); - return get.function.bind == duckdb_vx_table_function_bind ? &get : nullptr; + return is_vortex_scan(get.function) ? &get : nullptr; } diff --git a/vortex-duckdb/cpp/cast_pushdown.cpp b/vortex-duckdb/cpp/cast_pushdown.cpp index 3acd5b04161..d1a3b849c00 100644 --- a/vortex-duckdb/cpp/cast_pushdown.cpp +++ b/vortex-duckdb/cpp/cast_pushdown.cpp @@ -25,7 +25,7 @@ static bool ReachesPushdownGet(const LogicalOperator &op) { cur = cur->children[0].get(); switch (cur->type) { case LogicalOperatorType::LOGICAL_GET: - return cur->Cast().function.bind == duckdb_vx_table_function_bind; + return is_vortex_scan(cur->Cast().function); case LogicalOperatorType::LOGICAL_PROJECTION: continue; default: diff --git a/vortex-duckdb/cpp/include/multi_file_reader.hpp b/vortex-duckdb/cpp/include/multi_file_reader.hpp new file mode 100644 index 00000000000..3e0e0f541e1 --- /dev/null +++ b/vortex-duckdb/cpp/include/multi_file_reader.hpp @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "data.hpp" +#include "duckdb/common/multi_file/multi_file_function.hpp" + +using namespace duckdb; + +struct VortexBindData final : TableFunctionData { + VortexBindData() = default; + unique_ptr Copy() const override; + bool Equals(const FunctionData &other) const override; + + unique_ptr ffi_bind_data; +}; + +struct VortexBindResult { + vector &return_types; + vector &names; +}; + +struct VortexGlobalState final : GlobalTableFunctionState { + VortexGlobalState() = default; + ~VortexGlobalState() override = default; + + const void *ffi_bind_data = nullptr; // needed for local state partial accumulation + unique_ptr ffi_global_state; +}; + +struct VortexLocalState final : LocalTableFunctionState { + VortexLocalState() = default; + unique_ptr ffi_local_state; +}; + +struct VortexMultiFileReader final : MultiFileReader { + inline unique_ptr Copy() const override { + return make_uniq(); + } + + /* + * Called after InitializeGlobalState but before TryInitializeScan under + * file-local lock. Used to avoid opening the file for scanning if footer + * statistics prove false for pushed filter or file index is not present in + * file selection. + */ + ReaderInitializeType InitializeReader(MultiFileReaderData &reader_data, + const MultiFileBindData &bind_data, + const vector &global_columns, + const vector &global_column_ids, + optional_ptr table_filters, + ClientContext &context, + MultiFileGlobalState &gstate) override; +}; + +struct VortexReaderInterface final : MultiFileReaderInterface { + static unique_ptr CreateInterface(ClientContext &) { + return make_uniq(); + } + + inline unique_ptr InitializeOptions(ClientContext &, + optional_ptr) override { + return make_uniq(); + } + + inline bool ParseCopyOption(ClientContext &, + const string &, + const vector &, + BaseFileReaderOptions &, + vector &, + vector &) override { + return false; + }; + + inline bool ParseOption(ClientContext &, + const string &, + const Value &, + MultiFileOptions &, + BaseFileReaderOptions &) override { + return false; + } + + inline unique_ptr InitializeBindData(MultiFileBindData &, + unique_ptr) override { + return make_uniq(); + } + + // Open first file, populate types and names from it + void BindReader(ClientContext &context, + vector &types, + vector &names, + MultiFileBindData &bind_data) override; + + unique_ptr InitializeGlobalState(ClientContext &context, + MultiFileBindData &bind_data, + MultiFileGlobalState &global_state) override; + + unique_ptr InitializeLocalState(ExecutionContext &context, + GlobalTableFunctionState &global_state) override; + + inline shared_ptr CreateReader(ClientContext &, + GlobalTableFunctionState &, + BaseUnionData &, + const MultiFileBindData &) override { + throw BinderException("UNION BY NAME for Vortex files is not supported"); + } + + shared_ptr CreateReader(ClientContext &context, + GlobalTableFunctionState &gstate, + const OpenFileInfo &file, + idx_t file_idx, + const MultiFileBindData &bind_data) override; + + shared_ptr CreateReader(ClientContext &context, + const OpenFileInfo &file, + BaseFileReaderOptions &options, + const MultiFileOptions &file_options) override; + + unique_ptr GetCardinality(const MultiFileBindData &bind_data, idx_t file_count) override; + + inline FileGlobInput GetGlobInput() override { + return {FileGlobOptions::FALLBACK_GLOB, "vortex"}; + } + + inline unique_ptr Copy() override { + return make_uniq(); + } + + void GetVirtualColumns(ClientContext &, MultiFileBindData &, virtual_column_map_t &result) override; + bool FinalizeScan(ClientContext &, GlobalTableFunctionState &gstate, DataChunk &output) override; +}; + +struct VortexBaseReader final : BaseFileReader { + VortexBaseReader(OpenFileInfo file, unique_ptr ffi_file) + : BaseFileReader(file), ffi_file(std::move(ffi_file)) { + } + + unique_ptr ffi_file; + vector virtual_ids; + + inline void AddVirtualColumn(column_t id) override { + virtual_ids.push_back(id); + } + + /* + * Called by all threads on current file under global lock. Once + * TryInitializeScan returns false, first thread to receive it advances + * to next file and calls TryInitializeScan on it. + */ + bool TryInitializeScan(ClientContext &, + GlobalTableFunctionState &, + LocalTableFunctionState &local_state) override; + + AsyncResult Scan(ClientContext &, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state, + DataChunk &chunk) override; + + inline void FinishFile(ClientContext &, GlobalTableFunctionState &) override { + } + + double GetProgressInFile(ClientContext &) override; + + unique_ptr GetStatistics(ClientContext &context, const string &name) override; + + inline string GetReaderType() const override { + return "Vortex"; + } +}; diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index 65a1a3f1dd3..bdcf2da311f 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -10,16 +10,10 @@ extern "C" { #endif -// Info passed into the bind callback. The callback should set error or else add result columns. -typedef struct duckdb_vx_tfunc_bind_input_ *duckdb_vx_tfunc_bind_input; -typedef struct duckdb_vx_tfunc_bind_result_ *duckdb_vx_tfunc_bind_result; - -// Fetch a parameter from the bind info. -// The caller is responsible for freeing the value using duckdb_value_free. -duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, size_t index); +typedef struct duckdb_bind_result_ *duckdb_bind_result; // Add a result column to the bind info. -void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result, +void duckdb_vx_tfunc_bind_result_add_column(duckdb_bind_result ffi_result, const char *name_str, size_t name_len, duckdb_logical_type ffi_type); @@ -31,29 +25,8 @@ void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, cons // Input data passed into the init_global and init_local callbacks. typedef struct { const void *bind_data; - - /** - * Projected columns that are requested to be read. These are not - * all columns, only the ones DuckDB optimizer thinks we should read. - */ idx_t *column_ids; size_t column_ids_count; - - /** - * Post filter projected columns. Our table function implements filter - * pushdown so this list is a subset of columns referenced in column_ids - * after filter pushdown and filter pruning. May be empty, in which case - * column_ids should be used. - * Indices in this list reference values from column_ids. I.e. if - * column_ids=[1,5,6], projection_ids=[1], output column should be - * column_ids[1] = 5 - * - * Example usage: - * https://github.com/duckdb/duckdb/blob/dc11eadd8f0a7c600f0034810706605ebe10d5b9/src/include/duckdb/function/table_function.hpp#L147 - */ - const idx_t *projection_ids; - size_t projection_ids_count; - duckdb_vx_table_filter_set filters; duckdb_client_context client_context; } duckdb_vx_tfunc_init_input; @@ -74,18 +47,10 @@ typedef struct { // set only for strings uint64_t max_string_length; bool has_null; + // non-owned column type + duckdb_logical_type type; } duckdb_column_statistics; -const idx_t INVALID_IDX = UINT64_MAX; - -typedef struct { - idx_t partition_index; - // Either INVALID_IDX or position of column in output for file_index column - size_t file_index_column_pos; - // File index for the exported partition. - size_t file_index; -} duckdb_vx_partition_data; - duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db); typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input; diff --git a/vortex-duckdb/cpp/include/table_function.hpp b/vortex-duckdb/cpp/include/table_function.hpp index b5e1947ef60..54daff90d06 100644 --- a/vortex-duckdb/cpp/include/table_function.hpp +++ b/vortex-duckdb/cpp/include/table_function.hpp @@ -3,7 +3,6 @@ #pragma once -#include "data.hpp" #include "duckdb.h" #include "duckdb/function/function.hpp" #include "duckdb/function/table_function.hpp" @@ -12,11 +11,7 @@ using namespace duckdb; static_assert(sizeof(idx_t) == 8); -// We need this exposed to compare function addresses in optimizer.cpp -unique_ptr duckdb_vx_table_function_bind(ClientContext &context, - TableFunctionBindInput &input, - vector &return_types, - vector &names); +bool is_vortex_scan(const TableFunction &function); struct TableFunctionProjectionExpressionInput { const LogicalGet &get; @@ -35,36 +30,3 @@ struct TableFunctionUngroupedAggregateInput { }; bool aggregate_pushdown(ClientContext &context, const TableFunctionUngroupedAggregateInput &input); - -struct VortexBindData final : FunctionData { - VortexBindData(unique_ptr ffi_data, const vector &types) - : ffi_data(std::move(ffi_data)), types(types) { - } - unique_ptr Copy() const override; - bool Equals(const FunctionData &other) const override; - - unique_ptr ffi_data; - vector types; -}; - -struct VortexGlobalData final : GlobalTableFunctionState { - explicit VortexGlobalData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { - } - - idx_t MaxThreads() const override { - return GlobalTableFunctionState::MAX_THREADS; - } - - unique_ptr ffi_data; -}; - -struct VortexLocalData final : LocalTableFunctionState { - explicit VortexLocalData(unique_ptr ffi_data) : ffi_data(std::move(ffi_data)) { - } - unique_ptr ffi_data; -}; - -struct VortexBindResults { - vector &return_types; - vector &names; -}; diff --git a/vortex-duckdb/cpp/multi_file_reader.cpp b/vortex-duckdb/cpp/multi_file_reader.cpp new file mode 100644 index 00000000000..8325aa9e2ff --- /dev/null +++ b/vortex-duckdb/cpp/multi_file_reader.cpp @@ -0,0 +1,320 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "multi_file_reader.hpp" +#include "error.hpp" +#include "table_function.h" +#include "vortex_duckdb.h" +#include "vortex.h" + +// TODO (myrrc) remove NDEBUG in release builds + +unique_ptr VortexBindData::Copy() const { + auto result = make_uniq(); + if (ffi_bind_data) { + const duckdb_vx_data copy = duckdb_table_function_bind_data_clone(ffi_bind_data->DataPtr()); + result->ffi_bind_data = unique_ptr(reinterpret_cast(copy)); + } + return result; +} + +bool VortexBindData::Equals(const FunctionData &other_base) const { + const VortexBindData &other = other_base.Cast(); + return ffi_bind_data.get() == other.ffi_bind_data.get(); +} + +ReaderInitializeType +VortexMultiFileReader::InitializeReader(MultiFileReaderData &reader_data, + const MultiFileBindData &bind_data, + const vector &global_columns, + const vector &global_column_ids, + optional_ptr table_filters, + ClientContext &context, + MultiFileGlobalState &gstate) { + D_ASSERT(reader_data.reader != nullptr); + D_ASSERT(gstate.global_state != nullptr); + + VortexBaseReader &reader = reader_data.reader->Cast(); + + reader.columns = global_columns; // base InitializeReader requires columns to be set + const ReaderInitializeType base_skip = MultiFileReader::InitializeReader(reader_data, + bind_data, + global_columns, + global_column_ids, + table_filters, + context, + gstate); + if (base_skip == ReaderInitializeType::SKIP_READING_FILE) { + return base_skip; + } + + const VortexGlobalState &global = gstate.global_state->Cast(); + + duckdb_vx_error error = nullptr; + const void *const ffi_global = global.ffi_global_state->DataPtr(); + void *const ffi_file = reader.ffi_file->DataPtr(); + const bool skip = duckdb_reader_initialize(ffi_global, ffi_file, &error); + if (error) { + throw InvalidInputException(IntoErrString(error)); + } + + return skip ? ReaderInitializeType::SKIP_READING_FILE : ReaderInitializeType::INITIALIZED; +} + +void VortexReaderInterface::BindReader(ClientContext &context, + vector &types, + vector &names, + MultiFileBindData &bind_data) { + BaseFileReaderOptions options; + MultiFileOptions file_options; + VortexBindResult result = {types, names}; + + VortexBindData &bind = bind_data.bind_data->Cast(); + const OpenFileInfo first_file = bind_data.file_list->GetFirstFile(); + bind_data.initial_reader = CreateReader(context, first_file, options, file_options); + const VortexBaseReader &initial_reader = bind_data.initial_reader->Cast(); + + duckdb_vx_error error = nullptr; + const void *const ffi_file = initial_reader.ffi_file->DataPtr(); + duckdb_bind_result ffi_result = reinterpret_cast(&result); + + duckdb_vx_data ffi_bind_data = duckdb_reader_bind(ffi_file, ffi_result, &error); + if (error) { + throw BinderException(IntoErrString(error)); + } + + bind.ffi_bind_data = unique_ptr(reinterpret_cast(ffi_bind_data)); +} + +unique_ptr +VortexReaderInterface::InitializeGlobalState(ClientContext &context, + MultiFileBindData &bind_data, + MultiFileGlobalState &input) { + const VortexBindData &bind = bind_data.bind_data->Cast(); + + vector column_ids(input.column_indexes.size()); + for (size_t i = 0; i < input.column_indexes.size(); ++i) { + column_ids[i] = input.column_indexes[i].GetPrimaryIndex(); + } + + void *const ffi_bind = bind.ffi_bind_data->DataPtr(); + duckdb_vx_tfunc_init_input ffi_input = { + .bind_data = ffi_bind, + .column_ids = column_ids.data(), + .column_ids_count = column_ids.size(), + .filters = reinterpret_cast(input.filters.get()), + .client_context = reinterpret_cast(&context), + }; + + duckdb_vx_error error_out = nullptr; + duckdb_vx_data ffi_global_state = duckdb_table_function_init_global(&ffi_input, &error_out); + if (error_out) { + throw BinderException(IntoErrString(error_out)); + } + + auto result = make_uniq(); + result->ffi_bind_data = ffi_bind; + result->ffi_global_state = unique_ptr(reinterpret_cast(ffi_global_state)); + return result; +} + +unique_ptr +VortexReaderInterface::InitializeLocalState(ExecutionContext &, GlobalTableFunctionState &global_state) { + auto &global = global_state.Cast(); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + duckdb_vx_data ffi_local_state = duckdb_table_function_init_local(global.ffi_bind_data, ffi_global); + + auto result = make_uniq(); + result->ffi_local_state = unique_ptr(reinterpret_cast(ffi_local_state)); + return result; +} + +static shared_ptr OpenReader(const OpenFileInfo &file) { + duckdb_vx_error error = nullptr; + + const char *const ffi_file_path = file.path.c_str(); + const size_t ffi_file_size = file.path.size(); + + duckdb_vx_data ffi_file = duckdb_reader_open(ffi_file_path, ffi_file_size, &error); + if (error) { + throw IOException(IntoErrString(error)); + } + + auto cdata = unique_ptr(reinterpret_cast(ffi_file)); + return make_shared_ptr(file, std::move(cdata)); +} + +shared_ptr VortexReaderInterface::CreateReader(ClientContext &, + GlobalTableFunctionState &, + const OpenFileInfo &file, + idx_t, + const MultiFileBindData &) { + return OpenReader(file); +} + +shared_ptr VortexReaderInterface::CreateReader(ClientContext &, + const OpenFileInfo &file, + BaseFileReaderOptions &, + const MultiFileOptions &) { + return OpenReader(file); +} + +unique_ptr VortexReaderInterface::GetCardinality(const MultiFileBindData &data, + idx_t file_count) { + const VortexBindData &bind_data = data.bind_data->Cast(); + const void *const ffi_bind = bind_data.ffi_bind_data->DataPtr(); + + duckdb_vx_node_statistics stats = {}; + duckdb_table_function_cardinality(ffi_bind, file_count, &stats); + + auto out = make_uniq(); + out->has_estimated_cardinality = stats.has_estimated_cardinality; + out->estimated_cardinality = stats.estimated_cardinality; + out->has_max_cardinality = stats.has_max_cardinality; + out->max_cardinality = stats.max_cardinality; + return out; +} + +bool VortexBaseReader::TryInitializeScan(ClientContext &, + GlobalTableFunctionState &, + LocalTableFunctionState &local_state) { + VortexLocalState &local = local_state.Cast(); + + void *const ffi_local = local.ffi_local_state->DataPtr(); + void *const ffi_file_ptr = ffi_file->DataPtr(); + return duckdb_reader_try_initialize_scan(ffi_local, ffi_file_ptr); +} + +AsyncResult VortexBaseReader::Scan(ClientContext &, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state, + DataChunk &chunk) { + VortexGlobalState &global = global_state.Cast(); + VortexLocalState &local = local_state.Cast(); + + duckdb_vx_error error = nullptr; + duckdb_data_chunk ffi_chunk = reinterpret_cast(&chunk); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + void *const ffi_local = local.ffi_local_state->DataPtr(); + const void *const ffi_file_ptr = ffi_file->DataPtr(); + const bool has_more_data = duckdb_reader_scan(ffi_file_ptr, ffi_global, ffi_local, ffi_chunk, &error); + if (error) { + throw InvalidInputException(IntoErrString(error)); + } + return has_more_data ? SourceResultType::HAVE_MORE_OUTPUT : SourceResultType::FINISHED; +} + +void VortexReaderInterface::GetVirtualColumns(ClientContext &, + MultiFileBindData &, + virtual_column_map_t &result) { + // "filename", "file_index" and "empty" come from MultiFileReader + result.insert( + {MultiFileReader::COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::UBIGINT}}); +} + +bool VortexReaderInterface::FinalizeScan(ClientContext &, + GlobalTableFunctionState &global_state, + DataChunk &output) { + const VortexGlobalState &global = global_state.Cast(); + + duckdb_vx_error error = nullptr; + duckdb_data_chunk ffi_chunk = reinterpret_cast(&output); + const void *const ffi_global = global.ffi_global_state->DataPtr(); + const bool filled = duckdb_reader_finalize_scan(ffi_global, ffi_chunk, &error); + if (error) { + throw InvalidInputException(IntoErrString(error)); + } + return filled; +} + +static Value &UnwrapValue(duckdb_value value) { + return *(reinterpret_cast(value)); +} + +static unique_ptr numeric_stats(duckdb_column_statistics &stats, LogicalType type) { + BaseStatistics out = NumericStats::CreateUnknown(type); + if (stats.min) { + NumericStats::SetMin(out, UnwrapValue(stats.min)); + duckdb_destroy_value(&stats.min); + } + if (stats.max) { + NumericStats::SetMax(out, UnwrapValue(stats.max)); + duckdb_destroy_value(&stats.max); + } + if (!stats.has_null) { + out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); + } + return out.ToUnique(); +} + +static unique_ptr string_stats(duckdb_column_statistics &stats, LogicalType type) { + BaseStatistics out = StringStats::CreateUnknown(type); + if (stats.min) { + StringStats::SetMin(out, StringValue::Get(UnwrapValue(stats.min))); + duckdb_destroy_value(&stats.min); + } + if (stats.max) { + StringStats::SetMax(out, StringValue::Get(UnwrapValue(stats.max))); + duckdb_destroy_value(&stats.max); + } + if (stats.max_string_length >> 63) { + StringStats::SetMaxStringLength(out, uint32_t(stats.max_string_length)); + } + if (!stats.has_null) { + out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); + } + + return out.ToUnique(); +} + +static unique_ptr base_stats(duckdb_column_statistics &stats, LogicalType type) { + BaseStatistics out = BaseStatistics::CreateUnknown(type); + if (!stats.has_null) { + out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); + } + return out.ToUnique(); +} + +unique_ptr VortexBaseReader::GetStatistics(ClientContext &, const string &name) { + duckdb_column_statistics statistics = {}; + if (!duckdb_reader_get_statistics(ffi_file->DataPtr(), name.c_str(), name.size(), &statistics)) { + return {}; + } + + using enum LogicalTypeId; + const unique_ptr type(reinterpret_cast(statistics.type)); + switch (type->id()) { + case BOOLEAN: + case TINYINT: + case SMALLINT: + case INTEGER: + case BIGINT: + case FLOAT: + case DOUBLE: + case UTINYINT: + case USMALLINT: + case UINTEGER: + case UBIGINT: + case UHUGEINT: + case HUGEINT: { + return numeric_stats(statistics, *type); + } + case VARCHAR: + case BLOB: { + return string_stats(statistics, *type); + } + case STRUCT: { + // TODO(myrrc) + // Duckdb's has_null has a different semantics for structs. + // If we propagate our has_null, this breaks Duckdb optimizer. + // You can reproduce it in struct.slt test in vortex-sqllogictests: + return {}; + } + default: + return base_stats(statistics, *type); + } +} + +double VortexBaseReader::GetProgressInFile(ClientContext &) { + return duckdb_reader_get_progress_in_file(ffi_file->DataPtr()); +} diff --git a/vortex-duckdb/cpp/optimizer.cpp b/vortex-duckdb/cpp/optimizer.cpp index b6ba2bbf23b..9f799fb1e42 100644 --- a/vortex-duckdb/cpp/optimizer.cpp +++ b/vortex-duckdb/cpp/optimizer.cpp @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "optimizer.hpp" +#include "multi_file_reader.hpp" #include "table_function.hpp" #include "duckdb/planner/expression/bound_columnref_expression.hpp" @@ -12,7 +13,7 @@ void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections using enum LogicalOperatorType; switch (op.type) { case LOGICAL_GET: { - if (auto &get = op.Cast(); get.function.bind == duckdb_vx_table_function_bind) { + if (auto &get = op.Cast(); is_vortex_scan(get.function)) { analyses.emplace(get.table_index, GetAnalysis {get, {}}); } break; @@ -37,7 +38,7 @@ void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections get = &child.Cast(); } - if (get != nullptr && get->function.bind == duckdb_vx_table_function_bind) { + if (get != nullptr && is_vortex_scan(get->function)) { projections.emplace(projection.table_index, projection); } break; diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 3073bef9e88..0a823ec3ecd 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -3,11 +3,12 @@ #include "data.hpp" #include "error.hpp" -#include "table_function.hpp" #include "expr.h" -#include "vortex_duckdb.h" +#include "multi_file_reader.hpp" #include "table_function.h" +#include "table_function.hpp" #include "vortex.h" +#include "vortex_duckdb.h" #include "duckdb.h" #include "duckdb/catalog/catalog.hpp" @@ -23,19 +24,6 @@ using namespace std::string_literals; constexpr column_t COLUMN_IDENTIFIER_FILE_INDEX = MultiFileReader::COLUMN_IDENTIFIER_FILE_INDEX; constexpr column_t COLUMN_IDENTIFIER_FILE_ROW_NUMBER = MultiFileReader::COLUMN_IDENTIFIER_FILE_ROW_NUMBER; -unique_ptr VortexBindData::Copy() const { - const auto copied_ffi_data = duckdb_table_function_bind_data_clone(ffi_data->DataPtr()); - auto ffi_data_p = unique_ptr(reinterpret_cast(copied_ffi_data)); - return make_uniq(std::move(ffi_data_p), types); -} - -bool VortexBindData::Equals(const FunctionData &other_base) const { - const VortexBindData &other = other_base.Cast(); - // if "types" are different, "ffi_data" would also be different as it - // contains types inside, so omit "types" from comparison. - return ffi_data.get() == other.ffi_data.get(); -} - // This is a flaw of Duckdb API which doesn't allow passing non-const // expressions. We never modify the value on Rust side. static duckdb_vx_expr get_ffi_expr(const Expression &expr) { @@ -43,116 +31,7 @@ static duckdb_vx_expr get_ffi_expr(const Expression &expr) { } static void *get_ffi_bind(const FunctionData *bind_data) { - return bind_data->Cast().ffi_data->DataPtr(); -} - -static void *get_ffi_global(GlobalTableFunctionState *state) { - return state->Cast().ffi_data->DataPtr(); -} - -static void *get_ffi_local(LocalTableFunctionState *state) { - return state->Cast().ffi_data->DataPtr(); -} - -double -table_scan_progress(ClientContext &, const FunctionData *, const GlobalTableFunctionState *global_state) { - void *const c_global_state = global_state->Cast().ffi_data->DataPtr(); - return duckdb_table_function_scan_progress(c_global_state); -} - -static Value &UnwrapValue(duckdb_value value) { - return *(reinterpret_cast(value)); -} - -unique_ptr numeric_stats(duckdb_column_statistics &stats, LogicalType type) { - BaseStatistics out = StringStats::CreateUnknown(type); - if (stats.min) { - NumericStats::SetMin(out, UnwrapValue(stats.min)); - duckdb_destroy_value(&stats.min); - } - if (stats.max) { - NumericStats::SetMax(out, UnwrapValue(stats.max)); - duckdb_destroy_value(&stats.max); - } - if (!stats.has_null) { - out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); - } - return out.ToUnique(); -} - -unique_ptr string_stats(duckdb_column_statistics &stats, LogicalType type) { - BaseStatistics out = StringStats::CreateUnknown(type); - if (stats.min) { - StringStats::SetMin(out, StringValue::Get(UnwrapValue(stats.min))); - duckdb_destroy_value(&stats.min); - } - if (stats.max) { - StringStats::SetMax(out, StringValue::Get(UnwrapValue(stats.max))); - duckdb_destroy_value(&stats.max); - } - if (stats.max_string_length >> 63) { - StringStats::SetMaxStringLength(out, uint32_t(stats.max_string_length)); - } - if (!stats.has_null) { - out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); - } - - return out.ToUnique(); -} - -unique_ptr base_stats(duckdb_column_statistics &stats, LogicalType type) { - BaseStatistics out = StringStats::CreateUnknown(type); - if (!stats.has_null) { - out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES); - } - return out.ToUnique(); -} - -unique_ptr statistics(ClientContext &, const FunctionData *bind_data, column_t column_index) { - if (IsVirtualColumn(column_index)) { - return {}; - } - - const auto &bind = bind_data->Cast(); - const void *const ffi_bind = get_ffi_bind(bind_data); - - duckdb_column_statistics statistics = {}; - if (!duckdb_table_function_statistics(ffi_bind, column_index, &statistics)) { - return {}; - } - - const LogicalType type = bind.types[column_index]; - - switch (type.id()) { - case LogicalTypeId::BOOLEAN: - case LogicalTypeId::TINYINT: - case LogicalTypeId::SMALLINT: - case LogicalTypeId::INTEGER: - case LogicalTypeId::BIGINT: - case LogicalTypeId::FLOAT: - case LogicalTypeId::DOUBLE: - case LogicalTypeId::UTINYINT: - case LogicalTypeId::USMALLINT: - case LogicalTypeId::UINTEGER: - case LogicalTypeId::UBIGINT: - case LogicalTypeId::UHUGEINT: - case LogicalTypeId::HUGEINT: { - return numeric_stats(statistics, type); - } - case LogicalTypeId::VARCHAR: - case LogicalTypeId::BLOB: { - return string_stats(statistics, type); - } - case LogicalTypeId::STRUCT: { - // TODO(myrrc) - // Duckdb's has_null has a different semantics for structs. - // If we propagate our has_null, this breaks Duckdb optimizer. - // You can reproduce it in struct.slt test in vortex-sqllogictests: - return {}; - } - default: - return base_stats(statistics, type); - } + return bind_data->Cast().bind_data->Cast().ffi_bind_data->DataPtr(); } bool projection_expression_pushdown(ClientContext &, const TableFunctionProjectionExpressionInput &input) { @@ -197,69 +76,6 @@ bool aggregate_pushdown(ClientContext &, const TableFunctionUngroupedAggregateIn return res; } -unique_ptr duckdb_vx_table_function_bind(ClientContext &, - TableFunctionBindInput &input, - vector &return_types, - vector &names) { - VortexBindResults result = {return_types, names}; - - duckdb_vx_error error_out = nullptr; - duckdb_vx_tfunc_bind_input bind_input = reinterpret_cast(&input); - duckdb_vx_tfunc_bind_result bind_result = reinterpret_cast(&result); - duckdb_vx_data ffi_bind_data = duckdb_table_function_bind(bind_input, bind_result, &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } - - auto cdata = unique_ptr(reinterpret_cast(ffi_bind_data)); - return make_uniq(std::move(cdata), return_types); -} - -unique_ptr init_global(ClientContext &context, TableFunctionInitInput &input) { - const void *const ffi_bind = get_ffi_bind(input.bind_data.get()); - - duckdb_vx_tfunc_init_input ffi_input = { - .bind_data = ffi_bind, - .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; - duckdb_vx_data ffi_global_data = duckdb_table_function_init_global(&ffi_input, &error_out); - if (error_out) { - throw BinderException(IntoErrString(error_out)); - } - - auto cdata = unique_ptr(reinterpret_cast(ffi_global_data)); - return make_uniq(std::move(cdata)); -} - -unique_ptr -init_local(ExecutionContext &, TableFunctionInitInput &input, GlobalTableFunctionState *global_state) { - const void *const ffi_bind = get_ffi_bind(input.bind_data.get()); - void *const ffi_global = get_ffi_global(global_state); - - duckdb_vx_data ffi_local_data = duckdb_table_function_init_local(ffi_bind, ffi_global); - auto cdata = unique_ptr(reinterpret_cast(ffi_local_data)); - return make_uniq(std::move(cdata)); -} - -void function(ClientContext &, TableFunctionInput &input, DataChunk &output) { - void *const ffi_global = get_ffi_global(input.global_state.get()); - void *const ffi_local = get_ffi_local(input.local_state.get()); - - duckdb_data_chunk chunk = reinterpret_cast(&output); - duckdb_vx_error error_out = nullptr; - duckdb_table_function_scan(ffi_global, ffi_local, chunk, &error_out); - if (error_out) { - throw InvalidInputException(IntoErrString(error_out)); - } -} - using FilterVec = vector>; void pushdown_complex_filter(const FunctionData &bind_data, FilterVec &filters) { @@ -277,80 +93,20 @@ void pushdown_complex_filter(const FunctionData &bind_data, FilterVec &filters) } } -unique_ptr cardinality(ClientContext &, const FunctionData *bind_data) { - const void *const ffi_bind = get_ffi_bind(bind_data); - - duckdb_vx_node_statistics stats = {}; - duckdb_table_function_cardinality(ffi_bind, &stats); - - auto out = make_uniq(); - out->has_estimated_cardinality = stats.has_estimated_cardinality; - out->estimated_cardinality = stats.estimated_cardinality; - out->has_max_cardinality = stats.has_max_cardinality; - out->max_cardinality = stats.max_cardinality; - - return out; -} - -extern "C" duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, - size_t index) { - D_ASSERT(ffi_input); - const TableFunctionBindInput &input = *reinterpret_cast(ffi_input); - return reinterpret_cast(new Value(input.inputs[index])); -} - -extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result, +extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_bind_result ffi_result, const char *name_str, size_t name_len, duckdb_logical_type ffi_type) { D_ASSERT(ffi_result); D_ASSERT(name_str); D_ASSERT(ffi_type); - const VortexBindResults &result = *reinterpret_cast(ffi_result); + VortexBindResult &result = *reinterpret_cast(ffi_result); const LogicalType logical_type = *reinterpret_cast(ffi_type); result.names.emplace_back(name_str, name_len); result.return_types.emplace_back(logical_type); } -/** - * Called at planning time to determine whether data is partitioned by a - * given set of columns. Requested columns are GROUP BY parameters i.e. columns - * over which the query aggregates. - */ -TablePartitionInfo get_partition_info(ClientContext &, TableFunctionPartitionInput &input) { - const vector &ids = input.partition_ids; - // Our data is partitioned by array exporters. Each exporter processes a - // single Array which belongs to a single file. If data is partitioned only - // by file_index, there is one unique value for an Array. Otherwise there - // may be multiple values. - return (ids.size() == 1 && ids[0] == COLUMN_IDENTIFIER_FILE_INDEX) - ? TablePartitionInfo::SINGLE_VALUE_PARTITIONS - : TablePartitionInfo::NOT_PARTITIONED; -} - -OperatorPartitionData get_partition_data(ClientContext &, TableFunctionGetPartitionInput &input) { - void *const ffi_global = get_ffi_global(input.global_state.get()); - void *const ffi_local = get_ffi_local(input.local_state.get()); - duckdb_vx_partition_data partition_data; - duckdb_table_function_get_partition_data(ffi_global, ffi_local, &partition_data); - - OperatorPartitionData out(partition_data.partition_index); - - // file_index_column_pos may be INVALID_IDX, but column_index will never - // be INVALID_IDX, so we can compare directly - for (const column_t column_index : input.partition_info.partition_columns) { - if (column_index == partition_data.file_index_column_pos) { - out.partition_data.emplace_back(Value::UBIGINT(partition_data.file_index)); - } else { - throw InternalException(StringUtil::Format( - "get_partition_data: requested column_index %d is not constant for given partition", - column_index)); - } - } - return out; -} - extern "C" void duckdb_vx_string_map_insert(duckdb_vx_string_map map, const char *key, const char *value) { D_ASSERT(map); D_ASSERT(key); @@ -366,52 +122,45 @@ InsertionOrderPreservingMap to_string(TableFunctionToStringInput &input) return result; } +bool is_vortex_scan(const TableFunction &function) { + return function.bind == MultiFileFunction::MultiFileBind; +} + +unique_ptr get_multi_file_reader(const TableFunction &) { + return make_uniq(); +} + duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter, const std::string &name) { - TableFunction tf(name, {}, function, duckdb_vx_table_function_bind, init_global, init_local); + MultiFileFunction fn(name); + fn.arguments[0] = parameter; + // We neither support UNION BY NAME nor hive partitioning as for now + fn.named_parameters = {}; - tf.projection_pushdown = true; - tf.filter_pushdown = true; - tf.filter_prune = true; - tf.sampling_pushdown = false; + fn.filter_pushdown = true; + fn.filter_prune = true; - tf.pushdown_expression = [](auto &, const auto &, Expression &expression) { + fn.pushdown_expression = [](auto &, const auto &, Expression &expression) { return duckdb_table_function_pushdown_expression(reinterpret_cast(&expression)); }; - tf.pushdown_complex_filter = [](auto &, auto &, FunctionData *bind_data, FilterVec &filters) { + fn.pushdown_complex_filter = [](auto &, auto &, FunctionData *bind_data, FilterVec &filters) { pushdown_complex_filter(*bind_data, filters); }; - tf.cardinality = cardinality; - tf.get_partition_info = get_partition_info; - tf.get_partition_data = get_partition_data; - tf.to_string = to_string; - tf.table_scan_progress = table_scan_progress; - tf.statistics = statistics; + fn.to_string = to_string; - tf.late_materialization = true; + fn.late_materialization = true; // Columns that uniquely identify a row for deferred re-fetch in a multi // file scan: (file index, row number in file). - tf.get_row_id_columns = [](auto &, auto) -> vector { + fn.get_row_id_columns = [](auto &, auto) -> vector { return {COLUMN_IDENTIFIER_FILE_INDEX, COLUMN_IDENTIFIER_FILE_ROW_NUMBER}; }; - tf.get_virtual_columns = [](auto &, auto) -> virtual_column_map_t { - return { - {COLUMN_IDENTIFIER_EMPTY, {"", LogicalTypeId::BOOLEAN}}, - {COLUMN_IDENTIFIER_FILE_INDEX, {"file_index", LogicalType::UBIGINT}}, - // MultiFileReader's file_row_number column is BIGINT. - // row_idx() is UBIGINT. Use UBIGINT since there's no difference to - // Duckdb what to compare. - {COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::UBIGINT}}, - }; - }; - - tf.arguments.resize(1); - tf.arguments[0] = parameter; + fn.statistics = MultiFileFunction::MultiFileScanStats; + fn.get_multi_file_reader = get_multi_file_reader; try { auto &system_catalog = Catalog::GetSystemCatalog(db); auto data = CatalogTransaction::GetSystemTransaction(db); - CreateTableFunctionInfo tf_info(tf); + CreateTableFunctionInfo tf_info(fn); tf_info.on_conflict = OnCreateConflict::ALTER_ON_CONFLICT; system_catalog.CreateFunction(data, tf_info); } catch (const std::exception &e) { diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index 96f5804d3b4..47aa243a1fe 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -17,62 +17,72 @@ extern "C" { #endif // __cplusplus -extern void duckdb_table_function_to_string(const void *bind_data, duckdb_vx_string_map map); +extern void duckdb_table_function_to_string(const void *bind, duckdb_vx_string_map map); extern -bool duckdb_table_function_statistics(const void *bind_data, - size_t column_index, - duckdb_column_statistics *stats_out); - -extern double duckdb_table_function_scan_progress(void *global_state); - -extern -void duckdb_table_function_get_partition_data(void *global_init_data, - void *local_init_data, - duckdb_vx_partition_data *partition_data_out); - -extern -bool duckdb_table_function_pushdown_complex_filter(void *bind_data, +bool duckdb_table_function_pushdown_complex_filter(void *bind, duckdb_vx_expr expr, - duckdb_vx_error *error_out); + duckdb_vx_error *error); extern -bool duckdb_table_function_pushdown_projection_expression(void *bind_data, +bool duckdb_table_function_pushdown_projection_expression(void *bind, duckdb_vx_expr expr, size_t column_id, - duckdb_vx_error *error_out); + duckdb_vx_error *error); extern -bool duckdb_table_function_pushdown_projection_aggregates(void *bind_data, +bool duckdb_table_function_pushdown_projection_aggregates(void *bind, duckdb_vx_agg_input input, - duckdb_vx_error *error_out); - -extern -void duckdb_table_function_scan(void *global_init_data, - void *local_init_data, - duckdb_data_chunk output, - duckdb_vx_error *error_out); + duckdb_vx_error *error); extern bool duckdb_table_function_pushdown_expression(duckdb_vx_expr expr); extern -void duckdb_table_function_cardinality(const void *bind_data, - duckdb_vx_node_statistics *node_stats_out); +void duckdb_table_function_cardinality(const void *bind, + uint64_t file_count, + duckdb_vx_node_statistics *stats); extern duckdb_vx_data duckdb_table_function_init_global(const duckdb_vx_tfunc_init_input *init_input, - duckdb_vx_error *error_out); + duckdb_vx_error *error); + +extern duckdb_vx_data duckdb_table_function_init_local(const void *bind, const void *global); extern -duckdb_vx_data duckdb_table_function_init_local(const void *bind_data, - void *global_init_data); +duckdb_vx_data duckdb_reader_bind(const void *first_file, + duckdb_bind_result result, + duckdb_vx_error *error_out); + +extern +duckdb_vx_data duckdb_reader_open(const char *file_path, + size_t file_path_len, + duckdb_vx_error *error); + +extern +bool duckdb_reader_get_statistics(const void *file, + const char *column_name, + size_t column_name_len, + duckdb_column_statistics *stats_out); + +extern bool duckdb_reader_initialize(const void *global, void *file, duckdb_vx_error *error); + +extern bool duckdb_reader_try_initialize_scan(void *local, void *file); + +extern +bool duckdb_reader_scan(const void *file, + const void *global, + void *local, + duckdb_data_chunk chunk, + duckdb_vx_error *error); + +extern double duckdb_reader_get_progress_in_file(const void *file); extern -duckdb_vx_data duckdb_table_function_bind(duckdb_vx_tfunc_bind_input bind_input, - duckdb_vx_tfunc_bind_result bind_result, - duckdb_vx_error *error_out); +bool duckdb_reader_finalize_scan(const void *global, + duckdb_data_chunk chunk, + duckdb_vx_error *error); -extern duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind_data); +extern duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind); extern duckdb_vx_data duckdb_copy_function_copy_to_bind(const char *const *column_names, diff --git a/vortex-duckdb/patches/duckdb-mfr-finalize-scan.diff b/vortex-duckdb/patches/duckdb-mfr-finalize-scan.diff new file mode 100644 index 00000000000..d7214c7493f --- /dev/null +++ b/vortex-duckdb/patches/duckdb-mfr-finalize-scan.diff @@ -0,0 +1,69 @@ +--- a/src/include/duckdb/common/multi_file/multi_file_function.hpp ++++ b/src/include/duckdb/common/multi_file/multi_file_function.hpp +@@ -63,6 +63,9 @@ + virtual void GetVirtualColumns(ClientContext &context, MultiFileBindData &bind_data, virtual_column_map_t &result); + virtual unique_ptr Copy(); + virtual FileGlobInput GetGlobInput(); ++ virtual bool FinalizeScan(ClientContext &context, GlobalTableFunctionState &global_state, DataChunk &output) { ++ return false; ++ } + }; + + template +@@ -590,18 +593,30 @@ + + static OperatorPartitionData MultiFileGetPartitionData(ClientContext &context, + TableFunctionGetPartitionInput &input) { ++ if (!input.local_state) { ++ return OperatorPartitionData(0); ++ } + auto &bind_data = input.bind_data->CastNoConst(); + auto &data = input.local_state->Cast(); + auto &gstate = input.global_state->Cast(); + OperatorPartitionData partition_data(data.batch_index); +- bind_data.multi_file_reader->GetPartitionData(context, bind_data.reader_bind, *data.reader_data, +- gstate.multi_file_reader_state, input.partition_info, +- partition_data); ++ if (data.reader_data) { ++ bind_data.multi_file_reader->GetPartitionData(context, bind_data.reader_bind, *data.reader_data, ++ gstate.multi_file_reader_state, input.partition_info, ++ partition_data); ++ } + return partition_data; + } + + static void MultiFileScan(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) { + if (!data_p.local_state) { ++ auto &finalize_gstate = data_p.global_state->Cast(); ++ auto &finalize_bind_data = data_p.bind_data->CastNoConst(); ++ if (finalize_gstate.global_state && finalize_bind_data.interface && ++ finalize_bind_data.interface->FinalizeScan(context, *finalize_gstate.global_state, output)) { ++ data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; ++ return; ++ } + data_p.async_result = SourceResultType::FINISHED; + return; + } +@@ -610,6 +625,10 @@ + auto &bind_data = data_p.bind_data->CastNoConst(); + + if (gstate.finished) { ++ if (bind_data.interface->FinalizeScan(context, *gstate.global_state, output)) { ++ data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; ++ return; ++ } + data_p.async_result = SourceResultType::FINISHED; + return; + } +@@ -660,6 +679,11 @@ + } + + if (!TryInitializeNextBatch(context, bind_data, data, gstate)) { ++ if (output.size() == 0 && ++ bind_data.interface->FinalizeScan(context, *gstate.global_state, output)) { ++ data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; ++ return; ++ } + if (output.size() > 0 && data_p.results_execution_mode == AsyncResultsExecutionMode::SYNCHRONOUS) { + gstate.finished = true; + data_p.async_result = SourceResultType::HAVE_MORE_OUTPUT; diff --git a/vortex-duckdb/src/column_statistics.rs b/vortex-duckdb/src/column_statistics.rs index 2b1cba1f2bb..ce5b85a61ac 100644 --- a/vortex-duckdb/src/column_statistics.rs +++ b/vortex-duckdb/src/column_statistics.rs @@ -4,24 +4,27 @@ use vortex::array::stats::StatsSet; use vortex::dtype::DType; use vortex::error::VortexExpect as _; +use vortex::error::VortexResult; use vortex::expr::stats::Precision; use vortex::expr::stats::Stat; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; use crate::convert::ToDuckDBScalar as _; +use crate::duckdb::LogicalType; use crate::duckdb::Value; -#[derive(Debug, Default)] +#[derive(Debug)] pub struct ColumnStatistics { pub min: Option, pub max: Option, pub max_string_length: u64, pub has_null: bool, + pub logical_type: LogicalType, } impl ColumnStatistics { - pub fn from(stats: &ColumnStatisticsAggregate, dtype: DType) -> Self { + pub fn from(stats: &ColumnStatisticsAggregate, dtype: DType) -> VortexResult { let min = stats.min.as_ref().and_then(|value| { Scalar::try_new(dtype.clone(), Some(value.clone())) .and_then(|scalar| scalar.try_to_duckdb_scalar()) @@ -40,12 +43,15 @@ impl ColumnStatistics { // Useful estimate if we didn't get null count stats let has_null = stats.has_null && dtype.is_nullable(); - Self { + let logical_type = LogicalType::try_from(dtype)?; + + Ok(Self { min, max, max_string_length, has_null, - } + logical_type, + }) } } diff --git a/vortex-duckdb/src/duckdb/bind_input.rs b/vortex-duckdb/src/duckdb/bind_input.rs index 1049a8065f6..e8119ff5066 100644 --- a/vortex-duckdb/src/duckdb/bind_input.rs +++ b/vortex-duckdb/src/duckdb/bind_input.rs @@ -3,25 +3,9 @@ use crate::cpp; use crate::duckdb::LogicalTypeRef; -use crate::duckdb::Value; use crate::lifetime_wrapper; -lifetime_wrapper!(BindInput, cpp::duckdb_vx_tfunc_bind_input, |_| {}); - -impl BindInputRef { - /// Returns the parameter at the given index. - pub fn get_parameter(&self, index: usize) -> Option { - let value_ptr = - unsafe { cpp::duckdb_vx_tfunc_bind_input_get_parameter(self.as_ptr(), index as _) }; - if value_ptr.is_null() { - None - } else { - Some(unsafe { Value::own(value_ptr) }) - } - } -} - -lifetime_wrapper!(BindResult, cpp::duckdb_vx_tfunc_bind_result, |_| {}); +lifetime_wrapper!(BindResult, cpp::duckdb_bind_result, |_| {}); impl BindResultRef { pub fn add_result_column(&self, name: &str, logical_type: &LogicalTypeRef) { diff --git a/vortex-duckdb/src/duckdb/logical_type.rs b/vortex-duckdb/src/duckdb/logical_type.rs index 28c17cbcf02..edc71c302b2 100644 --- a/vortex-duckdb/src/duckdb/logical_type.rs +++ b/vortex-duckdb/src/duckdb/logical_type.rs @@ -233,6 +233,10 @@ impl LogicalType { } impl LogicalTypeRef { + pub fn to_owned(&self) -> LogicalType { + unsafe { LogicalType::own(duckdb_vx_logical_type_copy(self.as_ptr())) } + } + pub fn as_type_id(&self) -> DUCKDB_TYPE { unsafe { duckdb_get_type_id(self.as_ptr()) } } diff --git a/vortex-duckdb/src/duckdb/table_init_input.rs b/vortex-duckdb/src/duckdb/table_init_input.rs index f6ac05ae0b5..1d86cb0fded 100644 --- a/vortex-duckdb/src/duckdb/table_init_input.rs +++ b/vortex-duckdb/src/duckdb/table_init_input.rs @@ -17,7 +17,6 @@ impl Debug for TableInitInput<'_> { fn fmt(&self, f: &mut Formatter<'_>) -> Result { f.debug_struct("TableInitInput") .field("column_ids", &self.column_ids()) - .field("projection_ids", &self.projection_ids()) .field("table_filter_set", &self.table_filter_set()) .finish() } @@ -32,18 +31,6 @@ impl<'a> TableInitInput<'a> { unsafe { std::slice::from_raw_parts(self.input.column_ids, self.input.column_ids_count) } } - pub fn projection_ids(&self) -> Option<&[u64]> { - // Passed pointer is std::vector's .data(). However, C++ doesn't - // guarantee an empty vector's pointer is nullptr so we need to check - // both conditions - if self.input.projection_ids.is_null() || self.input.projection_ids_count == 0 { - return None; - } - Some(unsafe { - std::slice::from_raw_parts(self.input.projection_ids, self.input.projection_ids_count) - }) - } - /// Returns the table filter set for the table function. pub fn table_filter_set(&self) -> Option<&TableFilterSetRef> { let ptr = self.input.filters; diff --git a/vortex-duckdb/src/exporter/cache.rs b/vortex-duckdb/src/exporter/cache.rs index 2f495ba9608..3b0fd496360 100644 --- a/vortex-duckdb/src/exporter/cache.rs +++ b/vortex-duckdb/src/exporter/cache.rs @@ -21,5 +21,4 @@ pub struct ConversionCache { pub dict_cache: DashMap, pub values_cache: DashMap>)>, pub canonical_cache: DashMap, - pub file_index: usize, } diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index c75eaf474da..df8ad5e2b2c 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -79,13 +79,11 @@ impl ArrayExporter { }) } - /// Export the data into the next chunk. - /// - /// Returns `true` if a chunk was exported, `false` if all rows have been exported. + /// Export the data into next chunk. + /// Returns true if there's more data to export into next chunk. pub fn export( &mut self, chunk: &mut DataChunkRef, - file_index_column_pos: Option, file_row_number_column_pos: Option, ) -> VortexResult { chunk.reset(); @@ -96,7 +94,7 @@ impl ArrayExporter { let zero_projection = self.fields.is_empty(); // file_row_number column is already populated in scan construction - let expected_cols = self.fields.len() + file_index_column_pos.is_some() as usize; + let expected_cols = self.fields.len(); let chunk_cols = chunk.column_count(); if !zero_projection && chunk_cols != expected_cols { vortex_bail!("Expected {expected_cols} columns in output chunk, got {chunk_cols}"); @@ -143,14 +141,6 @@ impl ArrayExporter { } for i in 0..chunk_cols { - // file_index column: skip index - it will be filled after - // chunk export. - if let Some(pos) = file_index_column_pos - && i == pos - { - continue; - } - // file_row_number column: skip index, already filled if let Some(pos) = file_row_number_column_pos && i == pos diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 2a5e9316434..2921cd6ff16 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -18,7 +18,6 @@ use crate::copy::copy_to_initialize_global; use crate::copy::copy_to_sink; use crate::cpp; use crate::duckdb::AggregatePushdownInput; -use crate::duckdb::BindInput; use crate::duckdb::BindResult; use crate::duckdb::Data; use crate::duckdb::DataChunk; @@ -29,144 +28,71 @@ use crate::duckdb::LogicalTypeRef; use crate::duckdb::TableInitInput; use crate::duckdb::try_or; use crate::duckdb::try_or_null; +use crate::file_reader::File; +use crate::file_reader::reader_bind; +use crate::file_reader::reader_get_progress_in_file; +use crate::file_reader::reader_get_statistics; +use crate::file_reader::reader_initialize; +use crate::file_reader::reader_open; +use crate::file_reader::reader_scan; +use crate::file_reader::reader_try_initialize_scan; +use crate::table_function::BindState; use crate::table_function::Cardinality; -use crate::table_function::TableFunctionBind; -use crate::table_function::TableFunctionGlobal; -use crate::table_function::TableFunctionLocal; -use crate::table_function::bind; +use crate::table_function::GlobalState; +use crate::table_function::LocalState; use crate::table_function::cardinality; -use crate::table_function::get_partition_data; +use crate::table_function::finalize_scan; use crate::table_function::init_global; use crate::table_function::init_local; use crate::table_function::pushdown_complex_filter; use crate::table_function::pushdown_projection_aggregates; use crate::table_function::pushdown_projection_expression; -use crate::table_function::scan; -use crate::table_function::statistics; -use crate::table_function::table_scan_progress; use crate::table_function::to_string; #[unsafe(no_mangle)] unsafe extern "C-unwind" fn duckdb_table_function_to_string( - bind_data: *const c_void, + bind: *const c_void, map: cpp::duckdb_vx_string_map, ) { - let bind_data = unsafe { bind_data.cast::().as_ref() } - .vortex_expect("bind_data null pointer"); + let bind = unsafe { bind.cast::().as_ref() }.vortex_expect("null pointer"); let map = unsafe { DuckdbStringMap::borrow_mut(map) }; - to_string(bind_data, map); -} - -#[unsafe(no_mangle)] -unsafe extern "C-unwind" fn duckdb_table_function_statistics( - bind_data: *const c_void, - column_index: usize, - stats_out: *mut cpp::duckdb_column_statistics, -) -> bool { - let stats_out = unsafe { &mut *stats_out }; - let bind_data = unsafe { bind_data.cast::().as_ref() } - .vortex_expect("bind_data null pointer"); - let Some(stats) = statistics(bind_data, column_index) else { - return false; - }; - stats_out.min = stats.min.map_or(ptr::null_mut(), |v| v.into_ptr()); - stats_out.max = stats.max.map_or(ptr::null_mut(), |v| v.into_ptr()); - stats_out.max_string_length = stats.max_string_length; - stats_out.has_null = stats.has_null; - true -} - -#[unsafe(no_mangle)] -unsafe extern "C-unwind" fn duckdb_table_function_scan_progress(global_state: *mut c_void) -> f64 { - let global_state = unsafe { global_state.cast::().as_ref() } - .vortex_expect("global_init_data null pointer"); - table_scan_progress(global_state) -} - -#[unsafe(no_mangle)] -unsafe extern "C-unwind" fn duckdb_table_function_get_partition_data( - global_init_data: *mut c_void, - local_init_data: *mut c_void, - partition_data_out: *mut cpp::duckdb_vx_partition_data, -) { - let global_init_data = unsafe { global_init_data.cast::().as_ref() } - .vortex_expect("global_init_data null pointer"); - let local_init_data = unsafe { local_init_data.cast::().as_mut() } - .vortex_expect("local_init_data null pointer"); - let data = get_partition_data(global_init_data, local_init_data); - let out = unsafe { &mut *partition_data_out }; - - out.partition_index = data.partition_index; - out.file_index_column_pos = data.file_index_column_pos.unwrap_or(usize::MAX); - out.file_index = data.file_index; + to_string(bind, map); } #[unsafe(no_mangle)] unsafe extern "C-unwind" fn duckdb_table_function_pushdown_complex_filter( - bind_data: *mut c_void, + bind: *mut c_void, expr: cpp::duckdb_vx_expr, - error_out: *mut cpp::duckdb_vx_error, + error: *mut cpp::duckdb_vx_error, ) -> bool { - let bind_data = unsafe { bind_data.cast::().as_mut() } - .vortex_expect("bind_data null pointer"); + let bind = unsafe { bind.cast::().as_mut() }.vortex_expect("null pointer"); let expr = unsafe { Expression::borrow(expr) }; - try_or(error_out, || pushdown_complex_filter(bind_data, expr)) + try_or(error, || pushdown_complex_filter(bind, expr)) } #[unsafe(no_mangle)] unsafe extern "C-unwind" fn duckdb_table_function_pushdown_projection_expression( - bind_data: *mut c_void, + bind: *mut c_void, expr: cpp::duckdb_vx_expr, column_id: usize, - error_out: *mut cpp::duckdb_vx_error, + error: *mut cpp::duckdb_vx_error, ) -> bool { - let bind_data = unsafe { bind_data.cast::().as_mut() } - .vortex_expect("bind_data null pointer"); + let bind = unsafe { bind.cast::().as_mut() }.vortex_expect("null pointer"); let expr = unsafe { Expression::borrow(expr) }; - try_or(error_out, || { - pushdown_projection_expression(bind_data, expr, column_id) + try_or(error, || { + pushdown_projection_expression(bind, expr, column_id) }) } #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_pushdown_projection_aggregates( - bind_data: *mut c_void, + bind: *mut c_void, input: cpp::duckdb_vx_agg_input, - error_out: *mut cpp::duckdb_vx_error, + error: *mut cpp::duckdb_vx_error, ) -> bool { - let bind_data = unsafe { bind_data.cast::().as_mut() } - .vortex_expect("bind_data null pointer"); + let bind = unsafe { bind.cast::().as_mut() }.vortex_expect("null pointer"); let input = unsafe { AggregatePushdownInput::borrow(input) }; - try_or(error_out, || { - pushdown_projection_aggregates(bind_data, input) - }) -} - -#[unsafe(no_mangle)] -unsafe extern "C-unwind" fn duckdb_table_function_scan( - global_init_data: *mut c_void, - local_init_data: *mut c_void, - output: cpp::duckdb_data_chunk, - error_out: *mut cpp::duckdb_vx_error, -) { - let global_init_data = unsafe { global_init_data.cast::().as_ref() } - .vortex_expect("global_init_data null pointer"); - let local_init_data = unsafe { local_init_data.cast::().as_mut() } - .vortex_expect("local_init_data null pointer"); - let data_chunk = unsafe { DataChunk::borrow_mut(output) }; - - match scan(local_init_data, global_init_data, data_chunk) { - Ok(()) => { - // The data chunk is already filled by the function. - // No need to do anything here. - } - Err(e) => unsafe { - error_out.write(cpp::duckdb_vx_error_create( - e.to_string().as_ptr().cast(), - e.to_string().len(), - )); - }, - } + try_or(error, || pushdown_projection_aggregates(bind, input)) } #[unsafe(no_mangle)] @@ -178,25 +104,23 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_pushdown_expression( #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_cardinality( - bind_data: *const c_void, - node_stats_out: *mut cpp::duckdb_vx_node_statistics, + bind: *const c_void, + file_count: u64, + stats: *mut cpp::duckdb_vx_node_statistics, ) { - let bind_data = unsafe { bind_data.cast::().as_ref() } - .vortex_expect("bind_data null pointer"); - let node_stats = - unsafe { node_stats_out.as_mut() }.vortex_expect("node_stats_out null pointer"); + let bind = unsafe { bind.cast::().as_ref() }.vortex_expect("null pointer"); + let stats = unsafe { stats.as_mut() }.vortex_expect("null pointer"); - match cardinality(bind_data) { - Cardinality::Unknown => {} + match cardinality(bind, file_count) { Cardinality::Exact(c) => { - node_stats.has_estimated_cardinality = true; - node_stats.estimated_cardinality = c as _; - node_stats.has_max_cardinality = true; - node_stats.max_cardinality = c as _; + stats.has_estimated_cardinality = true; + stats.estimated_cardinality = c as _; + stats.has_max_cardinality = true; + stats.max_cardinality = c as _; } Cardinality::Estimate(c) => { - node_stats.has_estimated_cardinality = true; - node_stats.estimated_cardinality = c as _; + stats.has_estimated_cardinality = true; + stats.estimated_cardinality = c as _; } } } @@ -204,18 +128,17 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_cardinality( #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_init_global( init_input: *const cpp::duckdb_vx_tfunc_init_input, - error_out: *mut cpp::duckdb_vx_error, + error: *mut cpp::duckdb_vx_error, ) -> cpp::duckdb_vx_data { - let init_input = TableInitInput::new( - unsafe { init_input.as_ref() }.vortex_expect("init_input null pointer"), - ); + let init_input = + TableInitInput::new(unsafe { init_input.as_ref() }.vortex_expect("null pointer")); match init_global(&init_input) { Ok(init_data) => Data::from(Box::new(init_data)).as_ptr(), Err(e) => { // Set the error in the error output. let msg = e.to_string(); - unsafe { error_out.write(cpp::duckdb_vx_error_create(msg.as_ptr().cast(), msg.len())) }; + unsafe { error.write(cpp::duckdb_vx_error_create(msg.as_ptr().cast(), msg.len())) }; ptr::null_mut::().cast() } } @@ -223,40 +146,128 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_init_global( #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_init_local( - bind_data: *const c_void, - global_init_data: *mut c_void, + bind: *const c_void, + global: *const c_void, ) -> cpp::duckdb_vx_data { - let bind_data = unsafe { bind_data.cast::().as_ref() } - .vortex_expect("bind_data null pointer"); - let global_init_data = unsafe { global_init_data.cast::().as_ref() } - .vortex_expect("global_init_data null pointer"); - - let init_data = init_local(bind_data, global_init_data); - Data::from(Box::new(init_data)).as_ptr() + let bind = unsafe { bind.cast::().as_ref() }.vortex_expect("null pointer"); + let global = unsafe { global.cast::().as_ref() }.vortex_expect("null pointer"); + let local = init_local(bind, global); + Data::from(Box::new(local)).as_ptr() } #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn duckdb_table_function_bind( - bind_input: cpp::duckdb_vx_tfunc_bind_input, - bind_result: cpp::duckdb_vx_tfunc_bind_result, +pub unsafe extern "C-unwind" fn duckdb_reader_bind( + first_file: *const c_void, + result: cpp::duckdb_bind_result, error_out: *mut cpp::duckdb_vx_error, ) -> cpp::duckdb_vx_data { - let bind_input = unsafe { BindInput::own(bind_input) }; - let mut bind_result = unsafe { BindResult::own(bind_result) }; + let first_file = unsafe { first_file.cast::().as_ref() }.vortex_expect("null pointer"); + let mut result = unsafe { BindResult::own(result) }; try_or_null(error_out, || { - let bind_data = bind(&bind_input, &mut bind_result)?; + let bind_data = reader_bind(first_file, &mut result)?; Ok(Data::from(Box::new(bind_data)).as_ptr()) }) } +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_open( + file_path: *const c_char, + file_path_len: usize, + error: *mut cpp::duckdb_vx_error, +) -> cpp::duckdb_vx_data { + let path = unsafe { std::slice::from_raw_parts(file_path.cast::(), file_path_len) }; + let path = String::from_utf8_lossy(path).into_owned(); + + try_or_null(error, || { + let file = reader_open(&path)?; + Ok(Data::from(Box::new(file)).as_ptr()) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_get_statistics( + file: *const c_void, + column_name: *const c_char, + column_name_len: usize, + stats_out: *mut cpp::duckdb_column_statistics, +) -> bool { + let file = unsafe { file.cast::().as_ref() }.vortex_expect("null pointer"); + let name_bytes = + unsafe { std::slice::from_raw_parts(column_name.cast::(), column_name_len) }; + let column_name = String::from_utf8_lossy(name_bytes); + + let Some(stats) = reader_get_statistics(file, &column_name) else { + return false; + }; + let stats_out = unsafe { &mut *stats_out }; + stats_out.min = stats.min.map_or(ptr::null_mut(), |v| v.into_ptr()); + stats_out.max = stats.max.map_or(ptr::null_mut(), |v| v.into_ptr()); + stats_out.max_string_length = stats.max_string_length; + stats_out.has_null = stats.has_null; + stats_out.type_ = stats.logical_type.into_ptr(); + true +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_initialize( + global: *const c_void, + file: *mut c_void, + error: *mut cpp::duckdb_vx_error, +) -> bool { + let global = unsafe { global.cast::().as_ref() }.vortex_expect("null pointer"); + let file = unsafe { file.cast::().as_mut() }.vortex_expect("null pointer"); + try_or(error, || reader_initialize(file, global)) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_try_initialize_scan( + local: *mut c_void, + file: *mut c_void, +) -> bool { + let file = unsafe { file.cast::().as_mut() }.vortex_expect("null pointer"); + let local = unsafe { local.cast::().as_mut() }.vortex_expect("null pointer"); + reader_try_initialize_scan(file, local) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_scan( + file: *const c_void, + global: *const c_void, + local: *mut c_void, + chunk: cpp::duckdb_data_chunk, + error: *mut cpp::duckdb_vx_error, +) -> bool { + let file = unsafe { file.cast::().as_ref() }.vortex_expect("null pointer"); + let global = unsafe { global.cast::().as_ref() }.vortex_expect("null pointer"); + let local = unsafe { local.cast::().as_mut() }.vortex_expect("null pointer"); + let chunk = unsafe { DataChunk::borrow_mut(chunk) }; + try_or(error, || reader_scan(file, global, local, chunk)) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_get_progress_in_file(file: *const c_void) -> f64 { + let file = unsafe { file.cast::().as_ref() }.vortex_expect("null pointer"); + reader_get_progress_in_file(file) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_reader_finalize_scan( + global: *const c_void, + chunk: cpp::duckdb_data_chunk, + error: *mut cpp::duckdb_vx_error, +) -> bool { + let global = unsafe { global.cast::().as_ref() }.vortex_expect("null pointer"); + let chunk = unsafe { DataChunk::borrow_mut(chunk) }; + try_or(error, || finalize_scan(global, chunk)) +} + #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_bind_data_clone( - bind_data: *const c_void, + bind: *const c_void, ) -> cpp::duckdb_vx_data { - let bind_data = unsafe { bind_data.cast::().as_ref() } - .vortex_expect("bind_data null pointer"); - let copied_data = bind_data.clone(); + let bind = unsafe { bind.cast::().as_ref() }.vortex_expect("null pointer"); + let copied_data = bind.clone(); Data::from(Box::new(copied_data)).as_ptr() } diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs new file mode 100644 index 00000000000..23e5312dca7 --- /dev/null +++ b/vortex-duckdb/src/file_reader.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use futures::FutureExt; +use object_store::registry::ObjectStoreRegistry; +use url::Url; +use vortex::array::VortexSessionExecute as _; +use vortex::array::arrays::struct_::StructArrayExt as _; +use vortex::cloud::Registry; +use vortex::dtype::DType; +use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::error::vortex_panic; +use vortex::file::multi::open_cached; +use vortex::file::multi::parse_uri_or_path; +use vortex::file::v2::FileStatsLayoutReader; +use vortex::io::compat::Compat; +use vortex::io::filesystem::FileSystemRef; +use vortex::io::object_store::ObjectStoreFileSystem; +use vortex::io::runtime::BlockingRuntime as _; +use vortex::layout::LayoutReaderRef; +use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::mask::Mask; + +use crate::RUNTIME; +use crate::SESSION; +use crate::column_statistics::ColumnStatistics; +use crate::column_statistics::ColumnStatisticsAggregate; +use crate::duckdb::BindResultRef; +use crate::duckdb::DataChunkRef; +use crate::exporter::ArrayExporter; +use crate::exporter::ConversionCache; +use crate::projection::Filter; +use crate::projection::extract_schema_from_dtype; +use crate::table_function::BindState; +use crate::table_function::GlobalState; +use crate::table_function::LocalState; +use crate::table_function::Split; +use crate::table_function::convert_result; + +static REGISTRY: LazyLock = LazyLock::new(Registry::new); + +fn resolve_filesystem(url: &Url) -> VortexResult<(FileSystemRef, String)> { + // Compat makes us use tokio which is very bad for local reads on + // high-core machines because reads go into blocking pool + if url.scheme() == "file" { + return Ok(( + Arc::new(ObjectStoreFileSystem::local(RUNTIME.handle())), + url.path().to_string(), + )); + } + + let (object_store, path) = REGISTRY.resolve(url)?; + + Ok(( + Arc::new(ObjectStoreFileSystem::new( + Arc::new(Compat::new(object_store)), + RUNTIME.handle(), + )), + path.to_string(), + )) +} + +pub struct File { + pub reader: LayoutReaderRef, + /// File splits stored in inverse order + pub splits: Vec, + pub cache: ConversionCache, + total_splits: usize, +} + +async fn open_reader(file_path: String) -> VortexResult { + let url = parse_uri_or_path(&file_path)?; + let (fs, path) = resolve_filesystem(&url)?; + let file = fs.open_read(&path).await?; + let file = open_cached(&SESSION, file, &path, None, &|options| options).await?; + Ok(File { + reader: file.layout_reader()?, + cache: ConversionCache::default(), + splits: vec![], + total_splits: 0, + }) +} + +pub fn reader_open(file_path: &str) -> VortexResult { + RUNTIME.block_on(open_reader(file_path.to_owned())) +} + +pub fn reader_bind(file: &File, result: &mut BindResultRef) -> VortexResult { + let dtype = file.reader.dtype().clone(); + let columns = extract_schema_from_dtype(&dtype)?; + + for column in &columns { + result.add_result_column(&column.name, &column.logical_type); + } + + Ok(BindState { + dtype, + first_file_row_count: file.reader.row_count(), + filters: vec![], + columns, + has_non_optional_filter: AtomicBool::new(false), + aggregates: vec![], + }) +} + +/// Returns true if file should be skipped. +/// Called under file lock. +pub fn reader_initialize(file: &mut File, global: &GlobalState) -> VortexResult { + if reader_prune(file, &global.filter)? { + return Ok(true); + } + + let ordered = global.file_row_number_column_pos.is_some(); + let reader = Arc::clone(&file.reader); + let filter = &global.filter; + let mut builder = ScanBuilder::new(SESSION.clone(), reader) + .with_projection(global.projection.clone()) + .with_ordered(ordered) + .with_some_filter(filter.filter.clone()) + .with_selection(filter.row_selection.clone()); + if let Some(row_range) = filter.row_range.as_ref() { + builder = builder.with_row_range(row_range.clone()); + } + let mut splits = builder.build()?; + + // threads take last element of file.splits so we need to reverse + splits.reverse(); + file.total_splits = splits.len(); + file.splits = splits; + Ok(false) +} + +/// Returns false if file is exhausted. +/// Called from all threads under global lock. +pub fn reader_try_initialize_scan(file: &mut File, local: &mut LocalState) -> bool { + let Some(split) = file.splits.pop() else { + return false; + }; + local.split = Some(split); + true +} + +/// Returns false if file is exhausted +pub fn reader_scan( + file: &File, + global: &GlobalState, + local: &mut LocalState, + chunk: &mut DataChunkRef, +) -> VortexResult { + if !local.partials.is_empty() { + return reader_scan_aggregate(global, local); + } + + if local.exporter.is_none() { + let Some(split) = local.split.take() else { + return Ok(false); + }; + let Some(array) = RUNTIME.block_on(async move { split.await })? else { + // split is filtered + return Ok(true); + }; + let mut ctx = SESSION.create_execution_ctx(); + let array = convert_result(array, &mut ctx)?; + local.exporter = Some(ArrayExporter::try_new(&array, &file.cache, ctx)?); + } + let exporter = local.exporter.as_mut().vortex_expect("no exporter"); + + let has_more_data = exporter.export(chunk, global.file_row_number_column_pos)?; + if !has_more_data { + local.exporter = None; + } + Ok(true) +} + +fn reader_scan_aggregate(global: &GlobalState, local: &mut LocalState) -> VortexResult { + let Some(split) = local.split.take() else { + return Ok(false); + }; + let Some(array) = RUNTIME.block_on(async move { split.await })? else { + return Ok(true); + }; + global.pending.fetch_add(1, Ordering::Relaxed); + + let mut ctx = SESSION.create_execution_ctx(); + let array = convert_result(array, &mut ctx)?; + + for (position, partial) in global + .aggregate_positions + .iter() + .zip(local.partials.iter_mut()) + { + partial.accumulate(array.unmasked_field(*position), &mut ctx)?; + } + + { + let mut partials = global.partials.lock(); + for (global_partial, local_partial) in partials.iter_mut().zip(&mut local.partials) { + global_partial.combine_partials(local_partial.flush()?)?; + } + } + + let has_count_star = local.partials.len() < global.aggregates.len(); + if has_count_star { + let len = array.len() as u64; + global.row_count.fetch_add(len, Ordering::Relaxed); + } + + global.pending.fetch_sub(1, Ordering::Release); + Ok(true) +} + +pub fn reader_get_statistics(file: &File, column: &str) -> Option { + let reader = file + .reader + .as_any() + .downcast_ref::()?; + let stats_sets = reader.file_stats().stats_sets(); + + let DType::Struct(fields, _) = &file.reader.dtype() else { + return None; + }; + let index = fields.find(column)?; + let dtype = fields.field_by_index(index)?; + + let stats = ColumnStatisticsAggregate::new(stats_sets.get(index)?); + match ColumnStatistics::from(&stats, dtype) { + Ok(stats) => Some(stats), + Err(e) => vortex_panic!(e), + } +} + +fn reader_prune(file: &File, filter: &Filter) -> VortexResult { + let Some(filter) = &filter.filter else { + return Ok(false); + }; + let row_count = file.reader.row_count(); + let row_range = 0..row_count; + let mask = Mask::new_true(usize::try_from(row_count).unwrap_or(usize::MAX)); + let evaluation = file.reader.pruning_evaluation(&row_range, filter, mask)?; + match evaluation.now_or_never() { + Some(Ok(result_mask)) => Ok(result_mask.all_false()), + _ => Ok(false), + } +} + +pub fn reader_get_progress_in_file(file: &File) -> f64 { + let total = file.total_splits; + let left = file.splits.len(); + let denom = total + (total == 0) as usize; + 100.0 * (total - left) as f64 / denom as f64 +} diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index b16d41e9e75..84e31ff1440 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -27,7 +27,7 @@ mod convert; pub mod duckdb; mod exporter; mod ffi; -mod multi_file; +mod file_reader; mod projection; mod table_function; diff --git a/vortex-duckdb/src/multi_file.rs b/vortex-duckdb/src/multi_file.rs deleted file mode 100644 index f5bb3323e08..00000000000 --- a/vortex-duckdb/src/multi_file.rs +++ /dev/null @@ -1,100 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use itertools::Itertools; -use object_store::registry::ObjectStoreRegistry; -use url::Url; -use vortex::error::VortexResult; -use vortex::error::vortex_bail; -use vortex::error::vortex_err; -use vortex::file::multi::MultiFileDataSource; -use vortex::file::multi::parse_uri_or_path; -use vortex::io::compat::Compat; -use vortex::io::filesystem::FileSystemRef; -use vortex::io::object_store::ObjectStoreFileSystem; -use vortex::io::runtime::BlockingRuntime; -use vortex::layout::scan::multi::MultiLayoutDataSource; - -use crate::REGISTRY; -use crate::RUNTIME; -use crate::SESSION; -use crate::duckdb::BindInputRef; -use crate::duckdb::ExtractedValue; - -fn resolve_filesystem(glob_url: &Url) -> VortexResult<(FileSystemRef, String)> { - // Compat makes us use tokio which is very bad for local reads on - // high-core machines because reads go into blocking pool - if glob_url.scheme() == "file" { - return Ok(( - Arc::new(ObjectStoreFileSystem::local(RUNTIME.handle())), - glob_url.path().to_string(), - )); - } - - // The full URL goes through the shared registry, which reports the glob as a path *within* - // the store it returns. For most schemes the store is mounted at the URL authority, so the - // path is the whole URL path — but not for all of them: an `hf://` store is rooted at a - // repository and revision, which occupy path segments. Only the registry knows how deep the - // store is mounted, so globbing anything other than the path it reports would address the - // wrong keys. Going through the registry also means DuckDB resolves the same set of schemes - // as the Python and Java bindings, including the OpenDAL-backed ones when the `opendal` - // feature is on. The registry caches one client per store prefix, so repeated scans against - // the same bucket or repository share a client even though the filesystem wrapper is rebuilt. - let (object_store, path) = REGISTRY.resolve(glob_url)?; - - Ok(( - Arc::new(ObjectStoreFileSystem::new( - Arc::new(Compat::new(object_store)), - RUNTIME.handle(), - )), - path.to_string(), - )) -} - -/// Shared bind logic for both single-glob and multi-glob variants. -pub fn bind_multi_file_scan(input: &BindInputRef) -> VortexResult { - let glob_url_parameter = input - .get_parameter(0) - .ok_or_else(|| vortex_err!("Missing file glob parameter"))?; - - // The input to the table function can either be a single glob, or a List of glob patterns. - let glob_strings: Vec = match glob_url_parameter.extract() { - ExtractedValue::Varchar(glob) => { - vec![glob.to_string()] - } - ExtractedValue::List(globs) => globs - .into_iter() - .map(|glob| { - let ExtractedValue::Varchar(string) = glob.extract() else { - vortex_bail!("list element must be Varchar type") - }; - - Ok(string.to_string()) - }) - .try_collect()?, - _ => vortex_bail!("Invalid argument to read_vortex table function"), - }; - - // Parse each glob URL and resolve its filesystem. - let mut glob_urls: Vec = Vec::with_capacity(glob_strings.len()); - for glob_str in &glob_strings { - glob_urls.push(parse_uri_or_path(glob_str)?); - } - - let resolved = glob_urls - .iter() - .map(resolve_filesystem) - .collect::>>()?; - - RUNTIME.block_on(async { - let mut builder = MultiFileDataSource::new(SESSION.clone()); - - for (fs, glob) in resolved { - builder = builder.with_glob(&glob, Some(fs)); - } - - builder.build().await - }) -} diff --git a/vortex-duckdb/src/projection.rs b/vortex-duckdb/src/projection.rs index c1ae65f2c31..dad7060d1f1 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -7,6 +7,7 @@ use vortex::dtype::DType; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::expr::BoundExpression; use vortex::expr::Expression; use vortex::expr::and_collect; use vortex::expr::col; @@ -28,10 +29,8 @@ use crate::table_function::ColumnAggregate; // See MultiFileReader for constants -/// "file_index" virtual column -static FILE_INDEX_COLUMN_IDX: u64 = 9223372036854775810; /// "file_row_number" virtual column -static FILE_ROW_NUMBER_COLUMN_IDX: u64 = 9223372036854775809; +pub(crate) static FILE_ROW_NUMBER_COLUMN_IDX: u64 = 9223372036854775809; /// See duckdb/src/common/constants.cpp fn is_virtual_column(id: u64) -> bool { @@ -48,46 +47,19 @@ pub struct DuckdbField { pub projection_expr: Option, } -pub struct Projection { - pub projection: Expression, - pub file_index_column_pos: Option, - pub file_row_number_column_pos: Option, -} +pub struct Projection(pub Expression); impl Projection { - pub fn new( - projection_ids: Option<&[u64]>, - column_ids: &[u64], - column_fields: &[DuckdbField], - ) -> Self { - // If projection ids are empty, use column_ids. - // See duckdb/src/planner/operator/logical_get.cpp#L168 - let (ids, has_projection_ids) = match projection_ids { - Some(ids) => (ids, true), - None => (column_ids, false), - }; - - let mut file_index_column_pos = None; - let mut file_row_number_column_pos = None; + pub fn new(column_ids: &[u64], column_fields: &[DuckdbField]) -> Self { + let mut has_file_row_number = false; let mut is_star = true; let mut real_column_count = 0; let mut projected_col_count = 0; // DuckDB uses u64 as column indices but Rust uses usize - for (column_pos, &column_id) in ids.iter().enumerate() { - let column_id = if has_projection_ids { - let column_id: usize = column_id.as_(); - column_ids[column_id] - } else { - column_id - }; - - if column_id == FILE_INDEX_COLUMN_IDX { - file_index_column_pos = Some(column_pos); - continue; - } + for &column_id in column_ids { if column_id == FILE_ROW_NUMBER_COLUMN_IDX { - file_row_number_column_pos = Some(column_pos); + has_file_row_number = true; continue; } if is_virtual_column(column_id) { @@ -112,7 +84,6 @@ impl Projection { // 5 columns total. is_star &= real_column_count == column_fields.len() as u64; - let has_file_row_number = file_row_number_column_pos.is_some(); if is_star { let projection = if has_file_row_number { // row_idx will be moved to correct position in scan(), prepend here @@ -121,21 +92,17 @@ impl Projection { } else { root() }; - return Projection { - projection, - file_index_column_pos, - file_row_number_column_pos, - }; + return Projection(projection); } let has_columns_with_expr = projected_col_count > 0; let (mut all_exprs, mut named_fields) = if has_columns_with_expr { - let all = Vec::with_capacity(ids.len() + has_file_row_number as usize); + let all = Vec::with_capacity(column_ids.len() + has_file_row_number as usize); let named = Vec::new(); (all, named) } else { let all = Vec::new(); - let named = Vec::with_capacity(ids.len()); + let named = Vec::with_capacity(column_ids.len()); (all, named) }; @@ -144,13 +111,7 @@ impl Projection { all_exprs.push(("file_row_number", row_idx())); } - for &column_id in ids { - let column_id = if has_projection_ids { - let column_id: usize = column_id.as_(); - column_ids[column_id] - } else { - column_id - }; + for &column_id in column_ids { if is_virtual_column(column_id) { continue; } @@ -183,11 +144,7 @@ impl Projection { select(named_fields, root()) }; - Self { - projection, - file_index_column_pos, - file_row_number_column_pos, - } + Self(projection) } // Create a projection for aggregate scan @@ -219,20 +176,14 @@ impl Projection { let names = exprs.into_iter().map(|(name, _)| name).collect::>(); select(names, root()) }; - Projection { - projection, - file_index_column_pos: None, - file_row_number_column_pos: None, - } + Projection(projection) } } pub struct Filter { - pub filter: Option, + pub filter: Option, pub row_selection: Selection, pub row_range: Option>, - pub file_selection: Selection, - pub file_range: Option>, pub has_non_optional_filter: bool, } @@ -275,28 +226,25 @@ impl Filter { push_filter_expr(&mut table_filter_exprs, expr); } - let mut file_selection = Selection::All; let mut row_selection = Selection::All; let mut row_range = None; - let mut file_range = None; if let Some(filter) = table_filter_set { for (idx, expression) in filter.into_iter() { let idx: usize = idx.as_(); if column_ids[idx] == FILE_ROW_NUMBER_COLUMN_IDX { (row_selection, row_range) = try_from_virtual_column_filter(expression)?; } - if column_ids[idx] == FILE_INDEX_COLUMN_IDX { - (file_selection, file_range) = try_from_virtual_column_filter(expression)?; - } } }; + let filter = and_collect(table_filter_exprs) + .map(|expr| expr.optimize_recursive(dtype)?.bind(dtype)) + .transpose()?; + let out = Self { - filter: and_collect(table_filter_exprs), + filter, row_selection, row_range, - file_selection, - file_range, has_non_optional_filter, }; Ok(out) @@ -358,36 +306,28 @@ mod tests { }, ]; - assert_eq!(Projection::new(None, &ids, &fields).projection, root()); + assert_eq!(Projection::new(&ids, &fields).0, root()); - let ids = [FILE_ROW_NUMBER_COLUMN_IDX, 0, 1, FILE_INDEX_COLUMN_IDX, 2]; - let exprs = Projection::new(None, &ids, &fields); + let ids = [FILE_ROW_NUMBER_COLUMN_IDX, 0, 1, 2]; + let exprs = Projection::new(&ids, &fields); let row_idx_struct = pack([("file_row_number", row_idx())], false.into()); let root_with_virtual_cols = merge([row_idx_struct, root()]); - assert_eq!(exprs.projection, root_with_virtual_cols); - assert_eq!(exprs.file_index_column_pos, Some(3)); - assert_eq!(exprs.file_row_number_column_pos, Some(0)); - - // projections can't be set in SELECT *. - assert_ne!( - Projection::new(Some(&[0, 1]), &ids, &fields).projection, - root() - ); + assert_eq!(exprs.0, root_with_virtual_cols); let ids = [0, 1]; - assert_ne!(Projection::new(None, &ids, &fields).projection, root()); + assert_ne!(Projection::new(&ids, &fields).0, root()); let ids = [0, 2, 2]; - assert_ne!(Projection::new(None, &ids, &fields).projection, root()); + assert_ne!(Projection::new(&ids, &fields).0, root()); let ids = [2, 1, 0]; - assert_ne!(Projection::new(None, &ids, &fields).projection, root()); + assert_ne!(Projection::new(&ids, &fields).0, root()); // If any column has a projection expression, we can't use SELECT * fields[0].projection_expr = Some(lit(true)); let ids = [0, 1, 2]; - assert_ne!(Projection::new(None, &ids, &fields).projection, root()); + assert_ne!(Projection::new(&ids, &fields).0, root()); } #[test] diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 63bf76e2457..b448036d449 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -4,18 +4,12 @@ use std::cmp::max; use std::fmt::Formatter; use std::fmt::{self}; -use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; use custom_labels::CURRENT_LABELSET; -use futures::FutureExt; -use futures::Stream; -use futures::StreamExt; use futures::future::BoxFuture; use itertools::Itertools; use num_traits::AsPrimitive; @@ -26,85 +20,66 @@ use vortex::aggregate_fn::DynAccumulator; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::ExecutionCtx; -use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::ScalarFn; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; use vortex::array::arrays::scalar_fn::ScalarFnArrayExt; -use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::optimizer::ArrayOptimizer; use vortex::dtype::DType; use vortex::dtype::PType; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::expr::BoundExpression; use vortex::expr::Expression; -use vortex::expr::stats::Precision; -use vortex::file::v2::FileStatsLayoutReader; -use vortex::io::kanal_ext::KanalExt as _; -use vortex::io::runtime::BlockingRuntime as _; -use vortex::io::runtime::current::ThreadSafeIterator; -use vortex::layout::scan::multi::MultiLayoutChild; -use vortex::layout::scan::multi::MultiLayoutDataSource; use vortex::metrics::tracing::get_global_labels; use vortex::scalar::Scalar; use vortex::scalar_fn::fns::binary::Binary; use vortex::scalar_fn::fns::operators::Operator; use vortex::scalar_fn::fns::pack::Pack; -use vortex::scan::DataSource; -use vortex::scan::ScanRequest; use vortex_utils::aliases::hash_map::HashMap; -use vortex_utils::parallelism::get_available_parallelism; -use crate::RUNTIME; -use crate::SESSION; -use crate::column_statistics::ColumnStatistics; -use crate::column_statistics::ColumnStatisticsAggregate; use crate::convert::PushedAggregate; use crate::convert::try_from_bound_expression; use crate::convert::try_from_projection_aggregate; use crate::convert::try_from_projection_expression; use crate::cpp::DUCKDB_TYPE; -use crate::duckdb::AggregateExpression; -use crate::duckdb::AggregatePushdownInputRef; -use crate::duckdb::BindInputRef; -use crate::duckdb::BindResultRef; use crate::duckdb::DataChunkRef; use crate::duckdb::DuckdbStringMapRef; use crate::duckdb::ExpressionRef; use crate::duckdb::LogicalTypeRef; use crate::duckdb::TableInitInput; use crate::duckdb::Value; +use crate::duckdb::{AggregateExpression, TableFilterSetRef}; +use crate::duckdb::{AggregatePushdownInputRef, TableFilterSet}; use crate::exporter::ArrayExporter; -use crate::exporter::ConversionCache; -use crate::multi_file::bind_multi_file_scan; -use crate::projection::DuckdbField; -use crate::projection::Filter; +use crate::projection::FILE_ROW_NUMBER_COLUMN_IDX; use crate::projection::Projection; -use crate::projection::extract_schema_from_dtype; +use crate::projection::{DuckdbField, Filter}; // Aggregate projection index for count(*). See cpp/aggregate_fn_pushdown.cpp pub const COUNT_STAR_PROJ_IDX: u64 = u64::MAX; -pub struct TableFunctionBind { - data_source: Arc, - filter_exprs: Vec, - column_fields: Vec, +pub(crate) struct BindState { + pub dtype: DType, + pub first_file_row_count: u64, + pub filters: Vec, + pub columns: Vec, // There exists at least one non-optional table filter or at least one // complex filter is pushed down. - has_non_optional_filter: AtomicBool, + pub has_non_optional_filter: AtomicBool, // Non-empty iff this scan is aggregate - aggregates: Vec, + pub aggregates: Vec, } -assert_impl_all!(TableFunctionBind: Send, Clone); +assert_impl_all!(BindState: Send, Clone); -impl Clone for TableFunctionBind { +impl Clone for BindState { fn clone(&self) -> Self { Self { - data_source: Arc::clone(&self.data_source), - // filter_exprs are consumed once in `init_global`. - filter_exprs: vec![], - column_fields: self.column_fields.clone(), + dtype: self.dtype.clone(), + first_file_row_count: self.first_file_row_count, + filters: vec![], + columns: self.columns.clone(), has_non_optional_filter: AtomicBool::new( self.has_non_optional_filter.load(Ordering::Relaxed), ), @@ -113,14 +88,14 @@ impl Clone for TableFunctionBind { } } -impl fmt::Debug for TableFunctionBind { +impl fmt::Debug for BindState { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("DataSourceBindData") - .field("column_fields", &self.column_fields) + .field("columns", &self.columns) .field( - "filter_exprs", + "filters", &self - .filter_exprs + .filters .iter() .map(|e| e.to_string()) .collect::>(), @@ -130,48 +105,46 @@ impl fmt::Debug for TableFunctionBind { } impl<'a> TableInitInput<'a> { - pub fn bind_data(&self) -> &TableFunctionBind { - unsafe { &*self.input.bind_data.cast::() } + pub(crate) fn bind_data(&self) -> &BindState { + unsafe { &*self.input.bind_data.cast::() } + } + + pub fn filters(&self) -> Option<&TableFilterSetRef> { + let ptr = self.input.filters; + if ptr.is_null() { + None + } else { + Some(unsafe { TableFilterSet::borrow(ptr) }) + } } } -type ScanItem = VortexResult<(ArrayRef, Arc)>; -type DataSourceIterator = ThreadSafeIterator; - -pub struct TableFunctionGlobal { - iterator: DataSourceIterator, - batch_id: AtomicU64, - bytes_total: Arc, - bytes_read: AtomicU64, - file_index_column_pos: Option, - file_row_number_column_pos: Option, - - // Following 4 fields are used only in aggregate scans. - /// ArrayRef's scanned but not aggregated in "partials". - /// 0 means all arrays have been aggregated but output is not written. - /// u64::MAX means arrays have been aggregated and we've written output row - pending: Arc, - aggregates: Vec, - // Accumulated partials - partials: Mutex>>, - row_count: AtomicU64, +pub struct GlobalState { + pub projection: BoundExpression, + pub filter: Filter, + pub file_row_number_column_pos: Option, + + // Following fields are used only in aggregate scans. + /// Splits that are not merged into global partials + /// 0 means everything started is merged. + /// u64::MAX means output row is written. + pub pending: Arc, + pub aggregates: Vec, + pub aggregate_positions: Vec, + pub partials: Mutex>>, + pub row_count: AtomicU64, } -assert_impl_all!(TableFunctionGlobal: Send, Sync); +assert_impl_all!(GlobalState: Send, Sync); + +pub type Split = BoxFuture<'static, VortexResult>>; /// Per-thread scan state -pub struct TableFunctionLocal { - iterator: DataSourceIterator, - exporter: Option, - partition_index: u64, - file_index: usize, - // Aggregate scan accumulated partials. Empty for non-aggregate scan - partials: Vec>, -} +pub struct LocalState { + pub exporter: Option, + pub split: Option, -pub struct PartitionData { - pub partition_index: u64, - pub file_index_column_pos: Option, - pub file_index: usize, + /// Empty for non-aggregate scan + pub partials: Vec>, } #[derive(Clone)] @@ -185,207 +158,110 @@ pub(crate) enum ColumnAggregate { #[derive(Debug)] pub enum Cardinality { - /// Unknown number of rows - Unknown, /// The exact number of rows. Exact(u64), /// An estimate of the number of rows. Estimate(u64), } -// Called for every new query. For example, if there is a VIEW over *.vortex, -// and after a query another file is added matching the glob, for second query -// bind() will be called again. -pub fn bind(input: &BindInputRef, result: &mut BindResultRef) -> VortexResult { - let data_source = bind_multi_file_scan(input)?; - let column_fields = extract_schema_from_dtype(data_source.dtype())?; - for fields in &column_fields { - result.add_result_column(&fields.name, &fields.logical_type); +pub fn finalize_scan(global: &GlobalState, chunk: &mut DataChunkRef) -> VortexResult { + if global.aggregates.is_empty() { + return Ok(false); + } + // 0 means every produced array has been accumulated, u64::MAX means output is + // written. is_err() covers "still accumulating" and "already emitted" + if global + .pending + .compare_exchange(0, u64::MAX, Ordering::AcqRel, Ordering::Relaxed) + .is_err() + { + return Ok(false); } - Ok(TableFunctionBind { - data_source: Arc::new(data_source), - filter_exprs: vec![], - column_fields, - has_non_optional_filter: AtomicBool::new(false), - aggregates: vec![], - }) -} -pub fn init_global(init_input: &TableInitInput) -> VortexResult { - debug!(input=?init_input, "table function global input"); + let mut accumulators = global.partials.lock(); + let row_count = global.row_count.load(Ordering::Acquire) as i64; + let mut accum_iter = accumulators.iter_mut(); + for (idx, aggregate) in global.aggregates.iter().enumerate() { + let value = match aggregate { + ColumnAggregate::Real { .. } => { + let accum = accum_iter.next().vortex_expect("partial for real agg"); + let expected = chunk.get_vector_mut(idx).logical_type(); + aggregate_output_value(accum.finish()?, &expected)? + } + ColumnAggregate::CountStar => Value::from(row_count), + }; + chunk.get_vector_mut(idx).reference_value(&value); + } + chunk.set_len(1); + Ok(true) +} +pub fn init_global(init_input: &TableInitInput) -> VortexResult { let bind_data = init_input.bind_data(); + + let partials = build_partials(&bind_data.aggregates, &bind_data.columns, &bind_data.dtype)?; + + let mut file_row_number_column_pos = None; let column_ids = init_input.column_ids(); - let projection_ids = init_input.projection_ids(); + for (i, id) in column_ids.iter().enumerate() { + if *id == FILE_ROW_NUMBER_COLUMN_IDX { + file_row_number_column_pos = Some(i); + } + } - let Projection { - projection, - file_index_column_pos, - file_row_number_column_pos, - } = if bind_data.aggregates.is_empty() { - Projection::new(projection_ids, column_ids, &bind_data.column_fields) + let mut seen = HashMap::with_capacity(bind_data.aggregates.len()); + let mut aggregate_positions = Vec::with_capacity(bind_data.aggregates.len()); + for aggregate in &bind_data.aggregates { + let ColumnAggregate::Real { projection_id, .. } = aggregate else { + continue; + }; + let len = seen.len(); + let pos = *seen.entry(*projection_id).or_insert(len); + aggregate_positions.push(pos); + } + + let Projection(projection) = if bind_data.aggregates.is_empty() { + Projection::new(column_ids, &bind_data.columns) } else { - Projection::new_aggregate(&bind_data.aggregates, &bind_data.column_fields) + Projection::new_aggregate(&bind_data.aggregates, &bind_data.columns) }; - let Filter { - filter, - row_selection, - row_range, - file_selection, - file_range, - has_non_optional_filter, - } = Filter::new( - init_input.table_filter_set(), + let filter = Filter::new( + init_input.filters(), column_ids, - &bind_data.column_fields, - &bind_data.filter_exprs, - bind_data.data_source.dtype(), + &bind_data.columns, + &bind_data.filters, + &bind_data.dtype, )?; - - if has_non_optional_filter { - init_input - .bind_data() + if filter.has_non_optional_filter { + bind_data .has_non_optional_filter .store(true, Ordering::Relaxed); } debug!( %projection, - filter = filter + filter = filter.filter .as_ref() .map_or_else(|| "true".to_string(), |f| f.to_string()), - ?row_selection, - ?row_range, - ?file_selection, - ?file_range, + row_selection = ?filter.row_selection, + row_range = ?filter.row_range, "table function scan input" ); - let request = ScanRequest { + let projection = optimize_and_bind(projection, &bind_data.dtype)?; + Ok(GlobalState { projection, filter, - ordered: file_row_number_column_pos.is_some(), - selection: row_selection, - row_range, - partition_selection: file_selection, - partition_range: file_range, - limit: None, - }; - - let scan = RUNTIME.block_on(bind_data.data_source.scan(request))?; - - let num_workers = get_available_parallelism().unwrap_or(1); - - // We create an async bounded channel so that all thread-local workers can pull the next - // available array chunk regardless of which partition it came from. - let (tx, rx) = kanal::bounded_async(num_workers * 2); - - let pending = Arc::new(AtomicU64::new(0)); - let pending_producer = Arc::clone(&pending); - - // We drive one partition per worker thread. Each partition is driven as a spawned task - // that pushes array chunks into the shared channel as they are produced. This spawning - // allows all worker threads to drive the polling of all partitions, and then return the - // first available array chunk. - let stream = scan - .partitions() - .map(move |partition| { - let tx = tx.clone(); - let pending = Arc::clone(&pending_producer); - RUNTIME.handle().spawn(async move { - let partition = match partition { - Ok(partition) => partition, - Err(e) => { - let _ = tx.send(Err(e)).await; - return; - } - }; - - let cache = Arc::new(ConversionCache { - file_index: partition.index(), - ..Default::default() - }); - - let mut stream = match partition.execute() { - Ok(s) => s, - Err(e) => { - let _ = tx.send(Err(e)).await; - return; - } - }; - while let Some(item) = stream.next().await { - pending.fetch_add(1, Ordering::Relaxed); - if tx - .send(item.map(|a| (a, Arc::clone(&cache)))) - .await - .is_err() - { - // Exit early if the receiver has been dropped, which happens when the - // scan is complete or if an error has occurred in another partition. - return; - } - } - }) - }) - .buffer_unordered(num_workers); - - let iterator = RUNTIME.block_on_stream_thread_safe(|_handle| scan_driver_stream(stream, rx)); - - let aggregates = bind_data.aggregates.clone(); - let partials = build_partials( - &aggregates, - &bind_data.column_fields, - bind_data.data_source.dtype(), - )?; - - Ok(TableFunctionGlobal { - iterator, - batch_id: AtomicU64::new(0), - bytes_total: Arc::new(AtomicU64::new(0)), - bytes_read: AtomicU64::new(0), - file_index_column_pos, - file_row_number_column_pos, - pending, - aggregates, + aggregate_positions, + pending: Arc::new(AtomicU64::new(0)), + aggregates: bind_data.aggregates.clone(), partials: Mutex::new(partials), row_count: AtomicU64::new(0), + file_row_number_column_pos, }) } -fn scan_driver_stream(stream: S, rx: kanal::AsyncReceiver) -> ScanDriverStream -where - S: Stream + Send + 'static, -{ - ScanDriverStream { - driver: Some(stream.collect::<()>().boxed()), - rx: rx.into_stream().boxed(), - } -} - -struct ScanDriverStream { - driver: Option>, - rx: futures::stream::BoxStream<'static, ScanItem>, -} - -impl Stream for ScanDriverStream { - type Item = ScanItem; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - if let Some(driver) = this.driver.as_mut() - && driver.as_mut().poll(cx).is_ready() - { - this.driver = None; - } - - match this.rx.as_mut().poll_next(cx) { - Poll::Ready(None) if this.driver.is_some() => Poll::Pending, - poll => poll, - } - } -} - /// Dtype over which we accumulate fn aggregate_input_dtype(field: &DuckdbField, scope: &DType) -> VortexResult { match &field.projection_expr { @@ -417,10 +293,7 @@ fn build_partials( .collect() } -pub fn init_local( - bind_data: &TableFunctionBind, - global: &TableFunctionGlobal, -) -> TableFunctionLocal { +pub fn init_local(bind_data: &BindState, global: &GlobalState) -> LocalState { unsafe { use custom_labels::sys; @@ -436,25 +309,23 @@ pub fn init_local( CURRENT_LABELSET.set(key, value); } - let partials = build_partials( - &global.aggregates, - &bind_data.column_fields, - bind_data.data_source.dtype(), - ) - // if aggregate initialization produced an error, it would error in - // init_global, see "partials" initialization there - .vortex_expect("local state aggregate initialization failed"); - - TableFunctionLocal { - iterator: global.iterator.clone(), + let partials = build_partials(&global.aggregates, &bind_data.columns, &bind_data.dtype) + // if aggregate initialization produced an error, it would error in + // init_global, see "partials" initialization there + .vortex_expect("local state aggregate initialization failed"); + + LocalState { exporter: None, - partition_index: 0, - file_index: 0, partials, + split: None, } } -fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { +pub(crate) fn optimize_and_bind(expr: Expression, dtype: &DType) -> VortexResult { + expr.optimize_recursive(dtype)?.bind(dtype) +} + +pub(crate) fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let array_result = array.optimize_recursive(ctx.session())?; Ok(if let Some(array) = array_result.as_opt::() { array.into_owned() @@ -496,164 +367,13 @@ fn aggregate_output_value(scalar: Scalar, expected: &LogicalTypeRef) -> VortexRe } } -fn scan_aggregate( - local_state: &mut TableFunctionLocal, - global_state: &TableFunctionGlobal, - chunk: &mut DataChunkRef, -) -> VortexResult<()> { - let aggregates_len = global_state.aggregates.len(); - // seen[k] = output column for requested column k. - // If min(x), max(x), avg(y) are requested, seen = { 0: 0, 1: 1} - let mut seen: HashMap = HashMap::with_capacity(aggregates_len); - // positions[k] = column id for accumulator k - // If min(x), max(x), avg(y) are requested, positions = [0, 0, 1] - let mut positions: Vec = Vec::with_capacity(aggregates_len); - - for aggregate in &global_state.aggregates { - let ColumnAggregate::Real { projection_id, .. } = aggregate else { - continue; - }; - let len = seen.len(); - let pos = seen.entry_ref(projection_id).or_insert(len); - positions.push(*pos); - } - let has_count_star = local_state.partials.len() < aggregates_len; - - let mut ctx = SESSION.create_execution_ctx(); - loop { - let Some(result) = local_state.iterator.next() else { - // 0 means we're the last thread, u64::MAX means output is written. - // is_err() means CAS didn't succeed - if global_state - .pending - .compare_exchange(0, u64::MAX, Ordering::AcqRel, Ordering::Relaxed) - .is_err() - { - return Ok(()); - } - - let mut accumulators = global_state.partials.lock(); - let row_count = global_state.row_count.load(Ordering::Acquire) as i64; - let mut accum_iter = accumulators.iter_mut(); - for (idx, aggregate) in global_state.aggregates.iter().enumerate() { - let value = match aggregate { - ColumnAggregate::Real { .. } => { - let accum = accum_iter.next().vortex_expect("partial for real agg"); - let expected = chunk.get_vector_mut(idx).logical_type(); - aggregate_output_value(accum.finish()?, &expected)? - } - ColumnAggregate::CountStar => Value::from(row_count), - }; - chunk.get_vector_mut(idx).reference_value(&value); - } - chunk.set_len(1); - return Ok(()); - }; - let array = convert_result(result?.0, &mut ctx)?; - - for (i, partial) in positions.iter().zip(local_state.partials.iter_mut()) { - partial.accumulate(array.unmasked_field(*i), &mut ctx)?; - } - - { - let mut partials = global_state.partials.lock(); - for (global, local) in partials.iter_mut().zip(&mut local_state.partials) { - global.combine_partials(local.flush()?)?; - } - } - - if has_count_star { - global_state - .row_count - .fetch_add(array.len() as u64, Ordering::Relaxed); - } - global_state.pending.fetch_sub(1, Ordering::Release); - } -} - -pub fn scan( - local_state: &mut TableFunctionLocal, - global_state: &TableFunctionGlobal, - chunk: &mut DataChunkRef, -) -> VortexResult<()> { - if !local_state.partials.is_empty() { - return scan_aggregate(local_state, global_state, chunk); - } - - loop { - if local_state.exporter.is_none() { - let mut ctx = SESSION.create_execution_ctx(); - let Some(result) = local_state.iterator.next() else { - return Ok(()); - }; - let (array_result, conversion_cache) = result?; - local_state.file_index = conversion_cache.file_index; - let array_result = convert_result(array_result, &mut ctx)?; - - local_state.exporter = Some(ArrayExporter::try_new( - &array_result, - &conversion_cache, - ctx, - )?); - // Relaxed since there is no intra-instruction ordering required. - local_state.partition_index = global_state.batch_id.fetch_add(1, Ordering::Relaxed); - } - - let exporter = local_state - .exporter - .as_mut() - .vortex_expect("error: exporter missing"); - let has_more_data = exporter.export( - chunk, - global_state.file_index_column_pos, - global_state.file_row_number_column_pos, - )?; - - global_state - .bytes_read - .fetch_add(chunk.len(), Ordering::Relaxed); - - if !has_more_data { - // This exporter is fully consumed. - local_state.exporter = None; - local_state.partition_index = 0; - } else { - break; - } - } - - assert!(!chunk.is_empty()); - - if let Some(pos) = global_state.file_index_column_pos { - chunk - .get_vector_mut(pos) - .reference_value(&Value::from(local_state.file_index as u64)); - } - - Ok(()) -} - -/// Scan progress as a percentage (0.0–100.0). -pub fn table_scan_progress(global_state: &TableFunctionGlobal) -> f64 { - progress(&global_state.bytes_read, &global_state.bytes_total) -} - -/// Table filter pushdown is used for two tasks in duckdb: -/// -/// 1. Prune files based on filename or hive partitioning, see Parquet -/// filter pushdown. We don't use this because we do own file-level pruning -/// in FileStatsLayoutReader, and we don't support hive partitioning yet. -/// 2. Avoid reading unused file data. Filter expressions are pushed to Vortex, -/// converted to Vortex expressions and used during the scan. -/// Duckdb pushes a subset of expressions i.e. equality operators, and also -/// expressions which return true in pushdown_expression. pub fn pushdown_complex_filter( - bind_data: &mut TableFunctionBind, + bind_data: &mut BindState, expr: &ExpressionRef, ) -> VortexResult { debug!(%expr, "pushing down expression"); - let Some(expr) = try_from_bound_expression(expr, &bind_data.column_fields)? else { + let Some(expr) = try_from_bound_expression(expr, &bind_data.columns)? else { debug!(%expr, "failed to push down expression"); return Ok(false); }; @@ -685,16 +405,16 @@ pub fn pushdown_complex_filter( .store(true, Ordering::Relaxed); debug!(%expr, report_pushed, "pushed down expression"); - bind_data.filter_exprs.push(expr); + bind_data.filters.push(expr); Ok(report_pushed) } pub fn pushdown_projection_expression( - bind_data: &mut TableFunctionBind, + bind_data: &mut BindState, expr: &ExpressionRef, projection_id: usize, ) -> VortexResult { - let field = &bind_data.column_fields[projection_id]; + let field = &bind_data.columns[projection_id]; debug!(%expr, %projection_id, col_name=field.name, "pushing down projection expression"); match try_from_projection_expression(expr, field)? { None => { @@ -703,7 +423,13 @@ pub fn pushdown_projection_expression( } Some(vx_expr) => { debug!(%expr, "pushed down expression"); - bind_data.column_fields[projection_id].projection_expr = Some(vx_expr); + let Ok(out_dtype) = vx_expr.return_dtype(&bind_data.dtype) else { + return Ok(false); + }; + let field = &mut bind_data.columns[projection_id]; + field.logical_type = expr.return_type().to_owned(); + field.dtype = out_dtype; + field.projection_expr = Some(vx_expr); Ok(true) } } @@ -711,12 +437,12 @@ pub fn pushdown_projection_expression( fn can_push_projection_aggregate( aggregate: &PushedAggregate, - bind_data: &TableFunctionBind, + bind_data: &BindState, projection_id: u64, ) -> bool { let projection_id_usize: usize = projection_id.as_(); - let field = &bind_data.column_fields[projection_id_usize]; - let Ok(dtype) = aggregate_input_dtype(field, bind_data.data_source.dtype()) else { + let field = &bind_data.columns[projection_id_usize]; + let Ok(dtype) = aggregate_input_dtype(field, &bind_data.dtype) else { return false; }; @@ -760,18 +486,32 @@ fn can_push_projection_aggregate( /// same columns. If we return true, optimized pass expands output to N columns, /// e.g. min(x), max(x) turns into min(x0), max(x1), 2 columns in output. pub fn pushdown_projection_aggregates( - bind_data: &mut TableFunctionBind, + bind_data: &mut BindState, input: &AggregatePushdownInputRef, ) -> VortexResult { + // TODO + return Ok(false); + let len = input.len(); let mut aggregates = Vec::with_capacity(len); + let mut outputs = Vec::with_capacity(len); let mut has_non_count_star = false; debug!(%len, "pushing down projection aggregates"); for i in 0..len { - let Some(aggregate) = try_push_projection_aggregate(bind_data, input.get(i), i)? else { + let expression = input.get(i); + let output_type = expression.expr.return_type().to_owned(); + let Some(aggregate) = try_push_projection_aggregate(bind_data, expression, i)? else { return Ok(false); }; + let name = match &aggregate { + ColumnAggregate::CountStar => "count_star()".to_string(), + ColumnAggregate::Real { projection_id, .. } => { + let id: usize = projection_id.as_(); + bind_data.columns[id].name.clone() + } + }; + outputs.push((name, output_type)); has_non_count_star |= matches!(aggregate, ColumnAggregate::Real { .. }); aggregates.push(aggregate); } @@ -784,7 +524,7 @@ pub fn pushdown_projection_aggregates( } fn try_push_projection_aggregate( - bind_data: &TableFunctionBind, + bind_data: &BindState, aggregate: AggregateExpression<'_>, i: usize, ) -> VortexResult> { @@ -810,41 +550,6 @@ fn try_push_projection_aggregate( })) } -/// Get column-wise statistics. Available only if we're reading a single file. -pub fn statistics(bind_data: &TableFunctionBind, column_index: usize) -> Option { - // Aggregate output columns hold data we don't have in statistics - if !bind_data.aggregates.is_empty() { - return None; - } - let children = bind_data.data_source.children(); - // Otherwise we'd have to open all files eagerly which is a performance - // regression. Duckdb's Parquet reader only gets metadata for multiple - // files with a UNION BY NAME and we don't support it (yet) - // See duckdb/common/multi_file/multi_file_function.hpp#L691 - if children.len() != 1 { - return None; - } - let MultiLayoutChild::Opened { reader, .. } = &children[0] else { - return None; - }; - let stats_sets = reader - .as_any() - .downcast_ref::()? - .file_stats() - .stats_sets(); - // Columns with pushed projection expression output expression results, - // and not column values - if bind_data.column_fields[column_index] - .projection_expr - .is_some() - { - return None; - } - let dtype = bind_data.column_fields[column_index].dtype.clone(); - let stats_aggregate = ColumnStatisticsAggregate::new(&stats_sets[column_index]); - Some(ColumnStatistics::from(&stats_aggregate, dtype)) -} - /// Duckdb requires post-filter cardinality estimates, otherwise join planner /// may flip join sides which is a huge regression for some queries i.e. 1000x /// for tpcds 85. @@ -855,51 +560,30 @@ pub fn statistics(bind_data: &TableFunctionBind, column_index: usize) -> Option< /// duckdb uses is a 0.2 filter if there is any non-optional filter. We mimic it /// here. const DEFAULT_SELECTIVITY: f64 = 0.2; -pub fn cardinality(bind_data: &TableFunctionBind) -> Cardinality { +pub fn cardinality(bind_data: &BindState, file_count: u64) -> Cardinality { // If we're doing an aggregate scan, we don't change output cardinality to // 1 as we want duckdb to do our aggregation in parallel. That may look // counterintuitive in the plan, though. let has_non_optional_filter = bind_data.has_non_optional_filter.load(Ordering::Relaxed); - match bind_data.data_source.row_count() { - Precision::Exact(v) => { - if !has_non_optional_filter { - return Cardinality::Exact(v); - } - let post_cardinality = v as f64 * DEFAULT_SELECTIVITY; - let post_cardinality: u64 = post_cardinality.as_(); - Cardinality::Estimate(max(1, post_cardinality)) - } - Precision::Inexact(v) => { - if !has_non_optional_filter { - return Cardinality::Estimate(v); - } - let post_cardinality = v as f64 * DEFAULT_SELECTIVITY; - let post_cardinality: u64 = post_cardinality.as_(); - Cardinality::Estimate(max(1, post_cardinality)) - } - Precision::Absent => Cardinality::Unknown, - } -} - -/// Duckdb requests this function after exporting the chunk. We answer with -/// partition_index we have exported as well as information about constant -/// columns in this partition. As data is partitioned by array exporters, in -/// each partition ~ exported array file_index is constant. -pub fn get_partition_data( - global_init_data: &TableFunctionGlobal, - local_init_data: &mut TableFunctionLocal, -) -> PartitionData { - PartitionData { - partition_index: local_init_data.partition_index, - file_index_column_pos: global_init_data.file_index_column_pos, - file_index: local_init_data.file_index, + let total = bind_data + .first_file_row_count + .saturating_mul(max(file_count, 1)); + if !has_non_optional_filter { + return if file_count <= 1 { + Cardinality::Exact(total) + } else { + Cardinality::Estimate(total) + }; } + let post_cardinality = total as f64 * DEFAULT_SELECTIVITY; + let post_cardinality: u64 = post_cardinality.as_(); + Cardinality::Estimate(max(1, post_cardinality)) } -pub fn to_string(bind_data: &TableFunctionBind, map: &mut DuckdbStringMapRef) { +pub fn to_string(bind_data: &BindState, map: &mut DuckdbStringMapRef) { map.push("Function", "Vortex Scan"); - if !bind_data.filter_exprs.is_empty() { - let mut filters = bind_data.filter_exprs.iter().map(|f| format!("{f}")); + if !bind_data.filters.is_empty() { + let mut filters = bind_data.filters.iter().map(|f| format!("{f}")); map.push("Filters", &filters.join("\n")); } @@ -913,10 +597,7 @@ pub fn to_string(bind_data: &TableFunctionBind, map: &mut DuckdbStringMapRef) { aggregate, } => { let projection_id: usize = projection_id.as_(); - format!( - "{aggregate}({})", - bind_data.column_fields[projection_id].name - ) + format!("{aggregate}({})", bind_data.columns[projection_id].name) } ColumnAggregate::CountStar => "count(*)".to_string(), }) @@ -928,7 +609,7 @@ pub fn to_string(bind_data: &TableFunctionBind, map: &mut DuckdbStringMapRef) { } let projections = bind_data - .column_fields + .columns .iter() .filter_map(|field| { field @@ -941,57 +622,3 @@ pub fn to_string(bind_data: &TableFunctionBind, map: &mut DuckdbStringMapRef) { map.push("SELECT projections", &projections); } } - -fn progress(bytes_read: &AtomicU64, bytes_total: &AtomicU64) -> f64 { - let read = bytes_read.load(Ordering::Relaxed); - let mut total = bytes_total.load(Ordering::Relaxed); - total += (total == 0) as u64; - read as f64 / total as f64 * 100. -} - -#[cfg(test)] -mod tests { - use std::sync::atomic::AtomicU64; - use std::sync::atomic::Ordering::Relaxed; - use std::task::Poll; - - use crate::RUNTIME; - use crate::table_function::progress; - use crate::table_function::scan_driver_stream; - - #[test] - fn test_table_scan_progress() { - let bytes_total = AtomicU64::new(100); - let bytes_read = AtomicU64::new(0); - - assert_eq!(progress(&bytes_read, &bytes_total), 0.0); - - bytes_read.fetch_add(100, Relaxed); - assert_eq!(progress(&bytes_read, &bytes_total), 100.); - - bytes_total.fetch_add(100, Relaxed); - assert!((progress(&bytes_read, &bytes_total) - 50.).abs() < f64::EPSILON); - } - - #[test] - fn scan_driver_panic_propagates_through_iterator() { - let (tx, rx) = kanal::bounded_async(1); - let _tx = tx; - let stream = futures::stream::poll_fn(|_| -> Poll> { - panic!("duckdb scan driver panic"); - }); - - let mut iter = - RUNTIME.block_on_stream_thread_safe(|_handle| scan_driver_stream(stream, rx)); - let panic = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| iter.next())) { - Ok(_) => panic!("driver panic must propagate through iterator"), - Err(panic) => panic, - }; - let message = panic - .downcast_ref::<&'static str>() - .copied() - .or_else(|| panic.downcast_ref::().map(String::as_str)) - .unwrap_or(""); - assert!(message.contains("duckdb scan driver panic")); - } -}