PySpark dataframe agregar columna basada en otras columnas

En este artículo, veremos cómo agregar columnas basadas en otra columna al marco de datos de Pyspark.

Creando Dataframe para demostración:

Aquí vamos a crear un marco de datos a partir de una lista del conjunto de datos dado.

# Create a spark session
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('SparkExamples').getOrCreate()
  
# Create a spark dataframe
columns = ["Name", "Course_Name",
           "Months",
           "Course_Fees", "Discount",
           "Start_Date", "Payment_Done"]
data = [
    ("Amit Pathak", "Python", 3, 10000, 1000,
     "02-07-2021", True),
    ("Shikhar Mishra", "Soft skills", 2,
     8000, 800, "07-10-2021", False),
    ("Shivani Suvarna", "Accounting", 6,
     15000, 1500, "20-08-2021", True),
    ("Pooja Jain", "Data Science", 12,
     60000, 900, "02-12-2021", False),
]
  
df = spark.createDataFrame(data).toDF(*columns)
  
# View the dataframe
df.show()

Producción:

Método 1: Usando withColumns()

Se utiliza para cambiar el valor, convertir el tipo de datos de una columna existente, crear una nueva columna y mucho más.

Sintaxis: df.withColumn(colName, col)

Devuelve: una nueva :class:`DataFrame` agregando una columna o reemplazando la columna existente que tiene el mismo nombre. 

new_df = df.withColumn('After_discount',
                       df.Course_Fees - df.Discount)
new_df.show()

Producción:

Método 2: Usar consulta SQL

Aquí usaremos la consulta SQL dentro de Pyspark. Crearemos una vista temporal de la tabla con la ayuda de createTempView() y la vida de esta temperatura depende de la vida de la sparkSession. registerTempTable() creará la tabla temporal si no está disponible o si está disponible, entonces reemplácela.

Luego, después de crear la tabla, seleccione la tabla por cláusula SQL que tomará todos los valores como una string.

df.registerTempTable('table')
newDF = spark.sql('select *, Course_Fees - Discount as Total from table')
newDF.show()

Producción:

Método 3: Usar UDF

En este método, definiremos que el usuario defina una función que tomará dos parámetros y devolverá el precio total. Esta función nos permite crear una nueva función según nuestros requisitos.

Ahora definimos el tipo de datos de la función UDF y creamos las funciones que devolverán los valores, que es la suma de todos los valores en la fila.

# import the functions as F from pyspark.sql
import pyspark.sql.functions as F
from pyspark.sql.types import IntegerType
  
# define the sum_col
def Total(Course_Fees, Discount):
    res = Course_Fees - Discount
    return res
  
# integer datatype is defined
new_f = F.udf(Total, IntegerType())
  
# calling and creating the new
# col as udf_method_sum
new_df = df.withColumn(
  "Total_price", new_f("Course_Fees", "Discount"))
  
# Showing the Dataframe
new_df.show()

Producción: