You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2022/04/05 01:50:46 UTC
[spark] branch master updated: [SPARK-38759][PYTHON][SS] Add StreamingQueryListener support in PySpark
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 608f70d3f63 [SPARK-38759][PYTHON][SS] Add StreamingQueryListener support in PySpark
608f70d3f63 is described below
commit 608f70d3f63cb8d7705f4f2868923f448bf3a43c
Author: Hyukjin Kwon <gu...@apache.org>
AuthorDate: Tue Apr 5 10:50:11 2022 +0900
[SPARK-38759][PYTHON][SS] Add StreamingQueryListener support in PySpark
### What changes were proposed in this pull request?
This PR adds the support of `StreamingQueryListener` in PySpark. This PR adds the Python classes corresponding to Scala classes below:
- `StreamingQueryListener`
- `QueryStartedEvent`
- `QueryProgressEvent`
- `QueryTerminatedEvent`
- `StreamingQueryProgress`
- `StateOperatorProgress`
- `SourceProgress`
- `SinkProgress`
with the listener API, `StreamingQueryManager.addListener`
### Why are the changes needed?
This is a blocker of SPARK-38760. `DataFrame.observe` is an imporant feature for SS users to monitor their queries. However, this is missing in PySpark.
### Does this PR introduce _any_ user-facing change?
Yes, it adds a new API and class as described above. Users can do:
```python
from pyspark.sql.streaming import StreamingQueryListener
class Listener(StreamingQueryListener):
def onQueryStarted(self, event: "QueryStartedEvent") -> None:
print("Wow!")
def onQueryProgress(self, event: "QueryProgressEvent") -> None:
print("Hello!")
def onQueryTerminated(self, event: "QueryTerminatedEvent") -> None:
print("HI!")
spark.streams.addListener(Listener())
rates = spark.readStream.format("rate").option("rowsPerSecond", 1).load()
query = rates.writeStream.format("console").start()
query.awaitTermination()
```
### How was this patch tested?
Unittests added
Closes #36038 from HyukjinKwon/streaming-listener.
Authored-by: Hyukjin Kwon <gu...@apache.org>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
dev/sparktestsupport/modules.py | 2 +
python/docs/source/reference/pyspark.ss.rst | 2 +
python/pyspark/sql/streaming/__init__.py | 2 +
python/pyspark/sql/streaming/listener.py | 666 +++++++++++++++++++++
python/pyspark/sql/streaming/query.py | 25 +-
.../pyspark/sql/tests/test_streaming_listener.py | 283 +++++++++
.../spark/sql/api/python/PythonSQLUtils.scala | 14 +-
.../sql/streaming/StreamingQueryListener.scala | 23 +
8 files changed, 1013 insertions(+), 4 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index d8a35084fac..b0e09884035 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -425,6 +425,7 @@ pyspark_sql = Module(
"pyspark.sql.readwriter",
"pyspark.sql.streaming.query",
"pyspark.sql.streaming.readwriter",
+ "pyspark.sql.streaming.listener",
"pyspark.sql.udf",
"pyspark.sql.window",
"pyspark.sql.avro.functions",
@@ -460,6 +461,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.test_serde",
"pyspark.sql.tests.test_session",
"pyspark.sql.tests.test_streaming",
+ "pyspark.sql.tests.test_streaming_listener",
"pyspark.sql.tests.test_types",
"pyspark.sql.tests.test_udf",
"pyspark.sql.tests.test_udf_profiler",
diff --git a/python/docs/source/reference/pyspark.ss.rst b/python/docs/source/reference/pyspark.ss.rst
index cace2d5a6bb..3246afe50c2 100644
--- a/python/docs/source/reference/pyspark.ss.rst
+++ b/python/docs/source/reference/pyspark.ss.rst
@@ -32,6 +32,7 @@ Core Classes
DataStreamWriter
StreamingQuery
StreamingQueryManager
+ StreamingQueryListener
Input and Output
----------------
@@ -83,6 +84,7 @@ Query Management
StreamingQuery.status
StreamingQuery.stop
StreamingQueryManager.active
+ StreamingQueryManager.addListener
StreamingQueryManager.awaitAnyTermination
StreamingQueryManager.get
StreamingQueryManager.resetTerminated
diff --git a/python/pyspark/sql/streaming/__init__.py b/python/pyspark/sql/streaming/__init__.py
index 1c610e3b6c4..5c07687647a 100644
--- a/python/pyspark/sql/streaming/__init__.py
+++ b/python/pyspark/sql/streaming/__init__.py
@@ -17,3 +17,5 @@
from pyspark.sql.streaming.query import StreamingQuery, StreamingQueryManager # noqa: F401
from pyspark.sql.streaming.readwriter import DataStreamReader, DataStreamWriter # noqa: F401
+from pyspark.sql.streaming.listener import StreamingQueryListener # noqa: F401
+from pyspark.sql.utils import StreamingQueryException # noqa: F401
diff --git a/python/pyspark/sql/streaming/listener.py b/python/pyspark/sql/streaming/listener.py
new file mode 100644
index 00000000000..f8f4949bc6a
--- /dev/null
+++ b/python/pyspark/sql/streaming/listener.py
@@ -0,0 +1,666 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import pickle
+import uuid
+from typing import Optional, Dict, List
+from abc import ABC, abstractmethod
+
+from py4j.java_gateway import JavaObject
+
+from pyspark.sql import Row
+
+
+__all__ = ["StreamingQueryListener"]
+
+
+class StreamingQueryListener(ABC):
+ """
+ Interface for listening to events related to :class:`~pyspark.sql.streaming.StreamingQuery`.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ The methods are not thread-safe as they may be called from different threads.
+ The events received are identical with Scala API. Refer to its documentation.
+
+ This API is evolving.
+
+ Examples
+ --------
+ >>> class MyListener(StreamingQueryListener):
+ ... def onQueryStarted(self, event: QueryStartedEvent) -> None:
+ ... # Do something with event.
+ ... pass
+ ...
+ ... def onQueryProgress(self, event: QueryProgressEvent) -> None:
+ ... # Do something with event.
+ ... pass
+ ...
+ ... def onQueryTerminated(self, event: QueryTerminatedEvent) -> None:
+ ... # Do something with event.
+ ... pass
+ ...
+ >>> spark.streams.addListener(MyListener())
+ """
+
+ @abstractmethod
+ def onQueryStarted(self, event: "QueryStartedEvent") -> None:
+ """
+ Called when a query is started.
+
+ Notes
+ -----
+ This is called synchronously with :py:meth:`~pyspark.sql.streaming.DataStreamWriter.start`,
+ that is, `onQueryStart` will be called on all listeners before `DataStreamWriter.start()`
+ returns the corresponding :class:`~pyspark.streaming.StreamingQuery`.
+ Please don't block this method as it will block your query.
+ """
+ pass
+
+ @abstractmethod
+ def onQueryProgress(self, event: "QueryProgressEvent") -> None:
+ """
+ Called when there is some status update (ingestion rate updated, etc.)
+
+ Notes
+ -----
+ This method is asynchronous. The status in :class:`~pyspark.streaming.StreamingQuery`
+ will always be latest no matter when this method is called. Therefore, the status of
+ :class:`~pyspark.streaming.StreamingQuery`.
+ may be changed before/when you process the event. E.g., you may find
+ :class:`~pyspark.streaming.StreamingQuery` is terminated when you are
+ processing `QueryProgressEvent`.
+ """
+ pass
+
+ @abstractmethod
+ def onQueryTerminated(self, event: "QueryTerminatedEvent") -> None:
+ """
+ Called when a query is stopped, with or without error.
+ """
+ pass
+
+
+class JStreamingQueryListener:
+ """
+ Python class that implements Java interface by Py4J.
+ """
+
+ def __init__(self, pylistener: StreamingQueryListener) -> None:
+ self.pylistener = pylistener
+
+ def onQueryStarted(self, jevent: JavaObject) -> None:
+ self.pylistener.onQueryStarted(QueryStartedEvent(jevent))
+
+ def onQueryProgress(self, jevent: JavaObject) -> None:
+ self.pylistener.onQueryProgress(QueryProgressEvent(jevent))
+
+ def onQueryTerminated(self, jevent: JavaObject) -> None:
+ self.pylistener.onQueryTerminated(QueryTerminatedEvent(jevent))
+
+ class Java:
+ implements = ["org.apache.spark.sql.streaming.PythonStreamingQueryListener"]
+
+
+class QueryStartedEvent:
+ """
+ Event representing the start of a query.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jevent: JavaObject) -> None:
+ self._id: uuid.UUID = uuid.UUID(jevent.id().toString())
+ self._runId: uuid.UUID = uuid.UUID(jevent.runId().toString())
+ self._name: Optional[str] = jevent.name()
+ self._timestamp: str = jevent.timestamp()
+
+ @property
+ def id(self) -> uuid.UUID:
+ """
+ A unique query id that persists across restarts. See
+ py:meth:`~pyspark.streaming.StreamingQuery.id`.
+ """
+ return self._id
+
+ @property
+ def runId(self) -> uuid.UUID:
+ """
+ A query id that is unique for every start/restart. See
+ py:meth:`~pyspark.streaming.StreamingQuery.runId`.
+ """
+ return self._runId
+
+ @property
+ def name(self) -> Optional[str]:
+ """
+ User-specified name of the query, `None` if not specified.
+ """
+ return self._name
+
+ @property
+ def timestamp(self) -> str:
+ """
+ The timestamp to start a query.
+ """
+ return self._timestamp
+
+
+class QueryProgressEvent:
+ """
+ Event representing any progress updates in a query.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jevent: JavaObject) -> None:
+ self._progress: StreamingQueryProgress = StreamingQueryProgress(jevent.progress())
+
+ @property
+ def progress(self) -> "StreamingQueryProgress":
+ """
+ The query progress updates.
+ """
+ return self._progress
+
+
+class QueryTerminatedEvent:
+ """
+ Event representing that termination of a query.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jevent: JavaObject) -> None:
+ self._id: uuid.UUID = uuid.UUID(jevent.id().toString())
+ self._runId: uuid.UUID = uuid.UUID(jevent.runId().toString())
+ jexception = jevent.exception()
+ self._exception: Optional[str] = jexception.get() if jexception.isDefined() else None
+
+ @property
+ def id(self) -> uuid.UUID:
+ """
+ A unique query id that persists across restarts. See
+ py:meth:`~pyspark.streaming.StreamingQuery.id`.
+ """
+ return self._id
+
+ @property
+ def runId(self) -> uuid.UUID:
+ """
+ A query id that is unique for every start/restart. See
+ py:meth:`~pyspark.streaming.StreamingQuery.runId`.
+ """
+ return self._runId
+
+ @property
+ def exception(self) -> Optional[str]:
+ """
+ The exception message of the query if the query was terminated
+ with an exception. Otherwise, it will be `None`.
+ """
+ return self._exception
+
+
+class StreamingQueryProgress:
+ """
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jprogress: JavaObject) -> None:
+ from pyspark import SparkContext
+
+ self._jprogress: JavaObject = jprogress
+ self._id: uuid.UUID = uuid.UUID(jprogress.id().toString())
+ self._runId: uuid.UUID = uuid.UUID(jprogress.runId().toString())
+ self._name: Optional[str] = jprogress.name()
+ self._timestamp: str = jprogress.timestamp()
+ self._batchId: int = jprogress.batchId()
+ self._batchDuration: int = jprogress.batchDuration()
+ self._durationMs: Dict[str, int] = dict(jprogress.durationMs())
+ self._eventTime: Dict[str, str] = dict(jprogress.eventTime())
+ self._stateOperators: List[StateOperatorProgress] = [
+ StateOperatorProgress(js) for js in jprogress.stateOperators()
+ ]
+ self._sources: List[SourceProgress] = [SourceProgress(js) for js in jprogress.sources()]
+ self._sink: SinkProgress = SinkProgress(jprogress.sink())
+
+ # TODO(SPARK-38760): Write a test with DataFrame.observe API implementation.
+ self._observedMetrics: Dict[str, Row] = {
+ k: pickle.loads(
+ SparkContext._jvm.PythonSQLUtils.toPyRow(jr) # type: ignore[union-attr]
+ )
+ for k, jr in dict(jprogress.observedMetrics()).items()
+ }
+
+ @property
+ def id(self) -> uuid.UUID:
+ """
+ A unique query id that persists across restarts. See
+ py:meth:`~pyspark.streaming.StreamingQuery.id`.
+ """
+ return self._id
+
+ @property
+ def runId(self) -> uuid.UUID:
+ """
+ A query id that is unique for every start/restart. See
+ py:meth:`~pyspark.streaming.StreamingQuery.runId`.
+ """
+ return self._runId
+
+ @property
+ def name(self) -> Optional[str]:
+ """
+ User-specified name of the query, `None` if not specified.
+ """
+ return self._name
+
+ @property
+ def timestamp(self) -> str:
+ """
+ The timestamp to start a query.
+ """
+ return self._timestamp
+
+ @property
+ def batchId(self) -> int:
+ """
+ A unique id for the current batch of data being processed. Note that in the
+ case of retries after a failure a given batchId my be executed more than once.
+ Similarly, when there is no data to be processed, the batchId will not be
+ incremented.
+ """
+ return self._batchId
+
+ @property
+ def batchDuration(self) -> int:
+ """
+ The process duration of each batch.
+ """
+ return self._batchDuration
+
+ @property
+ def durationMs(self) -> Dict[str, int]:
+ """
+ The amount of time taken to perform various operations in milliseconds.
+ """
+ return self._durationMs
+
+ @property
+ def eventTime(self) -> Dict[str, str]:
+ """
+ Statistics of event time seen in this batch. It may contain the following keys:
+
+ .. code-block:: python
+
+ {
+ "max": "2016-12-05T20:54:20.827Z", # maximum event time seen in this trigger
+ "min": "2016-12-05T20:54:20.827Z", # minimum event time seen in this trigger
+ "avg": "2016-12-05T20:54:20.827Z", # average event time seen in this trigger
+ "watermark": "2016-12-05T20:54:20.827Z" # watermark used in this trigger
+ }
+
+ All timestamps are in ISO8601 format, i.e. UTC timestamps.
+ """
+ return self._eventTime
+
+ @property
+ def stateOperators(self) -> List["StateOperatorProgress"]:
+ """
+ Information about operators in the query that store state.
+ """
+ return self._stateOperators
+
+ @property
+ def sources(self) -> List["SourceProgress"]:
+ """
+ detailed statistics on data being read from each of the streaming sources.
+ """
+ return self._sources
+
+ @property
+ def sink(self) -> "SinkProgress":
+ """
+ A unique query id that persists across restarts. See
+ py:meth:`~pyspark.streaming.StreamingQuery.id`.
+ """
+ return self._sink
+
+ @property
+ def observedMetrics(self) -> Dict[str, Row]:
+ return self._observedMetrics
+
+ @property
+ def numInputRows(self) -> Optional[str]:
+ """
+ The aggregate (across all sources) number of records processed in a trigger.
+ """
+ return self._jprogress.numInputRows()
+
+ @property
+ def inputRowsPerSecond(self) -> str:
+ """
+ The aggregate (across all sources) rate of data arriving.
+ """
+ return self._jprogress.inputRowsPerSecond()
+
+ @property
+ def processedRowsPerSecond(self) -> str:
+ """
+ The aggregate (across all sources) rate at which Spark is processing data..
+ """
+ return self._jprogress.processedRowsPerSecond()
+
+ @property
+ def json(self) -> str:
+ """
+ The compact JSON representation of this progress.
+ """
+ return self._jprogress.json()
+
+ @property
+ def prettyJson(self) -> str:
+ """
+ The pretty (i.e. indented) JSON representation of this progress.
+ """
+ return self._jprogress.prettyJson()
+
+ def __str__(self) -> str:
+ return self.prettyJson
+
+
+class StateOperatorProgress:
+ """
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jprogress: JavaObject) -> None:
+ self._jprogress: JavaObject = jprogress
+ self._operatorName: str = jprogress.operatorName()
+ self._numRowsTotal: int = jprogress.numRowsTotal()
+ self._numRowsUpdated: int = jprogress.numRowsUpdated()
+ self._allUpdatesTimeMs: int = jprogress.allUpdatesTimeMs()
+ self._numRowsRemoved: int = jprogress.numRowsRemoved()
+ self._allRemovalsTimeMs: int = jprogress.allRemovalsTimeMs()
+ self._commitTimeMs: int = jprogress.commitTimeMs()
+ self._memoryUsedBytes: int = jprogress.memoryUsedBytes()
+ self._numRowsDroppedByWatermark: int = jprogress.numRowsDroppedByWatermark()
+ self._numShufflePartitions: int = jprogress.numShufflePartitions()
+ self._numStateStoreInstances: int = jprogress.numStateStoreInstances()
+ self._customMetrics: Dict[str, int] = dict(jprogress.customMetrics())
+
+ @property
+ def operatorName(self) -> str:
+ return self._operatorName
+
+ @property
+ def numRowsTotal(self) -> int:
+ return self._numRowsTotal
+
+ @property
+ def numRowsUpdated(self) -> int:
+ return self._numRowsUpdated
+
+ @property
+ def allUpdatesTimeMs(self) -> int:
+ return self._allUpdatesTimeMs
+
+ @property
+ def numRowsRemoved(self) -> int:
+ return self._numRowsRemoved
+
+ @property
+ def allRemovalsTimeMs(self) -> int:
+ return self._allRemovalsTimeMs
+
+ @property
+ def commitTimeMs(self) -> int:
+ return self._commitTimeMs
+
+ @property
+ def memoryUsedBytes(self) -> int:
+ return self._memoryUsedBytes
+
+ @property
+ def numRowsDroppedByWatermark(self) -> int:
+ return self._numRowsDroppedByWatermark
+
+ @property
+ def numShufflePartitions(self) -> int:
+ return self._numShufflePartitions
+
+ @property
+ def numStateStoreInstances(self) -> int:
+ return self._numStateStoreInstances
+
+ @property
+ def customMetrics(self) -> Dict[str, int]:
+ return self._customMetrics
+
+ @property
+ def json(self) -> str:
+ """
+ The compact JSON representation of this progress.
+ """
+ return self._jprogress.json()
+
+ @property
+ def prettyJson(self) -> str:
+ """
+ The pretty (i.e. indented) JSON representation of this progress.
+ """
+ return self._jprogress.prettyJson()
+
+ def __str__(self) -> str:
+ return self.prettyJson
+
+
+class SourceProgress:
+ """
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jprogress: JavaObject) -> None:
+ self._jprogress: JavaObject = jprogress
+ self._description: str = jprogress.description()
+ self._startOffset: str = jprogress.startOffset()
+ self._endOffset: str = jprogress.endOffset()
+ self._latestOffset: str = jprogress.latestOffset()
+ self._numInputRows: int = jprogress.numInputRows()
+ self._inputRowsPerSecond: float = jprogress.inputRowsPerSecond()
+ self._processedRowsPerSecond: float = jprogress.processedRowsPerSecond()
+ self._metrics: Dict[str, str] = dict(jprogress.metrics())
+
+ @property
+ def description(self) -> str:
+ """
+ Description of the source.
+ """
+ return self._description
+
+ @property
+ def startOffset(self) -> str:
+ """
+ The starting offset for data being read.
+ """
+ return self._startOffset
+
+ @property
+ def endOffset(self) -> str:
+ """
+ The ending offset for data being read.
+ """
+ return self._endOffset
+
+ @property
+ def latestOffset(self) -> str:
+ """
+ The latest offset from this source.
+ """
+ return self._latestOffset
+
+ @property
+ def numInputRows(self) -> int:
+ """
+ The number of records read from this source.
+ """
+ return self._numInputRows
+
+ @property
+ def inputRowsPerSecond(self) -> float:
+ """
+ The rate at which data is arriving from this source.
+ """
+ return self._inputRowsPerSecond
+
+ @property
+ def processedRowsPerSecond(self) -> float:
+ """
+ The rate at which data from this source is being processed by Spark.
+ """
+ return self._processedRowsPerSecond
+
+ @property
+ def metrics(self) -> Dict[str, str]:
+ return self._metrics
+
+ @property
+ def json(self) -> str:
+ """
+ The compact JSON representation of this progress.
+ """
+ return self._jprogress.json()
+
+ @property
+ def prettyJson(self) -> str:
+ """
+ The pretty (i.e. indented) JSON representation of this progress.
+ """
+ return self._jprogress.prettyJson()
+
+ def __str__(self) -> str:
+ return self.prettyJson
+
+
+class SinkProgress:
+ """
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ This API is evolving.
+ """
+
+ def __init__(self, jprogress: JavaObject) -> None:
+ self._jprogress: JavaObject = jprogress
+ self._description: str = jprogress.description()
+ self._numOutputRows: int = jprogress.numOutputRows()
+ self._metrics: Dict[str, str] = dict(jprogress.metrics())
+
+ @property
+ def description(self) -> str:
+ """
+ Description of the source.
+ """
+ return self._description
+
+ @property
+ def numOutputRows(self) -> int:
+ """
+ Number of rows written to the sink or -1 for Continuous Mode (temporarily)
+ or Sink V1 (until decommissioned).
+ """
+ return self._numOutputRows
+
+ @property
+ def metrics(self) -> Dict[str, str]:
+ return self._metrics
+
+ @property
+ def json(self) -> str:
+ """
+ The compact JSON representation of this progress.
+ """
+ return self._jprogress.json()
+
+ @property
+ def prettyJson(self) -> str:
+ """
+ The pretty (i.e. indented) JSON representation of this progress.
+ """
+ return self._jprogress.prettyJson()
+
+ def __str__(self) -> str:
+ return self.prettyJson
+
+
+def _test() -> None:
+ import sys
+ import doctest
+ import os
+ from pyspark.sql import SparkSession
+ import pyspark.sql.streaming.listener
+ from py4j.protocol import Py4JError
+
+ os.chdir(os.environ["SPARK_HOME"])
+
+ globs = pyspark.sql.streaming.listener.__dict__.copy()
+ try:
+ spark = SparkSession._getActiveSessionOrCreate()
+ except Py4JError: # noqa: F821
+ spark = SparkSession(sc) # type: ignore[name-defined] # noqa: F821
+
+ globs["spark"] = spark
+
+ (failure_count, test_count) = doctest.testmod(
+ pyspark.sql.streaming.listener,
+ globs=globs,
+ )
+ globs["spark"].stop()
+
+ if failure_count:
+ sys.exit(-1)
+
+
+if __name__ == "__main__":
+ _test()
diff --git a/python/pyspark/sql/streaming/query.py b/python/pyspark/sql/streaming/query.py
index 8565435da29..83965b955c0 100644
--- a/python/pyspark/sql/streaming/query.py
+++ b/python/pyspark/sql/streaming/query.py
@@ -19,11 +19,11 @@ import json
import sys
from typing import Any, Dict, List, Optional
-from py4j.java_gateway import JavaObject
+from py4j.java_gateway import JavaObject, java_import
from pyspark import since
from pyspark.sql.utils import StreamingQueryException
-
+from pyspark.sql.streaming.listener import StreamingQueryListener, JStreamingQueryListener
__all__ = ["StreamingQuery", "StreamingQueryManager"]
@@ -294,6 +294,27 @@ class StreamingQueryManager:
"""
self._jsqm.resetTerminated()
+ def addListener(self, listener: StreamingQueryListener) -> None:
+ """
+ Register a :class:`StreamingQueryListener` to receive up-calls for life cycle events of
+ :class:`~pyspark.sql.streaming.StreamingQuery`.
+
+ .. versionadded:: 3.4.0
+ """
+ from pyspark import SparkContext
+ from pyspark.java_gateway import ensure_callback_server_started
+
+ gw = SparkContext._gateway
+ assert gw is not None
+ java_import(gw.jvm, "org.apache.spark.sql.streaming.*")
+ ensure_callback_server_started(gw)
+
+ self._jsqm.addListener(
+ SparkContext._jvm.PythonStreamingQueryListenerWrapper( # type: ignore[union-attr]
+ JStreamingQueryListener(listener)
+ )
+ )
+
def _test() -> None:
import doctest
diff --git a/python/pyspark/sql/tests/test_streaming_listener.py b/python/pyspark/sql/tests/test_streaming_listener.py
new file mode 100644
index 00000000000..22a52e6ccab
--- /dev/null
+++ b/python/pyspark/sql/tests/test_streaming_listener.py
@@ -0,0 +1,283 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import json
+import time
+import uuid
+from datetime import datetime
+
+from pyspark.sql.streaming import StreamingQueryListener
+from pyspark.sql.streaming.listener import (
+ QueryStartedEvent,
+ QueryProgressEvent,
+ QueryTerminatedEvent,
+ SinkProgress,
+ SourceProgress,
+ StateOperatorProgress,
+ StreamingQueryProgress,
+)
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+
+
+class StreamingListenerTests(ReusedSQLTestCase):
+ def test_number_of_public_methods(self):
+ msg = (
+ "New field or method was detected in JVM side. If you added a new public "
+ "field or method, implement that in the corresponding Python class too."
+ "Otherwise, fix the number on the assert here."
+ )
+
+ def get_number_of_public_methods(clz):
+ return len(
+ self.spark.sparkContext._jvm.org.apache.spark.util.Utils.classForName(
+ clz, True, False
+ ).getMethods()
+ )
+
+ self.assertEquals(
+ get_number_of_public_methods(
+ "org.apache.spark.sql.streaming.StreamingQueryListener$QueryStartedEvent"
+ ),
+ 14,
+ msg,
+ )
+ self.assertEquals(
+ get_number_of_public_methods(
+ "org.apache.spark.sql.streaming.StreamingQueryListener$QueryProgressEvent"
+ ),
+ 11,
+ msg,
+ )
+ self.assertEquals(
+ get_number_of_public_methods(
+ "org.apache.spark.sql.streaming.StreamingQueryListener$QueryTerminatedEvent"
+ ),
+ 13,
+ msg,
+ )
+ self.assertEquals(
+ get_number_of_public_methods("org.apache.spark.sql.streaming.StreamingQueryProgress"),
+ 38,
+ msg,
+ )
+ self.assertEquals(
+ get_number_of_public_methods("org.apache.spark.sql.streaming.StateOperatorProgress"),
+ 27,
+ msg,
+ )
+ self.assertEquals(
+ get_number_of_public_methods("org.apache.spark.sql.streaming.SourceProgress"), 21, msg
+ )
+ self.assertEquals(
+ get_number_of_public_methods("org.apache.spark.sql.streaming.SinkProgress"), 19, msg
+ )
+
+ def test_listener_events(self):
+
+ start_event = None
+ progress_event = None
+ terminated_event = None
+
+ class TestListener(StreamingQueryListener):
+ def onQueryStarted(self, event) -> None:
+ nonlocal start_event
+ start_event = event
+
+ def onQueryProgress(self, event) -> None:
+ nonlocal progress_event
+ progress_event = event
+
+ def onQueryTerminated(self, event) -> None:
+ nonlocal terminated_event
+ terminated_event = event
+
+ self.spark.streams.addListener(TestListener())
+
+ df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load()
+ q = df.writeStream.format("noop").queryName("test").start()
+ self.assertTrue(q.isActive)
+ time.sleep(10)
+ q.stop()
+
+ # Make sure all events are empty
+ self.spark.sparkContext._jsc.sc().listenerBus().waitUntilEmpty()
+
+ self.check_start_event(start_event)
+ self.check_progress_event(progress_event)
+ self.check_terminated_event(terminated_event)
+
+ def check_start_event(self, event):
+ """Check QueryStartedEvent"""
+ self.assertTrue(isinstance(event, QueryStartedEvent))
+ self.assertTrue(isinstance(event.id, uuid.UUID))
+ self.assertTrue(isinstance(event.runId, uuid.UUID))
+ self.assertEquals(event.name, "test")
+ try:
+ datetime.strptime(event.timestamp, "%Y-%m-%dT%H:%M:%S.%fZ")
+ except ValueError:
+ self.fail("'%s' is not in ISO 8601 format.")
+
+ def check_progress_event(self, event):
+ """Check QueryProgressEvent"""
+ self.assertTrue(isinstance(event, QueryProgressEvent))
+ self.check_streaming_query_progress(event.progress)
+
+ def check_terminated_event(self, event):
+ """Check QueryTerminatedEvent"""
+ self.assertTrue(isinstance(event, QueryTerminatedEvent))
+ self.assertTrue(isinstance(event.id, uuid.UUID))
+ self.assertTrue(isinstance(event.runId, uuid.UUID))
+ # TODO: Needs a test for exception.
+ self.assertEquals(event.exception, None)
+
+ def check_streaming_query_progress(self, progress):
+ """Check StreamingQueryProgress"""
+ self.assertTrue(isinstance(progress, StreamingQueryProgress))
+ self.assertTrue(isinstance(progress.id, uuid.UUID))
+ self.assertTrue(isinstance(progress.runId, uuid.UUID))
+ self.assertEquals(progress.name, "test")
+ try:
+ json.loads(progress.json)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(progress.prettyJson)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(str(progress))
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ datetime.strptime(progress.timestamp, "%Y-%m-%dT%H:%M:%S.%fZ")
+ except Exception:
+ self.fail("'%s' is not in ISO 8601 format.")
+ self.assertTrue(isinstance(progress.batchId, int))
+ self.assertTrue(isinstance(progress.batchDuration, int))
+ self.assertTrue(isinstance(progress.durationMs, dict))
+ self.assertTrue(
+ set(progress.durationMs.keys()).issubset(
+ {
+ "triggerExecution",
+ "queryPlanning",
+ "getBatch",
+ "commitOffsets",
+ "latestOffset",
+ "addBatch",
+ "walCommit",
+ }
+ )
+ )
+ self.assertTrue(all(map(lambda v: isinstance(v, int), progress.durationMs.values())))
+
+ self.assertEquals(progress.eventTime, {})
+
+ self.assertTrue(isinstance(progress.stateOperators, list))
+ for so in progress.stateOperators:
+ self.check_state_operator_progress(so)
+
+ self.assertTrue(isinstance(progress.sources, list))
+ for so in progress.stateOperators:
+ self.check_source_progress(so)
+
+ self.assertTrue(isinstance(progress.sink, SinkProgress))
+ self.check_sink_progress(progress.sink)
+ self.assertTrue(isinstance(progress.observedMetrics, dict))
+
+ def check_state_operator_progress(self, progress):
+ """Check StateOperatorProgress"""
+ self.assertTrue(isinstance(progress, StateOperatorProgress))
+ try:
+ json.loads(progress.json)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(progress.prettyJson)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(str(progress))
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ self.assertTrue(isinstance(progress.operatorName, str))
+ self.assertTrue(isinstance(progress.numRowsTotal, int))
+ self.assertTrue(isinstance(progress.numRowsUpdated, int))
+ self.assertTrue(isinstance(progress.allUpdatesTimeMs, int))
+ self.assertTrue(isinstance(progress.numRowsRemoved, int))
+ self.assertTrue(isinstance(progress.allRemovalsTimeMs, int))
+ self.assertTrue(isinstance(progress.commitTimeMs, int))
+ self.assertTrue(isinstance(progress.memoryUsedBytes, int))
+ self.assertTrue(isinstance(progress.numRowsDroppedByWatermark, int))
+ self.assertTrue(isinstance(progress.numShufflePartitions, int))
+ self.assertTrue(isinstance(progress.numStateStoreInstances, int))
+ self.assertTrue(isinstance(progress.customMetrics, dict))
+
+ def check_source_progress(self, progress):
+ """Check SourceProgress"""
+ self.assertTrue(isinstance(progress, SourceProgress))
+ try:
+ json.loads(progress.json)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(progress.prettyJson)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(str(progress))
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ self.assertTrue(isinstance(progress.description, str))
+ self.assertTrue(isinstance(progress.startOffset, str))
+ self.assertTrue(isinstance(progress.endOffset, str))
+ self.assertTrue(isinstance(progress.latestOffset, str))
+ self.assertTrue(isinstance(progress.numInputRows, int))
+ self.assertTrue(isinstance(progress.inputRowsPerSecond, float))
+ self.assertTrue(isinstance(progress.processedRowsPerSecond, float))
+ self.assertTrue(isinstance(progress.metrics, dict))
+
+ def check_sink_progress(self, progress):
+ """Check SinkProgress"""
+ self.assertTrue(isinstance(progress, SinkProgress))
+ try:
+ json.loads(progress.json)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(progress.prettyJson)
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ try:
+ json.loads(str(progress))
+ except Exception:
+ self.fail("'%s' is not a valid JSON.")
+ self.assertTrue(isinstance(progress.description, str))
+ self.assertTrue(isinstance(progress.numOutputRows, int))
+ self.assertTrue(isinstance(progress.metrics, dict))
+
+
+if __name__ == "__main__":
+ import unittest
+
+ from pyspark.sql.tests.test_streaming_listener import * # noqa: F401
+
+ try:
+ import xmlrunner # type: ignore[import]
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala
index ab43aa49944..e1efcf81063 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala
@@ -20,20 +20,25 @@ package org.apache.spark.sql.api.python
import java.io.InputStream
import java.nio.channels.Channels
+import net.razorvine.pickle.Pickler
+
import org.apache.spark.api.java.JavaRDD
import org.apache.spark.api.python.PythonRDDServer
import org.apache.spark.internal.Logging
import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{Column, DataFrame, SparkSession}
+import org.apache.spark.sql.{Column, DataFrame, Row, SparkSession}
import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
-import org.apache.spark.sql.catalyst.expressions.{CastTimestampNTZToLong, ExpressionInfo}
+import org.apache.spark.sql.catalyst.expressions.{CastTimestampNTZToLong, ExpressionInfo, GenericRowWithSchema}
import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
import org.apache.spark.sql.execution.{ExplainMode, QueryExecution}
import org.apache.spark.sql.execution.arrow.ArrowConverters
+import org.apache.spark.sql.execution.python.EvaluatePython
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.DataType
private[sql] object PythonSQLUtils extends Logging {
+ private lazy val rowPickler = new Pickler(true, false)
+
def parseDataType(typeText: String): DataType = CatalystSqlParser.parseDataType(typeText)
// This is needed when generating SQL documentation for built-in functions.
@@ -78,6 +83,11 @@ private[sql] object PythonSQLUtils extends Logging {
queryExecution.explainString(ExplainMode.fromString(mode))
}
+ def toPyRow(row: Row): Array[Byte] = {
+ assert(row.isInstanceOf[GenericRowWithSchema])
+ rowPickler.dumps(EvaluatePython.toJava(row, row.schema))
+ }
+
def castTimestampNTZToLong(c: Column): Column = Column(CastTimestampNTZToLong(c.expr))
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryListener.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryListener.scala
index 7ae38c71a00..23a4c43709e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryListener.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryListener.scala
@@ -62,6 +62,29 @@ abstract class StreamingQueryListener {
def onQueryTerminated(event: QueryTerminatedEvent): Unit
}
+/**
+ * Py4J allows a pure interface so this proxy is required.
+ */
+private[spark] trait PythonStreamingQueryListener {
+ import StreamingQueryListener._
+
+ def onQueryStarted(event: QueryStartedEvent): Unit
+
+ def onQueryProgress(event: QueryProgressEvent): Unit
+
+ def onQueryTerminated(event: QueryTerminatedEvent): Unit
+}
+
+private[spark] class PythonStreamingQueryListenerWrapper(
+ listener: PythonStreamingQueryListener) extends StreamingQueryListener {
+ import StreamingQueryListener._
+
+ def onQueryStarted(event: QueryStartedEvent): Unit = listener.onQueryStarted(event)
+
+ def onQueryProgress(event: QueryProgressEvent): Unit = listener.onQueryProgress(event)
+
+ def onQueryTerminated(event: QueryTerminatedEvent): Unit = listener.onQueryTerminated(event)
+}
/**
* Companion object of [[StreamingQueryListener]] that defines the listener events.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org