如何在Pandas和PySpark中实现One Hot Encoding独热编码

直接使用内置的pd.get_dummies()方法即可，指定仅对country列编码，调整前缀参数让生成的列名直接为国家名称，强制输出整数类型的0/1值：

import pandas as pd

# 构造示例数据
df = pd.DataFrame({
    "id": [1, 2],
    "country": ["France", "Italy"],
    "amount": [4, 9],
    "city": ["Paris", "Naples"]
})

# 独热编码转换
df_encoded = pd.get_dummies(
    df,
    columns=["country"],
    prefix="",        # 去掉列名前缀
    prefix_sep="",    # 去掉前缀和类别名的分隔符
    dtype=int         # 输出0/1整数，而非布尔值
)

执行后输出的df_encoded结构和目标完全一致，id/amount/city列原样保留，新增France/Italy两列存储0/1标识。

根据使用场景可以选两种实现方式：

快速数据处理场景 #

直接用pivot透视实现，代码最简洁，不需要依赖MLlib组件：

from pyspark.sql import SparkSession
from pyspark.sql.functions import lit

# 初始化Spark会话、构造示例数据
spark = SparkSession.builder.appName("onehot_demo").getOrCreate()
data = [(1, "France", 4, "Paris"), (2, "Italy", 9, "Naples")]
df = spark.createDataFrame(data, schema=["id", "country", "amount", "city"])

# 透视实现独热编码
df_encoded = df.groupBy("id", "amount", "city")\
    .pivot("country")\
    .agg(lit(1))\
    .na.fill(0)  # 空值填充为0

生产ML管线场景 #

使用Spark MLlib原生的特征工程组件实现，适合封装到训练/预测统一管线，避免数据泄露：

from pyspark.ml.feature import StringIndexer, OneHotEncoder
from pyspark.ml import Pipeline
from pyspark.sql.functions import col

# 定义特征处理步骤
indexer = StringIndexer(inputCol="country", outputCol="country_idx", handleInvalid="keep")
# 注意dropLast=False，否则会默认删除最后一个类别列
encoder = OneHotEncoder(inputCol="country_idx", outputCol="country_vec", dropLast=False)
pipeline = Pipeline(stages=[indexer, encoder])

# 拟合模型、转换数据
model = pipeline.fit(df)
df_trans = model.transform(df)

# 拆分独热向量为单独列，按类别名重命名
country_cols = model.stages[0].labels
for idx, name in enumerate(country_cols):
    df_trans = df_trans.withColumn(name, col("country_vec")[idx].cast("int"))

# 筛选保留目标列
df_encoded = df_trans.select("id", "amount", "city", *country_cols)

两种方式最终输出的表结构完全匹配需求。

内容的提问来源于stack exchange，提问作者1stgear_learner