我正在尝试编写一个辅助函数,该函数采用任何类型的数据集Dataset[_]
,并返回一个新列"partitionId",这是单个数据单元所属分区的ID。
例如,如果我下面有一个数据集,默认情况下它有两个分区。
+-----+------+
| colA| colB|
+-----+------+
| 1 | a|
| 2 | b|
| 3 | c|
+-----+------+
在函数之后,它应该是下面的结果,其中前两个数据单元属于同一个分区,第三个属于另一个分区。
+-----+------+------------+
| colA| colB| partitionId|
+-----+------+------------+
| 1 | a| 1|
| 2 | b| 1|
| 3 | c| 2|
+-----+------+------------+
我尝试使用Column()和mapPartitions(),但它们都不适合我。 对于withColumn(),我无法获得数据单元所属分区的信息,例如withColumn("partitionId", {What should be here to add the partitionId?})
对于mapPartitions(),我尝试了:
dataset
.mapPartitions(iter => {
val partitionId = UUID.randomUUID().toString
iter.map(dataUnit => MyDataType.addPartitionId(partitionId))
})
但这仅适用于特定类型,如Dataset[MyDataType]
,不适用于Dataset[_]
如何为任何数据集添加分区 Id 列?
是否有理由需要每条记录的分区 ID?无论哪种方式,您都可以通过以下方式实现它:
import org.apache.spark.sql.functions.spark_partition_id
...
dataFrame.withColumn("partitionID", spark_partition_id)