从Spark导入数据至ClickHouse

代码示例如下。

package com.company.packageName

import java.util.Properties
import java.util.concurrent.ThreadLocalRandom

import scala.annotation.tailrec

import com.google.common.collect.ImmutableMap
import org.apache.spark.internal.Logging
import org.apache.spark.sql.{SaveMode, SparkSession}

case class Test(id: Int, key1: String, value1: Boolean, key2: Long, value2: Double)

object CKDataImporter extends Logging {

  private var dbName: String = "default"
  private var tableName: String = ""
  private var ckHost: String = ""
  private var ckPort: String = "8123"
  private var user: String = "default"
  private var password: String = ""
  private var local: Boolean = false

  def main(args: Array[String]): Unit = {
    parse(args.toList)
    checkArguments()

    val jdbcUrl = s"jdbc:clickhouse://$ckHost:$ckPort/$dbName"

    logInfo(s"Use jdbc: $jdbcUrl")
    logInfo(s"Use table: $tableName")

    val spark = getSparkSession

    // generate test data
    val rdd = spark.sparkContext.parallelize(1 to 1000).map(i => {
      val rand = ThreadLocalRandom.current()
      val randString = (0 until rand.nextInt(10, 20))
        .map(_ => rand.nextLong())
        .mkString("")

      Test(i, randString, rand.nextBoolean(), rand.nextLong(), rand.nextGaussian())
    })

    val df = spark.createDataFrame(rdd)

    df.write
      .mode(SaveMode.Append)
      .jdbc(jdbcUrl, tableName, getCKJdbcProperties(user, password))
  }

  private def printUsageAndExit(exitCode: Int = 0): Unit = {
    logError("Usage: java -jar /path/to/CKDataImporter.jar [options]")
    logError("  --dbName      设置ClickHouse数据库的名称，默认为default")
    logError("  --tableName   设置ClickHouse库中表的名称")
    logError("  --ckHost      设置ClickHouse地址")
    logError("  --ckPort      设置ClickHouse端口，默认为8123")
    logError("  --user        设置ClickHouse所使用的用户名")
    logError("  --password    设置ClickHouse用户的密码")
    logError("  --local       设置此程序使用Spark Local模式运行")
    System.exit(exitCode)
  }

  @tailrec
  private def parse(args: List[String]): Unit = args match {
    case ("--help" | "-h") :: _ =>
      printUsageAndExit()
    case "--dbName" :: value :: tail =>
      dbName = value
      parse(tail)
    case "--tableName" :: value :: tail =>
      tableName = value
      parse(tail)
    case "--ckHost" :: value :: tail =>
      ckHost = value
      parse(tail)
    case "--ckPort" :: value :: tail =>
      ckPort = value
      parse(tail)
    case "--user" :: value :: tail =>
      user = value
      parse(tail)
    case "--password" :: value :: tail =>
      password = value
      parse(tail)
    case "--local" :: tail =>
      local = true
      parse(tail)
    case Nil =>
    case _ =>
      printUsageAndExit(1)
  }

  private def checkArguments(): Unit = {
    if ("".equals(tableName) || "".equals(ckHost)) {
      printUsageAndExit(2)
    }
  }

  private def getCKJdbcProperties(
      user: String,
      password: String,
      batchSize: String = "1000",
      socketTimeout: String = "300000",
      numPartitions: String = "8",
      rewriteBatchedStatements: String = "true"): Properties = {
    val kvMap = ImmutableMap.builder()
      .put("driver", "ru.yandex.clickhouse.ClickHouseDriver")
      .put("user", user)
      .put("password", password)
      .put("batchsize", batchSize)
      .put("socket_timeout", socketTimeout)
      .put("numPartitions", numPartitions)
      .put("rewriteBatchedStatements", rewriteBatchedStatements)
      .build()
    val properties = new Properties
    properties.putAll(kvMap)
    properties
  }

  private def getSparkSession: SparkSession = {
    val builder = SparkSession.builder()
    if (local) {
      builder.master("local[*]")
    }
    builder.appName("ClickHouse-Data-Importer")
    builder.getOrCreate()
  }
}