Question

我有以下数据：

+-----+-----+----+
|Col1 |t0   |t1  |
+-----+-----+----+
| A   |null |20  |
| A   |20   |40  |
| B   |null |10  |
| B   |10   |20  |
| B   |20   |120 |
| B   |120  |140 |
| B   |140  |320 |
| B   |320  |340 |
| B   |340  |360 |
+-----+-----+----+

我想要的是这样的事情：

+-----+-----+----+----+
|Col1 |t0   |t1  |grp |
+-----+-----+----+----+
| A   |null |20  |1A  |
| A   |20   |40  |1A  |
| B   |null |10  |1B  |
| B   |10   |20  |1B  |
| B   |20   |120 |2B  |
| B   |120  |140 |2B  |
| B   |140  |320 |3B  |
| B   |320  |340 |3B  |
| B   |340  |360 |3B  |
+-----+-----+----+----+

说明：额外列基于Col1以及t1和t0之间的差异。当两者之间的差异太大时=>生成一个新数字。（当差异大于50时，在上面的数据集中）

我用：

构建t0

val windowSpec = Window.partitionBy($"Col1").orderBy("t1")
df = df.withColumn("t0", lag("t1", 1) over windowSpec)

有人可以帮我怎么做吗？我搜索但没有得到一个好主意。我有点迷失，因为我需要先前计算的grp行的值...

由于

Answer 1

我自己解决了

val grp =  (coalesce(
      ($"t" - lag($"t", 1).over(windowSpec)),
      lit(0)
    ) > 50).cast("bigint")

df = df.withColumn("grp", sum(grp).over(windowSpec))

有了这个，我不再需要colums（t0和t1），但只能使用t1（或t）而不使用compute t0。

（我只需要添加Col1的值，但最重要的部分是数字已完成且工作正常。）

我从以下方面得到了解决方案： Spark SQL window function with complex condition

感谢您的帮助

Answer 2

您可以使用udf功能生成grp column

def testUdf = udf((col1: String, t0: Int, t1: Int)=> (t1-t0) match {
  case x : Int if(x > 50) => 2+col1
  case _ => 1+col1
})

将udf功能调用为

df.withColumn("grp", testUdf($"Col1", $"t0", $"t1"))

上面的udf函数由于null中的t0值可以被0替换而无法正常工作

df.na.fill(0)

我希望这是您正在寻找的答案。

<强>被修改
这是使用udaf的完整解决方案。这个过程很复杂。你已经得到了简单的答案，但它可能会帮助那些可能使用它的人
首先定义udaf

class Boendal extends UserDefinedAggregateFunction {

  def inputSchema = new StructType().add("Col1", StringType).add("t0", IntegerType).add("t1", IntegerType).add("rank", IntegerType)
  def bufferSchema = new StructType().add("buff", StringType).add("buffer1", IntegerType)
  def dataType = StringType
  def deterministic = true

  def initialize(buffer: MutableAggregationBuffer) = {
    buffer.update(0, "")
    buffer.update(1, 0)
  }

  def update(buffer: MutableAggregationBuffer, input: Row) = {
    if (!input.isNullAt(0)) {
      val buff = buffer.getString(0)
      val col1 = input.getString(0)
      val t0 = input.getInt(1)
      val t1 = input.getInt(2)
      val rank = input.getInt(3)
      var value = 1
      if((t1-t0) < 50)
        value = 1
      else
        value = (t1-t0)/50

      val lastValue = buffer(1).asInstanceOf[Integer]
      // if(!buff.isEmpty) {
      if (value < lastValue)
        value = lastValue
      // }
      buffer.update(1, value)

      var finalString = ""
      if(buff.isEmpty){
        finalString = rank+";"+value+col1
      }
      else
        finalString = buff+"::"+rank+";"+value+col1

      buffer.update(0, finalString)
    }
  }

  def merge(buffer1: MutableAggregationBuffer, buffer2: Row) = {
    val buff1 = buffer1.getString(0)
    val buff2 = buffer2.getString(0)
    buffer1.update(0, buff1+buff2)
  }

  def evaluate(buffer: Row) : String = {
    buffer.getString(0)
  }
}

然后是一些udfs

def rankUdf = udf((grp: String)=> grp.split(";")(0))
def removeRankUdf = udf((grp: String) => grp.split(";")(1))

最后调用udaf和udfs

val windowSpec = Window.partitionBy($"Col1").orderBy($"t1")
df = df.withColumn("t0", lag("t1", 1) over windowSpec)
  .withColumn("rank", rank() over windowSpec)
df = df.na.fill(0)
val boendal = new Boendal
val df2 = df.groupBy("Col1").agg(boendal($"Col1", $"t0", $"t1", $"rank").as("grp2")).withColumnRenamed("Col1", "Col2")
    .withColumn("grp2", explode(split($"grp2", "::")))
    .withColumn("rank2", rankUdf($"grp2"))
    .withColumn("grp2", removeRankUdf($"grp2"))

df = df.join(df2, df("Col1") === df2("Col2") && df("rank") === df2("rank2"))
  .drop("Col2", "rank", "rank2")
df.show(false)

希望有所帮助

Spark Dataframe访问先前计算的行

2 个答案: