Sort the resulting DataFrame in pyspark using UDF

Question

I am working in SPARK with diamonds dataframe. Data as below:

+-----+-------+-----+-------+-----+-----+-----+----+----+----+
|carat|    cut|color|clarity|depth|table|price|   x|   y|   z|
+-----+-------+-----+-------+-----+-----+-----+----+----+----+
| 0.23|  Ideal|    E|    SI2| 61.5| 55.0|  326|3.95|3.98|2.43|
| 0.21|Premium|    E|    SI1| 59.8| 61.0|  326|3.89|3.84|2.31|
| 0.23|   Good|    E|    VS1| 56.9| 65.0|  327|4.05|4.07|2.31|
| 0.29|Premium|    I|    VS2| 62.4| 58.0|  334| 4.2|4.23|2.63|
| 0.31|   Good|    J|    SI2| 63.3| 58.0|  335|4.34|4.35|2.75|
+-----+-------+-----+-------+-----+-----+-----+----+----+----+

and schema:

root
|-- carat: double (nullable = true)
|-- cut: string (nullable = true)
|-- color: string (nullable = true)
|-- clarity: string (nullable = true)
|-- depth: double (nullable = true)
|-- table: double (nullable = true)
|-- price: integer (nullable = true)
|-- x: double (nullable = true)
|-- y: double (nullable = true)

I have created a custom function and registered as UDF:

def rank_cut(cut):
    cut_class_dict = {"Fair": 1, "Good": 2, "Very Good": 3, "Premium": 4, "Ideal": 5}
    for cut, v in cut_class_dict():
        x['cut'] = v
 
    return v

spark.udf.register('rank_cut', rank_cut)

I want to use this custom function to sort my dataframe as below:

( 
diamonds
.groupBy('cut')
.agg(
    expr('COUNT(*) AS n_diamonds'),
    expr('ROUND(AVG(price)) AS avg_price'),
    expr('ROUND(AVG(carat),2) AS avg_carat'),
    expr('ROUND(AVG(depth),2) AS avg_depth'),
    expr('ROUND(AVG(table),2) AS avg_table'),
  
  
)
.rank_cut('cut')
.show()
)

but it is not working. Anything I am missing?

Answer 1

Problem solved.

I changed my udf to:

cut_class_dict = {"Fair": 1, "Good": 2, "Very Good": 3, "Premium": 4, "Ideal": 5}
rank_cut =  udf(lambda cut: cut_class_dict.get(cut))
spark.udf.register('rank_cut', rank_cut)

and mapped it to cut column