fix: stddev_pop should not directly return 0.0 when count is 1.0 (#1184)

* add test

* fix

* fix

* fix

Author: viirya

native/spark-expr/src/variance.rs

@@ -245,10 +245,8 @@ impl Accumulator for VarianceAccumulator {
 
         Ok(ScalarValue::Float64(match self.count {
             count if count == 0.0 => None,
-            count if count == 1.0 => {
-                if let StatsType::Population = self.stats_type {
-                    Some(0.0)
-                } else if self.null_on_divide_by_zero {
+            count if count == 1.0 && StatsType::Sample == self.stats_type => {
+                if self.null_on_divide_by_zero {
                     None
                 } else {
                     Some(f64::NAN)

spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala

@@ -38,6 +38,23 @@ import org.apache.comet.CometSparkSessionExtensions.isSpark34Plus
 class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   import testImplicits._
 
+  test("stddev_pop should return NaN for some cases") {
+    withSQLConf(
+      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+      CometConf.COMET_EXPR_STDDEV_ENABLED.key -> "true") {
+      Seq(true, false).foreach { nullOnDivideByZero =>
+        withSQLConf("spark.sql.legacy.statisticalAggregate" -> nullOnDivideByZero.toString) {
+
+          val data: Seq[(Float, Int)] = Seq((Float.PositiveInfinity, 1))
+          withParquetTable(data, "tbl", false) {
+            val df = sql("SELECT stddev_pop(_1), stddev_pop(_2) FROM tbl")
+            checkSparkAnswer(df)
+          }
+        }
+      }
+    }
+  }
+
   test("count with aggregation filter") {
     withSQLConf(
       CometConf.COMET_ENABLED.key -> "true",