Spark与Scala：如何从onComplete成功回调返回值构建DataFrame？

问题核心分析 #

你现在拿不到预期值的核心原因是：onComplete方法的返回类型是Unit，所以你的executeQuery函数实际上返回的是()，而不是你想要的查询响应。Scala的Future是异步操作载体，不能直接从回调里"返回"值给同步代码，得用Future本身来承载结果并处理。

第一步：修复单个查询的返回值问题 #

先修改executeQuery，让它直接返回Future[YourResponseType]（把YourResponseType替换成Druid查询实际返回的结果类型），不要在方法内部提前消费onComplete：

import ing.wbaa.druid.SQLQuery
import scala.concurrent.ExecutionContext
import scala.util.{Failure, Success}
import scala.concurrent.Future

// 全局ExecutionContext，用于处理Future异步逻辑
implicit val executionContext: ExecutionContext = ExecutionContext.Implicits.global

// 修改后的executeQuery：直接返回execute的Future结果
def executeQuery(foo: SQLQuery): Future[YourResponseType] = {
  foo.execute // foo.execute本身就是Future[YourResponseType]，直接返回即可
}

// 调用示例
val query = SQLQuery(s"SELECT * FROM table1 WHERE col1 IN ('someString', 'someString2')")
val helloFuture: Future[YourResponseType] = executeQuery(query)

// 处理结果的两种方式：
// 1. 异步处理（推荐，非阻塞，适合生产环境）
helloFuture.onComplete {
  case Success(resp) => {
    println(resp)
    // 在这里可以直接处理resp，比如后续转换成Spark DataFrame的元素
  }
  case Failure(ex) => ex.printStackTrace()
}

// 2. 同步阻塞获取结果（仅测试或必须同步的场景用，不推荐生产代码）
import scala.concurrent.Await
import scala.concurrent.duration._
val hello = Await.result(helloFuture, 10.seconds) // 最多等待10秒
println(hello)

第二步：并行批量查询的正确实现 #

你要处理字符串列表并批量并行查询，这里要注意：par是针对同步操作的并行化，而Future本身已经是异步的，应该用Future的组合操作来实现高效并行。

修改后的代码示例（包含Spark DataFrame构建逻辑）：

import org.apache.spark.sql.{SparkSession, Row}
import org.apache.spark.sql.types._
import scala.concurrent.Future

// 初始化SparkSession（根据你的环境调整配置）
val spark = SparkSession.builder()
  .appName("DruidBatchQueryToSpark")
  .master("local[*]") // 生产环境请移除该行
  .getOrCreate()

// 待查询的字符串列表
val queryList = List("someString1", "someString2", ..., "someStringN")
// 按1000个一组拆分，避免单个IN子句过长
val splittedBatches: List[List[String]] = queryList.grouped(1000).toList

// 支持批量参数的executeQuery
def executeQuery(batch: List[String]): Future[YourResponseType] = {
  val inClause = batch.map(s => s"'$s'").mkString(",")
  val query = SQLQuery(s"SELECT * FROM table1 WHERE col1 IN ($inClause)")
  query.execute
}

// 生成所有批量查询的Future集合
val allBatchFutures: List[Future[YourResponseType]] = splittedBatches.map(executeQuery)

// 把多个Future合并成一个Future，包含所有批量查询的结果
val combinedResultsFuture: Future[Seq[YourResponseType]] = Future.sequence(allBatchFutures)

// 处理合并后的结果，构建Spark DataFrame
combinedResultsFuture.onComplete {
  case Success(allResponses) => {
    // 把所有Druid响应转换成Spark Row集合
    // 这里需要根据Druid返回的实际类型实现转换逻辑，比如每个resp是一行数据的集合
    val allRows = allResponses.flatMap(resp => {
      // 示例：如果resp是List[Map[String, Any]]，转换成Row
      resp.map(dataMap => Row(
        dataMap.get("col1").asInstanceOf[String],
        dataMap.get("col2").asInstanceOf[Int],
        // 其他字段按实际类型转换
      ))
    })

    // 定义DataFrame的Schema（和你的查询结果字段一一对应）
    val dfSchema = StructType(Array(
      StructField("col1", StringType, nullable = true),
      StructField("col2", IntegerType, nullable = true),
      // 添加其他字段...
    ))

    // 创建Spark DataFrame
    val resultDF = spark.createDataFrame(spark.sparkContext.parallelize(allRows), dfSchema)
    resultDF.show()
    // 后续可以对DF进行操作：保存、分析等
  }
  case Failure(ex) => {
    ex.printStackTrace()
    // 批量查询失败的错误处理逻辑
  }
}

// 如果是脚本运行，需要等待异步逻辑完成（生产环境Spark作业可忽略，根据生命周期调整）
Await.result(combinedResultsFuture, 5.minutes)

关键注意事项 #

1. 避免不必要的阻塞：尽量用onComplete、map、flatMap等异步方式处理Future，不要滥用Await.result，尤其是Spark作业中，阻塞线程会降低集群资源利用率。
2. 自定义线程池：如果需要调整并行查询的线程数，可以自定义ExecutionContext：

   import java.util.concurrent.ForkJoinPool
   implicit val ec: ExecutionContext = ExecutionContext.fromExecutor(new ForkJoinPool(10))
   

3. 响应类型适配：必须根据Druid查询返回的实际响应结构，实现到SparkRow的转换逻辑，确保字段类型和Schema匹配。

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