PySparkでビンする方法は？

Question

たとえば、人のDataFrameを年齢に応じて以下の4つのビンに分類したいと思います。私はpandas.cut()を使ってこれをパンダで行います。 PySparkでこれをどうやって行うのですか？

age_bins = [ 0, 6, 18, 60, np.Inf ] 
age_labels = [ 'infant', 'minor', 'adult', 'senior' ] 

Answer 1

sparkのmlライブラリからのBucketizer機能transfromを使用できます。

values = [("a", 23), ("b", 45), ("c", 10), ("d", 60), ("e", 56), ("f", 2), ("g", 25), ("h", 40), ("j", 33)] 


df = spark.createDataFrame(values, ["name", "ages"]) 


from pyspark.ml.feature import Bucketizer 
bucketizer = Bucketizer(splits=[ 0, 6, 18, 60, float('Inf') ],inputCol="ages", outputCol="buckets") 
df_buck = bucketizer.setHandleInvalid("keep").transform(df) 

df_buck.show() 

出力

+----+----+-------+ 
|name|ages|buckets| 
+----+----+-------+ 
| a| 23| 2.0| 
| b| 45| 2.0| 
| c| 10| 1.0| 
| d| 60| 3.0| 
| e| 56| 2.0| 
| f| 2| 0.0| 
| g| 25| 2.0| 
| h| 40| 2.0| 
| j| 33| 2.0| 
+----+----+-------+ 

あなたは、各バケットの名前をしたい場合は、バケット名で新しい列を作成するために、UDFを使用することができます

from pyspark.sql.functions import udf 
from pyspark.sql.types import * 

t = {0.0:"infant", 1.0: "minor", 2.0:"adult", 3.0: "senior"} 
udf_foo = udf(lambda x: t[x], StringType()) 
df_buck.withColumn("age_bucket", udf_foo("buckets")).show() 

出力

+----+----+-------+----------+ 
|name|ages|buckets|age_bucket| 
+----+----+-------+----------+ 
| a| 23| 2.0|  adult| 
| b| 45| 2.0|  adult| 
| c| 10| 1.0|  minor| 
| d| 60| 3.0| senior| 
| e| 56| 2.0|  adult| 
| f| 2| 0.0| infant| 
| g| 25| 2.0|  adult| 
| h| 40| 2.0|  adult| 
| j| 33| 2.0|  adult| 
+----+----+-------+----------+