What is the most elegant way to apply custom function to PySpark dataframe with multiple columns?

I need to create new fields based on three dataframe fields. This works but it seems inefficient:

def my_func(very_long_field_name_a, very_long_field_name_b, very_long_field_name_c):
  if very_long_field_name_a >= very_long_field_name_b and very_long_field_name_c <= very_long_field_name_b:
    return 'Y'
  elif very_long_field_name_a <= very_long_field_name_b and very_long_field_name_c >= very_long_field_name_b:
    return 'Y'
  else: 
    return 'N'

import pyspark.sql.functions as F
my_udf = F.udf(my_func)

df.withColumn('new_field', my_udf(df.very_long_field_name_a, df.very_long_field_name_b, df.very_long_field_name_c)).display()

Is it possible to just pass the dataframe like so? I tried but got an error:

def my_func(df):
  if df.very_long_field_name_a >= df.very_long_field_name_b and df.very_long_field_name_c <= df.very_long_field_name_b:
    return 'Y'
  df.elif very_long_field_name_a <= df.very_long_field_name_b and df.very_long_field_name_c >= df.very_long_field_name_b:
    return 'Y'
  else: 
    return 'N'

import pyspark.sql.functions as F
my_udf = F.udf(my_func)
df.withColumn('new_field', my_udf(df)).display()

Invalid argument, not a string or column:

The reason I want to shorten it is because I have create six new fields. It seems inefficient to copy and paste all the field names passed as arguments, so I’d like to know if there’s something cleaner.

>Solution :

To create new fields based on multiple columns in a DataFrame without explicitly passing each column as an argument to the UDF, you can use the struct function in PySpark. The struct function combines multiple columns into a single column of StructType. Here’s an example:

import pyspark.sql.functions as F

def my_func(row):
    if row.very_long_field_name_a >= row.very_long_field_name_b and row.very_long_field_name_c <= row.very_long_field_name_b:
        return 'Y'
    elif row.very_long_field_name_a <= row.very_long_field_name_b and row.very_long_field_name_c >= row.very_long_field_name_b:
        return 'Y'
    else:
        return 'N'

my_udf = F.udf(my_func)

# Use struct to combine the necessary columns into a single column
df = df.withColumn('combined_fields', F.struct('very_long_field_name_a', 'very_long_field_name_b', 'very_long_field_name_c'))

# Apply the UDF to the combined column
df = df.withColumn('new_field', my_udf(F.col('combined_fields')))

# Drop the temporary combined column
df = df.drop('combined_fields')

df.display()

In this approach, we use the struct function to combine the necessary columns (very_long_field_name_a, very_long_field_name_b, very_long_field_name_c) into a single column called combined_fields. Then, we apply the UDF to the combined_fields column using my_udf(F.col('combined_fields')). Finally, we drop the temporary combined column using df.drop('combined_fields').

By using struct, you can avoid explicitly passing each column as an argument to the UDF and make the code cleaner and more efficient.