在 PySpark 中創建 function 以返回列

Question

我有一個 function 我需要提供一個 dataframe 和一個專欄。 它使用正則表達式模式提取字符串並返回一個新的 dataframe。

但是，我希望返回只是正則表達式的結果，只是列。

def extract_strings(dataframe_selected, column_selected):
    dataframe2 = dataframe_selected.withColumn("strings", F.regexp_extract(dataframe_selected[f"{column_selected}"], r"([a-zA-Z]+)", 0))
    return dataframe2

我想創建一個 function 我可以像這樣使用

df['new_column'] = extract_strings(df, 'text')

Answer 1

您可以創建一個 function，它接受一列並返回一列。 但是您需要有一個 dataframe 來放入該列，因為列不會自行“存活”，它們需要一個 dataframe 來包含它們。

Function：

from pyspark.sql import functions as F

def extract_strings(column_selected):
    return F.regexp_extract(column_selected, r"([a-zA-Z]+)", 0)

測試：

df = spark.createDataFrame([("abc1",)], ["col_name"])

df.withColumn("new_column", extract_strings("col_name")).show()
# +--------+----------+
# |col_name|new_column|
# +--------+----------+
# |    abc1|       abc|
# +--------+----------+