python - pyspark 在分组的 applyInPandas 中添加多列(更改架构)

Question

我在 Pandas 中有一个简单的函数来处理时间序列。我的问题是:我想应用它的时间序列的数量非常大。所以我想使用 pySpark 来扩展它。
我面临两个问题:

模式必须显式传递，这可以使隐式传递更顺畅吗？

代码失败:Number of columns of the returned pandas.DataFrame doesn't match specified schema. Expected: 2 Actual: 3 -- 如何确保模式自动匹配？注意:我不想手动指定它(理想情况下)。在最坏的情况下 - 我可以将我的函数应用于单个时间序列并将 Pandas 数据帧 dtypes 的输出转换为 Spark 预期模式吗？

Pandas 功能

import pandas as pd
from pandas import Timestamp

df = pd.DataFrame({'time':['2020-01-01 00:00', '2020-01-01 03:00', '2020-01-01 04:00', '2020-01-06 00:00'], 'category':['1','1','1','1'], 'value':[5, 8, 7, 2]})
df['time'] = pd.to_datetime(df['time'])
display(df)
print(df.shape)

def my_complex_function(pdf):
    # fill missing hours
    pdf = pdf.set_index('time').resample('H').mean().fillna(0)
    
    # some complex computation/business logic which is adding columns/simplified here:
    pdf['foo'] = 'bar'
    pdf['baz'] = 123
    return pdf

print(df.groupby(['category']).apply(my_complex_function).reset_index().shape)

pyspark 函数

NOTICE: the version of Spark is: v3.0.1

from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("anomlydetection").master("local[4]").config("spark.driver.memory", "2G").getOrCreate()
sdf = spark.createDataFrame(df)
sdf.printSchema()

def my_complex_function_spark(pdf: pd.DataFrame)-> pd.DataFrame:
    # fill missing hours
    pdf = pdf.set_index('time').resample('H').mean().fillna(0)
    
    # some complex computation/business logic which is adding columns/simplified here:
    pdf['foo'] = 'bar'
    pdf['baz'] = 123
    return pdf

# 1) can I somehow implicitly get a reference to the schema?
# 2) this function fails due to schema mismatch
sdf.groupby("category").applyInPandas(my_complex_function_spark, schema=df.schema).show()

Answer 1

为 ApplyInPandas 定义 return_schema ，更健壮的方法是使用 p.s.t.StructType.fromJson方法和 df.schema.jsonValue方法可以保留所有现有的列属性 ( 可为空的 、 元数据 等)来自原始列。 ( 注意: 不匹配的 nullable 设置往往会导致不容易检测到的错误)
一些例子:

使用 StructType.add 追加新列方法，这是最常见的用例( example ):

from pyspark.sql.types import StructType

return_schema = StructType.fromJson(df.schema.jsonValue()) \
    .add('foo', 'string', False, "dummu string field") \
    .add('bar', 'integer')

删除现有列并附加新列:

return_schema = StructType.fromJson(df.drop("col1", "col2").schema.jsonValue()) \
    .add('foo', 'string')

设置与新列混合的列的任意顺序(仅适用于主要数据类型，注意数组、映射、结构和嵌套数据类型):

return_schema = df.selectExpr("col1", "cast('1' as string) as foo", "nested_col", "int(NULL) as bar").schema

新列通知，nullable可以通过指定 来推断空 或 非空 zero_value，见下面的例子:

cast('1' as string) as foo -> nullable is False
string(NULL) as foo        -> nullable is True

int(123) as bar            -> nullable is False
cast(NULL as int) as bar   -> nullable is True

所以如果 metadata不需要所有新列的属性，这种方法很好并且提供了灵活性。一个特殊情况是当 return_schema 不变时，我们可以使用 return_schema = df.schema .上面没有元数据的第一个示例的架构可以是:

return_schema = df.selectExpr("*", "string('') as foo", "int(NULL) as bar").schema

对于复杂的数据类型，下面是一个带有 nullable 的新 StructType 列的示例。环境:

# specify `nullable` attribute of the StructField
my_structs = StructType().add('f1', 'int', False).add('f2', 'string')

# specify `nullable=False` for the StructType
return_schema = df.select("*", F.struct(F.lit(0),F.lit('1')).cast(my_structs).alias('s1')).schema
# specify `nullable=True` for the StructType
return_schema = df.select("*", F.lit(None).cast(my_structs).alias('s2')).schema

我的建议:无需手动指定繁琐且容易出错的 return_schema，尽可能多地使用现有架构中的信息并执行 不是 完全依赖动态推理。