Selecting only numeric or string columns names from PySpark DataFrame - GeeksforGeeks

30 Apr, 2021
In this article, we will discuss how to select only numeric or string column names from a Spark DataFrame.
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 (columnNane,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 (columnNane, type).
Python3
from pyspark.sql import Rowfrom datetime import datefrom pyspark.sql import SparkSession    spark = 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’]
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 ?
Check if element exists in list in Python
How To Convert Python Dictionary To JSON?
Python Classes and Objects
How to drop one or multiple columns in Pandas Dataframe
Defaultdict in Python
Python | Get unique values from a list
Python | os.path.join() method
Create a directory in Python
Python | Pandas dataframe.groupby()