Question

假设我有一个spark列的列表和一个spark dataframe 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

我了解到，您只想从列表（A）中选择那些列，而不是数据框列。我有一个下面的示例，其中使用单独的列表选择名字和姓氏。检查一下

scala> val df = Seq((101,"Jack", "wright" , 27, "01976", "US")).toDF("id","fname","lname","age","zip","country")
df: org.apache.spark.sql.DataFrame = [id: int, fname: string ... 4 more fields]

scala> df.columns
res20: Array[String] = Array(id, fname, lname, age, zip, country)

scala> val needed =Seq("fname","lname")
needed: Seq[String] = List(fname, lname)

scala> val needed_df = needed.map( x=> col(x) )
needed_df: Seq[org.apache.spark.sql.Column] = List(fname, lname)

scala> df.select(needed_df:_*).show(false)
+-----+------+
|fname|lname |
+-----+------+
|Jack |wright|
+-----+------+


scala>

Answer 2

您的needed_columns类型为List[Column]，因此您可以简单地将needed_columns: _*用作select的参数：

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): _*)

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

2 个答案: