Spark根据单元格中的值添加相同的行号

Question

我有如下数据

//input data
df.show()
//+---+---+---+
//|  x|  y|  z|
//+---+---+---+
//|tes| 45| 34|
//|tes| 43| 67|
//|tes| 56| 43|
//|raj| 45| 43|
//|raj| 44| 67|
//+---+---+---+

我希望将其转换为不改变给定输入的顺序。

//output data
    df.show()
    //+---+---+---+---+
    //|  x|  y|  z|  n|
    //+---+---+---+---+
    //|tes| 45| 34|  1|
    //|tes| 43| 67|  1|
    //|tes| 56| 43|  1|
    //|raj| 45| 43|  2|
    //|raj| 44| 67|  2|
    //+---+---+---+---+

Answer 1

这里的想法是找到一种方法来添加一个包含“它是一个新的 x 吗？”信息的新列，然后进行滚动求和以“添加”这些信息。

为此，我们需要一个window函数和lag方法。

// Some imports
import org.apache.spark.sql.{functions => F}
import org.apache.spark.sql.expressions.Window

// inputs
val df = Seq(("tes", 45, 34), ("tes", 43, 67), ("tes", 56, 43), ("raj", 45, 43), ("raj", 44, 67)).toDF("x", "y", "z")

// adding new information
val windowSpec = Window.orderBy(F.monotonically_increasing_id())


val dfWithNewNameInfo = df
  .withColumn("n", (F.lag($"x", 1).over(windowSpec) =!= $"x").cast("bigint"))
  .na.fill(1, Seq("n"))
dfWithNewNameInfo.show
/*
+---+---+---+---+
|  x|  y|  z|  n|
+---+---+---+---+
|tes| 45| 34|  1|
|tes| 43| 67|  0|
|tes| 56| 43|  0|
|raj| 45| 43|  1|
|raj| 44| 67|  0|
+---+---+---+---+


*/
// We can see the "1" in the last column indicates whenever this is a new x

// Adding these 1
val resultDf = dfWithNewNameInfo.withColumn("n", F.sum("n").over(windowSpec))
resultDf.show
/*
+---+---+---+---+
|  x|  y|  z|  n|
+---+---+---+---+
|tes| 45| 34|  1|
|tes| 43| 67|  1|
|tes| 56| 43|  1|
|raj| 45| 43|  2|
|raj| 44| 67|  2|
+---+---+---+---+
*/

此方法为可以完全加载到内存中的小数据帧提供了所需的结果。

---

请注意windowSpec顺序，但这对于以下 DataFrame 会失败：

//+---+---+---+
//|  x|  y|  z|
//+---+---+---+
//|tes| 45| 34|
//|tes| 43| 67|
//|raj| 45| 43|
//|raj| 44| 67|
//|tes| 56| 43|
//+---+---+---+

给出结果：

//+---+---+---+---+
//|  x|  y|  z|  n|
//+---+---+---+---+
//|tes| 45| 34|  1|
//|tes| 43| 67|  1|
//|raj| 45| 43|  2|
//|raj| 44| 67|  2|
//|tes| 56| 43|  3|
//+---+---+---+---+

这就是为什么我强烈建议在windowSpec按"x" windowSpec 。

Answer 2

如果你在 Spark with Java 中需要这个，

df.withColumn("n", lag(col("x"), 1, "a").over(Window.orderBy(monotonicallyIncreasingId())).notEqual(col("x")).cast(DataTypes.LongType))
            .withColumn("n", sum("n").over(Window.orderBy(monotonicallyIncreasingId())));