zoopipe 2026.1.20__cp310-abi3-macosx_11_0_arm64.whl

zoopipe/engines/ray.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import os
+import re
+from datetime import datetime
+from importlib import metadata
+from typing import TYPE_CHECKING, Any
+
+import ray
+
+from zoopipe.engines.base import BaseEngine
+from zoopipe.engines.local import PipeReport
+from zoopipe.report import FlowReport, FlowStatus
+from zoopipe.utils.dependency import install_dependencies
+
+if TYPE_CHECKING:
+    from zoopipe.pipe import Pipe
+
+
+@ray.remote(memory=512 * 1024 * 1024)  # Limit actor memory to 512MB
+class RayPipeWorker:
+    """
+    Ray Actor that wraps a single Pipe execution.
+    """
+
+    def __init__(self, pipe: Pipe, index: int):
+        self.pipe = pipe
+        self.index = index
+        self.is_finished = False
+        self.has_error = False
+
+    def run(self) -> None:
+        try:
+            self.pipe.start(wait=True)
+        except Exception:
+            self.has_error = True
+        finally:
+            self.is_finished = True
+
+    def get_report(self) -> PipeReport:
+        report = self.pipe.report
+        return PipeReport(
+            pipe_index=self.index,
+            total_processed=report.total_processed,
+            success_count=report.success_count,
+            error_count=report.error_count,
+            ram_bytes=report.ram_bytes,
+            is_finished=self.is_finished or report.is_finished,
+            has_error=self.has_error,
+            is_alive=not self.is_finished,
+        )
+
+
+@ray.remote
+def _install_dependencies(packages: list[str]) -> None:
+    install_dependencies(packages)
+
+
+class RayEngine(BaseEngine):
+    """
+    Distributed execution engine using Ray.
+    """
+
+    def __init__(self, address: str | None = None, **kwargs: Any):
+        if not ray.is_initialized():
+            # Silence the accelerator visible devices warning for future Ray versions
+            os.environ.setdefault("RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO", "0")
+
+            # Prepare default runtime_env and get dependencies
+            runtime_env, deps = self._prepare_runtime_env(kwargs.pop("runtime_env", {}))
+
+            # Default to lean initialization
+            ray_args = {
+                "address": address,
+                "num_cpus": kwargs.pop("num_cpus", None),
+                "include_dashboard": kwargs.pop("include_dashboard", False),
+                "logging_level": kwargs.pop("logging_level", "error"),
+                "runtime_env": runtime_env,
+                **kwargs,
+            }
+            ray.init(**ray_args)
+
+            # Manually install dependencies on all nodes using our agnostic strategy
+            if deps:
+                self._install_deps_on_all_nodes(deps)
+
+        self._workers: list[Any] = []
+        self._futures: list[Any] = []
+        self._start_time: datetime | None = None
+        self._cached_report: FlowReport | None = None
+
+    def _install_deps_on_all_nodes(self, deps: list[str]) -> None:
+        """
+        Run the agnostic dependency installer on all connected Ray nodes.
+        This provides support for pip, uv, and poetry environments.
+        """
+        nodes = ray.nodes()
+        alive_nodes = [n for n in nodes if n.get("Alive")]
+
+        refs = []
+        for node in alive_nodes:
+            # We use resources placement group strategy to force execution
+            # on specific node. The 'node:<ip>' resource is automatically
+            # present on each node.
+            node_ip = node.get("NodeManagerAddress")
+            if node_ip:
+                refs.append(
+                    _install_dependencies.options(
+                        resources={f"node:{node_ip}": 0.001}
+                    ).remote(deps)
+                )
+
+        if refs:
+            ray.get(refs)
+
+    def _prepare_runtime_env(
+        self, runtime_env: dict[str, Any]
+    ) -> tuple[dict[str, Any], list[str]]:
+        """
+        Configure the Ray runtime environment based on whether we are in
+        development mode or being used as a library.
+        Returns modified runtime_env and a list of dependencies to install manually.
+        """
+        # 1. Detect environment and versions
+        is_dev_mode = False
+        try:
+            # heuristic: if we are in the zoopipe repo and have the ABI, it's dev mode
+            if (
+                os.path.exists("src/zoopipe")
+                and os.path.exists("pyproject.toml")
+                and any(f.endswith(".so") for f in os.listdir("src/zoopipe"))
+            ):
+                is_dev_mode = True
+        except Exception:
+            pass
+
+        # 2. Setup pip dependencies
+        deps = []
+        if "pip" not in runtime_env:
+            if is_dev_mode:
+                # Dev mode: Extract dependencies from pyproject.toml (Source of Truth)
+                try:
+                    with open("pyproject.toml", "r") as f:
+                        toml_content = f.read()
+                        # Find dependencies = [ ... ] block
+                        match = re.search(
+                            r"dependencies\s*=\s*\[(.*?)\]", toml_content, re.DOTALL
+                        )
+                        if match:
+                            dep_block = match.group(1)
+                            deps = re.findall(r'["\'](.*?)["\']', dep_block)
+                except Exception:
+                    pass
+            else:
+                # User mode: zoopipe package will pull its own dependencies
+                try:
+                    version = metadata.version("zoopipe")
+                    deps.append(f"zoopipe=={version}")
+                except metadata.PackageNotFoundError:
+                    # Fallback to hardcoded core if everything fails
+                    deps = ["pydantic>=2.0"]
+
+        # NOTE: We DO NOT set 'pip' in runtime_env because we want to use our
+        # agnostic installer.
+        # runtime_env["pip"] = deps  <-- REMOVED
+
+        # 3. Ship code and binaries
+        if "working_dir" not in runtime_env:
+            runtime_env["working_dir"] = "."
+
+            # In dev mode, we need src/ in PYTHONPATH to find the local zoopipe
+            if is_dev_mode:
+                env_vars = runtime_env.get("env_vars", {})
+                if "PYTHONPATH" not in env_vars:
+                    # Ray adds working_dir to sys.path,
+                    # but we need src/ for 'import zoopipe'
+                    env_vars["PYTHONPATH"] = "./src"
+                runtime_env["env_vars"] = env_vars
+
+        return runtime_env
+
+    def start(self, pipes: list[Pipe]) -> None:
+        if self.is_running:
+            raise RuntimeError("RayEngine is already running")
+
+        self._start_time = datetime.now()
+        self._workers = [RayPipeWorker.remote(pipe, i) for i, pipe in enumerate(pipes)]
+        self._futures = [w.run.remote() for w in self._workers]
+        self._cached_report = None
+
+    def wait(self, timeout: float | None = None) -> bool:
+        if not self._futures:
+            return True
+
+        ready, _ = ray.wait(
+            self._futures, num_returns=len(self._futures), timeout=timeout
+        )
+        return len(ready) == len(self._futures)
+
+    def shutdown(self, timeout: float = 5.0) -> None:
+        for worker in self._workers:
+            ray.kill(worker)
+        self._workers = []
+        self._futures = []
+        self._cached_report = None
+
+    @property
+    def is_running(self) -> bool:
+        if not self._futures:
+            return False
+        ready, _ = ray.wait(self._futures, num_returns=len(self._futures), timeout=0)
+        return len(ready) < len(self._futures)
+
+    @property
+    def report(self) -> FlowReport:
+        if self._cached_report and self._cached_report.is_finished:
+            return self._cached_report
+
+        report = FlowReport()
+        report.start_time = self._start_time
+
+        p_reports = self.pipe_reports
+        for pr in p_reports:
+            report.total_processed += pr.total_processed
+            report.success_count += pr.success_count
+            report.error_count += pr.error_count
+            report.ram_bytes += pr.ram_bytes
+
+        all_finished = all(pr.is_finished for pr in p_reports)
+        any_error = any(pr.has_error for pr in p_reports)
+
+        if all_finished:
+            report.status = FlowStatus.FAILED if any_error else FlowStatus.COMPLETED
+            report.end_time = datetime.now()
+            report._finished_event.set()
+            self._cached_report = report
+        else:
+            report.status = FlowStatus.RUNNING
+
+        return report
+
+    @property
+    def pipe_reports(self) -> list[PipeReport]:
+        if not self._workers:
+            return []
+        # Centralized collection from all actors in one pass
+        return ray.get([w.get_report.remote() for w in self._workers])
+
+    def get_pipe_report(self, index: int) -> PipeReport:
+        if not self._workers:
+            raise RuntimeError("Engine has not been started")
+        return ray.get(self._workers[index].get_report.remote())

