spark 2.4版本写SR怎么去优化

[专注数据者]

spark低版本我使用http方式写，但是这个写入速度不快，下面我我的代码，这个能优化吗 def writeDataFrameToStarRocksByLoad(proper: java.util.Properties, data: DataFrame, table: String, updateColumnsName: String = “”, isPartialUpdate: Boolean = false, BATCH_SIZE: Int = 50000): Unit = {

//生成待写入表的连接信息
val user = proper.getProperty("user", "")
val password = proper.getProperty("password", "")
val port = proper.getProperty("fe.http.port", "")
val feHost = proper.getProperty("fe.host", "")
val dbName = table.split("\\.")(0)
val tableName = table.split("\\.")(1)

// 拼接你这行的 Load URL
val loadUrl = String.format(
  "http://%s:%s/api/%s/%s/_stream_load",
  feHost, port, dbName, tableName
)
val tobeEncode = user + ":" + password
val encoded = Base64.encodeBase64(tobeEncode.getBytes(StandardCharsets.UTF_8))
val basicAuthHeader = "Basic " + new String(encoded)

// GZIP压缩工具
def gzipCompress(raw: String): Array[Byte] = {
  val bos = new ByteArrayOutputStream()
  val gz = new GZIPOutputStream(bos)
  try {
    gz.write(raw.getBytes(StandardCharsets.UTF_8))
    gz.finish()
    bos.toByteArray
  } finally {
    gz.close()
    bos.close()
  }
}

data.toJSON.rdd.foreachPartition(partition => {
  val partitionId = TaskContext.getPartitionId()
  // 批次缓存列表，控制内存峰值
  val batchBuffer = new ListBuffer[String]()
  var submit = 0;

  if (partition.nonEmpty) {
    // 单个分区共用1个http客户端，避免频繁创建销毁
    val httpClient = HttpClients.custom()
      .setRedirectStrategy(new DefaultRedirectStrategy() {
        override protected def isRedirectable(method: String): Boolean = true
      })
    val client = httpClient.build()

    try {
      val httpPut = new HttpPut(loadUrl)
      // 请求头
      httpPut.setHeader(HttpHeaders.EXPECT, "100-continue")
      httpPut.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader)
      // JSON 格式
      httpPut.removeHeaders(HttpHeaders.CONTENT_LENGTH)
      httpPut.removeHeaders(HttpHeaders.TRANSFER_ENCODING)
      httpPut.setHeader("format", "json")
      httpPut.setHeader("strip_outer_array", "true")
      httpPut.setHeader("ignore_json_size", "true")
      // 部分列更新
      if (isPartialUpdate) {
        httpPut.setHeader("partial_update", "true")
        httpPut.setHeader("columns", updateColumnsName)
      }

      while (partition.hasNext) {
        batchBuffer += partition.next()
        if (batchBuffer.size >= BATCH_SIZE) {
          httpPut.removeHeaders("label")
          submit += 1;
          // 把 DF 转成 JSON 字符串
          val jsonContent = batchBuffer.mkString("[", ",", "]")
          val entity = new StringEntity(jsonContent, StandardCharsets.UTF_8)
          val label = s"spark_load_${tableName}_${partitionId}_${System.currentTimeMillis()}" + "_" + submit
          httpPut.setHeader("label", label)
          httpPut.setEntity(entity)
          // 发送请求
          val response = client.execute(httpPut)
          try {
            val statusCode = response.getStatusLine.getStatusCode
            val loadResult = if (response.getEntity != null) {
              EntityUtils.toString(response.getEntity, StandardCharsets.UTF_8)
            } else {
              ""
            }
            if (statusCode != 200) {
              throw new Exception(s"Stream Load 失败 statusCode=$statusCode, result=$loadResult")
            }
            println(s"分区 ${TaskContext.get.partitionId()} 写入成功：$loadResult")
          } finally {
            response.close()
          }
          batchBuffer.clear()
        }
      }

      submit += 1;
      // 把 DF 转成 JSON 字符串
      val jsonContent = batchBuffer.mkString("[", ",", "]")
      val entity = new StringEntity(jsonContent, StandardCharsets.UTF_8)
      val label = s"spark_load_${tableName}_${partitionId}_${System.currentTimeMillis()}" + "_" + submit

      httpPut.setHeader("label", label + "_" + submit)
      httpPut.setEntity(entity)
      // 发送请求
      val response = client.execute(httpPut)
      try {
        val statusCode = response.getStatusLine.getStatusCode
        val loadResult = if (response.getEntity != null) {
          EntityUtils.toString(response.getEntity, StandardCharsets.UTF_8)
        } else {
          ""
        }
        if (statusCode != 200) {
          throw new Exception(s"Stream Load 失败 statusCode=$statusCode, result=$loadResult")
        }
        println(s"分区 ${TaskContext.get.partitionId()} 写入成功：$loadResult")
      } finally {
        response.close()
        batchBuffer.clear()
      }

    } finally {
      client.close()
    }

  }

})
logger.info("---写入的表明：" + table + "---写入的数据量是" + data.count())

}