Spark的函數array_contains
可用於檢查ArrayType
列的內容,但不幸的是,它似乎不能處理複雜類型的數組。這是可能的一個UDF做到這一點(用戶自定義函數)但是:
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)]