fix: use total ordering in the min & max accumulator for floats

datafusion/physical-expr/src/aggregate/min_max.rs

@@ -488,6 +488,20 @@ macro_rules! typed_min_max {
     }};
 }
 
+macro_rules! typed_min_max_float {
+    ($VALUE:expr, $DELTA:expr, $SCALAR:ident, $OP:ident) => {{
+        ScalarValue::$SCALAR(match ($VALUE, $DELTA) {
+            (None, None) => None,
+            (Some(a), None) => Some(*a),
+            (None, Some(b)) => Some(*b),
+            (Some(a), Some(b)) => match a.total_cmp(b) {
+                choose_min_max!($OP) => Some(*b),
+                _ => Some(*a),
+            },
+        })
+    }};
+}
+
 // min/max of two scalar string values.
 macro_rules! typed_min_max_string {
     ($VALUE:expr, $DELTA:expr, $SCALAR:ident, $OP:ident) => {{
@@ -500,7 +514,7 @@ macro_rules! typed_min_max_string {
     }};
 }
 
-macro_rules! interval_choose_min_max {
+macro_rules! choose_min_max {
     (min) => {
         std::cmp::Ordering::Greater
     };
@@ -512,7 +526,7 @@ macro_rules! interval_choose_min_max {
 macro_rules! interval_min_max {
     ($OP:tt, $LHS:expr, $RHS:expr) => {{
         match $LHS.partial_cmp(&$RHS) {
-            Some(interval_choose_min_max!($OP)) => $RHS.clone(),
+            Some(choose_min_max!($OP)) => $RHS.clone(),
             Some(_) => $LHS.clone(),
             None => {
                 return internal_err!("Comparison error while computing interval min/max")
@@ -555,10 +569,10 @@ macro_rules! min_max {
                 typed_min_max!(lhs, rhs, Boolean, $OP)
             }
             (ScalarValue::Float64(lhs), ScalarValue::Float64(rhs)) => {
-                typed_min_max!(lhs, rhs, Float64, $OP)
+                typed_min_max_float!(lhs, rhs, Float64, $OP)
             }
             (ScalarValue::Float32(lhs), ScalarValue::Float32(rhs)) => {
-                typed_min_max!(lhs, rhs, Float32, $OP)
+                typed_min_max_float!(lhs, rhs, Float32, $OP)
             }
             (ScalarValue::UInt64(lhs), ScalarValue::UInt64(rhs)) => {
                 typed_min_max!(lhs, rhs, UInt64, $OP)
@@ -1103,3 +1117,36 @@ impl Accumulator for SlidingMinAccumulator {
         std::mem::size_of_val(self) - std::mem::size_of_val(&self.min) + self.min.size()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn float_max_with_nans() {
+        let pos_nan = f32::NAN;
+        let zero = 0_f32;
+
+        let vals_float = Float32Array::from_iter_values([zero, pos_nan].iter().copied());
+
+        let vals: ArrayRef = Arc::new(vals_float.clone());
+
+        // We accumulate the above two rows in two different ways.  First, we pass in both as a single batch
+        let mut accumulator = MaxAccumulator::try_new(&DataType::Float32).unwrap();
+        accumulator.update_batch(&[vals]).unwrap();
+        let single_batch_result = &accumulator.evaluate().unwrap();
+
+        // Next we pass the two values in two different batches.
+        let vals_a: ArrayRef =
+            Arc::new(Float32Array::from_iter_values([pos_nan].iter().copied()));
+        let vals_b: ArrayRef =
+            Arc::new(Float32Array::from_iter_values([zero].iter().copied()));
+
+        let mut accumulator = MaxAccumulator::try_new(&DataType::Float32).unwrap();
+        accumulator.update_batch(&[vals_a]).unwrap();
+        accumulator.update_batch(&[vals_b]).unwrap();
+        let split_batch_result = &accumulator.evaluate().unwrap();
+
+        assert_eq!(single_batch_result, split_batch_result);
+    }
+}