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

Question

在Spark SQL的帮助下，我试图过滤掉属于特定组类别的所有业务项。

数据从JSON文件中加载:

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

文件的模式如下:

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

我正在尝试提取与餐厅业务相关的所有业务:

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

，但它不起作用，因为我理解列的预期类型应该是字符串，但在我的情况下，这是数组。About it告诉我一个异常:

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

你能建议其他方法来得到我想要的吗?

Answer 1

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|
## +-------------------------------------+