fix: account for NaN in float stats exported to DuckDB and DataFusion

vortex-datafusion/src/convert/stats/mod.rs

@@ -43,13 +43,17 @@ pub(crate) fn stats_set_to_df(
     })
 }
 
-/// Read one scalar-valued statistic and convert it to DataFusion, or `Absent` if it does not apply
-/// to `dtype` or cannot be represented.
+/// Read one NaN-skipping statistic (min, max or sum) and convert it to DataFusion, or `Absent` if it
+/// does not apply to `dtype`, cannot be represented, or the column may hold NaN.
 fn scalar_stat_to_df(
     stats_set: &StatsSet,
     stat: Stat,
     dtype: &DType,
 ) -> VortexPrecision<ScalarValue> {
+    if !is_nan_free(stats_set, dtype) {
+        return VortexPrecision::Absent;
+    }
+
     stats_set.get(stat).and_then(|stat_value| {
         Scalar::try_new(stat.dtype(dtype)?, Some(stat_value))
             .ok()?
@@ -58,6 +62,12 @@ fn scalar_stat_to_df(
     })
 }
 
+/// Whether `stats_set` proves the column holds no NaN.
+pub(crate) fn is_nan_free(stats_set: &StatsSet, dtype: &DType) -> bool {
+    !dtype.is_float()
+        || stats_set.get_as::<u64>(Stat::NaNCount, &PType::U64.into()) == VortexPrecision::Exact(0)
+}
+
 /// Infer a singleton from exact bounds, excluding nulls and, for floats, NaNs.
 pub(crate) fn constant_distinct_count(
     stats_set: &StatsSet,
@@ -90,9 +100,7 @@ pub(crate) fn constant_distinct_count(
     }
 
     // Float extrema skip NaNs, so equal extrema alone cannot prove a constant.
-    if dtype.is_float()
-        && stats_set.get_as::<u64>(Stat::NaNCount, &PType::U64.into()) != VortexPrecision::Exact(0)
-    {
+    if !is_nan_free(stats_set, dtype) {
         return Precision::Absent;
     }
 

vortex-datafusion/src/convert/stats/tests.rs

@@ -173,3 +173,34 @@ fn convertible_statistics_keep_their_precision(
     assert_eq!(stats.min_value, expected);
     Ok(())
 }
+
+#[rstest]
+#[case::nan_free(VortexPrecision::exact(0u64), true)]
+#[case::has_nan(VortexPrecision::exact(1u64), false)]
+#[case::inexact(VortexPrecision::inexact(0u64), false)]
+#[case::unknown(VortexPrecision::Absent, false)]
+fn float_statistics_require_no_nans(
+    #[case] nan_count: VortexPrecision<u64>,
+    #[case] reported: bool,
+) -> VortexResult<()> {
+    let stats_set = StatsSet::from_iter([
+        (Stat::Min, VortexPrecision::exact(ScalarValue::from(1.0f64))),
+        (Stat::Max, VortexPrecision::exact(ScalarValue::from(2.0f64))),
+        (Stat::Sum, VortexPrecision::exact(ScalarValue::from(3.0f64))),
+        (Stat::NaNCount, nan_count.map(ScalarValue::from)),
+    ]);
+    let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
+    let stats = stats_set_to_df(&stats_set, &dtype)?;
+
+    let expected = |value: f64| {
+        if reported {
+            Precision::Exact(DFScalarValue::Float64(Some(value)))
+        } else {
+            Precision::Absent
+        }
+    };
+    assert_eq!(stats.min_value, expected(1.0));
+    assert_eq!(stats.max_value, expected(2.0));
+    assert_eq!(stats.sum_value, expected(3.0));
+    Ok(())
+}

vortex-datafusion/src/persistent/format.rs

@@ -71,6 +71,7 @@ use crate::PrecisionExt as _;
 use crate::convert::ExpressionConverter;
 use crate::convert::TryToDataFusion;
 use crate::convert::stats::constant_distinct_count;
+use crate::convert::stats::is_nan_free;
 
 const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE;
 
@@ -671,19 +672,16 @@ impl FileFormat for VortexFormat {
                 let column_size =
                     stats_set.get_as::<usize>(Stat::UncompressedSizeInBytes, &PType::U64.into());
 
-                let min = scalar_stat_to_df(
-                    Stat::Min,
-                    stats_set.get(Stat::Min),
-                    stats_dtype,
-                    field.data_type(),
-                );
-
-                let max = scalar_stat_to_df(
-                    Stat::Max,
-                    stats_set.get(Stat::Max),
-                    stats_dtype,
-                    field.data_type(),
-                );
+                let nan_free = is_nan_free(stats_set, stats_dtype);
+                let bound = |stat| {
+                    if nan_free {
+                        scalar_stat_to_df(stat, stats_set.get(stat), stats_dtype, field.data_type())
+                    } else {
+                        DFPrecision::Absent
+                    }
+                };
+                let min = bound(Stat::Min);
+                let max = bound(Stat::Max);
 
                 let null_count = stats_set
                     .get_as::<usize>(Stat::NullCount, &PType::U64.into())

vortex-duckdb/src/column_statistics.rs

@@ -1,8 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex::array::match_each_float_ptype;
 use vortex::array::stats::StatsSet;
 use vortex::dtype::DType;
+use vortex::dtype::PType;
 use vortex::error::VortexExpect as _;
 use vortex::error::VortexResult;
 use vortex::expr::stats::Precision;
@@ -31,7 +33,16 @@ impl ColumnStatistics {
                 .ok()
         };
         let min = stats.min.and_then(to_value);
-        let max = stats.max.and_then(to_value);
+        // Vortex's max skips NaN, while DuckDB orders NaN above every other value.
+        let max = if !dtype.is_float() || stats.nan_count == Some(0) {
+            stats.max.and_then(to_value)
+        } else if stats.nan_count.is_some() {
+            to_value(match_each_float_ptype!(dtype.as_ptype(), |F| {
+                ScalarValue::from(F::NAN)
+            }))
+        } else {
+            None
+        };
 
         let max_string_length = stats
             .max_string_length
@@ -59,6 +70,7 @@ pub struct ColumnStatisticsAggregate {
     pub max_string_length: Option<u32>,
     /// May be true if null count stat isn't present
     pub has_null: bool,
+    pub nan_count: Option<u64>,
 }
 
 impl ColumnStatisticsAggregate {
@@ -86,11 +98,16 @@ impl ColumnStatisticsAggregate {
             _ => true,
         };
 
+        let nan_count = stats
+            .get_as::<u64>(Stat::NaNCount, &PType::U64.into())
+            .as_exact();
+
         Self {
             min,
             max,
             max_string_length,
             has_null,
+            nan_count,
         }
     }
 }

vortex-sqllogictest/slt/aggregates_edge_cases.slt

@@ -50,6 +50,25 @@ SELECT count(x), sum(x), min(x), max(x), avg(x) FROM '${WORK_DIR}/f-nan.vortex';
 ----
 3 NaN 1 NaN NaN
 
+query R
+SELECT max(x) FROM '${WORK_DIR}/f-nan.vortex';
+----
+NaN
+
+query I
+SELECT count(*) FROM '${WORK_DIR}/f-nan.vortex' WHERE x = 'NaN'::DOUBLE;
+----
+1
+
+statement ok
+COPY (SELECT -x AS x FROM (VALUES ('NaN'::DOUBLE),(-1.0),(-2.0)) AS t(x)) TO '${WORK_DIR}/f-negnan.vortex';
+
+onlyif datafusion
+query I
+SELECT count(*) FROM '${WORK_DIR}/f-negnan.vortex' WHERE x < 0.5;
+----
+1
+
 # all NaN
 query IRRRR
 SELECT count(x), sum(x), min(x), max(x), avg(x) FROM '${WORK_DIR}/f-allnan.vortex';

vortex-sqllogictest/slt/datafusion/float_min_max_stats.slt

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+include ../setup.slt.no
+
+query I
+COPY (SELECT CAST(column1 AS DOUBLE) AS x FROM (VALUES (1), (2), (3))) TO '${WORK_DIR}/float_min_max_stats.vortex';
+----
+3
+
+query TT
+EXPLAIN SELECT max(x) FROM '${WORK_DIR}/float_min_max_stats.vortex';
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[max(${WORK_DIR}/float_min_max_stats.vortex.x)]]
+02)--TableScan: ${WORK_DIR}/float_min_max_stats.vortex projection=[x]
+physical_plan
+01)ProjectionExec: expr=[3 as max(${WORK_DIR}/float_min_max_stats.vortex.x)]
+02)--PlaceholderRowExec

