From 4f43d7660cf40a1912f5c07369c947d39c9cc896 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 25 Aug 2026 19:37:41 -0700 Subject: [PATCH 1/3] fix: normalize scalar float sort and window rank keys --- .../latest/compatibility/floating-point.md | 23 +- native/core/src/execution/planner.rs | 272 +++++++++++++++++- .../src/math_funcs/internal/normalize_nan.rs | 33 ++- .../windows/window_group_limit_rank.sql | 15 +- .../comet/exec/CometWindowExecSuite.scala | 87 +++++- 5 files changed, 391 insertions(+), 39 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index a068e12a500..ec03e59b037 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -28,15 +28,18 @@ to Spark in some cases, especially when the data contains both positive and nega case that is not of concern for many users. If it is a concern, setting `spark.comet.exec.strictFloatingPoint=true` will make relevant operations fall back to Spark. -## Ordering: signed zero (`-0.0` vs `+0.0`) +## Ordering: NaN and signed zero (`-0.0` vs `+0.0`) Spark's `ORDER BY`, `RANK`, `DENSE_RANK`, and window frame comparisons route through -`SQLOrderingUtil.compareDoubles` / `compareFloats`, which explicitly define `-0.0 == 0.0`. Comet's -native sort and `WindowGroupLimitExec` use the `arrow-row` row-format encoder for `ORDER BY` keys, -which applies Rust's total-ordering transform to the raw IEEE-754 bits. Under that encoding `-0.0` -sorts strictly less than `+0.0`, so a partition that mixes the two zeros can produce a rank -distribution that differs from Spark. For example, `RANK() OVER (ORDER BY v ASC)` over -`[-0.0, 0.0, 1.0]` filtered to `rk <= 1` returns two rows in Spark (both zeros tied at rank 1) but -one row in Comet (`-0.0` at rank 1, `+0.0` at rank 2). If your workload materially mixes `-0.0` -and `+0.0` in a ranked column, prefer Spark for that stage or normalize the column to `+0.0` -upstream. +`SQLOrderingUtil.compareDoubles` / `compareFloats`, which equate all NaN representations and +define `-0.0 == 0.0`. NaN sorts above every non-NaN value. + +For scalar `FLOAT` and `DOUBLE` keys, Comet normalizes NaNs and signed zeros before native +sorting, window peer comparisons, and `WindowGroupLimitExec` rank comparisons. Native range +partitioning normalizes its keys and sampled boundaries in the same way. Only comparison keys +are normalized; returned values retain their original NaN representations and zero signs. + +Floating-point values nested in arrays or structs still use Arrow's raw total ordering and can +produce different ordering or rank results from Spark. The existing +`spark.comet.exec.strictFloatingPoint=true` fallback policy is unchanged, including its +conservative fallback for scalar floating-point sort keys. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 7fb1c89b962..b2121e68d13 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -933,7 +933,17 @@ impl PhysicalPlanner { ) -> Result { match spark_expr.expr_struct.as_ref().unwrap() { ExprStruct::SortOrder(expr) => { - let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?; + let child = + self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?; + let data_type = child.data_type(input_schema.as_ref())?; + // Spark treats every NaN as a peer and equates signed zeros. Normalize only + // comparison keys, preserving the original values in the output. Sort, Window, + // and WindowGroupLimit must agree so peers remain contiguous for compound keys. + let child = if matches!(data_type, DataType::Float32 | DataType::Float64) { + Arc::new(NormalizeNaNAndZero::new(data_type, child)) as Arc + } else { + child + }; let descending = expr.direction == 1; let nulls_first = expr.null_ordering == 0; @@ -3425,10 +3435,14 @@ impl PhysicalPlanner { } } - // Convert the collection of ScalarValues to collection of Arrow Arrays + // Normalize boundary arrays just like the incoming sort keys, so equal NaNs + // and signed zeros are assigned to the same range partition. let arrays: Vec = scalar_values .iter() - .map(|scalar_vec| ScalarValue::iter_to_array(scalar_vec.iter().cloned())) + .map(|scalar_vec| { + ScalarValue::iter_to_array(scalar_vec.iter().cloned()) + .map(|array| NormalizeNaNAndZero::normalize_array(&array)) + }) .collect::, _>>()?; // Create a RowConverter and use to create OwnedRows from the Arrays @@ -4747,19 +4761,25 @@ mod tests { use std::{sync::Arc, task::Poll}; use arrow::array::{ - Array, DictionaryArray, Int32Array, Int8Array, ListArray, RecordBatch, StringArray, + Array, ArrayRef, DictionaryArray, Float32Array, Float64Array, Int32Array, Int8Array, + ListArray, RecordBatch, StringArray, }; use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema}; use datafusion::catalog::memory::DataSourceExec; + use datafusion::common::ScalarValue; use datafusion::config::TableParquetOptions; use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfigBuilder, FileSource, ParquetSource, }; use datafusion::error::DataFusionError; use datafusion::logical_expr::ScalarUDF; + use datafusion::physical_expr::LexOrdering; + use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::ExecutionPlan; + use datafusion::prelude::SessionConfig; use datafusion::{assert_batches_eq, physical_plan::common::collect, prelude::SessionContext}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use tempfile::TempDir; @@ -4767,9 +4787,10 @@ mod tests { use crate::execution::{operators::InputBatch, planner::PhysicalPlanner}; - use crate::execution::operators::ExecutionError; + use crate::execution::operators::{ExecutionError, PartitionedRankLimitExec, WindowFnKind}; use crate::execution::planner::literal_to_array_ref; use crate::execution::planner::parse_file_scan_tasks_from_common; + use crate::execution::shuffle::CometPartitioning; use crate::parquet::parquet_support::SparkParquetOptions; use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; use datafusion_comet_proto::spark_expression::expr::ExprStruct; @@ -4780,6 +4801,9 @@ mod tests { spark_expression::{self, literal}, spark_operator, spark_operator::{operator::OpStruct, Operator}, + spark_partitioning::{ + partitioning::PartitioningStruct, BoundaryRow, Partitioning, RangePartition, + }, }; use datafusion_comet_spark_expr::EvalMode; @@ -4943,6 +4967,244 @@ mod tests { ); } + fn create_sort_order(index: i32, type_id: i32, descending: bool, nulls_first: bool) -> Expr { + Expr { + expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder { + child: Some(Box::new(Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index, + datatype: Some(spark_expression::DataType { + type_id, + type_info: None, + }), + })), + ..Default::default() + })), + direction: i32::from(descending), + null_ordering: i32::from(!nulls_first), + }))), + ..Default::default() + } + } + + fn floating_sort_batches() -> Vec<(i32, RecordBatch)> { + // Include distinct quiet/signaling NaNs of both signs. Construct them directly: + // round-tripping through Parquet or a string literal can canonicalize the payload. + let floats: ArrayRef = Arc::new(Float32Array::from(vec![ + Some(f32::NAN), + Some(f32::from_bits(0xffc0_0042)), + Some(f32::from_bits(0x7fc0_1234)), + Some(0.0), + Some(-0.0), + Some(f32::NEG_INFINITY), + Some(1.0), + Some(f32::INFINITY), + Some(f32::from_bits(0xff80_0001)), + Some(f32::from_bits(0x7f80_0001)), + Some(0.0), + Some(-0.0), + None, + Some(-1.0), + ])); + let doubles: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(f64::NAN), + Some(f64::from_bits(0xfff8_0000_0000_0042)), + Some(f64::from_bits(0x7ff8_0000_0000_1234)), + Some(0.0), + Some(-0.0), + Some(f64::NEG_INFINITY), + Some(1.0), + Some(f64::INFINITY), + Some(f64::from_bits(0xfff0_0000_0000_0001)), + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + Some(0.0), + Some(-0.0), + None, + Some(-1.0), + ])); + [(5, floats), (6, doubles)] + .into_iter() + .map(|(type_id, values)| { + let schema = Arc::new(Schema::new(vec![ + Field::new("ord", values.data_type().clone(), true), + Field::new("suffix", DataType::Int32, false), + Field::new("id", DataType::Int32, false), + ])); + let ids = Int32Array::from_iter_values(0..values.len() as i32); + let suffix = Int32Array::from(vec![1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1]); + let batch = + RecordBatch::try_new(schema, vec![values, Arc::new(suffix), Arc::new(ids)]) + .unwrap(); + (type_id, batch) + }) + .collect() + } + + #[tokio::test] + async fn floating_sort_keys_preserve_window_group_limit_peers() { + let planner = PhysicalPlanner::default(); + let context = SessionContext::new_with_config(SessionConfig::new().with_batch_size(3)); + for (type_id, batch) in floating_sort_batches() { + for (descending, kind, fetch, expected) in [ + (true, WindowFnKind::Rank, 3, vec![0, 1, 2, 8, 9]), + (true, WindowFnKind::DenseRank, 2, vec![0, 1, 2, 8, 9]), + (true, WindowFnKind::RowNumber, 2, vec![8, 9]), + (false, WindowFnKind::Rank, 3, vec![5, 10, 11, 13]), + (false, WindowFnKind::DenseRank, 3, vec![5, 10, 11, 13]), + (false, WindowFnKind::RowNumber, 4, vec![5, 10, 11, 13]), + ] { + let order_keys = [ + create_sort_order(0, type_id, descending, false), + create_sort_order(1, 3, false, false), + ] + .iter() + .map(|expr| planner.create_sort_expr(expr, batch.schema()).unwrap()) + .collect::>(); + let input = MemorySourceConfig::try_new_exec( + &[vec![ + batch.slice(0, 4), + batch.slice(4, 4), + batch.slice(8, 6), + ]], + batch.schema(), + None, + ) + .unwrap(); + let sorted = Arc::new(SortExec::new( + LexOrdering::new(order_keys.clone()).unwrap(), + input, + )); + let limit = + PartitionedRankLimitExec::try_new(sorted, vec![], order_keys, fetch, kind) + .unwrap(); + let output = collect(limit.execute(0, context.task_ctx()).unwrap()) + .await + .unwrap(); + let mut ids = vec![]; + for out in &output { + let out_ids = out.column(2).as_any().downcast_ref::().unwrap(); + for row in 0..out.num_rows() { + let id = out_ids.value(row); + ids.push(id); + // ScalarValue float equality compares raw bits, including NaN + // payloads and the zero sign. Only comparison keys may change. + assert_eq!( + ScalarValue::try_from_array(out.column(0), row).unwrap(), + ScalarValue::try_from_array(batch.column(0), id as usize).unwrap(), + ); + } + } + ids.sort_unstable(); + assert_eq!( + ids, expected, + "type={type_id}, descending={descending}, {kind:?}" + ); + // The zero peer group and the second NaN peer group straddle size-3 + // sort output batches. Tie state must survive those boundaries. + assert!(output.iter().all(|b| b.num_rows() <= 3)); + if expected.len() > 3 { + assert!(output.len() > 1); + } + } + } + } + + #[test] + fn floating_range_boundaries_match_sort_keys() { + let planner = PhysicalPlanner::default(); + for (type_id, batch) in floating_sort_batches() { + for descending in [false, true] { + for nulls_first in [false, true] { + // Use -0 and a negative payload NaN, while the input contains both + // zero signs and several NaN representations equivalent in Spark. + let boundary_indices = if descending { [1, 4] } else { [4, 1] }; + let boundary_rows = boundary_indices + .iter() + .map(|&index| { + let value = match ScalarValue::try_from_array(batch.column(0), index) + .unwrap() + { + ScalarValue::Float32(Some(value)) => { + literal::Value::FloatVal(value) + } + ScalarValue::Float64(Some(value)) => { + literal::Value::DoubleVal(value) + } + _ => unreachable!(), + }; + BoundaryRow { + partition_bounds: vec![Expr { + expr_struct: Some(Literal(spark_expression::Literal { + value: Some(value), + datatype: Some(spark_expression::DataType { + type_id, + type_info: None, + }), + is_null: false, + })), + ..Default::default() + }], + } + }) + .collect(); + let partitioning = Partitioning { + partitioning_struct: Some(PartitioningStruct::RangePartition( + RangePartition { + sort_orders: vec![create_sort_order( + 0, + type_id, + descending, + nulls_first, + )], + num_partitions: 3, + boundary_rows, + }, + )), + }; + let CometPartitioning::RangePartitioning(ordering, _, converter, boundaries) = + planner + .create_partitioning(&partitioning, batch.schema()) + .unwrap() + else { + panic!("expected range partitioning"); + }; + let columns = ordering + .iter() + .map(|expr| { + expr.expr + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap() + }) + .collect::>(); + let rows = converter.convert_columns(&columns).unwrap(); + for index in [3, 4, 10, 11] { + assert_eq!(rows.row(index), boundaries[usize::from(descending)].row()); + } + for index in [0, 1, 2, 8, 9] { + assert_eq!(rows.row(index), boundaries[usize::from(!descending)].row()); + } + // Match the range shuffle writer's binary-search routing. + let partitions = rows + .iter() + .map(|row| boundaries.partition_point(|bound| bound.row() <= row)) + .collect::>(); + let null_partition = if nulls_first { 0 } else { 2 }; + let expected = if descending { + vec![1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, null_partition, 2] + } else { + vec![2, 2, 2, 1, 1, 0, 1, 1, 2, 2, 1, 1, null_partition, 0] + }; + assert_eq!( + partitions, expected, + "type={type_id}, descending={descending}, nulls_first={nulls_first}", + ); + } + } + } + } + #[test] fn test_unpack_dictionary_primitive() { let op_scan = Operator { diff --git a/native/spark-expr/src/math_funcs/internal/normalize_nan.rs b/native/spark-expr/src/math_funcs/internal/normalize_nan.rs index bb22e575a4c..85afef5666e 100644 --- a/native/spark-expr/src/math_funcs/internal/normalize_nan.rs +++ b/native/spark-expr/src/math_funcs/internal/normalize_nan.rs @@ -18,7 +18,7 @@ use arrow::compute::unary; use arrow::datatypes::{DataType, Schema}; use arrow::{ - array::{as_primitive_array, Float32Array, Float64Array}, + array::{as_primitive_array, ArrayRef, Float32Array, Float64Array}, datatypes::{Float32Type, Float64Type}, record_batch::RecordBatch, }; @@ -53,6 +53,24 @@ impl NormalizeNaNAndZero { pub fn new(data_type: DataType, child: Arc) -> Self { Self { data_type, child } } + + /// Normalize scalar floating-point comparison keys, leaving other types unchanged. + /// Sorting and range-partition boundaries must use the same representation. + pub fn normalize_array(array: &ArrayRef) -> ArrayRef { + match array.data_type() { + DataType::Float32 => { + let input = as_primitive_array::(array); + let result: Float32Array = unary(input, normalize_float); + Arc::new(result) + } + DataType::Float64 => { + let input = as_primitive_array::(array); + let result: Float64Array = unary(input, normalize_float); + Arc::new(result) + } + _ => Arc::clone(array), + } + } } impl PhysicalExpr for NormalizeNaNAndZero { @@ -73,17 +91,8 @@ impl PhysicalExpr for NormalizeNaNAndZero { let array = cv.into_array(batch.num_rows())?; match &self.data_type { - DataType::Float32 => { - let input = as_primitive_array::(&array); - // Use unary which operates directly on values buffer without intermediate allocation - let result: Float32Array = unary(input, normalize_float); - Ok(ColumnarValue::Array(Arc::new(result))) - } - DataType::Float64 => { - let input = as_primitive_array::(&array); - // Use unary which operates directly on values buffer without intermediate allocation - let result: Float64Array = unary(input, normalize_float); - Ok(ColumnarValue::Array(Arc::new(result))) + DataType::Float32 | DataType::Float64 => { + Ok(ColumnarValue::Array(Self::normalize_array(&array))) } dt => panic!("Unexpected data type {dt:?}"), } diff --git a/spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql index cf214331dd1..6e3f3b76bb1 100644 --- a/spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_rank.sql @@ -178,12 +178,9 @@ SELECT a, b, score FROM ( -- ================================================================================ -- FP edge cases in the ORDER BY column (NaN / +Inf / -Inf / 0 / NULL). NaN sorts --- greater than any finite value in Spark, and Arrow's row-format total-ordering --- encoding matches -- so NaN gets rank 1 under DESC. -0.0 vs 0.0 is a known --- divergence: Spark's SQLOrderingUtil ties them, Arrow's bitwise total_cmp splits --- them. The primary test below keeps the cutoff (rk <= 3) above the 0-values so --- it runs as-is; a second test below with `query ignore(...)` pins the divergent --- shape directly. See docs/source/user-guide/latest/compatibility/floating-point.md. +-- greater than any finite value in Spark, so NaN gets rank 1 under DESC. +-- Native scalar sort and rank keys normalize NaNs and signed zeros before row +-- encoding. See docs/source/user-guide/latest/compatibility/floating-point.md. -- ================================================================================ statement @@ -209,9 +206,7 @@ SELECT part, v FROM ( ) t WHERE rk <= 3 ORDER BY rk, v DESC NULLS LAST -- -0.0 vs +0.0 in the ORDER BY column: Spark ties them at rank 1 (both rows --- survive `rk <= 1`); Comet's row encoder splits them (only -0.0 survives ASC, --- only +0.0 survives DESC). Kept as `ignore(...)` so the suite passes today and --- lands green once the divergence is closed. +-- survive `rk <= 1`). Native comparison-key normalization must preserve this tie. statement CREATE TABLE test_rank_fp_zero(part string, v double) USING parquet @@ -221,7 +216,7 @@ INSERT INTO test_rank_fp_zero VALUES ('p', -0.0), ('p', 1.0) -query ignore(signed-zero ORDER BY: Spark ties -0.0 with +0.0, Arrow row encoder splits them) +query SELECT part, v FROM ( SELECT part, v, RANK() OVER (PARTITION BY part ORDER BY v ASC) AS rk diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 7cb6816298d..a0257cfc9d8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -27,7 +27,7 @@ import org.scalatest.Tag import org.apache.hadoop.fs.Path import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Divide, Expression, MakeDecimal, WindowExpression} -import org.apache.spark.sql.comet.CometWindowExec +import org.apache.spark.sql.comet.{CometSortExec, CometWindowExec, CometWindowGroupLimitExec} import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.window.{WindowExec => SparkWindowExec} import org.apache.spark.sql.expressions.Window @@ -36,7 +36,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.DecimalType import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} class CometWindowExecSuite extends CometTestBase { @@ -117,6 +117,89 @@ class CometWindowExecSuite extends CometTestBase { digits.toString() } + for (orderColumn <- Seq("f", "d")) { + test(s"window group limit: $orderColumn NaN and signed zero peers at the cutoff") { + assume(isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") + + val positiveFloatNaN = java.lang.Float.intBitsToFloat(0x7fc00001) + val negativeFloatNaN = java.lang.Float.intBitsToFloat(0xffc00002) + val positiveDoubleNaN = java.lang.Double.longBitsToDouble(0x7ff8000000000001L) + val negativeDoubleNaN = java.lang.Double.longBitsToDouble(0xfff8000000000002L) + val data = Seq( + (0, 1, Some(positiveFloatNaN), Some(positiveDoubleNaN), 1), + (0, 2, Some(negativeFloatNaN), Some(negativeDoubleNaN), 0), + (0, 3, Some(positiveFloatNaN), Some(positiveDoubleNaN), 0), + (0, 4, Some(negativeFloatNaN), Some(negativeDoubleNaN), 1), + (0, 5, Some(1.0f), Some(1.0d), 0), + (0, 6, None, None, 0), + (1, 7, Some(0.0f), Some(0.0d), 1), + (1, 8, Some(-0.0f), Some(-0.0d), 0), + (1, 9, Some(0.0f), Some(0.0d), 0), + (1, 10, Some(-0.0f), Some(-0.0d), 1), + (1, 11, Some(-1.0f), Some(-1.0d), 0), + (1, 12, None, None, 0)) + val expectedBits = data.map { case (_, id, f, d, _) => + ( + id, + (f.map(java.lang.Float.floatToRawIntBits), d.map(java.lang.Double.doubleToRawLongBits))) + }.toMap + + def assertRawBits(rows: Seq[Row]): Unit = { + rows.foreach { row => + val floatBits = + if (row.isNullAt(2)) None + else Some(java.lang.Float.floatToRawIntBits(row.getFloat(2))) + val doubleBits = + if (row.isNullAt(3)) None + else Some(java.lang.Double.doubleToRawLongBits(row.getDouble(3))) + assert((floatBits, doubleBits) == expectedBits(row.getInt(1))) + } + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "false", + CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.key -> "true") { + withTempView("floating_window_peers") { + // LocalTableScan preserves NaN payloads, unlike a Parquet round trip. Check the input + // bits so that canonicalization before the native sort cannot hide the regression. + data + .toDF("p", "id", "f", "d", "secondary") + .createOrReplaceTempView("floating_window_peers") + assertRawBits(sql("SELECT * FROM floating_window_peers").collect().toSeq) + + for { + rankFunction <- Seq("RANK", "DENSE_RANK") + (ordering, cutoff) <- Seq("DESC NULLS LAST" -> 1, "ASC NULLS FIRST" -> 3) + secondaryKey <- Seq(false, true) + } { + // With a secondary key, sorting raw NaN/zero representations separates rows that + // must be peers. Normalizing only the limit operator's equality keys is insufficient. + val orderBy = s"$orderColumn $ordering" + (if (secondaryKey) ", secondary" else "") + val query = sql(s""" + |SELECT p, id, f, d, rnk FROM ( + | SELECT *, $rankFunction() OVER (PARTITION BY p ORDER BY $orderBy) AS rnk + | FROM floating_window_peers + |) WHERE rnk <= $cutoff + |""".stripMargin) + checkSparkAnswerAndOperator( + query, + Seq(classOf[CometSortExec], classOf[CometWindowGroupLimitExec])) + + val actual = query.collect().toSeq + val peerIds = if (secondaryKey) Seq(2, 3, 8, 9) else Seq(1, 2, 3, 4, 7, 8, 9, 10) + val preceding = if (cutoff == 3) Seq(6 -> 1, 12 -> 1, 5 -> 2, 11 -> 2) else Seq.empty + val expected = peerIds.map(_ -> cutoff) ++ preceding + assert(actual.map(row => row.getInt(1) -> row.getInt(4)).sorted == expected.sorted) + // Canonical comparison keys must not alter the selected rows' floating values. + assertRawBits(actual) + } + } + } + } + } + test("lead/lag should return the default value if the offset row does not exist") { withSQLConf( CometConf.COMET_ENABLED.key -> "true", From 96eafdfe14384976cb18b7364c4cf84d7d1fb723 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 25 Aug 2026 21:06:55 -0700 Subject: [PATCH 2/3] fix: preserve floating window partition ordering --- native/core/src/execution/planner.rs | 102 +++++++++++++++--- .../comet/exec/CometWindowExecSuite.scala | 61 +++++++++++ 2 files changed, 150 insertions(+), 13 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index b2121e68d13..536219d1253 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -925,6 +925,26 @@ impl PhysicalPlanner { } } + /// Normalize scalar floating-point comparison keys without changing output values. + /// Sort, Window, and WindowGroupLimit must use identical expressions so DataFusion + /// can recognize the ordering of window partition keys. + fn create_normalized_key_expr( + &self, + spark_expr: &Expr, + input_schema: SchemaRef, + ) -> Result, ExecutionError> { + let child = self.create_expr(spark_expr, Arc::clone(&input_schema))?; + let data_type = child.data_type(input_schema.as_ref())?; + // Spark may already have normalized a partition or join key. + if matches!(data_type, DataType::Float32 | DataType::Float64) + && child.downcast_ref::().is_none() + { + Ok(Arc::new(NormalizeNaNAndZero::new(data_type, child))) + } else { + Ok(child) + } + } + /// Create a DataFusion physical sort expression from Spark physical expression fn create_sort_expr<'a>( &'a self, @@ -934,16 +954,7 @@ impl PhysicalPlanner { match spark_expr.expr_struct.as_ref().unwrap() { ExprStruct::SortOrder(expr) => { let child = - self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?; - let data_type = child.data_type(input_schema.as_ref())?; - // Spark treats every NaN as a peer and equates signed zeros. Normalize only - // comparison keys, preserving the original values in the output. Sort, Window, - // and WindowGroupLimit must agree so peers remain contiguous for compound keys. - let child = if matches!(data_type, DataType::Float32 | DataType::Float64) { - Arc::new(NormalizeNaNAndZero::new(data_type, child)) as Arc - } else { - child - }; + self.create_normalized_key_expr(expr.child.as_ref().unwrap(), input_schema)?; let descending = expr.direction == 1; let nulls_first = expr.null_ordering == 0; @@ -2290,7 +2301,7 @@ impl PhysicalPlanner { let partition_exprs: Result>, ExecutionError> = wnd .partition_by_list .iter() - .map(|expr| self.create_expr(expr, Arc::clone(&input_schema))) + .map(|expr| self.create_normalized_key_expr(expr, Arc::clone(&input_schema))) .collect(); let sort_exprs = &sort_exprs?; @@ -2438,7 +2449,8 @@ impl PhysicalPlanner { let mut partition_keys: Vec = Vec::with_capacity(partition_prefix_len); for expr in &wgl.partition_by_list { - let phys = self.create_expr(expr, Arc::clone(&input_schema))?; + let phys = + self.create_normalized_key_expr(expr, Arc::clone(&input_schema))?; partition_keys.push(PhysicalSortExpr { expr: phys, options: SortOptions::default(), @@ -4776,8 +4788,9 @@ mod tests { }; use datafusion::error::DataFusionError; use datafusion::logical_expr::ScalarUDF; - use datafusion::physical_expr::LexOrdering; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion::physical_plan::sorts::sort::SortExec; + use datafusion::physical_plan::windows::get_ordered_partition_by_indices; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion::{assert_batches_eq, physical_plan::common::collect, prelude::SessionContext}; @@ -5040,6 +5053,69 @@ mod tests { .collect() } + #[test] + fn floating_window_partition_keys_preserve_ordering() { + let planner = PhysicalPlanner::default(); + for (type_id, batch) in floating_sort_batches() { + for already_normalized in [false, true] { + let mut partition_sort = create_sort_order(0, type_id, false, true); + let Some(SortOrder(sort_order)) = partition_sort.expr_struct.as_mut() else { + unreachable!(); + }; + if already_normalized { + sort_order.child = Some(Box::new(Expr { + expr_struct: Some(NormalizeNanAndZero(Box::new( + spark_expression::NormalizeNaNAndZero { + child: sort_order.child.take(), + datatype: Some(spark_expression::DataType { + type_id, + type_info: None, + }), + }, + ))), + ..Default::default() + })); + } + // A key can arrive as Spark's normalization expression or as a + // materialized floating column. Both must match the sort prefix. + let partition_expr = planner + .create_normalized_key_expr(sort_order.child.as_ref().unwrap(), batch.schema()) + .unwrap(); + let sort_expr = planner + .create_sort_expr(&partition_sort, batch.schema()) + .unwrap(); + let input = + MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None) + .unwrap(); + let sorted: Arc = Arc::new(SortExec::new( + LexOrdering::new(vec![sort_expr.clone()]).unwrap(), + input, + )); + let limited: Arc = Arc::new( + PartitionedRankLimitExec::try_new( + Arc::clone(&sorted), + vec![PhysicalSortExpr { + expr: Arc::clone(&partition_expr), + options: sort_expr.options, + }], + vec![], + 1, + WindowFnKind::RowNumber, + ) + .unwrap(), + ); + for input in [sorted, limited] { + assert_eq!( + get_ordered_partition_by_indices(&[Arc::clone(&partition_expr)], &input) + .unwrap(), + vec![0], + "type={type_id}, already_normalized={already_normalized}", + ); + } + } + } + } + #[tokio::test] async fn floating_sort_keys_preserve_window_group_limit_peers() { let planner = PhysicalPlanner::default(); diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index a0257cfc9d8..9d917d537e4 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -200,6 +200,67 @@ class CometWindowExecSuite extends CometTestBase { } } + for { + partitionColumn <- Seq("f", "d") + (function, groupLimit) <- Seq( + "RANK()" -> false, + "PERCENT_RANK()" -> false, + "NTILE(2)" -> false, + "RANK()" -> true) + } { + test( + s"window: floating partition keys ($partitionColumn, $function, group limit=$groupLimit)") { + assume(!groupLimit || isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") + + val partitionKeys = Seq( + (Some(1.0f), Some(1.0d)), + (Some(2.0f), Some(2.0d)), + (None, None), + (Some(0.0f), Some(0.0d)), + (Some(-0.0f), Some(-0.0d)), + (Some(Float.NaN), Some(Double.NaN)), + ( + Some(java.lang.Float.intBitsToFloat(0xffc00002)), + Some(java.lang.Double.longBitsToDouble(0xfff8000000000002L)))) + val data = partitionKeys.zipWithIndex.flatMap { case ((f, d), index) => + (0 until 4).map(i => (index * 4 + i, f, d, i % 2)) + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "false", + CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.key -> groupLimit.toString) { + withTempView("floating_window_partitions") { + data + .toDF("id", "f", "d", "secondary") + .createOrReplaceTempView("floating_window_partitions") + + for (partitionBy <- Seq(partitionColumn, s"secondary, $partitionColumn")) { + // Spark normalizes floating partition keys before constructing the child sort. + // Native sort and window expressions must still agree on these keys, including + // for bounded RANK and non-bounded PERCENT_RANK/NTILE without WindowGroupLimit. + val windowQuery = s""" + |SELECT id, $function OVER (PARTITION BY $partitionBy ORDER BY id) AS rnk + |FROM floating_window_partitions + |""".stripMargin + val query = if (groupLimit) { + s"SELECT * FROM ($windowQuery) WHERE rnk <= 1" + } else { + windowQuery + } + val operators = Seq(classOf[CometSortExec], classOf[CometWindowExec]) ++ + (if (groupLimit) Seq(classOf[CometWindowGroupLimitExec]) else Seq.empty) + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query), operators) + if (!groupLimit) { + assert(collect(cometPlan) { case w: CometWindowGroupLimitExec => w }.isEmpty) + } + } + } + } + } + } + test("lead/lag should return the default value if the offset row does not exist") { withSQLConf( CometConf.COMET_ENABLED.key -> "true", From 38825e26a1324297a46821a06fe5ce33e8b47779 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 27 Aug 2026 17:30:25 +0000 Subject: [PATCH 3/3] bench: measure normalized floating sort keys --- .../latest/compatibility/floating-point.md | 11 +- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/sort_float_keys.rs | 349 ++++++++++++++++++ 3 files changed, 360 insertions(+), 4 deletions(-) create mode 100644 native/spark-expr/benches/sort_float_keys.rs diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index ec03e59b037..b39ef7f645f 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -39,7 +39,10 @@ sorting, window peer comparisons, and `WindowGroupLimitExec` rank comparisons. N partitioning normalizes its keys and sampled boundaries in the same way. Only comparison keys are normalized; returned values retain their original NaN representations and zero signs. -Floating-point values nested in arrays or structs still use Arrow's raw total ordering and can -produce different ordering or rank results from Spark. The existing -`spark.comet.exec.strictFloatingPoint=true` fallback policy is unchanged, including its -conservative fallback for scalar floating-point sort keys. +Native sorting of floating-point values nested in arrays or structs still uses Arrow's raw total +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). diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 2b078268559..2c87debaf50 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -111,6 +111,10 @@ harness = false name = "normalize_nan" harness = false +[[bench]] +name = "sort_float_keys" +harness = false + [[bench]] name = "to_csv" harness = false diff --git a/native/spark-expr/benches/sort_float_keys.rs b/native/spark-expr/benches/sort_float_keys.rs new file mode 100644 index 00000000000..3c50a10816d --- /dev/null +++ b/native/spark-expr/benches/sort_float_keys.rs @@ -0,0 +1,349 @@ +// 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. + +//! Compare the scalar sort keys before/after #5469 in the same optimized binary. +//! +//! This runs DataFusion's full SortExec over an in-memory, single-partition input; +//! it does not time Spark planning, JNI, input generation, I/O, or a distributed +//! range shuffle. Both variants retain the original output values. The mixed +//! inputs intentionally have different ordering semantics before normalization. +//! +//! Run with `cargo bench -p datafusion-comet-spark-expr --bench sort_float_keys`. +//! Set COMET_SORT_BENCH_REVERSE_VARIANTS=1 for a second pass in the opposite order. + +use std::cmp::Ordering; +use std::hint::black_box; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{Array, ArrayRef, Float32Array, Float64Array, Int64Array}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput}; +use datafusion::common::Result; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::config::SessionConfig; +use datafusion::execution::memory_pool::{ + GreedyMemoryPool, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{PhysicalExpr, PhysicalSortExpr}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion_comet_spark_expr::NormalizeNaNAndZero; +use futures::TryStreamExt; +use tokio::runtime::Builder; + +const BATCH_ROWS: usize = 8192; +const BATCH_COUNT: usize = 32; +const ROWS: usize = BATCH_ROWS * BATCH_COUNT; +const MEMORY_LIMIT: usize = 128 * 1024 * 1024; + +fn input_batches(data_type: &DataType, mixed: bool) -> Vec { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", data_type.clone(), mixed), + Field::new("row_id", DataType::Int64, false), + ])); + let mut state = 0x1234_5678_9abc_def0_u64; + (0..BATCH_COUNT) + .map(|batch| { + let start = batch * BATCH_ROWS; + let values: Vec> = (start..start + BATCH_ROWS) + .map(|row| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + let finite = (state % 1_000_000) as f64 / 16.0 - 31_250.0; + match (mixed, row % 100) { + (true, 0..=19) => None, + (true, 20..=24) => Some(f64::from_bits(0xfff8_0000_0000_0000 | row as u64)), + (true, 25..=29) => Some(f64::from_bits(0x7ff8_0000_0000_0000 | row as u64)), + (true, 30..=34) => Some(-0.0), + (true, 35..=39) => Some(0.0), + _ => Some(finite), + } + }) + .collect(); + let keys: ArrayRef = match data_type { + DataType::Float32 => Arc::new(Float32Array::from( + values + .into_iter() + .enumerate() + .map(|(row, value)| { + value.map(|v| { + if v.is_nan() { + let sign = if v.is_sign_negative() { 0x8000_0000 } else { 0 }; + f32::from_bits(sign | 0x7fc0_0000 | (start + row) as u32) + } else { + v as f32 + } + }) + }) + .collect::>(), + )), + DataType::Float64 => Arc::new(Float64Array::from(values)), + _ => unreachable!(), + }; + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + keys, + Arc::new(Int64Array::from_iter_values( + (start..start + BATCH_ROWS).map(|row| row as i64), + )), + ], + ) + .unwrap() + }) + .collect() +} + +fn sort_plan(batches: &[RecordBatch], normalized: bool) -> Arc { + let input = + MemorySourceConfig::try_new_exec(&[batches.to_vec()], batches[0].schema(), None).unwrap(); + let mut key: Arc = Arc::new(Column::new("key", 0)); + if normalized { + key = Arc::new(NormalizeNaNAndZero::new( + batches[0].schema().field(0).data_type().clone(), + key, + )); + } + Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: key, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }] + .into(), + input, + )) +} + +fn task_context(pool: Arc) -> Arc { + let mut config = SessionConfig::new() + .with_batch_size(BATCH_ROWS) + .with_target_partitions(1); + config.options_mut().execution.sort_in_place_threshold_bytes = 1024 * 1024; + config.options_mut().execution.sort_spill_reservation_bytes = 10 * 1024 * 1024; + Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc() + .unwrap(), + ), + ) +} + +/// Only used by the untimed validation run. Pool reservations do not include +/// every Arrow allocation, including temporary normalized comparison keys. +#[derive(Debug)] +struct PeakPool { + inner: GreedyMemoryPool, + peak: AtomicUsize, +} + +impl PeakPool { + fn record_peak(&self) { + self.peak + .fetch_max(self.inner.reserved(), AtomicOrdering::Relaxed); + } +} + +impl std::fmt::Display for PeakPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Sort benchmark peak memory pool") + } +} + +impl MemoryPool for PeakPool { + fn name(&self) -> &str { + "sort_benchmark_peak" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record_peak(); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record_peak(); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +fn raw_value(array: &ArrayRef, row: usize) -> (Option, u64) { + if array.is_null(row) { + return (None, 0); + } + match array.data_type() { + DataType::Float32 => { + let value = array + .as_any() + .downcast_ref::() + .unwrap() + .value(row); + (Some(value as f64), u64::from(value.to_bits())) + } + DataType::Float64 => { + let value = array + .as_any() + .downcast_ref::() + .unwrap() + .value(row); + (Some(value), value.to_bits()) + } + _ => unreachable!(), + } +} + +fn compare(left: Option, right: Option, normalized: bool) -> Ordering { + let canonical = |v: f64| { + if normalized && v.is_nan() { + f64::NAN + } else if normalized && v == 0.0 { + 0.0 + } else { + v + } + }; + match (left, right) { + (None, None) => Ordering::Equal, + (None, Some(_)) => Ordering::Greater, + (Some(_), None) => Ordering::Less, + (Some(a), Some(b)) => canonical(a).total_cmp(&canonical(b)), + } +} + +async fn validate(batches: &[RecordBatch], normalized: bool) -> usize { + let pool = Arc::new(PeakPool { + inner: GreedyMemoryPool::new(MEMORY_LIMIT), + peak: AtomicUsize::new(0), + }); + let sort = sort_plan(batches, normalized); + let output = collect( + Arc::clone(&sort) as Arc, + task_context(Arc::clone(&pool) as Arc), + ) + .await + .unwrap(); + let mut seen = vec![false; ROWS]; + let mut previous = None; + for batch in output { + let ids = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + let id = ids.value(row) as usize; + assert!(!seen[id]); + seen[id] = true; + let (value, bits) = raw_value(batch.column(0), row); + let (original, original_bits) = + raw_value(batches[id / BATCH_ROWS].column(0), id % BATCH_ROWS); + assert_eq!(value.is_none(), original.is_none()); + assert_eq!( + bits, original_bits, + "sort changed the returned floating bits" + ); + if let Some(previous) = previous { + assert_ne!(compare(previous, value, normalized), Ordering::Greater); + } + previous = Some(value); + } + } + assert!(seen.into_iter().all(|value| value)); + assert_eq!(sort.metrics().unwrap().spill_count().unwrap_or(0), 0); + assert_eq!(pool.reserved(), 0); + pool.peak.load(AtomicOrdering::Relaxed) +} + +fn benchmark(c: &mut Criterion) { + let runtime = Builder::new_current_thread().enable_all().build().unwrap(); + let context = task_context(Arc::new(GreedyMemoryPool::new(MEMORY_LIMIT))); + let mut variants = [("bare", false), ("normalized", true)]; + if std::env::var_os("COMET_SORT_BENCH_REVERSE_VARIANTS").is_some() { + variants.reverse(); + } + let mut group = c.benchmark_group("sort_float_keys"); + group.sample_size(30); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + group.throughput(Throughput::Elements(ROWS as u64)); + for (name, data_type) in [ + ("float32", DataType::Float32), + ("float64", DataType::Float64), + ] { + for (shape, mixed) in [("finite", false), ("mixed", true)] { + let batches = input_batches(&data_type, mixed); + for (variant, normalized) in variants { + let peak = runtime.block_on(validate(&batches, normalized)); + eprintln!( + "{name}/{shape}/{variant}: rows={ROWS}, pool_peak_reserved_bytes={peak}, spills=0" + ); + group.bench_with_input( + BenchmarkId::new(format!("{name}_{shape}"), variant), + &normalized, + |b, &normalized| { + b.to_async(&runtime).iter_batched( + || sort_plan(&batches, normalized), + |sort| { + let context = Arc::clone(&context); + async move { + let mut stream = sort.execute(0, context).unwrap(); + let mut rows = 0; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += black_box(batch).num_rows(); + } + assert_eq!(rows, ROWS); + black_box(rows) + } + }, + BatchSize::LargeInput, + ); + }, + ); + } + } + } + group.finish(); +} + +criterion_group!(benches, benchmark); +criterion_main!(benches);