[PLT-817] Move to buffered reader (#1584)

Author: adrian-chang

libs/labelbox/src/labelbox/schema/export_task.py

@@ -16,9 +16,13 @@
     Union,
     TYPE_CHECKING,
     overload,
+    Any,
 )
 
 import requests
+import warnings
+import tempfile
+import os
 from labelbox import pydantic_compat
 
 from labelbox.schema.task import Task
@@ -106,7 +110,14 @@ class JsonConverterOutput:
 
 
 class JsonConverter(Converter[JsonConverterOutput]):  # pylint: disable=too-few-public-methods
-    """Converts JSON data."""
+    """Converts JSON data.
+    
+    Deprecated: This converter is deprecated and will be removed in a future release.
+    """
+    
+    def __init__(self) -> None:
+        warnings.warn("JSON converter is deprecated and will be removed in a future release")
+        super().__init__()
 
     def _find_json_object_offsets(self, data: str) -> List[Tuple[int, int]]:
         object_offsets: List[Tuple[int, int]] = []
@@ -166,7 +177,8 @@ class FileConverterOutput:
 
 
 class FileConverter(Converter[FileConverterOutput]):
-    """Converts data to a file."""
+    """Converts data to a file.
+    """
 
     def __init__(self, file_path: str) -> None:
         super().__init__()
@@ -248,6 +260,7 @@ def __init__(
 
     def _find_line_at_offset(self, file_content: str,
                              target_offset: int) -> int:
+        # TODO: Remove this, incorrect parsing of JSON to find braces
         stack = []
         line_number = 0
 
@@ -313,6 +326,7 @@ def __init__(
             )
 
     def _find_offset_of_line(self, file_content: str, target_line: int):
+        # TODO: Remove this, incorrect parsing of JSON to find braces
         start_offset = None
         stack = []
         line_number = 0
@@ -377,9 +391,13 @@ def read(self) -> Iterator[Tuple[_MetadataFileInfo, str]]:
 
 
 class _MultiGCSFileReader(_Reader):  # pylint: disable=too-few-public-methods
-    """Reads data from multiple GCS files in a seamless way."""
+    """Reads data from multiple GCS files in a seamless way.
+    
+    Deprecated: This reader is deprecated and will be removed in a future release.
+    """
 
     def __init__(self):
+        warnings.warn("_MultiGCSFileReader is deprecated and will be removed in a future release")
         super().__init__()
         self._retrieval_strategy = None
 
@@ -397,6 +415,54 @@ def read(self) -> Iterator[Tuple[_MetadataFileInfo, str]]:
             result = self._retrieval_strategy.get_next_chunk()
 
 
+@dataclass
+class BufferedJsonConverterOutput:
+    """Output with the JSON object"""
+    json: Any
+
+
+class _BufferedJsonConverter(Converter[BufferedJsonConverterOutput]):
+    """Converts JSON data in a buffered manner
+    """
+    def convert(
+        self, input_args: Converter.ConverterInputArgs
+    ) -> Iterator[BufferedJsonConverterOutput]:
+        yield BufferedJsonConverterOutput(json=json.loads(input_args.raw_data))
+
+
+class _BufferedGCSFileReader(_Reader): 
+    """Reads data from multiple GCS files and buffer them to disk"""
+
+    def __init__(self):
+        super().__init__()
+        self._retrieval_strategy = None
+
+    def set_retrieval_strategy(self, strategy: FileRetrieverStrategy) -> None:
+        """Sets the retrieval strategy."""
+        self._retrieval_strategy = strategy
+
+    def read(self) -> Iterator[Tuple[_MetadataFileInfo, str]]:
+        if not self._retrieval_strategy:
+            raise ValueError("retrieval strategy not set")
+        # create a buffer
+        with tempfile.NamedTemporaryFile(mode='w+', delete=False) as temp_file:
+            result = self._retrieval_strategy.get_next_chunk()
+            while result:
+                file_info, raw_data = result
+                temp_file.seek(file_info.offsets.start)
+                temp_file.write(raw_data)
+                result = self._retrieval_strategy.get_next_chunk()
+        # read buffer
+        with open(temp_file.name, 'r') as temp_file_reopened:
+            for idx, line in enumerate(temp_file_reopened):
+                yield _MetadataFileInfo(
+                    offsets=Range(start=0, end=len(line) - 1), 
+                    lines=Range(start=idx, end=idx + 1), 
+                    file=temp_file.name), line
+        # manually delete buffer
+        os.unlink(temp_file.name)
+
+
 class Stream(Generic[OutputT]):
     """Streams data from a Reader."""
 
@@ -582,12 +648,9 @@ def errors(self):
         Stream(
             _TaskContext(self._task.client, self._task.uid, StreamType.ERRORS,
                          metadata_header),
-            _MultiGCSFileReader(),
-            JsonConverter(),
-        ).start(stream_handler=lambda output: [
-            data.append(json.loads(row)) for row in output.json_str.split(
-                '\n') if row
-        ])
+            _BufferedGCSFileReader(),
+            _BufferedJsonConverter(),
+        ).start(stream_handler=lambda output: data.append(output.json))
         return data
 
     @property
@@ -607,12 +670,9 @@ def result(self):
             Stream(
                 _TaskContext(self._task.client, self._task.uid,
                              StreamType.RESULT, metadata_header),
-                _MultiGCSFileReader(),
-                JsonConverter(),
-            ).start(stream_handler=lambda output: [
-                data.append(json.loads(row)) for row in output.json_str.split(
-                    '\n') if row
-            ])
+                _BufferedGCSFileReader(),
+                _BufferedJsonConverter(),
+            ).start(stream_handler=lambda output: data.append(output.json))
             return data
         return self._task.result_url
 

libs/labelbox/tests/unit/export_task/test_export_task.py

@@ -0,0 +1,27 @@
+import json
+
+from unittest.mock import MagicMock, patch
+from labelbox.schema.export_task import ExportTask
+
+class TestExportTask:
+
+    def test_export_task(self):
+        with patch('requests.get') as mock_requests_get:
+            mock_task = MagicMock()
+            mock_task.client.execute.side_effect = [
+                {"task": {"exportMetadataHeader": { "total_size": 1, "total_lines": 1, "lines": { "start": 0, "end": 1 }, "offsets": { "start": 0, "end": 0 }, "file": "file" } } },
+                {"task": {"exportFileFromOffset": { "total_size": 1, "total_lines": 1, "lines": { "start": 0, "end": 1 }, "offsets": { "start": 0, "end": 0 }, "file": "file" } } },
+            ]
+            mock_task.status = "COMPLETE"
+            data = {
+                "data_row": {
+                    "raw_data": """
+                    {"raw_text":"}{"}
+                    {"raw_text":"\\nbad"}   
+                    """
+                }
+            }
+            mock_requests_get.return_value.text = json.dumps(data)
+            mock_requests_get.return_value.content = "b"
+            export_task = ExportTask(mock_task, is_export_v2=True)
+            assert export_task.result[0] == data