如何在ML pyspark管道中添加自己的功能作为自定义阶段?

时间:2018-07-19 06:30:45

标签: python apache-spark pyspark apache-spark-sql

Florian的示例代码

-----------+-----------+-----------+
|ball_column|keep_the   |hall_column|
+-----------+-----------+-----------+
|          0|          7|         14|
|          1|          8|         15|
|          2|          9|         16|
|          3|         10|         17|
|          4|         11|         18|
|          5|         12|         19|
|          6|         13|         20|
+-----------+-----------+-----------+

代码的第一部分在禁止列表中删除列名

#first part of the code

banned_list = ["ball","fall","hall"]
condition = lambda col: any(word in col for word in banned_list)
new_df = df.drop(*filter(condition, df.columns))

因此,上面的代码应删除ball_columnhall_column

代码的第二部分存储列表中的特定列。在此示例中,我们将仅存储剩余的keep_column

bagging = 
    Bucketizer(
        splits=[-float("inf"), 10, 100, float("inf")],
        inputCol='keep_the',
        outputCol='keep_the')

现在使用管道将列装袋如下

model = Pipeline(stages=bagging).fit(df)

bucketedData = model.transform(df)

如何将代码的第一块(banned listconditionnew_df)添加到ml管道作为阶段?

1 个答案:

答案 0 :(得分:6)

我相信这可以满足您的需求。您可以创建自定义Transformer,并将其添加到Pipeline中的阶段。请注意,由于您无法访问您提到的所有变量,所以我对您的功能进行了少许更改,但是概念保持不变。

希望这会有所帮助!

import pyspark.sql.functions as F
from pyspark.ml import Pipeline, Transformer
from pyspark.ml.feature import Bucketizer
from pyspark.sql import DataFrame
from typing import Iterable
import pandas as pd

# CUSTOM TRANSFORMER ----------------------------------------------------------------
class ColumnDropper(Transformer):
    """
    A custom Transformer which drops all columns that have at least one of the
    words from the banned_list in the name.
    """

    def __init__(self, banned_list: Iterable[str]):
        super(ColumnDropper, self).__init__()
        self.banned_list = banned_list

    def _transform(self, df: DataFrame) -> DataFrame:
        df = df.drop(*[x for x in df.columns if any(y in x for y in self.banned_list)])
        return df


# SAMPLE DATA -----------------------------------------------------------------------
df = pd.DataFrame({'ball_column': [0,1,2,3,4,5,6],
                   'keep_the': [6,5,4,3,2,1,0],
                   'hall_column': [2,2,2,2,2,2,2] })
df = spark.createDataFrame(df)


# EXAMPLE 1: USE THE TRANSFORMER WITHOUT PIPELINE -----------------------------------
column_dropper = ColumnDropper(banned_list = ["ball","fall","hall"])
df_example = column_dropper.transform(df)


# EXAMPLE 2: USE THE TRANSFORMER WITH PIPELINE --------------------------------------
column_dropper = ColumnDropper(banned_list = ["ball","fall","hall"])
bagging = Bucketizer(
        splits=[-float("inf"), 3, float("inf")],
        inputCol= 'keep_the',
        outputCol="keep_the_bucket")
model = Pipeline(stages=[column_dropper,bagging]).fit(df)
bucketedData = model.transform(df)
bucketedData.show()

输出:

+--------+---------------+
|keep_the|keep_the_bucket|
+--------+---------------+
|       6|            1.0|
|       5|            1.0|
|       4|            1.0|
|       3|            1.0|
|       2|            0.0|
|       1|            0.0|
|       0|            0.0|
+--------+---------------+

此外,请注意,如果需要使用自定义方法(例如,自定义StringIndexer),则还应创建自定义Estimator

class CustomTransformer(Transformer):

    def _transform(self, df) -> DataFrame:


class CustomEstimator(Estimator):

    def _fit(self, df) -> CustomTransformer: