fix: normalize floating-point values in native collect_set

spark/src/main/scala/org/apache/comet/serde/aggregates.scala

@@ -22,13 +22,14 @@ package org.apache.comet.serde
 import scala.jdk.CollectionConverters._
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Literal}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, HyperLogLogPlusPlus, Last, Max, Min, Percentile, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Complete, Corr, Count, Covariance, CovPopulation, CovSample, First, HyperLogLogPlusPlus, Last, Max, Min, Partial, Percentile, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
 import org.apache.spark.sql.catalyst.util.ArrayData
+import org.apache.spark.sql.comet.CometExecUtils
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, NumericType, ShortType, StringType, TimestampNTZType, TimestampType}
 
 import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT
-import org.apache.comet.CometSparkSessionExtensions.{isSpark41Plus, withFallbackReason}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark41Plus, isSpark42Plus, withFallbackReason}
 import org.apache.comet.expressions.CometEvalMode
 import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType}
 import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil}
@@ -851,11 +852,19 @@ object CometBloomFilterAggregate extends CometAggregateExpressionSerde[BloomFilt
 
 object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] {
 
-  override def getIncompatibleReasons(): Seq[String] = Seq(
-    "Comet deduplicates NaN values (treats `NaN == NaN`) while Spark treats each NaN as a" +
-      s" distinct value. When `${COMET_EXEC_STRICT_FLOATING_POINT.key}=true`, `collect_set`" +
-      " on floating-point types falls back to Spark unless" +
-      " `spark.comet.expression.CollectSet.allowIncompatible=true` is set.")
+  override def getIncompatibleReasons(): Seq[String] = {
+    if (isSpark42Plus) {
+      Nil
+    } else {
+      Seq(
+        "Before Spark 4.2, Comet deduplicates NaN values (treats `NaN == NaN`) while Spark" +
+          " treats each NaN as a distinct value. Comet treats -0.0 and 0.0 as distinct while" +
+          " Spark treats them as equal." +
+          s" When `${COMET_EXEC_STRICT_FLOATING_POINT.key}=true`, `collect_set` on" +
+          " floating-point types falls back to Spark on those versions unless" +
+          " `spark.comet.expression.CollectSet.allowIncompatible=true` is set.")
+    }
+  }
 
   override def getSupportLevel(expr: CollectSet): SupportLevel = {
     // The native path always drops null inputs. Spark 4.2 added an `ignoreNulls` field to
@@ -865,12 +874,14 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] {
     // analysis time, and CometCollectShim.ignoreNulls hardcodes true, making this a no-op.
     if (!CometCollectShim.ignoreNulls(expr)) {
       Unsupported(Some("collect_set with RESPECT NULLS (ignoreNulls = false) is not supported"))
+    } else if (isSpark42Plus) {
+      Compatible()
     } else {
       SupportLevel
         .strictFloatingPointReason(
           expr.children.head.dataType,
           "collect_set on floating-point types " +
-            "(Comet deduplicates NaN values while Spark treats each NaN as distinct)")
+            "(Comet deduplicates NaN values and distinguishes -0.0 from 0.0, unlike Spark)")
         .map(reason => Incompatible(Some(reason)))
         .getOrElse(Compatible())
     }
@@ -882,7 +893,12 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] {
       inputs: Seq[Attribute],
       binding: Boolean,
       conf: SQLConf): Option[ExprOuterClass.AggExpr] = {
-    val child = expr.children.head
+    val child = aggExpr.mode match {
+      case Partial | Complete if isSpark42Plus =>
+        CometExecUtils.normalizeFloatingNumbers(expr.children.head)
+      case _ =>
+        expr.children.head
+    }
     val childExpr = exprToProto(child, inputs, binding)
     val dataType = serializeDataType(expr.dataType)
 

spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala

@@ -24,7 +24,8 @@ import scala.reflect.ClassTag
 
 import org.apache.spark.{Partition, SparkContext, TaskContext}
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.catalyst.expressions.{Attribute, NamedExpression, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, NamedExpression, SortOrder}
+import org.apache.spark.sql.catalyst.optimizer.NormalizeFloatingNumbers
 import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.SparkPlan
@@ -36,6 +37,10 @@ import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType}
 
 object CometExecUtils {
 
+  /** Expose Spark's package-private recursive floating-point normalizer to Comet serde. */
+  def normalizeFloatingNumbers(expr: Expression): Expression =
+    NormalizeFloatingNumbers.normalize(expr)
+
   /**
    * Create an empty RDD with the given number of partitions.
    */

spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql

@@ -15,7 +15,6 @@
 -- specific language governing permissions and limitations
 -- under the License.
 
--- Config: spark.comet.exec.strictFloatingPoint=true
 -- ConfigMatrix: parquet.enable.dictionary=false,true
 
 -- ============================================================
@@ -147,42 +146,6 @@ SELECT grp, sort_array(collect_set(i)) FROM cs_src_intbig GROUP BY grp ORDER BY
 query
 SELECT grp, sort_array(collect_set(bi)) FROM cs_src_intbig GROUP BY grp ORDER BY grp
 
--- ============================================================
--- Float (with NULLs, NaN, Inf, -Inf, +0, -0)
--- Comet deduplicates NaN while Spark does not; with
--- strictFloatingPoint=true collect_set falls back to Spark.
--- ============================================================
-
-statement
-CREATE TABLE cs_src_float(v float, grp string) USING parquet
-
-statement
-INSERT INTO cs_src_float VALUES
-  (1.5,                'a'), (2.5,                'a'), (1.5, 'a'), (NULL, 'a'),
-  (CAST('NaN' AS FLOAT), 'b'), (CAST('NaN' AS FLOAT), 'b'), (1.0, 'b'),
-  (CAST('Infinity' AS FLOAT), 'c'), (CAST('-Infinity' AS FLOAT), 'c'), (CAST('Infinity' AS FLOAT), 'c'),
-  (CAST(0.0 AS FLOAT), 'd'), (CAST(-0.0 AS FLOAT), 'd'), (1.0, 'd'), (NULL, 'd')
-
-query expect_fallback(not fully compatible with Spark)
-SELECT grp, sort_array(collect_set(v)) FROM cs_src_float GROUP BY grp ORDER BY grp
-
--- ============================================================
--- Double (with NULLs, NaN, Inf, -Inf, +0, -0)
--- ============================================================
-
-statement
-CREATE TABLE cs_src_double(v double, grp string) USING parquet
-
-statement
-INSERT INTO cs_src_double VALUES
-  (1.1,   'a'), (2.2,   'a'), (1.1, 'a'), (NULL, 'a'),
-  (CAST('NaN' AS DOUBLE), 'b'), (CAST('NaN' AS DOUBLE), 'b'), (1.0, 'b'),
-  (CAST('Infinity' AS DOUBLE), 'c'), (CAST('-Infinity' AS DOUBLE), 'c'), (CAST('Infinity' AS DOUBLE), 'c'),
-  (0.0,  'd'), (-0.0,  'd'), (1.0, 'd'), (NULL, 'd')
-
-query expect_fallback(not fully compatible with Spark)
-SELECT grp, sort_array(collect_set(v)) FROM cs_src_double GROUP BY grp ORDER BY grp
-
 -- ============================================================
 -- String (with NULLs and empty string)
 -- ============================================================

spark/src/test/resources/sql-tests/expressions/aggregate/collect_set_floating_fallback.sql

@@ -0,0 +1,49 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements.  See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership.  The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License.  You may obtain a copy of the License at
+--
+--   http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied.  See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- MaxSparkVersion: 4.1
+-- Config: spark.comet.exec.strictFloatingPoint=true
+
+statement
+CREATE TABLE cs_fallback_float(v float, grp string) USING parquet
+
+statement
+INSERT INTO cs_fallback_float VALUES
+  (1.5, 'a'), (2.5, 'a'), (1.5, 'a'), (NULL, 'a'),
+  (CAST('NaN' AS FLOAT), 'b'), (CAST('NaN' AS FLOAT), 'b'), (1.0, 'b'),
+  (CAST('Infinity' AS FLOAT), 'c'), (CAST('-Infinity' AS FLOAT), 'c'),
+  (CAST('Infinity' AS FLOAT), 'c'),
+  (CAST(0.0 AS FLOAT), 'd'), (CAST('-0.0' AS FLOAT), 'd'), (1.0, 'd'), (NULL, 'd')
+
+query expect_fallback(not fully compatible with Spark)
+SELECT grp, sort_array(collect_set(v))
+FROM cs_fallback_float GROUP BY grp ORDER BY grp
+
+statement
+CREATE TABLE cs_fallback_double(v double, grp string) USING parquet
+
+statement
+INSERT INTO cs_fallback_double VALUES
+  (1.1, 'a'), (2.2, 'a'), (1.1, 'a'), (NULL, 'a'),
+  (CAST('NaN' AS DOUBLE), 'b'), (CAST('NaN' AS DOUBLE), 'b'), (1.0, 'b'),
+  (CAST('Infinity' AS DOUBLE), 'c'), (CAST('-Infinity' AS DOUBLE), 'c'),
+  (CAST('Infinity' AS DOUBLE), 'c'),
+  (0.0, 'd'), (CAST('-0.0' AS DOUBLE), 'd'), (1.0, 'd'), (NULL, 'd')
+
+query expect_fallback(not fully compatible with Spark)
+SELECT grp, sort_array(collect_set(v))
+FROM cs_fallback_double GROUP BY grp ORDER BY grp

spark/src/test/resources/sql-tests/expressions/aggregate/collect_set_normalization.sql

@@ -0,0 +1,112 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements.  See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership.  The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License.  You may obtain a copy of the License at
+--
+--   http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied.  See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- MinSparkVersion: 4.2
+-- Config: spark.comet.exec.strictFloatingPoint=true
+-- ConfigMatrix: parquet.enable.dictionary=false,true
+
+statement
+CREATE TABLE cs_norm_scalar(grp string, f float, d double) USING parquet
+
+statement
+INSERT INTO cs_norm_scalar VALUES
+  ('ordinary', CAST(1.5 AS FLOAT), CAST(1.1 AS DOUBLE)),
+  ('ordinary', CAST(2.5 AS FLOAT), CAST(2.2 AS DOUBLE)),
+  ('ordinary', CAST(1.5 AS FLOAT), CAST(1.1 AS DOUBLE)),
+  ('ordinary', CAST(NULL AS FLOAT), CAST(NULL AS DOUBLE)),
+  ('nan', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE)),
+  ('nan', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE)),
+  ('nan', CAST(1.0 AS FLOAT), CAST(1.0 AS DOUBLE)),
+  ('infinity', CAST('Infinity' AS FLOAT), CAST('Infinity' AS DOUBLE)),
+  ('infinity', CAST('-Infinity' AS FLOAT), CAST('-Infinity' AS DOUBLE)),
+  ('infinity', CAST('Infinity' AS FLOAT), CAST('Infinity' AS DOUBLE)),
+  ('zero', CAST('-0.0' AS FLOAT), CAST('-0.0' AS DOUBLE)),
+  ('zero', CAST(0.0 AS FLOAT), CAST(0.0 AS DOUBLE)),
+  ('zero', CAST(1.0 AS FLOAT), CAST(1.0 AS DOUBLE)),
+  ('zero', CAST(NULL AS FLOAT), CAST(NULL AS DOUBLE))
+
+-- Repartition so equal values must also deduplicate across partial buffers.
+query
+SELECT grp, sort_array(collect_set(f)), sort_array(collect_set(d))
+FROM (SELECT /*+ REPARTITION(3) */ * FROM cs_norm_scalar)
+GROUP BY grp
+ORDER BY grp
+
+-- Exercise the no-GROUP BY aggregate shape while merging partial buffers.
+query
+SELECT sort_array(collect_set(f)), sort_array(collect_set(d))
+FROM (
+  SELECT /*+ REPARTITION(3) */ f, d
+  FROM cs_norm_scalar
+  WHERE grp IN ('nan', 'zero')
+)
+
+statement
+CREATE TABLE cs_norm_nested(
+  grp string,
+  s struct<v:double>,
+  a array<float>,
+  deep_a array<struct<v:double>>,
+  deep_s struct<a:array<double>>) USING parquet
+
+statement
+INSERT INTO cs_norm_nested VALUES
+  ('nan',
+    named_struct('v', CAST('NaN' AS DOUBLE)),
+    array(CAST('NaN' AS FLOAT)),
+    array(named_struct('v', CAST('NaN' AS DOUBLE))),
+    named_struct('a', array(CAST('NaN' AS DOUBLE)))),
+  ('nan',
+    named_struct('v', CAST('NaN' AS DOUBLE)),
+    array(CAST('NaN' AS FLOAT)),
+    array(named_struct('v', CAST('NaN' AS DOUBLE))),
+    named_struct('a', array(CAST('NaN' AS DOUBLE)))),
+  ('zero',
+    named_struct('v', CAST('-0.0' AS DOUBLE)),
+    array(CAST('-0.0' AS FLOAT)),
+    array(named_struct('v', CAST('-0.0' AS DOUBLE))),
+    named_struct('a', array(CAST('-0.0' AS DOUBLE)))),
+  ('zero',
+    named_struct('v', CAST(0.0 AS DOUBLE)),
+    array(CAST(0.0 AS FLOAT)),
+    array(named_struct('v', CAST(0.0 AS DOUBLE))),
+    named_struct('a', array(CAST(0.0 AS DOUBLE)))),
+  ('null',
+    CAST(NULL AS STRUCT<v:DOUBLE>),
+    CAST(NULL AS ARRAY<FLOAT>),
+    CAST(NULL AS ARRAY<STRUCT<v:DOUBLE>>),
+    CAST(NULL AS STRUCT<a:ARRAY<DOUBLE>>)),
+  ('null',
+    named_struct('v', CAST(NULL AS DOUBLE)),
+    array(CAST(NULL AS FLOAT)),
+    array(CAST(NULL AS STRUCT<v:DOUBLE>)),
+    named_struct('a', array(CAST(NULL AS DOUBLE)))),
+  ('null',
+    named_struct('v', CAST(NULL AS DOUBLE)),
+    array(CAST(NULL AS FLOAT)),
+    array(CAST(NULL AS STRUCT<v:DOUBLE>)),
+    named_struct('a', array(CAST(NULL AS DOUBLE))))
+
+query
+SELECT grp,
+  sort_array(collect_set(s)),
+  sort_array(collect_set(a)),
+  sort_array(collect_set(deep_a)),
+  sort_array(collect_set(deep_s))
+FROM (SELECT /*+ REPARTITION(3) */ * FROM cs_norm_nested)
+GROUP BY grp
+ORDER BY grp