improve distinct count approximation to datafusion statistics

Author: JanKaul

datafusion_iceberg/src/statistics.rs

@@ -4,6 +4,7 @@ use datafusion::{
     scalar::ScalarValue,
 };
 use iceberg_rust::error::Error;
+use iceberg_rust::file_format::parquet::estimate_distinct_count;
 use iceberg_rust::spec::{
     manifest::{ManifestEntry, Status},
     schema::Schema,
@@ -45,12 +46,16 @@ pub(crate) fn statistics_from_datafiles(
                         .column_statistics
                         .into_iter()
                         .zip(column_stats)
-                        .map(|(acc, x)| ColumnStatistics {
-                            null_count: acc.null_count.add(&x.null_count),
-                            max_value: acc.max_value.max(&x.max_value),
-                            min_value: acc.min_value.min(&x.min_value),
-                            distinct_count: acc.distinct_count.add(&x.distinct_count),
-                            sum_value: acc.sum_value.add(&x.sum_value),
+                        .map(|(acc, x)| {
+                            let new_distinct_count = new_distinct_count(&acc, &x);
+
+                            ColumnStatistics {
+                                null_count: acc.null_count.add(&x.null_count),
+                                max_value: acc.max_value.max(&x.max_value),
+                                min_value: acc.min_value.min(&x.min_value),
+                                distinct_count: new_distinct_count,
+                                sum_value: acc.sum_value.add(&x.sum_value),
+                            }
                         })
                         .collect(),
                 }
@@ -134,3 +139,49 @@ fn convert_value_to_scalar_value(value: Value) -> Result<ScalarValue, Error> {
         )),
     }
 }
+
+fn new_distinct_count(acc: &ColumnStatistics, x: &ColumnStatistics) -> Precision<usize> {
+    match (
+        &acc.distinct_count,
+        &x.distinct_count,
+        &acc.min_value,
+        &acc.max_value,
+        &x.min_value,
+        &x.max_value,
+    ) {
+        (
+            Precision::Exact(old_count),
+            Precision::Exact(new_count),
+            Precision::Exact(ScalarValue::Int32(Some(old_min))),
+            Precision::Exact(ScalarValue::Int32(Some(old_max))),
+            Precision::Exact(ScalarValue::Int32(Some(new_min))),
+            Precision::Exact(ScalarValue::Int32(Some(new_max))),
+        ) => {
+            let estimated = estimate_distinct_count(
+                &[old_min, old_max],
+                &[new_min, new_max],
+                *old_count as i64,
+                *new_count as i64,
+            );
+            Precision::Inexact(*old_count + estimated as usize)
+        }
+        (
+            Precision::Exact(old_count),
+            Precision::Exact(new_count),
+            Precision::Exact(ScalarValue::Int64(Some(old_min))),
+            Precision::Exact(ScalarValue::Int64(Some(old_max))),
+            Precision::Exact(ScalarValue::Int64(Some(new_min))),
+            Precision::Exact(ScalarValue::Int64(Some(new_max))),
+        ) => {
+            let estimated = estimate_distinct_count(
+                &[old_min, old_max],
+                &[new_min, new_max],
+                *old_count as i64,
+                *new_count as i64,
+            );
+            Precision::Inexact(*old_count + estimated as usize)
+        }
+        (Precision::Absent, Precision::Exact(_), _, _, _, _) => x.distinct_count,
+        _ => acc.distinct_count.add(&x.distinct_count),
+    }
+}

iceberg-rust/src/file_format/parquet.rs

@@ -334,8 +334,20 @@ fn range_overlap<T: Ord + Sub + Copy>(
     overlap_end - overlap_start
 }
 
-/// Helper trait to convert numeric types to f64 for statistical calculations
-trait ToF64 {
+/// Helper trait to convert numeric types to f64 for statistical calculations.
+///
+/// This trait provides a uniform interface for converting integer types to f64,
+/// which is necessary for the statistical estimation algorithms. The conversion
+/// may be lossy for very large i64 values (beyond 2^53), but this is acceptable
+/// for statistical approximations.
+pub trait ToF64 {
+    /// Converts the value to f64.
+    ///
+    /// # Note
+    ///
+    /// For i64 values larger than 2^53, precision may be lost in the conversion.
+    /// This is acceptable for statistical calculations where exact precision is
+    /// not required.
     fn to_f64(self) -> f64;
 }
 
@@ -384,7 +396,7 @@ impl ToF64 for i64 {
 /// // New range [500, 1500] with 50 distinct values
 /// let new_count = estimate_distinct_count(&[&0, &1000], &[&500, &1500], 100, 50);
 /// ```
-fn estimate_distinct_count<T>(
+pub fn estimate_distinct_count<T>(
     old_range: &[&T; 2],
     new_range: &[&T; 2],
     old_distinct_count: i64,