initial

Signed-off-by: Mikhail Kot <to@myrrc.dev>

Author: myrrc

benchmarks/datafusion-bench/src/main.rs

@@ -316,7 +316,7 @@ async fn register_v2_tables<B: Benchmark + ?Sized>(
             None => format!("*.{}", format.ext()),
         };
 
-        let multi_ds = MultiFileDataSource::new(SESSION.clone())
+        let (multi_ds, _) = MultiFileDataSource::new(SESSION.clone())
             .with_glob(glob_pattern, Some(fs))
             .build()
             .await?;

vortex-array/src/stats/stats_set.rs

@@ -459,11 +459,28 @@ impl MutTypedStatsSetRef<'_, '_> {
         ) {
             (Some(m1), Some(m2)) => {
                 // If the combine sum is exact, then we can sum them.
-                if let Some(scalar_value) = m1.zip(m2).as_exact().and_then(|(s1, s2)| {
-                    s1.as_primitive()
-                        .checked_add(&s2.as_primitive())
-                        .and_then(|pscalar| pscalar.pvalue().map(ScalarValue::Primitive))
-                }) {
+                if let Some(scalar_value) =
+                    m1.zip(m2).as_exact().and_then(|(s1, s2)| match s1.dtype() {
+                        DType::Primitive(..) => s1
+                            .as_primitive()
+                            .checked_add(&s2.as_primitive())
+                            .and_then(|pscalar| pscalar.pvalue().map(ScalarValue::Primitive)),
+                        DType::Decimal(..) => s1
+                            .as_decimal()
+                            .checked_binary_numeric(
+                                &s2.as_decimal(),
+                                crate::scalar::NumericOperator::Add,
+                            )
+                            .map(|scalar| {
+                                ScalarValue::Decimal(
+                                    scalar
+                                        .decimal_value()
+                                        .vortex_expect("no decimal value in scalar"),
+                                )
+                            }),
+                        _ => None,
+                    })
+                {
                     self.set(Stat::Sum, Precision::Exact(scalar_value));
                 }
             }

vortex-duckdb/cpp/include/duckdb_vx/table_function.h

@@ -97,6 +97,15 @@ typedef struct {
     bool has_max_cardinality;
 } duckdb_vx_node_statistics;
 
+typedef struct {
+    duckdb_value min;
+    duckdb_value max;
+    // upper bit: "length is set". lower 32 bits: DuckDB's max string length.
+    uint64_t max_string_length;
+} duckdb_column_statistics;
+
+typedef idx_t column_t;
+
 // A transparent DuckDB table function vtable, which can be used to configure a table function.
 // See duckdb/include/function/tfunc.hpp for details on each field.
 typedef struct {
@@ -137,7 +146,11 @@ typedef struct {
 
     // void *in_out_function;
     // void *in_out_function_final;
-    void *statistics;
+
+    void (*statistics)(duckdb_client_context context,
+                       const void *bind_data,
+                       size_t column_index,
+                       duckdb_column_statistics *stats_out);
 
     // void *dependency;
     void (*cardinality)(void *bind_data, duckdb_vx_node_statistics *node_stats_out);

vortex-duckdb/cpp/table_function.cpp

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+#include "duckdb_vx/table_function.h"
 #include "duckdb_vx/duckdb_diagnostics.h"
 
 DUCKDB_INCLUDES_BEGIN
@@ -30,8 +31,10 @@ struct CTableFunctionInfo final : TableFunctionInfo {
 };
 
 struct CTableBindData final : TableFunctionData {
-    CTableBindData(unique_ptr<CTableFunctionInfo> info_p, unique_ptr<vortex::CData> ffi_data_p)
-        : info(std::move(info_p)), ffi_data(std::move(ffi_data_p)) {
+    CTableBindData(unique_ptr<CTableFunctionInfo> info_p,
+                   unique_ptr<vortex::CData> ffi_data_p,
+                   const vector<LogicalType> &types)
+        : info(std::move(info_p)), ffi_data(std::move(ffi_data_p)), types(types) {
     }
 
     unique_ptr<FunctionData> Copy() const override {
@@ -43,11 +46,13 @@ struct CTableBindData final : TableFunctionData {
             throw BinderException(IntoErrString(error_out));
         }
         return make_uniq<CTableBindData>(make_uniq<CTableFunctionInfo>(info->vtab),
-                                         unique_ptr<CData>(reinterpret_cast<CData *>(copied_ffi_data)));
+                                         unique_ptr<CData>(reinterpret_cast<CData *>(copied_ffi_data)),
+                                         types);
     }
 
     unique_ptr<CTableFunctionInfo> info;
     unique_ptr<CData> ffi_data;
+    vector<LogicalType> types;
 };
 
 struct CTableGlobalData final : GlobalTableFunctionState {
@@ -88,6 +93,79 @@ double c_table_scan_progress(ClientContext &context,
     return bind.info->vtab.table_scan_progress(c_ctx, c_bind_data, c_global_state);
 }
 
+static Value &UnwrapValue(duckdb_value value) {
+    return *(reinterpret_cast<Value *>(value));
+}
+
+unique_ptr<BaseStatistics> 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);
+    }
+    return out.ToUnique();
+}
+
+unique_ptr<BaseStatistics> 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));
+    }
+    return out.ToUnique();
+}
+
+unique_ptr<BaseStatistics>
+c_statistics(ClientContext &context, const FunctionData *bind_data, column_t column_index) {
+    if (column_index == COLUMN_IDENTIFIER_EMPTY) {
+        return BaseStatistics::CreateUnknown(LogicalTypeId::INVALID).ToUnique();
+    }
+
+    const auto &bind = bind_data->Cast<CTableBindData>();
+    void *const ffi_bind = bind.ffi_data->DataPtr();
+
+    duckdb_client_context c_ctx = reinterpret_cast<duckdb_client_context>(&context);
+    duckdb_column_statistics statistics = {};
+    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: {
+        bind.info->vtab.statistics(c_ctx, ffi_bind, column_index, &statistics);
+        return numeric_stats(statistics, type);
+    }
+    case LogicalTypeId::VARCHAR:
+    case LogicalTypeId::BLOB: {
+        bind.info->vtab.statistics(c_ctx, ffi_bind, column_index, &statistics);
+        return string_stats(statistics, type);
+    }
+    default:
+        return BaseStatistics::CreateUnknown(type).ToUnique();
+    }
+}
+
 unique_ptr<FunctionData> c_bind(ClientContext &context,
                                 TableFunctionBindInput &input,
                                 vector<LogicalType> &return_types,
@@ -111,7 +189,8 @@ unique_ptr<FunctionData> c_bind(ClientContext &context,
     }
 
     return make_uniq<CTableBindData>(make_uniq<CTableFunctionInfo>(info.vtab),
-                                     unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data)));
+                                     unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data)),
+                                     return_types);
 }
 
 unique_ptr<GlobalTableFunctionState> c_init_global(ClientContext &context, TableFunctionInitInput &input) {
@@ -363,6 +442,7 @@ extern "C" duckdb_state duckdb_vx_tfunc_register(duckdb_database ffi_db, const d
     tf.get_virtual_columns = c_get_virtual_columns;
     tf.to_string = c_to_string;
     tf.table_scan_progress = c_table_scan_progress;
+    tf.statistics = c_statistics;
 
     // Set up the parameters
     tf.arguments.reserve(vtab->parameter_count);

vortex-duckdb/src/datasource.rs

@@ -37,6 +37,7 @@ use vortex::expr::col;
 use vortex::expr::root;
 use vortex::expr::select;
 use vortex::expr::stats::Precision;
+use vortex::file::FileStatistics;
 use vortex::io::kanal_ext::KanalExt;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::runtime::current::ThreadSafeIterator;
@@ -54,6 +55,7 @@ use crate::duckdb::BindInputRef;
 use crate::duckdb::BindResultRef;
 use crate::duckdb::Cardinality;
 use crate::duckdb::ClientContextRef;
+use crate::duckdb::ColumnStatistics;
 use crate::duckdb::DataChunkRef;
 use crate::duckdb::ExpressionRef;
 use crate::duckdb::LogicalType;
@@ -80,6 +82,8 @@ use crate::exporter::ConversionCache;
 static EMPTY_COLUMN_IDX: u64 = 18446744073709551614;
 static EMPTY_COLUMN_NAME: &str = "";
 
+pub type DataSourceWithStats = (DataSourceRef, Option<FileStatistics>);
+
 /// A trait for table functions that resolve to a [`DataSourceRef`].
 ///
 /// Implementors only need to define how parameters are declared and how binding produces a
@@ -97,7 +101,14 @@ pub(crate) trait DataSourceTableFunction: Sized + Debug {
     }
 
     /// Bind the table function and return a [`DataSourceRef`].
-    fn bind(ctx: &ClientContextRef, input: &BindInputRef) -> VortexResult<DataSourceRef>;
+    fn bind(ctx: &ClientContextRef, input: &BindInputRef) -> VortexResult<DataSourceWithStats>;
+}
+
+#[derive(Clone)]
+pub enum DataSourceStatistics {
+    All(Vec<ColumnStatistics>),
+    /// Dummy column to return a reference to
+    None(ColumnStatistics),
 }
 
 /// Bind data produced by a [`DataSourceTableFunction`].
@@ -106,6 +117,7 @@ pub struct DataSourceBindData {
     filter_exprs: Vec<Expression>,
     column_names: Vec<String>,
     column_types: Vec<LogicalType>,
+    stats: DataSourceStatistics,
 }
 
 impl Clone for DataSourceBindData {
@@ -116,6 +128,7 @@ impl Clone for DataSourceBindData {
             filter_exprs: vec![],
             column_names: self.column_names.clone(),
             column_types: self.column_types.clone(),
+            stats: self.stats.clone(),
         }
     }
 }
