如何从另一列提供array_repeat函数的计数值？

Question

如何从另一列提供array_repeat函数计数值：

>>>import pyspark.sql.functions as F
>>>dftmp = spark.createDataFrame([('ab',)], ['data'])

>>>dftmp.select(F.array_repeat(dftmp.data, 3).alias('r'))
>>>dftmp.show()
+----+------------+
|data|      repeat|
+----+------------+
|  ab|[ab, ab, ab]|
+----+------------+

有没有一种方法可以使用基于另一列的重复计数值？例如

>>>dftmp.withColumn('len',  F.length(F.col('data')) )
>>>dftmp.withColumn('repeat', F.array_repeat(dftmp.data, F.col('len')))
TypeError: Column is not iterable

预期结果

+----+------------+---+
|data|      repeat|len|
+----+------------+---+
|  ab|[ab, ab]    |  2|
+----+------------+---+

Answer 1

您可以使用 .expr ：

from pyspark.sql import functions as F

dftmp.withColumn('repeat', F.expr("""array_repeat(data, len)"""))

或者您也可以在其中 calculate length ：

dftmp.withColumn('repeat', F.expr("""array_repeat(data, length(data))"""))