how to select all columns that starts with a common label

Question

I have a dataframe in Spark 1.6 and want to select just some columns out of it. The column names are like:

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

I know I can do like this to select specific columns:

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

but how to select, say "colA", "colB" and all the colF-* columns at once? Is there a way like in Pandas?

score 33 · Accepted Answer · edited Feb 06 '23 at 10:52

33

The process canbe broken down into following steps:

First grab the column names with df.columns,
then filter down to just the column names you want .filter(_.startsWith("colF")). This gives you an array of Strings.
But the select takes select(String, String*). Luckily select for columns is select(Column*), so finally convert the Strings into Columns with .map(df(_)),
and finally turn the Array of Columns into a var arg with : _*.

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

This filter could be made more complex (same as Pandas). It is however a rather ugly solution (IMO):

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

If the list of other columns is fixed you could also merge a fixed array of columns names with filtered array.

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

edited Feb 06 '23 at 10:52

Vivek Puurkayastha

466
1
9
18

answered Feb 11 '16 at 14:07

Michael Lloyd Lee mlk

14,561
3
44
81

thanks but how to select also other columns, like stated in the question? – user299791 Feb 11 '16 at 14:22
the exact solution to my problem comes out of a mixture of your second and third solution... thanks so much, you totally master the topic – user299791 Feb 11 '16 at 14:33
1

any ready-to-use solution in pyspark (python not scala) ? – Eugene Lycenok Mar 03 '21 at 02:02
Could you explain what this underscore does? – Vandit Goel Sep 20 '22 at 19:56
Also I get the error `AttributeError: 'list' object has no attribute 'filter'` – Vandit Goel Sep 20 '22 at 20:20

score 13 · Answer 2 · answered Mar 04 '21 at 02:39

13

Python (tested in Azure Databricks)

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

answered Mar 04 '21 at 02:39

Eugene Lycenok

603
6
14

score 1 · Answer 3 · answered Dec 22 '21 at 02:18

In PySpark, use: colRegex to select columns starting with colF Whit the sample:

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

Apply:

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

The result is:

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

score -1 · Answer 4 · answered Jan 13 '20 at 15:49

I wrote a function that does that. Read the comments to see how it works.

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

how to select all columns that starts with a common label

4 Answers4

Linked