jiwen624 commented on code in PR #58235:
URL: https://github.com/apache/spark/pull/58235#discussion_r3943229203
##########
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:
Added 👍
--
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]
