pyspark: Dataframe- UDF with multiple arguments

Question

I have a dataframe where column has an array and each element is a dictionary.

class	product
{"deleteDate": null, "class":"AB", "validFrom": "2022-09-01", "validTo": "2009-08-31"}, {"deleteDate": null, "class":"CD", "validFrom": "2009-09-01", "validTo": "2024-08-31"}	{"deleteDate": "2021-09-01", "class":"AB", "validFrom": "2003-09-01", "validTo": "2009-03-01"}, {"deleteDate": null, "class":"CD", "validFrom": "2009-09-01", "validTo": "2024-08-31"}

I am trying to filter an element base on a few conditions.

def getelement(value,entity):
    list_url = []
    for i in range(len(value)):
        if value[i] is not None and (value[i].deleteDate is None):
            if (value[i].validFrom <= (Date of Today)) & (value[i].validFrom >= (date of today)):

            list_url.append(value[i].entity)
    if list_url:
        return str(list_url[-1])
    if not list_url:
        return None

    
udfgeturl=F.udf(lambda z: getelement(z) if not z is None else "" , StringType() )



master = df.withColumn( 'ClassName', udfgeturl('Class'))

The function takes two elements, value and entity. where value refers to column name and entity refers to a key in dictionary for which I want to save the result. The code works with one element getelement(value) for UDF but I do not know how the UDF can take two arguments, any suggestion, please?

ARCrow · Accepted Answer · 2022-06-23 09:57:33Z

1

To improve the performance (Spark functions vs UDF performance?), you could use only spark transformations:
I'm assuming (value[i].validFrom >= (date of today)) is supposed to actually be (value[i].validTo >= (date of today))

import pyspark.sql.functions as f

def getelement(value, entity):

    df = (
        df
        .withColumn('output', f.expr(f'filter({value}, element -> (element.deleteDate is null) AND (element.validFrom <= current_date()) AND (element.validTo >= current_date()))')[entity][-1])
    )

    return df

edited Jun 23, 2022 at 9:57

answered Jun 21, 2022 at 20:38

ARCrow

1,8673 gold badges14 silver badges34 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

rooya sh Over a year ago

Thank you it worked, except I needed to change "&" with AND.

ARCrow Over a year ago

You’re right; fixed.

Matt Andruff · Accepted Answer · 2022-06-21 17:21:22Z

0

You can use a struct to bundle the parameters into 1 object. Then access the elements of the struct with . operator.

Code example:

def getelement(object):
  value = object.value
  entity = object.entity
  return str( entity + " " + value )

udfgeturl=f.udf(getelement , StringType() )
df.select( 
 udfgeturl( 
  f.struct( 
   f.col("col1").alias("value"), 
   f.col("col2").alias("entity")) 
  ) 
 ).show()

edited Jun 21, 2022 at 17:21

answered Jun 21, 2022 at 17:14

Matt Andruff

5,1901 gold badge7 silver badges25 bronze badges

Collectives™ on Stack Overflow

pyspark: Dataframe- UDF with multiple arguments

2 Answers 2

2 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related