PySpark ml.stat.Summarizer能否返回稀疏向量结果？如何强制其执行稀疏向量聚合操作及替代实现方案

Hey there! Let's tackle your questions about handling sparse vectors with pyspark.ml.stat.Summarizer:

1. Can pyspark.ml.stat.Summarizer return sparse vector results? #

Short answer: No. Out of the box, Summarizer's aggregation methods (like mean, sum) always return a DenseVector, even if all your input data is sparse. Your code example clearly shows this behavior—you start with a SparseVector but end up with a DenseVector for the mean.

2. Is there a way to force Summarizer to use sparse vectors during computation? #

Unfortunately, PySpark's native Summarizer doesn't support this. Under the hood, it converts sparse vectors to dense ones for aggregation, which wastes a ton of memory for high-dimensional sparse data. Converting the final DenseVector back to SparseVector is a band-aid solution—it doesn't fix the resource bloat from intermediate dense calculations.

Alternative Implementations for Sparse Vector Aggregation #

To avoid dense vector overhead entirely, you'll need to build custom aggregation logic. Here are two practical approaches, using mean calculation as an example:

Option 1: Manual Aggregation on RDDs #

This approach works directly with RDDs, keeping all operations on sparse vectors:

import pyspark
from pyspark.ml.linalg import SparseVector

sc = pyspark.SparkContext.getOrCreate()
sql_context = pyspark.SQLContext(sc)

# Test data with two sparse vectors
df = sc.parallelize([
    (SparseVector(100, {1: 1.0}),),
    (SparseVector(100, {3: 2.0}),)
]).toDF(['v'])

# Helper to add two sparse vectors
def add_sparse_vectors(v1, v2):
    merged_indices = sorted(set(v1.indices).union(v2.indices))
    idx_vals1 = dict(zip(v1.indices, v1.values))
    idx_vals2 = dict(zip(v2.indices, v2.values))
    merged_values = [idx_vals1.get(i, 0.0) + idx_vals2.get(i, 0.0) for i in merged_indices]
    return SparseVector(v1.size, merged_indices, merged_values)

# Accumulate vectors and count rows
sum_vec, total_count = df.rdd.map(lambda row: (row.v, 1)).reduce(
    lambda a, b: (add_sparse_vectors(a[0], b[0]), a[1] + b[1])
)

# Calculate sparse mean
mean_indices = sum_vec.indices
mean_values = [val / total_count for val in sum_vec.values]
sparse_mean = SparseVector(sum_vec.size, mean_indices, mean_values)

print(f"Sum of sparse vectors: {sum_vec}")
print(f"Sparse mean vector: {sparse_mean}")

Option 2: Custom User-Defined Aggregate Function (UDAF) #

If you prefer working with DataFrames, you can build a UDAF that operates exclusively on sparse vectors:

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, MapType, DoubleType
from pyspark.sql.expressions import UserDefinedAggregateFunction
from pyspark.ml.linalg import SparseVector, VectorUDT

class SparseMeanUDAF(UserDefinedAggregateFunction):
    def inputSchema(self):
        return StructType([StructField("v", VectorUDT())])
    
    def bufferSchema(self):
        return StructType([
            StructField("count", IntegerType()),
            StructField("sum_map", MapType(IntegerType(), DoubleType())),
            StructField("vec_size", IntegerType())
        ])
    
    def dataType(self):
        return VectorUDT()
    
    def deterministic(self):
        return True
    
    def initialize(self, buffer):
        buffer[0] = 0
        buffer[1] = {}
        buffer[2] = 0
    
    def update(self, buffer, input):
        vec = input[0]
        if isinstance(vec, SparseVector):
            buffer[0] += 1
            buffer[2] = vec.size
            for idx, val in zip(vec.indices, vec.values):
                buffer[1][idx] = buffer[1].get(idx, 0.0) + val
    
    def merge(self, buffer1, buffer2):
        buffer1[0] += buffer2[0]
        buffer1[2] = buffer2[2]  # Assumes all vectors have the same size
        for idx, val in buffer2[1].items():
            buffer1[1][idx] = buffer1[1].get(idx, 0.0) + val
    
    def evaluate(self, buffer):
        count = buffer[0]
        if count == 0:
            return SparseVector(buffer[2], [], [])
        sorted_indices = sorted(buffer[1].keys())
        mean_values = [buffer[1][idx] / count for idx in sorted_indices]
        return SparseVector(buffer[2], sorted_indices, mean_values)

# Use the custom UDAF
sparse_mean_udaf = SparseMeanUDAF()
result_df = df.select(sparse_mean_udaf(F.col('v')).alias('sparse_mean'))
print(result_df.head())

Quick Notes #

• Both solutions assume all input sparse vectors have the same dimension—add validation logic if your data varies
• RDD-based aggregation might perform better for extremely high-dimensional data, as it avoids UDAF framework overhead
• You can adapt these patterns to compute other metrics like sum, variance, or max for sparse vectors

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