¿Cómo obtener una substring de una columna en PySpark Dataframe?

En este artículo, veremos cómo obtener la substring de la columna PySpark Dataframe y cómo crear la nueva columna y colocar la substring en esa columna recién creada.

Podemos obtener la substring de la columna usando la función substring() y substr() .

Sintaxis: substring(str,pos,len)

df.col_name.substr(inicio, longitud)

Parámetro:

• str: puede ser una string o el nombre de la columna de la que obtenemos la substring.
• inicio y pos – A través de este parámetro podemos dar la posición de inicio desde donde se inicia la substring.
• length y len: es la longitud de la substring desde la posición inicial.

Vamos a crear un marco de datos.

# importing necessary libraries
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, substring
 
 
# function to create new SparkSession
def create_session():
    spk = SparkSession.builder \
        .master("local") \
        .appName("Substring.com") \
        .getOrCreate()
    return spk
 
def create_df(spark, data, schema):
 
    df1 = spark.createDataFrame(data, schema)
    return df1
 
if __name__ == "__main__":
 
    input_data = [("India", +91, 2701, 2020),
                  ("United States of America", +1, 1301, 2020),
                  ("Israel", +972, 3102, 2020),
                  ("Dubai", +971, 2901, 2020),
                  ("Russia", 7, 3101, 2020)]
 
    # calling function to create SparkSession
    spark = create_session()
 
    schema = ["Country", "Country Code",
              "Data", "Year"]
 
    # calling function to create dataframe
    df = create_df(spark, input_data, schema)
    df.show()

Producción:

Ejemplo 1: Usando substring() obteniendo la substring y creando una nueva columna usando la función withColumn().

if __name__ == "__main__":
   
    # creating Month column and get the
    # substring from the Data column
    # creating Date column and get the
    # substring from the Data column
    df = df.withColumn(
      "Month", substring("Data", 1, 2)).withColumn(
      "Date", substring("Data", 3, 4))
 
    # dropping the Data column from the
    # Dataframe
    df = df.drop("Data")
 
    # printing Dataframe schema to get the
    # column names
    df.printSchema()
 
    # visualizing the dataframe
    df.show(truncate=False)

Producción:

Ejemplo 2: Crear la columna New_Country obteniendo la substring usando la función substr().

if __name__ == "__main__":
 
    # Creating the new column New_Country
    # and store the substring using substr()
    df = df.withColumn("New_Country", df.Country.substr(0, 12))
 
    # printing Dataframe schema to get the
    # column names
    df.printSchema()
 
    # visualizing the dataframe
    df.show(truncate=False)

Producción:

Ejemplo 3: Usar substring() con la función select().

if __name__ == "__main__":
 
    input_data = [("India", +91, "AidanButler"),
                  ("United States of America", +1, "ConerFlores"),
                  ("Israel", +972, "RosseBryant"),
                  ("Dubai", +971, "JuliaSimmon"),
                  ("Russia", 7, "AliceBailey")]
 
    # calling function to create SparkSession
    spark = create_session()
 
    schema = ["Country", "Country Code", "Name"]
 
    # calling function to create dataframe
    df = create_df(spark, input_data, schema)
 
    # Selecting the column using select()
    # function and getting substring
    # using substring()
    df2 = df.select('Name', substring('Name',
                                      1, 5).alias('First Name'),
                    substring('Name', 6, 6).alias('Last Name'))
 
    # printing Dataframe schema to get the column names
    df2.printSchema()
 
    # visualizing the dataframe
    df2.show(truncate=False)

Producción:

Ejemplo 4: Usar substring() con la función selectExpr().

if __name__ == "__main__":
 
    input_data = [("India", +91, "AidanButler"),
                  ("United States of America", +1, "ConerFlores"),
                  ("Israel", +972, "RosseBryant"),
                  ("Dubai", +971, "JuliaSimmon"),
                  ("Russia", 7, "AliceBailey")]
 
    # calling function to create SparkSession
    spark = create_session()
 
    schema = ["Country", "Country Code", "Name"]
 
    # calling function to create dataframe
    df = create_df(spark, input_data, schema)
 
    # Selecting the column using selectExpr()
    # function and getting substring using substring()
    df2 = df.selectExpr('Name', 'substring(Name, 1,5) as First_Name',
                        'substring(Name, 6,6) as Last_Name')
 
    # printing Dataframe schema to get the column names
    df2.printSchema()
 
    # visualizing the dataframe
    df2.show(truncate=False)

Producción: