Github user rxin commented on a diff in the pull request:
https://github.com/apache/spark/pull/10968#discussion_r51090555
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala ---
@@ -374,21 +366,27 @@ final class DataFrameStatFunctions private[sql](df:
DataFrame) {
val singleCol = df.select(col)
val colType = singleCol.schema.head.dataType
- require(
- colType == StringType || colType.isInstanceOf[IntegralType],
- s"Count-min Sketch only supports string type and integral types, " +
- s"and does not support type $colType."
- )
+ val updater: (CountMinSketch, InternalRow) => Unit = colType match {
+ // For string type, we can get bytes of our `UTF8String` directly,
and call the `addBinary`
+ // instead of `addString` to avoid unnecessary conversion.
+ case StringType => (sketch, row) =>
sketch.addBinary(row.getUTF8String(0).getBytes)
+ case ByteType => (sketch, row) => sketch.addLong(row.getByte(0))
+ case ShortType => (sketch, row) => sketch.addLong(row.getShort(0))
+ case IntegerType => (sketch, row) => sketch.addLong(row.getInt(0))
+ case LongType => (sketch, row) => sketch.addLong(row.getLong(0))
+ case _ =>
+ throw new IllegalArgumentException(
+ s"Count-min Sketch only supports string type and integral types,
" +
+ s"and does not support type $colType."
+ )
+ }
- singleCol.rdd.aggregate(zero)(
- (sketch: CountMinSketch, row: Row) => {
- sketch.add(row.get(0))
+ singleCol.queryExecution.toRdd.aggregate(zero)(
+ (sketch: CountMinSketch, row: InternalRow) => {
+ updater(sketch, row)
sketch
},
-
- (sketch1: CountMinSketch, sketch2: CountMinSketch) => {
- sketch1.mergeInPlace(sketch2)
- }
+ _ mergeInPlace _
--- End diff --
no infix notation
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
