使用索引列查询数组列

Question

所以我有两栏

| col_arr | col_ind | 
|[1, 2, 3]| [0, 2]  |
|[5, 1]   | [1]     |

我希望我的结果是通过 col_ind 提取 col_arr 中的值，从而得到下面的 col_val：

| col_arr | col_ind | col_val |
|[1, 2, 3]| [0, 2]  | [1, 3]  |
|[5, 1]   | [1]     | [1]     |

最优雅、最有效的方法是什么？

---

我的第一个想法是使用 UDF，但感觉有点大材小用了

@udf
def sub_select(arr, inds):
    if (arr is not None) and (inds is not None):
        return [arr[ind] for ind in inds]

我也在考虑使用 expr 动态地使用 array_position 函数。我不清楚如何以灵活适应 col_ind 长度的方式做到这一点。

F.expr("array_position(col_arr, array_position(col_ind, 0))")

---

简化和复杂化：

1. 假设要查询的索引长度受一个小数字限制（例如 col_ind 列表的最大长度为 5）。
2. 如果我有多个 col_arr 列（col_arr1、col_arr2、col_arr3），而只有一个 col_ind，该怎么办？

Answer 1

在

col_ind

上应用转换函数，并为每个索引返回

col_arr

中的相应值

df = df.withColumn('col_val', F.expr("transform(col_ind, i -> col_arr[i])"))

# +---------+-------+-------+
# |  col_arr|col_ind|col_val|
# +---------+-------+-------+
# |[1, 2, 3]| [0, 2]| [1, 3]|
# |   [5, 1]|    [1]|    [1]|
# +---------+-------+-------+

奖励：如果我有多个 col_arr 列（col_arr1、col_arr2、col_arr3），而只有一个 col_ind 会怎么样？ 可以对每一列分别使用相同的方法