基于PySpark实现大型数据集行间余弦相似度计算方案问询

Absolutely, you can ditch the Pandas-based approach and leverage PySpark's distributed computing power to calculate pairwise cosine similarity at scale. The key is to break down the cosine similarity formula into distributed operations that play to Spark's strengths, instead of forcing all data onto a single node.

Here's a step-by-step solution tailored to your data structure (row_id, feature, score):

1. Understand the Cosine Similarity Formula #

Cosine similarity between two vectors A and B is calculated as:

cos(A,B) = (A · B) / (||A|| * ||B||)
Where:

• A · B = dot product of the two vectors
• ||A|| = L2 norm (magnitude) of vector A
• ||B|| = L2 norm of vector B

We'll compute each component in a distributed way to avoid bottlenecks.

2. Calculate L2 Norms for Each Row #

First, compute the L2 norm for every row_id—this is the square root of the sum of squared score values for each row's features:

from pyspark.sql import functions as F

# Compute L2 norm for each row_id
row_norms = df.groupBy("row_id").agg(
    F.sqrt(F.sum(F.pow("score", 2))).alias("row_norm")
)

3. Compute Pairwise Dot Products #

Next, we'll do a self-join on the feature column to pair every row_id with every other row_id, then calculate the dot product for each pair by summing the product of their score values for shared features:

# Self-join to get all row_id pairs, then calculate dot product
dot_products = df.alias("left").join(
    df.alias("right"),
    on="feature",
    how="inner"
).groupBy("left.row_id", "right.row_id").agg(
    F.sum(F.col("left.score") * F.col("right.score")).alias("dot_product")
)

This self-join is handled distributedly across your cluster—no single node has to hold all data at once, which fixes the scalability issue from your Pandas approach.

4. Combine Norms and Dot Products to Get Cosine Similarity #

Now join the dot product results with the row norms, then apply the cosine similarity formula. We'll also handle edge cases where a row has no features (norm = 0) by filling those values with 0:

# Join norms with dot products and compute cosine similarity
cosine_similarity = dot_products.join(
    row_norms.alias("left_norm"),
    on=F.col("left.row_id") == F.col("left_norm.row_id"),
    how="inner"
).join(
    row_norms.alias("right_norm"),
    on=F.col("right.row_id") == F.col("right_norm.row_id"),
    how="inner"
).withColumn(
    "cosine_similarity",
    # Avoid division by zero by handling zero norms
    F.when(
        F.col("left_norm.row_norm") * F.col("right_norm.row_norm") == 0,
        0.0
    ).otherwise(
        F.col("dot_product") / (F.col("left_norm.row_norm") * F.col("right_norm.row_norm"))
    )
).select(
    F.col("left.row_id").alias("row_id_1"),
    F.col("right.row_id").alias("row_id_2"),
    "cosine_similarity"
).fillna(0.0, subset=["cosine_similarity"])

5. Optional: Convert to Wide N×N Format #

If you need the exact N×N matrix format like your Pandas solution, you can pivot the results. Note that this can be memory-heavy for very large N, so prefer the long-form (row pairs) format unless you specifically need the wide matrix:

# Pivot to wide N×N matrix (use cautiously for large datasets)
wide_similarity_matrix = cosine_similarity.groupBy("row_id_1").pivot("row_id_2").agg(
    F.first("cosine_similarity")
).fillna(0.0)

Why This Works Better Than Pandas #

• Distributed Execution: All operations (groupbys, joins, aggregates) run across your Spark cluster, so you're not limited to a single node's memory or CPU.
• No Data Shuffling to Driver: Unlike the Pandas approach, you never pull the entire dataset into the driver node—data stays distributed throughout the process.
• Scalability: This will handle datasets that are orders of magnitude larger than what Pandas can process on a single machine.

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