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..fd7bf8a4ee3 --- /dev/null +++ b/vortex-duckdb/cpp/include/multi_file_reader.hpp @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "data.hpp" +#include "duckdb/common/multi_file/multi_file_function.hpp" + +using namespace duckdb; + +struct VortexBindData final : TableFunctionData { + VortexBindData() = default; + unique_ptr 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; + + 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(); + } + + // Prune reader if file statistics prove false for "table_filters" + 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(); + } + + void BindReader(ClientContext &, + vector &return_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; + unique_ptr ffi_file_scan; + vector virtual_ids; + + inline void AddVirtualColumn(column_t virtual_column_id) override { + virtual_ids.push_back(virtual_column_id); + } + + void StartScan(GlobalTableFunctionState &gstate); + + // Returns false when file is exhausted + bool TryInitializeScan(ClientContext &context, + GlobalTableFunctionState &gstate, + LocalTableFunctionState &lstate) override; + + AsyncResult Scan(ClientContext &context, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state, + DataChunk &chunk) override; + + inline void FinishFile(ClientContext &, GlobalTableFunctionState &) override { + } + + double GetProgressInFile(ClientContext &context) override; + + unique_ptr 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..05cf0ab122a 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -10,14 +10,8 @@ extern "C" { #endif -// Info passed into the bind callback. The callback should set error or else add result columns. -typedef struct duckdb_vx_tfunc_bind_input_ *duckdb_vx_tfunc_bind_input; typedef struct duckdb_vx_tfunc_bind_result_ *duckdb_vx_tfunc_bind_result; -// Fetch a parameter from the bind info. -// The caller is responsible for freeing the value using duckdb_value_free. -duckdb_value duckdb_vx_tfunc_bind_input_get_parameter(duckdb_vx_tfunc_bind_input ffi_input, size_t index); - // Add a result column to the bind info. void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_result ffi_result, const char *name_str, @@ -45,14 +39,11 @@ typedef struct { * after filter pushdown and filter pruning. May be empty, in which case * column_ids should be used. * Indices in this list reference values from column_ids. I.e. if - * column_ids=[1,5,6], projection_ids=[1], output column should be * column_ids[1] = 5 * * Example usage: * https://github.com/duckdb/duckdb/blob/dc11eadd8f0a7c600f0034810706605ebe10d5b9/src/include/duckdb/function/table_function.hpp#L147 */ - const idx_t *projection_ids; - size_t projection_ids_count; duckdb_vx_table_filter_set filters; duckdb_client_context client_context; @@ -76,16 +67,6 @@ typedef struct { bool has_null; } duckdb_column_statistics; -const idx_t INVALID_IDX = UINT64_MAX; - -typedef struct { - idx_t partition_index; - // Either INVALID_IDX or position of column in output for file_index column - size_t file_index_column_pos; - // File index for the exported partition. - size_t file_index; -} duckdb_vx_partition_data; - duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db); typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input; 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..cca002f9f76 --- /dev/null +++ b/vortex-duckdb/cpp/multi_file_reader.cpp @@ -0,0 +1,380 @@ +// 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" + +unique_ptr VortexBindData::Copy() const { + auto result = make_uniq(); + if (ffi_bind_data) { + const auto copied = duckdb_table_function_bind_data_clone(ffi_bind_data->DataPtr()); + result->ffi_bind_data = unique_ptr(reinterpret_cast(copied)); + } + 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) { + if (gstate.global_state) { + auto &reader = reader_data.reader->Cast(); + duckdb_vx_error error_out = nullptr; + const bool skip = duckdb_table_function_file_should_skip( + gstate.global_state->Cast().ffi_global_state->DataPtr(), + reader.ffi_file->DataPtr(), + &error_out); + if (error_out) { + throw InvalidInputException(IntoErrString(error_out)); + } + if (skip) { + return ReaderInitializeType::SKIP_READING_FILE; + } + } + + reader_data.reader->columns = global_columns; + + auto init = MultiFileReader::InitializeReader(reader_data, + bind_data, + global_columns, + global_column_ids, + table_filters, + context, + gstate); + /* + * Start the scan early. Default MultiFileReader model is "1 split per + * worker". This works poorly for queries on multiple files, think + * clickbench. What we do here is we launch prefetch/decode for the file + * before reader has been bound. This means by the time a worker starts + * processing next file, next file' splits are already decoded. + * + * Another performance degradation if we don't start the scan here is that + * TryInitializeScan is called under global lock, while InitializeReader is + * called under a file-local lock. + * + * With duckdb 2.0 this code will migrate to readahead queue. + */ + if (init != ReaderInitializeType::SKIP_READING_FILE && gstate.global_state) { + reader_data.reader->Cast().StartScan(*gstate.global_state); + } + return init; +} + +void VortexReaderInterface::BindReader(ClientContext &context, + vector &return_types, + vector &names, + MultiFileBindData &bind_data) { + auto &bind = bind_data.bind_data->Cast(); + BaseFileReaderOptions options; + bind_data.reader_bind = bind_data.multi_file_reader->BindReader(context, + return_types, + names, + *bind_data.file_list, + bind_data, + options, + bind_data.file_options); + + auto &initial_reader = bind_data.initial_reader->Cast(); + duckdb_vx_error error_out = nullptr; + duckdb_vx_data ffi_bind_data = duckdb_table_function_bind(initial_reader.ffi_file->DataPtr(), &error_out); + if (error_out) { + throw BinderException(IntoErrString(error_out)); + } + bind.ffi_bind_data = unique_ptr(reinterpret_cast(ffi_bind_data)); +} + +unique_ptr +VortexReaderInterface::InitializeGlobalState(ClientContext &context, + MultiFileBindData &bind_data, + MultiFileGlobalState &input) { + void *const ffi_bind = bind_data.bind_data->Cast().ffi_bind_data->DataPtr(); + + // Pushed projection expressions and casts change what the scan outputs + vector types; + vector names; + VortexBindResult schema = {types, names}; + duckdb_table_function_bind_schema(ffi_bind, reinterpret_cast(&schema)); + bind_data.columns = MultiFileColumnDefinition::ColumnsFromNamesAndTypes(names, types); + + 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(); + } + + 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(); + duckdb_vx_data ffi_local_state = + duckdb_table_function_init_local(global.ffi_bind_data, global.ffi_global_state->DataPtr()); + + auto result = make_uniq(); + result->ffi_local_state = unique_ptr(reinterpret_cast(ffi_local_state)); + return result; +} + +static shared_ptr OpenReader(const OpenFileInfo &file, idx_t file_idx) { + duckdb_vx_error error_out = nullptr; + duckdb_vx_data ffi_file = + duckdb_table_function_file_open(file.path.c_str(), file.path.size(), file_idx, &error_out); + if (error_out) { + throw IOException(IntoErrString(error_out)); + } + return make_shared_ptr(file, unique_ptr(reinterpret_cast(ffi_file))); +} + +shared_ptr VortexReaderInterface::CreateReader(ClientContext &, + GlobalTableFunctionState &, + const OpenFileInfo &file, + idx_t file_idx, + const MultiFileBindData &) { + return OpenReader(file, file_idx); +} + +shared_ptr VortexReaderInterface::CreateReader(ClientContext &, + const OpenFileInfo &file, + BaseFileReaderOptions &, + const MultiFileOptions &) { + auto reader = OpenReader(file, 0); + + vector types; + vector names; + VortexBindResult schema = {types, names}; + duckdb_vx_error error_out = nullptr; + duckdb_table_function_file_schema(reader->Cast().ffi_file->DataPtr(), + reinterpret_cast(&schema), + &error_out); + if (error_out) { + throw IOException(IntoErrString(error_out)); + } + reader->Cast().columns = + MultiFileColumnDefinition::ColumnsFromNamesAndTypes(names, types); + return reader; +} + +unique_ptr VortexReaderInterface::GetCardinality(const MultiFileBindData &bind_data, + idx_t file_count) { + const void *const ffi_bind = bind_data.bind_data->Cast().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; +} + +void VortexBaseReader::StartScan(GlobalTableFunctionState &gstate) { + if (ffi_file_scan) { + return; + } + auto &global = gstate.Cast(); + + const idx_t real_columns = columns.size() - virtual_ids.size(); + vector local_column_ids; + local_column_ids.reserve(column_ids.size()); + for (idx_t i = 0; i < column_ids.size(); i++) { + const idx_t local_id = column_ids[MultiFileLocalIndex(i)]; + local_column_ids.push_back(local_id >= real_columns ? virtual_ids[local_id - real_columns] + : local_id); + } + + duckdb_vx_error error_out = nullptr; + duckdb_vx_data ffi_scan = + duckdb_table_function_file_start_scan(global.ffi_bind_data, + global.ffi_global_state->DataPtr(), + ffi_file->DataPtr(), + local_column_ids.data(), + local_column_ids.size(), + reinterpret_cast(filters.get()), + &error_out); + if (error_out) { + throw InvalidInputException(IntoErrString(error_out)); + } + ffi_file_scan = unique_ptr(reinterpret_cast(ffi_scan)); +} + +bool VortexBaseReader::TryInitializeScan(ClientContext &, + GlobalTableFunctionState &gstate, + LocalTableFunctionState &) { + StartScan(gstate); + return duckdb_table_function_file_has_work(ffi_file_scan->DataPtr()); +} + +AsyncResult VortexBaseReader::Scan(ClientContext &, + GlobalTableFunctionState &global_state, + LocalTableFunctionState &local_state, + DataChunk &chunk) { + auto &local = local_state.Cast(); + + duckdb_data_chunk ffi_chunk = reinterpret_cast(&chunk); + duckdb_vx_error error_out = nullptr; + duckdb_table_function_file_scan(ffi_file_scan->DataPtr(), + global_state.Cast().ffi_global_state->DataPtr(), + local.ffi_local_state->DataPtr(), + ffi_chunk, + &error_out); + if (error_out) { + throw InvalidInputException(IntoErrString(error_out)); + } + + if (chunk.size() == 0) { + return SourceResultType::FINISHED; + } + return SourceResultType::HAVE_MORE_OUTPUT; +} + +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 &gstate, + DataChunk &output) { + auto &global = gstate.Cast(); + duckdb_data_chunk ffi_chunk = reinterpret_cast(&output); + duckdb_vx_error error_out = nullptr; + const bool filled = + duckdb_table_function_finalize_scan(global.ffi_global_state->DataPtr(), ffi_chunk, &error_out); + if (error_out) { + throw InvalidInputException(IntoErrString(error_out)); + } + 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_table_function_file_statistics(ffi_file->DataPtr(), name.c_str(), name.size(), &statistics)) { + return nullptr; + } + for (const auto &column : columns) { + if (column.name != name) { + continue; + } + const LogicalType &type = column.type; + 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 nullptr; +} + +double VortexBaseReader::GetProgressInFile(ClientContext &) { + if (!ffi_file_scan) { + return 0.0; + } + return duckdb_table_function_file_progress(ffi_file_scan->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..36b7ce785c6 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,28 +93,6 @@ 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, const char *name_str, size_t name_len, @@ -306,51 +100,13 @@ extern "C" void duckdb_vx_tfunc_bind_result_add_column(duckdb_vx_tfunc_bind_resu D_ASSERT(ffi_result); D_ASSERT(name_str); D_ASSERT(ffi_type); - const VortexBindResults &result = *reinterpret_cast(ffi_result); + const 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..87d707bac15 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -19,18 +19,6 @@ extern "C" { extern void duckdb_table_function_to_string(const void *bind_data, 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, duckdb_vx_expr expr, @@ -47,16 +35,11 @@ bool duckdb_table_function_pushdown_projection_aggregates(void *bind_data, 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); - extern bool duckdb_table_function_pushdown_expression(duckdb_vx_expr expr); extern void duckdb_table_function_cardinality(const void *bind_data, + uint64_t file_count, duckdb_vx_node_statistics *node_stats_out); extern @@ -68,10 +51,60 @@ duckdb_vx_data duckdb_table_function_init_local(const void *bind_data, void *global_init_data); extern -duckdb_vx_data duckdb_table_function_bind(duckdb_vx_tfunc_bind_input bind_input, - duckdb_vx_tfunc_bind_result bind_result, +duckdb_vx_data duckdb_table_function_bind(const void *first_file, duckdb_vx_error *error_out); +extern +void duckdb_table_function_bind_schema(const void *bind_data, + duckdb_vx_tfunc_bind_result schema_result); + +extern +duckdb_vx_data duckdb_table_function_file_open(const char *file_path, + size_t file_path_len, + uint64_t file_index, + duckdb_vx_error *error_out); + +extern +void duckdb_table_function_file_schema(const void *file, + duckdb_vx_tfunc_bind_result schema_result, + duckdb_vx_error *error_out); + +extern +bool duckdb_table_function_file_statistics(const void *file, + const char *column_name, + size_t column_name_len, + duckdb_column_statistics *stats_out); + +extern +bool duckdb_table_function_file_should_skip(const void *global_init_data, + const void *file, + duckdb_vx_error *error_out); + +extern +duckdb_vx_data duckdb_table_function_file_start_scan(const void *bind_data, + void *global_init_data, + const void *file, + const uint64_t *column_ids, + size_t column_ids_count, + duckdb_vx_table_filter_set filters, + duckdb_vx_error *error_out); + +extern bool duckdb_table_function_file_has_work(const void *file_scan_data); + +extern +void duckdb_table_function_file_scan(const void *file_scan_data, + void *global_init_data, + void *local_init_data, + duckdb_data_chunk output, + duckdb_vx_error *error_out); + +extern double duckdb_table_function_file_progress(const void *file_scan_data); + +extern +bool duckdb_table_function_finalize_scan(void *global_init_data, + duckdb_data_chunk output, + duckdb_vx_error *error_out); + extern duckdb_vx_data duckdb_table_function_bind_data_clone(const void *bind_data); extern 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/duckdb/bind_input.rs b/vortex-duckdb/src/duckdb/bind_input.rs index 1049a8065f6..fa6e24ca42f 100644 --- a/vortex-duckdb/src/duckdb/bind_input.rs +++ b/vortex-duckdb/src/duckdb/bind_input.rs @@ -3,24 +3,8 @@ 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, |_| {}); impl BindResultRef { 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/ffi.rs b/vortex-duckdb/src/ffi.rs index 2a5e9316434..930611da997 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,21 +28,29 @@ use crate::duckdb::LogicalTypeRef; use crate::duckdb::TableInitInput; use crate::duckdb::try_or; use crate::duckdb::try_or_null; +use crate::file_reader::FileReader; +use crate::file_reader::FileScan; +use crate::file_reader::file_has_work; +use crate::file_reader::file_open; +use crate::file_reader::file_progress; +use crate::file_reader::file_scan; +use crate::file_reader::file_schema; +use crate::file_reader::file_should_skip; +use crate::file_reader::file_start_scan; +use crate::file_reader::file_statistics; 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::bind_schema; 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)] @@ -57,50 +64,6 @@ unsafe extern "C-unwind" fn duckdb_table_function_to_string( 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; -} - #[unsafe(no_mangle)] unsafe extern "C-unwind" fn duckdb_table_function_pushdown_complex_filter( bind_data: *mut c_void, @@ -142,33 +105,6 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_pushdown_projection_aggreg }) } -#[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(), - )); - }, - } -} - #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_pushdown_expression( expr: cpp::duckdb_vx_expr, @@ -179,6 +115,7 @@ 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, + file_count: u64, node_stats_out: *mut cpp::duckdb_vx_node_statistics, ) { let bind_data = unsafe { bind_data.cast::().as_ref() } @@ -186,8 +123,7 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_cardinality( let node_stats = unsafe { node_stats_out.as_mut() }.vortex_expect("node_stats_out null pointer"); - match cardinality(bind_data) { - Cardinality::Unknown => {} + match cardinality(bind_data, file_count) { Cardinality::Exact(c) => { node_stats.has_estimated_cardinality = true; node_stats.estimated_cardinality = c as _; @@ -237,19 +173,168 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_init_local( #[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, + first_file: *const c_void, 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("file null pointer"); try_or_null(error_out, || { - let bind_data = bind(&bind_input, &mut bind_result)?; + let bind_data = bind(first_file)?; Ok(Data::from(Box::new(bind_data)).as_ptr()) }) } +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_bind_schema( + bind_data: *const c_void, + schema_result: cpp::duckdb_vx_tfunc_bind_result, +) { + let bind_data = unsafe { bind_data.cast::().as_ref() } + .vortex_expect("bind_data null pointer"); + let schema_result = unsafe { BindResult::borrow_mut(schema_result) }; + bind_schema(bind_data, schema_result); +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_open( + file_path: *const c_char, + file_path_len: usize, + file_index: u64, + error_out: *mut cpp::duckdb_vx_error, +) -> cpp::duckdb_vx_data { + let path_bytes = unsafe { std::slice::from_raw_parts(file_path.cast::(), file_path_len) }; + let file_path = String::from_utf8_lossy(path_bytes).into_owned(); + + try_or_null(error_out, || { + let file = file_open(&file_path, file_index)?; + Ok(Data::from(Box::new(file)).as_ptr()) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_schema( + file: *const c_void, + schema_result: cpp::duckdb_vx_tfunc_bind_result, + error_out: *mut cpp::duckdb_vx_error, +) { + let file = unsafe { file.cast::().as_ref() }.vortex_expect("file null pointer"); + let schema_result = unsafe { BindResult::borrow_mut(schema_result) }; + try_or(error_out, || file_schema(file, schema_result)) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_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("file 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) = file_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; + true +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_should_skip( + global_init_data: *const c_void, + file: *const c_void, + error_out: *mut cpp::duckdb_vx_error, +) -> bool { + let global_init_data = unsafe { global_init_data.cast::().as_ref() } + .vortex_expect("global_init_data null pointer"); + let file = unsafe { file.cast::().as_ref() }.vortex_expect("file null pointer"); + try_or(error_out, || file_should_skip(global_init_data, file)) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_start_scan( + bind_data: *const c_void, + global_init_data: *mut c_void, + file: *const c_void, + column_ids: *const u64, + column_ids_count: usize, + filters: cpp::duckdb_vx_table_filter_set, + error_out: *mut cpp::duckdb_vx_error, +) -> cpp::duckdb_vx_data { + let bind_data = unsafe { bind_data.cast::().as_ref() } + .vortex_expect("bind_data null pointer"); + let file = unsafe { file.cast::().as_ref() }.vortex_expect("file null pointer"); + let column_ids = if column_ids_count == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(column_ids, column_ids_count) } + }; + + let global_init_data = unsafe { global_init_data.cast::().as_ref() } + .vortex_expect("global_init_data null pointer"); + try_or_null(error_out, || { + let scan = file_start_scan(bind_data, global_init_data, file, column_ids, filters)?; + Ok(Data::from(Box::new(scan)).as_ptr()) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_has_work( + file_scan_data: *const c_void, +) -> bool { + let scan = unsafe { file_scan_data.cast::().as_ref() } + .vortex_expect("file_scan null pointer"); + file_has_work(scan) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_scan( + file_scan_data: *const c_void, + 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 scan = unsafe { file_scan_data.cast::().as_ref() } + .vortex_expect("file_scan null pointer"); + 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) }; + try_or(error_out, || { + file_scan(scan, global_init_data, local_init_data, data_chunk) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_file_progress( + file_scan_data: *const c_void, +) -> f64 { + let scan = unsafe { file_scan_data.cast::().as_ref() } + .vortex_expect("file_scan null pointer"); + file_progress(scan) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_finalize_scan( + global_init_data: *mut c_void, + output: cpp::duckdb_data_chunk, + error_out: *mut cpp::duckdb_vx_error, +) -> bool { + let global_init_data = unsafe { global_init_data.cast::().as_ref() } + .vortex_expect("global_init_data null pointer"); + let data_chunk = unsafe { DataChunk::borrow_mut(output) }; + try_or(error_out, || finalize_scan(global_init_data, data_chunk)) +} + #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_bind_data_clone( bind_data: *const c_void, diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs new file mode 100644 index 00000000000..4b10fc3722c --- /dev/null +++ b/vortex-duckdb/src/file_reader.rs @@ -0,0 +1,412 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use futures::FutureExt; +use kanal::AsyncReceiver; +use object_store::registry::ObjectStoreRegistry; +use url::Url; +use vortex::array::ArrayRef; +use vortex::array::VortexSessionExecute as _; +use vortex::array::arrays::struct_::StructArrayExt as _; +use vortex::cloud::Registry; +use vortex::dtype::DType; +use vortex::error::VortexResult; +use vortex::error::vortex_err; +use vortex::expr::BoundExpression; +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::io::runtime::Task; +use vortex::layout::LayoutReaderRef; +use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::mask::Mask; +use vortex::scan::selection::Selection; + +use crate::RUNTIME; +use crate::SESSION; +use crate::column_statistics::ColumnStatistics; +use crate::column_statistics::ColumnStatisticsAggregate; +use crate::cpp; +use crate::duckdb::BindResultRef; +use crate::duckdb::DataChunkRef; +use crate::duckdb::TableFilterSet; +use crate::exporter::ArrayExporter; +use crate::exporter::ConversionCache; +use crate::projection::FILE_INDEX_COLUMN_IDX; +use crate::projection::FILE_ROW_NUMBER_COLUMN_IDX; +use crate::projection::Filter; +use crate::projection::extract_schema_from_dtype; +use crate::table_function::TableFunctionBind; +use crate::table_function::TableFunctionGlobal; +use crate::table_function::TableFunctionLocal; +use crate::table_function::convert_result; +use crate::table_function::optimize_and_bind; + +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 FileReader { + reader: LayoutReaderRef, + pub(crate) dtype: DType, + pub(crate) row_count: u64, + file_index: u64, +} + +type ScanItem = VortexResult<(ArrayRef, Arc)>; + +pub struct FileScan { + receiver: AsyncReceiver, + exhausted: AtomicBool, + total_splits: u64, + delivered: AtomicU64, + file_row_number_column_pos: Option, + file_index_column_pos: Option, + aggregate_positions: Vec, + _driver: Task<()>, +} + +async fn open_reader(file_path: String, file_index: u64) -> VortexResult { + let url = parse_uri_or_path(&file_path)?; + let (fs, path) = resolve_filesystem(&url)?; + let source = fs.open_read(&path).await?; + let vortex_file = open_cached(&SESSION, source, &path, None, &|options| options).await?; + let reader = vortex_file.layout_reader()?; + Ok(FileReader { + dtype: reader.dtype().clone(), + row_count: reader.row_count(), + reader, + file_index, + }) +} + +pub fn file_open(file_path: &str, file_index: u64) -> VortexResult { + RUNTIME.block_on(open_reader(file_path.to_owned(), file_index)) +} + +pub fn file_schema(file: &FileReader, result: &mut BindResultRef) -> VortexResult<()> { + for field in extract_schema_from_dtype(&file.dtype)? { + result.add_result_column(&field.name, &field.logical_type); + } + Ok(()) +} + +pub fn file_statistics(file: &FileReader, column_name: &str) -> Option { + let stats_reader = file + .reader + .as_any() + .downcast_ref::()?; + let stats_sets = stats_reader.file_stats().stats_sets(); + + let DType::Struct(fields, _) = &file.dtype else { + return None; + }; + let index = fields + .names() + .iter() + .position(|name| name.as_ref() == column_name)?; + let dtype = fields.field_by_index(index)?; + + let stats_aggregate = ColumnStatisticsAggregate::new(stats_sets.get(index)?); + Some(ColumnStatistics::from(&stats_aggregate, dtype)) +} + +pub struct ScanPruning { + filter: Option, + file_selection: Selection, + file_range: Option>, +} + +fn convert_filter( + bind: &TableFunctionBind, + column_ids: &[u64], + filters: cpp::duckdb_vx_table_filter_set, +) -> VortexResult { + let table_filter_set = if filters.is_null() { + None + } else { + Some(unsafe { TableFilterSet::borrow(filters) }) + }; + Filter::new( + table_filter_set, + column_ids, + bind.column_fields.as_slice(), + &bind.filter_exprs, + &bind.dtype, + ) +} + +impl ScanPruning { + pub fn new( + bind: &TableFunctionBind, + column_ids: &[u64], + filters: cpp::duckdb_vx_table_filter_set, + ) -> VortexResult> { + if filters.is_null() && bind.filter_exprs.is_empty() { + return Ok(None); + } + let converted = convert_filter(bind, column_ids, filters)?; + let filter = converted + .filter + .map(|expr| optimize_and_bind(expr, &bind.dtype)) + .transpose()?; + if filter.is_none() + && matches!(converted.file_selection, Selection::All) + && converted.file_range.is_none() + { + return Ok(None); + } + Ok(Some(Self { + filter, + file_selection: converted.file_selection, + file_range: converted.file_range, + })) + } +} + +pub fn file_should_skip(global: &TableFunctionGlobal, file: &FileReader) -> VortexResult { + let Some(pruning) = global.pruning.as_ref() else { + return Ok(false); + }; + let index = file.file_index; + let excluded = match &pruning.file_selection { + Selection::IncludeByIndex(buffer) => buffer.as_slice().binary_search(&index).is_err(), + Selection::ExcludeByIndex(buffer) => buffer.as_slice().binary_search(&index).is_ok(), + _ => false, + }; + if excluded + || pruning + .file_range + .as_ref() + .is_some_and(|r| !r.contains(&index)) + { + return Ok(true); + } + + let Some(filter) = &pruning.filter else { + return Ok(false); + }; + let row_range = 0..file.row_count; + let mask = Mask::new_true(usize::try_from(file.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), + } +} + +const FILE_CHANNEL_CAPACITY: usize = 16; + +pub fn file_start_scan( + bind: &TableFunctionBind, + global: &TableFunctionGlobal, + file: &FileReader, + column_ids: &[u64], + filters: cpp::duckdb_vx_table_filter_set, +) -> VortexResult { + let file_row_number_column_pos = column_ids + .iter() + .position(|&id| id == FILE_ROW_NUMBER_COLUMN_IDX); + let file_index_column_pos = column_ids + .iter() + .position(|&id| id == FILE_INDEX_COLUMN_IDX); + + let Filter { + filter, + row_selection, + row_range, + has_non_optional_filter, + .. + } = convert_filter(bind, column_ids, filters)?; + if has_non_optional_filter { + bind.has_non_optional_filter.store(true, Ordering::Relaxed); + } + + let filter = filter + .map(|expr| optimize_and_bind(expr, &bind.dtype)) + .transpose()?; + + let mut builder = ScanBuilder::new(SESSION.clone(), Arc::clone(&file.reader)) + .with_projection(global.bound_projection.clone()) + .with_some_filter(filter) + .with_ordered(file_row_number_column_pos.is_some()) + .with_selection(row_selection); + if let Some(range) = row_range { + builder = builder.with_row_range(range); + } + let splits = builder.build()?; + + let handles = splits + .into_iter() + .map(|task| RUNTIME.handle().spawn(task)) + .collect::>(); + let total_splits = handles.len() as u64; + + let cache = Arc::new(ConversionCache::default()); + + let pending = (!global.aggregates.is_empty()).then(|| Arc::clone(&global.pending)); + let (sender, receiver) = kanal::bounded_async(FILE_CHANNEL_CAPACITY); + let driver = RUNTIME.handle().spawn(async move { + for handle in handles { + match handle.await { + Ok(Some(array)) => { + if let Some(pending) = &pending { + pending.fetch_add(1, Ordering::Relaxed); + } + if sender.send(Ok((array, Arc::clone(&cache)))).await.is_err() { + // The receiver is gone: the scan was cancelled or the query ended. + return; + } + } + // split is filtered + Ok(None) => {} + Err(e) => { + let _ = sender.send(Err(e)).await; + return; + } + } + } + }); + + Ok(FileScan { + receiver, + exhausted: AtomicBool::new(false), + total_splits, + delivered: AtomicU64::new(0), + file_row_number_column_pos, + file_index_column_pos, + aggregate_positions: global.aggregate_positions.clone(), + _driver: driver, + }) +} + +pub fn file_has_work(scan: &FileScan) -> bool { + !scan.exhausted.load(Ordering::Acquire) +} + +fn file_scan_aggregate( + scan: &FileScan, + global: &TableFunctionGlobal, + local: &mut TableFunctionLocal, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let has_count_star = local.partials.len() < global.aggregates.len(); + let mut accumulated = 0u64; + let mut rows = 0u64; + loop { + let Ok(item) = RUNTIME.block_on(scan.receiver.recv()) else { + scan.exhausted.store(true, Ordering::Release); + break; + }; + let (array, _cache) = item?; + scan.delivered.fetch_add(1, Ordering::Relaxed); + let array = convert_result(array, &mut ctx)?; + + for (position, partial) in scan + .aggregate_positions + .iter() + .zip(local.partials.iter_mut()) + { + partial.accumulate(array.unmasked_field(*position), &mut ctx)?; + } + rows += array.len() as u64; + accumulated += 1; + } + + if accumulated == 0 { + return Ok(()); + } + + { + 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()?)?; + } + } + if has_count_star { + global.row_count.fetch_add(rows, Ordering::Relaxed); + } + global.pending.fetch_sub(accumulated, Ordering::Release); + Ok(()) +} + +pub fn file_scan( + scan: &FileScan, + global: &TableFunctionGlobal, + local: &mut TableFunctionLocal, + output: &mut DataChunkRef, +) -> VortexResult<()> { + if !local.partials.is_empty() { + return file_scan_aggregate(scan, global, local); + } + loop { + if local.exporter.is_none() { + let Ok(item) = RUNTIME.block_on(scan.receiver.recv()) else { + scan.exhausted.store(true, Ordering::Release); + return Ok(()); + }; + let (array, cache) = item?; + scan.delivered.fetch_add(1, Ordering::Relaxed); + + let mut ctx = SESSION.create_execution_ctx(); + let array = convert_result(array, &mut ctx)?; + local.exporter = Some(ArrayExporter::try_new(&array, &cache, ctx)?); + } + + let exporter = local + .exporter + .as_mut() + .ok_or_else(|| vortex_err!("exporter missing"))?; + let has_more_data = exporter.export( + output, + scan.file_index_column_pos, + scan.file_row_number_column_pos, + )?; + + if !has_more_data { + // This exporter is fully consumed. + local.exporter = None; + } else { + break; + } + } + Ok(()) +} + +pub fn file_progress(scan: &FileScan) -> f64 { + if scan.total_splits == 0 { + return 100.0; + } + let delivered = scan.delivered.load(Ordering::Relaxed) as f64; + 100.0 * delivered / scan.total_splits as f64 +} diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index fe14c86f8f6..784b4da37b4 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -26,7 +26,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 b8ddc076b8e..00000000000 --- a/vortex-duckdb/src/multi_file.rs +++ /dev/null @@ -1,104 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; -use std::sync::LazyLock; - -use itertools::Itertools; -use object_store::registry::ObjectStoreRegistry; -use url::Url; -use vortex::cloud::Registry; -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::RUNTIME; -use crate::SESSION; -use crate::duckdb::BindInputRef; -use crate::duckdb::ExtractedValue; - -/// Process-wide registry, so repeated scans against the same bucket share one client. -static REGISTRY: LazyLock = LazyLock::new(Registry::new); - -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..1245278da03 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -29,9 +29,9 @@ use crate::table_function::ColumnAggregate; // See MultiFileReader for constants /// "file_index" virtual column -static FILE_INDEX_COLUMN_IDX: u64 = 9223372036854775810; +pub(crate) 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 +48,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 +85,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 +93,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 +112,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 +145,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,11 +177,7 @@ 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) } } @@ -358,36 +312,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 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..fc55c45eedb 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -4,19 +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; use parking_lot::Mutex; @@ -26,40 +19,25 @@ 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; @@ -67,19 +45,18 @@ 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::LogicalType; use crate::duckdb::LogicalTypeRef; use crate::duckdb::TableInitInput; use crate::duckdb::Value; use crate::exporter::ArrayExporter; -use crate::exporter::ConversionCache; -use crate::multi_file::bind_multi_file_scan; +use crate::file_reader::FileReader; +use crate::file_reader::ScanPruning; use crate::projection::DuckdbField; -use crate::projection::Filter; use crate::projection::Projection; use crate::projection::extract_schema_from_dtype; @@ -87,28 +64,32 @@ use crate::projection::extract_schema_from_dtype; pub const COUNT_STAR_PROJ_IDX: u64 = u64::MAX; pub struct TableFunctionBind { - data_source: Arc, - filter_exprs: Vec, - column_fields: Vec, + pub(crate) dtype: DType, + first_file_row_count: u64, + pub(crate) filter_exprs: Vec, + pub(crate) column_fields: 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(crate) has_non_optional_filter: AtomicBool, // Non-empty iff this scan is aggregate aggregates: Vec, + aggregate_outputs: Vec<(String, LogicalType)>, } assert_impl_all!(TableFunctionBind: Send, Clone); impl Clone for TableFunctionBind { fn clone(&self) -> Self { Self { - data_source: Arc::clone(&self.data_source), - // filter_exprs are consumed once in `init_global`. + dtype: self.dtype.clone(), + first_file_row_count: self.first_file_row_count, + // Cloning happens only for late materialization refetch filter_exprs: vec![], column_fields: self.column_fields.clone(), has_non_optional_filter: AtomicBool::new( self.has_non_optional_filter.load(Ordering::Relaxed), ), aggregates: self.aggregates.clone(), + aggregate_outputs: self.aggregate_outputs.clone(), } } } @@ -135,43 +116,27 @@ impl<'a> TableInitInput<'a> { } } -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, + pub(crate) pruning: Option, + pub(crate) bound_projection: BoundExpression, + // 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(crate) pending: Arc, + pub(crate) aggregates: Vec, + pub(crate) aggregate_positions: Vec, // Accumulated partials - partials: Mutex>>, - row_count: AtomicU64, + pub(crate) partials: Mutex>>, + pub(crate) row_count: AtomicU64, } assert_impl_all!(TableFunctionGlobal: Send, Sync); /// Per-thread scan state pub struct TableFunctionLocal { - iterator: DataSourceIterator, - exporter: Option, - partition_index: u64, - file_index: usize, + pub(crate) exporter: Option, // Aggregate scan accumulated partials. Empty for non-aggregate scan - partials: Vec>, -} - -pub struct PartitionData { - pub partition_index: u64, - pub file_index_column_pos: Option, - pub file_index: usize, + pub(crate) partials: Vec>, } #[derive(Clone)] @@ -185,8 +150,6 @@ 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. @@ -196,196 +159,104 @@ pub enum Cardinality { // 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 bind(first_file: &FileReader) -> VortexResult { + let dtype = first_file.dtype.clone(); + let column_fields = extract_schema_from_dtype(&dtype)?; Ok(TableFunctionBind { - data_source: Arc::new(data_source), + dtype, + first_file_row_count: first_file.row_count, filter_exprs: vec![], column_fields, has_non_optional_filter: AtomicBool::new(false), aggregates: vec![], + aggregate_outputs: vec![], }) } -pub fn init_global(init_input: &TableInitInput) -> VortexResult { - debug!(input=?init_input, "table function global input"); +pub fn bind_schema(bind_data: &TableFunctionBind, result: &mut BindResultRef) { + if !bind_data.aggregate_outputs.is_empty() { + for (name, logical_type) in &bind_data.aggregate_outputs { + result.add_result_column(name, logical_type); + } + return; + } + for field in &bind_data.column_fields { + result.add_result_column(&field.name, &field.logical_type); + } +} + +pub fn finalize_scan(global: &TableFunctionGlobal, 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); + } + + 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 column_ids = init_input.column_ids(); - let projection_ids = init_input.projection_ids(); - - 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) - } else { - Projection::new_aggregate(&bind_data.aggregates, &bind_data.column_fields) - }; - let Filter { - filter, - row_selection, - row_range, - file_selection, - file_range, - has_non_optional_filter, - } = Filter::new( - init_input.table_filter_set(), - column_ids, + let partials = build_partials( + &bind_data.aggregates, &bind_data.column_fields, - &bind_data.filter_exprs, - bind_data.data_source.dtype(), + &bind_data.dtype, )?; - if has_non_optional_filter { - init_input - .bind_data() - .has_non_optional_filter - .store(true, Ordering::Relaxed); + let pruning = ScanPruning::new(bind_data, init_input.column_ids(), init_input.input.filters)?; + + 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); } - debug!( - %projection, - filter = filter - .as_ref() - .map_or_else(|| "true".to_string(), |f| f.to_string()), - ?row_selection, - ?row_range, - ?file_selection, - ?file_range, - "table function scan input" - ); - - let request = ScanRequest { - 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 Projection(projection) = if bind_data.aggregates.is_empty() { + Projection::new(init_input.column_ids(), &bind_data.column_fields) + } else { + Projection::new_aggregate(&bind_data.aggregates, &bind_data.column_fields) }; - - 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(), - )?; + let bound_projection = optimize_and_bind(projection, &bind_data.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, + pruning, + aggregate_positions, + bound_projection, + pending: Arc::new(AtomicU64::new(0)), + aggregates: bind_data.aggregates.clone(), partials: Mutex::new(partials), row_count: AtomicU64::new(0), }) } -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 { @@ -439,22 +310,23 @@ pub fn init_local( let partials = build_partials( &global.aggregates, &bind_data.column_fields, - bind_data.data_source.dtype(), + &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"); TableFunctionLocal { - iterator: global.iterator.clone(), exporter: None, - partition_index: 0, - file_index: 0, partials, } } -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,157 +368,6 @@ 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, expr: &ExpressionRef, @@ -703,7 +424,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.column_fields[projection_id]; + field.logical_type = expr.return_type().to_owned(); + field.dtype = out_dtype; + field.projection_expr = Some(vx_expr); Ok(true) } } @@ -716,7 +443,7 @@ fn can_push_projection_aggregate( ) -> 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 Ok(dtype) = aggregate_input_dtype(field, &bind_data.dtype) else { return false; }; @@ -765,13 +492,24 @@ pub fn pushdown_projection_aggregates( ) -> VortexResult { 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.column_fields[id].name.clone() + } + }; + outputs.push((name, output_type)); has_non_count_star |= matches!(aggregate, ColumnAggregate::Real { .. }); aggregates.push(aggregate); } @@ -780,6 +518,7 @@ pub fn pushdown_projection_aggregates( return Ok(false); } bind_data.aggregates = aggregates; + bind_data.aggregate_outputs = outputs; Ok(true) } @@ -810,41 +549,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,45 +559,24 @@ 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: &TableFunctionBind, 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) { @@ -941,57 +624,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")); - } -}