检查 Spark DataFrame 是否为空

1. 引言

在使用 Apache Spark 时，我们经常会遇到 DataFrame 在经过转换、过滤或连接操作后变为空的情况。

在触发昂贵的操作或将数据保存到外部存储之前，最佳实践是验证 DataFrame 是否包含任何行。 虽然这看起来像是一个简单的检查，但不同的方法会触发不同的 Spark 操作，并且可能对性能产生显著影响，尤其是在处理大型数据集时。

在本教程中，我们将探讨检查 Spark DataFrame 是否为空的多种不同方法。

2. 环境准备

首先，让我们创建以下示例 DataFrame：

public static Dataset getDataFrame(SparkSession spark) {
    List<Row> players = List.of(
      RowFactory.create(1, "Messi", "Argentina"),
      RowFactory.create(2, "Ronaldo", "Portugal"),
      RowFactory.create(3, "Mbappe", "France"));
    return spark.createDataFrame(players, PLAYER_SCHEMA);
}

我们将在整个示例中使用这个 DataFrame。

3. 使用 isEmpty() 方法

自 Spark 2.4.0 起，检查 Dataset 是否为空的最简单方法是使用 isEmpty() 方法。

让我们通过过滤数据集中不存在的国家来创建一个空的 DataFrame：

Dataset<Row> englandPlayers = allPlayers.filter(col("country").equalTo("England"));
Assertions.assertTrue(englandPlayers.isEmpty());

由于数据集中没有英格兰球员，因此过滤操作会返回一个空的 DataFrame，并且 isEmpty() 会返回 true。

4. 使用 count() 方法

另一种检查 DataFrame 是否为空的方法是计算 DataFrame 中的行数，并检查它是否为零：

Dataset<Row> englandPlayers = allPlayers.filter(col("country").equalTo("England"));
Assertions.assertEquals(0, englandPlayers.count());

这里，我们调用了 count()，并验证了行数为 0。

5. 使用 takeAsList() 方法

另一种方法是使用 takeAsList() 获取至多一行：

Dataset<Row> englandPlayers = allPlayers.filter(col("country").equalTo("England"));
Assertions.assertEquals(0, englandPlayers.takeAsList(1).size());

这里我们调用了 takeAsList(1)，它最多返回一行。然后我们可以通过检查 list 的大小是否为 0 来验证返回的 list 是否为空。由于只请求了一行，Spark 在找到第一行匹配的数据后就可以停止处理，因此这种方法比统计所有行更加高效。

另外，我们也可以直接检查 list 是否为空，而不检查 size()：

Assertions.assertTrue(englandPlayers.takeAsList(1).isEmpty());

6. 比较各方法的性能

尽管所有这些方法都能判断 DataFrame 是否为空，但它们的性能特性并不相同。isEmpty() 方法是推荐使用的方法，因为它正是为此目的而设计的。在内部，Spark 只需要判断是否至少存在一行记录，从而在找到第一行时立即停止处理。

同样地，takeAsList(1) 只请求一行。一旦 Spark 获取到第一行，它就可以终止扫描，这使得这种方法比统计所有行要高效得多。相比之下，count() 会计算 DataFrame 中的总行数。由于 Spark 必须处理整个数据集才能给出准确的计数，因此对于大型 DataFrame 来说，这种方法代价相当高昂。

通常，只要可用，我们应优先使用 isEmpty()。如果我们使用的是不支持该方法的老版本 Spark，那么 takeAsList(1).isEmpty() 是一个不错的替代方案。我们应该将 count() 留到那些还需要总行数的情况再使用。

7. 结论

在本文中，我们探讨了几种判断 Spark DataFrame 是否为空的方法。我们从 isEmpty() 方法开始，这是最简单且最具表现力的解决方案，然后研究了基于 count() 和 takeAsList() 的替代方法。

在大多数情况下，isEmpty() 是首选方法，因为它既具表现力又高效。对于较旧的 Spark 版本，takeAsList(1).isEmpty() 提供了实用的替代方案，而 count() 则最适合用于还需要获取总行数的场景。

与往常一样，本文使用的示例代码可在 GitHub 上获取。