Spark-Scala项目：州码字符串格式转换适配Spark SQL问询

在Spark-Scala项目中，需从Parquet格式的states表获取每个州的州码，表中数据结构如下：

state state_cd
GA    AGAHUI,AGAUTY,AGAERE
CA    BCAHRT,CCAYTU,CCARTE

需要将state_cd字段的逗号分隔字符串，转换为Spark-SQL的IN子句格式，生成如下查询条件：

WHERE state = 'GA' AND state_cd IN ('AGAHUI','AGAUTY','AGAERE')
WHERE state = 'CA' AND state_cd IN ('BCAHRT','CCAYTU','CCARTE')

核心需求是实现将AGAHUI,AGAUTY,AGAERE转换为('AGAHUI','AGAUTY','AGAERE')的逻辑。

1. 纯Scala字符串处理方法 #

直接通过字符串分割、转换和拼接实现格式转换，适用于单个字符串的处理场景：

/**
 * 将逗号分隔的州码字符串转换为IN子句所需的带单引号的括号格式
 * @param stateCdStr 原始逗号分隔的州码字符串
 * @return 格式化后的字符串，如('AGAHUI','AGAUTY','AGAERE')
 */
def formatStateCd(stateCdStr: String): String = {
  // 处理空值、空字符串及带空格的情况
  if (stateCdStr == null || stateCdStr.trim.isEmpty) {
    "('')" // 可根据业务需求调整，如返回"()"或抛出异常
  } else {
    stateCdStr.split(",")
      .map(_.trim) // 去除每个州码前后的空格
      .filter(_.nonEmpty) // 过滤分割后产生的空元素
      .map(code => s"'$code'") // 给每个州码添加单引号
      .mkString("(", ",", ")") // 拼接成括号包裹的格式
  }
}

// 测试示例
val gaStateCd = "AGAHUI,AGAUTY,AGAERE"
println(formatStateCd(gaStateCd)) // 输出: ('AGAHUI','AGAUTY','AGAERE')

2. Spark DataFrame内置函数处理 #

如果需要直接从DataFrame中生成格式化后的字段，可使用Spark内置函数实现批量转换：

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

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

// 读取Parquet表
val statesDF = spark.read.parquet("path/to/your/states.parquet")

// 添加格式化后的字段
val formattedStatesDF = statesDF.withColumn(
  "formatted_state_cd",
  expr("""
    concat(
      '(',
      concat_ws(',', transform(split(state_cd, ','), x -> concat('\\'', trim(x), '\\''))),
      ')'
    )
  """)
)

// 查看结果
formattedStatesDF.select("state", "state_cd", "formatted_state_cd").show(false)

执行后输出结果：

+-----+---------------------+-------------------------------+
|state|state_cd             |formatted_state_cd             |
+-----+---------------------+-------------------------------+
|GA   |AGAHUI,AGAUTY,AGAERE|('AGAHUI','AGAUTY','AGAERE')   |
|CA   |BCAHRT,CCAYTU,CCARTE|('BCAHRT','CCAYTU','CCARTE')   |
+-----+---------------------+-------------------------------+

3. 动态生成Spark-SQL查询语句 #

结合上述方法，可遍历DataFrame生成对应的查询语句：

// 遍历DataFrame的每一行，生成查询语句
statesDF.collect().foreach { row =>
  val state = row.getAs[String]("state")
  val formattedCd = formatStateCd(row.getAs[String]("state_cd"))
  val sqlQuery = s"SELECT * FROM your_target_table WHERE state = '$state' AND state_cd IN $formattedCd"
  
  // 此处可执行查询、打印或保存语句
  println(sqlQuery)
}

输出的查询语句示例：

SELECT * FROM your_target_table WHERE state = 'GA' AND state_cd IN ('AGAHUI','AGAUTY','AGAERE')
SELECT * FROM your_target_table WHERE state = 'CA' AND state_cd IN ('BCAHRT','CCAYTU','CCARTE')

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