How to query JSON data according to JSON array's size with Spark SQL?

Question

I have a JSON like this:

{
     "uin":10000,
     "role":[
             {"role_id":1, "role_level": 10},
             {"role_id":2, "role_level": 1}
         ]
}
{     "uin":10001,
       "role":[
             {"role_id":1, "role_level": 1}, 
             {"role_id":2, "role_level": 1},
             {"role_id":3, "role_level": 1},
             {"role_id":4, "role_level": 20}
       ]

  } 

I want to query a uin which has more than two roles. How can I do it using Spark SQL?

Answer 1

You can use DataFrame and UserDefinedFunction to achieve what you want, as shown below. I've tried in spark-shell.

val jsonRdd = sc.parallelize(Seq("""{"uin":10000,"role":[{"role_id":1, "role_level": 10},{"role_id":2, "role_level": 1}]}"""))
val df = sqlContext.jsonRDD(jsonRdd)
val predict = udf((array: Seq[Any]) => if (array.length > 2) true else false)
val df1 = df.where( predict(df("role")) )
df1.show

Answer 2

Maybe size is what you need:

size(expr) - Returns the size of an array or a map.

In your case, "role" size must be bigger than 2.

If you have this JSON:

json = \
[
  {
    "uin":10000,
    "role":[
      {"role_id":1, "role_level": 10},
      {"role_id":2, "role_level": 1}
    ]
  },
  {
    "uin":10001,
    "role":[
      {"role_id":1, "role_level": 1}, 
      {"role_id":2, "role_level": 1},
      {"role_id":3, "role_level": 1},
      {"role_id":4, "role_level": 20}
    ]
  }
]

you can use this:

from pyspark.sql import functions as F
rdd = spark.sparkContext.parallelize([json])
df = spark.read.json(rdd)

df = df.filter(F.size('role') > 2)

df.show()
#+--------------------+-----+
#|                role|  uin|
#+--------------------+-----+
#|[{1, 1}, {2, 1}, ...|10001|
#+--------------------+-----+

Answer 3

Her is a simplified python version

r1 = ssc.jsonFile("role.json").select("uin","role.role_id")
r1.show()

slen = udf(lambda s: len(s), IntegerType())
r2 = r1.select(r1.uin,r1.role_id,slen(r1.role_id).alias("slen"))
res = r2.filter(r2.slen>1)
res.show()