Question

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

from pyspark.sql.types import *

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

因为我将返回100多个具有自动生成名称的列。是否有任何方法告诉PySpark隐式使用我的函数返回的Schema，并假定所有工作节点的Schema都相同？在运行期间，该架构也会发生变化，因为我必须尝试使用我要使用的预测变量，因此可以选择自动生成架构...

Answer 1

基于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

我要做的是获取一个示例df，并将其传递给函数，然后查看返回的结果：

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

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

Answer 2

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

如果内部流程某种程度上是基于代码生成的，那么最好的选择是集成逻辑和架构生成。例如

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|
# +---+---+

Answer 3

稍微修改@Thomas答案，我做了以下工作。由于df.types返回一个元组列表（至少在最新的熊猫版本中）而不是字典，因此我将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

PySpark中pandas_udf的隐式架构？

3 个答案: