添加/减去两个 pyspark CountVectorizer 稀疏向量列

Question

我想获取 CountVectorizer 转换的文档对的差异。换句话说，取两列稀疏向量之间的差异。我将相同的转换器应用于 df[doc1] 和 df[doc2]，因此结果向量对 (df['X1'] - df['X2']) 的维度将始终保持一致。

from pyspark.ml.feature import RegexTokenizer, CountVectorizer
from pyspark.ml import Pipeline
from pyspark.sql.functions import col

df = spark.createDataFrame([("homer likes donuts".split(" "), "donuts taste delicious".split(" "), 0),
                            ("five by five boss".split(" "), "five is a number".split(" "), 1)],
                            ["words1", "words2", "label"])

display(df)

cv = CountVectorizer()

union_words = df.select(col('words1').alias('words')).union(df.select(col('words2').alias('words')))

cv = CountVectorizer() \
      .setInputCol('words') \
      .fit(union_words)

df = cv.setInputCol('words1') \
        .setOutputCol('X1') \
        .transform(df)

df = cv.setInputCol('words2') \
        .setOutputCol('X2') \
        .transform(df)

display( df )

X1                         X2
[0,11,[1,2,9],[1,1,1]]     [0,11,[1,4,8],[1,1,1]]
[0,11,[0,3,10],[2,1,1]]    [0,11,[0,5,6,7],[1,1,1,1]]

我无法添加列（列类型不匹配，需要数字或日历间隔）。我尝试了@zero323 的add function，但在 isinstance(v1, SparseVector) 处遇到断言错误

df.withColumn("result", (col("X1") + col("X2"))
df.withColumn("result", add(col("X1"), col("X2"))

在稀疏向量格式中，我希望结果是：

[0,11,[2,4,8,9],[1,-1,-1,1]]
[0,11,[0,3,5,6,7,10],[1,1,-1,-1,-1,1]]

Answer 1

需要将函数转换为返回类型为 VectorUDT 的 udf。通过组合解决方案here 和here 解决。

import numpy as np
from pyspark.sql.functions import udf
from pyspark.ml.linalg import SparseVector, Vectors, VectorUDT
from pyspark.ml.feature import RegexTokenizer, CountVectorizer
from pyspark.ml import Pipeline
from pyspark.sql.functions import col

df = spark.createDataFrame(data=[("homer likes donuts".split(" "), ["donuts", "taste", "delicious"], 0),
                            (["five", "by", "five", "boss"], ["five", "is", "a", "number"], 1)],
                            schema=["words1", "words2", "label"])


cv = CountVectorizer()

union_words = df.select(col('words1').alias('words')).union(df.select(col('words2').alias('words')))

cv = CountVectorizer() \
      .setInputCol('words') \
      .fit(union_words)

df = cv.setInputCol('words1') \
          .setOutputCol('X1') \
          .transform(df)

df = cv.setInputCol('words2') \
          .setOutputCol('X2') \
          .transform(df)

@udf(VectorUDT())
def minus(v1, v2):
    # Sparse vector will become dense
    assert isinstance(v1, SparseVector) and isinstance(v2, SparseVector)
    assert v1.size == v2.size 
    # Compute union of indices
    indices = set(v1.indices).union(set(v2.indices))
    # Not particularly efficient but we are limited by SPARK-10973
    # Create index: value dicts
    v1d = dict(zip(v1.indices, v1.values))
    v2d = dict(zip(v2.indices, v2.values))
    zero = np.float64(0)
    # Create dictionary index: (v1[index] - v2[index])
    values =  {i: v1d.get(i, zero) - v2d.get(i, zero)
               for i in indices
               if v1d.get(i, zero) - v2d.get(i, zero) != zero}
    return Vectors.sparse(v1.size, values)

df = df.withColumn('NAME_X_DIFF', minus('X1', 'X2'))
display(df)

+----------------------+--------------------------+-----+---------------------------+--------------------------------+-----------------------------------------------+
|words1                |words2                    |label|X1                         |X2                              |X_DIFF                                         |
+----------------------+--------------------------+-----+---------------------------+--------------------------------+-----------------------------------------------+
|[homer, likes, donuts]|[donuts, taste, delicious]|0    |(11,[1,4,10],[1.0,1.0,1.0])|(11,[1,3,5],[1.0,1.0,1.0])      |(11,[3,4,5,10],[-1.0,1.0,-1.0,1.0])            |
|[five, by, five, boss]|[five, is, a, number]     |1    |(11,[0,6,9],[2.0,1.0,1.0]) |(11,[0,2,7,8],[1.0,1.0,1.0,1.0])|(11,[0,2,6,7,8,9],[1.0,-1.0,1.0,-1.0,-1.0,1.0])|
+----------------------+--------------------------+-----+---------------------------+--------------------------------+-----------------------------------------------+