基于另一列的长度，一列的pyspark子字符串

Question

使用pyspark 2.2

我有一个带有多列的火花数据框。我需要将2列输入UDF，然后返回第三列

输入：

+-----+------+
|col_A| col_B|
+-----+------+
|  abc|abcdef|
|  abc|     a|
+-----+------+

col_A和col_B都是StringType()

所需的输出：

+-----+------+-------+
|col_A| col_B|new_col|
+-----+------+-------+
|  abc|abcdef|    abc|
|  abc|     a|      a|
+-----+------+-------+

我希望new_col是col_B的col_A的子字符串。

我尝试了

udf_substring = F.udf(lambda x: F.substring(x[0],0,F.length(x[1])), StringType())
df.withColumn('new_col', udf_substring([F.col('col_A'),F.col('col_B')])).show()

，但它给出了TypeError: Column is not iterable。

任何想法如何进行此类操作？

Answer 1

这里有两个主要的事情。

• 首先，您定义了 udf以进行一个输入参数。
• 其次，您不能在udf中使用API功能。（调用udf序列化为Python，因此您需要使用Python语法和函数。）

这是此问题的适当udf实现：

import pyspark.sql.functions as F
def my_substring(a, b):
    # You should add in your own error checking
    return a[:len(b)]
udf_substring = F.udf(lambda x, y: my_substring(a, b), StringType())

，然后通过将两列以参数传递给它来调用：

df.withColumn('new_col', udf_substring(F.col('col_A'),F.col('col_B')))

但是，在这种情况下，您可以使用本文中描述的方法无需udf执行此操作。

df.withColumn(
    'new_col', 
    F.expr("substring(col_A,0,length(col_B))")
)