В Pyspark, как сопоставить столбец с другими значениями через словарь (ключом словаря является кортеж)

Question

Я пытаюсь преобразовать столбец с именем 'xyz' в кадре данных pyspark в некоторые другие значения.Например, в xyz col, если значения ячеек 1 или 3 или 5, сопоставьте их с «A» и т. Д. Я знаю, что вы могли бы просто написать dict как dict = {1: 'A', 3: 'A', 5:'A', ...} но я бы предпочел сделать это как dict = {(1,3,5): 'A', (2,4,6): 'B', ...} Я написалнекоторый код и чувствую, что я близко, но получил некоторые ошибки.

from pysapark.sql.functions import col, udf
import pyspark

df = sqlContext.table(table_a).limit(15)

dict_a = {(1,2): 'A', (3,4): 'B', (5,6): 'C'}
dict_b = {(3,6): 'A', (5,1): 'B', (2,4,7): 'C'}

mapping_lookup = {'map_a': dict_a, 'map_b': dict_b}

### udf to do map the column values with given dict name
def translate_(col):
some_dict = mapping_lookup['map_a'] 
dict_len = len(some_dict.keys())
for i in range(0, dict_len):
    if col in some_dict.keys()[i]:
return some_dict.values()[i]

###below works and return 'B' if use map_a
translate_(3)

### below does not work and gave error says "assert isinstance(col, Column), "col should be Column"" 
df1 = df.withColumn('new_col', translate_('xyz'))

Идея состоит в том, чтобы сопоставить значения столбца с другими значениями через заданные словари.поэтому моя конечная цель - дать имя dict (map_a или map_b) и имя столбца для сопоставления.см. udf, который я хочу получить ниже:

 def translate(dict_name):    
  def translate_(col):
     some_dict = mapping_lookup[dict_name] 
     dict_len = len(some_dict.keys())
     for i in range(0, dict_len):
         if col in some_dict.keys()[i]:
             kkk = some_dict.keys()[i]
     return some_dict.get(kkk)
 return udf(translate_) 

Я создал udf, который работает для pyspark, если ключи и значения словарей соответствуют карте 1 - 1, например, 1 to A, 2 to B.но не уверен, что ключи являются кортежами, как заставить его работать?

### code below works for pyspark dataframe and map the values to deisred

def translate(dict_name):    
  def translate_(col):
    some_dict = mapping_lookup[dict_name] 
    return some_dict.get(col)
return udf(translate_) 

спасибо за ваш вклад / помощь заранее!

Answer 1

понял, как показано ниже:

def a_to_b_map(dict_name):    
 def translate_(col):
     some_dict = mapping_lookup[dict_name] 
     dict_len = len(some_dict.keys())

     temp_key = (9999,)
     for i in range(0, dict_len):
         if col in some_dict.keys()[i]:
            temp_key = some_dict.keys()[i]
             break

     return some_dict.get(temp_key)
 return udf(translate_)