PySpark设置StringIndexer固定编码：空值为0非空值从1开始

问题背景 #

现有如下结构的Spark DataFrame：

a      b       c         d    
1      he     Null      fish
1      nb     canada     crab
2      ca     china     turtle
3      Null   chile      bird

需要实现两个编码目标：

• 对字段应用StringIndexer编码时，Null值始终被编码为0
• 若列中不包含Null值，编码值从1开始计数

初始方案尝试调用StringIndexer的stringOrderType参数，通过给Null填充空格前缀让其在字母排序时排到最前实现0编码，代码如下：

# added space to give lowest ascii code(32).
df = df.fillna(' unknown') 

cols = df.columns
indexers = {}
for col in cols:
    indexer = StringIndexer(inputCol=col, outputCol="ind_"+col, stringOrderType='alphabets')
    indexers[col] = indexer

该方案存在两个明显问题：实际数据集上可能出现填充值和真实值冲突，导致部分列编码结果不符合预期；即使排序逻辑生效，也无法满足「无Null列编码从1开始」的第二个需求。

后续尝试手动调整indexer的labels属性，将Null移到标签列表首位后直接传入transform方法，触发报错ValueError: Params must be a parammap but got list，原因是transform方法不支持直接传入labels参数修改映射规则。

正确实现方案 #

核心思路是在StringIndexer训练完成后，通过模型自带的setLabels()方法自定义标签和索引的映射关系（标签列表的下标即为对应编码值），不需要修改transform入参，从根源避免参数错误，具体实现代码如下：

from pyspark.ml.feature import StringIndexer
from pyspark.sql import functions as F

cols = df.columns
indexed_df = df

for col in cols:
    # 先判断当前列是否包含Null值
    null_cnt = indexed_df.select(F.count(F.when(F.isnull(col), 1))).first()[0]
    has_null = null_cnt > 0

    # 按指定排序规则训练基础模型，拿到非空值的默认排序结果
    base_indexer = StringIndexer(
        inputCol=col,
        outputCol=f"ind_{col}",
        stringOrderType="alphabetAsc" # 保持和原方案一致的字母排序规则，可按需替换为frequencyDesc等其他规则
    )
    base_model = base_indexer.fit(indexed_df)
    # 过滤掉默认标签里可能存在的Null值，拿到纯非空值的有序列表
    valid_labels = [label for label in base_model.labels if label is not None]

    # 按需求构建新的标签映射列表
    if has_null:
        # 有Null的列：Null放在列表第0位，对应编码0，其余有效值按原顺序顺延，对应编码从1开始
        new_labels = [None] + valid_labels
    else:
        # 无Null的列：用不会和现有值冲突的占位符占住0号索引位，真实值从1位开始，实现编码从1计数
        placeholder = f"__placeholder_for_{col}__"
        # 极端场景下校验占位符唯一性，避免和真实值冲突
        while placeholder in valid_labels:
            placeholder += "_tmp"
        new_labels = [placeholder] + valid_labels

    # 给模型设置自定义标签列表后执行转换
    final_model = base_model.setLabels(new_labels)
    indexed_df = final_model.transform(indexed_df)

方案说明 #

• 不需要提前对Null值做fillna填充，完全避免填充值和真实值冲突导致的编码错乱问题，稳定性远高于初始方案
• 存在Null的列天然保证Null对应编码0，其余非空值按指定规则排序后从1开始编码，完全满足第一个需求
• 无Null的列通过占位符占用0号索引位，所有真实值的编码都从1开始计数，满足第二个需求
• 排序规则可按需调整，不影响自定义编码的核心逻辑

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