如何通过多值列过滤JSON数据

Question

With help of Spark SQL I'm trying to filter out all business items from which belongs to a specific group category.

The data is loaded from JSON file:

businessJSON = os.path.join(targetDir, 'yelp_academic_dataset_business.json')
businessDF = sqlContext.read.json(businessJSON)

The schema of the file is following:

businessDF.printSchema()

root
  |-- business_id: string (nullable = true)
  |-- categories: array (nullable = true)
  |    |-- element: string (containsNull = true)
  ..
  |-- type: string (nullable = true)

I'm trying to extract all business connected to restaurant business:

restaurants = businessDF[businessDF.categories.inSet("Restaurants")]

but it doesn't work because as I understand the expected type of column should be a string, but in my case this is array. About it tells me an exception:

Py4JJavaError: An error occurred while calling o1589.filter.
: org.apache.spark.sql.AnalysisException: invalid cast from string to array<string>;

Can you please suggest any other way to get what I want?

Answer 1

How about an UDF?

from pyspark.sql.functions import udf, col, lit
from pyspark.sql.types import BooleanType

contains = udf(lambda xs, val: val in xs, BooleanType())
df = sqlContext.createDataFrame([Row(categories=["foo", "bar"])])

df.select(contains(df.categories, lit("foo"))).show()
## +----------------------------------+
## |PythonUDF#<lambda>(categories,foo)|
## +----------------------------------+
## |                              true|
## +----------------------------------+

df.select(contains(df.categories, lit("foobar"))).show()
## +-------------------------------------+
## |PythonUDF#<lambda>(categories,foobar)|
## +-------------------------------------+
## |                                false|
## +-------------------------------------+