[SPARK-31182][CORE][ML] PairRDD support aggregateByKeyWithinPartitions

[zhengruifeng]

What changes were proposed in this pull request?

1, impl aggregateByKeyWithinPartitions and reduceByKeyWithinPartitions
2, use aggregateByKeyWithinPartitions in RobustScaler

Why are the changes needed?

When implementing RobustScaler, I was looking for a way to guarantee that the QuantileSummaries in aggregateByKey are compressed at the map side.

(before merge and qurey, the QuantileSummaries must be compressed)

Then I only found a tricky method to work around (yet not applied), and there was no method for this.
previous discussions were here

Does this PR introduce any user-facing change?

Yes, add new methods for PairRDD

How was this patch tested?

added testsuites and existing ones

[zhengruifeng]

this impl follows combineByKeyWithClassTag (treat as if self.partitioner == Some(partitioner))

def combineByKeyWithClassTag[C](
      createCombiner: V => C,
      mergeValue: (C, V) => C,
      mergeCombiners: (C, C) => C,
      partitioner: Partitioner,
      mapSideCombine: Boolean = true,
      serializer: Serializer = null)(implicit ct: ClassTag[C]): RDD[(K, C)] = self.withScope {
    require(mergeCombiners != null, "mergeCombiners must be defined") // required as of Spark 0.9.0
    if (keyClass.isArray) {
      if (mapSideCombine) {
        throw new SparkException("Cannot use map-side combining with array keys.")
      }
      if (partitioner.isInstanceOf[HashPartitioner]) {
        throw new SparkException("HashPartitioner cannot partition array keys.")
      }
    }
    val aggregator = new Aggregator[K, V, C](
      self.context.clean(createCombiner),
      self.context.clean(mergeValue),
      self.context.clean(mergeCombiners))
    if (self.partitioner == Some(partitioner)) {
      self.mapPartitions(iter => {
        val context = TaskContext.get()
        new InterruptibleIterator(context, aggregator.combineValuesByKey(iter, context))
      }, preservesPartitioning = true)
    } else {
      new ShuffledRDD[K, V, C](self, partitioner)
        .setSerializer(serializer)
        .setAggregator(aggregator)
        .setMapSideCombine(mapSideCombine)
    }
  }

[zhengruifeng]

here we can trigger compression at the map side

[zhengruifeng]

then no longer to confirm compression like s1.compress.merge(s2.compress)

[SparkQA]

Test build #119982 has finished for PR 27947 at commit 964b9f8.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[zhengruifeng]

friendly ping @srowen

[wetneb]

Naive question from a newbie: you are introducing new methods for PairRDDs, so would it make sense to also expose them in the Java API (as methods of JavaPairRDD)? Or is it generally expected that the Java API is updated separately?

[srowen]

Well, I think that's part of the issue here - if it's public you kind of need to support it everywhere and for a long time. I don't know if it's worth it but I've lost the thread on this PR and would have to recall the motivation.

[zhengruifeng]

I tend to close it, since I can always workaround it. Maybe it is not necessary.