Sizing PBF Chunk Batches to a Memory Budget Jump to heading#

Replace the guessed chunk_size in a streaming OSM pipeline with a number you compute: measure the real per-element memory width of your rows, subtract fixed overhead from the RAM you are allowed to use, and let a runtime guard shrink the batch before resident memory reaches the ceiling.

Prerequisites Jump to heading#

Check these before trusting the sizing math; an unmeasured bytes_per_element is the reason a batch tuned on one extract crashes on another.

• Python 3.10+ for the type hints and float | None syntax used below.
• pandas ≥ 2.1 installed (pip install "pandas>=2.1") for DataFrame.memory_usage(deep=True).
• psutil ≥ 5.9 installed (pip install "psutil>=5.9") for the resident-set (RSS) guard.
• A representative sample of parsed OSM rows — the same shape produced upstream by Async PBF Parsing with Pyrosm — so the width you measure matches the width you will stream.
• The buffered-flush pattern from Memory-Efficient Chunk Processing, whose chunk_size this page computes instead of hard-codes.
• A known memory budget in bytes (container limit, cgroup memory.max, or a self-imposed fraction of total RAM).

Conceptual minimum Jump to heading#

Every streaming OSM job that flushes fixed-size batches has one dial that decides whether it survives a dense extract: how many rows it lets accumulate before spilling. Pick it too small and you drown in per-flush overhead and tiny row groups; too large and the buffer plus its transient copy blow past the RAM ceiling and the kernel OOM-killer ends the run. The mistake is treating that dial as a constant to be guessed, when it is really the output of a short calculation. If you know how many bytes one buffered element actually occupies, and how many bytes you are permitted to spend, the safe batch size follows directly. This is the same budgeting instinct that the parent Memory-Efficient Chunk Processing guide applies to the buffer as a whole — here it is made explicit and measured per element.

The width of an OSM row is dominated by its tag payload, and tags are contributor-defined free text, so the only reliable per-element figure is a measured one. DataFrame.memory_usage(deep=True) walks the actual Python objects behind each column — including the variable-length strings a shallow measurement misses — and dividing its total by the row count gives a faithful bytes-per-element for your real data rather than a schema guess. Reserve a fixed overhead for the interpreter, imported libraries, and any long-lived caches (the location store, for instance), and the batch size is a floor division:

$$\text{chunk\_size} = \left\lfloor \frac{B_{\text{budget}} - B_{\text{overhead}}}{\bar{w}_{\text{element}}} \right\rfloor$$

where $B_{\text{budget}}$ is the RAM you can spend, $B_{\text{overhead}}$ is the fixed resident floor, and $\bar{w}_{\text{element}}$ is the measured deep bytes per row. Because a flush briefly holds the buffer and its serialized copy at once, divide the numerator by a small safety factor (or the measured peak-to-buffer ratio) so the transient copy still fits. A static number computed this way is correct only for the data you measured, though — tag density drifts between rural and urban tiles — so pair it with a runtime psutil guard that samples RSS and shrinks the next batch when the process creeps toward the ceiling. The static formula sets the starting point; the adaptive feedback keeps a mis-measurement from becoming a crash.

Runnable solution Jump to heading#

The module below measures per-element width from a sample DataFrame, computes a starting chunk_size from the budget, and wraps accumulation in an adaptive guard that halves the batch target whenever RSS climbs past a high-water fraction of the budget. The measurement and the guard are independent, so you can use the sizing math alone, the guard alone, or both.

from __future__ import annotations

import logging
import sys

import pandas as pd
import psutil

logger = logging.getLogger(__name__)


def bytes_per_element(sample: pd.DataFrame) -> float:
    """Measured deep memory width of one row, in bytes.

    ``memory_usage(deep=True)`` follows the Python objects behind object and
    string columns (the variable-length tag strings a shallow count misses),
    so the per-row figure reflects real data rather than a schema estimate.
    """
    if sample.empty:
        raise ValueError("need a non-empty sample to measure element width")
    total = int(sample.memory_usage(deep=True).sum())
    return total / len(sample)


def sizing_overhead() -> int:
    """A conservative fixed floor: current process RSS at measurement time."""
    return psutil.Process().memory_info().rss


