背景

    data = [
        (1, "北京", "上海", ["北京", "上海", "深圳"], 1),
        (2, "上海", "北京", ["上海", "北京"], 0),
        (3, "广州", "深圳", ["广州", "深圳", "北京"], 1),
        (4, "深圳", "广州", ["深圳"], 0),
        (5, "北京", "广州", ["广州", "北京"], 1),
        (6, "上海", "深圳", ["上海", "深圳"], 0),
    ]
    columns = ["id", "origin_city", "dest_city", "love_citys", "label"]
    df = spark.createDataFrame(data, schema=columns)

如上spark数据,要求origin_city dest_city love_citys 中的城市公用一套labelEncoder编码,即三列中同一个城市编码后需要时同一个值.

对比之前博客的区别是该版本支持传入数组类型(love_citys)

列类型如下:

df.schema
# StructType([StructField('id', LongType(), True)
#             , StructField('origin_city', StringType(), True)
#             , StructField('dest_city', StringType(), True)
#             , StructField('love_citys', ArrayType(StringType(), True), True)
#             , StructField('label', LongType(), True)])

代码

"""
需求说明
- 两列城市字段（origin_city、dest_city）表达同一语义，需要共享一套 Label 编码映射。
- 使用 PySpark 框架实现，且编码器可复用，并可整合进 Spark ML Pipeline。
"""

from __future__ import annotations

from typing import Dict, Iterable, List, Optional, Tuple

from pyspark.sql import SparkSession, functions as F, types as T
from pyspark.ml import Estimator, Model, Pipeline, PipelineModel
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.param.shared import Param, Params, TypeConverters
from pyspark.ml.util import DefaultParamsReadable, DefaultParamsWritable


class SharedStringIndexer(Estimator, DefaultParamsReadable, DefaultParamsWritable):
    """一个 Estimator：基于多列拟合一次 StringIndexer 的 labels，并产出可对多列统一编码的 Model。

    Params
    - inputCols: List[str] 要共享映射的输入列
    - outputCols: List[str] 对应输出列名（与 inputCols 等长）
    - handleInvalid: keep/skip/error（含义同 StringIndexer）
    """

    inputCols = Param(Params._dummy(), "inputCols", "input columns", typeConverter=TypeConverters.toListString)
    outputCols = Param(Params._dummy(), "outputCols", "output columns", typeConverter=TypeConverters.toListString)
    handleInvalid = Param(Params._dummy(), "handleInvalid", "how to handle invalid labels", typeConverter=TypeConverters.toString)

    def __init__(self, inputCols: List[str], outputCols: List[str], handleInvalid: str = "keep"):
        super().__init__()
        if len(inputCols) != len(outputCols):
            raise ValueError("inputCols 与 outputCols 长度需一致")
        self._set(inputCols=inputCols, outputCols=outputCols, handleInvalid=handleInvalid)

    def _fit(self, dataset):
        # 将多列（含 string 与 array<string>）堆叠为单列 value，用 StringIndexer 拟合一次，得到统一 labels
        stacked = None
        schema = {f.name: f.dataType for f in dataset.schema.fields}
        for c in self.getOrDefault(self.inputCols):
            dt = schema.get(c)
            if isinstance(dt, T.ArrayType):
                # 将数组元素展开为多行
                col_df = dataset.select(F.explode_outer(F.col(c)).alias("value")).select(F.col("value").cast(T.StringType()))
            else:
                col_df = dataset.select(F.col(c).cast(T.StringType()).alias("value"))
            col_df = col_df.na.fill({"value": ""})
            stacked = col_df if stacked is None else stacked.unionByName(col_df)
        indexer = StringIndexer(inputCol="value", outputCol="value_idx", handleInvalid="keep")
        model = indexer.fit(stacked)
        labels = list(model.labels)
        return SharedStringIndexerModel().setParams(
            inputCols=self.getOrDefault(self.inputCols),
            outputCols=self.getOrDefault(self.outputCols),
            handleInvalid=self.getOrDefault(self.handleInvalid),
            labels=labels,
        )


