Создать новый столбец с целым диапазоном, используя существующий столбец целых чисел в Spark Scala Dataframe - PullRequest
0 голосов
/ 12 октября 2019

Предположим, у меня есть объект Spark Scala DataFrame, например:

+--------+
|col1    |
+--------+
|1       |
|3       |
+--------+

И мне нужен объект DataFrame, подобный:

+-----------------+
|col1  |col2      |
+-----------------+
|1     |[0,1]     |
|3     |[0,1,2,3] |
+-----------------+

Ответы [ 2 ]

0 голосов
/ 12 октября 2019
You can achieve it with the below approach

    val input_df = spark.sparkContext.parallelize(List(1, 2, 3, 4, 5)).toDF("col1")
    input_df.show(false)
Input:
+----+
|col1|
+----+
|1   |
|2   |
|3   |
|4   |
|5   |
+----+

    val output_df = input_df.rdd.map(x => x(0).toString()).map(x => (x, Range(0, x.toInt + 1).mkString(","))).toDF("col1", "col2")
    output_df.withColumn("col2", split($"col2", ",")).show(false)

Output:
+----+------------------+
|col1|col2              |
+----+------------------+
|1   |[0, 1]            |
|2   |[0, 1, 2]         |
|3   |[0, 1, 2, 3]      |
|4   |[0, 1, 2, 3, 4]   |
|5   |[0, 1, 2, 3, 4, 5]|
+----+------------------+

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

0 голосов
/ 12 октября 2019

Spark предлагает множество API / функций, которые можно поиграть, но большинство функций по умолчанию пригодятся, однако для конкретной задачи могут быть написаны пользовательские функции UserDefinedFunctions.

Ссылка https://jaceklaskowski.gitbooks.io/mastering-spark-sql/spark-sql-udfs.html

import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.functions.col
import spark.implicits._

val df = spark.sparkContext.parallelize(Seq(1,3)).toDF("index")
val rangeDF = df.withColumn("range", indexToRange(col("index")))
rangeDF.show(10)

def indexToRange: UserDefinedFunction = udf((index: Integer) => for (i <- 0 to index) yield i)
...