[SPARK-6328] [Python] Python API for StreamingListener

python/pyspark/streaming/__init__.py

@@ -17,5 +17,6 @@
 
 from pyspark.streaming.context import StreamingContext
 from pyspark.streaming.dstream import DStream
+from pyspark.streaming.listener import StreamingListener
 
-__all__ = ['StreamingContext', 'DStream']
+__all__ = ['StreamingContext', 'DStream', 'StreamingListener']

python/pyspark/streaming/context.py

@@ -28,6 +28,7 @@
 from pyspark.storagelevel import StorageLevel
 from pyspark.streaming.dstream import DStream
 from pyspark.streaming.util import TransformFunction, TransformFunctionSerializer
+from pyspark.streaming.listener import StreamingListener
 
 __all__ = ["StreamingContext"]
 
@@ -363,3 +364,10 @@ def union(self, *dstreams):
         first = dstreams[0]
         jrest = [d._jdstream for d in dstreams[1:]]
         return DStream(self._jssc.union(first._jdstream, jrest), self, first._jrdd_deserializer)
+
+    def addStreamingListener(self, streamingListener):
+        """
+        Add a [[org.apache.spark.streaming.scheduler.StreamingListener]] object for
+        receiving system events related to streaming.
+        """
+        self._jssc.addStreamingListener(streamingListener)

python/pyspark/streaming/listener.py

@@ -0,0 +1,192 @@
+#
+# 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.
+#
+
+__all__ = ["StreamingListener"]
+
+
+class StreamingListener(object):
+
+    def __init__(self):
+        pass
+
+    def onReceiverStarted(self, receiverStarted):
+        """
+        Called when a receiver has been started
+        """
+        pass
+
+    def onReceiverError(self, receiverError):
+        """
+        Called when a receiver has reported an error
+        """
+        pass
+
+    def onReceiverStopped(self, receiverStopped):
+        """
+        Called when a receiver has been stopped
+        """
+        pass
+
+    def onBatchSubmitted(self, batchSubmitted):
+        """
+        Called when a batch of jobs has been submitted for processing.
+        """
+        pass
+
+    def onBatchStarted(self, batchStarted):
+        """
+        Called when processing of a batch of jobs has started.
+        """
+        pass
+
+    def onBatchCompleted(self, batchCompleted):
+        """
+        Called when processing of a batch of jobs has completed.
+        """
+        pass
+
+    def onOutputOperationStarted(self, outputOperationStarted):
+        """
+        Called when processing of a job of a batch has started.
+        """
+        pass
+
+    def onOutputOperationCompleted(self, outputOperationCompleted):
+        """
+        Called when processing of a job of a batch has completed
+        """
+        pass
+
+    def getEventInfo(self, event):
+        """
+        :param event: StreamingListenerEvent
+        :return Returns a BatchInfo, OutputOperationInfo, or ReceiverInfo based on
+                event passed.
+        """
+        event_name = event.getClass().getSimpleName()
+        if 'Batch' in event_name:
+            return BatchInfo(event.batchInfo())
+
+        elif 'Output' in event_name:
+            return OutputOperationInfo(event.outputOperationInfo())
+
+        elif 'Receiver' in event_name:
+            return ReceiverInfo(event.receiverInfo())
+
+    class Java:
+        implements = ["org.apache.spark.streaming.scheduler.StreamingListener"]
+
+
+class BatchInfo(object):
+
+    def __init__(self, javaBatchInfo):
+
+        self.processingStartTime = None
+        self.processingEndTime = None
+
+        self.batchTime = javaBatchInfo.batchTime()
+        self.streamIdToInputInfo = self._map2dict(javaBatchInfo.streamIdToInputInfo())
+
+        self.submissionTime = javaBatchInfo.submissionTime()
+        if javaBatchInfo.processingStartTime().isEmpty() is False:
+            self.processingStartTime = javaBatchInfo.processingStartTime().get()
+        if javaBatchInfo.processingEndTime().isEmpty() is False:
+            self.processingEndTime = javaBatchInfo.processingEndTime().get()
+
+        self.outputOperationInfos = self._map2dict(javaBatchInfo.outputOperationInfos())
+
+    def schedulingDelay(self):
+        """
+        Time taken for the first job of this batch to start processing from the time this batch
+        was submitted to the streaming scheduler.
+        """
+        if self.processingStartTime is None:
+            return None
+        else:
+            return self.processingStartTime - self.submissionTime
+
+    def processingDelay(self):
+        """
+        Time taken for the all jobs of this batch to finish processing from the time they started
+        processing.
+        """
+        if self.processingEndTime is None or self.processingStartTime is None:
+            return None
+        else:
+            return self.processingEndTime - self.processingStartTime
+
+    def totalDelay(self):
+        """
+        Time taken for all the jobs of this batch to finish processing from the time they
+        were submitted
+        """
+        if self.processingEndTime is None or self.processingStartTime is None:
+            return None
+        else:
+            return self.processingDelay() + self.schedulingDelay()
+
+    def numRecords(self):
+        """
+        The number of recorders received by the receivers in this batch.
+        """
+        return len(self.streamIdToInputInfo)
+
+    def _map2dict(self, javaMap):
+        """
+        Converts a scala.collection.immutable.Map to a Python dict
+        """
+        mapping = dict()
+        map_iterator = javaMap.iterator()
+        while map_iterator.hasNext():
+            entry = map_iterator.next()
+            mapping[entry._1()] = entry._2()
+        return mapping
+
+
+class OutputOperationInfo(object):
+
+    def __init__(self, outputOperationInfo):
+        self.batchTime = outputOperationInfo.batchTime()
+        self.id = outputOperationInfo.id()
+        self.name = outputOperationInfo.name()
+        self.startTime = None
+        if outputOperationInfo.startTime().isEmpty() is False:
+            self.startTime = outputOperationInfo.startTime().get()
+        self.endTime = None
+        if outputOperationInfo.endTime().isEmpty() is False:
+            self.endTime = outputOperationInfo.endTime().get()
+        self.failureReason = None
+        if outputOperationInfo.failureReason().isEmpty() is False:
+            self.failureReason = outputOperationInfo.failureReason().get()
+
+    def duration(self):
+        if self.endTime is None or self.startTime is None:
+            return None
+        else:
+            return self.endTime - self.startTime
+
+
+class ReceiverInfo(object):
+
+    def __init__(self, receiverInfo):
+        self.streamId = receiverInfo.streamId()
+        self.name = receiverInfo.name()
+        self.active = receiverInfo.active()
+        self.location = receiverInfo.location()
+        self.lastErrorMessage = receiverInfo.lastErrorMessage()
+        self.lastError = receiverInfo.lastError()
+        self.lastErrorTime = receiverInfo.lastErrorTime()