zoopipe/hooks/__init__.py

@@ -0,0 +1,4 @@
+from zoopipe.hooks.base import BaseHook, HookPriority, HookStore
+from zoopipe.hooks.sql import SQLExpansionHook
+
+__all__ = ["BaseHook", "HookStore", "HookPriority", "SQLExpansionHook"]

zoopipe/hooks/base.py

@@ -0,0 +1,70 @@
+import typing
+
+from zoopipe.report import EntryTypedDict
+
+#: Type alias for the shared state between hooks.
+HookStore = dict[str, typing.Any]
+
+
+class HookPriority:
+    """
+    Standard priority levels for hooks.
+
+    Lower values correspond to higher priority (run earlier).
+    """
+
+    VERY_HIGH = 0
+    HIGH = 25
+    NORMAL = 50
+    LOW = 75
+    VERY_LOW = 100
+
+
+class BaseHook:
+    """
+    Abstract base class for pipeline lifecycle hooks.
+
+    Hooks allow executing custom Python logic at different stages of
+    the pipeline (setup, batch execution, and teardown). They can maintain
+    state between batches using the shared HookStore.
+    """
+
+    def __init__(self, priority: int = HookPriority.NORMAL):
+        """
+        Initialize the hook with a specific priority.
+
+        Args:
+            priority: Execution order (lower values run first).
+        """
+        self.priority = priority
+
+    def setup(self, store: HookStore) -> None:
+        """
+        Called once before the pipeline starts processing data.
+
+        Use this to initialize connections, resources, or shared state.
+        """
+        pass
+
+    def execute(
+        self, entries: list[EntryTypedDict], store: HookStore
+    ) -> list[EntryTypedDict]:
+        """
+        Process a batch of entries.
+
+        This method is where transformations or decorations happen.
+        It can modify the entries in-place or return a new list.
+
+        Args:
+            entries: List of dictionaries representing pipeline items.
+            store: Shared state between different hooks and different batches.
+        """
+        return entries
+
+    def teardown(self, store: HookStore) -> None:
+        """
+        Called once after the pipeline finish or if an error occurs.
+
+        Use this to release resources, close connections, or log final stats.
+        """
+        pass

