Buffered stream

Author: Adrian Chang

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

@@ -414,59 +414,6 @@ def read(self) -> Iterator[Tuple[_MetadataFileInfo, str]]:
             yield file_info, raw_data
             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:
-                _, raw_data = result
-                # there is something wrong with the way the offsets are being calculated
-                # so just write all of the chunks as is too the file, with pointer initially
-                # pointed to the start of the file (like what is in GCS) and do not
-                # rely on offsets for file location
-                # 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."""
 
@@ -524,6 +471,142 @@ def start(
                 stream_handler(output)
 
 
+class _BufferedFileRetrieverByOffset(FileRetrieverStrategy):  # pylint: disable=too-few-public-methods
+    """Retrieves files by offset."""
+
+    def __init__(
+        self,
+        ctx: _TaskContext,
+        offset: int,
+    ) -> None:
+        super().__init__(ctx)
+        self._current_offset = offset
+        self._current_line: Optional[int] = None
+        if self._current_offset >= self._ctx.metadata_header.total_size:
+            raise ValueError(
+                f"offset is out of range, max offset is {self._ctx.metadata_header.total_size - 1}"
+            )
+
+    def get_next_chunk(self) -> Optional[Tuple[_MetadataFileInfo, str]]:
+        if self._current_offset >= self._ctx.metadata_header.total_size:
+            return None
+        query = (
+            f"query GetExportFileFromOffsetPyApi"
+            f"($where: WhereUniqueIdInput, $streamType: TaskStreamType!, $offset: UInt64!)"
+            f"{{task(where: $where)"
+            f"{{{'exportFileFromOffset'}(streamType: $streamType, offset: $offset)"
+            f"{{offsets {{start end}} lines {{start end}} file}}"
+            f"}}}}")
+        variables = {
+            "where": {
+                "id": self._ctx.task_id
+            },
+            "streamType": self._ctx.stream_type.value,
+            "offset": str(self._current_offset),
+        }
+        file_info, file_content = self._get_file_content(
+            query, variables, "exportFileFromOffset")
+        file_info.offsets.start = self._current_offset
+        file_info.lines.start = self._current_line
+        self._current_offset = file_info.offsets.end + 1
+        self._current_line = file_info.lines.end + 1
+        return file_info, file_content
+
+
+class BufferedStream(Generic[OutputT]):
+    """Streams data from a Reader."""
+
+    def __init__(
+        self,
+        ctx: _TaskContext,
+    ):
+        self._ctx = ctx
+        self._reader = _BufferedGCSFileReader()
+        self._converter = _BufferedJsonConverter()
+        self._reader.set_retrieval_strategy(_BufferedFileRetrieverByOffset(self._ctx, 0))
+
+    def __iter__(self):
+        yield from self._fetch()
+
+    def _fetch(self,) -> Iterator[OutputT]:
+        """Fetches the result data.
+        Returns an iterator that yields the offset and the data.
+        """
+        if self._ctx.metadata_header.total_size is None:
+            return
+
+        stream = self._reader.read()
+        with self._converter as converter:
+            for file_info, raw_data in stream:
+                for output in converter.convert(
+                        Converter.ConverterInputArgs(self._ctx, file_info,
+                                                     raw_data)):
+                    yield output
+
+    def start(
+            self,
+            stream_handler: Optional[Callable[[OutputT], None]] = None) -> None:
+        """Starts streaming the result data.
+        Calls the stream_handler for each result.
+        """
+        # this calls the __iter__ method, which in turn calls the _fetch method
+        for output in self:
+            if stream_handler:
+                stream_handler(output)
+
+
+@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:
+                _, raw_data = result
+                # there is something wrong with the way the offsets are being calculated
+                # so just write all of the chunks as is too the file, with pointer initially
+                # pointed to the start of the file (like what is in GCS) and do not
+                # rely on offsets for file location
+                # 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 ExportTask:
     """
     An adapter class for working with task objects, providing extended functionality
@@ -649,11 +732,9 @@ def errors(self):
             self._task.client, self._task.uid, StreamType.ERRORS)
         if metadata_header is None:
             return None
-        Stream(
+        BufferedStream(
             _TaskContext(self._task.client, self._task.uid, StreamType.ERRORS,
                          metadata_header),
-            _BufferedGCSFileReader(),
-            _BufferedJsonConverter(),
         ).start(stream_handler=lambda output: data.append(output.json))
         return data
 
@@ -671,11 +752,9 @@ def result(self):
                 self._task.client, self._task.uid, StreamType.RESULT)
             if metadata_header is None:
                 return []
-            Stream(
+            BufferedStream(
                 _TaskContext(self._task.client, self._task.uid,
                              StreamType.RESULT, metadata_header),
-                _BufferedGCSFileReader(),
-                _BufferedJsonConverter(),
             ).start(stream_handler=lambda output: data.append(output.json))
             return data
         return self._task.result_url
@@ -767,11 +846,33 @@ def get_stream(
     ) -> Stream[FileConverterOutput]:
         """Overload for getting the right typing hints when using a FileConverter."""
 
+    def get_buffered_stream(
+        self,
+        stream_type: StreamType = StreamType.RESULT,
+    ) -> Stream:
+        """Returns the result of the task."""
+        if self._task.status == "FAILED":
+            raise ExportTask.ExportTaskException("Task failed")
+        if self._task.status != "COMPLETE":
+            raise ExportTask.ExportTaskException("Task is not ready yet")
+
+        metadata_header = self._get_metadata_header(self._task.client,
+                                                    self._task.uid, stream_type)
+        if metadata_header is None:
+            raise ValueError(
+                f"Task {self._task.uid} does not have a {stream_type.value} stream"
+            )
+        return BufferedStream(
+            _TaskContext(self._task.client, self._task.uid, stream_type,
+                         metadata_header),
+        )
+
     def get_stream(
         self,
         converter: Optional[Converter] = None,
         stream_type: StreamType = StreamType.RESULT,
     ) -> Stream:
+        warnings.warn("get_stream is deprecated and will be removed in a future release, use get_buffered_stream")
         if converter is None:
             converter = JsonConverter()
         """Returns the result of the task."""