Pyspark：根据 Array(String) 长度或 CountVectorizer 计数过滤 DF

Question

我将 URL 数据聚合到一个字符串数组中。这种形式。[xyz.com,abc.com,efg.com]

我最终在 pyspark 中使用计数向量化器将其转换为类似 (262144,[3,20,83721],[1.0,1.0,1.0]) 的向量。

向量从 262144 中说出来；对于某一行，有 3 个 URL 索引在 3,20 和 83721。所有这些数据都是二进制的，因此是 1 的数组。

我想过滤并仅使用包含一定数量条目的行。因此，如果我说只想使用带有 4 个示例的行；我希望删除上面的行。

我很好地对由 Count Vectorizer 返回的 Vector 形式的 String Array 执行此过滤。

在我的数据中，我有数千万行，只是不知道如何有效地做到这一点。

这是编辑以生成示例的文档中的代码

from pyspark.ml.feature import CountVectorizer

# Input data: Each row is a bag of words with a ID.
df = spark.createDataFrame([
    (0, "a".split(" ")),
    (1, "a b c".split(" "))
], ["id", "words"])

# fit a CountVectorizerModel from the corpus.
cv = CountVectorizer(inputCol="words", outputCol="features",         vocabSize=3)

model = cv.fit(df)

result = model.transform(df)

所以这里说我们只想要长度为 2 或更长的结果。同样，我不介意在 CountVectorizer 或之前的 String 数组生成的向量上执行此操作，只要它对我的数据大小有效。

Answer 1

你可以只使用 DataFrame.filter 和 size 函数：

from pyspark.sql.functions import size

df.filter(size('words') > 2).show()

+---+---------+
| id|    words|
+---+---------+
|  1|[a, b, c]|
+---+---------+

我会在 CountVectorizer 之前做这件事，以避免让它做不需要做的工作。如果 Spark 可以确定这样做是安全的，它会在执行计划中更早地提取过滤器操作，但显式总是更好。