如何创建一列数组，其值来自一列并且它们的长度来自pyspark数据帧中的另一列？

Question

在 pyspark 数据框中，我需要创建一个新的 ArrayType(StringType()) 列，其值来自 StringType() 列，它们的长度来自另一个 ArrayType(StringType()) 列的长度。有点像具有动态长度的 array_repeat。

输入：

+-------------+-------------+
|col1         |col2         |
+-------------+-------------+
|[1,2]        |‘a’          |
|[1,2,3]      |‘b’          |
+-------------+-------------+

输出：

+-------------+-------------+----------------+
|col1         |col2         |col3            |
+-------------+-------------+----------------+
|[1,2]        |‘a’          |['a’,‘a’]       |
|[1,2,3]      |‘b’          |['b’,’b’,’b’]   |
+-------------+----------- -+----------------+

谢谢

Answer 1

另一种选择-

加载提供的测试数据

val df = spark.sql(
      """
        |select col1, col2
        |from values
        | (array(1, 2), 'a'),
        | (array(1, 2, 3), 'b')
        | T(col1, col2)
      """.stripMargin)
    df.show(false)
    df.printSchema()

    /**
      * +---------+----+
      * |col1     |col2|
      * +---------+----+
      * |[1, 2]   |a   |
      * |[1, 2, 3]|b   |
      * +---------+----+
      *
      * root
      * |-- col1: array (nullable = false)
      * |    |-- element: integer (containsNull = false)
      * |-- col2: string (nullable = false)
      */

Alternative-1

    // alternative-1
    df.withColumn("col3", expr("array_repeat(col2, size(col1))"))
      .show(false)

    /**
      * +---------+----+---------+
      * |col1     |col2|col3     |
      * +---------+----+---------+
      * |[1, 2]   |a   |[a, a]   |
      * |[1, 2, 3]|b   |[b, b, b]|
      * +---------+----+---------+
      */

备选方案 2

    // alternative-2
    df.withColumn("col3", expr(s"TRANSFORM(col1, x -> col2)"))
      .show(false)

    /**
      * +---------+----+---------+
      * |col1     |col2|col3     |
      * +---------+----+---------+
      * |[1, 2]   |a   |[a, a]   |
      * |[1, 2, 3]|b   |[b, b, b]|
      * +---------+----+---------+
      */

Answer 2

使用array_repeat + size：

import pyspark.sql.functions as f

df = spark.createDataFrame([[[1,2],'a'], [[1,2,3], 'b']], ['col1', 'col2'])

df.withColumn('col3', f.array_repeat('col2', f.size('col1'))).show()

+---------+----+---------+
|     col1|col2|     col3|
+---------+----+---------+
|   [1, 2]|   a|   [a, a]|
|[1, 2, 3]|   b|[b, b, b]|
+---------+----+---------+

如果在任何情况下这不起作用，您可以编写一个 udf 来执行此操作：

from pyspark.sql.types import StringType, ArrayType
import pyspark.sql.functions as f

@f.udf(ArrayType(StringType()))
def repeat_sizeof(col1, col2):
    return [col1] * len(col2)

df.withColumn('col3', repeat_sizeof('col2', 'col1')).show()

+---------+----+---------+
|     col1|col2|     col3|
+---------+----+---------+
|   [1, 2]|   a|   [a, a]|
|[1, 2, 3]|   b|[b, b, b]|
+---------+----+---------+