zoopipe/hooks/sql.py

@@ -0,0 +1,94 @@
+import typing
+
+from zoopipe.hooks.base import BaseHook, HookStore
+from zoopipe.report import EntryStatus, get_logger
+
+if typing.TYPE_CHECKING:
+    from zoopipe.report import EntryTypedDict
+
+
+class SQLExpansionHook(BaseHook):
+    """
+    Expands anchor records (e.g., ID ranges) into full records by querying a SQL table.
+
+    This hook is designed to work with chunked data ingestion. It takes minimal
+    identifying information (anchors) and performs a bulk fetch from the database
+    to retrieve the complete rows.
+    """
+
+    def __init__(
+        self, connection_factory: typing.Callable[[], typing.Any], table_name: str
+    ):
+        """
+        Initialize the SQLExpansionHook.
+
+        Args:
+            connection_factory: Callable that returns a database connection.
+            table_name: Name of the SQL table to fetch data from.
+        """
+        super().__init__()
+        self.connection_factory = connection_factory
+        self.table_name = table_name
+        self.logger = get_logger()
+
+    def setup(self, store: HookStore) -> None:
+        pass
+
+    def execute(
+        self, entries: list["EntryTypedDict"], store: HookStore
+    ) -> list["EntryTypedDict"]:
+        expanded = []
+        conn = self.connection_factory()
+
+        try:
+            cursor = conn.cursor()
+            self.logger.debug(
+                f"SQLExpansionHook: Expanding batch of {len(entries)} anchor(s)"
+            )
+
+            for anchor in entries:
+                raw = anchor["raw_data"]
+                min_id = raw.get("min_id")
+                max_id = raw.get("max_id")
+
+                if min_id is None or max_id is None:
+                    continue
+
+                cursor.execute(
+                    f"SELECT * FROM {self.table_name} WHERE id BETWEEN ? AND ?",
+                    (min_id, max_id),
+                )
+
+                columns = (
+                    [column[0] for column in cursor.description]
+                    if cursor.description
+                    else []
+                )
+
+                rows = cursor.fetchall()
+
+                for row in rows:
+                    if columns:
+                        data = dict(zip(columns, row))
+                    else:
+                        data = dict(row)
+
+                    expanded.append(
+                        {
+                            "id": None,
+                            "position": None,
+                            "status": EntryStatus.PENDING,
+                            "raw_data": data,
+                            "validated_data": None,
+                            "metadata": anchor["metadata"],
+                            "errors": [],
+                        }
+                    )
+            cursor.close()
+        finally:
+            conn.close()
+
+        return expanded
+
+    def teardown(self, store: HookStore) -> None:
+        pass

