PySpark中CountVectorizer稀疏向量的高效特征哈希实现

针对你需要将CountVectorizer生成的特征向量哈希到指定维度（如50）的需求，完全可以利用Spark内置的分布式操作替代嵌套循环的UDF，避免大数据集下的性能瓶颈。以下是具体实现方案：

核心思路 #

1. 拆解原始SparseVector的索引与对应值，转换成扁平的索引-值对
2. 对每个原始索引进行哈希计算，映射到目标哈希维度范围内
3. 聚合哈希冲突的索引值（标准特征哈希逻辑：冲突时求和）
4. 重新构建哈希后的SparseVector

Python 实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, pmod, hash, posexplode, collect_list, sum as spark_sum
from pyspark.ml.linalg import SparseVector, VectorUDT
from pyspark.sql.functions import udf

# 初始化Spark会话
spark = SparkSession.builder.appName("VectorFeatureHash").getOrCreate()
hash_size = 50

# 假设你的输入DataFrame为df，包含id和features列（SparseVector类型）
# df = spark.read... 加载你的数据

# 1. 拆解SparseVector的索引和对应值
df_exploded = df.select(
    col("id"),
    posexplode(col("features").getItem("indices")).alias("pos", "index"),
    col("features").getItem("values")[col("pos")].alias("value")
)

# 2. 计算哈希后的索引（映射到0~hash_size-1范围）
df_hashed = df_exploded.withColumn(
    "hashed_index",
    pmod(hash(col("index")), hash_size)
)

# 3. 聚合哈希冲突的索引值（冲突时求和）
df_agg = df_hashed.groupBy("id", "hashed_index")\
    .agg(spark_sum("value").alias("hashed_value"))

# 4. 收集哈希索引与值，构建新的SparseVector
df_grouped = df_agg.groupBy("id")\
    .agg(
        collect_list("hashed_index").alias("hashed_indices"),
        collect_list("hashed_value").alias("hashed_values")
    )

# 定义轻量UDF构建SparseVector（无嵌套循环，仅调用构造器）
build_vector = udf(
    lambda indices, values: SparseVector(hash_size, indices, values),
    VectorUDT()
)

# 生成最终的featuresHashed列
result_df = df_grouped.withColumn("featuresHashed", build_vector(col("hashed_indices"), col("hashed_values")))\
    .join(df.select("id", "features"), on="id")

# 展示结果
result_df.select("id", "features", "featuresHashed").show(truncate=False)

关键优势 #

• 无嵌套循环：全程使用Spark内置的分布式SQL操作，避免Python UDF的序列化开销与循环低效
• JVM层面执行：除最后构建向量的轻量UDF外，其余步骤均在JVM中执行，性能远优于自定义嵌套循环UDF
• 自动处理冲突：通过聚合操作实现特征哈希的标准冲突处理（相同哈希位置的值求和）

关于示例中的零值说明 #

你提供的示例中，空向量生成了含零值的哈希向量，这不符合SparseVector的存储逻辑（SparseVector仅存储非零值）。实际执行上述代码时，空向量会生成(50,[],[])格式的哈希向量，若需要保留示例中的零值展示，可额外添加逻辑填充，但不建议在生产环境使用（会浪费存储与计算资源）。

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