Apache Spark calculating column value on the basis of distinct value of columns

Viewed 56

I am processing the following tables and I would like to compute a new column (outcome) based on the distinct value of 2 other columns.

| id1  | id2 | outcome
|  1   |  1  |  1
|  1   |  1  |  1
|  1   |  3  |  2
|  2   |  5  |  1 
|  3   |  1  |  1  
|  3   |  2  |  2
|  3   |  3  |  3

The outcome should begin in incremental order starting from 1 based on the combined value of id1 and id2. Any hints how this can be accomplished in Scala. row_number doesn't seem to be useful here in this case.

The logic here is that for each unique value of id1 we will start numbering the outcome with min(id2) for corresponding id1 being assigned a value of 1.

2 Answers

You could try dense_rank()

with your example

      val df = sqlContext
        .read
        .option("sep","|")
        .option("header", true)
        .option("inferSchema",true)
        .csv("/home/cloudera/files/tests/ids.csv") // Here we read the .csv files
        .cache()

      df.show()
      df.printSchema()

      df.createOrReplaceTempView("table")
      sqlContext.sql(
        """
          |SELECT id1, id2, DENSE_RANK() OVER(PARTITION BY id1 ORDER BY id2) AS outcome
          |FROM table
          |""".stripMargin).show()

output

+---+---+-------+
|id1|id2|outcome|
+---+---+-------+
|  2|  5|      1|
|  1|  1|      1|
|  1|  1|      1|
|  1|  3|      2|
|  3|  1|      1|
|  3|  2|      2|
|  3|  3|      3|
+---+---+-------+

Use Window function to club(partition) them by first id and then order each partition based on second id.

Now you just need to assign a rank (dense_rank) over each Window partition.

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

df
.withColumn("outcome", dense_rank().over(Window.partitionBy("id1").orderBy("id2")))

Related