py: implement foreach_chunk method for streaming HTTP output

abhizer · gz · commit cb41db924076 · 2024-05-28T12:42:39.000-07:00
Signed-off-by: Abhinav Gyawali &lt;22275402+abhizer@users.noreply.github.com&gt;
diff --git a/python/feldera/_callback_runner.py b/python/feldera/_callback_runner.py
@@ -0,0 +1,64 @@
+from threading import Thread
+from typing import Callable
+from queue import Queue, Empty
+
+import pandas as pd
+from feldera import FelderaClient
+from feldera._helpers import dataframe_from_response
+from feldera.output_handler import _OutputHandlerInstruction
+
+
+class CallbackRunner(Thread):
+    def __init__(
+            self,
+            client: FelderaClient,
+            pipeline_name: str,
+            view_name: str,
+            callback: Callable[[pd.DataFrame, int], None],
+            queue: Queue,
+    ):
+        super().__init__()
+        self.client: FelderaClient = client
+        self.pipeline_name: str = pipeline_name
+        self.view_name: str = view_name
+        self.callback: Callable[[pd.DataFrame, int], None] = callback
+        self.queue: Queue = queue
+
+    def run(self):
+        """
+        The main loop of the thread. Listens for data and calls the callback function on each chunk of data received.
+
+        :meta private:
+        """
+
+        ack: _OutputHandlerInstruction = self.queue.get()
+
+        match ack:
+            case _OutputHandlerInstruction.PipelineStarted:
+                gen_obj = self.client.listen_to_pipeline(self.pipeline_name, self.view_name, format="json")
+                self.queue.task_done()
+
+                for chunk in gen_obj:
+                    chunk: dict = chunk
+                    data: list[dict] = chunk.get("json_data")
+                    seq_no: int = chunk.get("sequence_number")
+
+                    if data is not None:
+                        self.callback(dataframe_from_response([data]), seq_no)
+
+                    try:
+                        again_ack = self.queue.get_nowait()
+                        if again_ack:
+                            match again_ack:
+                                case _OutputHandlerInstruction.RanToCompletion:
+                                    self.queue.task_done()
+                                    return
+                                case _OutputHandlerInstruction.PipelineStarted:
+                                    self.queue.task_done()
+                                    continue
+                    except Empty:
+                        continue
+
+            case _OutputHandlerInstruction.RanToCompletion:
+                self.queue.task_done()
+                return
diff --git a/python/feldera/_helpers.py b/python/feldera/_helpers.py
@@ -0,0 +1,11 @@
+import pandas as pd
+
+
+def dataframe_from_response(buffer: list[list[dict]]):
+    """
+    Converts the response from Feldera to a pandas DataFrame.
+    """
+    return pd.DataFrame([
+        {**item['insert'], 'insert_delete': 1} if 'insert' in item else {**item['delete'], 'insert_delete': -1}
+        for sublist in buffer for item in sublist
+    ])
diff --git a/python/feldera/output_handler.py b/python/feldera/output_handler.py
@@ -4,6 +4,7 @@
 from threading import Thread
 from queue import Queue, Empty
 from feldera import FelderaClient
+from feldera._helpers import dataframe_from_response
 from enum import Enum
 
 
@@ -66,7 +67,4 @@ def to_pandas(self):
         Converts the output of the pipeline to a pandas DataFrame
         """
         self.join()
-        return pd.DataFrame([
-            {**item['insert'], 'insert_delete': 1} if 'insert' in item else {**item['delete'], 'insert_delete': -1}
-            for sublist in self.buffer for item in sublist
-        ])
+        return dataframe_from_response(self.buffer)
diff --git a/python/feldera/sql_context.py b/python/feldera/sql_context.py
@@ -2,7 +2,7 @@
 import pandas
 import re
 
-from typing import Optional, Dict
+from typing import Optional, Dict, Callable
 from typing_extensions import Self
 from queue import Queue
 
@@ -14,6 +14,7 @@
 from feldera.sql_schema import SQLSchema
 from feldera.output_handler import OutputHandler
 from feldera.output_handler import _OutputHandlerInstruction
+from feldera._callback_runner import CallbackRunner
 from enum import Enum
 
 
@@ -241,6 +242,9 @@ def listen(self, view_name: str) -> OutputHandler:
         Listens to the output of the provided view so that it is available in the notebook / python code.
 
         :param view_name: The name of the view to listen to.
+
+        .. note::
+            - This method must be called before calling :meth:`.run_to_completion`, or :meth:`.start`.
         """
 
         queue = Queue(maxsize=1)
@@ -310,6 +314,32 @@ def connect_sink_delta_table(self, view_name: str, connector_name: str, config:
         else:
             self.output_connectors_buffer[view_name] = [connector]
 
+    def foreach_chunk(self, view_name: str, callback: Callable[[pandas.DataFrame, int], None]):
+        """
+        Runs the given callback on each chunk of the output of the specified view.
+
+        :param view_name: The name of the view.
+        :param callback: The callback to run on each chunk. The callback should take two arguments:
+
+                - **chunk**  -> The chunk as a pandas DataFrame
+                - **seq_no** -> The sequence number. The sequence number is a monotonically increasing integer that
+                  starts from 0. Note that the sequence number is unique for each chunk, but not necessarily contiguous.
+
+        Please note that the callback is run in a separate thread, so it should be thread-safe.
+
+        .. note::
+            - The callback must be thread-safe as it will be run in a separate thread.
+            - This method must be called before calling :meth:`.run_to_completion`, or :meth:`.start`.
+
+        """
+
+        queue = Queue(maxsize=1)
+
+        self.views_tx.append({view_name: queue})
+
+        handler = CallbackRunner(self.client, self.pipeline_name, view_name, callback, queue)
+        handler.start()
+
     def run_to_completion(self):
         """
         .. _run_to_completion:
diff --git a/python/tests/test_wireframes.py b/python/tests/test_wireframes.py
@@ -74,6 +74,36 @@ def test_two_SQLContexts(self):
 
         assert df.columns.tolist() not in df2.columns.tolist()
 
+    def test_foreach_chunk(self):
+        def callback(df: pd.DataFrame, seq_no: int):
+            print(f"\nSeq No: {seq_no}, DF size: {df.shape[0]}\n")
+
+        sql = SQLContext('foreach_chunk', TEST_CLIENT).get_or_create()
+
+        TBL_NAMES = ['students', 'grades']
+        view_name = "average_scores"
+
+        df_students = pd.read_csv('students.csv')
+        df_grades = pd.read_csv('grades.csv')
+
+        sql.register_table(TBL_NAMES[0], SQLSchema({"name": "STRING", "id": "INT"}))
+        sql.register_table(TBL_NAMES[1], SQLSchema({
+            "student_id": "INT",
+            "science": "INT",
+            "maths": "INT",
+            "art": "INT"
+        }))
+
+        query = f"SELECT name, ((science + maths + art) / 3) as average FROM {TBL_NAMES[0]} JOIN {TBL_NAMES[1]} on id = student_id ORDER BY average DESC"
+        sql.register_view(view_name, query)
+
+        sql.connect_source_pandas(TBL_NAMES[0], df_students)
+        sql.connect_source_pandas(TBL_NAMES[1], df_grades)
+
+        sql.foreach_chunk(view_name, callback)
+
+        sql.run_to_completion()
+
 
 if __name__ == '__main__':
     unittest.main()
