feat: ability to use queue chunks

Author: ringoldsdev

laygo/__init__.py

@@ -14,7 +14,9 @@
 from laygo.transformers.threaded import ThreadedTransformer
 from laygo.transformers.threaded import createThreadedTransformer
 from laygo.transformers.transformer import Transformer
+from laygo.transformers.transformer import build_chunk_generator
 from laygo.transformers.transformer import createTransformer
+from laygo.transformers.transformer import passthrough_chunks
 
 __all__ = [
   "Pipeline",
@@ -28,4 +30,6 @@
   "createHTTPTransformer",
   "PipelineContext",
   "ErrorHandler",
+  "passthrough_chunks",
+  "build_chunk_generator",
 ]

laygo/pipeline.py

@@ -14,6 +14,7 @@
 from laygo.helpers import PipelineContext
 from laygo.helpers import is_context_aware
 from laygo.transformers.transformer import Transformer
+from laygo.transformers.transformer import passthrough_chunks
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -154,6 +155,7 @@ def branch(
     branches: dict[str, Transformer[T, Any]],
     batch_size: int = 1000,
     max_batch_buffer: int = 1,
+    use_queue_chunks: bool = True,
   ) -> dict[str, list[Any]]:
     """Forks the pipeline into multiple branches for concurrent, parallel processing."""
     if not branches:
@@ -185,7 +187,10 @@ def consumer(transformer: Transformer, queue: Queue) -> list[Any]:
 
       def stream_from_queue() -> Iterator[T]:
         while (batch := queue.get()) is not None:
-          yield from batch
+          yield batch
+
+      if use_queue_chunks:
+        transformer = transformer.set_chunker(passthrough_chunks)
 
       result_iterator = transformer(stream_from_queue(), self.ctx)  # type: ignore
       return list(result_iterator)

laygo/transformers/transformer.py

@@ -60,6 +60,20 @@ def chunk_generator(data: Iterable[T]) -> Iterator[list[T]]:
   return chunk_generator
 
 
+def passthrough_chunks[T](data: Iterable[list[T]]) -> Iterator[list[T]]:
+  """A chunk generator that yields the entire input as a single chunk.
+
+  This is useful for transformers that do not require chunking.
+
+  Args:
+      data: The input data to process.
+
+  Returns:
+      An iterator yielding the entire input as a single chunk.
+  """
+  yield from iter(data)
+
+
 class Transformer[In, Out]:
   """Define and compose data transformations by passing context explicitly.