How To Select Multiple Columns From PySpark DataFrames | Towards Data Science

When working with Spark, we typically need to deal with a fairly large number of rows and columns and thus, we sometimes have to work only with a small subset of columns.
In today’s short guide we will explore different ways for selecting columns from PySpark DataFrames. Specifically, we will discuss how to select multiple columns
by column name
by index
with the use of regular expressions
First, let’s create an example DataFrame that we’ll reference throughout this article to demonstrate a few concepts.
from pyspark.sql import SparkSession# Create an instance of spark sessionspark_session = SparkSession.builder \    .master('local[1]') \    .appName('Example') \    .getOrCreate()# Create an example DataFramedf = spark_session.createDataFrame(    [       (1, 'a', True, 1.0, 5),       (2, 'b', False, 2.0, None),       (3, 'c', False, 3.0, 4),       (4, 'd', True, 4.0, 3),    ],    ['colA', 'colB', 'colC', 'colD', 'E'])df.show()+----+----+-----+----+----+|colA|colB| colC|colD|   E|+----+----+-----+----+----+|   1|   a| true| 1.0|   5||   2|   b|false| 2.0|null||   3|   c|false| 3.0|   4||   4|   d| true| 4.0|   3|+----+----+-----+----+----+
In order to select multiple column from an existing PySpark DataFrame you can simply specify the column names you wish to retrieve to the pyspark.sql.DataFrame.select method. For example,
df.select('colA', 'colC').show()+----+-----+|colA| colC|+----+-----+|   1| true||   2|false||   3|false||   4| true|+----+-----+
Alternatively, if the columns you wish to retrieve are stored in a list, you can use the following notation:
col_names = ['colA', 'colC']df.select(*col_names).show()+----+-----+|colA| colC|+----+-----+|   1| true||   2|false||   3|false||   4| true|+----+-----+
Now if you want to select columns based on their index, then you can simply slice the result from df.columns that returns a list of column names. For example, in order to retrieve the first three columns then the following expression should do the trick:
df.select(df.columns[:3]).show()+----+----+-----+|colA|colB| colC|+----+----+-----+|   1|   a| true||   2|   b|false||   3|   c|false||   4|   d| true|+----+----+-----+
Alternatively, if you want to fetch the 2nd and 3rd column only, then df.columns[1:3] will do the trick:
df.select(df.columns[1:3]).show()+----+-----+|colB| colC|+----+-----+|   a| true||   b|false||   c|false||   d| true|+----+-----+
Finally, in order to select multiple columns that match a specific regular expression then you can make use of pyspark.sql.DataFrame.colRegex method. For instance, in order to fetch all the columns that start with or contain col, then the following will do the trick:
df.select(df.colRegex("`(col)+?.+`")).show()+----+----+-----+----+|colA|colB| colC|colD|+----+----+-----+----+|   1|   a| true| 1.0||   2|   b|false| 2.0||   3|   c|false| 3.0||   4|   d| true| 4.0|+----+----+-----+----+
Similarly, we can use the following regex in order to select all columns but colA :
df.select(df.colRegex("`(colA)?+.+`")).show()+----+-----+----+----+|colB| colC|colD|   E|+----+-----+----+----+|   a| true| 1.0|   5||   b|false| 2.0|null||   c|false| 3.0|   4||   d| true| 4.0|   3|+----+-----+----+----+
In today’s short guide we discussed how to perform column selection in PySpark DataFrames. We explored how to select multiple columns by specifying the column name or index. Additionally, we saw in action how to perform column selection with the use of regular expressions.
Become a member and read every story on Medium. Your membership fee directly supports me and other writers you read.
You may also like