chaoqin-li1123 commented on code in PR #45977:
URL: https://github.com/apache/spark/pull/45977#discussion_r1563131511
##########
python/pyspark/sql/datasource.py:
##########
@@ -469,6 +501,192 @@ def stop(self) -> None:
...
+class SimpleInputPartition(InputPartition):
+ def __init__(self, start: dict, end: dict):
+ self.start = start
+ self.end = end
+
+
+class SimpleDataSourceStreamReader(ABC):
+ """
+ A base class for simplified streaming data source readers.
+ Compared to :class:`DataSourceStreamReader`,
:class:`SimpleDataSourceStreamReader` doesn't
+ require planning data partition. Also, the read api of
:class:`SimpleDataSourceStreamReader`
+ allows reading data and planning the latest offset at the same time.
+
+ .. versionadded: 4.0.0
+ """
+
+ def initialOffset(self) -> dict:
+ """
+ Return the initial offset of the streaming data source.
+ A new streaming query starts reading data from the initial offset.
+ If Spark is restarting an existing query, it will restart from the
check-pointed offset
+ rather than the initial one.
+
+ Returns
+ -------
+ dict
+ A dict or recursive dict whose key and value are primitive types,
which includes
+ Integer, String and Boolean.
+
+ Examples
+ --------
+ >>> def initialOffset(self):
+ ... return {"parititon-1": {"index": 3, "closed": True},
"partition-2": {"index": 5}}
+ """
+ raise PySparkNotImplementedError(
+ error_class="NOT_IMPLEMENTED",
+ message_parameters={"feature": "initialOffset"},
+ )
+
+ def read(self, start: dict) -> Tuple[Iterator[Tuple], dict]:
+ """
+ Read all available data from start offset and return the offset that
next read attempt
+ starts from.
+
+ Parameters
+ ----------
+ start : dict
+ The start offset to start reading from.
+
+ Returns
+ -------
+ A :class:`Tuple` of an iterator of :class:`Tuple` and a dict\\s
+ The iterator contains all the available records after start offset.
+ The dict is the end offset of this read attempt and the start of
next read attempt.
+ """
+ raise PySparkNotImplementedError(
+ error_class="NOT_IMPLEMENTED",
+ message_parameters={"feature": "read"},
+ )
+
+ def readBetweenOffsets(self, start: dict, end: dict) -> Iterator[Tuple]:
+ """
+ Read all available data from specific start offset and end offset.
+ This is invoked during failure recovery to re-read a batch
deterministically
+ in order to achieve exactly once.
+
+ Parameters
+ ----------
+ start : dict
+ The start offset to start reading from.
+
+ end : dict
+ The offset where the reading stop.
+
+ Returns
+ -------
+ iterator of :class:`Tuple`\\s
+ All the records between start offset and end offset.
+ """
+ raise PySparkNotImplementedError(
+ error_class="NOT_IMPLEMENTED",
+ message_parameters={"feature": "read2"},
+ )
+
+ def commit(self, end: dict) -> None:
+ """
+ Informs the source that Spark has completed processing all data for
offsets less than or
+ equal to `end` and will only request offsets greater than `end` in the
future.
+
+ Parameters
+ ----------
+ end : dict
+ The latest offset that the streaming query has processed for this
source.
+ """
+ ...
+
+
+class _SimpleStreamReaderWrapper(DataSourceStreamReader):
+ """
+ A private class that wrap :class:`SimpleDataSourceStreamReader` in
prefetch and cache pattern,
+ so that :class:`SimpleDataSourceStreamReader` can integrate with streaming
engine like an
+ ordinary :class:`DataSourceStreamReader`.
+
+ current_offset track the latest progress of the record prefetching, it is
initialized to be
+ initialOffset() when query start for the first time or initialized to be
the end offset of
+ the last committed batch when query restarts.
+
+ When streaming engine call latestOffset(), the wrapper calls read() that
start from
+ current_offset, prefetch and cache the data, then update the
current_offset to be
+ the end offset of the new data.
+
+ When streaming engine call planInputPartitions(start, end), the wrapper
get the prefetched data
+ from cache and send it to JVM along with the input partitions.
+
+ When query restart, batches in write ahead offset log that has not been
committed will be
+ replayed by reading data between start and end offset through read2(start,
end).
+ """
+
+ def __init__(self, simple_reader: SimpleDataSourceStreamReader):
+ self.simple_reader = simple_reader
+ self.initial_offset = None
+ self.current_offset = None
+ self.cache: List[Tuple[dict, dict, Iterator[Tuple]]] = []
+
+ def initialOffset(self) -> dict:
+ if self.initial_offset is None:
+ self.initial_offset = self.simple_reader.initialOffset() # type:
ignore[assignment]
+ return self.initial_offset
+
+ def latestOffset(self) -> dict:
+ # when query start for the first time, use initial offset as the start
offset.
+ if self.current_offset is None:
+ self.current_offset = self.initialOffset() # type:
ignore[assignment]
+ (iter, end) = self.simple_reader.read(self.current_offset) # type:
ignore[arg-type]
+ self.cache.append((self.current_offset, end, iter)) # type:
ignore[arg-type]
+ self.current_offset = end # type: ignore[assignment]
+ return end
+
+ def commit(self, end: dict) -> None:
+ if self.current_offset is None:
+ self.current_offset = end # type: ignore[assignment]
+
+ end_idx = -1
+ for i in range(len(self.cache)):
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
