How to update value inside a nested array in spark

Question

I am currently trying to update the value of a column in a pyspark dataframe.

Here is the schema :

root
|-- date: string
|-- ticket: struct
|    |-- money: array
|    |    |-- element: struct
|    |    |    |-- currency: string
|    |    |    |-- total: double

I want to change the value of 'total' if we have a value of -9999 to None

Here is what I have in :

+-----------------------++----------------------------------------+
|date                   ||ticket                                  |
+-----------------------++----------------------------------------+
|2024-02-02T04:31:06    ||[{CAD, -9999}, {CAD, -9999}]            |
|2024-02-02T04:31:06    ||[{CAD, -9999}, {CAD, -9999}]            |
+-----------------------++----------------------------------------+

And here is the end result I try to have :

+-----------------------++----------------------------------------+
|date                   ||ticket                                  |
+-----------------------++----------------------------------------+
|2024-02-02T04:31:06    ||[{CAD, None}, {CAD, None}]              |
|2024-02-02T04:31:06    ||[{CAD, None}, {CAD, None}]              |
+-----------------------++----------------------------------------+

Edit : I can't have a new column, I need to keep this schema as is

Answer 1

Here's a solution using withField

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, ArrayType
from pyspark.sql import functions as F

# Create a SparkSession
spark = SparkSession.builder \
    .appName("CreateDataFrame") \
    .getOrCreate()

# Define the schema
schema = StructType([
    StructField("date", StringType(), True),
    StructField("ticket", StructType([
        StructField("money", ArrayType(StructType([
            StructField("currency", StringType(), True),
            StructField("total", DoubleType(), True)
        ])), True)
    ]), True)
])


# Sample data
data = [
    ("2024-04-18", [[{"currency": "USD", "total": -9999.0}, {"currency": "EUR", "total": 80.0}]]),
    ("2024-04-19", [[{"currency": "GBP", "total": -9999.0}, {"currency": "JPY", "total": 1000.0}]])
]

# Create DataFrame
df = spark.createDataFrame(data, schema)

# Show DataFrame
df.show(truncate=False)
# +----------+---------------------------------+
# |date      |ticket                           |
# +----------+---------------------------------+
# |2024-04-18|{[{USD, -9999.0}, {EUR, 80.0}]}  |
# |2024-04-19|{[{GBP, -9999.0}, {JPY, 1000.0}]}|
# +----------+---------------------------------+

# Show schema
df.printSchema()
# root
#  |-- date: string (nullable = true)
#  |-- ticket: struct (nullable = true)
#  |    |-- money: array (nullable = true)
#  |    |    |-- element: struct (containsNull = true)
#  |    |    |    |-- currency: string (nullable = true)
#  |    |    |    |-- total: double (nullable = true)

Here's how the sample dataframe looks:

Change all values equal to -9999 to None

df.withColumn('ticket',
 F.col('ticket')
    .withField('money', 
       F.transform("ticket.money",
            lambda c: c.withField("total",F.when(c.getField("total")==-9999, None)
                                           .otherwise(c.getField("total")))
                  )
               )
              ).show(truncate=False)
    
# +----------+------------------------------+
# |date      |ticket                        |
# +----------+------------------------------+
# |2024-04-18|{[{USD, null}, {EUR, 80.0}]}  |
# |2024-04-19|{[{GBP, null}, {JPY, 1000.0}]}|
# +----------+------------------------------+

There's a also a library for such nested fields transformations: pyspark-nested-functions.

Just figured out how to use pyspark-nested-functions, this library indeed simplifies the code:

from nestedfunctions.functions.terminal_operations import apply_terminal_operation
from pyspark.sql.functions import when
processed = apply_terminal_operation(
      df,
      field="ticket.money.total",
      f=lambda x, type: when(x==-9999, None).otherwise(x),
  )
processed.show(truncate=False)
# +----------+------------------------------+
# |date      |ticket                        |
# +----------+------------------------------+
# |2024-04-18|{[{USD, null}, {EUR, 80.0}]}  |
# |2024-04-19|{[{GBP, null}, {JPY, 1000.0}]}|
# +----------+------------------------------+