def compute_chunk_size(
    budget_bytes: int,
    element_bytes: float,
    overhead_bytes: int,
    safety: float = 2.0,
) -> int:
    """chunk_size = floor((budget - overhead) / (element_bytes * safety)).

    ``safety`` reserves room for the transient copy a flush holds alongside
    the live buffer; 2.0 mirrors the buffer-plus-DataFrame peak.
    """
    spendable = budget_bytes - overhead_bytes
    if spendable <= 0:
        raise ValueError("overhead already exceeds the budget; raise the budget")
    size = int(spendable // (element_bytes * safety))
    if size < 1:
        raise ValueError("budget too small for even one row at this width")
    logger.info(
        "sized chunk: %d rows (%.1f B/elem, %.2f GiB spendable)",
        size, element_bytes, spendable / 2**30,
    )
    return size


class AdaptiveBatcher:
    """Accumulate rows to a computed target, shrinking it as RSS rises.

    When resident memory crosses ``high_water`` of the budget the target is
    halved before the next batch, so a per-element under-estimate (dense
    urban tags) cannot walk the process into an OOM kill.
    """

    def __init__(self, budget_bytes: int, target: int, high_water: float = 0.85) -> None:
        self.budget_bytes = budget_bytes
        self.target = max(1, target)
        self.high_water = high_water
        self._proc = psutil.Process()

    def _rss_fraction(self) -> float:
        return self._proc.memory_info().rss / self.budget_bytes

    def should_flush(self, buffered_rows: int) -> bool:
        if buffered_rows >= self.target:
            return True
        # Pre-emptive flush if we are already close to the ceiling.
        return self._rss_fraction() >= self.high_water and buffered_rows > 0

    def note_after_flush(self) -> None:
        """Shrink the target if the last cycle pushed RSS toward the ceiling."""
        frac = self._rss_fraction()
        if frac >= self.high_water and self.target > 1:
            self.target = max(1, self.target // 2)
            logger.warning(
                "RSS at %.0f%% of budget; target shrunk to %d rows",
                frac * 100, self.target,
            )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    # 1. Measure width from a representative sample of parsed OSM rows.
    sample = pd.read_parquet("sample_rows.parquet")  # a few thousand real rows
    width = bytes_per_element(sample)
    logger.info("cross-check: sys.getsizeof(sample)=%d", sys.getsizeof(sample))

    # 2. Compute the starting batch size against an 8 GiB budget.
    budget = 8 * 2**30
    target = compute_chunk_size(budget, width, sizing_overhead())

    # 3. Drive accumulation with the adaptive guard.
    batcher = AdaptiveBatcher(budget, target)
    buffer: list[dict] = []
    for row in stream_parsed_rows():   # your pyosmium / pyrosm producer
        buffer.append(row)
        if batcher.should_flush(len(buffer)):
            flush_to_parquet(buffer)   # your writer from the chunk-processing stage
            buffer.clear()
            batcher.note_after_flush()
    if buffer:
        flush_to_parquet(buffer)

Step-by-step walkthrough Jump to heading#

1. Measure, do not guess — bytes_per_element sums memory_usage(deep=True) and divides by row count, so variable-length tag strings are counted at their true cost rather than as an opaque pointer.
2. Overhead is the resident floor — sizing_overhead() snapshots the current process RSS (interpreter, imports, any warm caches) so the budget the batch draws on is what is actually left, not the raw total.
3. The safety factor covers the flush copy — dividing by element_bytes * 2.0 reserves headroom for the moment the live buffer and its serialized DataFrame coexist, the same factor-of-two peak the parent stage describes.
4. compute_chunk_size fails loud — if overhead already exceeds the budget, or the budget cannot hold a single row, it raises rather than returning a zero that would spin forever.
5. should_flush has two triggers — the normal one is reaching the row target; the second is a pre-emptive flush when RSS is already past high_water, which catches a batch whose rows turned out wider than the sample.
6. note_after_flush closes the loop — if a cycle ended near the ceiling it halves the target, so the process converges to a sustainable batch size instead of repeatedly brushing the limit.
7. sys.getsizeof is a sanity cross-check — it reports the shallow container size only, so a large gap between it and the deep figure confirms the tag strings dominate and the deep measurement was necessary.

Verification Jump to heading#

Prove the size is safe before committing it to a long run:

• Recompute is stable. Measure bytes_per_element on two disjoint samples; the figures should agree within a few percent, or your sample is not representative.
• The formula matches observed peak. Run one batch at the computed chunk_size and read RSS at the flush; peak should land near overhead + chunk_size × element_bytes × safety, not far above it.
• The guard actually shrinks. Feed a deliberately dense (urban) tile and confirm a target shrunk to warning appears and the run still completes without an OOM kill.
• No premature flushing on sparse data. On a rural tile, should_flush should fire on the row target, not the RSS trigger — if it fires early, high_water is set too low.
• Determinism of the static size. For a fixed sample, budget, and overhead, compute_chunk_size must return the same integer every time.

Common errors and fixes Jump to heading#

Specification reference Jump to heading#

Per-element width must be measured deeply, because the default is shallow: pandas documents that DataFrame.memory_usage(deep=True) will “introspect the data deeply by interrogating object dtypes for system-level memory consumption, and include it in the returned values” — see the pandas memory_usage documentation. The runtime ceiling is read from the operating system: the resident set size returned by psutil.Process().memory_info() is “the non-swapped physical memory a process has used,” which is the figure a container’s cgroup limit and the OOM-killer both watch, making it the correct signal for the adaptive guard.

Frequently Asked Questions Jump to heading#

Related Jump to heading#

• Memory-Efficient Chunk Processing — the streaming buffer whose chunk_size this page computes and guards.
• A Bounded LRU Node Cache for OSM Streaming — the companion lever that caps the location store rather than the record batch.
• Async PBF Parsing with Pyrosm — the parallel producer whose row shape you should sample when measuring element width.
• Batch Attribute Mapping Strategies — the tag registries whose payload width drives most of the per-element cost.
• Error Handling in Large OSM Extracts — where a batch that overruns its budget routes its overflow and defective rows.
• PBF File Structure Deep Dive — why a block-streamed source lets a computed batch size hold across the whole file.

Up one level: Memory-Efficient Chunk Processing.