zoopipe 2026.1.20__cp310-abi3-macosx_11_0_arm64.whl

zoopipe/input_adapter/duckdb.py

@@ -0,0 +1,54 @@
+import pathlib
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import DuckDBReader
+
+
+class DuckDBInputAdapter(BaseInputAdapter):
+    """
+    Executes SQL queries against DuckDB database files.
+
+    Directly interfaces with DuckDB to stream query results, enabling
+    efficient processing of large datasets stored in analytical databases.
+    """
+
+    def __init__(
+        self,
+        source: typing.Union[str, pathlib.Path],
+        query: str | None = None,
+        table_name: str | None = None,
+        generate_ids: bool = True,
+    ):
+        """
+        Initialize the DuckDBInputAdapter.
+
+        Args:
+            source: Path to the DuckDB database file.
+            query: SQL query to execute.
+            table_name: Or name of the table to read (equiv to SELECT * FROM table).
+            generate_ids: Whether to generate unique IDs for each record.
+        """
+        self.source_path = str(source)
+        self.generate_ids = generate_ids
+
+        if query is None and table_name is None:
+            raise ValueError("Either query or table_name must be provided")
+
+        if query is not None and table_name is not None:
+            raise ValueError("Only one of query or table_name should be provided")
+
+        if query is not None:
+            self.query = query
+        else:
+            self.query = f"SELECT * FROM {table_name}"
+
+    def get_native_reader(self) -> DuckDBReader:
+        return DuckDBReader(
+            self.source_path,
+            self.query,
+            generate_ids=self.generate_ids,
+        )
+
+
+__all__ = ["DuckDBInputAdapter"]

zoopipe/input_adapter/excel.py

@@ -0,0 +1,51 @@
+import pathlib
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import ExcelReader
+
+
+class ExcelInputAdapter(BaseInputAdapter):
+    """
+    Reads Excel files (.xlsx, .xls, .ods, .xlsb) using the Calamine engine.
+
+    Provides high-performance, memory-efficient parsing of various spreadsheet
+    formats directly from Rust. Supports sheet selection by name or index,
+    skipping header rows, and custom field mapping.
+    """
+
+    def __init__(
+        self,
+        source: typing.Union[str, pathlib.Path],
+        sheet: typing.Union[str, int, None] = None,
+        skip_rows: int = 0,
+        fieldnames: typing.Optional[typing.List[str]] = None,
+        generate_ids: bool = True,
+    ):
+        """
+        Initialize the ExcelInputAdapter.
+
+        Args:
+            source: Path to the Excel file.
+            sheet: Sheet name (str) or index (int) to read. Defaults to the first sheet.
+            skip_rows: Number of rows to skip at the beginning.
+            fieldnames: Optional list of column names.
+            generate_ids: Whether to generate unique IDs for each record.
+        """
+        self.source_path = str(source)
+        self.sheet = sheet
+        self.skip_rows = skip_rows
+        self.fieldnames = fieldnames
+        self.generate_ids = generate_ids
+
+    def get_native_reader(self) -> ExcelReader:
+        return ExcelReader(
+            self.source_path,
+            sheet=self.sheet,
+            skip_rows=self.skip_rows,
+            fieldnames=self.fieldnames,
+            generate_ids=self.generate_ids,
+        )
+
+
+__all__ = ["ExcelInputAdapter"]

zoopipe/input_adapter/json.py

@@ -0,0 +1,73 @@
+import pathlib
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import JSONReader
+
+
+class JSONInputAdapter(BaseInputAdapter):
+    """
+    Reads data from JSON or JSONLines (.jsonl) files.
+
+    It supports both standard JSON arrays and line-delimited records.
+    The adapter uses a fast Rust-based parser that streams data efficiently,
+    making it suitable for very large datasets.
+    """
+
+    def __init__(
+        self,
+        source: typing.Union[str, pathlib.Path],
+        start_byte: int = 0,
+        end_byte: int | None = None,
+    ):
+        """
+        Initialize the JSONInputAdapter.
+
+        Args:
+            source: Path to the JSONLines file.
+            start_byte: Byte offset to start reading from.
+            end_byte: Byte offset to stop reading at.
+        """
+        self.source_path = str(source)
+        self.start_byte = start_byte
+        self.end_byte = end_byte
+
+    @property
+    def can_split(self) -> bool:
+        """Only allow splitting for JSONLines/NDJSON formats."""
+        path = self.source_path.lower()
+        return path.endswith(".jsonl") or path.endswith(".ndjson")
+
+    def split(self, workers: int) -> typing.List["JSONInputAdapter"]:
+        """
+        Split the JSON input into `workers` byte-range shards.
+        """
+        from zoopipe.zoopipe_rust_core import get_file_size
+
+        file_size = get_file_size(self.source_path)
+
+        chunk_size = file_size // workers
+        shards = []
+        for i in range(workers):
+            start = i * chunk_size
+            # Last worker takes rest of file
+            end = (i + 1) * chunk_size if i < workers - 1 else None
+
+            shards.append(
+                self.__class__(
+                    source=self.source_path,
+                    start_byte=start,
+                    end_byte=end,
+                )
+            )
+        return shards
+
+    def get_native_reader(self) -> JSONReader:
+        return JSONReader(
+            self.source_path,
+            start_byte=self.start_byte,
+            end_byte=self.end_byte,
+        )
+
+
+__all__ = ["JSONInputAdapter"]

zoopipe/input_adapter/kafka.py

