Pyspark: Reemplazando valor en una columna buscando en un diccionario con valores de .

Tengo una situación en la que tengo un diccionario de elementos en PySpark como este:

swap={‘A’: 0.07677341668184234,
: 0.1497896460766734,
‘B’: 0.07186667210628232}

Tengan en cuenta el objeto “pandas.NA” definido como una de las claves.

También tengo una tabla de pandas configurada con varios valores que pueden estar o no en el diccionario “swap” anterior:

índice columna
1 C
2 B
3
4 A

Por otro código que encontré en stackoverflow aquí: https://stackoverflow.com/questions/43976237/pyspark-replacing-value-in-a-column-by-searching-a-dictionary, he estado usando la siguiente función para intercambiar la columna anterior:

def recode(col_name, map_dict, default=None):
if not isinstance(col_name, Column): # Permite pasar tanto el nombre de la columna como la instancia de la columna
col_name = col(col_name)
mapping_expr = create_map([lit(x) for x in chain(*map_dict.items())])
if default is None:
return mapping_expr.getItem(col_name)
else:
return when(~isnull(mapping_expr.getItem(col_name)), mapping_expr.getItem(col_name)).otherwise(default)

y el siguiente comando:

data=df.withColumn(“column”, recode(‘column’, swap, default=0))

donde ‘df’ sería el marco de datos de ejemplo anterior. La salida esperada debería ser:

índice columna
1 0
2 0.07186667210628232
3 0.1497896460766734
4 0.07677341668184234

Sin embargo, obtengo el error:

AttributeError: ‘NAType’ object has no attribute ‘_get_object_id’

Este error es causado por el objeto “pandas.NA” en el diccionario “swap”. ¿Cómo puedo hacer que este código funcione como se espera y deje de fallar?