Added to_chunked_dataframes generator method in the Job class

Author: voonhous

sdk/python/feast/job.py

@@ -1,12 +1,11 @@
 import tempfile
 import time
 from datetime import datetime, timedelta
-from typing import List
+from typing import Iterable
 from urllib.parse import urlparse
 
 import fastavro
 import pandas as pd
-from fastavro import reader as fastavro_reader
 from google.cloud import storage
 
 from feast.serving.ServingService_pb2 import GetJobRequest
@@ -96,6 +95,7 @@ def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC):
             )
 
         uris = [urlparse(uri) for uri in self.job_proto.file_uris]
+        print(uris)
         for file_uri in uris:
             if file_uri.scheme == "gs":
                 file_obj = tempfile.TemporaryFile()
@@ -113,16 +113,53 @@ def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC):
             for record in avro_reader:
                 yield record
 
-    def to_dataframe(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC):
+    def to_dataframe(
+            self, timeout_sec: int = DEFAULT_TIMEOUT_SEC
+    ) -> pd.DataFrame:
         """
-        Wait until job is done to get an interable rows of result
+        Wait until job is done to get an iterable rows of result.
 
         Args:
-            timeout_sec: max no of seconds to wait until job is done. If "timeout_sec" is exceeded, an exception will be raised.
-        Returns: pandas Dataframe of the feature values
+            timeout_sec (int):
+                Max no of seconds to wait until job is done. If "timeout_sec"
+                is exceeded, an exception will be raised.
+        Returns:
+            pd.DataFrame:
+                Pandas DataFrame of the feature values.
         """
         records = [r for r in self.result(timeout_sec=timeout_sec)]
         return pd.DataFrame.from_records(records)
 
+    def to_chunked_dataframes(
+            self,
+            max_chunk_size: int,
+            timeout_sec: int = DEFAULT_TIMEOUT_SEC
+    ) -> Iterable[pd.DataFrame]:
+        """
+        Wait until a job is done to get an iterable rows of result. This method
+        will split the response into chunked DataFrame of a specified size to
+        to be yielded to the instance calling it.
+
+        Args:
+            max_chunk_size (int):
+                Maximum number of rows that the DataFrame should contain.
+
+            timeout_sec (int):
+                Max no of seconds to wait until job is done. If "timeout_sec"
+                is exceeded, an exception will be raised.
+
+        Returns:
+            pd.DataFrame:
+                Pandas DataFrame of the feature values.
+        """
+
+        records = []
+        for result in self.result(timeout_sec=timeout_sec):
+            result.append(records)
+            if len(records) == max_chunk_size:
+                df = pd.DataFrame.from_records(records)
+                records.clear()  # Empty records array
+                yield df
+
     def __iter__(self):
         return iter(self.result())