PySpark 中 pandas_udf 的隐式模式？

Question

This answer 很好地解释了如何使用 pyspark 的 groupby 和 pandas_udf 进行自定义聚合。但是，我不可能像示例的这一部分所示手动声明我的架构

from pyspark.sql.types import *

schema = StructType([
    StructField("key", StringType()),
    StructField("avg_min", DoubleType())
])

因为我将返回 100 多列自动生成的名称。有没有办法告诉 PySpark 只是隐含地使用我的函数返回的 Schema 并假设它对于所有工作节点都是相同的？这个模式也会在运行过程中发生变化，因为我将不得不使用我想要使用的预测器，所以一个用于模式生成的自动化过程可能是一个选项......

Answer 1

稍微修改@Thomas 回答我做了以下。由于df.types 返回的是元组列表（至少在最新的pandas 版本中）而不是字典，所以我将str(dfp.dtypes[key]) 替换为dict(df.dtypes)[key]

def udf_schema_from_pandas(df):
  column_types  = [StructField(key, mapping[dict(df.dtypes)[key]]()) for key in df.columns]
  schema = StructType(column_types)
  return schema

Answer 2

基于Sanxofons comment,，我对如何自己实现这一点有了一个想法：

from pyspark.sql.types import *

mapping = {"float64": DoubleType,
           "object":StringType,
           "int64":IntegerType} # Incomplete - extend with your types.

def createUDFSchemaFromPandas(dfp):
  column_types  = [StructField(key, mapping[str(dfp.dtypes[key])]()) for key in dfp.columns]
  schema = StructType(column_types)
  return schema

我所做的是获取一个示例 pandas df，将其传递给函数，然后查看返回什么：

dfp = df_total.limit(100).toPandas()
df_return = my_UDF_function(dfp)
schema = createUDFSchemaFromPandas(df_return)

这似乎对我有用。问题在于它是一种递归（需要定义函数来获取模式，将模式定义为 udf）。我通过创建一个简单地传递数据帧的“包装器”UDF 解决了这个问题。

Answer 3

不幸的是，没有这样的选择。在评估任何组件之前，必须静态地知道模式，因此任何基于实际数据的形式推断都根本不在讨论范围内。

如果内部流程以某种方式基于代码生成，那么您最好的选择是集成逻辑和模式生成。例如

def describe(cols, fun):
    schema = StructType([StructField(c, DoubleType()) for c in cols])
    @pandas_udf(schema, PandasUDFType, PandasUDFType.GROUPED_MAP)
    def _(df):
        return df[cols].agg([fun])
    return _

df = spark.createDataFrame([(1, 2.0, 1.0, 3.0), (1, 4.0, 2.0, 5.0)], ("id", "x", "y", "z"))

df.groupBy("id").apply(describe(["x"], "mean")).show()                                         
# +---+                                                                           
# |  x|
# +---+
# |3.0|
#+---+


df.groupBy("id").apply(describe(["x", "y"], "mean")).show()                                    
# +---+---+                                                                       
# |  x|  y|
# +---+---+
# |3.0|1.5|
# +---+---+