Spark Dataframes：将条件列添加到数据框

Question

我想将条件列Flag添加到数据帧A.当满足以下两个条件时，将1添加到Flag，否则为0：

1. 来自数据帧A的num位于数据帧B的numStart和numEnd之间。
2. 如果满足上述条件，请检查include是否为1。

DataFrame A（它是一个非常大的数据帧，包含数百万行）：

+----+------+-----+------------------------+
|num |food  |price|timestamp               |
+----+------+-----+------------------------+
|1275|tomato|1.99 |2018-07-21T00:00:00.683Z|
|145 |carrot|0.45 |2018-07-21T00:00:03.346Z|
|2678|apple |0.99 |2018-07-21T01:00:05.731Z|
|6578|banana|1.29 |2018-07-20T01:11:59.957Z|
|1001|taco  |2.59 |2018-07-21T01:00:07.961Z|
+----+------+-----+------------------------+

DataFrame B（它是一个非常小的DF，只包含100行）：

+----------+-----------+-------+
|numStart  |numEnd     |include|
+----------+-----------+-------+
|0         |200        |1      |
|250       |1050       |0      |
|2000      |3000       |1      |
|10001     |15001      |1      |
+----------+-----------+-------+

预期产量：

+----+------+-----+------------------------+----------+
|num |food  |price|timestamp               |Flag      |
+----+------+-----+------------------------+----------+
|1275|tomato|1.99 |2018-07-21T00:00:00.683Z|0         |
|145 |carrot|0.45 |2018-07-21T00:00:03.346Z|1         |
|2678|apple |0.99 |2018-07-21T01:00:05.731Z|1         |
|6578|banana|1.29 |2018-07-20T01:11:59.957Z|0         |
|1001|taco  |2.59 |2018-07-21T01:00:07.961Z|0         |
+----+------+-----+------------------------+----------+

Answer 1

你可以根据你在（i）中描述的条件将dfB左连接到dfA，然后使用Flag和withColumn函数构建一个coalesce列，使其“默认”为0：

• 找到匹配项的记录将使用匹配的include记录的dfB值
• 没有匹配的记录会有include=null，根据你的要求，这样的记录应该得到Flag=0，所以我们使用coalesce，在null的情况下返回默认值和文字lit(0)

最后，摆脱对你不感兴趣的dfB专栏：

import org.apache.spark.sql.functions._
import spark.implicits._ // assuming "spark" is your SparkSession

dfA.join(dfB, $"num".between($"numStart", $"numEnd"), "left")
  .withColumn("Flag", coalesce($"include", lit(0)))
  .drop(dfB.columns: _*)
  .show()

// +----+------+-----+--------------------+----+
// | num|  food|price|           timestamp|Flag|
// +----+------+-----+--------------------+----+
// |1275|tomato| 1.99|2018-07-21T00:00:...|   0|
// | 145|carrot| 0.45|2018-07-21T00:00:...|   1|
// |2678| apple| 0.99|2018-07-21T01:00:...|   1|
// |6578|banana| 1.29|2018-07-20T01:11:...|   0|
// |1001|  taco| 2.59|2018-07-21T01:00:...|   0|
// +----+------+-----+--------------------+----+

Answer 2

在第一个条件下将两个数据帧连接在一起，同时保留数据帧A中的所有行（即使用左连接，请参阅下面的代码）。在连接之后，include列可以重命名为Flag，其中的任何NaN值都设置为0.删除了两个额外的列numStart和numEnd。

因此代码可以写成如下：

A.join(B, $"num" >= $"numStart" && $"num" <= $"numEnd", "left")
  .withColumnRenamed("include", "Flag")
  .drop("numStart", "numEnd")
  .na.fill(Map("Flag" -> 0))