PySpark - Select columns by type - GeeksforGeeks

19 Dec, 2021
In this article, we will discuss how to select columns by type in PySpark using Python.
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # import data field typesfrom pyspark.sql.types import StringType, DoubleType, IntegerType, StructType, StructField, FloatType    # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of student  datadata = [(1, "sravan", 9.8, 4500.00), (2, "ojsawi",                                      9.2, 6789.00),        (3, "bobby", 8.9, 988.000)]  # specify column names with data typescolumns = StructType([    StructField("ID", IntegerType(), True),    StructField("NAME", StringType(), True),    StructField("GPA", FloatType(), True),    StructField("FEE", DoubleType(), True),  ])  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)    # displaydataframe.show()
Output:
We can select the column by name using the following keywords:
Integer: int
String : string
Float: float
Double: double
Here we are using dtypes followed by startswith() method to get the columns of a particular type.
Syntax: dataframe[[item[0] for item in dataframe.dtypes if item[1].startswith(‘datatype’)]]
where,
dataframe is the input dataframe
datatype refers the keyword types
item defines the values in the column
And finally, we are using collect() method to display column data
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # import data field typesfrom pyspark.sql.types import (StringType,DoubleType, IntegerType, StructType,StructField, FloatType)    # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of student  datadata = [(1, "sravan", 9.8, 4500.00), (2, "ojsawi",                                      9.2, 6789.00),         (3, "bobby", 8.9, 988.000)]  # specify column names with data typescolumns = StructType([    StructField("ID", IntegerType(), True),    StructField("NAME", StringType(), True),    StructField("GPA", FloatType(), True),    StructField("FEE", DoubleType(), True),  ])  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)  # select columns that are integer typeprint(dataframe[[item[0]                 for item in dataframe.dtypes if item[                   1].startswith('int')]].collect())  # select columns that are string typeprint(dataframe[[item[0]                 for item in dataframe.dtypes if item[                   1].startswith('string')]].collect())  # select columns that are float typeprint(dataframe[[item[0]                 for item in dataframe.dtypes if item[                   1].startswith('float')]].collect())  # select columns that are double typeprint(dataframe[[item[0]                 for item in dataframe.dtypes if item[                   1].startswith('double')]].collect())
Output:
[Row(ID=1), Row(ID=2), Row(ID=3)]
[Row(NAME=’sravan’), Row(NAME=’ojsawi’), Row(NAME=’bobby’)]
[Row(GPA=9.800000190734863), Row(GPA=9.199999809265137), Row(GPA=8.899999618530273)]
[Row(FEE=4500.0), Row(FEE=6789.0), Row(FEE=988.0)]
Here we are using schema.fields method to get the type of the columns. We are checking the particular type using methods that are available in pyspark.sql.types module.
Let’s check one by one:
Integer – IntegerType
Float-FloatType
Double – DoubleType
String- StringType
We are using isinstance() operator to check with these data types.
Syntax: dataframe[[f.name for f in dataframe.schema.fields if isinstance(f.dataType, datatype)]]
where,
dataframe is the input dataframe
name is the values
datatype refers to above types
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # import data field typesfrom pyspark.sql.types import StringType, DoubleType,IntegerType, StructType, StructField, FloatType    # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of student  datadata = [(1, "sravan", 9.8, 4500.00), (2, "ojsawi",                                      9.2, 6789.00),         (3, "bobby", 8.9, 988.000)]  # specify column names with data typescolumns = StructType([    StructField("ID", IntegerType(), True),    StructField("NAME", StringType(), True),    StructField("GPA", FloatType(), True),    StructField("FEE", DoubleType(), True),  ])  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)  # select columns that are integer typeprint(dataframe[[f.name for f in dataframe.schema.fields if isinstance(    f.dataType, IntegerType)]].collect())  # select columns that are string typeprint(dataframe[[f.name for f in dataframe.schema.fields if isinstance(    f.dataType, StringType)]].collect())  # select columns that are float typeprint(dataframe[[f.name for f in dataframe.schema.fields if isinstance(    f.dataType, FloatType)]].collect())  # select columns that are double typeprint(dataframe[[f.name for f in dataframe.schema.fields if isinstance(    f.dataType, DoubleType)]].collect())
Output:
[Row(ID=1), Row(ID=2), Row(ID=3)]
[Row(NAME=’sravan’), Row(NAME=’ojsawi’), Row(NAME=’bobby’)]
[Row(GPA=9.800000190734863), Row(GPA=9.199999809265137), Row(GPA=8.899999618530273)]
[Row(FEE=4500.0), Row(FEE=6789.0), Row(FEE=988.0)]
Picked
Python-Pyspark
Python
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
How to Install PIP on Windows ?
How To Convert Python Dictionary To JSON?
How to drop one or multiple columns in Pandas Dataframe
Check if element exists in list in Python
Python | os.path.join() method
Selecting rows in pandas DataFrame based on conditions
Defaultdict in Python
Python | Get unique values from a list
Create a directory in Python
Python | Pandas dataframe.groupby()