In this article, we will discuss how to select only numeric or string column names from a Spark DataFrame.
Methods Used:
- createDataFrame: This method is used to create a spark DataFrame.
- isinstance: This is a Python function used to check if the specified object is of the specified type.
- dtypes: It returns a list of tuple (columnName,type). The returned list contains all columns present in DataFrame with their data types.
- schema.fields: It is used to access DataFrame fields metadata.
Method #1:
In this method, dtypes function is used to get a list of tuple (columnName, type).
Python3
from pyspark.sql import Rowfrom datetime import datefrom pyspark.sql import SparkSessionspark = SparkSession.builder.getOrCreate()# Creating dataframe from list of Rowdf = spark.createDataFrame([ Row(a=1, b='string1', c=date(2021, 1, 1)), Row(a=2, b='string2', c=date(2021, 2, 1)), Row(a=4, b='string3', c=date(2021, 3, 1))])# Printing DataFrame structureprint("DataFrame structure:", df)# Getting list of columns and printing# resultdt = df.dtypesprint("dtypes result:", dt)# Getting list of columns having type# string or bigint# This statement will loop over all the # tuples present in dt list# item[0] will contain column name and# item[1] will contain column typecolumnList = [item[0] for item in dt if item[1].startswith( 'string') or item[1].startswith('bigint')]print("Result: ", columnList) |
Output:
DataFrame structure: DataFrame[a: bigint, b: string, c: date]
dtypes result: [('a', 'bigint'), ('b', 'string'), ('c', 'date')]
Result: ['a', 'b']
Method #2:
In this method schema.fields is used to get fields metadata then column data type is extracted from metadata and compared with the desired data type.
Python3
from pyspark.sql.types import StringType, LongTypefrom pyspark.sql import Rowfrom datetime import datefrom pyspark.sql import SparkSession# Initializing spark sessionspark = SparkSession.builder.getOrCreate()# Creating dataframe from list of Rowdf = spark.createDataFrame([ Row(a=1, b='string1', c=date(2021, 1, 1)), Row(a=2, b='string2', c=date(2021, 2, 1)), Row(a=4, b='string3', c=date(2021, 3, 1))])# Printing DataFrame structureprint("DataFrame structure:", df)# Getting and printing metadatameta = df.schema.fieldsprint("Metadata: ", meta)# Getting list of columns having type # string or int# This statement will loop over all the fields# field.name will return column name and# field.dataType will return column typecolumnList = [field.name for field in df.schema.fields if isinstance( field.dataType, StringType) or isinstance(field.dataType, LongType)]print("Result: ", columnList) |
Output:
DataFrame structure: DataFrame[a: bigint, b: string, c: date]
Metadata: [StructField(a,LongType,true), StructField(b,StringType,true), StructField(c,DateType,true)]
Result: [‘a’, ‘b’]
