python - 在 Spark Dataframe (Pyspark) 中提取与特定条件匹配的第一个 "set of rows"

Question

我有一个 Spark DataFrame，其中包含如下数据:

ID | UseCase
-----------------
0  | Unidentified
1  | Unidentified
2  | Unidentified
3  | Unidentified
4  | UseCase1
5  | UseCase1
6  | Unidentified
7  | Unidentified
8  | UseCase2
9  | UseCase2
10 | UseCase2
11 | Unidentified
12 | Unidentified

我必须提取 UseCase 列中具有值 Unidentified 的前 4 行，并对它们进行进一步处理。此时我不想让中间和最后两行具有 Unidentified 值。

我想避免使用 ID 列，因为它们不固定。上述数据只是一个示例。当我使用映射函数(将其转换为 RDD 后)或 UDF 时，我最终在输出 DataFrame 中得到 8 行(这是这些函数所期望的)。

如何实现这一目标？我在 PySpark 工作。我不想在 DataFrame 上使用收集并将其作为列表进行迭代。这将违背 Spark 的目的。 DataFrame 大小可达 4-5 GB。

您能建议一下如何做到这一点吗？提前致谢!

Answer 1

只需做一个过滤器和一个限制。下面的代码是Scala，但你会明白这一点。

假设您的数据框名为 df，那么:

df.filter($"UseCase"==="Unidentified").limit(4).collect()