Github user jose-torres commented on a diff in the pull request:
https://github.com/apache/spark/pull/20382#discussion_r164943885
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/socket.scala
---
@@ -47,130 +48,141 @@ object TextSocketSource {
* This source will *not* work in production applications due to multiple
reasons, including no
* support for fault recovery and keeping all of the text read in memory
forever.
*/
-class TextSocketSource(host: String, port: Int, includeTimestamp: Boolean,
sqlContext: SQLContext)
- extends Source with Logging {
-
- @GuardedBy("this")
- private var socket: Socket = null
-
- @GuardedBy("this")
- private var readThread: Thread = null
-
- /**
- * All batches from `lastCommittedOffset + 1` to `currentOffset`,
inclusive.
- * Stored in a ListBuffer to facilitate removing committed batches.
- */
- @GuardedBy("this")
- protected val batches = new ListBuffer[(String, Timestamp)]
-
- @GuardedBy("this")
- protected var currentOffset: LongOffset = new LongOffset(-1)
-
- @GuardedBy("this")
- protected var lastOffsetCommitted : LongOffset = new LongOffset(-1)
+class TextSocketSource(
+ protected val host: String,
+ protected val port: Int,
+ includeTimestamp: Boolean,
+ sqlContext: SQLContext)
+ extends Source with TextSocketReader with Logging {
initialize()
- private def initialize(): Unit = synchronized {
- socket = new Socket(host, port)
- val reader = new BufferedReader(new
InputStreamReader(socket.getInputStream))
- readThread = new Thread(s"TextSocketSource($host, $port)") {
- setDaemon(true)
-
- override def run(): Unit = {
- try {
- while (true) {
- val line = reader.readLine()
- if (line == null) {
- // End of file reached
- logWarning(s"Stream closed by $host:$port")
- return
- }
- TextSocketSource.this.synchronized {
- val newData = (line,
- Timestamp.valueOf(
-
TextSocketSource.DATE_FORMAT.format(Calendar.getInstance().getTime()))
- )
- currentOffset = currentOffset + 1
- batches.append(newData)
- }
- }
- } catch {
- case e: IOException =>
- }
- }
- }
- readThread.start()
- }
-
/** Returns the schema of the data from this source */
- override def schema: StructType = if (includeTimestamp)
TextSocketSource.SCHEMA_TIMESTAMP
- else TextSocketSource.SCHEMA_REGULAR
-
- override def getOffset: Option[Offset] = synchronized {
- if (currentOffset.offset == -1) {
- None
- } else {
- Some(currentOffset)
- }
- }
+ override def schema: StructType =
+ if (includeTimestamp) TextSocketSource.SCHEMA_TIMESTAMP else
TextSocketSource.SCHEMA_REGULAR
+
+ override def getOffset: Option[Offset] =
getOffsetInternal.map(LongOffset(_))
/** Returns the data that is between the offsets (`start`, `end`]. */
- override def getBatch(start: Option[Offset], end: Offset): DataFrame =
synchronized {
- val startOrdinal =
-
start.flatMap(LongOffset.convert).getOrElse(LongOffset(-1)).offset.toInt + 1
- val endOrdinal =
LongOffset.convert(end).getOrElse(LongOffset(-1)).offset.toInt + 1
-
- // Internal buffer only holds the batches after lastOffsetCommitted
- val rawList = synchronized {
- val sliceStart = startOrdinal - lastOffsetCommitted.offset.toInt - 1
- val sliceEnd = endOrdinal - lastOffsetCommitted.offset.toInt - 1
- batches.slice(sliceStart, sliceEnd)
- }
+ override def getBatch(start: Option[Offset], end: Offset): DataFrame = {
+ val rawList =
getBatchInternal(start.flatMap(LongOffset.convert).map(_.offset),
+ LongOffset.convert(end).map(_.offset))
val rdd = sqlContext.sparkContext
.parallelize(rawList)
.map { case (v, ts) => InternalRow(UTF8String.fromString(v),
ts.getTime) }
sqlContext.internalCreateDataFrame(rdd, schema, isStreaming = true)
}
- override def commit(end: Offset): Unit = synchronized {
+ override def commit(end: Offset): Unit = {
val newOffset = LongOffset.convert(end).getOrElse(
sys.error(s"TextSocketStream.commit() received an offset ($end) that
did not " +
s"originate with an instance of this class")
)
- val offsetDiff = (newOffset.offset - lastOffsetCommitted.offset).toInt
+ commitInternal(newOffset.offset)
+ }
- if (offsetDiff < 0) {
- sys.error(s"Offsets committed out of order: $lastOffsetCommitted
followed by $end")
- }
+ override def toString: String = s"TextSocketSource[host: $host, port:
$port]"
+}
+
+case class TextSocketOffset(offset: Long) extends V2Offset {
+ override def json(): String = offset.toString
+}
+
+class TextSocketMicroBatchReader(options: DataSourceV2Options)
+ extends MicroBatchReader with TextSocketReader with Logging {
+
+ private var startOffset: TextSocketOffset = _
+ private var endOffset: TextSocketOffset = _
+
+ protected val host: String = options.get("host").get()
+ protected val port: Int = options.get("port").get().toInt
- batches.trimStart(offsetDiff)
- lastOffsetCommitted = newOffset
+ initialize()
+
+ override def setOffsetRange(start: Optional[V2Offset], end:
Optional[V2Offset]): Unit = {
+ startOffset =
start.orElse(TextSocketOffset(-1L)).asInstanceOf[TextSocketOffset]
+ endOffset = end.orElse(
+
TextSocketOffset(getOffsetInternal.getOrElse(-1L))).asInstanceOf[TextSocketOffset]
}
- /** Stop this source. */
- override def stop(): Unit = synchronized {
- if (socket != null) {
- try {
- // Unfortunately, BufferedReader.readLine() cannot be interrupted,
so the only way to
- // stop the readThread is to close the socket.
- socket.close()
- } catch {
- case e: IOException =>
- }
- socket = null
+ override def getStartOffset(): V2Offset = {
+ Option(startOffset).getOrElse(throw new IllegalStateException("start
offset not set"))
+ }
+
+ override def getEndOffset(): V2Offset = {
+ Option(endOffset).getOrElse(throw new IllegalStateException("end
offset not set"))
+ }
+
+ override def deserializeOffset(json: String): V2Offset = {
+ TextSocketOffset(json.toLong)
+ }
+
+ override def readSchema(): StructType = {
+ val includeTimestamp = options.getBoolean("includeTimestamp", false)
+ if (includeTimestamp) TextSocketSource.SCHEMA_TIMESTAMP else
TextSocketSource.SCHEMA_REGULAR
+ }
+
+ override def createReadTasks(): JList[ReadTask[Row]] = {
+ val rawList = getBatchInternal(Option(startOffset.offset),
Option(endOffset.offset))
+
+ assert(SparkSession.getActiveSession.isDefined)
+ val spark = SparkSession.getActiveSession.get
+ val numPartitions = spark.sparkContext.defaultParallelism
+
+ val slices = Array.fill(numPartitions)(new ListBuffer[(String,
Timestamp)])
+ rawList.zipWithIndex.foreach { case (r, idx) =>
+ slices(idx % numPartitions).append(r)
}
+
+ (0 until numPartitions).map { i =>
+ val slice = slices(i)
+ new ReadTask[Row] {
+ override def createDataReader(): DataReader[Row] = new
DataReader[Row] {
+ private var currentIdx = -1
+
+ override def next(): Boolean = {
+ currentIdx += 1
+ currentIdx < slice.size
+ }
+
+ override def get(): Row = {
+ Row(slice(currentIdx)._1, slice(currentIdx)._2)
+ }
+
+ override def close(): Unit = {}
+ }
+ }
+ }.toList.asJava
}
- override def toString: String = s"TextSocketSource[host: $host, port:
$port]"
+ override def commit(end: V2Offset): Unit = {
+ val newOffset = end.asInstanceOf[TextSocketOffset]
+ commitInternal(newOffset.offset)
+ }
+
+ override def toString: String = s"TextSocketMicroBatchReader[host:
$host, port: $port]"
}
-class TextSocketSourceProvider extends StreamSourceProvider with
DataSourceRegister with Logging {
- private def parseIncludeTimestamp(params: Map[String, String]): Boolean
= {
- Try(params.getOrElse("includeTimestamp", "false").toBoolean) match {
- case Success(bool) => bool
+class TextSocketSourceProvider extends DataSourceV2
+ with MicroBatchReadSupport with StreamSourceProvider with
DataSourceRegister with Logging {
--- End diff --
TD and I discussed this offline. It should be fine to remove the V1
StreamSourceProvider implementation, because:
* this isn't a production-quality source, so users shouldn't need to fall
back to it
* this source won't be particularly useful at exercising the V1 execution
pipeline once we transition all sources to V2
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
