feat: narrow strict floating-point admission for scalar sort keys

docs/source/user-guide/latest/compatibility/floating-point.md

@@ -43,6 +43,8 @@ Native sorting of floating-point values nested in arrays or structs still uses A
 ordering. Nested keys can therefore produce different ordering or rank results from Spark; see
 [#5507](https://github.com/apache/datafusion-comet/issues/5507).
 
-The existing `spark.comet.exec.strictFloatingPoint=true` fallback policy is unchanged, including
-its conservative fallback for scalar floating-point sort keys. Narrowing that scalar-sort
-admission policy is tracked in [#5506](https://github.com/apache/datafusion-comet/issues/5506).
+Because those scalar comparison keys match Spark, `spark.comet.exec.strictFloatingPoint=true` no
+longer forces a fallback for them: scalar `FLOAT` and `DOUBLE` sort keys, window and rank order
+keys, and range partitioning keys all stay native under strict mode. Floating-point values nested
+in arrays, structs, or maps still fall back under strict mode, because their ordering is the raw
+total ordering described above.

docs/source/user-guide/latest/tuning.md

@@ -216,9 +216,20 @@ suggested) to overlap I/O across files at the cost of extra memory.
 
 ## Optimizing Sorting on Floating-Point Values
 
-Sorting on floating-point data types (or complex types containing floating-point values) is not compatible with
-Spark if the data contains both zero and negative zero. This is likely an edge case that is not of concern for many users
-and sorting on floating-point data can be enabled by setting `spark.comet.expression.SortOrder.allowIncompatible=true`.
+Comet normalizes NaN payloads and signed zeros in scalar `FLOAT` and `DOUBLE` ordering keys, so `ORDER BY`, window
+ordering and range partitioning on them match Spark and stay native even with
+`spark.comet.exec.strictFloatingPoint=true`. Only the comparison key is normalized; returned values keep their original
+NaN representation and zero sign.
+
+Floating-point values nested in arrays, structs, or maps are compared with Arrow's raw total ordering instead, which can
+differ from Spark when the data contains both zero and negative zero, or more than one NaN representation. This is likely
+an edge case that is not of concern for many users. Setting `spark.comet.exec.strictFloatingPoint=true` makes those
+nested cases fall back to Spark, and they can be forced back onto the native path with
+`spark.comet.expression.SortOrder.allowIncompatible=true`.
+
+`sort_array` is separate. It sorts array elements rather than ordering rows, and its elements are compared with Arrow's
+raw total ordering, so `spark.comet.exec.strictFloatingPoint=true` makes it fall back even for a scalar floating-point
+element type. Use `spark.comet.expression.SortArray.allowIncompatible=true` to keep it native.
 
 ## Optimizing Joins
 

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -933,7 +933,11 @@ object CometConf extends ShimCometConf {
       .category(CATEGORY_EXEC)
       .doc(
         "When enabled, fall back to Spark for floating-point operations that may differ from " +
-          s"Spark, such as when comparing or sorting -0.0 and 0.0. $COMPAT_GUIDE.")
+          "Spark, such as comparing -0.0 and 0.0, sorting floating-point values nested in " +
+          "arrays, structs, or maps, or sorting the elements of a floating-point array with " +
+          "`sort_array`. Scalar `ORDER BY`, window ordering and range partitioning keys are " +
+          "unaffected, because Comet normalizes those comparison keys to match Spark. " +
+          s"$COMPAT_GUIDE.")
       .booleanConf
       .createWithDefault(false)
 

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

@@ -20,22 +20,38 @@
 package org.apache.comet.serde
 
 import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Descending, NullsFirst, NullsLast, SortOrder}
+import org.apache.spark.sql.types.{DoubleType, FloatType}
 
 import org.apache.comet.CometConf
 import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
 
 object CometSortOrder extends CometExpressionSerde[SortOrder] {
 
+  /**
+   * Subject of both the runtime fallback reason and the generated compatibility docs. Shared so
+   * the two cannot describe the policy differently.
+   */
+  private val nestedFloatingPointSort =
+    "Sorting on floating-point values nested in arrays, structs, or maps"
+
   override def getIncompatibleReasons(): Seq[String] = Seq(
-    "When `" + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key + "=true`, sorting on" +
-      " floating-point types is not 100% compatible with Spark")
+    s"$nestedFloatingPointSort is not 100% compatible with Spark when " +
+      s"`${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true`")
 