zoopipe/input_adapter/__init__.py

@@ -0,0 +1,24 @@
+from zoopipe.input_adapter.arrow import ArrowInputAdapter
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.input_adapter.csv import CSVInputAdapter
+from zoopipe.input_adapter.duckdb import DuckDBInputAdapter
+from zoopipe.input_adapter.excel import ExcelInputAdapter
+from zoopipe.input_adapter.json import JSONInputAdapter
+from zoopipe.input_adapter.kafka import KafkaInputAdapter
+from zoopipe.input_adapter.parquet import ParquetInputAdapter
+from zoopipe.input_adapter.pygen import PyGeneratorInputAdapter
+from zoopipe.input_adapter.sql import SQLInputAdapter, SQLPaginationInputAdapter
+
+__all__ = [
+    "BaseInputAdapter",
+    "CSVInputAdapter",
+    "JSONInputAdapter",
+    "DuckDBInputAdapter",
+    "ArrowInputAdapter",
+    "ExcelInputAdapter",
+    "SQLInputAdapter",
+    "SQLPaginationInputAdapter",
+    "ParquetInputAdapter",
+    "PyGeneratorInputAdapter",
+    "KafkaInputAdapter",
+]

zoopipe/input_adapter/arrow.py

@@ -0,0 +1,38 @@
+import pathlib
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import ArrowReader
+
+
+class ArrowInputAdapter(BaseInputAdapter):
+    """
+    Reads records from Apache Arrow IPC (feather) files.
+
+    Provides high-speed sequential access to Arrow data with minimal
+    serialization overhead.
+    """
+
+    def __init__(
+        self,
+        source: typing.Union[str, pathlib.Path],
+        generate_ids: bool = True,
+    ):
+        """
+        Initialize the ArrowInputAdapter.
+
+        Args:
+            source: Path to the Arrow file.
+            generate_ids: Whether to generate unique IDs for each record.
+        """
+        self.source_path = str(source)
+        self.generate_ids = generate_ids
+
+    def get_native_reader(self) -> ArrowReader:
+        return ArrowReader(
+            self.source_path,
+            generate_ids=self.generate_ids,
+        )
+
+
+__all__ = ["ArrowInputAdapter"]

zoopipe/input_adapter/base.py

@@ -0,0 +1,48 @@
+import abc
+import typing
+
+
+class BaseInputAdapter(abc.ABC):
+    """
+    Abstract base class for all input adapters.
+
+    Input adapters are responsible for providing a native Rust reader
+    and optional hooks that are specific to the data source.
+    """
+
+    @abc.abstractmethod
+    def get_native_reader(self) -> typing.Any:
+        """
+        Return the underlying Rust reader instance.
+
+        This reader must implement the common reader interface in Rust
+        to be compatible with the NativePipe.
+        """
+        raise NotImplementedError
+
+    def get_hooks(self) -> list[typing.Any]:
+        """
+        Return a list of hooks to be executed by the pipeline.
+
+        Typically used for pre-fetching data or expanding anchor records
+        before they reach the main processing stage.
+        """
+        return []
+
+    @property
+    def can_split(self) -> bool:
+        """Return True if this adapter supports parallel splitting."""
+        return type(self).split != BaseInputAdapter.split
+
+    def split(self, workers: int) -> typing.List["BaseInputAdapter"]:
+        """
+        Split the input adapter into `workers` shards for parallel processing.
+
+        Args:
+            workers: Number of partitions to create.
+
+        Returns:
+            A list of input adapters, each responsible for a subset of the data.
+            Default implementation returns [self] (no splitting).
+        """
+        return [self]

