uros-b commented on code in PR #58235: URL: https://github.com/apache/spark/pull/58235#discussion_r3933121751
########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/PivotFirst.scala: ########## @@ -88,11 +89,27 @@ case class PivotFirst( private val usesTreeMap: Boolean = !TypeUtils.typeWithProperEquals(pivotColumn.dataType) Review Comment: On this PR’s base, DoubleType still uses HashMap, and java.lang.Double.equals treats +0.0 and -0.0 as different. IN (0.0D, -0.0D) does not shrink the buffer today. They do compare equal under SQLOrderingUtil.compareDoubles (if (x == y) 0), so they will collapse once [#58234](https://github.com/apache/spark/pull/58234) puts floats/doubles on the TreeMap path. That is why this PR should land first — 58234 even calls out the duplicate-NaN regression as SPARK-58959. Please tighten the description to: identical literals, and values that compare equal on the current index (UTF8_LCASE 'a'/'A', equal structs/binaries). Leave signed zeros to 58234, or add a test there. ########## sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala: ########## @@ -517,4 +521,100 @@ class DataFramePivotSuite extends SharedSparkSession { matchPVals = true) } } + + test("SPARK-58959: duplicate pivot values each get their own output column") { + withTempView("dup_pivot") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_pivot AS + |SELECT * FROM VALUES + | (1, 1, 10), + | (1, 2, 20), + | (2, 1, 30) + |AS t(id, k, v)""".stripMargin) + + // The optimized PivotFirst path deduplicates the pivot values into a map, so a repeated + // value must not shrink the aggregation buffer below the number of output columns. + val df = sql( + """SELECT * FROM dup_pivot + |PIVOT (SUM(v) FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(usesPivotFirst(df)) + checkAnswer(df, Row(1, 10L, 10L, 20L) :: Row(2, 30L, 30L, null) :: Nil) + + // The path that does not use PivotFirst answers the same query the same way. + val arrayDf = sql( + """SELECT * FROM (SELECT id, k, ARRAY(v) AS v FROM dup_pivot) + |PIVOT (MIN(v) FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(!usesPivotFirst(arrayDf)) + checkAnswer( + arrayDf, + Row(1, Seq(10), Seq(10), Seq(20)) :: Row(2, Seq(30), Seq(30), null) :: Nil) + } + } + + test("SPARK-58959: duplicate pivot values do not corrupt a neighbouring aggregate buffer") { + withTempView("dup_multi_agg") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_multi_agg AS + |SELECT * FROM VALUES + | (1, 1, 10), + | (1, 1, 40), + | (1, 2, 20), + | (2, 1, 30) + |AS t(id, k, v)""".stripMargin) + + // With more than one aggregate the PivotFirst buffers sit back to back in a single row, so + // an out-of-range index reaches into the next aggregate's slots: SUM writes over MAX's + // buffer, and MAX runs off the end of the row entirely. + val df = sql( + """SELECT * FROM dup_multi_agg + |PIVOT (SUM(v) AS s, MAX(v) AS m FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(usesPivotFirst(df)) + checkAnswer( + df, + Row(1, 50L, 40, 50L, 40, 20L, 20) :: Row(2, 30L, 30, 30L, 30, null, null) :: Nil) + } + } + + test("SPARK-58959: duplicate pivot values of every PivotFirst datatype") { + withTempView("dup_types") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_types AS + |SELECT * FROM VALUES + | ('a', 1Y, 1S, 1, 1L, 1.0F, 1.0D, 1.0BD, TRUE, 10), + | ('b', 2Y, 2S, 2, 2L, 2.0F, 2.0D, 2.0BD, FALSE, 20) + |AS t(s, b, sh, i, l, f, d, dec, bool, v)""".stripMargin) + + Seq("s" -> "'a'", "b" -> "1Y", "sh" -> "1S", "i" -> "1", "l" -> "1L", + "f" -> "1.0F", "d" -> "1.0D", "dec" -> "1.0BD", "bool" -> "TRUE").foreach { + case (column, value) => + checkAnswer( + sql( + s"""SELECT * FROM (SELECT $column AS k, v FROM dup_types) + |PIVOT (SUM(v) FOR k IN ($value AS x, $value AS y))""".stripMargin), + Row(10L, 10L)) + } + } + } + + test("SPARK-58959: pivot values that compare as equal share an output value") { + // The pivot column is collated, so PivotFirst indexes the values with a comparison-based + // TreeMap: 'a' and 'A' are distinct strings that compare as equal under UTF8_LCASE. + withTable("dup_lcase_pivot") { + sql( + """CREATE TABLE dup_lcase_pivot ( + | key STRING COLLATE UTF8_LCASE, + | amount INT + |) USING PARQUET""".stripMargin) + sql( + """INSERT INTO dup_lcase_pivot VALUES + | ('a', 10), + | ('b', 20)""".stripMargin) + + checkAnswer( + sql( + """SELECT * FROM dup_lcase_pivot + |PIVOT (SUM(amount) FOR key IN ('a' AS x, 'A' AS y, 'b' AS z))""".stripMargin), + Row(10L, 10L, 20L)) + } + } Review Comment: Duplicate NaN is an untested HashMap overflow. Double.equals treats two NaNs as equal, so IN (double('NaN') AS x, double('NaN') AS y) already collapses on master and should hit the same OOB write. Worth one case next to the integer duplicate test. (Input-row NaN lookup is SPARK-39031 / 58234; duplicate literals in the IN list are this ticket.) ########## sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala: ########## @@ -517,4 +521,100 @@ class DataFramePivotSuite extends SharedSparkSession { matchPVals = true) } } + + test("SPARK-58959: duplicate pivot values each get their own output column") { + withTempView("dup_pivot") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_pivot AS + |SELECT * FROM VALUES + | (1, 1, 10), + | (1, 2, 20), + | (2, 1, 30) + |AS t(id, k, v)""".stripMargin) + + // The optimized PivotFirst path deduplicates the pivot values into a map, so a repeated + // value must not shrink the aggregation buffer below the number of output columns. + val df = sql( + """SELECT * FROM dup_pivot + |PIVOT (SUM(v) FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(usesPivotFirst(df)) + checkAnswer(df, Row(1, 10L, 10L, 20L) :: Row(2, 30L, 30L, null) :: Nil) + + // The path that does not use PivotFirst answers the same query the same way. + val arrayDf = sql( + """SELECT * FROM (SELECT id, k, ARRAY(v) AS v FROM dup_pivot) + |PIVOT (MIN(v) FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(!usesPivotFirst(arrayDf)) + checkAnswer( + arrayDf, + Row(1, Seq(10), Seq(10), Seq(20)) :: Row(2, Seq(30), Seq(30), null) :: Nil) + } + } + + test("SPARK-58959: duplicate pivot values do not corrupt a neighbouring aggregate buffer") { + withTempView("dup_multi_agg") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_multi_agg AS + |SELECT * FROM VALUES + | (1, 1, 10), + | (1, 1, 40), + | (1, 2, 20), + | (2, 1, 30) + |AS t(id, k, v)""".stripMargin) + + // With more than one aggregate the PivotFirst buffers sit back to back in a single row, so + // an out-of-range index reaches into the next aggregate's slots: SUM writes over MAX's + // buffer, and MAX runs off the end of the row entirely. + val df = sql( + """SELECT * FROM dup_multi_agg + |PIVOT (SUM(v) AS s, MAX(v) AS m FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(usesPivotFirst(df)) + checkAnswer( + df, + Row(1, 50L, 40, 50L, 40, 20L, 20) :: Row(2, 30L, 30, 30L, 30, null, null) :: Nil) + } + } + + test("SPARK-58959: duplicate pivot values of every PivotFirst datatype") { + withTempView("dup_types") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_types AS + |SELECT * FROM VALUES + | ('a', 1Y, 1S, 1, 1L, 1.0F, 1.0D, 1.0BD, TRUE, 10), + | ('b', 2Y, 2S, 2, 2L, 2.0F, 2.0D, 2.0BD, FALSE, 20) + |AS t(s, b, sh, i, l, f, d, dec, bool, v)""".stripMargin) + + Seq("s" -> "'a'", "b" -> "1Y", "sh" -> "1S", "i" -> "1", "l" -> "1L", + "f" -> "1.0F", "d" -> "1.0D", "dec" -> "1.0BD", "bool" -> "TRUE").foreach { + case (column, value) => + checkAnswer( + sql( + s"""SELECT * FROM (SELECT $column AS k, v FROM dup_types) + |PIVOT (SUM(v) FOR k IN ($value AS x, $value AS y))""".stripMargin), + Row(10L, 10L)) + } + } + } + + test("SPARK-58959: pivot values that compare as equal share an output value") { + // The pivot column is collated, so PivotFirst indexes the values with a comparison-based Review Comment: The collation test is the important TreeMap case and matches the existing UTF8_LCASE tests in that suite. Optional extras: duplicate struct IN-list values (same TreeMap as SPARK-55483), and DataFrame pivot("k", Seq(1, 1, 2)). ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/PivotFirst.scala: ########## @@ -104,7 +121,7 @@ case class PivotFirst( case _ => pivotIndex.getOrElse(key, -1) } - val indexSize = pivotIndex.size + val indexSize = slotValues.length Review Comment: Please consider adding a one-line comment that indexSize is distinct buffer slots and slotOfValue.length is output arity would help; mixing those two was the bug. ########## sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala: ########## @@ -517,4 +521,100 @@ class DataFramePivotSuite extends SharedSparkSession { matchPVals = true) } } + + test("SPARK-58959: duplicate pivot values each get their own output column") { + withTempView("dup_pivot") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_pivot AS + |SELECT * FROM VALUES + | (1, 1, 10), + | (1, 2, 20), + | (2, 1, 30) + |AS t(id, k, v)""".stripMargin) + + // The optimized PivotFirst path deduplicates the pivot values into a map, so a repeated + // value must not shrink the aggregation buffer below the number of output columns. + val df = sql( + """SELECT * FROM dup_pivot + |PIVOT (SUM(v) FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(usesPivotFirst(df)) + checkAnswer(df, Row(1, 10L, 10L, 20L) :: Row(2, 30L, 30L, null) :: Nil) + + // The path that does not use PivotFirst answers the same query the same way. + val arrayDf = sql( + """SELECT * FROM (SELECT id, k, ARRAY(v) AS v FROM dup_pivot) + |PIVOT (MIN(v) FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(!usesPivotFirst(arrayDf)) + checkAnswer( + arrayDf, + Row(1, Seq(10), Seq(10), Seq(20)) :: Row(2, Seq(30), Seq(30), null) :: Nil) + } + } + + test("SPARK-58959: duplicate pivot values do not corrupt a neighbouring aggregate buffer") { + withTempView("dup_multi_agg") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_multi_agg AS + |SELECT * FROM VALUES + | (1, 1, 10), + | (1, 1, 40), + | (1, 2, 20), + | (2, 1, 30) + |AS t(id, k, v)""".stripMargin) + + // With more than one aggregate the PivotFirst buffers sit back to back in a single row, so + // an out-of-range index reaches into the next aggregate's slots: SUM writes over MAX's + // buffer, and MAX runs off the end of the row entirely. + val df = sql( + """SELECT * FROM dup_multi_agg + |PIVOT (SUM(v) AS s, MAX(v) AS m FOR k IN (1 AS x, 1 AS y, 2 AS z))""".stripMargin) + assert(usesPivotFirst(df)) + checkAnswer( + df, + Row(1, 50L, 40, 50L, 40, 20L, 20) :: Row(2, 30L, 30, 30L, 30, null, null) :: Nil) + } + } + + test("SPARK-58959: duplicate pivot values of every PivotFirst datatype") { + withTempView("dup_types") { + sql( + """CREATE OR REPLACE TEMP VIEW dup_types AS + |SELECT * FROM VALUES + | ('a', 1Y, 1S, 1, 1L, 1.0F, 1.0D, 1.0BD, TRUE, 10), + | ('b', 2Y, 2S, 2, 2L, 2.0F, 2.0D, 2.0BD, FALSE, 20) + |AS t(s, b, sh, i, l, f, d, dec, bool, v)""".stripMargin) + + Seq("s" -> "'a'", "b" -> "1Y", "sh" -> "1S", "i" -> "1", "l" -> "1L", + "f" -> "1.0F", "d" -> "1.0D", "dec" -> "1.0BD", "bool" -> "TRUE").foreach { + case (column, value) => + checkAnswer( + sql( + s"""SELECT * FROM (SELECT $column AS k, v FROM dup_types) + |PIVOT (SUM(v) FOR k IN ($value AS x, $value AS y))""".stripMargin), + Row(10L, 10L)) + } + } + } + Review Comment: Two tests never check they are on the fast path. usesPivotFirst is the whole point of the first tests, but it is missing from: duplicate pivot values of every PivotFirst datatype pivot values that compare as equal share an output value If those queries ever fall off PivotFirst, they would still pass on the slow path and stop guarding this bug. Same assert(usesPivotFirst(df)) as the other new tests. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/PivotFirst.scala: ########## @@ -88,11 +89,27 @@ case class PivotFirst( private val usesTreeMap: Boolean = !TypeUtils.typeWithProperEquals(pivotColumn.dataType) - val pivotIndex: Map[Any, Int] = if (usesTreeMap) { - TreeMap(pivotColumnValues.zipWithIndex: _*)( - TypeUtils.getInterpretedOrdering(pivotColumn.dataType)) - } else { - HashMap(pivotColumnValues.zipWithIndex: _*) + // Every pivot value gets a buffer slot, shared between the values that compare as equal, so a + // repeated value does not shrink the buffer below the number of output columns. + private val (pivotIndex, slotOfValue, slotValues) = { + val slots = new Array[Int](pivotColumnValues.length) + val values = mutable.ArrayBuffer.empty[Any] + var index: Map[Any, Int] = if (usesTreeMap) { + TreeMap.empty[Any, Int](TypeUtils.getInterpretedOrdering(pivotColumn.dataType)) + } else { + HashMap.empty[Any, Int] + } + pivotColumnValues.zipWithIndex.foreach { case (value, i) => + val existingSlot = index.getOrElse(value, -1) Review Comment: ```suggestion val existingSlot = index.get(value) ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
