duckdb: propagate per-file "partition" stats

vortex-duckdb/cpp/include/multi_file_reader.hpp

@@ -3,10 +3,14 @@
 #pragma once
 
 #include "data.hpp"
+#include "table_function.h"
 #include "duckdb/common/multi_file/multi_file_function.hpp"
 
 using namespace duckdb;
 
+// Convert Vortex statistics to DuckDB statistics
+unique_ptr<BaseStatistics> to_duckdb_statistics(duckdb_column_statistics &statistics);
+
 struct VortexBindData final : TableFunctionData {
     VortexBindData() = default;
     unique_ptr<FunctionData> Copy() const override;

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,28 @@ struct TableFunctionUngroupedAggregateInput {
 };
 
 bool aggregate_pushdown(ClientContext &context, const TableFunctionUngroupedAggregateInput &input);
+
+/*
+ * DuckDB uses partition row groups for two purposes:
+ *
+ * 1. If optimizer proves query (e.g. SELECT min(col)) can be answered from
+ *    metadata, it replaces a real scan with a call to GetColumnStatistics.
+ * 2. If optimizer proves a row group can be pruned because of statistics, it
+ *    removes row group's read. We don't use this as we implement own prunung.
+ *
+ * For (1) we care about providing statistics fast, so we report a file as
+ * a "row group".
+ */
+struct VortexRowGroup final : PartitionRowGroup {
+    explicit VortexRowGroup(unique_ptr<CData> ffi_footer) : ffi_footer(std::move(ffi_footer)) {
+    }
+
+    unique_ptr<CData> ffi_footer;
+
+    unique_ptr<BaseStatistics> 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;
+    }
+};

vortex-duckdb/cpp/multi_file_reader.cpp

@@ -344,17 +344,7 @@ static unique_ptr<BaseStatistics> base_stats(duckdb_column_statistics &stats, Lo
     return out.ToUnique();
 }
 
-unique_ptr<BaseStatistics> 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<BaseStatistics> to_duckdb_statistics(duckdb_column_statistics &statistics) {
     using enum LogicalTypeId;
     const unique_ptr<LogicalType> type(reinterpret_cast<LogicalType *>(statistics.type));
     switch (type->id()) {
@@ -370,7 +360,15 @@ unique_ptr<BaseStatistics> 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<BaseStatistics> VortexBaseReader::GetStatistics(ClientContext &, cons
     }
 }
 
+unique_ptr<BaseStatistics> 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());
 }

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,51 @@ unique_ptr<MultiFileReader> get_multi_file_reader(const TableFunction &) {
     return make_uniq<VortexMultiFileReader>();
 }
 
+unique_ptr<BaseStatistics> 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<PartitionStatistics> get_partition_stats(ClientContext &, GetPartitionStatsInput &input) {
+    const MultiFileBindData &bind_data = input.bind_data->Cast<MultiFileBindData>();
+    VortexBindData &bind = bind_data.bind_data->Cast<VortexBindData>();
+    void *const ffi_bind = bind.ffi_bind_data->DataPtr();
+
+    if (!duckdb_table_function_can_get_partition_stats(ffi_bind)) {
+        return {};
+    }
+
+    vector<OpenFileInfo> files = bind_data.file_list->GetAllFiles();
+    vector<PartitionStatistics> result(files.size());
+    idx_t row_start = 0;
+    uint64_t count = 0;
+    duckdb_vx_error error = nullptr;
+    for (size_t i = 0; i < files.size(); ++i) {
+        const std::string_view path = files[i].path;
+        duckdb_vx_data raw = duckdb_footer_get_cached(ffi_bind, path.data(), path.size(), &count, &error);
+        unique_ptr<CData> cdata(reinterpret_cast<CData *>(raw));
+        if (error) {
+            throw BinderException(IntoErrString(error));
+        }
+        if (!cdata) {
+            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<VortexRowGroup>(std::move(cdata));
+    }
+    return result;
+}
+
 duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter, const std::string &name) {
     MultiFileFunction<VortexReaderInterface> fn(name);
     fn.arguments[0] = parameter;
@@ -156,6 +203,7 @@ duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter
     };
 
     fn.statistics = MultiFileFunction<VortexReaderInterface>::MultiFileScanStats;
+    fn.get_partition_stats = get_partition_stats;
     fn.get_multi_file_reader = get_multi_file_reader;
 
     try {

vortex-duckdb/include/vortex.h

@@ -65,6 +65,20 @@ bool duckdb_reader_get_statistics(const void *file,
                                   size_t column_name_len,
                                   duckdb_column_statistics *stats_out);
 
+extern bool duckdb_table_function_can_get_partition_stats(const void *bind);
+
+extern
+duckdb_vx_data duckdb_footer_get_cached(void *bind,
+                                        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);

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,9 @@ use crate::duckdb::TableInitInput;
 use crate::duckdb::try_or;
 use crate::duckdb::try_or_null;
 use crate::file_reader::OpenFileReader;
+use crate::file_reader::can_get_partition_stats;
+use crate::file_reader::footer_get_cached;
+use crate::file_reader::footer_get_statistics;
 use crate::file_reader::reader_bind;
 use crate::file_reader::reader_get_progress_in_file;
 use crate::file_reader::reader_get_statistics;
@@ -217,6 +221,55 @@ pub unsafe extern "C-unwind" fn duckdb_reader_get_statistics(
     true
 }
 
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn duckdb_table_function_can_get_partition_stats(
+    bind: *const c_void,
+) -> bool {
+    let bind = unsafe { bind.cast::<BindState>().as_ref() }.vortex_expect("null pointer");
+    can_get_partition_stats(bind)
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn duckdb_footer_get_cached(
+    bind: *mut c_void,
+    path: *const c_char,
+    len: usize,
+    row_count_out: *mut u64,
+    error: *mut cpp::duckdb_vx_error,
+) -> cpp::duckdb_vx_data {
+    let bind = unsafe { bind.cast::<BindState>().as_mut() }.vortex_expect("null pointer");
+    let path = unsafe { std::slice::from_raw_parts(path.cast::<u8>(), len) };
+    try_or_null(error, || {
+        let path = str::from_utf8(path).map_err(|_| vortex_err!("invalid utf-8"))?;
+        Ok(match footer_get_cached(bind, 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::<Footer>().as_ref() }.vortex_expect("null pointer");
+    let Some(stats) = footer_get_statistics(footer, column_index) else {
+        return false;
+    };
+    let stats_out = unsafe { &mut *stats_out };
+    stats_out.min = stats.min.map_or(ptr::null_mut(), |v| v.into_ptr());
+    stats_out.max = stats.max.map_or(ptr::null_mut(), |v| v.into_ptr());
+    stats_out.max_string_length = stats.max_string_length;
+    stats_out.has_null = stats.has_null;
+    stats_out.type_ = stats.logical_type.into_ptr();
+    true
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C-unwind" fn duckdb_reader_initialize(
     global: *const c_void,