Skip to content

[SPARK-52738][SQL] Support aggregating the TIME type with a UDAF when the underlying buffer is an UnsafeRow #51430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

[SPARK-52738][SQL] Support aggregating the TIME type with a UDAF when the underlying buffer is an UnsafeRow #51430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4eee734
Testing
bersprockets Jul 9, 2025
f59e75f
Put back AggregationQuerySuite more like it used to be pre-SPARK-41359
bersprockets Jul 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ sealed trait BufferSetterGetterUtils {
(row: InternalRow, ordinal: Int) =>
if (row.isNullAt(ordinal)) null else row.getInt(ordinal)

case TimestampType | TimestampNTZType =>
case TimestampType | TimestampNTZType | _: TimeType =>
(row: InternalRow, ordinal: Int) =>
if (row.isNullAt(ordinal)) null else row.getLong(ordinal)

Expand Down Expand Up @@ -188,7 +188,7 @@ sealed trait BufferSetterGetterUtils {
row.setNullAt(ordinal)
}

case TimestampType | TimestampNTZType =>
case TimestampType | TimestampNTZType | _: TimeType =>
(row: InternalRow, ordinal: Int, value: Any) =>
if (value != null) {
row.setLong(ordinal, value.asInstanceOf[Long])
Expand Down
16 changes: 10 additions & 6 deletions sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import test.org.apache.spark.sql.MyDoubleAvg
import test.org.apache.spark.sql.MyDoubleSum

import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, RandomDataGenerator, Row}
import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode
import org.apache.spark.sql.catalyst.expressions.{CodegenObjectFactoryMode, UnsafeRow}
import org.apache.spark.sql.classic.ClassicConversions.castToImpl
import org.apache.spark.sql.classic.Dataset
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
Expand Down Expand Up @@ -899,11 +899,15 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
ArrayType(IntegerType), MapType(StringType, LongType), struct,
new TestUDT.MyDenseVectorUDT()) ++ dayTimeIntervalTypes ++ unsafeRowMutableFieldTypes ++
timeTypes
// Right now, we will use SortAggregate to handle UDAFs.
// UnsafeRow.mutableFieldTypes.asScala.toSeq will trigger SortAggregate to use
// UnsafeRow as the aggregation buffer. While, dataTypes will trigger
// SortAggregate to use a safe row as the aggregation buffer.
Seq(dataTypes).foreach { dataTypes =>
// A schema that contains only data types where UnsafeRow.isMutable is true
// will trigger the aggregator to use unsafe row as the aggregation buffer.
// Other dataTypes will trigger the aggregator to use a safe row as the
// aggregation buffer.
//
// Below we want to test with *both* UnsafeRow and safe row as the underlying
// buffer.
val mutableDataTypes = dataTypes.filter(UnsafeRow.isMutable)
Seq(dataTypes, mutableDataTypes).foreach { dataTypes =>
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that even without the change to BufferSetterGetterUtils, the unsafe buffer test will sometimes succeed. That's because the inappropriate setter function works fine for nulls and RandomDataGenerator will occasionally give us a null TIME value for the row where id == 50.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, this test was succeeding even with TIME types because SPARK-41359 accidentally removed the unsafe buffer part of the test.

val fields = dataTypes.zipWithIndex.map { case (dataType, index) =>
StructField(s"col$index", dataType, nullable = true)
}
Expand Down