Add a column to a struct nested in an array

Question

I have a PySpark DataFrame with an array of structs, containing two columns (colorcode and name). I want to add a new column to the struct, newcol.

This question answered "how to add a column to a nested struct", but I'm failing to transfer it to my case, where the struct is further nested inside an array. I can't seem to reference/recreate the array-struct schema.

My schema:

 |-- Id: string (nullable = true)
 |-- values: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- Dep: long (nullable = true)
 |    |    |-- ABC: string (nullable = true)

What is should become:

 |-- Id: string (nullable = true)
 |-- values: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- Dep: long (nullable = true)
 |    |    |-- ABC: string (nullable = true)
 |    |    |-- newcol: string (nullable = true)

How do I transfer the solution to my nested struct?

Reproducible code to get a df of the above schema:

data = [
    (10, [{"Dep": 10, "ABC": 1}, {"Dep": 10, "ABC": 1}]),
    (20, [{"Dep": 20, "ABC": 1}, {"Dep": 20, "ABC": 1}]),
    (30, [{"Dep": 30, "ABC": 1}, {"Dep": 30, "ABC": 1}]),
    (40, [{"Dep": 40, "ABC": 1}, {"Dep": 40, "ABC": 1}])
  ]
myschema = StructType(
[
    StructField("id", IntegerType(), True),
    StructField("values",
                ArrayType(
                    StructType([
                        StructField("Dep", StringType(), True),
                        StructField("ABC", StringType(), True)
                    ])
    ))
]
)
df = spark.createDataFrame(data=data, schema=myschema)
df.printSchema()
df.show(10, False)

Please add some reproducible code

Mahesh Gupta
– Mahesh Gupta

2022-03-31 06:41:17 +00:00
Commented Mar 31, 2022 at 6:41 — Mahesh Gupta
– Mahesh Gupta, Commented Mar 31, 2022 at 6:41
what is your spark version?

过过招
– 过过招

2022-03-31 06:49:27 +00:00
Commented Mar 31, 2022 at 6:49 — 过过招
– 过过招, Commented Mar 31, 2022 at 6:49
pyspark 3.2.1, I am working on a reproducible code

Cribber
– Cribber

2022-03-31 06:49:52 +00:00
Commented Mar 31, 2022 at 6:49 — Cribber
– Cribber, Commented Mar 31, 2022 at 6:49

过过招 · Accepted Answer · 2022-03-31 08:05:58Z

9

For spark version >= 3.1, you can use the transform function and withField method to achieve this.

transform performs the transformation calculation according to the provided function for each element (struct(Dep, ABC) here) in the array (values column here). withField adds/replaces a field in StructType by name.

df = df.withColumn('values', F.transform('values', lambda x: x.withField('newcol', F.lit(1))))

edited Mar 31, 2022 at 8:05

answered Mar 31, 2022 at 6:52

过过招

4,3372 gold badges7 silver badges14 bronze badges

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

3 Comments

过过招 Over a year ago

try: df = df.withColumn('values', F.transform('values', lambda x: x.withField('Dep', x['Dep'].cast('int'))))

wwnde Over a year ago

This looks neater df.withColumn("values",F.expr("transform(values, x -> struct(cast((x.Dep) as integer) as Dep, x.ABC))"))

过过招 Over a year ago

It depends on personal habits and familiarity. At the beginning, I was used to answering questions using spark sql expressions, but I found that many people are more used to the dataframe API.

wwnde · Accepted Answer · 2022-03-31 11:42:06Z

1

Another way, of doing it would be using sql expressions.

df = df.withColumn("values",F.expr("transform(values, x -> struct(COALESCE('1') as newcol,x.Dep,x.ABC))"))

answered Mar 31, 2022 at 11:42

wwnde

26.7k6 gold badges22 silver badges38 bronze badges

Collectives™ on Stack Overflow

Add a column to a struct nested in an array

2 Answers 2

3 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related