scala - 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访问( constant time access for each index)，在 String或基于 Column的 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]