SparkSQL自定义DataSourceV2源实现指南
ckckck2025年1月10日1 浏览基于SparkSQL DataSourceV2实现自定义数据源
基于SparkSQL DataSourceV2实现输入源
SparkSQL的DataSourceV2实现与StructuredStreaming自定义数据源类似，思想一致，但具体实现有所不同。主要步骤如下：
-
继承DataSourceV2和ReadSupport创建XXXDataSource类
重写ReadSupport的createReader方法，返回自定义的DataSourceReader类，如返回自定义的XXXDataSourceReader实例。 -
继承DataSourceReader创建XXXDataSourceReader类
重写DataSourceReader的readSchema方法，返回数据源的schema；重写createDataReaderFactories方法，返回多个自定义的DataReaderFactory实例。 -
继承DataReaderFactory创建DataReader工厂类
如XXXDataReaderFactory，重写createDataReader方法，返回自定义的DataReader实例。 -
继承DataReader类创建自定义的DataReader
如XXXDataReader，重写next()方法，告诉Spark是否有下一条数据，触发get()方法；重写get()方法获取数据；重写close()方法关闭资源。
基于DataSourceV2实现输出源
基于DataSourceV2实现自定义输出源的步骤如下：
-
继承DataSourceV2和WriteSupport创建XXXDataSource
重写createWriter方法，返回自定义的DataSourceWriter。 -
继承DataSourceWriter创建XXXDataSourceWriter类
重写createWriterFactory方法，返回自定义的DataWriterFactory；重写commit方法提交整个事务；重写abort方法进行事务回滚。 -
继承DataWriterFactory创建XXXDataWriterFactory类
重写createWriter方法，返回自定义的DataWriter。 -
继承DataWriter创建XXXDataWriter类
重写write方法，将数据写出；重写commit方法提交事务；重写abort方法进行事务回滚。
SparkSQL自定义RestDataSource代码
scala
package com.hollysys.spark.sql.datasource.rest
import java.math.BigDecimal
import java.util
import java.util.Optional
import com.alibaba.fastjson.{JSONArray, JSONObject, JSONPath}
import org.apache.http.client.fluent.Request
import org.apache.http.entity.ContentType
import org.apache.spark.sql.{Row, SaveMode, SparkSession}
import org.apache.spark.sql.sources.v2.reader.{DataReader, DataReaderFactory, DataSourceReader, SupportsPushDownRequiredColumns}
import org.apache.spark.sql.sources.v2.writer.{DataSourceWriter, DataWriter, DataWriterFactory, WriterCommitMessage}
import org.apache.spark.sql.sources.v2.{DataSourceOptions, DataSourceV2, ReadSupport, WriteSupport}
import org.apache.spark.sql.types.StructType
/**
* 基于Rest的Spark SQL DataSource
*/
class RestDataSource extends DataSourceV2 with ReadSupport with WriteSupport {
override def createReader(options: DataSourceOptions): DataSourceReader =
new RestDataSourceReader(
options.get("url").get(),
options.get("params").get(),
options.get("xPath").get(),
options.get("schema").get()
)
override def createWriter(jobId: String,
schema: StructType,
mode: SaveMode,
options: DataSourceOptions): Optional[DataSourceWriter] = Optional.of(new RestDataSourceWriter)
}
/**
* 创建RestDataSourceReader
*
* @param url REST服务的API
* @param params 请求需要的参数
* @param xPath JSON数据的xPath
* @param schemaString 用户传入的schema字符串
*/
class RestDataSourceReader(url: String, params: String, xPath: String, schemaString: String)
extends DataSourceReader {
// 使用StructType.fromDDL方法将schema字符串转成StructType类型
var requiredSchema: StructType = StructType.fromDDL(schemaString)
/**
* 生成schema
*
* @return schema
*/
override def readSchema(): StructType = requiredSchema
/**
* 创建工厂类
*
* @return List[实例]
*/
override def createDataReaderFactories(): util.List[DataReaderFactory[Row]] = {
import collection.JavaConverters._
Seq(
new RestDataReaderFactory(url, params, xPath).asInstanceOf[DataReaderFactory[Row]]
).asJava
}
}
/**
* RestDataReaderFactory工厂类
*
* @param url REST服务的API
* @param params 请求需要的参数
* @param xPath JSON数据的xPath
*/
class RestDataReaderFactory(url: String, params: String, xPath: String) extends DataReaderFactory[Row] {
override def createDataReader(): DataReader[Row] = new RestDataReader(url, params, xPath)
}
/**
* RestDataReader类
*
* @param url REST服务的API
* @param params 请求需要的参数
* @param xPath JSON数据的xPath
*/
class RestDataReader(url: String, params: String, xPath: String) extends DataReader[Row] {
// 使用Iterator模拟数据
val data: Iterator[Seq[AnyRef]] = getIterator
override def next(): Boolean = {
data.hasNext
}
override def get(): Row = {
val seq = data.next().map {
// 浮点类型会自动转为BigDecimal，导致Spark无法转换
case decimal: BigDecimal =>
decimal.doubleValue()
case x => x
}
Row(seq: _*)
}
override def close(): Unit = {
println("close source")
}
def getIterator: Iterator[Seq[AnyRef]] = {
import scala.collection.JavaConverters._
val res: List[AnyRef] = RestDataSource.requestData(url, params, xPath)
res.map(r => {
r.asInstanceOf[JSONObject].asScala.values.toList
}).toIterator
}
}
/**
* RestDataSourceWriter
*/
class RestDataSourceWriter extends DataSourceWriter {
/**
* 创建RestDataWriter工厂类
*
* @return RestDataWriterFactory
*/
override def createWriterFactory(): DataWriterFactory[Row] = new RestDataWriterFactory
/**
* commit
*
* @param messages 所有分区提交的commit信息
* 触发一次
*/
override def commit(messages: Array[WriterCommitMessage]): Unit = ???
/**
* abort
*
* @param messages 当write异常时调用
*/
override def abort(messages: Array[WriterCommitMessage]): Unit = ???
}
/**
* DataWriterFactory工厂类
*/
class RestDataWriterFactory extends DataWriterFactory[Row] {
/**
* 创建DataWriter
*
* @param partitionId 分区ID
* @param attemptNumber 重试次数
* @return DataWriter
* 每个分区创建一个RestDataWriter实例
*/
override def createDataWriter(partitionId: Int, attemptNumber: Int): DataWriter[Row] = new RestDataWriter(partitionId, attemptNumber)
}
/**
* RestDataWriter
*
* @param partitionId 分区ID
* @param attemptNumber 重试次数
*/
class RestDataWriter(partitionId: Int, attemptNumber: Int) extends DataWriter[Row] {
/**
* write
*
* @param record 单条记录
* 每条记录都会触发该方法
*/
override def write(record: Row): Unit = {
println(record)
}
/**
* commit
*
* @return commit message
* 每个分区触发一次
*/
override def commit(): WriterCommitMessage = {
RestWriterCommitMessage(partitionId, attemptNumber)
}
/**
* 回滚：当write发生异常时触发该方法
*/
override def abort(): Unit = {
println("abort 方法被出发了")
}
}
case class RestWriterCommitMessage(partitionId: Int, attemptNumber: Int) extends WriterCommitMessage
object RestDataSource {
def requestData(url: String, params: String, xPath: String): List[AnyRef] = {
import scala.collection.JavaConverters._
val response = Request.Post(url).bodyString(params, ContentType.APPLICATION_JSON).execute()
JSONPath.read(response.returnContent().asString(), xPath)
.asInstanceOf[JSONArray].asScala.toList
}
}
object RestDataSourceTest {
def main(args: Array[String]): Unit = {
val spark = SparkSession
.builder()
.master("local[2]")
.appName(this.getClass.getSimpleName)
.getOrCreate()
val df = spark.read
.format("com.hollysys.spark.sql.datasource.rest.RestDataSource")
.option("url", "http://model-opcua-hollysysdigital-test.hiacloud.net.cn/aggquery/query/queryPointHistoryData")
.option("params", "{
"startTime": "1543887720000",
"endTime": "1543891320000",
"maxSizePerNode": 1000,
"nodes": [
{
"uri": "/SymLink-10000012030100000-device/5c174da007a54e0001035ddd"
}
]
}")
.option("xPath", "$.result.historyData")
.option("schema", "`s` INT,`t` LONG,`v` DOUBLE")
.load()
df.printSchema()
df.show(false)
}
}SparkSQL自定义HBaseSource
scala
import java.util
import java.util.Optional
import com.travel.utils.HbaseTools
import org.apache.hadoop.hbase.TableName
import org.apache.hadoop.hbase.client._
import org.apache.hadoop.hbase.util.Bytes
import org.apache.spark.sql.sources.v2.reader._
import org.apache.spark.sql.sources.v2.writer.{DataSourceWriter, DataWriter, DataWriterFactory, WriterCommitMessage}
import org.apache.spark.sql.sources.v2.{DataSourceOptions, DataSourceV2, ReadSupport, WriteSupport}
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}
object HBaseSourceAndSink {
def main(args: Array[String]): Unit = {
val spark = SparkSession
.builder()
.master("local[2]")
.getOrCreate()
val df = spark.read
.format("com.travel.programApp.HBaseSource")
.option("hbase.table.name", "spark_hbase_sql")
.option("schema", "`name` STRING,`score` STRING")
.option("cf.cc","cf:name,cf:score")
.load()
df.explain(true)
df.createOrReplaceTempView("sparkHBaseSQL")
df.printSchema()
val frame: DataFrame = spark.sql("select * from sparkHBaseSQL where score > 60")
frame.write.format("com.travel.programApp.HBaseSource")
.mode(SaveMode.Overwrite)
.option("hbase.table.name","spark_hbase_write")
.save()
}
}
class HBaseSource extends DataSourceV2 with ReadSupport with WriteSupport{
override def createReader(options: DataSourceOptions): DataSourceReader = {
new HBaseDataSourceReader(options.get("hbase.table.name").get(),options.get("schema").get(),options.get("cf.cc").get())
}
override def createWriter(jobId: String, schema: StructType, mode: SaveMode, options: DataSourceOptions): Optional[DataSourceWriter] = {
Optional.of(new HBaseDataSourceWrite)
}
}
class HBaseDataSourceWrite extends DataSourceWriter{
override def createWriterFactory(): DataWriterFactory[Row] = {
new HBaseDataWriterFactory
}
override def commit(messages: Array[WriterCommitMessage]): Unit = {}
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
class HBaseDataWriterFactory extends DataWriterFactory[Row]{
override def createDataWriter(partitionId: Int, attemptNumber: Int): DataWriter[Row] = {
new HBaseDataWriter
}
}
class HBaseDataWriter extends DataWriter[Row]{
private val conn: Connection = HbaseTools.getHbaseConn
private val table: Table = conn.getTable(TableName.valueOf("spark_hbase_write"))
override def write(record: Row): Unit = {
val name: String = record.getString(0)
val score: String = record.getString(1)
val put = new Put("0001".getBytes())
put.addColumn("cf".getBytes(),"name".getBytes(),name.getBytes())
put.addColumn("cf".getBytes(),"score".getBytes(),score.getBytes())
table.put(put)
}
override def commit(): WriterCommitMessage = {
table.close()
conn.close()
null
}
override def abort(): Unit = {
null
}
}
class HBaseDataSourceReader(tableName:String,schema:String,cfcc:String) extends DataSourceReader {
//定义HBase的schema
private val structType: StructType = StructType.fromDDL(schema)
override def readSchema(): StructType = {
structType
}
//返回DataReaderFactory
override def createDataReaderFactories(): util.List[DataReaderFactory[Row]] = {
import collection.JavaConverters._
Seq(
new HBaseReaderFactory(tableName,cfcc).asInstanceOf[DataReaderFactory[Row]]
).asJava
}
}
class HBaseReaderFactory(tableName:String,cfcc:String) extends DataReaderFactory[Row] {
override def createDataReader(): DataReader[Row] = {
new HBaseReader(tableName,cfcc)
}
}
class HBaseReader(tableName:String,cfcc:String) extends DataReader[Row] {
private var hbaseConnection:Connection = null
private var resultScanner:ResultScanner = null
private var nextResult:Result = null
// 获取HBase当中的数
val data: Iterator[Seq[AnyRef]] = getIterator
def getIterator: Iterator[Seq[AnyRef]] = {
import scala.collection.JavaConverters._
hbaseConnection = HbaseTools.getHbaseConn
val table: Table = hbaseConnection.getTable(TableName.valueOf(tableName))
resultScanner = table.getScanner(new Scan())
val iterator: Iterator[Seq[AnyRef]] = resultScanner.iterator().asScala.map(eachResult => {
val str: String = Bytes.toString(eachResult.getValue("cf".getBytes(), "name".getBytes()))
val score: String = Bytes.toString(eachResult.getValue("cf".getBytes(), "score".getBytes()))
Seq(str, score)
})
iterator
}
override def next(): Boolean = {
data.hasNext
}
override def get(): Row = {
val seq: Seq[Any] = data.next()
Row.fromSeq(seq)
}
override def close(): Unit = {
hbaseConnection.close()
}
}SparkSQL基于ClickHouse扩展DataSourceV2源
scala
package com.mengyao.spark.datasourcev2.ext.example1
import java.io.Serializable
import java.sql.{Connection, Date, PreparedStatement, ResultSet, SQLException, Statement}
import java.text.SimpleDateFormat
import java.util
import java.util.Optional
import cn.itcast.logistics.etl.Configure
import org.apache.commons.lang3.{StringUtils, SystemUtils}
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.{SaveMode, SparkSession}
import org.apache.spark.sql.sources.{DataSourceRegister, EqualTo, Filter}
import org.apache.spark.sql.sources.v2.reader.{DataSourceReader, InputPartition, InputPartitionReader, SupportsPushDownFilters, SupportsPushDownRequiredColumns}
import org.apache.spark.sql.sources.v2.writer.streaming.StreamWriter
import org.apache.spark.sql.sources.v2.writer.{DataSourceWriter, DataWriter, DataWriterFactory, WriterCommitMessage}
import org.apache.spark.sql.sources.v2.{DataSourceOptions, DataSourceV2, ReadSupport, StreamWriteSupport, WriteSupport}
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.types.{StructType, _}
import org.apache.spark.unsafe.types.UTF8String
import org.javatuples.Triplet
import ru.yandex.clickhouse.domain.ClickHouseDataType
import ru.yandex.clickhouse.response.{ClickHouseResultSet, ClickHouseResultSetMetaData}
import ru.yandex.clickhouse.settings.ClickHouseProperties
import ru.yandex.clickhouse.{ClickHouseConnection, ClickHouseDataSource, ClickHouseStatement}
import scala.collection.mutable.ArrayBuffer
/**
* @ClassName CKTest
* @Description 测试ClickHouse的DataSourceV2实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
object CKTest {
private val APP_NAME: String = CKTest.getClass.getSimpleName
private val master: String = "local[2]"
def main(args: Array[String]) {
if (SystemUtils.IS_OS_WINDOWS) System.setProperty("hadoop.home.dir", Configure.LOCAL_HADOOP_HOME)
val spark = SparkSession.builder()
.master(master)
.appName(APP_NAME).getOrCreate();
val df = spark.read.format(Configure.SPARK_CLICKHOUSE_FORMAT)
.option("driver", Configure.clickhouseDriver)
.option("url", Configure.clickhouseUrl)
.option("user", Configure.clickhouseUser)
.option("password", Configure.clickhousePassword)
.option("table", "tbl_address")
.option("use_server_time_zone", "false")
.option("use_time_zone", "Asia/Shanghai")
.option("max_memory_usage", "2000000000")
.option("max_bytes_before_external_group_by", "1000000000")
.load().coalesce(1)
df.show(1000, false)
import spark.implicits._
df.where($"id"===328).distinct().coalesce(1).write.format(Configure.SPARK_CLICKHOUSE_FORMAT)
.option("driver", Configure.clickhouseDriver)
.option("url", Configure.clickhouseUrl)
.option("user", Configure.clickhouseUser)
.option("password", Configure.clickhousePassword)
.option("table", "tbl_address")
.option("use_server_time_zone", "false")
.option("use_time_zone", "Asia/Shanghai")
.option("max_memory_usage", "2000000000")
.option("max_bytes_before_external_group_by", "1000000000")
.mode(SaveMode.Append)
.save();
}
}
/**
* @ClassName ClickHouseDataSourceV2
* @Description 扩展SparkSQL DataSourceV2的ClickHouse数据源实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class ClickHouseDataSourceV2 extends DataSourceV2 with DataSourceRegister with ReadSupport with WriteSupport with StreamWriteSupport {
/** 声明ClickHouse数据源的简称，使用方式为spark.read.format("clickhouse")... */
override def shortName(): String = "clickhouse"
/** 批处理方式下的数据读取 */
override def createReader(options: DataSourceOptions): DataSourceReader = new CKReader(new CKOptions(options.asMap()))
/** 批处理方式下的数据写入 */
override def createWriter(writeUUID: String, schema: StructType, mode: SaveMode, options: DataSourceOptions): Optional[DataSourceWriter] = Optional.of(new CKWriter(writeUUID, schema, mode, null, new CKOptions(options.asMap())))
/** 流处理方式下的数据写入 */
override def createStreamWriter(queryId: String, schema: StructType, mode: OutputMode, options: DataSourceOptions): StreamWriter = new CKWriter(queryId, schema, null, mode, new CKOptions(options.asMap()))
}
/**
* @ClassName CKReader
* @Description 基于批处理方式的ClickHouse数据读取（此处只使用1个分区实现）
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKReader(options: CKOptions) extends DataSourceReader {
//with SupportsPushDownRequiredColumns with SupportsPushDownFilters {
private val customSchema: java.lang.String = options.getCustomSchema
private val helper = new CKHelper(options)
import collection.JavaConversions._
private val schema = if(StringUtils.isEmpty(customSchema)) {
helper.getSparkTableSchema()
} else {
helper.getSparkTableSchema(new util.LinkedList[String](asJavaCollection(customSchema.split(","))))
}
override def readSchema(): StructType = schema
override def planInputPartitions(): util.List[InputPartition[InternalRow]] = util.Arrays.asList(new CKInputPartition(schema, options))
}
/**
* @ClassName CKInputPartition
* @Description 基于批处理方式的ClickHouse分区实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKInputPartition(schema: StructType, options: CKOptions) extends InputPartition[InternalRow] {
override def createPartitionReader(): InputPartitionReader[InternalRow] = new CKInputPartitionReader(schema, options)
}
/**
* @ClassName CKInputPartitionReader
* @Description 基于批处理方式的ClickHouse分区读取数据实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKInputPartitionReader(schema: StructType, options: CKOptions) extends InputPartitionReader[InternalRow] with Logging with Serializable{
val helper = new CKHelper(options)
var connection: ClickHouseConnection = null
var st: ClickHouseStatement = null
var rs: ResultSet = null
override def next(): Boolean = {
if (null == connection || connection.isClosed && null == st || st.isClosed && null == rs || rs.isClosed){
connection = helper.getConnection
st = connection.createStatement()
rs = st.executeQuery(helper.getSelectStatement(schema))
println(/**logInfo**/s"初始化ClickHouse连接.")
}
if(null != rs && !rs.isClosed) rs.next() else false
}
override def get(): InternalRow = {
val fields = schema.fields
val length = fields.length
val record = new Array[Any](length)
for (i <- 0 until length) {
val field = fields(i)
val name = field.name
val dataType = field.dataType
try {
dataType match {
case DataTypes.BooleanType => record(i) = rs.getBoolean(name)
case DataTypes.DateType => record(i) = DateTimeUtils.fromJavaDate(rs.getDate(name))
case DataTypes.DoubleType => record(i) = rs.getDouble(name)
case DataTypes.FloatType => record(i) = rs.getFloat(name)
case DataTypes.IntegerType => record(i) = rs.getInt(name)
case DataTypes.LongType => record(i) = rs.getLong(name)
case DataTypes.ShortType => record(i) = rs.getShort(name)
case DataTypes.StringType => record(i) = UTF8String.fromString(rs.getString(name))
case DataTypes.TimestampType => record(i) = DateTimeUtils.fromJavaTimestamp(rs.getTimestamp(name))
case DataTypes.BinaryType => record(i) = rs.getBytes(name)
case DataTypes.NullType => record(i) = StringUtils.EMPTY
}
} catch {
case e: SQLException => logError(e.getStackTrace.mkString("", scala.util.Properties.lineSeparator, scala.util.Properties.lineSeparator))
}
}
new GenericInternalRow(record)
}
override def close(): Unit = {helper.closeAll(connection, st, null, rs)}
}
/**
* @ClassName CKWriter
* @Description 支持Batch和Stream的数据写实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKWriter(writeUuidOrQueryId: String, schema: StructType, batchMode: SaveMode, streamMode: OutputMode, options: CKOptions) extends StreamWriter {
private val isStreamMode:Boolean = if (null!=streamMode&&null==batchMode) true else false
override def useCommitCoordinator(): Boolean = true
override def onDataWriterCommit(message: WriterCommitMessage): Unit = {}
override def createWriterFactory(): DataWriterFactory[InternalRow] = new CKDataWriterFactory(writeUuidOrQueryId, schema, batchMode, streamMode, options)
/** Batch writer commit */
override def commit(messages: Array[WriterCommitMessage]): Unit = {}
/** Batch writer abort */
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
/** Streaming writer commit */
override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {}
/** Streaming writer abort */
override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {}
}
/**
* @ClassName CKDataWriterFactory
* @Description 写数据工厂，用来实例化CKDataWriter
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKDataWriterFactory(writeUUID: String, schema: StructType, batchMode: SaveMode, streamMode: OutputMode, options: CKOptions) extends DataWriterFactory[InternalRow] {
override def createDataWriter(partitionId: Int, taskId: Long, epochId: Long): DataWriter[InternalRow] = new CKDataWriter(writeUUID, schema, batchMode, streamMode, options)
}
/**
* @ClassName CKDataWriter
* @Description ClickHouse的数据写实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKDataWriter(writeUUID: String, schema: StructType, batchMode: SaveMode, streamMode: OutputMode, options: CKOptions) extends DataWriter[InternalRow] with Logging with Serializable {
val helper = new CKHelper(options)
val opType = options.getOpTypeField
private val sqls = ArrayBuffer[String]()
private val autoCreateTable: Boolean = options.autoCreateTable
private val init = if (autoCreateTable) {
val createSQL = helper.createTable(options.getFullTable, schema)
println(/**logInfo**/s"==== 初始化表SQL：$createSQL")
helper.executeUpdate(createSQL)
}
val fields = schema.fields
override def commit(): WriterCommitMessage = {
helper.executeUpdateBatch(sqls)
val batchSQL = sqls.mkString("
")
// logDebug(batchSQL)
println(batchSQL)
new WriterCommitMessage{override def toString: String = s"批量插入SQL: $batchSQL"}
}
override def write(record: InternalRow): Unit = {
if(StringUtils.isEmpty(opType)) {
throw new RuntimeException("未传入opTypeField字段名称，无法确定数据持久化类型！")
}
var sqlStr: String = helper.getStatement(options.getFullTable, schema, record)
logDebug(s"==== 拼接完成的INSERT SQL语句为：$sqlStr")
try {
if (StringUtils.isEmpty(sqlStr)) {
val msg = "==== 拼接INSERT SQL语句失败，因为该语句为NULL或EMPTY！"
logError(msg)
throw new RuntimeException(msg)
}
Thread.sleep(options.getInterval())
// 在流处理模式下操作
if (null == batchMode) {
if (streamMode == OutputMode.Append) {
sqls += sqlStr
// val state = helper.executeUpdate(sqlStr)
// println(s"==== 在OutputMode.Append模式下执行：$sqlStr
状态：$state")
}
else if(streamMode == OutputMode.Complete) {logError("==== 未实现OutputMode.Complete模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
else if(streamMode == OutputMode.Update) {logError("==== 未实现OutputMode.Update模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
else {logError(s"==== 未知模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
// 在批处理模式下操作
} else {
if (batchMode == SaveMode.Append) {
sqls += sqlStr
//val state = helper.executeUpdate(sqlStr)
//println(s"==== 在SaveMode.Append模式下执行：$sqlStr
状态：$state")
}
else if(batchMode == SaveMode.Overwrite) {logError("==== 未实现SaveMode.Overwrite模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
else if(batchMode == SaveMode.ErrorIfExists) {logError("==== 未实现SaveMode.ErrorIfExists模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
else if(batchMode == SaveMode.Ignore) {logError("==== 未实现SaveMode.Ignore模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
else {logError(s"==== 未知模式下的写入操作，请在CKDataWriter.write方法中添加相关实现！")}
}
} catch {
case e: Exception => logError(e.getMessage)
}
}
override def abort(): Unit = {}
}
/**
* @ClassName CKOptions
* @Description 从SparkSQL中DataSourceOptions中提取适用于ClickHouse的参数（spark.[read/write].options参数）
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKOptions(var originalMap: util.Map[String, String]) extends Logging with Serializable {
val DRIVER_KEY: String = "driver"
val URL_KEY: String = "url"
val USER_KEY: String = "user"
val PASSWORD_KEY: String = "password"
val DATABASE_KEY: String = "database"
val TABLE_KEY: String = "table"
val AUTO_CREATE_TABLE = "autoCreateTable".toLowerCase
val PATH_KEY = "path"
val INTERVAL = "interval"
val CUSTOM_SCHEMA_KEY: String = "customSchema".toLowerCase
val WHERE_KEY: String = "where"
val OP_TYPE_FIELD = "opTypeField".toLowerCase
val PRIMARY_KEY = "primaryKey".toLowerCase
def getValue[T](key: String, `type`: T): T = (if (originalMap.containsKey(key)) originalMap.get(key) else null).asInstanceOf[T]
def getDriver: String = getValue(DRIVER_KEY, new String)
def getURL: String = getValue(URL_KEY, new String)
def getUser: String = getValue(USER_KEY, new String)
def getPassword: String = getValue(PASSWORD_KEY, new String)
def getDatabase: String = getValue(DATABASE_KEY, new String)
def getTable: String = getValue(TABLE_KEY, new String)
def autoCreateTable: Boolean = {
originalMap.getOrDefault(AUTO_CREATE_TABLE, "false").toLowerCase match {
case "true" => true
case "false" => false
case _ => false
}
}
def getInterval(): Long = {originalMap.getOrDefault(INTERVAL, "200").toLong}
def getPath: String = if(StringUtils.isEmpty(getValue(PATH_KEY, new String))) getTable else getValue(PATH_KEY, new String)
def getWhere: String = getValue(WHERE_KEY, new String)
def getCustomSchema: String = getValue(CUSTOM_SCHEMA_KEY, new String)
def getOpTypeField: String = getValue(OP_TYPE_FIELD, new String)
def getPrimaryKey: String = getValue(PRIMARY_KEY, new String)
def getFullTable: String = {
val database = getDatabase
val table = getTable
if (StringUtils.isEmpty(database) && !StringUtils.isEmpty(table)) table else if (!StringUtils.isEmpty(database) && !StringUtils.isEmpty(table)) database+"."+table else table
}
def asMap(): util.Map[String, String] = this.originalMap
override def toString: String = originalMap.toString
}
/**
* @ClassName CKHelper
* @Description ClickHouse的JDBCHelper实现
* @Created by MengYao
* @Date 2020/5/17 16:34
* @Version V1.0
*/
class CKHelper(options: CKOptions) extends Logging with Serializable {
private val opType: String = options.getOpTypeField
private val id: String = options.getPrimaryKey
private var connection: ClickHouseConnection = getConnection
def getConnection: ClickHouseConnection = {
val url = options.getURL
val ds = new ClickHouseDataSource(url, new ClickHouseProperties())
ds.getConnection(options.getUser, options.getPassword)
}
def createTable(table: String, schema: StructType): String = {
val cols = ArrayBuffer[String]()
for (field <- schema.fields) {
val dataType = field.dataType
val ckColName = field.name
if (ckColName!=opType) {
var ckColType = getClickhouseSqlType(dataType)
if (!StringUtils.isEmpty(ckColType)) {
if (ckColType.toLowerCase=="string") {ckColType="Nullable("+ckColType+")"}
}
cols += ckColName+" "+ ckColType
}
}
s"CREATE TABLE IF NOT EXISTS $table(${cols.mkString(",")},sign Int8,version UInt64) ENGINE=VersionedCollapsingMergeTree(sign, version) ORDER BY $id"
}
def getSparkTableSchema(customFields: util.LinkedList[String] = null): StructType = {
import collection.JavaConversions._
val list: util.LinkedList[Triplet[String, String, String]] = getCKTableSchema(customFields)
var fields = ArrayBuffer[StructField]()
for(trp <- list) {
fields += StructField(trp.getValue0, getSparkSqlType(trp.getValue1))
}
StructType(fields)
}
private def getFieldValue(fieldName: String, schema: StructType, data:InternalRow): Any = {
var flag = true
var fieldValue:String = null
val fields = schema.fields
for(i <- 0 until fields.length if flag) {
val field = fields(i)
if(fieldName==field.name) {
fieldValue = field.dataType match {
case DataTypes.BooleanType => if (data.isNullAt(i)) "NULL" else s"${data.getBoolean(i)}"
case DataTypes.DoubleType => if (data.isNullAt(i)) "NULL" else s"${data.getDouble(i)}"
case DataTypes.FloatType => if (data.isNullAt(i)) "NULL" else s"${data.getFloat(i)}"
case DataTypes.IntegerType => if (data.isNullAt(i)) "NULL" else s"${data.getInt(i)}"
case DataTypes.LongType => if (data.isNullAt(i)) "NULL" else s"${data.getLong(i)}"
case DataTypes.ShortType => if (data
end