pyspark：通过arraytype列过滤和提取struct

c9x0cxw0 于 2021-07-12 发布在 Spark

关注(0)|答案(2)|浏览(349)

我使用的是pyspark2.2，有以下模式

root
 |-- col1: string (nullable = true)
 |-- col2: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- id: string (nullable = true)
 |    |    |-- metadata: map (nullable = true)
 |    |    |    |-- key: string
 |    |    |    |-- value: string (valueContainsNull = true)

和数据

+----+----------------------------------------------+
|col1|col2                                          |
+----+----------------------------------------------+
|A   |[[id1, [k -> v1]], [id2, [k2 -> v5, k -> v2]]]|
|B   |[[id3, [k -> v3]], [id4, [k3 -> v6, k -> v4]]]|
+----+----------------------------------------------+
``` `col2` 是一个复杂的结构。它是一个结构数组，每个结构都有两个元素 `id` 字符串和 `metadata` Map(这是一个简化的数据集，真正的数据集在struct中有10+个元素，在struct中有10+个键值对 `metadata` 字段）。
我想形成一个查询，返回一个与我的过滤逻辑匹配的Dataframe（比如 `col1 == 'A'` 以及 `col2.id == 'id2'` 以及 `col2.metadata.k == 'v2'` ).
结果是这样的，过滤逻辑最多可以匹配数组中的一个结构，因此在第二列中，它只是一个结构，而不是一个结构的数组

+----+--------------------------+
|col1|col2_filtered |
+----+--------------------------+
|A |[id2, [k2 -> v5, k -> v2]]|
+----+--------------------------+

我知道如何通过 `explode` ，但问题是 `col2` 通常有超过100+结构，最多会有一个匹配我的过滤逻辑，所以我不认为 `explode` 是一个可扩展的解决方案。
有人能告诉我怎么做吗，提前谢谢！
下面是设置的代码块。

from pyspark.sql.types import StructType, StructField, StringType, ArrayType, MapType

schema = StructType([
StructField('col1', StringType(), True),
StructField('col2', ArrayType(
StructType([
StructField('id', StringType(), True),
StructField('metadata', MapType(StringType(), StringType()), True)
])
))
])

data = [
('A', [('id1', {'k': 'v1'}), ('id2', {'k': 'v2', 'k2': 'v5'})]),
('B', [('id3', {'k': 'v3'}), ('id4', {'k': 'v4', 'k3': 'v6'})])
]

df = spark.createDataFrame(data=data, schema=schema)

JSON python apache-spark pyspark apache-spark-sql

来源：https://stackoverflow.com/questions/66577318/pyspark-filtering-and-extract-struct-through-arraytype-column

2条答案

按热度按时间

68de4m5k1#

编辑：您可以尝试自定义项：

import pyspark.sql.functions as F

df2 = df.filter(
    F.udf(lambda x: any([y.id == 'id2' and 'k' in y.metadata.keys() for y in x]), 'boolean')('col2')
).withColumn(
    'col2',
    F.udf(lambda x: [y for y in x if y.id == 'id2' and 'k' in y.metadata.keys()][0], 'struct<id:string,metadata:map<string,string>>')('col2')
)

df2.show(truncate=False)
+----+--------------------------+
|col1|col2                      |
+----+--------------------------+
|A   |[id2, [k2 -> v5, k -> v2]]|
+----+--------------------------+

您可以将这些列强制转换为json，并检查col2是否包含所需的json：

import pyspark.sql.functions as F

df2 = df.filter(
    (F.col('col1') == 'A') &
    F.to_json('col2').contains(
        F.to_json(
            F.struct(
                F.lit('id2').alias('id'),
                F.create_map(F.lit('k'), F.lit('v2')).alias('metadata')
            )
        )
    )
)

df2.show(truncate=False)
+----+------------------------------------+
|col1|col2                                |
+----+------------------------------------+
|A   |[[id1, [k -> v1]], [id2, [k -> v2]]]|
+----+------------------------------------+

如果您只想在col2中保留匹配的结构，可以使用 withColumn :

df3 = df2.withColumn(
    'col2', 
    F.struct(
        F.lit('id2').alias('id'),
        F.create_map(F.lit('k'), F.lit('v2')).alias('metadata')
    )
)

df3.show()
+----+----------------+
|col1|            col2|
+----+----------------+
|   A|[id2, [k -> v2]]|
+----+----------------+

赞(0）回复(0）举报 2021-07-12

kdfy810k2#

除了@mck提供的解决方案，我在搜索之后尝试了另外三种方法，都得到了想要的结果。
过滤使用 udf 并返回匹配的结构

df.filter(df.col1 == 'A') \
  .select(df.col1, udf(lambda a: [s for s in a if s.id == 'id2' and s.metadata['k'] == 'v2'], df.schema['col2'].dataType)('col2')[0].alias('col2_filtered')) \
  .na.drop('any')

过滤使用 udf 并得到匹配结构的索引

df.filter(df.col1 == 'A') \
  .select(df.col1, df.col2.getItem(udf(lambda a: [i for i, s in enumerate(a) if s.id == 'id2' and s.metadata['k'] == 'v2'], ArrayType(IntegerType(), True))(df.col2)[0]).alias('col2_filtered')) \
  .na.drop('any')

过滤使用 expr ，这是spark 2.4中的一项功能，因此可以作为未来升级的候选

df.filter(df.col1 == 'A') \
  .select(df.col1, expr("filter(col2, s -> s.id == 'id2' AND s.metadata['k'] == 'v2')").getItem(0).alias('col2_filtered')) \
  .na.drop('any')

赞(0）回复(0）举报 2021-07-12

我来回答

pyspark：通过arraytype列过滤和提取struct

2条答案

相关问题

热门标签

最新问答