使用 udf 从 python 应用函数到 Pyspark 代码的等效代码

Question

df = spark.createDataFrame(
        [
            (1, "AxtTR"),  # create your data here, be consistent in the types.
            (2, "HdyOP"),
            (3, "EqoPIC"),
            (4, "OkTEic"),
        ], ["id", "label"] )# add your column names here]
df.show()

下面的代码在 python 中，我使用 apply 函数并尝试提取每行的前 2 个字母。我想在 pyspark 中复制相同的代码。其中一个函数用于对每一行应用并获取输出。

def get_string(lst):    
    lst = str(lst)
    lst = lst.lower
    lst=  lst[0:2]
    return(lst)
df['firt_2letter'] = df['label'].apply(get_string)

如下图所示的黄色标记是预期的输出。

Answer 1

您可以使用相关的 Spark SQL 函数：

import pyspark.sql.functions as F

df2 = df.withColumn('first_2letter', F.lower('label')[0:2])

df2.show()
+---+------+-------------+
| id| label|first_2letter|
+---+------+-------------+
|  1| AxtTR|           ax|
|  2| HdyOP|           hd|
|  3|EqoPIC|           eq|
|  4|OkTEic|           ok|
+---+------+-------------+

如果要使用用户自定义函数，可以定义为：

def get_string(lst):    
    lst = str(lst)
    lst = lst.lower()
    lst = lst[0:2]
    return lst

import pyspark.sql.functions as F

df2 = df.withColumn('first_2letter', F.udf(get_string)('label'))

df2.show()
+---+------+-------------+
| id| label|first_2letter|
+---+------+-------------+
|  1| AxtTR|           ax|
|  2| HdyOP|           hd|
|  3|EqoPIC|           eq|
|  4|OkTEic|           ok|
+---+------+-------------+