Dataframe 如何创建新列'；计数'；在一定条件下，在Spark数据帧中

我有一个关于连接日志的数据框，其中包含

Id

，

targetIP

，

Time

。此数据帧中的每个记录都是到一个系统的连接事件。Id表示此连接，

targetIP

表示此次的目标IP地址，time表示连接时间。价值观：

身份证件 时间 目标 1. 1. 192.163.0.1 2. 2. 192.163.0.2 3. 3. 192.163.0.1 4. 5. 192.163.0.1 5. 6. 192.163.0.2 6. 7. 192.163.0.2 7. 8. 192.163.0.2

---

您可以在范围介于-2和当前行之间的窗口上使用

count

，以获取最近2个时间单位中的IP计数

使用Spark SQL，您可以执行以下操作：

df.createOrReplaceTempView("connection_logs")

df1 = spark.sql("""
    SELECT  *,
            COUNT(*) OVER(PARTITION BY targetIP ORDER BY Time 
                          RANGE BETWEEN 2 PRECEDING AND CURRENT ROW
                          ) -1 AS count
    FROM    connection_logs
    ORDER BY ID
""")

df1.show()

#+---+----+-----------+-----+
#| ID|Time|   targetIP|count|
#+---+----+-----------+-----+
#|  1|   1|192.163.0.1|    0|
#|  2|   2|192.163.0.2|    0|
#|  3|   3|192.163.0.1|    1|
#|  4|   5|192.163.0.1|    1|
#|  5|   6|192.163.0.2|    0|
#|  6|   7|192.163.0.2|    1|
#|  7|   8|192.163.0.2|    2|
#+---+----+-----------+-----+

或使用数据帧API：

from pyspark.sql import Window
from pyspark.sql import functions as F

time_unit = lambda x: x

w = Window.partitionBy("targetIP").orderBy(col("Time").cast("int")).rangeBetween(-time_unit(2), 0)

df1 = df.withColumn("count", F.count("*").over(w) - 1).orderBy("ID")

df1.show()

---

所以，您基本上需要的是一个窗口函数

让我们从您的初始数据开始

import org.apache.spark.sql.expressions.Window
import spark.implicits._

case class Event(ID: Int, Time: Int, targetIP: String)

val events = Seq(
    Event(1, 1, "192.163.0.1"),
    Event(2, 2, "192.163.0.2"),
    Event(3, 3, "192.163.0.1"),
    Event(4, 5, "192.163.0.1"),
    Event(5, 6, "192.163.0.2"),
    Event(6, 7, "192.163.0.2"),
    Event(7, 8, "192.163.0.2")
).toDS()

现在我们需要定义一个窗口函数本身

val timeWindow = Window.orderBy($"Time").rowsBetween(-2, -1)

现在最有趣的部分是：如何在窗户上数东西？没有简单的方法，因此我们将执行以下操作

将所有targetIp聚合到列表中

筛选列表以仅查找所需的IP

计算列表的大小

结果将包含我们需要的新列“count”

UPD：

还有一个没有聚合的简短版本，由@Blackishop编写

val timeWindow = Window.partitionBy($"targetIP").orderBy($"Time").rangeBetween(-2, Window.currentRow)
val df = events
        .withColumn("count", count("*").over(timeWindow) - lit(1))
        .explain(true)

---

请注意，您应该使用

介于

之间的范围，而不是

介于

之间的行。如果向数据帧添加一行，比如说

（8,11,192.163.0.2

），会怎么样？您的解决方案将给出

count=2

，它应该是0，因为在时间单元11-2没有连接。

val timeWindow = Window.partitionBy($"targetIP").orderBy($"Time").rangeBetween(-2, Window.currentRow)
val df = events
        .withColumn("count", count("*").over(timeWindow) - lit(1))
        .explain(true)