neo4j-etl-lib 0.3.1__tar.gz → 0.3.2__tar.gz

{neo4j_etl_lib-0.3.1 → neo4j_etl_lib-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-etl-lib
-Version: 0.3.1
+Version: 0.3.2
 Summary: Building blocks for ETL pipelines.
 Keywords: etl,graph,database
 Author-email: Bert Radke <bert.radke@pm.me>

{neo4j_etl_lib-0.3.1 → neo4j_etl_lib-0.3.2}/src/etl_lib/__init__.py

@@ -1,4 +1,4 @@
 """
 Building blocks for ETL pipelines.
 """
-__version__ = "0.3.1"
+__version__ = "0.3.2"

{neo4j_etl_lib-0.3.1 → neo4j_etl_lib-0.3.2}/src/etl_lib/core/ETLContext.py

@@ -172,15 +172,8 @@ if sqlalchemy_available:
                 database_url,
                 pool_pre_ping=True,
                 pool_size=pool_size,
-                max_overflow=max_overflow,
-                pool_recycle=1800,  # recycle connections older than 30m
-                connect_args={
-                    # turn on TCP keepalives on the client socket:
-                    "keepalives": 1,
-                    "keepalives_idle": 60,  # after 60s of idle
-                    "keepalives_interval": 10,  # probe every 10s
-                    "keepalives_count": 5,  # give up after 5 failed probes
-                })
+                max_overflow=max_overflow
+            )
 
 
 class ETLContext:

neo4j_etl_lib-0.3.2/src/etl_lib/core/ParallelBatchProcessor.py

@@ -0,0 +1,150 @@
+import queue
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Generator, List
+
+from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults
+from etl_lib.core.utils import merge_summery
+
+
+class ParallelBatchResult(BatchResults):
+    """
+    Represents one *wave* produced by the splitter.
+
+    `chunk` is a list of bucket-batches. Each sub-list is processed by one worker instance.
+    """
+    pass
+
+
+class ParallelBatchProcessor(BatchProcessor):
+    """
+    BatchProcessor that runs a worker over the bucket-batches of each ParallelBatchResult
+    in parallel threads, while prefetching the next ParallelBatchResult from its predecessor.
+
+    Note:
+        - The predecessor must emit `ParallelBatchResult` instances (waves).
+        - This processor collects the BatchResults from all workers for one wave and merges them
+          into one BatchResults.
+        - The returned BatchResults will not obey the max_batch_size from get_batch() because
+          it represents the full wave.
+
+    Args:
+        context: ETL context.
+        worker_factory: A zero-arg callable that returns a new BatchProcessor each time it's called.
+        task: optional Task for reporting.
+        predecessor: upstream BatchProcessor that must emit ParallelBatchResult See
+                :class:`~etl_lib.core.SplittingBatchProcessor.SplittingBatchProcessor`.
+        max_workers: number of parallel threads for bucket processing.
+        prefetch: number of waves to prefetch.
+
+    Behavior:
+        - For every wave, spins up `max_workers` threads.
+        - Each thread processes one bucket-batch using a fresh worker from `worker_factory()`.
+        - Collects and merges worker results in a fail-fast manner.
+    """
+
+    def __init__(
+            self,
+            context,
+            worker_factory: Callable[[], BatchProcessor],
+            task=None,
+            predecessor=None,
+            max_workers: int = 4,
+            prefetch: int = 4,
+    ):
+        super().__init__(context, task, predecessor)
+        self.worker_factory = worker_factory
+        self.max_workers = max_workers
+        self.prefetch = prefetch
+
+    def _process_wave(self, wave: ParallelBatchResult) -> BatchResults:
+        """
+        Process one wave: run one worker per bucket-batch and merge their BatchResults.
+
+        Statistics:
+            `wave.statistics` is used as the initial merged stats, then merged with each worker's stats.
+        """
+        merged_stats = dict(wave.statistics or {})
+        merged_chunk = []
+        total = 0
+
+        self.logger.debug(f"Processing wave with {len(wave.chunk)} buckets")
+
+        with ThreadPoolExecutor(max_workers=self.max_workers, thread_name_prefix="PBP_worker_") as pool:
+            futures = [pool.submit(self._process_bucket_batch, bucket_batch) for bucket_batch in wave.chunk]
+            try:
+                for f in as_completed(futures):
+                    out = f.result()
+                    merged_stats = merge_summery(merged_stats, out.statistics or {})
+                    total += out.batch_size
+                    merged_chunk.extend(out.chunk if isinstance(out.chunk, list) else [out.chunk])
+            except Exception as e:
+                for g in futures:
+                    g.cancel()
+                pool.shutdown(cancel_futures=True)
+                raise RuntimeError("bucket processing failed") from e
+
+        self.logger.debug(f"Finished wave with stats={merged_stats}")
+        return BatchResults(chunk=merged_chunk, statistics=merged_stats, batch_size=total)
+
+    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
+        """
+        Pull waves from the predecessor (prefetching up to `prefetch` ahead), process each wave's
+        buckets in parallel, and yield one flattened BatchResults per wave.
+        """
+        wave_queue: queue.Queue[ParallelBatchResult | object] = queue.Queue(self.prefetch)
+        SENTINEL = object()
+        exc: BaseException | None = None
+
+        def producer():
+            nonlocal exc
+            try:
+                for wave in self.predecessor.get_batch(max_batch_size):
+                    self.logger.debug(
+                        f"adding wave stats={wave.statistics} buckets={len(wave.chunk)} to queue size={wave_queue.qsize()}"
+                    )
+                    wave_queue.put(wave)
+            except BaseException as e:
+                exc = e
+            finally:
+                wave_queue.put(SENTINEL)
+
+        threading.Thread(target=producer, daemon=True, name="prefetcher").start()
+
+        while True:
+            wave = wave_queue.get()
+            if wave is SENTINEL:
+                if exc is not None:
+                    self.logger.error("Upstream producer failed", exc_info=True)
+                    raise exc
+                break
+            yield self._process_wave(wave)
+
+    class SingleBatchWrapper(BatchProcessor):
+        """
+        Simple BatchProcessor that returns exactly one batch (the bucket-batch passed in via init).
+        Used as predecessor for the per-bucket worker.
+        """
+
+        def __init__(self, context, batch: List[Any]):
+            super().__init__(context=context, predecessor=None)
+            self._batch = batch
+
+        def get_batch(self, max_size: int) -> Generator[BatchResults, None, None]:
+            yield BatchResults(
+                chunk=self._batch,
+                statistics={},
+                batch_size=len(self._batch),
+            )
+
+    def _process_bucket_batch(self, bucket_batch):
+        """
+        Process one bucket-batch by running a fresh worker over it.
+        """
+        self.logger.debug(f"Processing batch w/ size {len(bucket_batch)}")
+        wrapper = self.SingleBatchWrapper(self.context, bucket_batch)
+        worker = self.worker_factory()
+        worker.predecessor = wrapper
+        result = next(worker.get_batch(len(bucket_batch)))
+        self.logger.debug(f"Finished bucket batch stats={result.statistics}")
+        return result