-  override def getSupportLevel(expr: SortOrder): SupportLevel = {
-    // https://github.com/apache/datafusion-comet/issues/2626
-    SupportLevel
-      .strictFloatingPointReason(expr.child.dataType, "Sorting on floating-point")
-      .map(reason => Incompatible(Some(reason)))
-      .getOrElse(Compatible())
+  override def getSupportLevel(expr: SortOrder): SupportLevel = expr.child.dataType match {
+    // Scalar FLOAT/DOUBLE comparison keys are normalized natively (NaN payloads folded together,
+    // signed zeros tied) for Sort, TopK, Window, WindowGroupLimit, and range partitioning, which
+    // matches Spark's SQLOrderingUtil. Only the comparison key is normalized; returned values keep
+    // their original NaN representation and zero sign. So these are compatible even in strict mode.
+    case _: FloatType | _: DoubleType => Compatible()
+    // Floating-point values nested in arrays, structs, or maps are still compared with Arrow's raw
+    // total ordering, under which -0.0 sorts below 0.0 and a sign-bit NaN sorts below -Infinity.
+    // https://github.com/apache/datafusion-comet/issues/5507
+    case dt =>
+      SupportLevel
+        .strictFloatingPointReason(dt, nestedFloatingPointSort)
+        .map(reason => Incompatible(Some(reason)))
+        .getOrElse(Compatible())
   }
 
   override def convert(

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala

@@ -525,8 +525,6 @@ object CometShuffleExchangeExec
       case SinglePartition =>
       // we already checked that the input types are supported
       case RangePartitioning(orderings, _) =>
-        val strictFloatingPoint = CometConf.COMET_EXEC_STRICT_FLOATING_POINT.get(conf)
-
         /**
          * Determine which data types are supported as partition columns in native shuffle.
          *
@@ -537,8 +535,10 @@ object CometShuffleExchangeExec
         def supportedRangePartitioningDataType(dt: DataType): Boolean = dt match {
           // Collated strings require collation-aware ordering; Comet only compares raw bytes.
           case st: StringType if isStringCollationType(st) => false
-          case _: FloatType | _: DoubleType =>
-            !strictFloatingPoint
+          // The native range partitioner normalizes its comparison keys and its sampled boundary
+          // rows the same way the native sort does, so scalar floats match Spark's ordering even
+          // under spark.comet.exec.strictFloatingPoint=true.
+          case _: FloatType | _: DoubleType => true
           case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType |
               _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType |
               _: DecimalType | _: DateType =>
@@ -562,15 +562,7 @@ object CometShuffleExchangeExec
         }
         for (dt <- orderings.map(_.dataType).distinct) {
           if (!supportedRangePartitioningDataType(dt)) {
-            val reason = dt match {
-              case _: FloatType | _: DoubleType if strictFloatingPoint =>
-                s"Range partitioning on $dt is not 100% compatible with Spark, and Comet is " +
-                  s"running with ${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true. " +
-                  s"${CometConf.COMPAT_GUIDE}"
-              case _ =>
-                s"unsupported range partitioning data type for native shuffle: $dt"
-            }
-            reasons += reason
+            reasons += s"unsupported range partitioning data type for native shuffle: $dt"
           }
         }
       case RoundRobinPartitioning(_) =>

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -25,8 +25,8 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row}
 import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, FromUnixTime, InSet, Literal, StructsToJson, TruncDate, TruncTimestamp}
 import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn, SimplifyExtractValueOps}
-import org.apache.spark.sql.comet.CometProjectExec
-import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
+import org.apache.spark.sql.comet.{CometProjectExec, CometSortExec, CometTakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.{LocalTableScanExec, ProjectExec, SparkPlan}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
@@ -60,16 +60,142 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       schema,
       1000,
       DataGenOptions(generateNegativeZero = true))
-    df.createOrReplaceTempView("tbl")
+
+    withTempDir { dir =>
+      // `-0.0` and `+0.0` are peers under Spark's comparison, so `ORDER BY c0, c1` alone leaves
+      // genuine ties whose relative order neither engine promises. Materialize a unique `id` to
+      // Parquet and sort on it last, making the ordering total so the row-by-row comparison below
+      // tests the sort keys rather than an unspecified tie order.
+      val path = new Path(dir.toString, "tbl").toString
+      df.withColumn("id", monotonically_increasing_id()).write.parquet(path)
+      spark.read.parquet(path).createOrReplaceTempView("tbl")
+
+      // Scalar floating-point sort keys are normalized natively, so strict mode admits them even
+      // with allowIncompatible off. Nested floating-point keys are not, and the two tests below
+      // still assert the strict-mode fallback.
+      withSQLConf(
+        CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false",
+        CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") {
+        checkSparkAnswerAndOperator(
+          sql("select * from tbl order by 1, 2, 3"),
+          Seq(classOf[CometSortExec]))
+      }
+    }
+  }
+  // https://github.com/apache/datafusion-comet/issues/5506
+  //
+  // Scalar floating-point sort keys are admitted under strict floating point because their
+  // comparison keys are normalized natively. Signed zeros survive a Parquet round trip, so these
+  // cases are materialized to Parquet, which also keeps constant folding from deleting the sort.
+  // NaN payloads do NOT survive that round trip (every payload is canonicalized on write), so
+  // payload fidelity is covered separately below over a local relation.
+  //
+  // Every query ends with a unique `id` so the ordering is total: -0.0 and +0.0 are peers under
+  // Spark's comparison, and neither engine promises an order within a tie.
+  private val strictFpSortRows = Seq(
+    (0, Some(-0.0f), Some(-0.0d), 1),
+    (1, Some(0.0f), Some(0.0d), 0),
+    (2, Some(0.0f), Some(0.0d), 1),
+    (3, Some(-0.0f), Some(-0.0d), 0),
+    (4, Some(1.0f), Some(1.0d), 0),
+    (5, Some(-1.0f), Some(-1.0d), 1),
+    (6, None, None, 0),
+    (7, None, None, 1),
+    (8, Some(Float.NaN), Some(Double.NaN), 0),
+    (9, Some(Float.PositiveInfinity), Some(Double.PositiveInfinity), 1),
+    (10, Some(Float.NegativeInfinity), Some(Double.NegativeInfinity), 0))
+
+  for {
+    col <- Seq("f", "d")
+    direction <- Seq("ASC", "DESC")
+    nullOrder <- Seq("NULLS FIRST", "NULLS LAST")
+    compound <- Seq(false, true)
+  } {
+    val label = s"$col $direction $nullOrder" + (if (compound) " with compound key" else "")
+    test(s"strict floating point: scalar sort on $label") {
+      withTempDir { dir =>
+        val path = new Path(dir.toString, "strict_fp_sort").toString
+        strictFpSortRows.toDF("id", "f", "d", "s").write.parquet(path)
+        spark.read.parquet(path).createOrReplaceTempView("strict_fp_sort")
+
+        val secondary = if (compound) ", s DESC" else ""
+        val query =
+          s"SELECT id, f, d, s FROM strict_fp_sort ORDER BY $col $direction $nullOrder$secondary, id"
+
+        withSQLConf(
+          CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false",
+          CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") {
+          checkSparkAnswerAndOperator(sql(query), Seq(classOf[CometSortExec]))
+        }
+      }
+    }
+  }
+
+  test("strict floating point: scalar TopK sort key") {
+    // TakeOrderedAndProject reaches the native path through the same CometSortOrder gate as a
+    // plain sort, but lands on a different operator (Sort with a fetch). Widening admission
+    // therefore opens this path too, so assert it rather than inferring it.
+    withTempDir { dir =>
+      val path = new Path(dir.toString, "strict_fp_topk").toString
+      strictFpSortRows.toDF("id", "f", "d", "s").write.parquet(path)
+      spark.read.parquet(path).createOrReplaceTempView("strict_fp_topk")
+
+      withSQLConf(
+        CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false",
+        CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") {
+        for (col <- Seq("f", "d")) {
+          checkSparkAnswerAndOperator(
+            sql(s"SELECT id, f, d FROM strict_fp_topk ORDER BY $col ASC NULLS LAST, id LIMIT 4"),
+            Seq(classOf[CometTakeOrderedAndProjectExec]))
+        }
+      }
+    }
+  }
+
+  test("strict floating point: scalar sort keeps NaN payloads and zero signs unchanged") {
+    // A local relation preserves the raw bits that a Parquet round trip would canonicalize, so
+    // this is where "only the comparison key is normalized" can actually be asserted.
+    val negNan = java.lang.Double.longBitsToDouble(0xfff8000000000002L)
+    val posNan = java.lang.Double.longBitsToDouble(0x7ff8000000000002L)
+    val negNanF = java.lang.Float.intBitsToFloat(0xffc00002)
+    val posNanF = java.lang.Float.intBitsToFloat(0x7fc00002)
+    val rows = Seq(
+      (0, negNanF, negNan),
+      (1, posNanF, posNan),
+      (2, -0.0f, -0.0d),
+      (3, 0.0f, 0.0d),
+      (4, 1.0f, 1.0d))
+    val expected = rows.map { case (id, f, d) =>
+      id -> (java.lang.Float.floatToRawIntBits(f), java.lang.Double.doubleToRawLongBits(d))
+    }.toMap
+
+    rows.toDF("id", "f", "d").createOrReplaceTempView("strict_fp_bits")
 
     withSQLConf(
       CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false",
       CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") {
-      checkSparkAnswerAndFallbackReasons(
-        "select * from tbl order by 1, 2",
-        Set(
-          "unsupported range partitioning sort order",
-          "Sorting on floating-point is not 100% compatible with Spark"))
+      for (col <- Seq("f", "d")) {
+        val query = s"SELECT id, f, d FROM strict_fp_bits ORDER BY $col, id"
+        // LocalTableScanExec has no native counterpart enabled by default; it is the source of
+        // the unmodified bits, not part of what this test asserts.
+        checkSparkAnswerAndOperator(
+          sql(query),
+          Seq(classOf[CometSortExec]),
+          classOf[LocalTableScanExec])
+
+        val actual = sql(query).collect().toSeq
+        actual.foreach { row =>
+          val bits = (
+            java.lang.Float.floatToRawIntBits(row.getFloat(1)),
+            java.lang.Double.doubleToRawLongBits(row.getDouble(2)))
+          assert(
+            bits == expected(row.getInt(0)),
+            s"row ${row.getInt(0)} had its floating-point bits rewritten by the sort")
+        }
+        // The two NaN payloads are peers, so they must be adjacent and ordered by the tiebreaker.
+        val nanIds = actual.map(_.getInt(0)).filter(id => id == 0 || id == 1)
+        assert(nanIds == Seq(0, 1), s"NaN peers were not tied and ordered by id: $nanIds")
+      }
     }
   }
 

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

@@ -792,6 +792,28 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
    * Checks that `df` produces the same answer as Spark does, and has the `expectedNum` Comet
    * exchange operators.
    */
+  test("range partitioning on floating-point uses columnar shuffle under strictFloatingPoint") {
+    // The columnar path partitions on the JVM with Spark's own RangePartitioner, but it still
+    // probes whether Comet can serialize the sort order. Before #5506 that probe reported
+    // Incompatible for scalar float and double under strict mode, so this exchange fell back to
+    // Spark's shuffle for no compatibility reason. Consulting a native-serde gate on a path that
+    // never goes native is tracked separately in #5971; this test pins the floating-point half of
+    // it, which #5506 fixed.
+    withSQLConf(
+      CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true",
+      CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false") {
+      withParquetTable(
+        (0 until 20).map(i => (if (i % 3 == 0) -0.0d else i.toDouble, i)),
+        "range_fp_tbl") {
+        Seq("FLOAT", "DOUBLE").foreach { sqlType =>
+          val df = sql(s"SELECT CAST(_1 AS $sqlType) AS c, _2 FROM range_fp_tbl")
+            .repartitionByRange(4, col("c"))
+          checkShuffleAnswer(df, 1)
+        }
+      }
+    }
+  }
+
   private def checkShuffleAnswer(df: DataFrame, expectedNum: Int): Unit = {
     checkCometExchange(df, expectedNum, false)
     checkSparkAnswer(df)