StructuredStreaming自定义DataSourceV2源
ckckck2025年1月10日2 浏览自定义输入源
在某些应用场景下，我们可能需要自定义数据源。例如，在获取KafkaSource的同时，动态从缓存或HTTP请求中加载业务数据，或其他数据源等。自定义输入源需要以下步骤：
- 继承DataSourceRegister和StreamSourceProvider创建自定义Provider类
- 重写DataSourceRegister类中的shortName和StreamSourceProvider中的createSource以及sourceSchema方法
- 继承Source创建自定义Source类
- 重写Source中的schema方法指定输入源的schema
- 重写Source中的getOffset方法监听流数据
- 重写Source中的getBatch方法获取数据
- 重写Source中的stop方法用来关闭资源
创建CustomDataSourceProvider类
- 继承
DataSourceRegister和StreamSourceProvider - 重写
DataSourceRegister的shortName方法 - 重写
StreamSourceProvider中的sourceSchema方法 - 重写
StreamSourceProvider中的createSource方法
scala
import org.apache.spark.internal.Logging
import org.apache.spark.sql.SQLContext
import org.apache.spark.sql.execution.streaming.{Sink, Source}
import org.apache.spark.sql.sources.{DataSourceRegister, StreamSinkProvider, StreamSourceProvider}
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.types.StructType
/**
* 自定义Structured Streaming数据源
*
* （1）继承DataSourceRegister类
* 需要重写shortName方法，用来向Spark注册该组件
* （2）继承StreamSourceProvider类
* 需要重写createSource以及sourceSchema方法，用来创建数据输入源
* （3）继承StreamSinkProvider类
* 需要重写createSink方法，用来创建数据输出源
*/
class CustomDataSourceProvider extends DataSourceRegister
with StreamSourceProvider
with StreamSinkProvider
with Logging {
/**
* 数据源的描述名字，如：kafka、socket
* @return 字符串shortName
* 该方法用来指定一个数据源的名字，用来向Spark注册该数据源。如Spark内置的数据源的shortName:kafka、socket、rate等，该方法返回一个字符串
*/
override def shortName(): String = "custom"
/**
* 定义数据源的Schema
*
* @param sqlContext Spark SQL 上下文
* @param schema 通过.schema()方法传入的schema
* @param providerName Provider的名称，包名+类名
* @param parameters 通过.option()方法传入的参数
* @return 元组，(shortName, schema)
*/
override def sourceSchema(sqlContext: SQLContext,
schema: Option[StructType],
providerName: String,
parameters: Map[String, String]): (String, StructType) = (shortName(), schema.get)
/**
* 创建输入源
*
* @param sqlContext Spark SQL 上下文
* @param metadataPath 元数据Path
* @param schema 通过.schema()方法传入的schema
* @param providerName Provider的名称，包名+类名
* @param parameters 通过.option()方法传入的参数
* @return 自定义source，需要继承Source接口实现
*/
override def createSource(sqlContext: SQLContext,
metadataPath: String,
schema: Option[StructType],
providerName: String,
parameters: Map[String, String]): Source = new CustomDataSource(sqlContext, parameters, schema)
/**
* 创建输出源
*
* @param sqlContext Spark SQL 上下文
* @param parameters 通过.option()方法传入的参数
* @param partitionColumns 分区列名
* @param outputMode 输出模式
* @return
*/
override def createSink(sqlContext: SQLContext,
parameters: Map[String, String],
partitionColumns: Seq[String],
outputMode: OutputMode): Sink = new CustomDataSink(sqlContext, parameters, outputMode)
}创建CustomDataSource类
- 继承
Source创建CustomDataSource类 - 重写
Source的schema方法 - 重写
Source的getOffset方法 - 重写
Source的getBatch方法 - 重写
Source的stop方法
scala
package org.apache.spark.sql.structured.datasource.custom
import org.apache.spark.internal.Logging
import org.apache.spark.sql.execution.streaming.{Offset, Source}
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.{DataFrame, SQLContext}
/**
* 自定义数据输入源：需要继承Source接口
* 实现思路：
* （1）通过重写schema方法来指定数据输入源的schema，这个schema需要与Provider中指定的schema保持一致
* （2）通过重写getOffset方法来获取数据的偏移量，这个方法会一直被轮询调用，不断的获取偏移量
* （3）通过重写getBatch方法，来获取数据，这个方法是在偏移量发生改变后被触发
* （4）通过stop方法，来进行一下关闭资源的操作
*/
class CustomDataSource(sqlContext: SQLContext,
parameters: Map[String, String],
schemaOption: Option[StructType]) extends Source
with Logging {
/**
* 指定数据源的schema，需要与Provider中sourceSchema中指定的schema保持一致，否则会报异常
* 触发机制：当创建数据源的时候被触发执行
*
* @return schema
*/
override def schema: StructType = schemaOption.get
/**
* 获取offset，用来监控数据的变化情况
* 触发机制：不断轮询调用
* 实现要点：
* （1）Offset的实现：
* 由函数返回值可以看出，我们需要提供一个标准的返回值Option[Offset]
* 我们可以通过继承 org.apache.spark.sql.sources.v2.reader.streaming.Offset实现，这里面其实就是保存了个json字符串
*
* （2）JSON转化
* 因为Offset里实现的是一个json字符串，所以我们需要将我们存放offset的集合或者case class转化成json字符串
* spark里是通过org.json4s.jackson这个包来实现case class 集合类（Map、List、Seq、Set等）与json字符串的相互转化
*
* @return Offset
*/
override def getOffset: Option[Offset] = ???
/**
* 获取数据
*
* @param start 上一个批次的end offset
* @param end 通过getOffset获取的新的offset
* 触发机制：当不断轮询的getOffset方法，获取的offset发生改变时，会触发该方法
*
* 实现要点：
* （1）DataFrame的创建：
* 可以通过生成RDD，然后使用RDD创建DataFrame
* RDD创建：sqlContext.sparkContext.parallelize(rows.toSeq)
* DataFrame创建：sqlContext.internalCreateDataFrame(rdd, schema, isStreaming = true)
* @return DataFrame
*/
override def getBatch(start: Option[Offset], end: Offset): DataFrame = ???
/**
* 关闭资源
* 将一些需要关闭的资源放到这里来关闭，如MySQL的数据库连接等
*/
override def stop(): Unit = ???
}自定义输出源
相比较输入源的自定义性，输出源自定义的应用场景更为常用。例如：数据写入关系型数据库、数据写入HBase、数据写入Redis等。Structured Streaming提供的foreach以及2.4版本的foreachBatch方法已经可以实现绝大多数应用场景，几乎可以将数据写入任何地方。但为了更优雅的实现，我们可以参考Spark SQL Sink规范，通过自定义Sink的方式来实现。实现自定义Sink需要以下四个步骤：
- 继承DataSourceRegister和StreamSinkProvider创建自定义SinkProvider类
- 重写DataSourceRegister类中的shortName和StreamSinkProvider中的createSink方法
- 继承Sink创建自定义Sink类
- 重写Sink中的addBatch方法
创建CustomDataSink类
scala
import org.apache.spark.internal.Logging
import org.apache.spark.sql.execution.streaming.Sink
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.{DataFrame, SQLContext}
/**
* 自定义数据输出源
*/
class CustomDataSink(sqlContext: SQLContext,
parameters: Map[String, String],
outputMode: OutputMode) extends Sink with Logging {
/**
* 添加Batch，即数据写出
*
* @param batchId batchId
* @param data DataFrame
* 触发机制：当发生计算时，会触发该方法，并且得到要输出的DataFrame
* 实现摘要：
* 1. 数据写入方式：
* （1）通过SparkSQL内置的数据源写出
* 我们拿到DataFrame之后可以通过SparkSQL内置的数据源将数据写出，如：
* JSON数据源、CSV数据源、Text数据源、Parquet数据源、JDBC数据源等。
* （2）通过自定义SparkSQL的数据源进行写出
* （3）通过foreachPartition 将数据写出
*/
override def addBatch(batchId: Long, data: DataFrame): Unit = ???
}实现MySQL/PG自定义输入源输出源完整Demo
1. 创建MySQLSourceProvider类
scala
package org.apache.spark.sql.structured.datasource
import org.apache.spark.internal.Logging
import org.apache.spark.sql.SQLContext
import org.apache.spark.sql.execution.streaming.{Sink, Source}
import org.apache.spark.sql.sources.{DataSourceRegister, StreamSinkProvider, StreamSourceProvider}
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.types.StructType
class MySQLSourceProvider extends DataSourceRegister
with StreamSourceProvider
with StreamSinkProvider
with Logging {
/**
* 数据源的描述名字，如：kafka、socket
*
* @return 字符串shortName
*/
override def shortName(): String = "mysql"
/**
* 定义数据源的Schema
*
* @param sqlContext Spark SQL 上下文
* @param schema 通过.schema()方法传入的schema
* @param providerName Provider的名称，包名+类名
* @param parameters 通过.option()方法传入的参数
* @return 元组，(shortName, schema)
*/
override def sourceSchema(
sqlContext: SQLContext,
schema: Option[StructType],
providerName: String,
parameters: Map[String, String]): (String, StructType) = {
(providerName, schema.get)
}
/**
* 创建输入源
*
* @param sqlContext Spark SQL 上下文
* @param metadataPath 元数据Path
* @param schema 通过.schema()方法传入的schema
* @param providerName Provider的名称，包名+类名
* @param parameters 通过.option()方法传入的参数
* @return 自定义source，需要继承Source接口实现
*/
override def createSource(
sqlContext: SQLContext,
metadataPath: String, schema: Option[StructType],
providerName: String, parameters: Map[String, String]): Source = new MySQLSource(sqlContext, parameters, schema)
/**
* 创建输出源
*
* @param sqlContext Spark SQL 上下文
* @param parameters 通过.option()方法传入的参数
* @param partitionColumns 分区列名
* @param outputMode 输出模式
* @return
*/
override def createSink(
sqlContext: SQLContext,
parameters: Map[String, String],
partitionColumns: Seq[String], outputMode: OutputMode): Sink = new MySQLSink(sqlContext: SQLContext, parameters, outputMode)
}2. 创建MySQLSource类
scala
package org.apache.spark.sql.structured.datasource
import java.sql.{Connection, ResultSet}
import org.apache.spark.executor.InputMetrics
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.streaming.{Offset, Source}
import org.apache.spark.sql.types.{DoubleType, IntegerType, LongType, StringType, StructField, StructType, TimestampType}
import org.apache.spark.sql.{DataFrame, SQLContext}
import org.apache.spark.unsafe.types.UTF8String
import org.json4s.jackson.Serialization
import org.json4s.{Formats, NoTypeHints}
class MySQLSource(sqlContext: SQLContext,
options: Map[String, String],
schemaOption: Option[StructType]) extends Source with Logging {
lazy val conn: Connection = C3p0Utils.getDataSource(options).getConnection
val tableName: String = options("tableName")
var currentOffset: Map[String, Long] = Map[String, Long](tableName -> 0)
val maxOffsetPerBatch: Option[Long] = Option(100)
val inputMetrics = new InputMetrics()
override def schema: StructType = schemaOption.get
/**
* 获取Offset
* 这里监控MySQL数据库表中条数变化情况
* @return Option[Offset]
*/
override def getOffset: Option[Offset] = {
val latest = getLatestOffset //获取表中数据条数
val offsets = maxOffsetPerBatch match {
case None => MySQLSourceOffset(latest)
case Some(limit) =>
MySQLSourceOffset(rateLimit(limit, currentOffset, latest))
}
Option(offsets)
}
/**
* 获取数据
* @param start 上一次的offset
* @param end 最新的offset
* @return df
*/
override def getBatch(start: Option[Offset], end: Offset): DataFrame = {
var offset: Long = 0
if (start.isDefined) {
offset = offset2Map(start.get)(tableName)
}
val limit = offset2Map(end)(tableName) - offset
val sql = s"SELECT * FROM $tableName limit $limit offset $offset"
val st = conn.prepareStatement(sql)
val rs = st.executeQuery()
val rows = getInternalRow1(rs)
val rdd = sqlContext.sparkContext.parallelize(rows.toSeq)
currentOffset = offset2Map(end)
sqlContext.internalCreateDataFrame(rdd, schema, isStreaming = true)
}
def getInternalRow1(rs: ResultSet): Iterator[InternalRow] = {
//获取StructType
val schema: StructType = schemaOption.get
val fields: Array[StructField] = schema.fields
var rows: List[InternalRow] = List[InternalRow]()
while (rs.next()) {
var s = Seq[Any]()
//遍历字段 根据字段类型 去获取对应的数据
//注意类型 基本数值类型要转成包装器类型 string转成utf8
for (field <- fields) {
val value = field.dataType match {
case IntegerType => rs.getInt(field.name)
case StringType => UTF8String.fromString(rs.getString(field.name))
case LongType => rs.getLong(field.name)
case TimestampType => {
val t: java.lang.Long = rs.getTimestamp(field.name) match {
case null => null
//时间戳在InternalRow里面必须是微秒单位的
case value => value.getTime * 1000
}
t
}
case DoubleType => rs.getDouble(field.name)
}
s = s.:+(value)
}
//将Seq转成InternalRow
val internalRow = InternalRow.fromSeq(s)
rows = rows.:+(internalRow)
}
rows.toIterator
}
override def stop(): Unit = {
conn.close()
}
def rateLimit(limit: Long, currentOffset: Map[String, Long], latestOffset: Map[String, Long]): Map[String, Long] = {
val co = currentOffset(tableName)
val lo = latestOffset(tableName)
if (co + limit > lo) {
Map[String, Long](tableName -> lo)
} else {
Map[String, Long](tableName -> (co + limit))
}
}
// 获取最新条数
def getLatestOffset: Map[String, Long] = {
var offset: Long = 0
val sql = s"SELECT COUNT(1) FROM $tableName"
val st = conn.prepareStatement(sql)
val rs = st.executeQuery()
while (rs.next()) {
offset = rs.getLong(1)
}
Map[String, Long](tableName -> offset)
}
def offset2Map(offset: Offset): Map[String, Long] = {
implicit val formats: AnyRef with Formats = Serialization.formats(NoTypeHints)
Serialization.read[Map[String, Long]](offset.json())
}
}
case class MySQLSourceOffset(offset: Map[String, Long]) extends Offset {
implicit val formats: AnyRef with Formats = Serialization.formats(NoTypeHints)
override def json(): String = Serialization.write(offset)
}3. 创建MySQLSink类
scala
package org.apache.spark.sql.structured.datasource
import org.apache.spark.internal.Logging
import org.apache.spark.sql.execution.streaming.Sink
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
class MySQLSink(sqlContext: SQLContext, parameters: Map[String, String], outputMode: OutputMode) extends Sink with Logging {
override def addBatch(batchId: Long, data: DataFrame): Unit = {
val query = data.queryExecution
val rdd = query.toRdd
val df = sqlContext.internalCreateDataFrame(rdd, data.schema)
df.show(false)
df.write.format("jdbc").options(parameters).mode(SaveMode.Append).save()
}
}4. 运行主类
scala
package org.apache.spark.sql.structured.datasource.onlineIndex
import java.sql.{Connection, DriverManager, PreparedStatement}
import org.apache.spark.sql.types._
import org.apache.spark.sql.{DataFrame, ForeachWriter, SparkSession}
object TotalNumOfIncoming {
def main(args: Array[String]): Unit = {
val spark = SparkSession
.builder()
.appName(this.getClass.getSimpleName)
.master("local[*]")
.getOrCreate()
val schema = StructType(List(
StructField("tiff_id", IntegerType),
StructField("order_id", IntegerType),
StructField("user_name", StringType),
StructField("tiff", StringType),
StructField("tar_gz", StringType),
StructField("directory", StringType),
StructField("tiff_size", LongType),
StructField("xml", StringType),
StructField("rpb", StringType),
StructField("product_level", StringType),
StructField("satellite", StringType),
StructField("sensor", StringType),
StructField("receive_time", TimestampType),
StructField("upload_time", TimestampType),
StructField("bands", StringType), //将ArrayType(StringType) => StringType 后面读取数据时如果是Array会报pg序列化异常
StructField("top_left_longitude", DoubleType),
StructField("top_left_latitude", DoubleType),
StructField("bottom_right_longitude", DoubleType),
StructField("bottom_right_latitude", DoubleType),
StructField("geom", StringType),
StructField("pixel_spacing", IntegerType),
StructField("adcode", StringType),
StructField("cloud_cover", DoubleType),
StructField("province", StringType),
StructField("city", StringType),
StructField("district", StringType),
StructField("status", IntegerType),
StructField("downNum", IntegerType)
))
val options = Map[String, String](
"driverClass" -> "org.postgresql.Driver",
"jdbcUrl" -> "jdbc:postgresql://192.168.0.214:25433/test1",
"user" -> "postgres",
"password" -> "bjsh",
"tableName" -> "tiff_entry_test")
val source: DataFrame = spark
.readStream
.format("org.apache.spark.sql.structured.datasource.MySQLSourceProvider")
.options(options)
.schema(schema)
.load()
import org.apache.spark.sql.functions._
source.createOrReplaceTempView("tiff_entry_test")
//总入库
val res1: DataFrame = spark.sql("select count(1) ct from tiff_entry_test")
res1.toJSON.withColumn("index", lit("总入库数"))
.select("index", "ct")
.writeStream
.outputMode("update")
.format("org.apache.spark.sql.structured.datasource.MySQLSourceProvider")
.option("truncate", "true")
.option("checkpointLocation", "./tmp/MySQLSourceProvider11")
.option("user", "postgres")
.option("password", "bjsh")
.option("dbtable", "result")
.option("url", "jdbc:postgresql://192.168.0.214:25433/test1")
.start()
.awaitTermination()
}
}end