diff --git a/vortex-duckdb/cpp/include/multi_file_reader.hpp b/vortex-duckdb/cpp/include/multi_file_reader.hpp index c0dce58c5f6..a6e28fb75e9 100644 --- a/vortex-duckdb/cpp/include/multi_file_reader.hpp +++ b/vortex-duckdb/cpp/include/multi_file_reader.hpp @@ -3,16 +3,20 @@ #pragma once #include "data.hpp" +#include "table_function.h" #include "duckdb/common/multi_file/multi_file_function.hpp" using namespace duckdb; +unique_ptr to_duckdb_statistics(duckdb_column_statistics &statistics); + struct VortexBindData final : TableFunctionData { VortexBindData() = default; unique_ptr Copy() const override; bool Equals(const FunctionData &other) const override; unique_ptr ffi_bind_data; + bool no_footer_caches = false; }; struct VortexBindResult { diff --git a/vortex-duckdb/cpp/include/table_function.hpp b/vortex-duckdb/cpp/include/table_function.hpp index 54daff90d06..9459d23f0ec 100644 --- a/vortex-duckdb/cpp/include/table_function.hpp +++ b/vortex-duckdb/cpp/include/table_function.hpp @@ -3,6 +3,7 @@ #pragma once +#include "data.hpp" #include "duckdb.h" #include "duckdb/function/function.hpp" #include "duckdb/function/table_function.hpp" @@ -30,3 +31,18 @@ struct TableFunctionUngroupedAggregateInput { }; bool aggregate_pushdown(ClientContext &context, const TableFunctionUngroupedAggregateInput &input); + +// Vortex "row group" is a file +struct VortexRowGroup final : PartitionRowGroup { + explicit VortexRowGroup(unique_ptr ffi_footer) : ffi_footer(std::move(ffi_footer)) { + } + + unique_ptr ffi_footer; + + unique_ptr GetColumnStatistics(const StorageIndex &storage_index) override; + bool MinMaxIsExact(const BaseStatistics &, const StorageIndex &) override { + // TODO(myrrc): in duckdb 2.0 we should report false for strings and + // also add TRUNCATED_STATS type for them + return true; + } +}; diff --git a/vortex-duckdb/cpp/multi_file_reader.cpp b/vortex-duckdb/cpp/multi_file_reader.cpp index 9184bb01239..02d4da3caaa 100644 --- a/vortex-duckdb/cpp/multi_file_reader.cpp +++ b/vortex-duckdb/cpp/multi_file_reader.cpp @@ -344,17 +344,7 @@ static unique_ptr base_stats(duckdb_column_statistics &stats, Lo return out.ToUnique(); } -unique_ptr VortexBaseReader::GetStatistics(ClientContext &, const string &name) { - D_ASSERT(ffi_bind); - duckdb_column_statistics statistics = {}; - if (!duckdb_reader_get_statistics(ffi_file->DataPtr(), - ffi_bind, - name.c_str(), - name.size(), - &statistics)) { - return {}; - } - +unique_ptr to_duckdb_statistics(duckdb_column_statistics &statistics) { using enum LogicalTypeId; const unique_ptr type(reinterpret_cast(statistics.type)); switch (type->id()) { @@ -370,7 +360,15 @@ unique_ptr VortexBaseReader::GetStatistics(ClientContext &, cons case UINTEGER: case UBIGINT: case UHUGEINT: - case HUGEINT: { + case HUGEINT: + case DATE: + case TIME: + case TIME_TZ: + case TIMESTAMP_SEC: + case TIMESTAMP_MS: + case TIMESTAMP: + case TIMESTAMP_NS: + case TIMESTAMP_TZ: { return numeric_stats(statistics, *type); } case VARCHAR: @@ -395,6 +393,20 @@ unique_ptr VortexBaseReader::GetStatistics(ClientContext &, cons } } +unique_ptr VortexBaseReader::GetStatistics(ClientContext &, const string &name) { + D_ASSERT(ffi_bind); + duckdb_column_statistics statistics = {}; + if (!duckdb_reader_get_statistics(ffi_file->DataPtr(), + ffi_bind, + name.c_str(), + name.size(), + &statistics)) { + return {}; + } + + return to_duckdb_statistics(statistics); +} + double VortexBaseReader::GetProgressInFile(ClientContext &) { return duckdb_reader_get_progress_in_file(ffi_file->DataPtr()); } diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 7b6e118d7c6..c6442f6de80 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -17,8 +17,10 @@ #include "duckdb/function/table_function.hpp" #include "duckdb/main/capi/capi_internal.hpp" #include "duckdb/main/connection.hpp" +#include "duckdb/function/partition_stats.hpp" #include "duckdb/parser/parsed_data/create_table_function_info.hpp" #include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/storage/storage_index.hpp" using namespace std::string_literals; constexpr column_t COLUMN_IDENTIFIER_FILE_INDEX = MultiFileReader::COLUMN_IDENTIFIER_FILE_INDEX; @@ -130,6 +132,53 @@ unique_ptr get_multi_file_reader(const TableFunction &) { return make_uniq(); } +unique_ptr VortexRowGroup::GetColumnStatistics(const StorageIndex &storage_index) { + duckdb_column_statistics statistics = {}; + const idx_t idx = storage_index.GetPrimaryIndex(); + const void *const ffi_footer_ptr = ffi_footer->DataPtr(); + if (!duckdb_footer_get_statistics(ffi_footer_ptr, idx, &statistics)) { + return {}; + } + return to_duckdb_statistics(statistics); +} + +static vector get_partition_stats(ClientContext &, GetPartitionStatsInput &input) { + const MultiFileBindData &bind_data = input.bind_data->Cast(); + VortexBindData &bind = bind_data.bind_data->Cast(); + if (bind.no_footer_caches) { + return {}; + } + if (duckdb_table_function_has_pushed_filters(bind.ffi_bind_data->DataPtr())) { + return {}; + } + + vector files = bind_data.file_list->GetAllFiles(); + vector result(files.size()); + idx_t row_start = 0; + for (size_t i = 0; i < files.size(); ++i) { + const std::string_view path = files[i].path; + duckdb_vx_error error = nullptr; + uint64_t count = 0; + duckdb_vx_data raw = duckdb_footer_open(path.data(), path.size(), &count, &error); + unique_ptr cdata(reinterpret_cast(raw)); + if (error) { + throw BinderException(IntoErrString(error)); + } + if (!cdata) { + bind.no_footer_caches = true; + return {}; + } + + PartitionStatistics &stats = result[i]; + stats.row_start = row_start; + stats.count = count; + stats.count_type = CountType::COUNT_EXACT; + row_start += count; + stats.partition_row_group = make_shared_ptr(std::move(cdata)); + } + return result; +} + duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter, const std::string &name) { MultiFileFunction fn(name); fn.arguments[0] = parameter; @@ -156,6 +205,7 @@ duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter }; fn.statistics = MultiFileFunction::MultiFileScanStats; + fn.get_partition_stats = get_partition_stats; fn.get_multi_file_reader = get_multi_file_reader; try { diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index e096926d454..13f90fccb98 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -65,6 +65,19 @@ bool duckdb_reader_get_statistics(const void *file, size_t column_name_len, duckdb_column_statistics *stats_out); +extern bool duckdb_table_function_has_pushed_filters(const void *bind); + +extern +duckdb_vx_data duckdb_footer_open(const char *path, + size_t len, + uint64_t *row_count_out, + duckdb_vx_error *error); + +extern +bool duckdb_footer_get_statistics(const void *footer, + size_t column_index, + duckdb_column_statistics *stats_out); + extern bool duckdb_reader_initialize(const void *global, void *file, duckdb_vx_error *error); extern duckdb_logical_type duckdb_reader_bind_column_type(const void *bind, size_t index); diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 160f0217d71..765550b9468 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -9,6 +9,7 @@ use std::ptr; use num_traits::AsPrimitive; use vortex::error::VortexExpect; use vortex::error::vortex_err; +use vortex::file::Footer; use crate::convert::can_push_expression; use crate::copy::CopyFunctionBind; @@ -33,6 +34,8 @@ use crate::duckdb::TableInitInput; use crate::duckdb::try_or; use crate::duckdb::try_or_null; use crate::file_reader::OpenFileReader; +use crate::file_reader::footer_get_statistics; +use crate::file_reader::footer_open; use crate::file_reader::reader_bind; use crate::file_reader::reader_get_progress_in_file; use crate::file_reader::reader_get_statistics; @@ -217,6 +220,53 @@ pub unsafe extern "C-unwind" fn duckdb_reader_get_statistics( true } +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_has_pushed_filters( + bind: *const c_void, +) -> bool { + let bind = unsafe { bind.cast::().as_ref() }.vortex_expect("null pointer"); + !bind.filters.is_empty() +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_footer_open( + path: *const c_char, + len: usize, + row_count_out: *mut u64, + error: *mut cpp::duckdb_vx_error, +) -> cpp::duckdb_vx_data { + let path = unsafe { std::slice::from_raw_parts(path.cast::(), len) }; + try_or_null(error, || { + let path = str::from_utf8(path).map_err(|_| vortex_err!("invalid utf-8"))?; + Ok(match footer_open(path)? { + Some(footer) => { + unsafe { *row_count_out = footer.row_count() }; + Data::from(Box::new(footer)).as_ptr() + } + None => ptr::null_mut(), + }) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_footer_get_statistics( + footer: *const c_void, + column_index: usize, + stats_out: *mut cpp::duckdb_column_statistics, +) -> bool { + let footer = unsafe { footer.cast::