@@ -0,0 +1,39 @@
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import KafkaReader
+
+
+class KafkaInputAdapter(BaseInputAdapter):
+    """
+    Consumes messages from Apache Kafka topics.
+
+    Acts as a Kafka consumer, streaming messages into the pipeline with
+    support for consumer groups and offset management.
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        group_id: str | None = None,
+        generate_ids: bool = True,
+    ):
+        """
+        Kafka Input Adapter.
+
+        Args:
+            uri: Kafka URI (e.g., 'kafka://localhost:9092/topic')
+            group_id: Optional consumer group ID.
+            generate_ids: Whether to generate unique IDs for each message.
+        """
+        self.uri = uri
+        self.group_id = group_id
+        self.generate_ids = generate_ids
+
+    def get_native_reader(self) -> KafkaReader:
+        return KafkaReader(
+            self.uri,
+            group_id=self.group_id,
+            generate_ids=self.generate_ids,
+        )
+
+
+__all__ = ["KafkaInputAdapter"]

zoopipe/input_adapter/parquet.py

@@ -0,0 +1,85 @@
+import pathlib
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import ParquetReader
+
+
+class ParquetInputAdapter(BaseInputAdapter):
+    """
+    Reads records from Apache Parquet files.
+
+    Utilizes the Arrow ecosystem for efficient columnar data reading and
+    multi-threaded loading.
+    """
+
+    def __init__(
+        self,
+        source: typing.Union[str, pathlib.Path],
+        generate_ids: bool = True,
+        batch_size: int = 1024,
+        limit: int | None = None,
+        offset: int = 0,
+        row_groups: typing.List[int] | None = None,
+    ):
+        """
+        Initialize the ParquetInputAdapter.
+
+        Args:
+            source: Path to the Parquet file.
+            generate_ids: Whether to generate unique IDs for each record.
+            batch_size: Number of records to read at once from the file.
+            limit: Maximum number of rows to read.
+            offset: Number of rows to skip.
+        """
+        self.source_path = str(source)
+        self.generate_ids = generate_ids
+        self.batch_size = batch_size
+        self.limit = limit
+        self.offset = offset
+        self.row_groups = row_groups
+
+    def split(self, workers: int) -> typing.List["ParquetInputAdapter"]:
+        """
+        Split the Parquet input into `workers` shards based on Row Groups.
+        """
+        row_group_rows = ParquetReader.get_row_groups_info(self.source_path)
+        num_groups = len(row_group_rows)
+
+        if num_groups < workers:
+            workers = num_groups
+
+        if workers <= 1:
+            return [self]
+
+        # Distribute row groups among workers
+        groups_per_worker = num_groups // workers
+        shards = []
+        for i in range(workers):
+            start_idx = i * groups_per_worker
+            end_idx = (i + 1) * groups_per_worker if i < workers - 1 else num_groups
+
+            assigned_groups = list(range(start_idx, end_idx))
+
+            shards.append(
+                self.__class__(
+                    source=self.source_path,
+                    generate_ids=self.generate_ids,
+                    batch_size=self.batch_size,
+                    row_groups=assigned_groups,
+                )
+            )
+        return shards
+
+    def get_native_reader(self) -> ParquetReader:
+        return ParquetReader(
+            self.source_path,
+            generate_ids=self.generate_ids,
+            batch_size=self.batch_size,
+            limit=self.limit,
+            offset=self.offset,
+            row_groups=self.row_groups,
+        )
+
+
+__all__ = ["ParquetInputAdapter"]

zoopipe/input_adapter/pygen.py

@@ -0,0 +1,37 @@
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import PyGeneratorReader
+
+
+class PyGeneratorInputAdapter(BaseInputAdapter):
+    """
+    Bridges Python iterables and generators into the pipeline.
+
+    Allows using any custom Python logic or in-memory data as a source
+    for the pipeline.
+    """
+
+    def __init__(
+        self,
+        iterable: typing.Iterable[typing.Any],
+        generate_ids: bool = True,
+    ):
+        """
+        Initialize the PyGeneratorInputAdapter.
+
+        Args:
+            iterable: Any Python iterable or generator yielding dictionaries.
+            generate_ids: Whether to generate unique IDs for each record.
+        """
+        self.iterable = iterable
+        self.generate_ids = generate_ids
+
+    def get_native_reader(self) -> PyGeneratorReader:
+        return PyGeneratorReader(
+            self.iterable,
+            generate_ids=self.generate_ids,
+        )
+
+
+__all__ = ["PyGeneratorInputAdapter"]

zoopipe/input_adapter/sql.py

@@ -0,0 +1,103 @@
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import SQLReader
+
+
+class SQLInputAdapter(BaseInputAdapter):
+    """
+    Streams records from SQL databases using standard queries.
+
+    Supports any database compatible with SQLAlchemy URIs. It executes a
+    provided query or fetches a whole table using a native Rust executor
+    for optimal performance.
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        query: str | None = None,
+        table_name: str | None = None,
+        generate_ids: bool = True,
+    ):
+        """
+        Initialize the SQLInputAdapter.
+
+        Args:
+            uri: Database URI (e.g., 'sqlite:///data.db').
+            query: SQL query to execute.
+            table_name: Or name of the table to read (equiv to SELECT * FROM table).
+            generate_ids: Whether to generate unique IDs for each record.
+        """
+        self.uri = uri
+        self.generate_ids = generate_ids
+
+        if query is None and table_name is None:
+            raise ValueError("Either query or table_name must be provided")
+
+        if query is not None and table_name is not None:
+            raise ValueError("Only one of query or table_name should be provided")
+
+        if query is not None:
+            self.query = query
+        else:
+            self.query = f"SELECT * FROM {table_name}"
+
+    def get_native_reader(self) -> SQLReader:
+        return SQLReader(
+            self.uri,
+            self.query,
+            generate_ids=self.generate_ids,
+        )
+
+
+class SQLPaginationInputAdapter(SQLInputAdapter):
+    """
+    Input adapter for SQL databases using anchor-based pagination.
+
+    This adapter generates ID ranges (anchors) and utilizes SQLExpansionHook
+    to fetch full records in chunks, which is more efficient for very large tables.
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        table_name: str,
+        id_column: str,
+        chunk_size: int,
+        connection_factory: typing.Callable[[], typing.Any],
+    ):
+        """
+        Initialize the SQLPaginationInputAdapter.
+
+        Args:
+            uri: Database URI.
+            table_name: Name of the table to read.
+            id_column: Primary key or indexed column used for pagination.
+            chunk_size: Number of records to fetch per chunk.
+            connection_factory: Callable that returns a database connection
+                for the hook.
+        """
+        self.table_name = table_name
+        self.id_column = id_column
+        self.chunk_size = chunk_size
+        self.connection_factory = connection_factory
+
+        query = f"""
+        WITH RECURSIVE ranges(n) AS (
+            SELECT MIN({id_column}) FROM {table_name}
+            UNION ALL
+            SELECT n + {chunk_size} FROM ranges 
+            WHERE n + {chunk_size} <= (SELECT MAX({id_column}) FROM {table_name})
+        )
+        SELECT n as min_id, n + {chunk_size} - 1 as max_id FROM ranges
+        """
+        super().__init__(uri, query=query)
+
+    def get_hooks(self):
+        from zoopipe.hooks.sql import SQLExpansionHook
+
+        return [SQLExpansionHook(self.connection_factory, self.table_name)]
+
+
+__all__ = ["SQLInputAdapter", "SQLPaginationInputAdapter"]

zoopipe/manager.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import os
+import shutil
+from typing import TYPE_CHECKING, Any
+
+from zoopipe.engines import MultiProcessEngine
+from zoopipe.engines.local import PipeReport
+from zoopipe.zoopipe_rust_core import MultiThreadExecutor, SingleThreadExecutor
+
+if TYPE_CHECKING:
+    from zoopipe.engines.base import BaseEngine
+    from zoopipe.pipe import Pipe
+    from zoopipe.report import FlowReport
+
+
+class PipeManager:
+    """
+    Manages one or more Pipes using an execution Engine.
+
+    PipeManager acts as the high-level orchestrator. It handles the sharding
+    of data sources across multiple workers and coordinates their execution
+    through a pluggable Engine (e.g., Local Multiprocessing, Ray, Dask).
+    """
+
+    def __init__(self, pipes: list[Pipe], engine: BaseEngine | None = None):
+        """
+        Initialize PipeManager with a list of Pipe instances.
+
+        Args:
+            pipes: List of Pipe objects to manage.
+            engine: Optional execution engine. Defaults to MultiProcessEngine.
+        """
+        self.pipes = pipes
+        self.engine = engine or MultiProcessEngine()
+        self._merge_info: dict[str, Any] = {}
+        self.should_merge = False
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the execution is currently running."""
+        return self.engine.is_running
+
+    @property
+    def pipe_count(self) -> int:
+        """Get the number of pipes being managed."""
+        return len(self.pipes)
+
+    def start(self) -> None:
+        """
+        Start all managed pipes using the configured engine.
+        """
+        self.engine.start(self.pipes)
+
+    def wait(self, timeout: float | None = None) -> bool:
+        """
+        Wait for execution to finish.
+
+        Args:
+            timeout: Optional maximum time to wait.
+        Returns:
+            True if execution finished.
+        """
+        return self.engine.wait(timeout)
+
+    def shutdown(self, timeout: float = 5.0) -> None:
+        """
+        Forcibly stop all running pipes.
+
+        Args:
+            timeout: Maximum time to wait for termination.
+        """
+        self.engine.shutdown(timeout)
+
+    @property
+    def report(self) -> FlowReport:
+        """Get an aggregated report of all running pipes."""
+        return self.engine.report
+
+    def get_pipe_report(self, index: int) -> PipeReport:
+        """
+        Get the current report for a specific pipe.
+
+        Args:
+            index: The index of the pipe in the original list.
+        """
+        if hasattr(self.engine, "get_pipe_report"):
+            return self.engine.get_pipe_report(index)
+        raise AttributeError(
+            f"Engine {self.engine.__class__.__name__} does not support per-pipe reports"
+        )
+
+    @property
+    def pipe_reports(self) -> list[PipeReport]:
+        """Get reports for all managed pipes."""
+        if hasattr(self.engine, "pipe_reports"):
+            return self.engine.pipe_reports
+        # Fallback if the engine doesn't have the property but has the method
+        if hasattr(self.engine, "get_pipe_report"):
+            return [self.engine.get_pipe_report(i) for i in range(self.pipe_count)]
+        raise AttributeError(
+            f"Engine {self.engine.__class__.__name__} does not support per-pipe reports"
+        )
+
+    def __enter__(self) -> PipeManager:
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.is_running:
+            self.shutdown()
+
+    @classmethod
+    def parallelize_pipe(
+        cls,
+        pipe: Pipe,
+        workers: int,
+        should_merge: bool = False,
+        executor: SingleThreadExecutor | MultiThreadExecutor | None = None,
+        engine: BaseEngine | None = None,
+    ) -> PipeManager:
+        """
+        Create a PipeManager that runs the given pipe in parallel across
+        `workers` shards.
+
+        Automatically splits the input and output adapters to ensure safe
+        parallel execution.
+
+        Args:
+            pipe: The source pipe to parallelize.
+            workers: Number of shards to use.
+            should_merge: Whether to merge the output shards automatically.
+            executor: Internal batch executor for each shard.
+            engine: Optional execution engine.
+
+        Returns:
+            A configured PipeManager instance.
+        """
+        if not pipe.input_adapter.can_split or not pipe.output_adapter.can_split:
+            workers = 1
+
+        input_shards = pipe.input_adapter.split(workers)
+        output_shards = pipe.output_adapter.split(workers)
+
+        if len(input_shards) != workers or len(output_shards) != workers:
+            raise ValueError(
+                f"Adapters failed to split into {workers} shards. "
+                f"Got {len(input_shards)} inputs and {len(output_shards)} outputs."
+            )
+
+        exec_strategy = executor or pipe.executor
+
+        pipes = []
+        for i in range(workers):
+            sharded_pipe = type(pipe)(
+                input_adapter=input_shards[i],
+                output_adapter=output_shards[i],
+                schema_model=pipe.schema_model,
+                pre_validation_hooks=pipe.pre_validation_hooks,
+                post_validation_hooks=pipe.post_validation_hooks,
+                report_update_interval=pipe.report_update_interval,
+                executor=exec_strategy,
+            )
+            pipes.append(sharded_pipe)
+
+        manager = cls(pipes, engine=engine)
+        manager.should_merge = should_merge
+        manager._merge_info = {
+            "target": getattr(pipe.output_adapter, "output_path", None),
+            "sources": [getattr(shard, "output_path", None) for shard in output_shards],
+        }
+        return manager
+
+    def merge(self) -> None:
+        """
+        Merge the output files from all pipes into the final destination.
+        """
+        if not self._should_merge():
+            return
+
+        target = self._merge_info["target"]
+        sources = [s for s in self._merge_info["sources"] if s and os.path.exists(s)]
+
+        with open(target, "wb") as dest:
+            for src_path in sources:
+                with open(src_path, "rb") as src:
+                    self._append_file(dest, src)
+
+    def _should_merge(self) -> bool:
+        if not self.should_merge or not self._merge_info.get("target"):
+            return False
+        sources = [s for s in self._merge_info.get("sources", []) if s]
+        return len(sources) > 1
+
+    def _append_file(self, dest, src) -> None:
+        """Append file content using zero-copy where available."""
+        try:
+            offset, size = 0, os.fstat(src.fileno()).st_size
+            while offset < size:
+                sent = os.sendfile(dest.fileno(), src.fileno(), offset, size - offset)
+                if sent == 0:
+                    break
+                offset += sent
+        except (OSError, AttributeError):
+            src.seek(0)
+            shutil.copyfileobj(src, dest)
+
+    def __repr__(self) -> str:
+        status = "running" if self.is_running else "stopped"
+        return f"<PipeManager pipes={self.pipe_count} status={status} "
+        f"engine={self.engine.__class__.__name__}>"

zoopipe/output_adapter/__init__.py

@@ -0,0 +1,23 @@
+from zoopipe.output_adapter.arrow import ArrowOutputAdapter
+from zoopipe.output_adapter.base import BaseOutputAdapter
+from zoopipe.output_adapter.csv import CSVOutputAdapter
+from zoopipe.output_adapter.duckdb import DuckDBOutputAdapter
+from zoopipe.output_adapter.excel import ExcelOutputAdapter
+from zoopipe.output_adapter.json import JSONOutputAdapter
+from zoopipe.output_adapter.kafka import KafkaOutputAdapter
+from zoopipe.output_adapter.parquet import ParquetOutputAdapter
+from zoopipe.output_adapter.pygen import PyGeneratorOutputAdapter
+from zoopipe.output_adapter.sql import SQLOutputAdapter
+
+__all__ = [
+    "BaseOutputAdapter",
+    "CSVOutputAdapter",
+    "JSONOutputAdapter",
+    "DuckDBOutputAdapter",
+    "ArrowOutputAdapter",
+    "ExcelOutputAdapter",
+    "SQLOutputAdapter",
+    "ParquetOutputAdapter",
+    "PyGeneratorOutputAdapter",
+    "KafkaOutputAdapter",
+]