Pyspark 不带udf的动态子字符串

我有这样的数据框

+------------------+
|            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小时，没有任何进展迹象，尽管只处理了0.00001%的数据。我知道udf很贵。我可以不使用udf执行此操作吗？

只需使用子字符串索引并将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)])

---

检查子字符串_索引'domain'，'-1@jxc谢谢所以我很快就得到了这个域名。但是如果我想查字典呢？我还需要自定义项吗？只需从dict创建映射，查看我的帖子