Python 如何在Spark中动态切片数组列？

Spark 2.4引入了新的SQL函数

slice

，可用于从数组列中提取一定范围的元素。 我想根据一个整数列动态定义每行的范围，该整数列包含我想从该列中选取的元素数

但是，简单地将列传递给slice函数失败，函数似乎希望开始值和结束值为整数。有没有一种不用编写UDF就可以做到这一点的方法

要通过示例将问题可视化，请执行以下操作： 我有一个数据帧，它有一个数组列

arr

，每行中都有一个数组，看起来像

['a'，'b'，'c']

。还有一个

end\u idx

列，其中包含元素

3

、

1

和

2

：

+---------+-------+
|arr      |end_idx|
+---------+-------+
|[a, b, c]|3      |
|[a, b, c]|1      |
|[a, b, c]|2      |
+---------+-------+

我尝试创建一个新列

arr\u trimmed

，如下所示：

导入pyspark.sql.F函数
l=[（['a'，'b'，'c']，3），（['a'，'b'，'c']，1），（['a'，'b'，'c']，2）]
df=spark.createDataFrame（l，[“arr”，“end_idx”]）
df=df.withColumn（“arr_-trimmed”）、F.slice（F.col（“arr”）、1、F.col（“end_-idx”））

我希望这段代码能够创建一个包含元素的新列

['a'，b'，c']

，

['a']

，

['a'，b']

---

相反，我得到了一个错误

TypeError:Column不可编辑

，您可以通过如下方式传递SQL表达式来执行此操作：

df.withColumn("arr_trimmed", F.expr("slice(arr, 1, end_idx)"))

以下是整个工作示例：

import pyspark.sql.functions as F

l = [(['a', 'b', 'c'], 3), (['a', 'b', 'c'], 1), (['a', 'b', 'c'], 2)]

df = spark.createDataFrame(l, ["arr", "end_idx"])

df.withColumn("arr_trimmed", F.expr("slice(arr, 1, end_idx)")).show(truncate=False)

+---------+-------+-----------+
|arr      |end_idx|arr_trimmed|
+---------+-------+-----------+
|[a, b, c]|3      |[a, b, c]  |
|[a, b, c]|1      |[a]        |
|[a, b, c]|2      |[a, b]     |
+---------+-------+-----------+

可能重复的