Seleccionar solo nombres de columnas numéricas o de string de PySpark DataFrame

En este artículo, analizaremos cómo seleccionar solo nombres de columna numéricos o de string de un Spark DataFrame.

Métodos utilizados:

• createDataFrame: este método se utiliza para crear un Spark DataFrame.
• isinstance: esta es una función de Python utilizada para verificar si el objeto especificado es del tipo especificado.
• dtypes: Devuelve una lista de tuplas (columnNane,type). La lista devuelta contiene todas las columnas presentes en DataFrame con sus tipos de datos.
• schema.fields: Se utiliza para acceder a los metadatos de los campos de DataFrame.

Método 1: 

En este método, la función dtypes se usa para obtener una lista de tuplas (columnNane, type).

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)

Producción:

DataFrame structure: DataFrame[a: bigint, b: string, c: date]
dtypes result: [('a', 'bigint'), ('b', 'string'), ('c', 'date')]
Result:  ['a', 'b']

Método #2: 

En este esquema de método. los campos se utilizan para obtener los metadatos de los campos, luego el tipo de datos de la columna se extrae de los metadatos y se compara con el tipo de datos deseado.

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)

Producción:

Estructura de DataFrame: DataFrame[a: bigint, b: string, c: fecha]

Metadatos: [StructField(a,LongType,true), StructField(b,StringType,true), StructField(c,DateType,true)]

Resultado: [‘a’, ‘b’]