在Pyspark数据框中访问数组的字段

Question

我正在开发针对一组ORC文件的Spark DataFrame的SQL查询语句。程序如下：

from pyspark.sql import SparkSession
spark_session = SparkSession.builder.appName("test").getOrCreate()
sdf = spark_session.read.orc("../data/")
sdf.createOrReplaceTempView("test")

现在我有一个名为“test”的表。如果我做以下操作：

spark_session.sql("select count(*) from test")

如果只需要查询的结果是正确的，那么就没有问题。但我需要在查询中获取更多的列，包括数组中的一些字段。

In [8]: sdf.take(1)[0]["person"]
Out[8]:
[Row(name='name', value='tom'),
 Row(name='age', value='20'),
 Row(name='gender', value='m')]

我尝试过类似以下的方法：

spark_session.sql("select person.age, count(*) from test group by person.age")

但是这样做不起作用。我的问题是：如何访问“person”数组中的字段？

谢谢！

编辑：

sdf.printSchema() 的结果为：

In [3]: sdf.printSchema()
root
 |-- person: integer (nullable = true)
 |-- customtags: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- name: string (nullable = true)
 |    |    |-- value: string (nullable = true)

错误信息：

AnalysisException: 'No such struct field age in name, value; line 16 pos 8'

- Bo Qiang

---

df = spark.sql("select * from test") func = udf(lambda x: x['age'], IntegerType()) df.withColumn('age', func(df.person)).groupBy('age').count().show() - pratiklodha

Answer 1

我不知道如何仅使用PySpark-SQL完成此操作，但是以下是使用PySpark数据框架的一种方法。基本上，我们可以使用create_map()函数将结构列转换为MapType()。然后，我们可以使用字符串索引直接访问字段。

考虑以下示例：

定义模式

schema = StructType([
        StructField('person', IntegerType()),
        StructField(
            'customtags',
            ArrayType(
                StructType(
                    [
                        StructField('name', StringType()),
                        StructField('value', StringType())
                    ]
                )
            )
        )
    ]
)

创建示例数据框

data = [
    (
        1, 
        [
            {'name': 'name', 'value': 'tom'},
            {'name': 'age', 'value': '20'},
            {'name': 'gender', 'value': 'm'}
        ]
    ),
    (
        2,
        [
            {'name': 'name', 'value': 'jerry'},
            {'name': 'age', 'value': '20'},
            {'name': 'gender', 'value': 'm'}
        ]
    ),
    (
        3,
        [
            {'name': 'name', 'value': 'ann'},
            {'name': 'age', 'value': '20'},
            {'name': 'gender', 'value': 'f'}
        ]
    )
]
df = sqlCtx.createDataFrame(data, schema)
df.show(truncate=False)
#+------+------------------------------------+
#|person|customtags                          |
#+------+------------------------------------+
#|1     |[[name,tom], [age,20], [gender,m]]  |
#|2     |[[name,jerry], [age,20], [gender,m]]|
#|3     |[[name,ann], [age,20], [gender,f]]  |
#+------+------------------------------------+

将结构列转换为映射

from operator import add
import pyspark.sql.functions as f

df = df.withColumn(
        'customtags',
        f.create_map(
            *reduce(
                add, 
                [
                    [f.col('customtags')['name'][i],
                     f.col('customtags')['value'][i]] for i in range(3)
                ]
            )
        )
    )\
    .select('person', 'customtags')

df.show(truncate=False)
#+------+------------------------------------------+
#|person|customtags                                |
#+------+------------------------------------------+
#|1     |Map(name -> tom, age -> 20, gender -> m)  |
#|2     |Map(name -> jerry, age -> 20, gender -> m)|
#|3     |Map(name -> ann, age -> 20, gender -> f)  |
#+------+------------------------------------------+

这里的关键是您必须预先知道ArrayType()的长度(在本例中为3)，因为我不知道动态循环的方法。这也假设数组在所有行中具有相同的长度。

我需要在这里使用reduce(add, ...)，因为create_map()希望以(key, value)形式的元素对为对。

按照映射列中的字段分组。

df.groupBy((f.col('customtags')['name']).alias('name')).count().show()
#+-----+-----+
#| name|count|
#+-----+-----+
#|  ann|    1|
#|jerry|    1|
#|  tom|    1|
#+-----+-----+

df.groupBy((f.col('customtags')['gender']).alias('gender')).count().show()
#+------+-----+
#|gender|count|
#+------+-----+
#|     m|    2|
#|     f|    1|
#+------+-----+

- pault

---

现在有什么建议可以动态获取数组的长度吗？ - Saharsh Kislaya