выбрать не более 3 первых элементов после разделения

Question

у меня в кадре есть датафрейм:

Col_A | Col_B
--------------
  A    |    1,6,5,3,7      
  B    |   17,18,78     
  C    |    90  
  D    |    56,57

и Col_A, и Col_B имеют строковый тип.

как выбрать (максимум) первые 3 элемента (или только 1 или 2, если имеется менее 3 элементов) после разделения определенной строки разделителем (",") и преобразования указанного выше фрейма данных в новый фрейм данных ниже:

Col_A  |  Col_B
----------------    
A      |  1    
A      |  6    
A      |  5    
B      |  17    
B      |  18    
B      |  78    
C      |   90    
D      |   56
D      |   57

здесь также оба col_A и col_B имеют строковый тип.

Answer 1

Вы можете получить ваш окончательный требуемый кадр данных, используя maps и flatMaps как

//get the column names
val columns = df.columns
//transform the dataframe using dataset way
df.flatMap(row => row.getAs[String]("Col_B").trim.split(",").take(3).map(value => (row.getAs[String]("Col_A"), value))).toDF(columns: _*).show(false)

, что должно дать вам

+-----+-----+
|Col_A|Col_B|
+-----+-----+
|A    |1    |
|A    |6    |
|A    |5    |
|B    |17   |
|B    |18   |
|B    |78   |
|C    |90   |
|D    |56   |
|D    |57   |
+-----+-----+

Answer 2

Вы можете написать простое udf, которое splits String и возвращает самое большее 3 значения из array как

val origDF = Seq(
  ("A", "1,6,5,3,7"),
  ("B", "17,18,78"),
  ("C", "90"),
  ("D", "56,57")
).toDF("Col_A", "Col_B")

val valueAsArray = udf((value: String) => {
  value.split(",").map(_.trim).take(3)
})

Теперь используйте udf с withColumn иexplode

origDF.withColumn("Col_B", explode(valueAsArray($"Col_B"))).show(false)

Вывод:

+-----+-----+
|Col_A|Col_B|
+-----+-----+
|A    |1    |
|A    |6    |
|A    |5    |
|B    |17   |
|B    |18   |
|B    |78   |
|C    |90   |
|D    |56   |
|D    |57   |
+-----+-----+

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