Add new column to dataframe with modified schema or drop nested columns of array type in Scala

user_8761446 0
Viewed 114

json:-

{"ID": "500", "Data": [{"field2": 308, "field3": 346, "field1": 40.36582609126494, "field7": 3, "field4": 1583057346.0, "field5": -80.03243596528726, "field6": 16.0517578125, "field8": 5}, {"field2": 307, "field3": 348, "field1": 40.36591421686625, "field7": 3, "field4": 1583057347.0, "field5": -80.03259684675493, "field6": 16.234375, "field8": 5}]}

schema:-

   val MySchema: StructType = 
      StructType( Array(
        StructField("ID",StringType,true), 
        StructField("Data", ArrayType(
          StructType( Array(
              StructField("field1",DoubleType,true),
              StructField("field2",LongType,true),
              StructField("field3",LongType,true),
              StructField("field4",DoubleType,true),
              StructField("field5",DoubleType,true),
              StructField("field6",DoubleType,true),
              StructField("field7",LongType,true),          
              StructField("field8",LongType,true)
              )),true),true)))

Load json into dataframe:-

 val MyDF = spark.readStream
                    .schema(MySchema)
                    .json(input)

where 'input' is a file that contains above json

How can I add a new column "Data_New" to the above dataframe 'MyDF' with schema as

   val Data_New_Schema: StructType = 
      StructType( Array(
        StructField("Data", ArrayType(
          StructType( Array(
              StructField("field1",DoubleType,true),
              StructField("field4",DoubleType,true),
              StructField("field5",DoubleType,true),
              StructField("field6",DoubleType,true)
              )),true),true)))

Please note a huge volume of such json files will be loaded in the source and so performing an explode followed by a collect_list will crash the driver

apache-spark dataframe scala schema json
1 Answers

You can try one of the following two methods:

for Spark 2.4+, use transform:

import org.apache.spark.sql.functions._

val df_new = df.withColumn("Data_New", expr("struct(transform(Data, x -> (x.field1 as f1, x.field4 as f4, x.field5 as f5, x.field6 as f6)))").cast(Data_New_Schema))

scala> df_new.printSchema
root
 |-- ID: string (nullable = true)
 |-- Data: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- field1: double (nullable = true)
 |    |    |-- field2: long (nullable = true)
 |    |    |-- field3: long (nullable = true)
 |    |    |-- field4: double (nullable = true)
 |    |    |-- field5: double (nullable = true)
 |    |    |-- field6: double (nullable = true)
 |    |    |-- field7: long (nullable = true)
 |    |    |-- field8: long (nullable = true)
 |-- Data_New: struct (nullable = false)
 |    |-- Data: array (nullable = true)
 |    |    |-- element: struct (containsNull = false)
 |    |    |    |-- field1: double (nullable = true)
 |    |    |    |-- field4: double (nullable = true)
 |    |    |    |-- field5: double (nullable = true)
 |    |    |    |-- field6: double (nullable = true)

Notice that nullable = false on the top level schema of Data_New, if you want to make it true, add nullif function to the SQL expression: nullif(struct(transform(Data, x -> (...))), null), or a more efficient way if(true, struct(transform(Data, x -> (...))), null).

Prior to Spark 2.4, use from_json + to_json:

val df_new = df.withColumn("Data_New", from_json(to_json(struct('Data)), Data_New_Schema))

Edit: Per comment, if you want Data_New to be an array of structs, just remove struct function, for example:

val Data_New_Schema: ArrayType = ArrayType(
    StructType( Array(
    StructField("field1",DoubleType,true),
    StructField("field4",DoubleType,true),
    StructField("field5",DoubleType,true),
    StructField("field6",DoubleType,true)
)),true)

// if you need `containsNull=true`, then cast the above Type definition
val df_new = df.withColumn("Data_New", expr("transform(Data, x -> (x.field1 as field1, x.field4 as field4, x.field5 as field5, x.field6 as field6))"))

Or

val df_new = df.withColumn("Data_New", from_json(to_json('Data), Data_New_Schema))
user_95107290
Related