vortex-sqllogictest/slt/duckdb/agg_partition_stats.slt

@@ -69,6 +69,34 @@ SELECT min(d), max(d) FROM '${WORK_DIR}/agg_partition_stats-date.vortex';
 ----
 2013-07-01 2013-07-30
 
+statement ok
+COPY (SELECT i::DOUBLE AS x FROM generate_series(1, 5000) t(i))
+TO '${WORK_DIR}/agg_partition_stats-double.vortex';
+
+query TT
+EXPLAIN SELECT max(x) FROM '${WORK_DIR}/agg_partition_stats-double.vortex';
+----
+<REGEX>:.*COLUMN_DATA_SCAN.*
+
+query R
+SELECT max(x) FROM '${WORK_DIR}/agg_partition_stats-double.vortex';
+----
+5000
+
+statement ok
+COPY (SELECT CASE WHEN i = 2500 THEN 'nan'::DOUBLE ELSE i::DOUBLE END AS x FROM generate_series(1, 5000) t(i))
+TO '${WORK_DIR}/agg_partition_stats-double-nan.vortex';
+
+query TT
+EXPLAIN SELECT min(x), max(x) FROM '${WORK_DIR}/agg_partition_stats-double-nan.vortex';
+----
+<REGEX>:.*COLUMN_DATA_SCAN.*
+
+query RR
+SELECT min(x), max(x) FROM '${WORK_DIR}/agg_partition_stats-double-nan.vortex';
+----
+1 NaN
+
 query TT
 EXPLAIN SELECT count(*) FROM '${WORK_DIR}/agg_partition_stats.vortex'
 WHERE i > 50000;