Spark SQL - Generate array of arrays from the sql function

user_6043615 0
Viewed 10567

I want to create an array of arrays. This is my data table:

// A case class for our sample table
case class Testing(name: String, age: Int, salary: Int)

// Create an RDD with some data
val x = sc.parallelize(Array(
    Testing(null, 21, 905),
    Testing("Noelia", 26, 1130),
    Testing("Pilar", 52,  1890),
    Testing("Roberto", 31, 1450)
 ))

// Convert RDD to a DataFrame 
val df = sqlContext.createDataFrame(x) 

// For SQL usage we need to register the table
df.registerTempTable("df")

I want to create an array of integer column "age". For that I use "collect_list":

sqlContext.sql("SELECT collect_list(age) as age from df").show

But now I want to generate an array containing multiple arrays as created above:

 sqlContext.sql("SELECT collect_list(collect_list(age), collect_list(salary)) as arrayInt from df").show

But this does not work , or use the function org.apache.spark.sql.functions.array. Any ideas?

apache-spark-sql apache-spark scala
2 Answers

Let's create the DataFrame the way have created above.

// A case class for our sample table
import org.apache.spark.sql.functions._

case class Testing(name: String, age: Int, salary: Int)

// Create an RDD with some data
val x = sc.parallelize(Array(
    Testing(null, 21, 905),
    Testing("Noelia", 26, 1130),
    Testing("Pilar", 52,  1890),
    Testing("Roberto", 31, 1450)
 ))

// Convert RDD to a DataFrame 
val df = spark.createDataFrame(x)

Here we can use array_union function to achieve the desired result. array_unionfunction will return the union of all elements from the input arrays. This function is available since spark 2.4.0

// Scala Ref : https://spark.apache.org/docs/latest/api/scala/index.html#org.apache.spark.sql.functions$

// Pyspark Ref : https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.functions.array_union

df.select(collect_list("age").as("age"), collect_list("salary").as("salary"))
  .withColumn("new_col", array_union($"age", $"salary")).show(truncate=false)

// Output

+----------------+-----------------------+---------------------------------------+
|age             |salary                 |new_col                                |
+----------------+-----------------------+---------------------------------------+
|[21, 26, 52, 31]|[905, 1130, 1890, 1450]|[21, 26, 52, 31, 905, 1130, 1890, 1450]|
+----------------+-----------------------+---------------------------------------+

I hope this helps.

user_72203860
Related