如何從 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   |
+-----------+-------+------------------------------------------------+-----------+