python - 如何在 PySpark ML 中创建自定义 SQLTransformer 来透视数据

Question

我有一个类似于以下结构的数据框:

# Prepare training data
training = spark.createDataFrame([
    (990011, 1001, 01, "Salary", 1000, 0.0),
    (990011, 1002, 02, "POS Purchase", 50, 0.0),
    (990022, 1003, 01, "Cash Withdrawl", 500, 1.0),
    (990022, 1004, 02, "Interest Charge", 35, 1.0)
], ["customer_id", "transaction_id", "week_of_year", "category", "amount", "label"])

我能够使用 PySpark 动态地转换这些数据，这消除了每周和类别的硬编码 case 语句:

# Attempt 1
tx_pivot = training \
    .withColumn("week_of_year", sf.concat(sf.lit("T"), sf.col("week_of_year"))) \
    .groupBy("customer_id") \
    .pivot("week_of_year") \
    .sum("amount")

tx_pivot.show(20)

我想开发一个自定义 Transformer 来动态转换数据，以便我可以将此自定义 Transform 阶段合并到 Spark ML Pipeline 中。不幸的是，当前 Spark/PySpark 中的 SQLTransfomer 仅支持 SQL，例如 E.g. “从此中选择...”(请参阅​​https://github.com/apache/spark/blob/master/python/pyspark/ml/feature.py)。

任何关于如何创建自定义 Transformer 来动态转换数据的指导将不胜感激。

Answer 1

实现一个接受一个数据帧并返回另一个数据帧的自定义转换器是非常简单的。对于您的情况:

import pyspark.ml.pipeline.Transformer as Transformer

class PivotTransformer(Transformer):

    def _transform(self, data):           
        return data.withColumn("week_of_year",sf.concat(sf.lit("T"),\
                    sf.col("week_of_year"))) \
                   .groupBy("customer_id") \
                   .pivot("week_of_year") \
                   .sum("amount")