neo4j_etl_lib-0.3.2/src/etl_lib/core/SplittingBatchProcessor.py

@@ -0,0 +1,391 @@
+import logging
+from typing import Any, Callable, Dict, Generator, List, Tuple
+
+from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults
+from etl_lib.core.utils import merge_summery
+from tabulate import tabulate
+
+
+def tuple_id_extractor(table_size: int = 10) -> Callable[[Tuple[str | int, str | int]], Tuple[int, int]]:
+    """
+    Create an ID extractor function for tuple items, using the last decimal digit of each element.
+    The output is a `(row, col)` tuple within a `table_size x table_size` grid (default 10x10).
+
+    Args:
+        table_size: The dimension of the grid (number of rows/cols). Defaults to 10.
+
+    Returns:
+        A callable that maps a tuple `(a, b)` to a tuple `(row, col)` using the last digit of `a` and `b`.
+
+    Notes:
+        This extractor does not validate the returned indices. Range validation is performed by
+        :class:`~etl_lib.core.SplittingBatchProcessor.SplittingBatchProcessor`.
+    """
+
+    def extractor(item: Tuple[Any, Any]) -> Tuple[int, int]:
+        a, b = item
+        try:
+            row = int(str(a)[-1])
+            col = int(str(b)[-1])
+        except Exception as e:
+            raise ValueError(f"Failed to extract ID from item {item}: {e}")
+        return row, col
+
+    extractor.table_size = table_size
+    return extractor
+
+
+def dict_id_extractor(
+        table_size: int = 10,
+        start_key: str = "start",
+        end_key: str = "end",
+) -> Callable[[Dict[str, Any]], Tuple[int, int]]:
+    """
+    Build an ID extractor for dict rows. The extractor reads two fields (configurable via
+    `start_key` and `end_key`) and returns (row, col) based on the last decimal digit of each.
+
+    Args:
+        table_size: Informational hint carried on the extractor.
+        start_key: Field name for the start node identifier.
+        end_key: Field name for the end node identifier.
+
+    Returns:
+        Callable that maps {start_key, end_key} → (row, col).
+    """
+
+    def extractor(item: Dict[str, Any]) -> Tuple[int, int]:
+        missing = [k for k in (start_key, end_key) if k not in item]
+        if missing:
+            raise KeyError(f"Item missing required keys: {', '.join(missing)}")
+        try:
+            row = int(str(item[start_key])[-1])
+            col = int(str(item[end_key])[-1])
+        except Exception as e:
+            raise ValueError(f"Failed to extract ID from item {item}: {e}")
+        return row, col
+
+    extractor.table_size = table_size
+    return extractor
+
+
+def canonical_integer_id_extractor(
+        table_size: int = 10,
+        start_key: str = "start",
+        end_key: str = "end",
+) -> Callable[[Dict[str, Any]], Tuple[int, int]]:
+    """
+    ID extractor for integer IDs with canonical folding.
+
+    - Uses Knuth's multiplicative hashing to scatter sequential integers across the grid.
+    - Canonical folding ensures (A,B) and (B,A) land in the same bucket by folding the lower
+      triangle into the upper triangle (row <= col).
+
+    The extractor marks itself as mono-partite by setting `extractor.monopartite = True`.
+    """
+    MAGIC = 2654435761
+
+    def extractor(item: Dict[str, Any]) -> Tuple[int, int]:
+        try:
+            s_val = item[start_key]
+            e_val = item[end_key]
+
+            s_hash = (s_val * MAGIC) & 0xffffffff
+            e_hash = (e_val * MAGIC) & 0xffffffff
+
+            row = s_hash % table_size
+            col = e_hash % table_size
+
+            if row > col:
+                row, col = col, row
+
+            return row, col
+        except KeyError:
+            raise KeyError(f"Item missing keys: {start_key} or {end_key}")
+        except Exception as e:
+            raise ValueError(f"Failed to extract ID: {e}")
+
+    extractor.table_size = table_size
+    extractor.monopartite = True
+    return extractor
+
+
+class SplittingBatchProcessor(BatchProcessor):
+    """
+    Streaming wave scheduler for mix-and-batch style loading.
+
+    Incoming rows are assigned to buckets via an `id_extractor(item) -> (row, col)` inside a
+    `table_size x table_size` grid. The processor emits waves; each wave contains bucket-batches
+    that are safe to process concurrently under the configured non-overlap rule.
+
+    Non-overlap rules
+    -----------------
+    - Bi-partite (default): within a wave, no two buckets share a row index and no two buckets share a col index.
+    - Mono-partite: within a wave, no node index is used more than once (row/col indices are the same domain).
+      Enable by setting `id_extractor.monopartite = True` (as done by `canonical_integer_id_extractor`).
+
+    Emission strategy
+    -----------------
+    - During streaming: emit a wave when at least one bucket is full (>= max_batch_size).
+      The wave is then filled with additional non-overlapping buckets that are near-full to
+      keep parallelism high without producing tiny batches.
+    - If a bucket backlog grows beyond a burst threshold, emit a burst wave to bound memory.
+    - After source exhaustion: flush leftovers in capped waves (max_batch_size per bucket).
+
+    Statistics policy
+    -----------------
+    - Every emission except the last carries {}.
+    - The last emission carries the accumulated upstream statistics (unfiltered).
+    """
+
+    def __init__(
+            self,
+            context,
+            table_size: int,
+            id_extractor: Callable[[Any], Tuple[int, int]],
+            task=None,
+            predecessor=None,
+            near_full_ratio: float = 0.85,
+            burst_multiplier: int = 25,
+    ):
+        super().__init__(context, task, predecessor)
+
+        if hasattr(id_extractor, "table_size"):
+            expected_size = id_extractor.table_size
+            if table_size is None:
+                table_size = expected_size
+            elif table_size != expected_size:
+                raise ValueError(
+                    f"Mismatch between provided table_size ({table_size}) and id_extractor table_size ({expected_size})."
+                )
+        elif table_size is None:
+            raise ValueError("table_size must be specified if id_extractor has no defined table_size")
+
+        if not (0 < near_full_ratio <= 1.0):
+            raise ValueError(f"near_full_ratio must be in (0, 1], got {near_full_ratio}")
+        if burst_multiplier < 1:
+            raise ValueError(f"burst_multiplier must be >= 1, got {burst_multiplier}")
+
+        self.table_size = table_size
+        self._id_extractor = id_extractor
+        self._monopartite = bool(getattr(id_extractor, "monopartite", False))
+
+        self.near_full_ratio = float(near_full_ratio)
+        self.burst_multiplier = int(burst_multiplier)
+
+        self.buffer: Dict[int, Dict[int, List[Any]]] = {
+            r: {c: [] for c in range(self.table_size)}
+            for r in range(self.table_size)
+        }
+        self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
+
+    def _bucket_claims(self, row: int, col: int) -> Tuple[Any, ...]:
+        """
+        Return the resource claims a bucket consumes within a wave.
+
+        - Bi-partite: claims (row-slot, col-slot)
+        - Mono-partite: claims node indices touched by the bucket
+        """
+        if self._monopartite:
+            return (row,) if row == col else (row, col)
+        return ("row", row), ("col", col)
+
+    def _all_bucket_sizes(self) -> List[Tuple[int, int, int]]:
+        """
+        Return all non-empty buckets as (size, row, col).
+        """
+        out: List[Tuple[int, int, int]] = []
+        for r in range(self.table_size):
+            for c in range(self.table_size):
+                n = len(self.buffer[r][c])
+                if n > 0:
+                    out.append((n, r, c))
+        return out
+
+    def _select_wave(self, *, min_bucket_len: int, seed: List[Tuple[int, int]] | None = None) -> List[Tuple[int, int]]:
+        """
+        Greedy wave scheduler: pick a non-overlapping set of buckets with len >= min_bucket_len.
+
+        If `seed` is provided, it is taken as fixed and the wave is extended greedily.
+        """
+        candidates: List[Tuple[int, int, int]] = []
+        for r in range(self.table_size):
+            for c in range(self.table_size):
+                n = len(self.buffer[r][c])
+                if n >= min_bucket_len:
+                    candidates.append((n, r, c))
+
+        if not candidates and not seed:
+            return []
+
+        candidates.sort(key=lambda x: (-x[0], x[1], x[2]))
+
+        used: set[Any] = set()
+        wave: List[Tuple[int, int]] = []
+
+        if seed:
+            for r, c in seed:
+                claims = self._bucket_claims(r, c)
+                used.update(claims)
+                wave.append((r, c))
+
+        for _, r, c in candidates:
+            if (r, c) in wave:
+                continue
+            claims = self._bucket_claims(r, c)
+            if any(claim in used for claim in claims):
+                continue
+            wave.append((r, c))
+            used.update(claims)
+            if len(wave) >= self.table_size:
+                break
+
+        return wave
+
+    def _find_hottest_bucket(self, *, threshold: int) -> Tuple[int, int, int] | None:
+        """
+        Find the single hottest bucket (largest backlog) whose size is >= threshold.
+        Returns (row, col, size) or None.
+        """
+        best: Tuple[int, int, int] | None = None
+        for r in range(self.table_size):
+            for c in range(self.table_size):
+                n = len(self.buffer[r][c])
+                if n < threshold:
+                    continue
+                if best is None or n > best[2]:
+                    best = (r, c, n)
+        return best
+
+    def _flush_wave(
+            self,
+            wave: List[Tuple[int, int]],
+            max_batch_size: int,
+            statistics: Dict[str, Any] | None = None,
+    ) -> BatchResults:
+        """
+        Extract up to `max_batch_size` items from each bucket in `wave`, remove them from the buffer,
+        and return a BatchResults whose chunk is a list of per-bucket lists (aligned with `wave`).
+        """
+        self._log_buffer_matrix(wave=wave)
+
+        bucket_batches: List[List[Any]] = []
+        for r, c in wave:
+            q = self.buffer[r][c]
+            take = min(max_batch_size, len(q))
+            bucket_batches.append(q[:take])
+            self.buffer[r][c] = q[take:]
+
+        emitted = sum(len(b) for b in bucket_batches)
+
+        return BatchResults(
+            chunk=bucket_batches,
+            statistics=statistics or {},
+            batch_size=emitted,
+        )
+
+    def _log_buffer_matrix(self, *, wave: List[Tuple[int, int]]) -> None:
+        """
+        Dumps a compact 2D matrix of per-bucket sizes (len of each buffer) when DEBUG is enabled.
+        """
+        if not self.logger.isEnabledFor(logging.DEBUG):
+            return
+
+        counts = [
+            [len(self.buffer[r][c]) for c in range(self.table_size)]
+            for r in range(self.table_size)
+        ]
+        marks = set(wave)
+
+        pad = max(2, len(str(self.table_size - 1)))
+        col_headers = [f"c{c:0{pad}d}" for c in range(self.table_size)]
+
+        rows = []
+        for r in range(self.table_size):
+            row_label = f"r{r:0{pad}d}"
+            row_vals = [f"[{v}]" if (r, c) in marks else f"{v}" for c, v in enumerate(counts[r])]
+            rows.append([row_label, *row_vals])
+
+        table = tabulate(
+            rows,
+            headers=["", *col_headers],
+            tablefmt="psql",
+            stralign="right",
+            disable_numparse=True,
+        )
+        self.logger.debug("buffer matrix:\n%s", table)
+
+    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
+        """
+        Consume upstream batches, bucket incoming rows, and emit waves of non-overlapping buckets.
+
+        Streaming behavior:
+          - If at least one bucket is full (>= max_batch_size), emit a wave seeded with full buckets
+            and extended with near-full buckets to keep parallelism high.
+          - If a bucket exceeds a burst threshold, emit a burst wave (seeded by the hottest bucket)
+            and extended with near-full buckets.
+
+        End-of-stream behavior:
+          - Flush leftovers in capped waves (max_batch_size per bucket).
+
+        Statistics policy:
+          * Every emission except the last carries {}.
+          * The last emission carries the accumulated upstream stats (unfiltered).
+        """
+        if self.predecessor is None:
+            return
+
+        accum_stats: Dict[str, Any] = {}
+        pending: BatchResults | None = None
+
+        near_full_threshold = max(1, int(max_batch_size * self.near_full_ratio))
+        burst_threshold = self.burst_multiplier * max_batch_size
+
+        for upstream in self.predecessor.get_batch(max_batch_size):
+            if upstream.statistics:
+                accum_stats = merge_summery(accum_stats, upstream.statistics)
+
+            for item in upstream.chunk:
+                r, c = self._id_extractor(item)
+                if self._monopartite and r > c:
+                    r, c = c, r
+                if not (0 <= r < self.table_size and 0 <= c < self.table_size):
+                    raise ValueError(f"bucket id out of range: {(r, c)} for table_size={self.table_size}")
+                self.buffer[r][c].append(item)
+
+            while True:
+                full_seed = self._select_wave(min_bucket_len=max_batch_size)
+                if not full_seed:
+                    break
+                wave = self._select_wave(min_bucket_len=near_full_threshold, seed=full_seed)
+                br = self._flush_wave(wave, max_batch_size, statistics={})
+                if pending is not None:
+                    yield pending
+                pending = br
+
+            while True:
+                hot = self._find_hottest_bucket(threshold=burst_threshold)
+                if hot is None:
+                    break
+                hot_r, hot_c, hot_n = hot
+                wave = self._select_wave(min_bucket_len=near_full_threshold, seed=[(hot_r, hot_c)])
+                self.logger.debug(
+                    "burst flush: hottest_bucket=(%d,%d len=%d) threshold=%d near_full_threshold=%d wave_size=%d",
+                    hot_r, hot_c, hot_n, burst_threshold, near_full_threshold, len(wave)
+                )
+                br = self._flush_wave(wave, max_batch_size, statistics={})
+                if pending is not None:
+                    yield pending
+                pending = br
+
+        self.logger.debug("start flushing leftovers")
+        while True:
+            wave = self._select_wave(min_bucket_len=1)
+            if not wave:
+                break
+            br = self._flush_wave(wave, max_batch_size, statistics={})
+            if pending is not None:
+                yield pending
+            pending = br
+
+        if pending is not None:
+            yield BatchResults(chunk=pending.chunk, statistics=accum_stats, batch_size=pending.batch_size)

neo4j_etl_lib-0.3.2/src/etl_lib/data_source/SQLBatchSource.py

@@ -0,0 +1,84 @@
+import logging
+from typing import Generator, Callable, Optional
+
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError as SAOperationalError, DBAPIError
+
+# Conditional import for psycopg2 to avoid crashing if not installed
+try:
+    from psycopg2 import OperationalError as PsycopgOperationalError
+except ImportError:
+    class PsycopgOperationalError(Exception):
+        pass
+
+from etl_lib.core.BatchProcessor import BatchResults, BatchProcessor
+from etl_lib.core.ETLContext import ETLContext
+from etl_lib.core.Task import Task
+
+
+class SQLBatchSource(BatchProcessor):
+    def __init__(
+            self,
+            context: ETLContext,
+            task: Task,
+            query: str,
+            record_transformer: Optional[Callable[[dict], dict]] = None,
+            **kwargs
+    ):
+        """
+        Constructs a new SQLBatchSource that streams results instead of paging them.
+        """
+        super().__init__(context, task)
+        # Remove any trailing semicolons to prevent SQL syntax errors
+        self.query = query.strip().rstrip(";")
+        self.record_transformer = record_transformer
+        self.kwargs = kwargs
+        self.logger = logging.getLogger(__name__)
+
+    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
+        """
+        Yield successive batches using a Server-Side Cursor (Streaming).
+
+        This avoids 'LIMIT/OFFSET' pagination, which causes performance degradation
+        on large tables. Instead, it holds a cursor open and fetches rows incrementally.
+        """
+
+        with self.context.sql.engine.connect() as conn:
+
+            conn = conn.execution_options(stream_results=True)
+
+            try:
+                self.logger.info("Starting SQL Result Stream...")
+
+                result_proxy = conn.execute(text(self.query), self.kwargs)
+
+                chunk = []
+                count = 0
+
+                for row in result_proxy.mappings():
+                    item = self.record_transformer(dict(row)) if self.record_transformer else dict(row)
+                    chunk.append(item)
+                    count += 1
+
+                    # Yield when batch is full
+                    if len(chunk) >= max_batch_size:
+                        yield BatchResults(
+                            chunk=chunk,
+                            statistics={"sql_rows_read": len(chunk)},
+                            batch_size=len(chunk),
+                        )
+                        chunk = []  # Clear memory
+
+                # Yield any remaining rows
+                if chunk:
+                    yield BatchResults(
+                        chunk=chunk,
+                        statistics={"sql_rows_read": len(chunk)},
+                        batch_size=len(chunk),
+                    )
+
+                self.logger.info(f"SQL Stream finished. Total rows read: {count}")
+
+            except (PsycopgOperationalError, SAOperationalError, DBAPIError) as err:
+                self.logger.error(f"Stream failed: {err}")
+                raise

neo4j_etl_lib-0.3.1/src/etl_lib/core/ParallelBatchProcessor.py

@@ -1,180 +0,0 @@
-import queue
-import threading
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Any, Callable, Generator, List
-
-from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults
-from etl_lib.core.utils import merge_summery
-
-
-class ParallelBatchResult(BatchResults):
-    """
-    Represents a batch split into parallelizable partitions.
-
-    `chunk` is a list of lists, each sub-list is a partition.
-    """
-    pass
-
-
-class ParallelBatchProcessor(BatchProcessor):
-    """
-    BatchProcessor that runs worker threads over partitions of batches.
-
-    Receives a special BatchResult (:py:class:`ParallelBatchResult`) from the predecessor.
-    All chunks in a ParallelBatchResult it receives can be processed in parallel.
-    See :py:class:`etl_lib.core.SplittingBatchProcessor` on how to produce them.
-    Prefetches the next ParallelBatchResults from its predecessor.
-    The actual processing of the batches is deferred to the configured worker.
-
-    Note:
-        - The predecessor must emit `ParallelBatchResult` instances.
-
-    Args:
-        context: ETL context.
-        worker_factory: A zero-arg callable that returns a new BatchProcessor
-                        each time it's called. This processor is responsible for the processing pf the batches.
-        task: optional Task for reporting.
-        predecessor: upstream BatchProcessor that must emit ParallelBatchResult.
-        max_workers: number of parallel threads for partitions.
-        prefetch: number of ParallelBatchResults to prefetch from the predecessor.
-
-    Behavior:
-        - For every ParallelBatchResult, spins up `max_workers` threads.
-        - Each thread calls its own worker from `worker_factory()`, with its
-          partition wrapped by `SingleBatchWrapper`.
-        - Collects and merges their BatchResults in a fail-fast manner: on first
-          exception, logs the error, cancels remaining threads, and raises an exception.
-    """
-
-    def __init__(
-            self,
-            context,
-            worker_factory: Callable[[], BatchProcessor],
-            task=None,
-            predecessor=None,
-            max_workers: int = 4,
-            prefetch: int = 4,
-    ):
-        super().__init__(context, task, predecessor)
-        self.worker_factory = worker_factory
-        self.max_workers = max_workers
-        self.prefetch = prefetch
-        self._batches_done = 0
-
-    def _process_parallel(self, pbr: ParallelBatchResult) -> BatchResults:
-        """
-        Run one worker per partition in `pbr.chunk`, merge their outputs, and include upstream
-        statistics from `pbr.statistics` so counters (e.g., valid/invalid rows from validation)
-        are preserved through the parallel stage.
-
-        Progress reporting:
-          - After each partition completes, report batch count only
-        """
-        merged_stats = dict(pbr.statistics or {})
-        merged_chunk = []
-        total = 0
-
-        parts_total = len(pbr.chunk)
-        partitions_done = 0
-
-        self.logger.debug(f"Processing pbr of len {parts_total}")
-        with ThreadPoolExecutor(max_workers=self.max_workers, thread_name_prefix='PBP_worker_') as pool:
-            futures = [pool.submit(self._process_partition, part) for part in pbr.chunk]
-            try:
-                for f in as_completed(futures):
-                    out = f.result()
-
-                    # Merge into this PBR's cumulative result (returned downstream)
-                    merged_stats = merge_summery(merged_stats, out.statistics or {})
-                    total += out.batch_size
-                    merged_chunk.extend(out.chunk if isinstance(out.chunk, list) else [out.chunk])
-
-                    partitions_done += 1
-                    self.context.reporter.report_progress(
-                        task=self.task,
-                        batches=self._batches_done,
-                        expected_batches=None,
-                        stats={},
-                    )
-
-            except Exception as e:
-                for g in futures:
-                    g.cancel()
-                pool.shutdown(cancel_futures=True)
-                raise RuntimeError("partition processing failed") from e
-
-        self.logger.debug(f"Finished processing pbr with {merged_stats}")
-        return BatchResults(chunk=merged_chunk, statistics=merged_stats, batch_size=total)
-
-    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
-        """
-        Pulls ParallelBatchResult batches from the predecessor, prefetching
-        up to `prefetch` ahead, processes each batch's partitions in
-        parallel threads, and yields a flattened BatchResults. The predecessor
-        can run ahead while the current batch is processed.
-        """
-        pbr_queue: queue.Queue[ParallelBatchResult | object] = queue.Queue(self.prefetch)
-        SENTINEL = object()
-        exc: BaseException | None = None
-
-        def producer():
-            nonlocal exc
-            try:
-                for pbr in self.predecessor.get_batch(max_batch_size):
-                    self.logger.debug(
-                        f"adding pgr {pbr.statistics} / {len(pbr.chunk)} to queue of size {pbr_queue.qsize()}"
-                    )
-                    pbr_queue.put(pbr)
-            except BaseException as e:
-                exc = e
-            finally:
-                pbr_queue.put(SENTINEL)
-
-        threading.Thread(target=producer, daemon=True, name='prefetcher').start()
-
-        while True:
-            pbr = pbr_queue.get()
-            if pbr is SENTINEL:
-                if exc is not None:
-                    self.logger.error("Upstream producer failed", exc_info=True)
-                    raise exc
-                break
-            result = self._process_parallel(pbr)
-            yield result
-
-    class SingleBatchWrapper(BatchProcessor):
-        """
-        Simple BatchProcessor that returns the batch it receives via init.
-        Will be used as predecessor for the worker
-        """
-
-        def __init__(self, context, batch: List[Any]):
-            super().__init__(context=context, predecessor=None)
-            self._batch = batch
-
-        def get_batch(self, max_batch__size: int) -> Generator[BatchResults, None, None]:
-            # Ignores max_size; yields exactly one BatchResults containing the whole batch
-            yield BatchResults(
-                chunk=self._batch,
-                statistics={},
-                batch_size=len(self._batch)
-            )
-
-    def _process_partition(self, partition):
-        """
-        Processes one partition of items by:
-          1. Wrapping it in SingleBatchWrapper
-          2. Instantiating a fresh worker via worker_factory()
-          3. Setting the worker's predecessor to the wrapper
-          4. Running exactly one batch and returning its BatchResults
-
-        Raises whatever exception the worker raises, allowing _process_parallel
-        to handle fail-fast behavior.
-        """
-        self.logger.debug("Processing partition")
-        wrapper = self.SingleBatchWrapper(self.context, partition)
-        worker = self.worker_factory()
-        worker.predecessor = wrapper
-        result = next(worker.get_batch(len(partition)))
-        self.logger.debug(f"finished processing partition with {result.statistics}")
-        return result

neo4j_etl_lib-0.3.1/src/etl_lib/core/SplittingBatchProcessor.py

@@ -1,268 +0,0 @@
-import logging
-from typing import Any, Callable, Dict, Generator, List, Tuple
-
-from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults
-from etl_lib.core.utils import merge_summery
-from tabulate import tabulate
-
-
-def tuple_id_extractor(table_size: int = 10) -> Callable[[Tuple[str | int, str | int]], Tuple[int, int]]:
-    """
-    Create an ID extractor function for tuple items, using the last decimal digit of each element.
-    The output is a `(row, col)` tuple within a `table_size x table_size` grid (default 10x10).
-
-    Args:
-        table_size: The dimension of the grid (number of rows/cols). Defaults to 10.
-
-    Returns:
-        A callable that maps a tuple `(a, b)` to a tuple `(row, col)` using the last digit of `a` and `b`.
-    """
-
-    def extractor(item: Tuple[Any, Any]) -> Tuple[int, int]:
-        a, b = item
-        try:
-            row = int(str(a)[-1])
-            col = int(str(b)[-1])
-        except Exception as e:
-            raise ValueError(f"Failed to extract ID from item {item}: {e}")
-        return row, col
-
-    extractor.table_size = table_size
-    return extractor
-
-
-def dict_id_extractor(
-        table_size: int = 10,
-        start_key: str = "start",
-        end_key: str = "end",
-) -> Callable[[Dict[str, Any]], Tuple[int, int]]:
-    """
-    Build an ID extractor for dict rows. The extractor reads two fields (configurable via
-    `start_key` and `end_key`) and returns (row, col) based on the last decimal digit of each.
-    Range validation remains the responsibility of the SplittingBatchProcessor.
-
-    Args:
-        table_size: Informational hint carried on the extractor; used by callers to sanity-check.
-        start_key: Field name for the start node identifier.
-        end_key: Field name for the end node identifier.
-
-    Returns:
-        Callable[[Mapping[str, Any]], tuple[int, int]]: Maps {start_key, end_key} → (row, col).
-    """
-
-    def extractor(item: Dict[str, Any]) -> Tuple[int, int]:
-        missing = [k for k in (start_key, end_key) if k not in item]
-        if missing:
-            raise KeyError(f"Item missing required keys: {', '.join(missing)}")
-        try:
-            row = int(str(item[start_key])[-1])
-            col = int(str(item[end_key])[-1])
-        except Exception as e:
-            raise ValueError(f"Failed to extract ID from item {item}: {e}")
-        return row, col
-
-    extractor.table_size = table_size
-    return extractor
-
-
-class SplittingBatchProcessor(BatchProcessor):
-    """
-    BatchProcessor that splits incoming BatchResults into non-overlapping partitions based
-    on row/col indices extracted by the id_extractor, and emits full or remaining batches
-    using the mix-and-batch algorithm from https://neo4j.com/blog/developer/mix-and-batch-relationship-load/
-    Each emitted batch is a list of per-cell lists (array of arrays), so callers
-    can process each partition (other name for a cell) in parallel.
-
-    A batch for a schedule group is  emitted when all cells in that group have at least `batch_size` items.
-    In addition, when a cell/partition reaches 3x the configured max_batch_size, the group is emitted to avoid
-    overflowing the buffer when the distribution per cell is uneven.
-    Leftovers are flushed after source exhaustion.
-    Emitted batches never exceed the configured max_batch_size.
-    """
-
-    def __init__(self, context, table_size: int, id_extractor: Callable[[Any], Tuple[int, int]],
-                 task=None, predecessor=None):
-        super().__init__(context, task, predecessor)
-
-        # If the extractor carries an expected table size, use or validate it
-        if hasattr(id_extractor, "table_size"):
-            expected_size = id_extractor.table_size
-            if table_size is None:
-                table_size = expected_size  # determine table size from extractor if not provided
-            elif table_size != expected_size:
-                raise ValueError(f"Mismatch between provided table_size ({table_size}) and "
-                                 f"id_extractor table_size ({expected_size}).")
-        elif table_size is None:
-            raise ValueError("table_size must be specified if id_extractor has no defined table_size")
-        self.table_size = table_size
-        self._id_extractor = id_extractor
-
-        # Initialize 2D buffer for partitions
-        self.buffer: Dict[int, Dict[int, List[Any]]] = {
-            i: {j: [] for j in range(self.table_size)} for i in range(self.table_size)
-        }
-        self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
-
-    def _generate_batch_schedule(self) -> List[List[Tuple[int, int]]]:
-        """
-        Create diagonal stripes (row, col) partitions to ensure no overlapping IDs
-        across emitted batches.
-        Example grid:
-                 ||  0  |  1  |  2
-            =====++=====+=====+=====
-              0  ||  0  |  1  |  2
-            -----++-----+-----+----
-              1  ||  2  |  0  |  1
-            -----++-----+-----+-----
-              2  ||  1  |  2  |  0
-
-        would return [[(0, 0), (1, 1), (2, 2)], [(0, 1), (1, 2), (2, 0)], [(0, 2), (1, 0), (2, 1)]]
-        """
-        schedule: List[List[Tuple[int, int]]] = []
-        for shift in range(self.table_size):
-            partition: List[Tuple[int, int]] = [
-                (i, (i + shift) % self.table_size)
-                for i in range(self.table_size)
-            ]
-            schedule.append(partition)
-        return schedule
-
-    def _flush_group(
-            self,
-            partitions: List[Tuple[int, int]],
-            batch_size: int,
-            statistics: Dict[str, Any] | None = None,
-    ) -> Generator[BatchResults, None, None]:
-        """
-        Extract up to `batch_size` items from each cell in `partitions`, remove them from the buffer,
-        and yield a BatchResults whose chunks is a list of per-cell lists from these partitions.
-
-        Args:
-            partitions: Cell coordinates forming a diagonal group from the schedule.
-            batch_size: Max number of items to take from each cell.
-            statistics: Stats dict to attach to this emission (use {} for interim waves).
-                        The "final" emission will pass the accumulated stats here.
-
-        Notes:
-            - Debug-only: prints a 2D matrix of cell sizes when logger is in DEBUG.
-            - batch_size in the returned BatchResults equals the number of emitted items.
-        """
-        self._log_buffer_matrix(partition=partitions)
-
-        cell_chunks: List[List[Any]] = []
-        for row, col in partitions:
-            q = self.buffer[row][col]
-            take = min(batch_size, len(q))
-            part = q[:take]
-            cell_chunks.append(part)
-            # remove flushed items
-            self.buffer[row][col] = q[take:]
-
-        emitted = sum(len(c) for c in cell_chunks)
-
-        result = BatchResults(
-            chunk=cell_chunks,
-            statistics=statistics or {},
-            batch_size=emitted,
-        )
-        yield result
-
-    def get_batch(self, max_batch__size: int) -> Generator[BatchResults, None, None]:
-        """
-        Consume upstream batches, split data across cells, and emit diagonal partitions:
-          - During streaming: emit a full partition when all its cells have >= max_batch__size.
-          - Also during streaming: if any cell in a partition builds up beyond a 'burst' threshold
-            (3 * of max_batch__size), emit that partition early, taking up to max_batch__size
-            from each cell.
-          - After source exhaustion: flush leftovers in waves capped at max_batch__size per cell.
-
-        Statistics policy:
-          * Every emission except the last carries {}.
-          * The last emission carries the accumulated upstream stats (unfiltered).
-        """
-        schedule = self._generate_batch_schedule()
-
-        accum_stats: Dict[str, Any] = {}
-        pending: BatchResults | None = None  # hold back the most recent emission so we know what's final
-
-        burst_threshold = 3 * max_batch__size
-
-        for upstream in self.predecessor.get_batch(max_batch__size):
-            # accumulate upstream stats (unfiltered)
-            if upstream.statistics:
-                accum_stats = merge_summery(accum_stats, upstream.statistics)
-
-            # add data to cells
-            for item in upstream.chunk:
-                r, c = self._id_extractor(item)
-                if not (0 <= r < self.table_size and 0 <= c < self.table_size):
-                    raise ValueError(f"partition id out of range: {(r, c)} for table_size={self.table_size}")
-                self.buffer[r][c].append(item)
-
-            # process partitions
-            for partition in schedule:
-                # normal full flush when all cells are ready
-                if all(len(self.buffer[r][c]) >= max_batch__size for r, c in partition):
-                    br = next(self._flush_group(partition, max_batch__size, statistics={}))
-                    if pending is not None:
-                        yield pending
-                    pending = br
-                    continue
-
-                # burst flush if any cell backlog explodes
-                hot_cells = [(r, c, len(self.buffer[r][c])) for r, c in partition if
-                             len(self.buffer[r][c]) >= burst_threshold]
-                if hot_cells:
-                    top_r, top_c, top_len = max(hot_cells, key=lambda x: x[2])
-                    self.logger.debug(
-                        "burst flush: partition=%s threshold=%d top_cell=(%d,%d len=%d)",
-                        partition, burst_threshold, top_r, top_c, top_len
-                    )
-                    br = next(self._flush_group(partition, max_batch__size, statistics={}))
-                    if pending is not None:
-                        yield pending
-                    pending = br
-
-        # source exhausted: drain leftovers in capped waves (respecting batch size)
-        self.logger.debug("start flushing leftovers")
-        for partition in (p for p in schedule if any(self.buffer[r][c] for r, c in p)):
-            while any(self.buffer[r][c] for r, c in partition):
-                br = next(self._flush_group(partition, max_batch__size, statistics={}))
-                if pending is not None:
-                    yield pending
-                pending = br
-
-        # final emission carries accumulated stats once
-        if pending is not None:
-            yield BatchResults(chunk=pending.chunk, statistics=accum_stats, batch_size=pending.batch_size)
-
-    def _log_buffer_matrix(self, *, partition: List[Tuple[int, int]]) -> None:
-        """
-        Dumps a compact 2D matrix of per-cell sizes (len of each buffer) when DEBUG is enabled.
-        """
-        if not self.logger.isEnabledFor(logging.DEBUG):
-            return
-
-        counts = [
-            [len(self.buffer[r][c]) for c in range(self.table_size)]
-            for r in range(self.table_size)
-        ]
-        marks = set(partition)
-
-        pad = max(2, len(str(self.table_size - 1)))
-        col_headers = [f"c{c:0{pad}d}" for c in range(self.table_size)]
-
-        rows = []
-        for r in range(self.table_size):
-            row_label = f"r{r:0{pad}d}"
-            row_vals = [f"[{v}]" if (r, c) in marks else f"{v}" for c, v in enumerate(counts[r])]
-            rows.append([row_label, *row_vals])
-
-        table = tabulate(
-            rows,
-            headers=["", *col_headers],
-            tablefmt="psql",
-            stralign="right",
-            disable_numparse=True,
-        )
-        self.logger.debug("buffer matrix:\n%s", table)

