How to bin in PySpark?

Question

For example, I'd like to classify a DataFrame of people into the following 4 bins according to age.

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

I would use pandas.cut() to do this in pandas. How do I do this in PySpark?

Answer 1

You can use Bucketizer feature transfrom from ml library in spark.

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() 

output

+----+----+-------+ |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| +----+----+-------+ 

If you want names for each bucket you can use udf to create a new column with bucket names

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() 

output

+----+----+-------+----------+ |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| +----+----+-------+----------+