如何用Spark SQL过滤正标签数量不足的ID对应的DataFrame行

你的问题核心是只统计每个ID下label=1的正标签数量，然后过滤掉正标签数不足2的ID的所有行。之前的SQL有两个问题：一是count(label)统计的是每个ID的总记录数（而非正标签数），二是SELECT *配合GROUP BY id在Spark SQL的默认模式下不允许，因为非聚合列没有出现在GROUP BY中。

下面提供两种可行的解决方案，分别用Spark SQL和DataFrame API实现：

方案一：Spark SQL分步实现（先统计再关联） #

先单独计算每个ID的正标签数量，再关联回原表过滤出符合条件的行：

datalake_spark_dataframe_downsampled.createOrReplaceTempView("df_filtered")

-- 第一步：统计每个ID的正标签数量，筛选出数量≥2的ID
CREATE OR REPLACE TEMP VIEW valid_ids AS
SELECT 
    id,
    SUM(CASE WHEN label = 1 THEN 1 ELSE 0 END) AS positive_label_count
FROM df_filtered
GROUP BY id
HAVING positive_label_count >= 2;

-- 第二步：关联原表，保留有效ID的所有行
SELECT df.*
FROM df_filtered df
JOIN valid_ids vi ON df.id = vi.id;

方案二：Spark SQL窗口函数一步实现 #

用窗口函数直接在原表中添加每个ID的正标签统计列，再过滤：

datalake_spark_dataframe_downsampled.createOrReplaceTempView("df_filtered")

SELECT *
FROM (
    SELECT 
        *,
        -- 按ID分区，统计当前ID下的正标签总数
        SUM(CASE WHEN label = 1 THEN 1 ELSE 0 END) OVER (PARTITION BY id) AS positive_label_count
    FROM df_filtered
) temp_table
WHERE positive_label_count >= 2;

方案三：PySpark DataFrame API实现（更简洁） #

如果用DataFrame API操作，代码会更直观：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 定义按ID分区的窗口
id_window = Window.partitionBy("id")

# 添加正标签统计列，然后过滤
spark_dataset_filtered = datalake_spark_dataframe_downsampled \
    .withColumn(
        "positive_label_count",
        F.sum(F.when(F.col("label") == 1, 1).otherwise(0)).over(id_window)
    ) \
    .filter(F.col("positive_label_count") >= 2)

关键说明 #

• 统计正标签数量时，用SUM(CASE WHEN label=1 THEN 1 ELSE 0 END)（或COUNT(CASE WHEN label=1 THEN 1 END)，因为COUNT会忽略NULL值），这两种写法都能准确统计每个ID下label=1的行数。
• 窗口函数的优势是不用额外创建临时表，直接在原数据上计算统计值，代码更紧凑。

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