from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType
from pyspark.sql import functions as F
import numpy as np

# UDF to calculate discounted amount
def calculate_discount(state, amount):
  if state == "CA":
    return amount * 0.90  # 10% off
  else:
    return amount * 0.85  # 15% off

discount_udf = udf(calculate_discount, DoubleType())

@time_decorator
def have_udf(data):
  # Use the UDF
  discounted_data = data.withColumn("discountedTotalAmt", discount_udf("state", "totalAmt"))

  # Show the results
  return discounted_data.select("customerID", "totalAmt", "state", "discountedTotalAmt").show()

display(have_udf(df))