scala - 从 Spark 数据框中选择几列，并以列列表作为开始

Question

假设我有一个 spark 列列表和一个 spark 数据框 df，为了选择仅包含列表中的列的子数据框，合适的代码片段是什么？

类似于可能的东西:

var needed_column: List[Column]=List[Column](new Column("a"),new Column("b"))

df(needed_columns)

我想获取列名，然后使用以下代码行选择它们。

不幸的是，列名似乎只处于写入模式。

df.select(needed_columns.head.as(String),needed_columns.tail: _*)

Answer 1

您的 needed_columns 属于 List[Column] 类型，因此您可以简单地使用 needed_columns: _* 作为 的参数选择:

val df = Seq((1, "x", 10.0), (2, "y", 20.0)).toDF("a", "b", "c")

import org.apache.spark.sql.Column

val needed_columns: List[Column] = List(new Column("a"), new Column("b"))

df.select(needed_columns: _*)
// +---+---+
// |  a|  b|
// +---+---+
// |  1|  x|
// |  2|  y|
// +---+---+

请注意 select接受两种类型的参数:

def select(cols: Column*): DataFrame

def select(col: String, cols: String*): DataFrame

如果你有String类型的列名列表，可以使用后面的select:

val needed_col_names: List[String] = List("a", "b")

df.select(needed_col_names.head, needed_col_names.tail: _*)

或者，您可以将 String 列表映射到 Column 以使用前一个 select

df.select(needed_col_names.map(col): _*)