Question

我有一个记录每个人姓名和地址的数据框，表示为

case class Person(name:String, addr:String)

数据框看起来像这样，

+----+-----+
|name| addr|
+----+-----+
|  u1|addr1|
|  u1|addr2|
|  u2|addr1|
+----+-----+

但现在我需要为此数据框中的每个元素分配一个Long类型唯一ID，可以表示为

case class PersonX(name:String, name_id:Long, addr:String, addr_id:Long)

和dataframe看起来像这样，

+----+-------+-----+------+
|name|name_id| addr|addr_id|
+----+-------+-----+------+
|  u1|      1|addr1|     2|
|  u1|      1|addr2|     3|
|  u2|      4|addr1|     2|
+----+-------+-----+------+

请注意，两列中的元素（name和addr）共享相同的ID空间，这意味着name_id不应该有重复项，并addr_id也不应该，name_id s＆amp; addr_id不应相互重叠。

如何实现这一目标？

Answer 1

分配id的最简单方法是使用spark sql中的dense_rank函数。

为确保ID不会在名称和地址之间重叠，您可以制作一个技巧：

分别计算名称和地址的ID
通过添加名称的最大ID

通过这种方式，地址ID将位于名称ID之后

val input = spark.sparkContext.parallelize(List(
  Person("u1", "addr1"),
  Person("u1", "addr2"),
  Person("u2", "addr1")
)).toDF("name", "addr")

input.createOrReplaceTempView("people")
val people = spark.sql(
  """select name,
    |       dense_rank() over(partition by 1 order by name) as name_id,
    |       addr,
    |       dense_rank() over(partition by 1 order by addr) as addr_id
    |   from people
  """.stripMargin)
people.show()

//+----+-------+-----+-------+
//|name|name_id| addr|addr_id|
//+----+-------+-----+-------+
//|  u1|      1|addr1|      1|
//|  u2|      2|addr1|      1|
//|  u1|      1|addr2|      2|
//+----+-------+-----+-------+    

val name = people.col("name")
val nameId = people.col("name_id")
val addr = people.col("addr")
val addrId = people.col("addr_id")

val maxNameId = people.select(max(nameId)).first().getInt(0)
val shiftedAddrId = (addrId + maxNameId).as("addr_id")
people.select(name, addr, nameId, shiftedAddrId).as[PersonX].show()

//+----+-----+-------+-------+
//|name| addr|name_id|addr_id|
//+----+-----+-------+-------+
//|  u1|addr1|      1|      3|
//|  u2|addr1|      2|      3|
//|  u1|addr2|      1|      4|
//+----+-----+-------+-------+

为数据框元素分配唯一ID

1 个答案: