python - pyspark 中的分桶和一种热编码

Question

我有一个包含以下列的 PySpark 数据框:

id        Age
1         30
2         25
3         21

我有以下年龄段:[20, 24, 27, 30]。

我的预期结果:

id    Age    age_bucket     age_27_30     age_24_27   age_20_24
1     30      (27-30]           1            0           0
2     25      (24-27]           0            1           0
3     21      (20-24]           0            0           1

我当前的代码:

from pyspark.ml.feature import Bucketizer
bucketizer = Bucketizer(splits=[ 20,24,27,30 ],inputCol="Age", outputCol="age_bucket")
df1 = bucketizer.setHandleInvalid("keep").transform(df)

Answer 1

如果您想要一个与您的问题完全相同的结果，如果没有其他一些奇特的映射技巧，OneHotEstimatorEncoder 将无法工作。

我会在这里使用连接:

age_buckets = [20, 24, 27, 30]
bins = list(zip(age_buckets, age_buckets[1:]))

data = [[i] + ['({0}-{1}]'.format(*bin_endpoints)] + [0] * i + [1] + [0] * (len(bins) - i - 1) 
        for i, bin_endpoints in enumerate(bins)]
schema = ', '.join('age_bucket_{}_{}: int'.format(start, end) 
                   for start, end in zip(age_buckets, age_buckets[1:]))

join_df = spark.createDataFrame(data, 'age_bucket: int, age_bucket_string: string, ' + schema)

result = (df1.join(join_df, on='age_bucket', how='left')
             .drop('age_bucket')
             .withColumnRenamed('age_bucket_string', 'age_bucket')
             .orderBy('id'))
result.show()

输出:

+---+---+----------+----------------+----------------+----------------+
| id|Age|age_bucket|age_bucket_20_24|age_bucket_24_27|age_bucket_27_30|
+---+---+----------+----------------+----------------+----------------+
|  1| 30|   (27-30]|               0|               0|               1|
|  2| 25|   (24-27]|               0|               1|               0|
|  3| 21|   (20-24]|               1|               0|               0|
+---+---+----------+----------------+----------------+----------------+