@@ -189,19 +202,39 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         input: &BindInputRef,
         result: &mut BindResultRef,
     ) -> VortexResult<Self::BindData> {
-        let data_source = T::bind(ctx, input)?;
-
+        let (data_source, file_stats) = T::bind(ctx, input)?;
         let (column_names, column_types) = extract_schema_from_dtype(data_source.dtype())?;
 
-        for (column_name, column_type) in column_names.iter().zip(&column_types) {
-            result.add_result_column(column_name, column_type);
+        if file_stats.is_none() {
+            for i in 0..column_names.len() {
+                result.add_result_column(&column_names[i], &column_types[i]);
+            }
+
+            let stats = DataSourceStatistics::None(ColumnStatistics::default());
+            return Ok(DataSourceBindData {
+                data_source,
+                filter_exprs: vec![],
+                column_names,
+                column_types,
+                stats,
+            });
         }
+        let file_stats = file_stats.vortex_expect("no stats");
+
+        let mut stats = Vec::with_capacity(column_names.len());
+        for i in 0..column_names.len() {
+            result.add_result_column(&column_names[i], &column_types[i]);
+            let (stats_set, dtype) = file_stats.get(i);
+            stats.push(ColumnStatistics::new(stats_set, dtype.clone()));
+        }
+        let stats = DataSourceStatistics::All(stats);
 
         Ok(DataSourceBindData {
             data_source,
             filter_exprs: vec![],
             column_names,
             column_types,
+            stats,
         })
     }
 
@@ -412,6 +445,17 @@ impl<T: DataSourceTableFunction> TableFunction for T {
         Ok(false)
     }
 
+    fn statistics<'a>(
+        _client_context: &ClientContextRef,
+        bind_data: &'a Self::BindData,
+        column_index: usize,
+    ) -> &'a ColumnStatistics {
+        match &bind_data.stats {
+            DataSourceStatistics::All(items) => &items[column_index],
+            DataSourceStatistics::None(dummy) => dummy,
+        }
+    }
+
     fn cardinality(bind_data: &Self::BindData) -> Cardinality {
         match bind_data.data_source.row_count() {
             Some(Precision::Exact(v)) => Cardinality::Maximum(v),