Как создать столбец массивов, значения которого поступают из одного столбца, а их длина - из другого столбца в фреймах данных pyspark? - PullRequest
Новая
       11

Как создать столбец массивов, значения которого поступают из одного столбца, а их длина - из другого столбца в фреймах данных pyspark?

1 голос
/ 03 августа 2020

В фрейме данных pyspark мне нужно создать новый столбец ArrayType (StringType ()), значения которого поступают из столбца StringType (), а их длина - из длины другого столбца ArrayType (StringType ()). Вроде как array_repeat с динамической c длиной.

ввод : 

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

вывод : 

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

Спасибо

Ответы [ 2 ]

2 голосов
/ 03 августа 2020

Другая альтернатива-

Загрузить предоставленные тестовые данные 

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)
      */

Альтернатива-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]|
      * +---------+----+---------+
      */
2 голосов
/ 03 августа 2020

Используйте 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]|
+---------+----+---------+
...