Spark map lookup is O(n)

[DanSkells]

In my day job, we encountered performance issues when we were performing lookups in a map with > 1,000,000 element. After profiling, we realised that the map lookup is a linear scan, and looking to either optimise this or understand if some other library or other programming metaphor exists within Spark. The use case is insertion and lookup, without any need to remove elements. Happy to discuss / work with maintainers if appropriate to address the issue. I have seen the OpenHashMap implementation, but this seems to be a Spark internal implementation, so not sure if this would be an appropriate base to work from.

[gaogaotiantian]

Could you share a minimum example demonstrating the issue? Which map specifically you are talking about?

[DanSkells]

So we do something like:

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

val spark = SparkSession.builder().appName("testApp").getOrCreate()

import spark.implicits._

val map = spark.read.parquet("somePath") // Some path with 1 million + records
   .select($"key", $"value").as[(String, String)].collect().toMap()

val mapCol = typedLit(map)

val dataframe = spark.read.parquet("otherPath").select(element_at($"key", mapCol))

This then does a linear scan through the map.

[HyukjinKwon]

What specifically does the linear scan? Can you share the profile result?

[DanSkells]

I don't have the profile to hand, but this is the code that it ends up hitting, and there's an 11 year old TODO stating that this is a linear scan: https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala#L444. The original PR where this TODO was created was the PR that added the MapData type to begin with, so it seems that this hasn't been addressed since then.

[LuciferYang]

I submitted a pr with an optimization approach of "trading space for time" and added microbenchmark tests. The results from the microbenchmark data look quite promising:

• [SPARK-55959][SQL] Optimize Map Key Lookup for GetMapValue and ElementAt #54748