How to split comma separated string and get n values in Spark Scala dataframe? - scala

Question

How to take only 2 data from arraytype column in Spark Scala?
I got the data like val df = spark.sqlContext.sql("select col1, col2 from test_tbl").
I have data like following:
col1 | col2
--- | ---
a | [test1,test2,test3,test4,.....]
b | [a1,a2,a3,a4,a5,.....]
I want to get data like following:
col1| col2
----|----
a | test1,test2
b | a1,a2
When I am doing df.withColumn("test", col("col2").take(5)) it is not working. It give this error:
value take is not a member of org.apache.spark.sql.ColumnName
How can I get the data in above order?

Answer 1

Inside withColumn you can call udf getPartialstring for that you can use slice or take method like below example snippet untested.
import sqlContext.implicits._
import org.apache.spark.sql.functions._
val getPartialstring = udf((array : Seq[String], fromIndex : Int, toIndex : Int)
=> array.slice(fromIndex ,toIndex ).mkString(","))
your caller will appear like
df.withColumn("test",getPartialstring(col("col2"))
col("col2").take(5) is failing because column doesn't have a method take(..) that's why your error message says
error: value take is not a member of org.apache.spark.sql.ColumnName
You can use udf approach to tackle this.

Answer 2

You can use the array Column's apply function to get each individual item up to a certain index, and then build a new array using the array function:
import spark.implicits._
import org.apache.spark.sql.functions._
// Sample data:
val df = Seq(
("a", Array("a1", "a2", "a3", "a4", "a5", "a6")),
("a", Array("b1", "b2", "b3", "b4", "b5")),
("c", Array("c1", "c2"))
).toDF("col1", "col2")
val n = 4
val result = df.withColumn("col2", array((0 until n).map($"col2"(_)): _*))
result.show(false)
// +----+--------------------+
// |col1|col2 |
// +----+--------------------+
// |a |[a1, a2, a3, a4] |
// |a |[b1, b2, b3, b4] |
// |c |[c1, c2, null, null]|
// +----+--------------------+
Note that this will "pad" results with nulls for records with arrays smaller than n.