Pyspark SQL Pandas 分組 Map 無分組

Question

I have a dataset that I want to map over using several Pyspark SQL Grouped Map UDFs, at different stages of a larger ETL process that runs on ephemeral clusters in AWS EMR. 分組的 Map API 要求 Pyspark Z6A8064B5DF479455500553C47C5505 實際已分組，但實際上不需要分組。

目前，我正在使用任意分組，它有效，但結果是：

1. 不必要的洗牌。

2. 每個作業中任意 groupby 的 Hacky 代碼。

我的理想解決方案允許矢量化 Pandas UDF 在沒有任意分組的情況下應用，但如果我可以保存至少可以消除隨機播放的任意分組。

編輯：

這是我的代碼的樣子。 我最初使用的是任意分組，但目前正在根據@pault 下面的評論嘗試spark_partition_id() 。

@pandas_udf(b_schema, PandasUDFType.GROUPED_MAP)
def transform(a_partition):
  b = a_partition.drop("pid", axis=1)
  # Some other transform stuff
  return b

(sql
  .read.parquet(a_path)
  .withColumn("pid", spark_partition_id())
  .groupBy("pid")
  .apply(transform)
  .write.parquet(b_path))

使用spark_partition_id()似乎仍然會導致洗牌。 我得到以下 DAG：

階段1

1. 掃描鑲木地板
2. 項目
3. 項目
4. 交換

第二階段

1. 交換
2. 種類
3. FlatMapGroupsInPandas

Answer 1

要支持大致等效的邏輯（函數(pandas.core.frame.DataFrame) -> pandas.core.frame.DataFrame ），您必須切換到 Spark 3.0.0 並使用MAP_ITER轉換。

在最新的預覽版 (3.0.0-preview2) 中，您需要一個 UDF：

@pandas_udf(b_schema, PandasUDFType.MAP_ITER)
def transform(dfs):
    for df in dfs:
        b = df.drop("pid", axis=1)
        ...
        yield b

df.mapInPandas(transform)

在即將發布的 3.0.0 版本（ SPARK-28264 ）中，只是一個普通的 function：

def transform(dfs):
    for df in dfs:
        b = df.drop("pid", axis=1)
        # Some other transform stuff
        ...
        yield b

df.mapInPandas(transform, b_schema)

2.x 上可能的解決方法是使用普通的SCALAR UDF，將結果的每一行序列化為 JSON，然后在另一側反序列化，即

import json
from pyspark.sql.functions import from_json

@pandas_udf("string", PandasUDFType.SCALAR)
def transform(col1, col2):
    b = pd.DataFrame({"x": col1, "y": col2})
    ...
    return b.apply(lambda x: json.dumps(dict(zip(df.columns, x))), axis=1)


(df
    .withColumn("json_result", transform("col1", "col2"))
    .withColumn("a_struct", from_json("json_result", b_schema)))