java - Spark 数据集使用 agg() 方法计算与条件匹配的行数(在 Java 中)

Question

我在 Java 中使用 Apache Spark 2.3.1。我想通过使用 Dataset 类的 agg() 方法来计算数据集中匹配给定条件的行数。

例如，我想计算以下数据集中 label 等于 1.0 的行数:

SparkSession spark = ...

List<Row> rows = new ArrayList<>();
rows.add(RowFactory.create(0, 0.0));
rows.add(RowFactory.create(1, 1.0));
rows.add(RowFactory.create(2, 1.0));

Dataset<Row> ds =
    spark.sqlContext().createDataFrame(rows,
        new StructType(new StructField[] {
            new StructField("id", DataTypes.LongType, false, Metadata.empty()),
            new StructField("label", DataTypes.DoubleType, false, Metadata.empty())}));

我的猜测是使用以下代码:

ds.agg(functions.count(ds.col("label").equalTo(1.0))).show();

但是，显示错误的结果:

+--------------------+
|count((label = 1.0))|
+--------------------+
|                   3|
+--------------------+

正确的结果当然应该是2。

agg() 方法不应该以这种方式工作吗？

Answer 1

count in agg() 只会计算非空值，因此可以这样做:

 import org.apache.spark.sql.functions._
 ds.agg(count(when('label.equalTo(1.0),1).otherwise(null))).show()

我在这里找到了这个解决方案https://stackoverflow.com/a/1400115/9687910