Спарка поиск SQL внутри массива для структуры

Question

структуры

Моих данных определяются примерно следующим образом:

schema = StructType([ 
# ... fields skipped 
StructField("extra_features", 
ArrayType(StructType([ 
    StructField("key", StringType(), False), 
    StructField("value", StringType(), True) 
])), nullable = False)], 
) 

Теперь я хотел бы искать записи в кадре данных, где структура {"key": "somekey", "value": "somevalue"} существует в массиве колонка. Как мне это сделать?

Answer 1

Spark имеет функцию array_contains, которая может использоваться для проверки содержимого столбца ArrayType, но, к сожалению, похоже, что он не может обрабатывать массивы сложных типов. Можно сделать это с помощью UDF (User Defined Function), однако:

from pyspark.sql.types import * 
from pyspark.sql import Row 
import pyspark.sql.functions as F 

schema = StructType([StructField("extra_features", ArrayType(StructType([ 
    StructField("key", StringType(), False), 
    StructField("value", StringType(), True)])), 
    False)]) 

df = spark.createDataFrame([ 
    Row([{'key': 'a', 'value': '1'}]), 
    Row([{'key': 'b', 'value': '2'}])], schema) 

# UDF to check whether {'key': 'a', 'value': '1'} is in an array 
# The actual data of a (nested) StructType value is a Row 
contains_keyval = F.udf(lambda extra_features: Row(key='a', value='1') in extra_features, BooleanType()) 

df.where(contains_keyval(df.extra_features)).collect() 

Это приводит к:

[Row(extra_features=[Row(key=u'a', value=u'1')])] 

Вы также можете использовать UDF, чтобы добавить еще один столбец, который указывает, является ли ключ-значение пары присутствует:

df.withColumn('contains_it', contains_keyval(df.extra_features)).collect() 

приводит к:

[Row(extra_features=[Row(key=u'a', value=u'1')], contains_it=True), 
Row(extra_features=[Row(key=u'b', value=u'2')], contains_it=False)]