Sum aggregate doesn't include nan values

vortex-array/src/aggregate_fn/fns/sum/mod.rs

@@ -203,7 +203,11 @@ impl AggregateFnVTable for Sum {
 
     #[inline]
     fn is_saturated(&self, partial: &Self::Partial) -> bool {
-        partial.current.is_none()
+        match partial.current.as_ref() {
+            None => true,
+            Some(SumState::Float(v)) => v.is_nan(),
+            Some(_) => false,
+        }
     }
 
     fn accumulate(

vortex-array/src/aggregate_fn/fns/sum/primitive.rs

@@ -54,7 +54,9 @@ fn accumulate_primitive_all(inner: &mut SumState, p: &PrimitiveArray) -> VortexR
             signed: |_T| { vortex_panic!("float sum state with signed input") },
             floating: |T| {
                 for &v in p.as_slice::<T>() {
-                    *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64");
+                    if !v.is_nan() {
+                        *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64");
+                    }
                 }
                 Ok(false)
             }
@@ -98,7 +100,7 @@ fn accumulate_primitive_valid(
             signed: |_T| { vortex_panic!("float sum state with signed input") },
             floating: |T| {
                 for (&v, valid) in p.as_slice::<T>().iter().zip_eq(validity.iter()) {
-                    if valid {
+                    if valid && !v.is_nan() {
                         *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64");
                     }
                 }
@@ -213,6 +215,66 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn sum_f64_with_nan() -> VortexResult<()> {
+        let arr = PrimitiveArray::new(
+            buffer![1.0f64, f64::NAN, 2.0, f64::NAN, 3.0],
+            Validity::NonNullable,
+        )
+        .into_array();
+        let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?;
+        assert_eq!(result.as_primitive().typed_value::<f64>(), Some(6.0));
+        Ok(())
+    }
+
+    #[test]
+    fn sum_f32_with_nan() -> VortexResult<()> {
+        let arr =
+            PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array();
+        let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?;
+        assert_eq!(result.as_primitive().typed_value::<f64>(), Some(5.0));
+        Ok(())
+    }
+
+    #[test]
+    fn sum_f64_with_nan_and_nulls() -> VortexResult<()> {
+        let arr = PrimitiveArray::from_option_iter([Some(1.0f64), None, Some(f64::NAN), Some(3.0)])
+            .into_array();
+        let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?;
+        assert_eq!(result.as_primitive().typed_value::<f64>(), Some(4.0));
+        Ok(())
+    }
+
+    #[test]
+    fn sum_all_nan() -> VortexResult<()> {
+        let arr =
+            PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array();
+        let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?;
+        assert_eq!(result.as_primitive().typed_value::<f64>(), Some(0.0));
+        Ok(())
+    }
+
+    #[test]
+    fn sum_f64_with_infinity() -> VortexResult<()> {
+        let batch = PrimitiveArray::new(
+            buffer![1.0f64, f64::INFINITY, f64::NEG_INFINITY, 2.0],
+            Validity::NonNullable,
+        )
+        .into_array();
+        let acc = sum(&batch, &mut LEGACY_SESSION.create_execution_ctx())?;
+        // INFINITY + NEG_INFINITY = NaN, which is treated as saturated
+        assert!(acc.as_primitive().typed_value::<f64>().unwrap().is_nan());
+
+        let mut acc = Accumulator::try_new(
+            Sum,
+            EmptyOptions,
+            DType::Primitive(PType::F64, Nullability::NonNullable),
+        )?;
+        acc.accumulate(&batch, &mut LEGACY_SESSION.create_execution_ctx())?;
+        assert!(acc.is_saturated());
+        Ok(())
+    }
+
     #[test]
     fn sum_checked_overflow() -> VortexResult<()> {
         let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array();