neo4j_etl_lib-0.3.1/src/etl_lib/data_source/SQLBatchSource.py

@@ -1,114 +0,0 @@
-import time
-from typing import Generator, Callable, Optional, List, Dict
-
-from psycopg2 import OperationalError as PsycopgOperationalError
-from sqlalchemy import text
-from sqlalchemy.exc import OperationalError as SAOperationalError, DBAPIError
-
-from etl_lib.core.BatchProcessor import BatchResults, BatchProcessor
-from etl_lib.core.ETLContext import ETLContext
-from etl_lib.core.Task import Task
-
-
-class SQLBatchSource(BatchProcessor):
-    def __init__(
-            self,
-            context: ETLContext,
-            task: Task,
-            query: str,
-            record_transformer: Optional[Callable[[dict], dict]] = None,
-            **kwargs
-    ):
-        """
-        Constructs a new SQLBatchSource.
-
-        Args:
-            context: :class:`etl_lib.core.ETLContext.ETLContext` instance.
-            task: :class:`etl_lib.core.Task.Task` instance owning this batchProcessor.
-            query: SQL query to execute.
-            record_transformer: Optional function to transform each row (dict format).
-            kwargs: Arguments passed as parameters with the query.
-        """
-        super().__init__(context, task)
-        self.query = query.strip().rstrip(";")
-        self.record_transformer = record_transformer
-        self.kwargs = kwargs
-
-    def _fetch_page(self, limit: int, offset: int) -> Optional[List[Dict]]:
-        """
-        Fetch a single batch of rows using LIMIT/OFFSET, with retry/backoff.
-
-        Each page is executed in its own transaction. On transient
-        disconnects or DB errors, it retries up to 3 times with exponential backoff.
-
-        Args:
-            limit: maximum number of rows to return.
-            offset: number of rows to skip before starting this page.
-
-        Returns:
-            A list of row dicts (after applying record_transformer, if any),
-            or None if no rows are returned.
-
-        Raises:
-            Exception: re-raises the last caught error on final failure.
-        """
-        paged_sql = f"{self.query} LIMIT :limit OFFSET :offset"
-        params = {**self.kwargs, "limit": limit, "offset": offset}
-        max_retries = 5
-        backoff = 2.0
-
-        for attempt in range(1, max_retries + 1):
-            try:
-                with self.context.sql.engine.connect() as conn:
-                    with conn.begin():
-                        rows = conn.execute(text(paged_sql), params).mappings().all()
-                result = [
-                    self.record_transformer(dict(r)) if self.record_transformer else dict(r)
-                    for r in rows
-                ]
-                return result if result else None
-
-            except (PsycopgOperationalError, SAOperationalError, DBAPIError) as err:
-
-                if attempt == max_retries:
-                    self.logger.error(
-                        f"Page fetch failed after {max_retries} attempts "
-                        f"(limit={limit}, offset={offset}): {err}"
-                    )
-                    raise
-
-                self.logger.warning(
-                    f"Transient DB error on page fetch {attempt}/{max_retries}: {err!r}, "
-                    f"retrying in {backoff:.1f}s"
-                )
-                time.sleep(backoff)
-                backoff *= 2
-
-        return None
-
-    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
-        """
-        Yield successive batches until the query is exhausted.
-
-        Calls _fetch_page() repeatedly, advancing the offset by the
-        number of rows returned. Stops when no more rows are returned.
-
-        Args:
-            max_batch_size: upper limit on rows per batch.
-
-        Yields:
-            BatchResults for each non-empty page.
-        """
-        offset = 0
-        while True:
-            chunk = self._fetch_page(max_batch_size, offset)
-            if not chunk:
-                break
-
-            yield BatchResults(
-                chunk=chunk,
-                statistics={"sql_rows_read": len(chunk)},
-                batch_size=len(chunk),
-            )
-
-            offset += len(chunk)