如何选择以公共标签开头的所有列

Posted 2023-03-23

【中文标题】如何选择以公共标签开头的所有列

【英文标题】：how to select all columns that starts with a common label

【发布时间】：2016-05-22 06:59:57

【问题描述】：

我在 Spark 1.6 中有一个数据框，想从中选择一些列。列名如下：

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

我知道我可以这样做来选择特定的列：

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

但是如何同时选择“colA”、“colB”和所有 colF-* 列？有没有像Pandas这样的方法？

【参考方案1】：

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

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

【讨论】：

谢谢，但如何选择其他列，如问题中所述？

我的问题的确切解决方案来自您的第二个和第三个解决方案的混合...非常感谢，您完全掌握了这个主题

pyspark 中的任何即用型解决方案（python 不是 scala）？

【参考方案2】：

Python（在 Azure Databricks 中测试）

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

【参考方案3】：

在 PySpark 中，使用：colRegex 选择以 colF 开头的列 打样：

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

申请：

df.select(col("colA"), col("colB"), df.colRegex("`(colF)+?.+`")).show()

结果是：

colA, colB, colF-0, colF-1, colF-2

【参考方案4】：

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

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