Добавить новый столбец в фрейме данных Pyspark, основываясь на условии, где находится другой столбец

Question

У меня есть фрейм данных Pyspark следующим образом:

+------------+-------------+--------------------+
|package_id  | location    | package_scan_code  | 
+------------+-------------+--------------------+
|123         | Denver      |05                  |  
|123         | LosAngeles  |03                  |  
|123         | Dallas      |09                  |  
|123         | Vail        |02                  | 
|456         | Jacksonville|05                  |  
|456         | Nashville   |09                  |
|456         | Memphis     |03                  |

"package_scan_code" 03 представляет источник пакета.

Я хочу добавить столбец «origin» к этому фрейму данных таким образом, чтобы для каждого пакета (определяемого как «package_id») значения во вновь добавленном столбце источника были в том же месте, которое соответствует «package_scan_code»03.

В приведенном выше случае есть два уникальных пакета 123 и 456, и они имеют происхождение как LosAngeles и Memphis соответственно (соответствует package_scan_code 03).

Поэтому я хочу, чтобы мой вывод былбыть следующим:

+------------+-------------+--------------------+------------+
| package_id |location     | package_scan_code  |origin      |
+------------+-------------+--------------------+------------+
|123         | Denver      |05                  | LosAngeles |
|123         | LosAngeles  |03                  | LosAngeles |
|123         | Dallas      |09                  | LosAngeles |
|123         | Vail        |02                  | LosAngeles |
|456         | Jacksonville|05                  |  Memphis   |
|456         | Nashville   |09                  |  Memphis   |
|456         | Memphis     |03                  |  Memphis   |

Как мне добиться этого в Pyspark?Я попытался .withColumn метод, но я не смог получить правильное условие.

Answer 1

Этот код должен работать независимо от того, сколько раз package_scan_code=03 встречается для каждого package_id в кадре данных.Я добавил еще один (123,'LosAngeles','03'), чтобы продемонстрировать, что -

Шаг 1: Создание кадра данных

values = [(123,'Denver','05'),(123,'LosAngeles','03'),(123,'Dallas','09'),(123,'Vail','02'),(123,'LosAngeles','03'),
          (456,'Jacksonville','05'),(456,'Nashville','09'),(456,'Memphis','03')]
df = sqlContext.createDataFrame(values,['package_id','location','package_scan_code'])

Шаг 2: Создание словаряиз package_id и location.

df_count = df.where(col('package_scan_code')=='03').groupby('package_id','location').count()
dict_location_scan_code = dict(df_count.rdd.map(lambda x: (x['package_id'], x['location'])).collect())
print(dict_location_scan_code)
    {456: 'Memphis', 123: 'LosAngeles'}

Шаг 3: Создание столбца, отображение словаря.

from pyspark.sql.functions import col, create_map, lit
from itertools import chain
mapping_expr = create_map([lit(x) for x in chain(*dict_location_scan_code.items())])
df = df.withColumn('origin', mapping_expr.getItem(col('package_id')))
df.show()
+----------+------------+-----------------+----------+
|package_id|    location|package_scan_code|    origin|
+----------+------------+-----------------+----------+
|       123|      Denver|               05|LosAngeles|
|       123|  LosAngeles|               03|LosAngeles|
|       123|      Dallas|               09|LosAngeles|
|       123|        Vail|               02|LosAngeles|
|       123|  LosAngeles|               03|LosAngeles|
|       456|Jacksonville|               05|   Memphis|
|       456|   Nashville|               09|   Memphis|
|       456|     Memphis|               03|   Memphis|
+----------+------------+-----------------+----------+

Answer 2

Отфильтруйте фрейм данных по package_scan_code == '03', а затем присоединитесь к исходному фрейму данных:

(df.filter(df.package_scan_code == '03')
   .selectExpr('package_id', 'location as origin')
   .join(df, ['package_id'], how='right')
   .show())
+----------+----------+------------+-----------------+
|package_id|    origin|    location|package_scan_code|
+----------+----------+------------+-----------------+
|       123|LosAngeles|      Denver|               05|
|       123|LosAngeles|  LosAngeles|               03|
|       123|LosAngeles|      Dallas|               09|
|       123|LosAngeles|        Vail|               02|
|       456|   Memphis|Jacksonville|               05|
|       456|   Memphis|   Nashville|               09|
|       456|   Memphis|     Memphis|               03|
+----------+----------+------------+-----------------+

Примечание. Предполагается, что у вас есть максимум один package_scan_code, равный 03 на * 1007.* в противном случае логика не будет правильной, и вам нужно переосмыслить способ определения origin.