Добавить значение после каждого элемента в Dataframe списка PySpark

Question

У меня есть такой фрейм данных

 Data           ID   

[1,2,3,4]        22    

Я хочу создать новый столбец, и каждая запись в новом столбце будет значением из поля данных, к которому добавлен идентификатор с помощью символа ~, как показано ниже

 Data         ID               New_Column

[1,2,3,4]     22               [1|22~2|22~3|22~4|22]

Примечание : В поле Data размер массива не фиксированный. В нем может не быть записи или там будет N записей. Может ли кто-нибудь помочь мне решить!

Answer 1

package spark

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

object DF extends App {

  val spark = SparkSession.builder()
    .master("local")
    .appName("DataFrame-example")
    .getOrCreate()

  import spark.implicits._

  val df = Seq(
    (22, Seq(1,2,3,4)),
    (23, Seq(1,2,3,4,5,6,7,8)),
    (24, Seq())
  ).toDF("ID", "Data")

  val arrUDF = udf((id: Long, array: Seq[Long]) => {
    val r = array.size match {
      case 0 => ""
      case _ => array.map(x => s"$x|$id").mkString("~")
    }

    s"[$r]"
  })

  val resDF = df.withColumn("New_Column", lit(arrUDF('ID, 'Data)))

  resDF.show(false)
  //+---+------------------------+-----------------------------------------+
  //|ID |Data                    |New_Column                               |
  //+---+------------------------+-----------------------------------------+
  //|22 |[1, 2, 3, 4]            |[1|22~2|22~3|22~4|22]                    |
  //|23 |[1, 2, 3, 4, 5, 6, 7, 8]|[1|23~2|23~3|23~4|23~5|23~6|23~7|23~8|23]|
  //|24 |[]                      |[]                                       |
  //+---+------------------------+-----------------------------------------+




}

Answer 2

Spark 2.4 +

Эквивалент Pyspark для того же самого выглядит как

df = spark.createDataFrame([(22, [1,2,3,4]),(23, [1,2,3,4,5,6,7,8]),(24, [])],['Id','Data'])

df.show()

+---+--------------------+
| Id|                Data|
+---+--------------------+
| 22|        [1, 2, 3, 4]|
| 23|[1, 2, 3, 4, 5, 6...|
| 24|                  []|
+---+--------------------+

df.withColumn('ff', f.when(f.size('Data')==0,'').otherwise(f.expr('''concat_ws('~',transform(Data, x->concat(x,'|',Id)))'''))).show(20,False)

+---+------------------------+---------------------------------------+
|Id |Data                    |ff                                     |
+---+------------------------+---------------------------------------+
|22 |[1, 2, 3, 4]            |1|22~2|22~3|22~4|22                    |
|23 |[1, 2, 3, 4, 5, 6, 7, 8]|1|23~2|23~3|23~4|23~5|23~6|23~7|23~8|23|
|24 |[]                      |                                       |
+---+------------------------+---------------------------------------+

Если вы хотите получить окончательный результат в виде массива

df.withColumn('ff',f.array(f.when(f.size('Data')==0,'').otherwise(f.expr('''concat_ws('~',transform(Data, x->concat(x,'|',Id)))''')))).show(20,False)

+---+------------------------+-----------------------------------------+
|Id |Data                    |ff                                       |
+---+------------------------+-----------------------------------------+
|22 |[1, 2, 3, 4]            |[1|22~2|22~3|22~4|22]                    |
|23 |[1, 2, 3, 4, 5, 6, 7, 8]|[1|23~2|23~3|23~4|23~5|23~6|23~7|23~8|23]|
|24 |[]                      |[]                                       |
+---+------------------------+-----------------------------------------+

Надеюсь, это поможет

Answer 3

udf может помочь:

def func(array, suffix):
    return '~'.join([str(x) + '|' + str(suffix) for x in array])

from pyspark.sql.types import StringType
from pyspark.sql import functions as F
my_udf = F.udf(func, StringType())

df.withColumn("New_Column", my_udf("Data", "ID")).show()

распечатывает

+------------+---+-------------------+
|        Data| ID|      New_Column   |
+------------+---+-------------------+
|[1, 2, 3, 4]| 22|22~1|22~2|22~3|22~4|
+------------+---+-------------------+