Spark DataFrame中将年龄列不同数值区间映射为整数的实现方法

数据老陈

2026-5-28

解决方案

针对你提出的Spark DataFrame年龄映射需求，我整理了两种实用的实现方式，能完美满足列类型转换、区间映射和列名重命名的要求，一起来看看吧：

方式一：使用`when()`条件函数（直观易读，适合少量区间）

这种方式通过Spark内置的条件判断函数直接处理，逻辑清晰，容易理解和修改。

步骤1：导入必要的依赖

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when
from pyspark.sql.types import IntegerType

步骤2：初始化SparkSession并创建示例数据

spark = SparkSession.builder.appName("AgeCategoryMapping").getOrCreate()

# 模拟你的原始DataFrame结构
raw_data = [("A", "25"), ("B", "17"), ("C", "8")]
df = spark.createDataFrame(raw_data, ["customer", "age"])

# 查看原始数据结构，确认age是string类型
df.printSchema()
# root
#  |-- customer: string (nullable = true)
#  |-- age: string (nullable = true)

步骤3：实现映射逻辑并生成目标DataFrame

# 先将string类型的age转成整数，再按区间映射，最后重命名列并指定输出顺序
result_df = df.withColumn(
    "age_new",
    # 按区间匹配映射值，同时确保最终类型是Integer
    when((col("age").cast(IntegerType()) >= 0) & (col("age").cast(IntegerType()) <= 10), 0)
    .when((col("age").cast(IntegerType()) >= 11) & (col("age").cast(IntegerType()) <= 20), 1)
    .when((col("age").cast(IntegerType()) >= 21) & (col("age").cast(IntegerType()) <= 30), 2)
    .cast(IntegerType())
).withColumnRenamed("customer", "customer_new") \
 .select("customer_new", "age_new")

# 查看结果
result_df.show()
# +-------------+-------+
# |customer_new|age_new|
# +-------------+-------+
# |            A|      2|
# |            B|      1|
# |            C|      0|
# +-------------+-------+

# 验证结果结构
result_df.printSchema()
# root
#  |-- customer_new: string (nullable = true)
#  |-- age_new: integer (nullable = true)

方式二：使用`Bucketizer`分桶工具（适合多区间场景）

如果后续需要扩展更多年龄区间，使用Spark ML的Bucketizer会更简洁，减少重复代码。

from pyspark.ml.feature import Bucketizer

# 定义分桶边界：注意Bucketizer是左闭右开规则，所以要匹配你的区间需求
splits = [-float("inf"), 10.0, 20.0, 30.0, float("inf")]
bucketizer = Bucketizer(splits=splits, inputCol="age_int", outputCol="age_new")

# 先将string类型的age转成整数
df_with_int_age = df.withColumn("age_int", col("age").cast(IntegerType()))

# 执行分桶映射
bucketed_df = bucketizer.transform(df_with_int_age)

# 调整列名和输出结构
result_df_bucket = bucketed_df.withColumnRenamed("customer", "customer_new") \
                              .select("customer_new", "age_new") \
                              .withColumn("age_new", col("age_new").cast(IntegerType()))

# 输出结果和方式一完全一致
result_df_bucket.show()

补充说明

如果你的数据中存在超出0-30区间的年龄值，上面的代码会返回null。可以在when()链最后添加.otherwise(-1)（或其他你需要的默认值）来处理这种情况。
两种方式都确保了最终age_new列是Integer类型，customer_new保持String类型，完全符合你的需求。

内容的提问来源于stack exchange，提问作者Dipanjan Das