neo4j-etl-lib 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

etl_lib/__init__.py

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

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:

etl_lib/core/ParallelBatchProcessor.py

@@ -9,41 +9,38 @@ from etl_lib.core.utils import merge_summery
 
 class ParallelBatchResult(BatchResults):
     """
-    Represents a batch split into parallelizable partitions.
+    Represents one *wave* produced by the splitter.
 
-    `chunk` is a list of lists, each sub-list is a partition.
+    `chunk` is a list of bucket-batches. Each sub-list is processed by one worker instance.
     """
     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.
+    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.
+        - 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. This processor is responsible for the processing pf the batches.
+        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.
-        max_workers: number of parallel threads for partitions.
-        prefetch: number of ParallelBatchResults to prefetch from the predecessor.
+        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 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.
+        - 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__(
@@ -59,122 +56,95 @@ class ParallelBatchProcessor(BatchProcessor):
         self.worker_factory = worker_factory
         self.max_workers = max_workers
         self.prefetch = prefetch
-        self._batches_done = 0
 
-    def _process_parallel(self, pbr: ParallelBatchResult) -> BatchResults:
+    def _process_wave(self, wave: 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.
+        Process one wave: run one worker per bucket-batch and merge their BatchResults.
 
-        Progress reporting:
-          - After each partition completes, report batch count only
+        Statistics:
+            `wave.statistics` is used as the initial merged stats, then merged with each worker's stats.
         """
-        merged_stats = dict(pbr.statistics or {})
+        merged_stats = dict(wave.statistics or {})
         merged_chunk = []
         total = 0
 
-        parts_total = len(pbr.chunk)
-        partitions_done = 0
+        self.logger.debug(f"Processing wave with {len(wave.chunk)} buckets")
 
-        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]
+        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()
-
-                    # 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
+                raise RuntimeError("bucket processing failed") from e
 
-        self.logger.debug(f"Finished processing pbr with {merged_stats}")
+        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]:
         """
-        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.
+        Pull waves from the predecessor (prefetching up to `prefetch` ahead), process each wave's
+        buckets in parallel, and yield one flattened BatchResults per wave.
         """
-        pbr_queue: queue.Queue[ParallelBatchResult | object] = queue.Queue(self.prefetch)
+        wave_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):
+                for wave 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()}"
+                        f"adding wave stats={wave.statistics} buckets={len(wave.chunk)} to queue size={wave_queue.qsize()}"
                     )
-                    pbr_queue.put(pbr)
+                    wave_queue.put(wave)
             except BaseException as e:
                 exc = e
             finally:
-                pbr_queue.put(SENTINEL)
+                wave_queue.put(SENTINEL)
 
-        threading.Thread(target=producer, daemon=True, name='prefetcher').start()
+        threading.Thread(target=producer, daemon=True, name="prefetcher").start()
 
         while True:
-            pbr = pbr_queue.get()
-            if pbr is SENTINEL:
+            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
-            result = self._process_parallel(pbr)
-            yield result
+            yield self._process_wave(wave)
 
     class SingleBatchWrapper(BatchProcessor):
         """
-        Simple BatchProcessor that returns the batch it receives via init.
-        Will be used as predecessor for the worker
+        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_batch__size: int) -> Generator[BatchResults, None, None]:
-            # Ignores max_size; yields exactly one BatchResults containing the whole batch
+        def get_batch(self, max_size: int) -> Generator[BatchResults, None, None]:
             yield BatchResults(
                 chunk=self._batch,
                 statistics={},
-                batch_size=len(self._batch)
+                batch_size=len(self._batch),
             )
 
-    def _process_partition(self, partition):
+    def _process_bucket_batch(self, bucket_batch):
         """
-        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.
+        Process one bucket-batch by running a fresh worker over it.
         """
-        self.logger.debug("Processing partition")
-        wrapper = self.SingleBatchWrapper(self.context, partition)
+        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(partition)))
-        self.logger.debug(f"finished processing partition with {result.statistics}")
+        result = next(worker.get_batch(len(bucket_batch)))
+        self.logger.debug(f"Finished bucket batch stats={result.statistics}")
         return result

etl_lib/core/SplittingBatchProcessor.py

@@ -16,6 +16,10 @@ def tuple_id_extractor(table_size: int = 10) -> Callable[[Tuple[str | int, str |
 
     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]:
@@ -39,15 +43,14 @@ def dict_id_extractor(
     """
     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.
+        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[[Mapping[str, Any]], tuple[int, int]]: Maps {start_key, end_key} → (row, col).
+        Callable that maps {start_key, end_key} → (row, col).
     """
 
     def extractor(item: Dict[str, Any]) -> Tuple[int, int]:
@@ -65,180 +68,224 @@ def dict_id_extractor(
     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):
     """
-    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.
+    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):
+    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 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
+                table_size = expected_size
             elif table_size != expected_size:
-                raise ValueError(f"Mismatch between provided table_size ({table_size}) and "
-                                 f"id_extractor 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)
 
-        # 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)
+            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 _generate_batch_schedule(self) -> List[List[Tuple[int, int]]]:
+    def _bucket_claims(self, row: int, col: int) -> Tuple[Any, ...]:
         """
-        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)]]
+        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
         """
-        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]:
+        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]]:
         """
-        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.
+        Return all non-empty buckets as (size, row, col).
         """
-        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
+        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 get_batch(self, max_batch__size: int) -> Generator[BatchResults, None, None]:
+    def _select_wave(self, *, min_bucket_len: int, seed: List[Tuple[int, int]] | None = None) -> List[Tuple[int, int]]:
         """
-        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.
+        Greedy wave scheduler: pick a non-overlapping set of buckets with len >= min_bucket_len.
 
-        Statistics policy:
-          * Every emission except the last carries {}.
-          * The last emission carries the accumulated upstream stats (unfiltered).
+        If `seed` is provided, it is taken as fixed and the wave is extended greedily.
         """
-        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)
+        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
 
-            # 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)
+    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)
 
-            # 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
+        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:]
 
-                # 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
+        emitted = sum(len(b) for b in bucket_batches)
 
-        # final emission carries accumulated stats once
-        if pending is not None:
-            yield BatchResults(chunk=pending.chunk, statistics=accum_stats, batch_size=pending.batch_size)
+        return BatchResults(
+            chunk=bucket_batches,
+            statistics=statistics or {},
+            batch_size=emitted,
+        )
 
-    def _log_buffer_matrix(self, *, partition: List[Tuple[int, int]]) -> None:
+    def _log_buffer_matrix(self, *, wave: List[Tuple[int, int]]) -> None:
         """
-        Dumps a compact 2D matrix of per-cell sizes (len of each buffer) when DEBUG is enabled.
+        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
@@ -247,7 +294,7 @@ class SplittingBatchProcessor(BatchProcessor):
             [len(self.buffer[r][c]) for c in range(self.table_size)]
             for r in range(self.table_size)
         ]
-        marks = set(partition)
+        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)]
@@ -266,3 +313,79 @@ class SplittingBatchProcessor(BatchProcessor):
             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)

etl_lib/data_source/SQLBatchSource.py

@@ -1,10 +1,16 @@
-import time
-from typing import Generator, Callable, Optional, List, Dict
+import logging
+from typing import Generator, Callable, Optional
 
-from psycopg2 import OperationalError as PsycopgOperationalError
 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
@@ -20,95 +26,59 @@ class SQLBatchSource(BatchProcessor):
             **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.
+        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 _fetch_page(self, limit: int, offset: int) -> Optional[List[Dict]]:
+    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
         """
-        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.
+        Yield successive batches using a Server-Side Cursor (Streaming).
 
-        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.
+        This avoids 'LIMIT/OFFSET' pagination, which causes performance degradation
+        on large tables. Instead, it holds a cursor open and fetches rows incrementally.
         """
-        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
+        with self.context.sql.engine.connect() as conn:
 
-            except (PsycopgOperationalError, SAOperationalError, DBAPIError) as err:
+            conn = conn.execution_options(stream_results=True)
 
-                if attempt == max_retries:
-                    self.logger.error(
-                        f"Page fetch failed after {max_retries} attempts "
-                        f"(limit={limit}, offset={offset}): {err}"
+            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),
                     )
-                    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
+                self.logger.info(f"SQL Stream finished. Total rows read: {count}")
 
-        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)
+            except (PsycopgOperationalError, SAOperationalError, DBAPIError) as err:
+                self.logger.error(f"Stream failed: {err}")
+                raise

{neo4j_etl_lib-0.3.1.dist-info → neo4j_etl_lib-0.3.2.dist-info}/METADATA

@@ -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.dist-info → neo4j_etl_lib-0.3.2.dist-info}/RECORD

@@ -1,12 +1,12 @@
-etl_lib/__init__.py,sha256=x6coFV38ytJ_wPhR3c0UEzX65oTz2ouKwygkC_tyRLM,65
+etl_lib/__init__.py,sha256=q7f7YqfTzmaIBlhdhwum8vxg-YfqXBBdyqS-LS5Bq9U,65
 etl_lib/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/cli/run_tools.py,sha256=KIar-y22P4kKm-yoJjecYsPwqC7U76M71dEgFO5-ZBo,8561
 etl_lib/core/BatchProcessor.py,sha256=mRpdxZ6ZMKI8XsY3TPuy4dVcvRqLKCO-p63KeOhFyKE,3417
 etl_lib/core/ClosedLoopBatchProcessor.py,sha256=WzML1nldhZRbP8fhlD6utuK5SBYRl1cJgEobVDIdBP4,1626
-etl_lib/core/ETLContext.py,sha256=wmEnbs3n_80B6La9Py_-MHG8BN0FajE9MjGPej0A3To,8045
-etl_lib/core/ParallelBatchProcessor.py,sha256=jNo1Xv1Ts34UZIseoQLDZOhHOVeEr8dUibKUt0FJ4Hw,7318
+etl_lib/core/ETLContext.py,sha256=7dVU3qc8tK-8vbrF-pofKVih8yey7tJUiUjKKjOS28o,7625
+etl_lib/core/ParallelBatchProcessor.py,sha256=mR6N3C75NgYOVQvmioGmEkQo7RXQ0tDj13-IqH4TkeY,6067
 etl_lib/core/ProgressReporter.py,sha256=tkE-W6qlR25nU8nUoECcxZDnjnG8AtQH9s9s5WBh_-Q,9377
-etl_lib/core/SplittingBatchProcessor.py,sha256=OIRMUVFpUoZc0w__JJjUr7B9QC3sBlqQp41xghrQzC0,11616
+etl_lib/core/SplittingBatchProcessor.py,sha256=ZxLYoY41D9f1wEmiGgqSO0q1iqdnMyU98e0B9iycWh8,14909
 etl_lib/core/Task.py,sha256=muQFY5qj2n-ZVV8F6vlHqo2lVSvB3wtGdIgkSXVpOFM,9365
 etl_lib/core/ValidationBatchProcessor.py,sha256=U1M2Qp9Ledt8qFiHAg8zMxE9lLRkBrr51NKs_Y8skK8,3400
 etl_lib/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -17,7 +17,7 @@ etl_lib/data_sink/SQLBatchSink.py,sha256=vyGrrxpdmCLUZMI2_W2ORej3FLGbwN9-b2GMYHd
 etl_lib/data_sink/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/data_source/CSVBatchSource.py,sha256=0q1XdPhAIKw1HcTpnp_F4WxRUzk-24Q8Qd-WeIo5OZ0,2779
 etl_lib/data_source/CypherBatchSource.py,sha256=06WuW11BqYjAXBZqL96Qr9MR8JrcjujDpxXe8cI-SYY,2238
-etl_lib/data_source/SQLBatchSource.py,sha256=O3ZA2GXvo5j_KGwOILzguYZMPY_FJkV5j8FIa3-d9oM,4067
+etl_lib/data_source/SQLBatchSource.py,sha256=99NF0-H1tRqYY8yQBlfuF0ORMToysXpOQXrS62_KIeQ,3038
 etl_lib/data_source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/task/CreateReportingConstraintsTask.py,sha256=nTcHLBIgXz_h2OQg-SHjQr68bhH974u0MwrtWPnVwng,762
 etl_lib/task/ExecuteCypherTask.py,sha256=thE8YTZzv1abxNhhDcb4p4ke6qmI6kWR4XQ-GrCBBBU,1284
@@ -30,7 +30,7 @@ etl_lib/task/data_loading/SQLLoad2Neo4jTask.py,sha256=HR3DcjOUkQN4SbCkgQYzljQCYh
 etl_lib/task/data_loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/test_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/test_utils/utils.py,sha256=CgYOCXcUyndOdRAmGyPLoCIuEik0yzy6FLV2k16cpDM,5673
-neo4j_etl_lib-0.3.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-neo4j_etl_lib-0.3.1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-neo4j_etl_lib-0.3.1.dist-info/METADATA,sha256=Pm921qyxL36Ed_Ppp2cW3OFPxUGMv7IyRTmtba3n96o,2580
-neo4j_etl_lib-0.3.1.dist-info/RECORD,,
+neo4j_etl_lib-0.3.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+neo4j_etl_lib-0.3.2.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+neo4j_etl_lib-0.3.2.dist-info/METADATA,sha256=nLCy9Iu1V3_WTKoo551tN__hZEbcM8Aeg1YxVv2eU1M,2580
+neo4j_etl_lib-0.3.2.dist-info/RECORD,,