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

Question

使用 Pyspark 2.2

我有一个包含多列的 spark DataFrame。我需要向 UDF 输入 2 列并返回第 3 列

输入：

+-----+------+
|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_A 的子字符串，长度为col_B。

我试过了

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 定义为在应该采用 2 时采用一个输入参数。
• 其次，您不能使用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')))

但是，在这种情况下，您可以使用this post 中描述的方法在没有udf 的情况下执行此操作。

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