PySpark：基于字符串数组长度或CountVectorizer统计过滤DataFrame

嘿，这事儿其实挺 straightforward 的，咱们一步步来搞定它：

首先得明确一点：你用CountVectorizer生成的稀疏向量里，非零元素的个数正好对应该行的URL数量——因为你设置了二进制模式（binary=True），每个存在的URL都会对应一个值为1.0的非零项，而这些项的索引存在indices数组里，所以数组的长度就是URL的数量。

步骤1：准备示例数据（替换成你的真实数据集即可） #

先创建一个带URL数组列的DataFrame，模拟你的场景：

from pyspark.sql import SparkSession
from pyspark.ml.feature import CountVectorizer
from pyspark.sql.functions import udf
from pyspark.ml.linalg import SparseVector
from pyspark.sql.types import IntegerType

spark = SparkSession.builder.appName("URLFilterDemo").getOrCreate()

# 模拟数据：有的行3个URL，有的4个，有的2个
sample_data = [
    (["xyz.com", "abc.com", "efg.com"],),
    (["a.com", "b.com", "c.com", "d.com"],),
    (["x.com", "y.com"],),
    (["p.com", "q.com", "r.com", "s.com"],)
]
df = spark.createDataFrame(sample_data, ["url_array"])

步骤2：用CountVectorizer生成二进制向量列 #

记得开启binary=True来匹配你需要的二进制向量需求：

# 初始化CountVectorizer，指定输入输出列，开启二进制模式
cv = CountVectorizer(inputCol="url_array", outputCol="url_vector", binary=True)
# 训练模型并转换数据
cv_model = cv.fit(df)
vectorized_df = cv_model.transform(df)

# 可以先看看转换后的结果
vectorized_df.show(truncate=False)

步骤3：计算每行的URL数量 #

我们写一个简单的UDF来提取稀疏向量的非零元素个数（也就是URL的数量）：

# 定义UDF：获取稀疏向量的indices数组长度，即该行的URL数量
count_urls = udf(lambda vec: len(vec.indices), IntegerType())

# 添加新列存储URL数量
df_with_count = vectorized_df.withColumn("url_count", count_urls("url_vector"))
df_with_count.show(truncate=False)

步骤4：过滤出指定数量的行 #

比如你要保留正好4个URL的行，直接过滤就行：

# 过滤出URL数量等于4的行
filtered_df = df_with_count.filter(df_with_count.url_count == 4)
filtered_df.show(truncate=False)

小提示 #

如果你的场景中特征特别多（比如你提到的262144个），用稀疏向量的indices数组长度来计数是最高效的，不用遍历整个向量。要是万一遇到DenseVector的情况，也可以改成len([val for val in vec if val != 0])，不过CountVectorizer默认会输出SparseVector，所以不用担心。

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