CODEWITHSUNDEEP
apache-sparkdataframeapache-spark-sql
samst
samst

Reputation: 556

extract or filter MapType of Spark DataFrame

I have a DataFrame that contains various columns. One column contains a Map[Integer,Integer[]]. It looks like { 2345 -> [1,34,2]; 543 -> [12,3,2,5]; 2 -> [3,4]} Now what I need to do is filter out some keys. I have a Set of Integers (javaIntSet) in Java with which I should filter such that

col(x).keySet.isin(javaIntSet)

ie. the above map should only contain the key 2 and 543 but not the other two and should look like {543 -> [12,3,2,5]; 2 -> [3,4]} after filtering.

Documentation of how to use the Java Column Class is sparse. How do I extract the col(x) such that I can just filter it in java and then replace the cell data with a filtered map. Or are there any useful functions of columns I am overlooking. Can I write an UDF2<Map<Integer, Integer[]>,Set<Integer>,Map<Integer,Integer[]> I can write an UDF1<String,String> but I am not so sure how it works with more complex parameters.

Generally the javaIntSet is only a dozen and usually less than a 100 values. The Map usually also has only a handful entries (0-5 usually).

I have to do this in Java (unfortunately) but I am familiar with Scala. A Scala answer that I translate myself to Java would already be very helpful.

Upvotes: 1

Views: 3183

Answers (1)

David Griffin
David Griffin

Reputation: 13927

You don't need a UDF. Might be cleaner with one, but you could just as easily do it with DataFrame.explode:

case class MapTest(id: Int, map: Map[Int,Int])
val mapDf = Seq(
  MapTest(1, Map((1,3),(2,10),(3,2)) ),
  MapTest(2, Map((1,12),(2,333),(3,543)) )
).toDF("id", "map")

mapDf.show
+---+--------------------+
| id|                 map|
+---+--------------------+
|  1|Map(1 -> 3, 2 -> ...|
|  2|Map(1 -> 12, 2 ->...|
+---+--------------------+

Then you can use explode:

mapDf.explode($"map"){
  case Row(map: Map[Int,Int] @unchecked) => {
    val newMap = map.filter(m => m._1 != 1)   // <-- do filtering here
    Seq(Tuple1(newMap)) 
  }
}.show
+---+--------------------+--------------------+
| id|                 map|                  _1|
+---+--------------------+--------------------+
|  1|Map(1 -> 3, 2 -> ...|Map(2 -> 10, 3 -> 2)|
|  2|Map(1 -> 12, 2 ->...|Map(2 -> 333, 3 -...|
+---+--------------------+--------------------+

If you did want to do the UDF, it would look like this:

val mapFilter = udf[Map[Int,Int],Map[Int,Int]](map => {
  val newMap = map.filter(m => m._1 != 1)   // <-- do filtering here
  newMap
})

mapDf.withColumn("newMap", mapFilter($"map")).show
+---+--------------------+--------------------+
| id|                 map|              newMap|
+---+--------------------+--------------------+
|  1|Map(1 -> 3, 2 -> ...|Map(2 -> 10, 3 -> 2)|
|  2|Map(1 -> 12, 2 ->...|Map(2 -> 333, 3 -...|
+---+--------------------+--------------------+

DataFrame.explode is a little more complicated, but ultimately more flexible. For example, you could divide the original row into two rows -- one containing the map with the elements filtered out, the other a map with the reverse -- the elements that were filtered.

Upvotes: 4

Related Questions