有条件地将`filter` /`where`应用于Spark`Dataset` /`Dataframe`

Question

嗨伙计我有一个函数从一些S3位置加载数据集并返回有趣的数据

private def filterBrowseIndex(spark: SparkSession, s3BrowseIndex: String, mids: Seq[String] = Seq(), indices: Seq[String] = Seq()): Dataset[BrowseIndex] = {
import spark.implicits._

spark
  .sparkContext.textFile(s3BrowseIndex)
  // split text dataset
  .map(line => line.split("\\s+"))
  // get types for attributes
  .map(BrowseIndex.strAttributesToBrowseIndex)
  // cast it to a dataset (requires implicit conversions)
  .toDS()
  // pick rows for the given marketplaces
  .where($"mid".isin(mids: _*))
  // pick rows for the given indices
  .where($"index".isin(indices: _*))

}

如果有人提供mids = Seq()或indices = Seq() ，则此实现将过滤掉所有内容。 另一方面，我希望语义为“仅当mids不为空时才应用此where子句”（ indices相同），以便在函数的用户提供空序列时不会发生过滤。

有一个很好的功能方法吗？

Answer 1

如果您不介意稍微复杂的逻辑，Raphael Roth的答案对于应用滤镜的具体问题是一个很好的选择。 适用于任何条件转换（不仅仅是过滤而不仅仅在其中一个决策分支上执行任何操作）的通用解决方案是使用transform ，例如，

spark
  .sparkContext.textFile(s3BrowseIndex)
  // split text dataset
  .map(line => line.split("\\s+"))
  // get types for attributes
  .map(BrowseIndex.strAttributesToBrowseIndex)
  // cast it to a dataset (requires implicit conversions)
  .toDS()
  .transform { ds =>
    // pick rows for the given marketplaces
    if (mids.isEmpty) ds
    else ds.where($"mid".isin(mids: _*))
  }
  .transform { ds =>
    // pick rows for the given indices
    if (indices.isEmpty) ds
    else ds.where($"index".isin(indices: _*))
  }

如果您正在使用稳定类型的数据集（或数据框，即Dataset[Row] ），则transform非常有用，因为您可以构建转换函数序列，然后应用它们：

transformations.foldLeft(ds)(_ transform _)

在许多情况下，这种方法有助于代码重用和可测试性。

Answer 2

您可以使用短路评估，如果提供的Seq不为空，则只应用过滤器：

import org.apache.spark.sql.functions.lit

spark
    .sparkContext.textFile(s3BrowseIndex)
    // split text dataset
    .map(line => line.split("\\s+"))
    // get types for attributes
    .map(BrowseIndex.strAttributesToBrowseIndex)
    // cast it to a dataset (requires implicit conversions)
    .toDS()
    // pick rows for the given marketplaces
    .where(lit(mids.isEmpty) or $"mid".isin(mids: _*))
    // pick rows for the given indices
    .where(lit(indices.isEmpty) or $"index".isin(indices: _*))