How can I add an Index grouped by ID in Dataframe Spark

Question

I have this Dataframe

+---------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+
|_id      |details__line_items                                                                                                                                                  |searchable_tags|
+---------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+
|307131663|[[line_item_1345678912345678M, {}, {},, loan, 1,, 116000], [line_item_123456789, {}, {},, Test, 1,, 1234567], [line_item_2kZgNnPXvEgnKCAaM, {}, {},, loan, 1,, 1234]]|[]             |
|040013496|[[line_item_1345678912345678M, {}, {},, loan, 1,, 116000], [line_item_123456789, {}, {},, Test, 1,, 1234567], [line_item_2kZgNnPXvEgnKCAaM, {}, {},, loan, 1,, 1234]]|[]             |
+---------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+

I'm exploding the details__line_items column using this function:

 def getArrayDataFrame(df: DataFrame): ListBuffer[DataFrame] = {
    df.schema
      .filter(field => {
        field.dataType.typeName == "array"
      })
      .map(field => {
        val explodeColumn = (colsName: String) =>
          df.withColumn("items", explode(df.col(s"${field.name}")))
            .select("_id", colsName)
        field.dataType match {
          case arrayType: ArrayType => {
            arrayType.elementType.typeName match {
              case "struct" => explodeColumn("items.*")
              case _        => explodeColumn(s"${field.name}")
            }
          }
        }
      })
      .to[ListBuffer]
  }

I'm getting this Dataframe:

+---------+---------------------------+--------------+---------------+-----------+----+--------+----+----------+
|_id      |_id                        |antifraud_info|contextual_data|description|name|quantity|sku |unit_price|
+---------+---------------------------+--------------+---------------+-----------+----+--------+----+----------+
|307131663|line_item_1345678912345678M|{}            |{}             |null       |loan|1       |null|116000    |
|307131663|line_item_123456789        |{}            |{}             |null       |Test|1       |null|1234567   |
|307131663|line_item_2kZgNnPXvEgnKCAaM|{}            |{}             |null       |loan|1       |null|1234      |
|040013496|line_item_1345678912345678M|{}            |{}             |null       |loan|1       |null|116000    |
|040013496|line_item_123456789        |{}            |{}             |null       |Test|1       |null|1234567   |
|040013496|line_item_2kZgNnPXvEgnKCAaM|{}            |{}             |null       |loan|1       |null|1234      |
+---------+---------------------------+--------------+---------------+-----------+----+--------+----+----------+

How could I get a new Dataframe like this ?

+---------+---+---------------------------+-------------------+--------------------+----------------+---------+-------------+--------+---------------+
|_id      |index|_id                  |antifraud_info|contextual_data|description|name|quantity|sku|unit_price|
+---------+---+---------------------------+-------------------+--------------------+----------------+---------+-------------+--------+---------------+
|307131663|0  |line_item_1345678912345678M|{}                 |{}                  |null            |loan     |1            |null    |116000         |
|307131663|1  |line_item_123456789        |{}                 |{}                  |null            |Test     |1            |null    |1234567        |
|307131663|2  |line_item_2kZgNnPXvEgnKCAaM|{}                 |{}                  |null            |loan     |1            |null    |1234           |
|040013496|0  |line_item_1345678912345678M|{}                 |{}                  |null            |loan     |1            |null    |116000         |
|040013496|1  |line_item_123456789        |{}                 |{}                  |null            |Test     |1            |null    |1234567        |
|040013496|2  |line_item_2kZgNnPXvEgnKCAaM|{}                 |{}                  |null            |loan     |1            |null    |1234           |
+---------+---+---------------------------+-------------------+--------------------+----------------+---------+-------------+--------+---------------+

I have already tried using posexplode but it changes my dataframe schema adding the col and pos columns, I modified my function like this.

def getArrayDataFrame(df: DataFrame): ListBuffer[DataFrame] = {
    df.schema
      .filter(field => {
        field.dataType.typeName == "array"
      })
      .map{ (field) => {

        println(s"This is the name of the field ${field.name}")

        val testDF =  df.select($"_id", posexplode(df.col(s"${field.name}") ))

        testDF.printSchema()
        val newDF = testDF.select(flattenSchema(testDF.schema): _*)
        newDF.printSchema()
        newDF
      }}
      .to[ListBuffer]
  }

So, How could I get the index of the exploded column without changing my Dataframe schema ?

user_11702900 · Answer

In order to add an index column for each group use Window function partitionBy() and row_number function from spark sql functions.

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

df.withColumn("index", row_number(Window.partitionBy("col_to_groupBy").orderBy("some_col")))

The row_number() function will label each row with increasing index, to do it for each group we use Window function partitioBy() to group the DF based on some column(col_to_groupBy). Each group needs to be ordered in itself - so we use orderBy to order by some column(some_col). In your example the order has no meaning so you can choose any column you want.

How can I add an Index grouped by ID in Dataframe Spark

1 Answers