Skip to main content
 首页 » 编程设计

python之将函数应用于 Spark DataFrame 中的所有单元格

2025年05月04日29bluestorm

我正在尝试将一些 Pandas 代码转换为 Spark 以进行缩放。 myfunc 是一个复杂 API 的包装器,它接受一个字符串并返回一个新字符串(这意味着我不能使用矢量化函数)。

def myfunc(ds): 
    for attribute, value in ds.items(): 
        value = api_function(attribute, value) 
        ds[attribute] = value 
    return ds 
 
df = df.apply(myfunc, axis='columns') 

myfunc 获取一个 DataSeries,将其分解为单独的单元格,为每个单元格调用 API,并构建一个具有相同列名的新 D​​ataSeries。这有效地修改了 DataFrame 中的所有单元格。

我是 Spark 的新手,我想使用 pyspark 翻译这个逻辑。我已将我的 pandas DataFrame 转换为 Spark:

spark = SparkSession.builder.appName('My app').getOrCreate() 
spark_schema = StructType([StructField(c, StringType(), True) for c in df.columns]) 
spark_df = spark.createDataFrame(df, schema=spark_schema) 

这是我迷路的地方。我需要一个 UDF,一个 pandas_udf 吗?如何遍历所有单元格并使用 myfunc 为每个单元格返回一个新字符串? spark_df.foreach() 不返回任何内容,也没有 map() 函数。

我可以将 myfunc 从 DataSeries -> DataSeries 修改为 string -> string 如有必要。

请您参考如下方法:

选项 1:一次在一列上使用 UDF

最简单的方法是重写您的函数以将字符串作为参数(因此它是 string -> string)并使用 UDF。有一个很好的例子 here .这一次适用于一列。因此,如果您的 DataFrame 具有合理数量的列,您可以一次将 UDF 应用于每一列:

from pyspark.sql.functions import col 
new_df = df.select(udf(col("col1")), udf(col("col2")), ...) 

例子

df = sc.parallelize([[1, 4], [2,5], [3,6]]).toDF(["col1", "col2"]) 
df.show() 
+----+----+ 
|col1|col2| 
+----+----+ 
|   1|   4| 
|   2|   5| 
|   3|   6| 
+----+----+ 
 
def plus1_udf(x): 
    return x + 1 
plus1 = spark.udf.register("plus1", plus1_udf) 
 
new_df = df.select(plus1(col("col1")), plus1(col("col2"))) 
new_df.show() 
+-----------+-----------+ 
|plus1(col1)|plus1(col2)| 
+-----------+-----------+ 
|          2|          5| 
|          3|          6| 
|          4|          7| 
+-----------+-----------+ 
 

选项 2:一次映射整个 DataFrame

map 可用于 Scala DataFrame,但目前在 PySpark 中不可用。 下级RDD API 在 PySpark 中确实有一个 map 函数。因此,如果您有太多列而无法一次转换一个,您可以像这样对 DataFrame 中的每个单元格进行操作:

def map_fn(row): 
    return [api_function(x) for (column, x) in row.asDict().items() 
 
column_names = df.columns 
new_df = df.rdd.map(map_fn).toDF(df.columns) 

例子

df = sc.parallelize([[1, 4], [2,5], [3,6]]).toDF(["col1", "col2"]) 
def map_fn(row): 
   return [value + 1 for (_, value) in row.asDict().items()] 
 
columns = df.columns 
new_df = df.rdd.map(map_fn).toDF(columns) 
new_df.show() 
+----+----+ 
|col1|col2| 
+----+----+ 
|   2|   5| 
|   3|   6| 
|   4|   7| 
+----+----+ 
 

上下文

documentation foreach 只给出了打印的例子,但是我们可以通过查看 code 来验证它确实不返回任何东西。

您可以在 this post 中阅读有关 pandas_udf 的信息,但它似乎最适合向量化函数,正如您所指出的,由于 api_function 而无法使用。