在pyspark中制作具有值范围和特定值的容器

Question

我想在pyspark dataframe列上创建具有值范围和特定值的容器

在这篇文章How to bin in PySpark?中，它解释了如何制作箱子，但它仅适用于值的范围，从0到6，从6到18，等等，我只希望有一个用于特定值的箱子，例如，我想要从0到5，5到9，10，11到20，21和更多的箱子

+-----------+
| Age | Bin |
+-----------+
|  4  |  1  |
| 10  |  3  |
|  6  |  2  |
| 40  |  5  |
|  2  |  1  |
|  8  |  2  |
+-----------+

我试过了

splits = [ 0, 5, 9, 10, 10, 11, float('Inf') ]

但它给出了一个错误

我使用了一个公差，使垃圾箱仅为10人，但我想知道是否有更好的方法

tol= 1.0e-10
splits = [ 0, 5, 9, 10, 10+tol, 11, float('Inf') ]

Answer 1

您可以手动指定用于装箱的CASE WHEN列：

from functools import reduce

splits = [ 0, 5, 9, 10, 11 ]
splits = list(enumerate(splits))
# >>> splits
# [(0, 0), (1, 5), (2, 9), (3, 10), (4, 11)]

bins = reduce(lambda c, i: c.when(F.col('Age') <= i[1], i[0]), splits, F.when(F.col('Age') < splits[0][0], None)).otherwise(splits[-1][0] + 1).alias('bins')
# >>> bins
# Column<b'CASE WHEN (Age < 0) THEN NULL WHEN (Age <= 0) THEN 0 WHEN (Age <= 5) THEN 1 WHEN (Age <= 9) THEN 2 WHEN (Age <= 10) THEN 3 WHEN (Age <= 11) THEN 4 ELSE 5 END AS `bins`'>

df = df.select('age', bins)

df.show()
+ -+  +
|age|bins|
+ -+  +
|  4|   1|
| 10|   3|
|  6|   2|
| 40|   5|
|  2|   1|
|  8|   2|
+ -+  +