Skip to content
Related Articles

Related Articles

Improve Article

Selecting only numeric or string columns names from PySpark DataFrame

  • Last Updated : 30 Apr, 2021

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 (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 Row
from datetime import date
from pyspark.sql import SparkSession
  
  
spark = SparkSession.builder.getOrCreate()
  
# Creating dataframe from list of Row
df = 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 structure
print("DataFrame structure:", df)
  
# Getting list of columns and printing
# result
dt = df.dtypes
print("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 type
columnList = [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, LongType
from pyspark.sql import Row
from datetime import date
from pyspark.sql import SparkSession
  
  
# Initializing spark session
spark = SparkSession.builder.getOrCreate()
  
# Creating dataframe from list of Row
df = 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 structure
print("DataFrame structure:", df)
  
# Getting and printing metadata
meta = df.schema.fields
print("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 type
columnList = [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’]

 Attention geek! Strengthen your foundations with the Python Programming Foundation Course and learn the basics.  

To begin with, your interview preparations Enhance your Data Structures concepts with the Python DS Course. And to begin with your Machine Learning Journey, join the Machine Learning – Basic Level Course




My Personal Notes arrow_drop_up
Recommended Articles
Page :