class SharedStringIndexerModel(Model, DefaultParamsReadable, DefaultParamsWritable):
    """Transformer：将拟合得到的 labels 作为共享映射，对多列输出统一索引。

    为了能够被 PipelineModel.save/load 序列化，labels 作为一个 Param 保存。
    """

    inputCols = Param(Params._dummy(), "inputCols", "input columns", typeConverter=TypeConverters.toListString)
    outputCols = Param(Params._dummy(), "outputCols", "output columns", typeConverter=TypeConverters.toListString)
    handleInvalid = Param(Params._dummy(), "handleInvalid", "how to handle invalid labels", typeConverter=TypeConverters.toString)
    labels = Param(Params._dummy(), "labels", "shared label list", typeConverter=TypeConverters.toListString)

    def __init__(self):
        # 必须无参构造以支持反序列化
        super().__init__()

    def setParams(self, **kwargs):
        self._set(**kwargs)
        return self

    def _transform(self, dataset):
        label_list = self.getOrDefault(self.labels) or []
        mapping = {v: i for i, v in enumerate(label_list)}
        unknown_index = len(label_list)
        handle_invalid = self.getOrDefault(self.handleInvalid)
        bmap = dataset.sparkSession.sparkContext.broadcast(mapping)

        def map_value(v: Optional[str]) -> Optional[int]:
            if v is None:
                return None if handle_invalid == "skip" else unknown_index if handle_invalid == "keep" else None
            idx = bmap.value.get(v)
            if idx is not None:
                return idx
            if handle_invalid == "keep":
                return unknown_index
            if handle_invalid == "skip":
                return None
            raise ValueError(f"未知标签: {v}")

        def map_array(arr: Optional[List[Optional[str]]]) -> Optional[List[Optional[int]]]:
            if arr is None:
                return None if handle_invalid == "skip" else [] if handle_invalid == "keep" else None
            return [map_value(v) for v in arr]

        enc_scalar = F.udf(map_value, T.IntegerType())
        enc_array = F.udf(map_array, T.ArrayType(T.IntegerType()))

        schema = {f.name: f.dataType for f in dataset.schema.fields}
        out = dataset
        for src, dst in zip(self.getOrDefault(self.inputCols), self.getOrDefault(self.outputCols)):
            dt = schema.get(src)
            if isinstance(dt, T.ArrayType):
                out = out.withColumn(dst, enc_array(F.col(src)))
            else:
                out = out.withColumn(dst, enc_scalar(F.col(src).cast(T.StringType())))
        return out


def main():
    spark = SparkSession.builder.appName("shared_string_indexer_pipeline").getOrCreate()
    spark.sparkContext.setLogLevel("ERROR")

    data = [
        (1, "北京", "上海", ["北京", "上海", "深圳"], 1),
        (2, "上海", "北京", ["上海", "北京"], 0),
        (3, "广州", "深圳", ["广州", "深圳", "北京"], 1),
        (4, "深圳", "广州", ["深圳"], 0),
        (5, "北京", "广州", ["广州", "北京"], 1),
        (6, "上海", "深圳", ["上海", "深圳"], 0),
    ]
    columns = ["id", "origin_city", "dest_city", "love_citys", "label"]
    df = spark.createDataFrame(data, schema=columns)
    df_test = spark.createDataFrame([
        (7, "北京111", "上海", ["北京", "杭州"], 1)
    ], schema=columns)

    shared_indexer = SharedStringIndexer(
        inputCols=["origin_city", "dest_city", "love_citys"],
        outputCols=["origin_city_idx", "dest_city_idx", "love_citys_idx"],
        handleInvalid="keep",
    )

    # love_citys_idx 为 array<int>，不直接进 VectorAssembler，这里仅演示编码输出
    pipeline = Pipeline(stages=[shared_indexer])
    model = pipeline.fit(df)
    out = model.transform(df)
    out.select("id", "origin_city", "dest_city", "love_citys", "origin_city_idx", "dest_city_idx", "love_citys_idx").show(truncate=False)

    # 复用：保存/加载整个 PipelineModel（包含共享映射）
    # 也可以仅保存 shared_indexer 的 model（通过 pipeline.stages[0] 的写接口）
    model.write().overwrite().save("./shared_indexer_pipeline_model")
    
    print('新数据转换:')   # handleInvalid="keep" 所以这里新枚举值不报错 
    model.transform(df_test).show(truncate=False)
    
    print("加载导出后的模型 新数据转换")
    loaded_model = PipelineModel.load("./shared_indexer_pipeline_model")
    loaded_model.transform(df_test).show(truncate=False)

    # spark.stop()


main()

输出

+---+-----------+---------+------------------+---------------+-------------+--------------+
|id |origin_city|dest_city|love_citys        |origin_city_idx|dest_city_idx|love_citys_idx|
+---+-----------+---------+------------------+---------------+-------------+--------------+
|1  |北京       |上海     |[北京, 上海, 深圳]|0              |2            |[0, 2, 1]     |
|2  |上海       |北京     |[上海, 北京]      |2              |0            |[2, 0]        |
|3  |广州       |深圳     |[广州, 深圳, 北京]|3              |1            |[3, 1, 0]     |
|4  |深圳       |广州     |[深圳]            |1              |3            |[1]           |
|5  |北京       |广州     |[广州, 北京]      |0              |3            |[3, 0]        |
|6  |上海       |深圳     |[上海, 深圳]      |2              |1            |[2, 1]        |
+---+-----------+---------+------------------+---------------+-------------+--------------+

新数据转换:
+---+-----------+---------+------------+-----+---------------+-------------+--------------+
|id |origin_city|dest_city|love_citys  |label|origin_city_idx|dest_city_idx|love_citys_idx|
+---+-----------+---------+------------+-----+---------------+-------------+--------------+
|7  |北京111    |上海     |[北京, 杭州]|1    |4              |2            |[0, 4]        |
+---+-----------+---------+------------+-----+---------------+-------------+--------------+

加载导出后的模型 新数据转换
+---+-----------+---------+------------+-----+---------------+-------------+--------------+
|id |origin_city|dest_city|love_citys  |label|origin_city_idx|dest_city_idx|love_citys_idx|
+---+-----------+---------+------------+-----+---------------+-------------+--------------+
|7  |北京111    |上海     |[北京, 杭州]|1    |4              |2            |[0, 4]        |
+---+-----------+---------+------------+-----+---------------+-------------+--------------+

-