duckdb: propagate per-file "partition" stats

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<BaseStatistics> to_duckdb_statistics(duckdb_column_statistics &statistics);
+
 struct VortexBindData final : TableFunctionData {
     VortexBindData() = default;
     unique_ptr<FunctionData> Copy() const override;
     bool Equals(const FunctionData &other) const override;
 
     unique_ptr<CData> ffi_bind_data;
+    bool no_footer_caches = false;
 };
 
 struct VortexBindResult {

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<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,53 @@ 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>();
+    if (bind.no_footer_caches) {
+        return {};
+    }
+    if (duckdb_table_function_has_pushed_filters(bind.ffi_bind_data->DataPtr())) {
+        return {};
+    }
+
+    vector<OpenFileInfo> files = bind_data.file_list->GetAllFiles();
+    vector<PartitionStatistics> 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> cdata(reinterpret_cast<CData *>(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<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 +205,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,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);

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::<BindState>().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::<u8>(), 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::<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,

vortex-duckdb/src/file_reader.rs

@@ -16,16 +16,20 @@ use vortex::dtype::DType;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::error::vortex_panic;
+use vortex::file::Footer;
+use vortex::file::multi::MultiFileSession;
 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::object_store::object_path_from_literal;
 use vortex::io::runtime::BlockingRuntime as _;
 use vortex::layout::LayoutReaderRef;
 use vortex::layout::scan::scan_builder::ScanBuilder;
 use vortex::mask::Mask;
+use vortex::session::SessionExt as _;
 
 use crate::RUNTIME;
 use crate::SESSION;
@@ -91,6 +95,14 @@ fn resolve_filesystem(url: &Url) -> VortexResult<(FileSystemRef, String)> {
     ))
 }
 
+/// Same as resolve_filesystem, but doesn't create filesystem object
+fn resolve_path(url: &Url) -> VortexResult<String> {
+    if url.scheme() == "file" {
+        return Ok(url.path().to_string());
+    }
+    Ok(REGISTRY.resolve(url)?.1.to_string())
+}
+
 pub struct OpenFileReader {
     pub reader: LayoutReaderRef,
     /// File splits stored in inverse order
@@ -271,12 +283,13 @@ pub fn reader_get_statistics(
         .reader
         .as_any()
         .downcast_ref::<FileStatsLayoutReader>()?;
-    let stats_sets = reader.file_stats().stats_sets();
 
     let DType::Struct(fields, _) = &file.reader.dtype() else {
         return None;
     };
     let index = fields.find(column)?;
+    let stats_sets = reader.file_stats().stats_sets();
+
     let dtype = fields.field_by_index(index)?;
 
     let stats = ColumnStatisticsAggregate::new(stats_sets.get(index)?);
@@ -294,3 +307,24 @@ pub fn reader_get_progress_in_file(file: &OpenFileReader) -> f64 {
     let denom = total + (total == 0) as usize;
     100.0 * (total - left) as f64 / denom as f64
 }
+
+pub fn footer_open(path: &str) -> VortexResult<Option<Footer>> {
+    let url = parse_uri_or_path(path)?;
+    let path = resolve_path(&url)?;
+    let key = object_path_from_literal(&path).to_string();
+    Ok(SESSION.get::<MultiFileSession>().get_footer(&key))
+}
+
+pub fn footer_get_statistics(footer: &Footer, index: usize) -> Option<ColumnStatistics> {
+    let DType::Struct(fields, _) = footer.dtype() else {
+        return None;
+    };
+    let stats = footer.statistics()?;
+    let dtype = fields.field_by_index(index)?;
+    let stats = stats.stats_sets().get(index)?;
+    let stats = ColumnStatisticsAggregate::new(stats);
+    match ColumnStatistics::try_from(&stats, dtype) {
+        Ok(stats) => Some(stats),
+        Err(e) => vortex_panic!(e),
+    }
+}

vortex-file/src/multi/session.rs

@@ -55,7 +55,7 @@ impl Debug for MultiFileSession {
 
 impl MultiFileSession {
     /// Retrieve a cached footer for the given file path.
-    pub(crate) fn get_footer(&self, path: &str) -> Option<Footer> {
+    pub fn get_footer(&self, path: &str) -> Option<Footer> {
         self.footer_cache.get(path)
     }