如何在pyspark中对一组列进行分桶？

Posted 2023-04-15

【中文标题】如何在pyspark中对一组列进行分桶？

【英文标题】：How to bucketize a group of columns in pyspark?

【发布时间】：2018-07-18 12:44:12

【问题描述】：

我正在尝试对 5k 数据集中包含“道路”一词的列进行分桶。并创建一个新的数据框。

我不知道该怎么做，这是我迄今为止尝试过的：

from pyspark.ml.feature import Bucketizer

spike_cols = [col for col in df.columns if "road" in col]

for x in spike_cols :

    bucketizer = Bucketizer(splits=[-float("inf"), 10, 100, float("inf")],
                        inputCol=x, outputCol=x + "bucket")

bucketedData = bucketizer.transform(df)

【参考方案1】：

在循环中修改df：

from pyspark.ml.feature import Bucketizer

for x in spike_cols :
    bucketizer = Bucketizer(splits=[-float("inf"), 10, 100, float("inf")],
                    inputCol=x, outputCol=x + "bucket")
    df = bucketizer.transform(df)

或使用Pipeline:

from pyspark.ml import Pipeline
from pyspark.ml.feature import Bucketizer 

model = Pipeline(stages=[
    Bucketizer(
        splits=[-float("inf"), 10, 100, float("inf")],
        inputCol=x, outputCol=x + "bucket") for x in spike_cols
]).fit(df)

model.transform(df)

【参考方案2】：

从3.0.0开始，Bucketizer 可以通过设置 inputCols 参数一次映射多个列。

所以这变得更容易了：

from pyspark.ml.feature import Bucketizer

splits = [-float("inf"), 10, 100, float("inf")]
params = [(col, col+'bucket', splits) for col in df.columns if "road" in col]
input_cols, output_cols, splits_array = zip(*params)

bucketizer = Bucketizer(inputCols=input_cols, outputCols=output_cols,
                        splitsArray=splits_array)

bucketedData = bucketizer.transform(df)