如何选择以通用标签开头的所有列

时间:2016-02-11 13:15:53

标签: scala apache-spark spark-dataframe

我在Spark 1.6中有一个数据框,并希望从中选择一些列。列名称如下:

colA, colB, colC, colD, colE, colF-0, colF-1, colF-2

我知道我可以这样选择特定列:

df.select("colA", "colB", "colE")

但如何选择,说" colA"," colB"和所有colF- *列一次?有没有像Pandas那样的方式?

3 个答案:

答案 0 :(得分:17)

首先使用df.columns获取列名称,然后向下过滤到您想要的列名称.filter(_.startsWith("colF"))。这为您提供了一系列字符串。但是选择需要select(String, String*)。幸运的是,选择列为select(Column*),最后将字符串转换为.map(df(_))的列,最后将列数组转换为带: _*的变量。

df.select(df.columns.filter(_.startsWith("colF")).map(df(_)) : _*).show

此过滤器可能更复杂(与Pandas相同)。然而,它是一个相当丑陋的解决方案(IMO):

df.select(df.columns.filter(x => (x.equals("colA") || x.startsWith("colF"))).map(df(_)) : _*).show 

如果修复了其他列的列表,您还可以将固定的列名称数组与已过滤的数组合并。

df.select((Array("colA", "colB") ++ df.columns.filter(_.startsWith("colF"))).map(df(_)) : _*).show

答案 1 :(得分:2)

Python(在 Azure Databricks 中测试)

selected_columns = [column for column in df.columns if column.startswith("colF")]
df2 = df.select(selected_columns)

答案 2 :(得分:-1)

我写了一个函数来做到这一点。阅读评论以了解其工作原理。

  /**
    * Given a sequence of prefixes, select suitable columns from [[DataFrame]]
    * @param columnPrefixes Sequence of prefixes
    * @param dF Incoming [[DataFrame]]
    * @return [[DataFrame]] with prefixed columns selected
    */
  def selectPrefixedColumns(columnPrefixes: Seq[String], dF: DataFrame): DataFrame = {
    // Find out if given column name matches any of the provided prefixes
    def colNameStartsWith: String => Boolean = (colName: String) =>
        columnsPrefix.map(prefix => colName.startsWith(prefix)).reduce(_ || _)
    // Filter columns list by checking against given prefixes sequence
    val columns = dF.columns.filter(colNameStartsWith)
    // Select filtered columns list
    dF.select(columns.head, columns.tail:_*)
  }