基于Scala字典实现DataFrame州代码到州名称的映射并新增列

First, let's fix a tiny syntax error in your stateCodes map — you missed a closing quote for the "tx" value. The corrected map should be:

val stateCodes = Map("100" -> "ca", "22" -> "ny", "12" -> "tx", "71" -> "va")

There are two clean ways to achieve your goal of adding the stateName column, with fallback to "unknown" when the stateId isn't in the map:

Method 1: Custom UDF (Intuitive for Custom Logic) #

If you prefer a straightforward approach that's easy to tweak for more complex rules, use a User-Defined Function (UDF). For better performance with large datasets, broadcast the map to avoid duplicating it across every executor task.

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

// Broadcast the map to optimize performance
val broadcastStateCodes = spark.sparkContext.broadcast(stateCodes)

// Define the UDF to look up state names
val getStateName = udf((stateId: String) => broadcastStateCodes.value.getOrElse(stateId, "unknown"))

// Transform the original DataFrame
val new_df = df.withColumn("stateName", getStateName($"stateId"))
  .drop("stateId") // Optional: remove the original stateId column if you don't need it

Method 2: Spark Built-in Functions (No UDF Needed) #

If you want to avoid UDFs entirely, use Spark's built-in create_map and coalesce functions to build a lookup directly in the DataFrame API:

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

// Convert the Scala map into a Spark map column
val stateMap = create_map(
  stateCodes.flatMap { case (k, v) => Seq(lit(k), lit(v)) }.toSeq: _*
)

// Add the stateName column, using "unknown" when no match is found
val new_df = df.withColumn("stateName", coalesce(stateMap($"stateId"), lit("unknown")))
  .drop("stateId") // Optional: remove original stateId column

Output Verification #

Both methods will produce your desired output DataFrame:

+---------+--------+
|stateName|   count|
+---------+--------+
|       ca|17243160|
|       va|12398011|
|       tx|12302708|
|  unknown| 9734593|
+---------+--------+

Quick Notes #

• If your stateId is an integer type instead of string, cast it to string in the lookup (e.g., $"stateId".cast(StringType)).
• The broadcast variable in Method 1 is especially useful for large maps, as it cuts down on memory usage across executors.

内容的提问来源于stack exchange，提问作者user1579557