Replicate values in a column using Scala/Spark

Question

I have a specific problem and needs to be solved using Scala/SPARK.

I have a column in a Dataframe as shown below

ColA  

Tag1  
Tag2  
Tag3  
Tag1  
Tag2  
Tag3  
Tag1  
Tag2  
Tag3  

Now I want to include a new column in the dataframe as shown below in the required format.

ColA  ColB   

Tag1  1    
Tag2  1  
Tag3  1  
Tag1  2  
Tag2  2  
Tag3  2  
Tag1  3  
Tag2  3  
Tag3  3  

Can this be done in Scala/Spark. I am newbie in Scala/Spark and really stuck with it for a long time now. Please help me on this.

Answer 1

Imports:

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

Solution:

df.groupBy("ColA").agg(collect_list("ColA").alias("tmp"))
  .select(posexplode($"tmp"))

Answer 2

You can use Window function (row_number()) to generate ColB as below

import org.apache.spark.sql.expressions.Window
val windowSpec = Window.partitionBy("ColA").orderBy("ColA")
import org.apache.spark.sql.functions._
df.withColumn("ColB", row_number().over(windowSpec)).show(false)

You should have the following output

+----+----+
|ColA|ColB|
+----+----+
|Tag1|1   |
|Tag1|2   |
|Tag1|3   |
|Tag3|1   |
|Tag3|2   |
|Tag3|3   |
|Tag2|1   |
|Tag2|2   |
|Tag2|3   |
+----+----+

Answer 3

Try this:

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

df.createOrReplaceTempView("tab")

val q = """
SELECT ColA, ROW_NUMBER() OVER(PARTITION BY ColA ORDER BY ColA) as ColB
FROM tab
ORDER BY ColB, ColA
"""

val res = spark.sql(q)

Result:

scala> res.show
+----+----+
|ColA|ColB|
+----+----+
|Tag1|   1|
|Tag2|   1|
|Tag3|   1|
|Tag1|   2|
|Tag2|   2|
|Tag3|   2|
|Tag1|   3|
|Tag2|   3|
|Tag3|   3|
+----+----+