Объедините данные из нескольких строк в одну, а затем вложите данные

Question

Я относительно новичок в программировании scala и spark.

У меня есть сценарий использования, когда мне нужно сгруппировать данные, основанные на определенных столбцах, и иметь счетчик определенного столбца (используя сводную таблицу), а затем, наконец, мне нужно создать вложенный кадр данных из моего плоского кадра данных.

Одна из основных проблем, с которыми я сталкиваюсь, заключается в том, что мне нужно сохранить и некоторые другие столбцы (не те, на которых я разворачиваюсь).

Я не могу найти эффективный способ сделать это.

ВХОД

ID ID2 ID3 country items_purchased quantity
1 1    1     UK    apple           1      
1 1    1     USA   mango           1       
1 2    3     China banana          3       
2 1    1     UK    mango           1       

Теперь, скажем, я хочу повернуть «страну» и сгруппировать по («ID», «ID2», «ID3») Но я также хочу сохранить другие столбцы в виде списка.

Например,

ВЫХОД-1:

ID ID2 ID3 UK USA China items_purchased quantity
1  1   1   1  1    0    [apple,mango]   [1,1] 
1  2   3   0  0    1    [banana]        [3]
2  1   1   1  0    0    [mango]         [1]

Как только я достигну этого,

Я хочу вложить его во вложенную структуру, чтобы моя схема выглядела так:

{
"ID"  : 1,
"ID2" : 1,
"ID3" : 1,
"countries" : 
{
"UK" : 1,
"USA" : 1,
"China" : 0,
},
"items_purchased" : ["apple", "mango"],
"quantity" : [1,1]
}

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

Я имею в виду что-то вроде этих строк:

dataframe.map(row => myCaseClass(row.getAs[Long]("ID"),
row.getAs[Long]("ID2"),
row.getAs[Long]("ID3"),
CountriesCaseClass(
row.getAs[String]("UK")
)
)

и так далее ...

Answer 1

Я думаю, что это должно работать для вашего случая.Количество разделов рассчитывается по формуле partitions_num = data_size / 500MB.

import org.apache.spark.sql.functions.{collect_list, count, col, lit, map}

val data = Seq(
(1, 1, 1, "UK", "apple", 1),
(1, 1, 1, "USA","mango", 1),
(1, 2, 3, "China", "banana", 3),
(2, 1, 1, "UK", "mango", 1))

// e.g: partitions_num = 100GB / 500MB = 200, adjust it according to the size of your data
val partitions_num = 250
val df = data.toDF("ID", "ID2", "ID3", "country", "items_purchased", "quantity")
              .repartition(partitions_num, $"ID", $"ID2", $"ID3") //the partition should remain the same for all the operations
              .persist()

//get countries, we will need it to fill with 0 the null values after pivoting, for the mapping and for the drop
val countries = df.select("country").distinct.collect.map{_.getString(0)}

//creates a sequence of key/value which should be the input for the map function
val countryMapping = countries.flatMap{c => Seq(lit(c), col(c))}
val pivotCountriesDF = df.select("ID", "ID2", "ID3", "country")
                          .groupBy("ID", "ID2", "ID3")
                          .pivot($"country")
                          .count()
                          .na.fill(0, countries)
                          .withColumn("countries", map(countryMapping:_*))//i.e map("UK", col("UK"), "China", col("China")) -> {"UK":0, "China":1}
                          .drop(countries:_*)

// pivotCountriesDF.rdd.getNumPartitions == 250, Spark will retain the partition number since we didnt change the partition key

// +---+---+---+-------------------------------+
// |ID |ID2|ID3|countries                      |
// +---+---+---+-------------------------------+
// |1  |2  |3  |[China -> 1, USA -> 0, UK -> 0]|
// |1  |1  |1  |[China -> 0, USA -> 1, UK -> 1]|
// |2  |1  |1  |[China -> 0, USA -> 0, UK -> 1]|
// +---+---+---+-------------------------------+

val listDF = df.select("ID", "ID2", "ID3", "items_purchased", "quantity")
                .groupBy("ID", "ID2", "ID3")
                .agg(
                  collect_list("items_purchased").as("items_purchased"), 
                  collect_list("quantity").as("quantity"))

// +---+---+---+---------------+--------+
// |ID |ID2|ID3|items_purchased|quantity|
// +---+---+---+---------------+--------+
// |1  |2  |3  |[banana]       |[3]     |
// |1  |1  |1  |[apple, mango] |[1, 1]  |
// |2  |1  |1  |[mango]        |[1]     |
// +---+---+---+---------------+--------+


// listDF.rdd.getNumPartitions == 250, to validate this try to change the partition key with .groupBy("ID", "ID2") it will fall back to the default 200 value of spark.sql.shuffle.partitions setting 

val joinedDF = pivotCountriesDF.join(listDF, Seq("ID", "ID2", "ID3"))

// joinedDF.rdd.getNumPartitions == 250, the same partitions will be used for the join as well.

// +---+---+---+-------------------------------+---------------+--------+
// |ID |ID2|ID3|countries                      |items_purchased|quantity|
// +---+---+---+-------------------------------+---------------+--------+
// |1  |2  |3  |[China -> 1, USA -> 0, UK -> 0]|[banana]       |[3]     |
// |1  |1  |1  |[China -> 0, USA -> 1, UK -> 1]|[apple, mango] |[1, 1]  |
// |2  |1  |1  |[China -> 0, USA -> 0, UK -> 1]|[mango]        |[1]     |
// +---+---+---+-------------------------------+---------------+--------+

joinedDF.toJSON.show(false)

// +--------------------------------------------------------------------------------------------------------------------+
// |value                                                                                                               |
// +--------------------------------------------------------------------------------------------------------------------+
// |{"ID":1,"ID2":2,"ID3":3,"countries":{"China":1,"USA":0,"UK":0},"items_purchased":["banana"],"quantity":[3]}         |
// |{"ID":1,"ID2":1,"ID3":1,"countries":{"China":0,"USA":1,"UK":1},"items_purchased":["apple","mango"],"quantity":[1,1]}|
// |{"ID":2,"ID2":1,"ID3":1,"countries":{"China":0,"USA":0,"UK":1},"items_purchased":["mango"],"quantity":[1]}          |
// +--------------------------------------------------------------------------------------------------------------------++

Удачи и дайте мне знать, если вам нужны какие-либо разъяснения.

Answer 2

Я не вижу никаких проблем, это хорошее решение.В любом случае, я бы создал 'Dataset for your final dataframe`.Работать проще.

val ds: Dataset[myCaseClass] = dataframe.map(row => myCaseClass(row.getAs[Long]("ID"),
  ...

РЕДАКТИРОВАТЬ Вы просите что-то вроде этого.

input.groupby("ID","ID2","ID3")
  .withColumn("UK", col("country").contains("UK"))
  .withColumn("China", col("country").contains("China"))
  .withColumnRenamed("country","USA", col("country").contains("USA"))