如何在 Spark-Scala 数据框中应用布尔索引？答案

【问题标题】：How can I apply boolean indexing in a Spark-Scala dataframe?如何在 Spark-Scala 数据框中应用布尔索引？
【发布时间】：2020-12-27 03:32:30
【问题描述】：

我有两个 Spark-Scala 数据帧，我需要使用一个数据帧中的一个布尔列来过滤第二个数据帧。两个数据框的行数相同。

在熊猫中我会这样：

import pandas as pd

df1 = pd.DataFrame({"col1": ["A", "B", "A", "C"], "boolean_column": [True, False, True, False]})
df2 = pd.DataFrame({"col1": ["Z", "X", "Y", "W"], "col2": [1, 2, 3, 4]})

filtered_df2 = df2[df1['boolean_column']]

// Expected filtered_df2 should be this:
// df2 = pd.DataFrame({"col1": ["Z", "Y"], "col2": [1, 3]})

如何以最省时的方式在 Spark-Scala 中执行相同的操作？

我目前的解决方案是将"boolean_column"从df1添加到df2，然后通过在新添加的列中仅选择具有true值的行来过滤df2，最后从@中删除"boolean_column" 987654328@，但我不确定这是不是最好的解决方案。

欢迎提出任何建议。

编辑：

预期输出是 Spark-Scala 数据帧（不是列表或列），其架构与第二个数据帧相同，并且只有来自 df2 的行子集满足来自 "boolean_column" 的布尔掩码df1。
上述df2 的架构只是一个示例。我希望收到 df2 作为参数，其中包含任意数量的不同（而非固定）模式的列。

【问题讨论】：

他们有共同的钥匙吗？否则你不能合并/加入你的数据框。我建议过滤 df1（只保留 true），然后使用 left-semi 加入
我同意，如果没有连接条件，您不能保证两个表中的行将以相同的顺序连接
不，数据集没有任何公共键，即使第一个数据集的每一行都引用第二个数据集的每一行（第一行与第一行，第二行与第二行，等等。 )
如果他们没有公用键，你怎么知道df1[row1]对应df2[row1]？每次加载两个数据集的顺序是否一致？一般来说，使用/应用值从一个数据集到另一个数据集的唯一方法是通过连接，正如@RaphaelRoth 已经提到的那样。如果两个数据集的顺序是确定的，那么您可以使用 rowNumber 之类的东西为两个数据集添加一个公共 id，最后加入它们。
正如你所说，顺序是确定性的

标签： scala apache-spark apache-spark-sql

【解决方案1】：

您可以压缩DataFrames 并过滤这些元组。

val ints     = sparkSession.sparkContext.parallelize(List(1, 2, 3, 4, 5, 6, 7, 8, 9, 10))
val bools    = sparkSession.sparkContext.parallelize(List(true, false, true, false, true, false, true, false, true, false))
val filtered = ints.zip(bools).filter { case (int, bool) => bool }.map { case (int, bool) => int }
println(filtered.collect().toList) //List(1, 3, 5, 7, 9)

【讨论】：

这可以用ints.zip(bools).collect{ case (int, bool) if bool => int }更简洁地完成
我编辑了这个问题，对预期的输出和输入进行了一些额外的规范，因为建议的答案没有返回与df2具有相同架构的数据框

【解决方案2】：

我设法用以下代码解决了这个问题：

import org.apache.spark.sql.types.{LongType, StructField, StructType}
import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.sql.{DataFrame, Row, SQLContext, SparkSession}

val spark = SparkSession.builder().appName(sc.appName).master(sc.master).getOrCreate()
val sqlContext = spark.sqlContext


def addColumnIndex(df: DataFrame, sqlContext: SQLContext) = sqlContext.createDataFrame(
  // Add Column index
  df.rdd.zipWithIndex.map{case (row, columnindex) => Row.fromSeq(row.toSeq :+ columnindex)},
  // Create schema
  StructType(df.schema.fields :+ StructField("columnindex", LongType, nullable = false))
)

import spark.implicits._

val DF1 = Seq(
  ("A", true),
  ("B", false),
  ("A", true),
  ("C", false)
).toDF("col1", "boolean_column")

val DF2 = Seq(
  ("Z", 1),
  ("X", 2),
  ("Y", 3),
  ("W", 4)
).toDF("col_1", "col_2")

// Add index
val DF1WithIndex = addColumnIndex(DF1, sqlContext)
val DF2WithIndex = addColumnIndex(DF2, sqlContext)

// Join
val joinDF  = DF2WithIndex
  .join(DF1WithIndex, Seq("columnindex"))
  .drop("columnindex", "col1")

// Filter
val filteredDF2 = joinDF.filter(joinDF("boolean_column")).drop("boolean_column")

过滤后的数据框如下：

+-----+-----+
|col_1|col_2|
+-----+-----+
|    Z|    1|
|    Y|    3|
+-----+-----+

【讨论】：