zoopipe/input_adapter/csv.py

@@ -0,0 +1,144 @@
+import csv
+import pathlib
+import typing
+
+from zoopipe.input_adapter.base import BaseInputAdapter
+from zoopipe.zoopipe_rust_core import CSVReader, get_file_size
+
+
+class CSVInputAdapter(BaseInputAdapter):
+    """
+    A high-performance CSV reader supporting both local and S3 sources.
+
+    Uses a multi-threaded parser in the Rust core to ensure fast data ingestion
+    without blocking the Python GIL.
+    """
+
+    def __init__(
+        self,
+        source: typing.Union[str, pathlib.Path],
+        delimiter: str = ",",
+        quotechar: str = '"',
+        skip_rows: int = 0,
+        fieldnames: list[str] | None = None,
+        generate_ids: bool = True,
+        limit: int | None = None,
+        start_byte: int = 0,
+        end_byte: int | None = None,
+    ):
+        """
+        Initialize the CSVInputAdapter.
+
+        Args:
+            source: Path to the CSV file or S3 URI.
+            delimiter: Column separator.
+            quotechar: Character used for quoting fields.
+            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.
+            limit: Maximum number of rows to read (optional).
+            start_byte: Byte offset to start reading from.
+            end_byte: Byte offset to stop reading at.
+        """
+        self.source_path = str(source)
+        self.delimiter = delimiter
+        self.quotechar = quotechar
+        self.skip_rows = skip_rows
+        self.fieldnames = fieldnames
+        self.generate_ids = generate_ids
+        self.limit = limit
+        self.start_byte = start_byte
+        self.end_byte = end_byte
+
+    def split(self, workers: int) -> typing.List["CSVInputAdapter"]:
+        """
+        Split the CSV input into `workers` byte-range shards.
+        """
+
+        file_size = get_file_size(self.source_path)
+
+        chunk_size = file_size // workers
+
+        # Ensure we have fieldnames if not explicitly provided
+        # This is CRITICAL for partial reads (start_byte > 0)
+        final_fieldnames = self.fieldnames
+        if final_fieldnames is None:
+            if self.source_path.startswith("s3://"):
+                # Use Rust reader to discover headers from S3
+                final_fieldnames = self.get_native_reader().headers
+            else:
+                with open(self.source_path, "r") as f:
+                    reader = csv.reader(
+                        f, delimiter=self.delimiter, quotechar=self.quotechar
+                    )
+                    try:
+                        final_fieldnames = next(reader)
+                    except StopIteration:
+                        final_fieldnames = []
+
+        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,
+                    delimiter=self.delimiter,
+                    quotechar=self.quotechar,
+                    skip_rows=self.skip_rows,
+                    fieldnames=final_fieldnames,
+                    generate_ids=self.generate_ids,
+                    limit=self.limit,
+                    start_byte=start,
+                    end_byte=end,
+                )
+            )
+        return shards
+
+    def get_native_reader(self) -> CSVReader:
+        # Pass start_byte and end_byte
+        return CSVReader(
+            self.source_path,
+            delimiter=ord(self.delimiter),
+            quote=ord(self.quotechar),
+            skip_rows=self.skip_rows,
+            fieldnames=self.fieldnames,
+            generate_ids=self.generate_ids,
+            limit=self.limit,
+            start_byte=self.start_byte,
+            end_byte=self.end_byte,
+        )
+
+    @staticmethod
+    def count_rows(
+        source: str | pathlib.Path,
+        delimiter: str = ",",
+        quotechar: str = '"',
+        has_header: bool = True,
+    ) -> int:
+        """
+        Efficiently count the number of rows in a CSV file using the Rust core.
+
+        Args:
+            source: Path to the CSV file.
+            delimiter: Column separator. (Default: ',')
+            quotechar: Character used for quoting. (Default: '"')
+            has_header: Whether the file has a header row to ignore in user count
+                (if used in context where header matters, though CSVReader.count_rows
+                name implies all records).
+                Actually pass this to rust to decide if first row is record or header.
+
+        Returns:
+            Number of rows (records).
+        """
+        return CSVReader.count_rows(
+            str(source),
+            ord(delimiter),
+            ord(quotechar),
+            has_header,
+        )
+
+
+__all__ = ["CSVInputAdapter"]