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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-etl-lib
-Version: 0.3.0
+Version: 0.3.2
 Summary: Building blocks for ETL pipelines.
 Keywords: etl,graph,database
 Author-email: Bert Radke <bert.radke@pm.me>
@@ -14,11 +14,11 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Database
 Classifier: Development Status :: 4 - Beta
 License-File: LICENSE
-Requires-Dist: pydantic>=2.10.5; python_version >= '3.8'
-Requires-Dist: neo4j-rust-ext>=5.27.0; python_version >= '3.7'
-Requires-Dist: python-dotenv>=1.0.1; python_version >= '3.8'
-Requires-Dist: tabulate>=0.9.0; python_version >= '3.7'
-Requires-Dist: click>=8.1.8; python_version >= '3.7'
+Requires-Dist: pydantic>=2.10.5; python_version >= '3.10'
+Requires-Dist: neo4j-rust-ext>=5.27.0,<6; python_version >= '3.10'
+Requires-Dist: python-dotenv>=1.0.1; python_version >= '3.10'
+Requires-Dist: tabulate>=0.9.0; python_version >= '3.10'
+Requires-Dist: click>=8.1.8; python_version >= '3.10'
 Requires-Dist: pydantic[email-validator]
 Requires-Dist: pytest>=8.3.0 ; extra == "dev" and ( python_version >= '3.8')
 Requires-Dist: testcontainers[neo4j]==4.9.0 ; extra == "dev" and ( python_version >= '3.9' and python_version < '4.0')
@@ -35,11 +35,13 @@ Requires-Dist: sphinx-autoapi ; extra == "dev"
 Requires-Dist: sqlalchemy ; extra == "dev"
 Requires-Dist: psycopg2-binary ; extra == "dev"
 Requires-Dist: graphdatascience>=1.13 ; extra == "gds" and ( python_version >= '3.9')
+Requires-Dist: nox>=2024.0.0 ; extra == "nox"
 Requires-Dist: sqlalchemy ; extra == "sql"
 Project-URL: Documentation, https://neo-technology-field.github.io/python-etl-lib/index.html
 Project-URL: Home, https://github.com/neo-technology-field/python-etl-lib
 Provides-Extra: dev
 Provides-Extra: gds
+Provides-Extra: nox
 Provides-Extra: sql
 
 # Neo4j ETL Toolbox

{neo4j_etl_lib-0.3.0 → neo4j_etl_lib-0.3.2}/pyproject.toml

@@ -22,11 +22,11 @@ dynamic = ["version", "description"]
 keywords = ["etl", "graph", "database"]
 
 dependencies = [
-    "pydantic>=2.10.5; python_version >= '3.8'",
-    "neo4j-rust-ext>=5.27.0; python_version >= '3.7'",
-    "python-dotenv>=1.0.1; python_version >= '3.8'",
-    "tabulate>=0.9.0; python_version >= '3.7'",
-    "click>=8.1.8; python_version >= '3.7'",
+    "pydantic>=2.10.5; python_version >= '3.10'",
+    "neo4j-rust-ext>=5.27.0,<6; python_version >= '3.10'",
+    "python-dotenv>=1.0.1; python_version >= '3.10'",
+    "tabulate>=0.9.0; python_version >= '3.10'",
+    "click>=8.1.8; python_version >= '3.10'",
     "pydantic[email_validator]"
 ]
 
@@ -41,6 +41,11 @@ dev = [
 gds = ["graphdatascience>=1.13; python_version >= '3.9'"]
 sql = ["sqlalchemy"]
 
+# Local-only multy-version testing, install via `pip install ".[dev,nox]"`
+nox = [
+    "nox>=2024.0.0"
+]
+
 [project.urls]
 Home = "https://github.com/neo-technology-field/python-etl-lib"
 Documentation = "https://neo-technology-field.github.io/python-etl-lib/index.html"

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

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

{neo4j_etl_lib-0.3.0 → 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.0 → neo4j_etl_lib-0.3.2}/src/etl_lib/core/ProgressReporter.py

@@ -45,7 +45,7 @@ class ProgressReporter:
             The task that was provided.
         """
         task.start_time = datetime.now()
-        self.logger.info(f"{'\t' * task.depth}starting {task.task_name()}")
+        self.logger.info(f"{' ' * (4 * task.depth)}starting {task.task_name()}")
         return task
 
     def finished_task(self, task: Task, result: TaskReturn) -> Task:

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.0 → neo4j_etl_lib-0.3.2}/src/etl_lib/data_source/CSVBatchSource.py

@@ -30,7 +30,7 @@ class CSVBatchSource(BatchProcessor):
         self.csv_file = csv_file
         self.kwargs = kwargs
 
-    def get_batch(self, max_batch__size: int) -> Generator[BatchResults]:
+    def get_batch(self, max_batch__size: int) -> Generator[BatchResults, None, None]:
         for batch_size, chunks_ in self.__read_csv(self.csv_file, batch_size=max_batch__size, **self.kwargs):
             yield BatchResults(chunk=chunks_, statistics={"csv_lines_read": batch_size}, batch_size=batch_size)