Usando Spark 1.6, tengo una DataFrame column (llamada digamos col1 ) con valores A, B, C, DS, DNS, E, F, G y H y quiero crear una nueva columna (digamos col2 ) con los valores del dict aquí abajo, ¿cómo mapeo esto? (así que fi 'A' debe asignarse a 'S', etc.)
dict = {'A': 'S', 'B': 'S', 'C': 'S', 'DS': 'S', 'DNS': 'S', 'E': 'NS', 'F': 'NS', 'G': 'NS', 'H': 'NS'}Solución ineficiente con UDF (independiente de la versión):
from pyspark.sql.types import StringType from pyspark.sql.functions import udf def translate(mapping): def translate_(col): return mapping.get(col) return udf(translate_, StringType()) df = sc.parallelize([('DS', ), ('G', ), ('INVALID', )]).toDF(['key']) mapping = { 'A': 'S', 'B': 'S', 'C': 'S', 'DS': 'S', 'DNS': 'S', 'E': 'NS', 'F': 'NS', 'G': 'NS', 'H': 'NS'} df.withColumn("value", translate(mapping)("key"))con el resultado:
+-------+-----+ | key|value| +-------+-----+ | DS| S| | G| NS| |INVALID| null| +-------+-----+ Mucho más eficiente ( Spark >= 2.0, Spark < 3.0 ) es crear un literal MapType :
from pyspark.sql.functions import col, create_map, lit from itertools import chain mapping_expr = create_map([lit(x) for x in chain(*mapping.items())]) df.withColumn("value", mapping_expr.getItem(col("key")))con el mismo resultado:
+-------+-----+ | key|value| +-------+-----+ | DS| S| | G| NS| |INVALID| null| +-------+-----+pero plan de ejecución más eficiente:
== Physical Plan == *Project [key#15, keys: [B,DNS,DS,F,E,H,C,G,A], values: [S,S,S,NS,NS,NS,S,NS,S][key#15] AS value#53] +- Scan ExistingRDD[key#15]en comparación con la versión UDF:
== Physical Plan == *Project [key#15, pythonUDF0#61 AS value#57] +- BatchEvalPython [translate_(key#15)], [key#15, pythonUDF0#61] +- Scan ExistingRDD[key#15] En Spark >= 3.0 getItem debe reemplazarse con __getitem__ ( [] ), es decir:
df.withColumn("value", mapping_expr[col("key")]).show()Parece que la solución más simple sería usar la función de reemplazo: http://spark.apache.org/docs/2.4.0/api/python/pyspark.sql.html#pyspark.sql.DataFrame.replace
mapping= { 'A': '1', 'B': '2' } df2 = df.replace(to_replace=mapping, subset=['yourColName'])Si desea crear una columna de mapa a partir de un diccionario anidado, puede usar esto:
def create_map(d,): if type(d) != dict: return F.lit(d) level_map = [] for k in d: level_map.append(F.lit(k)) level_map.append(create_map(d[k])) return F.create_map(level_map) d = {'a': 1, 'b': {'c': 2, 'd': 'blah'}} print(create_map(d)) # <- Column<b'map(a, 1, b, map(c, 2, d, blah))'>