Динамическая подстрока без udf - PullRequest
Новая
       16

Динамическая подстрока без udf

0 голосов
/ 10 октября 2019

У меня есть фрейм данных, который выглядит следующим образом 

+------------------+
|            domain|
+------------------+
|      domain.co.uk|
|    somedomain.net|
|someotherdomain.fi|
+------------------+

Теперь я хочу проверить окончание домена, посмотреть его в словаре, где расположены домены, и добавить эту информацию в дополнительный столбец. Мой словарь выглядел бы так в следующем простом случае: 

d = {
     'eu': ['fi', 'uk'],
     'us': ['net']
    }

, и, следовательно, тот кадр данных, который я хотел бы иметь, должен выглядеть следующим образом: 

+------------------+--------+
|            domain|location|
+------------------+--------+
|      domain.co.uk|      eu|
|    somedomain.net|      us|
|someotherdomain.fi|      eu|
+------------------+--------+

Сначала я попытался получитьподстрока домена после последней точки, в python это было бы что-то вроде 

domain[domain.rfind('.')+1:]

, но в pyspark я не знаю, как это сделать динамически, за исключением использования udf. Однако, если я реализую простой udf, такой как 

udf1 = udf(lambda x: x[1:], StringType())
df = df.withColumn('xxx', udf1(domain))

, он зависает (24 часа без признаков прогресса, хотя обрабатывает только .00001% данных). Я понимаю, что UDF стоят дорого. Могу ли я сделать это без использования udf?

1 Ответ

1 голос
/ 10 октября 2019

Просто используйте substring_index и конвертируйте dict в отображение: 

from pyspark.sql.functions import create_map, lit, substring_index
from itertools import chain

# convert dict into a list of tuple
d_new = [ (k,v) for v,ks in d.items() for k in ks ]
# [('fi', 'eu'), ('uk', 'eu'), ('net', 'us')]

# create mapping from the list of tuple 
mapping = create_map([lit(i) for i in chain.from_iterable(d_new)])
# Column<b'map(fi, eu, uk, eu, net, us)'>

df_new = df.withColumn('location', mapping[substring_index('domain', '.', -1)])
...