从Spark DataFrame中选择特定列

Question

我已将CSV数据加载到Spark DataFrame中。

我需要将此数据帧切成两个不同的数据帧，其中每个数据帧都包含来自原始数据帧的一组列。

如何基于列选择Spark数据框中的子集？

Answer 1

如果要将数据帧分为两个不同的数据帧，请使用所需的不同列对其进行两个选择。

 val sourceDf = spark.read.csv(...)
 val df1 = sourceDF.select("first column", "second column", "third column")
 val df2 = sourceDF.select("first column", "second column", "third column")

请注意，这当然意味着对sourceDf进行两次评估，因此，如果它可以放入分布式内存中并且您在两个数据帧中使用大多数列，则缓存它可能是一个好主意。它有许多不需要的多余列，然后您可以先对其进行选择，然后再选择需要的列，这样它将所有这些多余的数据存储在内存中。

Answer 2

有多种选项（尤其是在 Scala 中）可以选择该 Dataframe 的列的子集。以下几行显示了选项，其中大部分都记录在Column的 ScalaDocs 中：

import spark.implicits._
import org.apache.spark.sql.functions.{col, column, expr}

inputDf.select(col("colA"), col("colB"))
inputDf.select(inputDf.col("colA"), inputDf.col("colB"))
inputDf.select(column("colA"), column("colB"))
inputDf.select(expr("colA"), expr("colB"))

// only available in Scala
inputDf.select($"colA", $"colB")
inputDf.select('colA, 'colB) // makes use of Scala's Symbol

// selecting columns based on a given iterable of Strings
val selectedColumns: Seq[Column] = Seq("colA", "colB").map(c => col(c))
inputDf.select(selectedColumns: _*)

// Special cases
col("columnName.field")     // Extracting a struct field
col("`a.column.with.dots`") // Escape `.` in column names.

// select the first or last 2 columns
inputDf.selectExpr(inputDf.columns.take(2): _*)
inputDf.selectExpr(inputDf.columns.takeRight(2): _*)

的使用$是可能的，因为 Scala 提供了一个隐式类，该类使用以下方法将字符串转换为列$：

implicit class StringToColumn(val sc : scala.StringContext) extends scala.AnyRef {
  def $(args : scala.Any*) : org.apache.spark.sql.ColumnName = { /* compiled code */ }
}

通常，当您想将一个 DataFrame 派生为多个 DataFrame 时，如果您persist在创建其他 DataFrame 之前使用原始 DataFrame，则可能会提高您的性能。最后，您可以unpersist使用原始 DataFrame。

请记住，列不会在编译时解析，而只会在与在查询执行的分析器阶段发生的目录的列名进行比较时进行解析。如果您需要更强的类型安全性，您可以创建一个Dataset.

为了完整起见，这里是用于尝试上述代码的 csv：

// csv file:
// colA,colB,colC
// 1,"foo","bar"

val inputDf = spark.read.format("csv").option("header", "true").load(csvFilePath)

// resulting DataFrame schema
root
 |-- colA: string (nullable = true)
 |-- colB: string (nullable = true)
 |-- colC: string (nullable = true)

Answer 3

假设我们的父 Dataframe 有“n”列

我们可以创建“x”个子数据帧（在我们的例子中让我们考虑 2）。

可以根据需要从任何父 Dataframe 列中选择子 Dataframe 的列。

考虑源有10 列，我们希望拆分为2 个数据帧，其中包含从父数据帧引用的列。

可以使用select Dataframe API决定子 Dataframe 的列

val parentDF = spark.read.format("csv").load("/path of the CSV file")

val Child1_DF = parentDF.select("col1","col2","col3","col9","col10").show()

val child2_DF = parentDF.select("col5", "col6","col7","col8","col1","col2").show()

请注意，子数据帧中的列数的长度可能不同，并且会小于父数据帧的列数。

我们还可以使用父数据帧中所需列的位置索引来引用列名，而无需提及真实名称

导入 spark 隐式首先充当帮助类，用于使用 $-notation 来访问使用位置索引的列

import spark.implicits._
import org.apache.spark.sql.functions._

val child3_DF  = parentDF.select("_c0","_c1","_c2","_c8","_c9").show()

我们也可以根据特定条件选择列。假设我们只想在子数据框中选择偶数列。通过 even 我们指的是偶数索引列和索引从“0”开始

val parentColumns = parentDF.columns.toList


res0: List[String] = List(_c0, _c1, _c2, _c3, _c4, _c5, _c6, _c7,_c8,_c9)

val evenParentColumns =  res0.zipWithIndex.filter(_._2 % 2 == 0).map( _._1).toSeq

res1: scala.collection.immutable.Seq[String] = List(_c0, _c2, _c4, _c6,_c8)

现在提供这些要从 parentDF 中选择的列。注意选择 API 需要 seq 类型参数。所以我们将“evenParentColumns”转换为 Seq 集合

val child4_DF = parentDF.select(res1.head, res1.tail:_*).show()

这将显示来自父 Dataframe 的偶数索引列。

| _c0 | _c2 | _c4 |_c6 |_c8 |

|ITE00100554|TMAX|空| E| 1 |

|TE00100554 |TMIN|空| E| 4 |

|GM000010962|PRCP|空| E| 7 |

所以现在我们在数据框中留下偶数列

同样，我们也可以对 Dataframe 列应用其他操作，如下所示

val child5_DF = parentDF.select($"_c0", $"_c8" + 1).show()

因此，通过前面提到的多种方式，我们可以选择 Dataframe 中的列。