python/pyspark/streaming/tests.py

@@ -43,6 +43,7 @@
 from pyspark.streaming.flume import FlumeUtils
 from pyspark.streaming.mqtt import MQTTUtils
 from pyspark.streaming.kinesis import KinesisUtils, InitialPositionInStream
+from pyspark.streaming.listener import StreamingListener, BatchInfo
 
 
 class PySparkStreamingTestCase(unittest.TestCase):
@@ -398,6 +399,69 @@ def func(dstream):
         self._test_func(input, func, expected)
 
 
+class StreamingListenerTests(PySparkStreamingTestCase):
+
+    duration = .5
+
+    class BatchInfoCollector(StreamingListener):
+
+        def __init__(self):
+            super(StreamingListener, self).__init__()
+            self.batchInfosCompleted = []
+            self.batchInfosStarted = []
+            self.batchInfosSubmitted = []
+
+        def onBatchSubmitted(self, batchSubmitted):
+            self.batchInfosSubmitted.append(self.getEventInfo(batchSubmitted))
+
+        def onBatchStarted(self, batchStarted):
+            self.batchInfosStarted.append(self.getEventInfo(batchStarted))
+
+        def onBatchCompleted(self, batchCompleted):
+            self.batchInfosCompleted.append(self.getEventInfo(batchCompleted))
+
+    def test_batch_info_reports(self):
+        batch_collector = self.BatchInfoCollector()
+        self.ssc.addStreamingListener(batch_collector)
+        input = [[1], [2], [3], [4]]
+
+        def func(dstream):
+            return dstream.map(int)
+        expected = [[1], [2], [3], [4]]
+        self._test_func(input, func, expected)
+
+        # Test occasionally fails without a delay
+        time.sleep(.1)
+
+        batchInfosSubmitted = batch_collector.batchInfosSubmitted
+        self.assertEqual(len(batchInfosSubmitted), 4)
+
+        for info in batchInfosSubmitted:
+
+            self.assertIsNone(info.schedulingDelay())
+            self.assertIsNone(info.processingDelay())
+            self.assertIsNone(info.totalDelay())
+
+        batchInfosStarted = batch_collector.batchInfosStarted
+        self.assertEqual(len(batchInfosStarted), 4)
+        for info in batchInfosStarted:
+            self.assertIsNotNone(info.schedulingDelay())
+            self.assertGreaterEqual(info.schedulingDelay(), 0)
+            self.assertIsNone(info.processingDelay())
+            self.assertIsNone(info.totalDelay())
+
+        batchInfosCompleted = batch_collector.batchInfosCompleted
+        self.assertEqual(len(batchInfosCompleted), 4)
+
+        for info in batchInfosCompleted:
+            self.assertIsNotNone(info.schedulingDelay())
+            self.assertIsNotNone(info.processingDelay())
+            self.assertIsNotNone(info.totalDelay())
+            self.assertGreaterEqual(info.schedulingDelay(), 0)
+            self.assertGreaterEqual(info.processingDelay(), 0)
+            self.assertGreaterEqual(info.totalDelay(), 0)
+
+
 class WindowFunctionTests(PySparkStreamingTestCase):
 
     timeout = 15
@@ -1281,7 +1345,8 @@ def search_kinesis_asl_assembly_jar():
 
     os.environ["PYSPARK_SUBMIT_ARGS"] = "--jars %s pyspark-shell" % jars
     testcases = [BasicOperationTests, WindowFunctionTests, StreamingContextTests, CheckpointTests,
-                 KafkaStreamTests, FlumeStreamTests, FlumePollingStreamTests, MQTTStreamTests]
+                 KafkaStreamTests, FlumeStreamTests, FlumePollingStreamTests, MQTTStreamTests,
+                 StreamingListenerTests]
 
     if kinesis_jar_present is True:
         testcases.append(KinesisStreamTests)