Spark Dataframes: добавить условный столбец в dataframe - PullRequest
Новая
       17

Spark Dataframes: добавить условный столбец в dataframe

0 голосов
/ 08 апреля 2019

Я хочу добавить условный столбец Flag к кадру данных A. Если выполнены следующие два условия, добавьте 1 к Flag, в противном случае 0:

  1. num от кадра данных A находится между numStart и numEnd от кадра данных B.

  2. Если вышеуказанное условие удовлетворяет, проверьте, равно ли include 1.

DataFrame A (это очень большой dataframe, содержащий миллионы строк): 

+----+------+-----+------------------------+
|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         |
+----+------+-----+------------------------+----------+

Ответы [ 2 ]

1 голос
/ 09 апреля 2019

Вы можете присоединиться влево dfB к dfA на основе условия, которое вы описали в (i), затем построить столбец Flag, используя withColumn и функцию coalesce для "default" в 0:

  • Записи, для которых найдено совпадение, будут использовать значение include соответствующей dfB записи
  • Записи, для которых не было совпадения, имели бы include=null, и по вашему требованию такие записи должны получить Flag=0, поэтому мы используем coalesce, который в случае нуля возвращает значение по умолчанию с литералом 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|
// +----+------+-----+--------------------+----+
0 голосов
/ 09 апреля 2019

Соедините два кадра данных вместе при первом условии, сохраняя все строки в кадре данных 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))
...