如何从 pyspark dataframe 中查询/提取数组元素

Question

我正在尝试构建一个将数据合并到数组中的数据框。 该数组是具有索引和值对的 2 值数组。 并非数据框：数组中的每一行都存在每个索引值。 这是架构的样子

root
|-- visitNumber: string (nullable = true)
|-- visitId: string (nullable = true)
|-- customDimensions: array (nullable = true)
|    |-- element: struct (containsNull = true)
|    |    |-- index: string (nullable = true)
|    |    |-- value: string (nullable = true)

还有许多其他列，但问题不涉及这些列。

以下是数据方面的 customDimensions 数组示例：

[[1, ],[2, ],[3,"apple"],[6,"1-111-32"],[42, ],[5, ]]

我想要完成的是合并包含特定索引值的列。 例如：

df = df.withColumn("index6", *stuff to get the value at index 6*)

这将是一个可重复的迭代，因为整个“customDimensions”中都有数据，其中包含我们可以“展平”并表示为单独列的所需数据。

Answer 1

这可以使用 udf 来实现。 以下应该工作。

# import dependencies, start spark-session
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, lit
spark = SparkSession.builder.appName('my-app').master('local[2]').getOrCreate()

# data preparation
df = spark.createDataFrame(
        [
            (1, 'id-1', [['1', ],['2', ],['3', 'apple'],[6, '1-111-32'],[42, ],[5, ]]),
            (2, 'id-2', [['1', ],['2', ],['3', 'apple'],[6, ],          [42, ],[5, ]]),
            (3, 'id-3', [['1', ],['2', ],['3', 'apple'],[6, '2-111-32'],[42, ],[5, ]]),
        ]
        , schema=(
            'visitNumber', 'visitId', 'customDimensions'))

df.show(3, False)
+-----------+-------+------------------------------------------------+
|visitNumber|visitId|customDimensions                                |
+-----------+-------+------------------------------------------------+
|1          |id-1   |[[1], [2], [3, apple], [6, 1-111-32], [42], [5]]|
|2          |id-2   |[[1], [2], [3, apple], [6],           [42], [5]]|
|3          |id-3   |[[1], [2], [3, apple], [6, 2-111-32], [42], [5]]|
+-----------+-------+------------------------------------------------+

现在，让我们准备 udf

def user_func(x, y):
    elem_set = {elem[1] if elem[0] == y and len(elem) == 2 else None for elem in x} - {None}
    return None if not elem_set else elem_set.pop()

user_func_udf = udf(lambda x, y: user_func(x, y))
spark.udf.register("user_func_udf", user_func_udf)

假设我们需要一列来获取索引 6 处的值，以下应该可以工作。

df_2 = df.withColumn('valAtIndex6', user_func_udf(df.customDimensions, lit('6')))
df_2.show(3, False)
+-----------+-------+------------------------------------------------+-----------+
|visitNumber|visitId|customDimensions                                |valAtIndex6|
+-----------+-------+------------------------------------------------+-----------+
|1          |id-1   |[[1], [2], [3, apple], [6, 1-111-32], [42], [5]]|1-111-32   |
|2          |id-2   |[[1], [2], [3, apple], [6], [42], [5]]          |null       |
|3          |id-3   |[[1], [2], [3, apple], [6, 2-111-32], [42], [5]]|2-111-32   |
+-----------+-------+------------------------------------------------+-----------+