GideonPotok commented on code in PR #47154:
URL: https://github.com/apache/spark/pull/47154#discussion_r1770554191
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala:
##########
@@ -86,6 +91,53 @@ case class Mode(
buffer
}
+ private def getCollationAwareBuffer(
+ childDataType: DataType,
+ buffer: OpenHashMap[AnyRef, Long]): Iterable[(AnyRef, Long)] = {
+ def groupAndReduceBuffer(groupingFunction: AnyRef => _): Iterable[(AnyRef,
Long)] = {
+ buffer.groupMapReduce(t =>
+ groupingFunction(t._1))(x => x)((x, y) => (x._1, x._2 + y._2)).values
+ }
+ def determineBufferingFunction(
+ childDataType: DataType): Option[AnyRef => _] = {
+ childDataType match {
+ case _ if UnsafeRowUtils.isBinaryStable(child.dataType) => None
+ case _ => Some(collationAwareTransform(_, childDataType))
+ }
+ }
+
determineBufferingFunction(childDataType).map(groupAndReduceBuffer).getOrElse(buffer)
+ }
+
+ private def collationAwareTransform(data: AnyRef, dataType: DataType):
AnyRef = {
+ dataType match {
+ case _ if UnsafeRowUtils.isBinaryStable(dataType) => data
+ case st: StructType =>
+
processStructTypeWithBuffer(data.asInstanceOf[InternalRow].toSeq(st).zip(st.fields))
+ case at: ArrayType => processArrayTypeWithBuffer(at,
data.asInstanceOf[ArrayData])
+ case st: StringType =>
+ CollationFactory.getCollationKey(data.asInstanceOf[UTF8String],
st.collationId)
+ case _ =>
+ throw new SparkUnsupportedOperationException(
+ "UNSUPPORTED_MODE_DATA_TYPE",
Review Comment:
@MaxGekk For the test, I suppose the situation where a type is not binary
stable and is not covered by checkInputDataTypes would include some UDTs?
checkInputType just confirms it is not a MapType (a blacklist/blocklist),
whereas `collationAwareTransform` is an allowlist/whitelist).
--
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]