使用一系列数字范围创建新列

Question

所以我需要创建一个从 1 到 100 枚举的数字数组，作为每行的值作为一个额外的列。

使用array()带有一堆文字值的函数是可行的，但是肯定有一种方法可以使用/转换 ScalaRange(a to b)而不是单独列出每个数字？

spark.sql("SELECT key FROM schema.table")
  .otherCommands
  .withColumn("range", array(lit(1), lit(2), ..., lit(100)))

类似于：

withColumn("range", array(1 to 100))

Answer 1

从 Spark 2.4 开始，您可以使用 [sequence][1] 函数如果您有以下数据框：

df.show()
+--------+
|column_1|
+--------+
|       1|
|       2|
|       3|
|       0|
+--------+

如果您使用从 0 到 column_1 的序列函数，您会得到以下结果：

df.withColumn("range", sequence(lit(0), col("column_1"))).show()
+--------+------------+
|column_1|       range|
+--------+------------+
|       1|      [0, 1]|
|       2|   [0, 1, 2]|
|       3|[0, 1, 2, 3]|
|       0|         [0]|
+--------+------------+

对于这种情况，请使用以下命令设置这两个值lit：

df.withColumn("range", sequence(lit(0), lit(100)))

Answer 2

您可以使用函数内的内置函数作为map函数使用litarray

df.withColumn("range", array((1 to 100).map(lit(_)): _*))