基于列索引的 Spark Dataframe 选择

Question

如何选择在 Scala 中具有特定索引的数据帧的所有列？

例如，如果一个数据框有 100 列，而我只想提取列（10、12、13、14、15），该怎么做？

下面从数据框df选择所有列，这些列具有 Array colNames 中提到的列名：

df = df.select(colNames.head,colNames.tail: _*)

如果有类似的，colNos 数组，其中有

colNos = Array(10,20,25,45)

如何转换上面的df.select以仅获取特定索引处的那些列。

Answer 1

您可以map columns ：

import org.apache.spark.sql.functions.col

df.select(colNos map df.columns map col: _*)

或者：

df.select(colNos map (df.columns andThen col): _*)

或者：

df.select(colNos map (col _ compose df.columns): _*)

上面显示的所有方法都是等效的，不会造成性能损失。 以下映射：

colNos map df.columns 

只是一个本地Array访问（ 每个索引的恒定时间访问）并且在基于String或Column的select变体之间进行select不会影响执行计划：

val df = Seq((1, 2, 3 ,4, 5, 6)).toDF

val colNos = Seq(0, 3, 5)

df.select(colNos map df.columns map col: _*).explain

== Physical Plan ==
LocalTableScan [_1#46, _4#49, _6#51]

df.select("_1", "_4", "_6").explain

== Physical Plan ==
LocalTableScan [_1#46, _4#49, _6#51]

Answer 2

@ user6910411 上面的回答很有魅力，任务/逻辑计划的数量与我下面的方法相似。 但是我的方法要快一些。
所以，
我建议你使用column names而不是column numbers 。 Column names比使用numbers更安全、更轻松。 您可以使用以下解决方案：

val colNames = Seq("col1", "col2" ...... "col99", "col100")

val selectColNames = Seq("col1", "col3", .... selected column names ... )

val selectCols = selectColNames.map(name => df.col(name))

df = df.select(selectCols:_*)

如果您对写所有 100 个列名犹豫不决，那么也有一个快捷方法

val colNames = df.schema.fieldNames

Answer 3

示例：使用 Scala 按索引抓取 Spark Dataframe 的前 14 列。

import org.apache.spark.sql.functions.col

// Gives array of names by index (first 14 cols for example)
val sliceCols = df.columns.slice(0, 14)
// Maps names & selects columns in dataframe
val subset_df = df.select(sliceCols.map(name=>col(name)):_*)

你不能简单地这样做（因为我尝试过但失败了）：

// Gives array of names by index (first 14 cols for example)
val sliceCols = df.columns.slice(0, 14)
// Maps names & selects columns in dataframe
val subset_df = df.select(sliceCols)

原因是您必须将 Array[String] 的数据类型转换为 Array[org.apache.spark.sql.Column] 才能进行切片。

或使用 Currying 将其包装在一个函数中（为此向我的同事致敬）：

// Subsets Dataframe to using beg_val & end_val index.
def subset_frame(beg_val:Int=0, end_val:Int)(df: DataFrame): DataFrame = {
  val sliceCols = df.columns.slice(beg_val, end_val)
  return df.select(sliceCols.map(name => col(name)):_*)
}

// Get first 25 columns as subsetted dataframe
val subset_df:DataFrame = df_.transform(subset_frame(0, 25))