Spark Structured Streaming - groupByKey 按分区单独

Question

我的 Kafka 生产者正在根据给定的键将消息分发到主题分区中。

所以，在 Spark 端，我已经有需要在同一个分区中一起处理的消息。

现在，我需要做一个 groupByKey 以在每个分区中通过键在列表中聚合值，但不需要合并分区，因为没有机会在多个分区中拥有给定的键。

我怎样才能仅在分区级别执行此 groupByKey ？

|topic-partition1| ---> |spark-partition1| -- groupByKey --> |spark-partition1.1| -- mapGroupsWithState --> ...
|topic-partition2| ---> |spark-partition2| -- groupByKey --> |spark-partition2.1| -- mapGroupsWithState --> ...
|topic-partition3| ---> |spark-partition3| -- groupByKey --> |spark-partition3.1| -- mapGroupsWithState --> ...

Answer 1

如果您知道所有事件都将进入给定的分区，您可以在数据集上使用DataSet.mapPartitions：

val dataSet: DataSet[(String, String)] = ???
dataSet.mapPartitions { iter =>
  val res: Map[String, List[(String, String)] =
    iter.toList.groupBy { case (key, _) => key }

  // Do additional processing on res, which is now grouped by each key
  // present in the partition.
}

否则，如果您需要mapGroupsWithState，则可以避免使用groupByKey，因为您需要KeyValueGroupedDataset[K, V]。

如果您关心性能，请不要担心，除非您在分析时发现这是一个瓶颈。