Сопоставить значения столбца с ArrayType на основе значений из другого фрейма данных в PySpark

Question

Что у меня есть:

| ids.   |items   |item_id|value|timestamp|
+--------+--------+-------+-----+---------+
|[A,B,C] |1.0     |1      |5    |100      | 
|[A,B,D] |1.0     |2      |6    |90       | 
|[D]     |0.0.    |3      |7    |80       |
|[C]     |0.0.    |4      |8    |80       |
+--------+--------+-------+-----+----------

| ids    |id_num  |
+--------+--------+
|A       |1       |
|B       |2       |
|C       |3       |
|D       |4       |
+---+----+--------+

Что я хочу:

| ids    |
+--------+
|[1,2,3] |      
|[1,2,4] |    
|[3]     | 
|[4]     | 
+--------+

Есть ли способ сделать это без взрыва? Спасибо за вашу помощь!

Answer 1

Вы можете использовать UDF:

from pyspark.sql.functions import udf, col
from pyspark.sql.types import ArrayType

# Suppose this is the dictionary you want to map
map_dict = {'A':1, 'B':2,'C':3,'D':4}

def array_map(array_col):
    return list(map(map_dict.get, array_col))
"""
If you prefer list comprehension, you can return [map_dict[k] for k in array_col]
"""

array_map_udf = udf(array_map, ArrayType())

df = df.withColumn("mapped_array", array_map_udf(col("ids")))

Я не могу придумать другой метод, но чтобы получить распараллеленный словарь, вы можете просто использовать метод toJSON. Это потребует дальнейшей обработки в зависимости от типа ссылки:

import json
df_json = df.toJSON().map(lambda x: json.loads(x))