PySpark 中 pandas.cut() 的等价物是什么？

Question

pandas.cut() 用于将值分箱为离散区间。例如，

pd.cut(
    np.array([0.2, 0.25, 0.36, 0.55, 0.67, 0.78]), 
    3, 
    include_lowest=True, 
    right=False
)

Out[9]: 
[[0.2, 0.393), [0.2, 0.393), [0.2, 0.393), [0.393, 0.587), [0.587, 0.781), [0.587, 0.781)]
Categories (3, interval[float64]): [[0.2, 0.393) < [0.393, 0.587) < [0.587, 0.781)]

如何在 PySpark 中实现同样的效果？我查看了QuantileDiscretizer，但它绝对不等同于pd.cut()，因为它不返回间隔。

Answer 1

RDD.histogram 是 Spark 中类似的函数。

假设数据包含在具有col1 列的数据框中。

+----+
|col1|
+----+
| 0.2|
|0.25|
|0.36|
|0.55|
|0.67|
|0.78|
+----+

h = df.rdd.flatMap(lambda x: x).histogram(3) #change 3 to the number of expected intervals

bins = [ (x, h[0][i+1]) for i,x in enumerate(h[0][:-1])]

def label(bin):
    return f"'{bin[0]:5.2f} - {bin[1]:5.2f}'"

e = "case "
for bin in bins[:-1]:
    e += f"when col1 >= {bin[0]} and col1 < {bin[1]} then {label(bin)} "
e += f"else {label(bins[-1])} end as bin"

df.selectExpr("col1", e).show()

输出：

+----+-------------+
|col1|          bin|
+----+-------------+
| 0.2| 0.20 -  0.39|
|0.25| 0.20 -  0.39|
|0.36| 0.20 -  0.39|
|0.55| 0.39 -  0.59|
|0.67| 0.59 -  0.78|
|0.78| 0.59 -  0.78|
+----+-------------+

bins 包含作为元组的区间：

[(0.2, 0.39333333333333337),
 (0.39333333333333337, 0.5866666666666667),
 (0.5866666666666667, 0.78)]