• 8 min read
  • 1 code sample

Polars iter_batches memory profiling for HTSUS releases

A full HTSUS release from USITC arrives as a wide, multi-hundred-thousand-row export, and calling pl.read_csv(...).collect() on it can spike resident memory to several gigabytes — enough to OOM a modest ingestion worker mid-load and strand a tariff update that clearance depends on. This page shows how to measure the true peak footprint of loading an HTSUS release with Polars, compare eager LazyFrame.collect() against streaming via iter_batches and sink_parquet, pick a batch size from the numbers rather than folklore, and hand fixed-size batches to an idempotent ingestion step so a re-run never doubles memory or duplicates rows. Every measurement uses tracemalloc for Python-object allocation plus resident-set sampling for the real process footprint, because the two disagree in ways that matter.

Prerequisites

This procedure assumes a specific toolchain and a known upstream shape. Confirm each before profiling:

  • Python 3.10+ and Polars 0.20+. The streaming engine, LazyFrame.sink_parquet, and DataFrame.iter_batches behave differently across minor versions; pin the exact build in your lockfile so a profile is reproducible.
  • A real HTSUS release file, not a sample. Peak memory scales with the widest column and the row count together. Profiling a 500-row excerpt tells you nothing about the release-day footprint. Download the current USITC HTS export (CSV or JSON) and profile that.
  • A defined column schema. Let Polars infer types once, then pin an explicit schema_overrides mapping. Duty-rate and value columns must land as Decimal or Utf8 for later exact math — never Float64, which cannot represent ISO 4217 minor units. Inference that guesses Float64 for a rate column is itself a memory and correctness bug.
  • A downstream idempotent sink. This page produces batches; it does not upsert them. The receiving step is the point-in-time schema described in HTS Schedule Database Design, and the ingestion contract that consumes these batches is owned by Tariff Update Ingestion Pipelines.
  • Structured logging configured (stdlib logging with the %(asctime)s | %(levelname)s | %(name)s | %(message)s format, or structlog) so every batch’s byte count is greppable when a profile is later questioned in an audit.

Implementation

The profiler wraps three load strategies behind one interface and records, for each, the peak Python-object allocation from tracemalloc and the peak resident-set size sampled on a background thread. Eager collection materializes the entire frame; streamed collection lets the Polars engine bound working memory; batched iteration yields fixed-row chunks you can feed straight to an upsert.

import logging
import threading
import time
import tracemalloc
from dataclasses import dataclass
from decimal import Decimal
from pathlib import Path
from typing import Callable, Iterator, Optional

import polars as pl

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger("htsus_mem_profiler")

# Pin duty/value columns to exact types. Float64 cannot hold ISO 4217 minor
# units and inflates width; Utf8 keeps the raw rate for later Decimal parsing.
HTSUS_SCHEMA: dict[str, pl.DataType] = {
    "hts_number": pl.Utf8,
    "description": pl.Utf8,
    "unit_of_quantity": pl.Utf8,
    "general_rate": pl.Utf8,
    "special_rate": pl.Utf8,
    "column_2_rate": pl.Utf8,
    "effective_date": pl.Date,
}


@dataclass(frozen=True)
class ProfileResult:
    strategy: str
    rows: int
    batches: int
    py_peak_mib: Decimal      # tracemalloc peak, Python objects only
    rss_peak_mib: Decimal     # sampled resident-set peak, whole process
    wall_seconds: Decimal


class _RssSampler(threading.Thread):
    """Sample process resident-set size every interval on a daemon thread."""

    def __init__(self, interval: float = 0.05) -> None:
        super().__init__(daemon=True)
        self._interval = interval
        self._stop = threading.Event()
        self.peak_bytes = 0

    def run(self) -> None:
        import resource
        while not self._stop.is_set():
            # ru_maxrss is KiB on Linux, bytes on macOS; normalize on Linux.
            rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024
            self.peak_bytes = max(self.peak_bytes, rss)
            time.sleep(self._interval)

    def stop(self) -> int:
        self._stop.set()
        self.join(timeout=1.0)
        return self.peak_bytes


def _mib(num_bytes: int) -> Decimal:
    return (Decimal(num_bytes) / Decimal(1024 * 1024)).quantize(Decimal("0.1"))


def _profile(strategy: str, run: Callable[[], tuple[int, int]]) -> ProfileResult:
    """Run a load strategy under tracemalloc + RSS sampling; return peaks."""
    sampler = _RssSampler()
    sampler.start()
    tracemalloc.start()
    start = time.perf_counter()

    rows, batches = run()

    wall = Decimal(str(time.perf_counter() - start)).quantize(Decimal("0.01"))
    _, py_peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    rss_peak = sampler.stop()

    result = ProfileResult(
        strategy=strategy,
        rows=rows,
        batches=batches,
        py_peak_mib=_mib(py_peak),
        rss_peak_mib=_mib(rss_peak),
        wall_seconds=wall,
    )
    logger.info(
        "profile %-14s rows=%d batches=%d py_peak=%sMiB rss_peak=%sMiB wall=%ss",
        result.strategy, result.rows, result.batches,
        result.py_peak_mib, result.rss_peak_mib, result.wall_seconds,
    )
    return result


def load_eager(path: Path) -> tuple[int, int]:
    """Baseline: whole release materialized in one frame. Worst-case peak."""
    frame = pl.read_csv(path, schema_overrides=HTSUS_SCHEMA)
    # Touch the frame so lazy columns are actually realized.
    return frame.height, 1


def load_streamed(path: Path, out: Path) -> tuple[int, int]:
    """Lazy scan sunk to Parquet via the streaming engine; bounded memory."""
    lf = pl.scan_csv(path, schema_overrides=HTSUS_SCHEMA)
    lf.sink_parquet(out)                     # never holds the full frame in RAM
    rows = pl.scan_parquet(out).select(pl.len()).collect().item()
    return rows, 1


def iter_release_batches(
    path: Path, batch_rows: int
) -> Iterator[pl.DataFrame]:
    """Yield fixed-size batches for an idempotent upsert. Peak ~= one batch."""
    reader = pl.read_csv_batched(
        path, schema_overrides=HTSUS_SCHEMA, batch_size=batch_rows
    )
    while (chunk := reader.next_batches(1)) is not None:
        yield chunk[0]


def load_batched(path: Path, batch_rows: int) -> tuple[int, int]:
    rows = 0
    batches = 0
    for batch in iter_release_batches(path, batch_rows):
        rows += batch.height
        batches += 1
        # Downstream upsert would consume `batch` here; we only count.
    return rows, batches


def compare(path: Path, work: Path, batch_rows: int) -> list[ProfileResult]:
    results = [
        _profile("eager", lambda: load_eager(path)),
        _profile("streamed", lambda: load_streamed(path, work / "hts.parquet")),
        _profile("batched", lambda: load_batched(path, batch_rows)),
    ]
    baseline = results[0].rss_peak_mib
    for r in results:
        ratio = (r.rss_peak_mib / baseline).quantize(Decimal("0.01"))
        logger.info("strategy=%s rss_peak_ratio_vs_eager=%s", r.strategy, ratio)
    return results

The important contrast is between py_peak_mib and rss_peak_mib. tracemalloc sees only Python-object allocations, so it under-reports Polars, whose columnar buffers live in Rust-owned Arrow memory outside the Python allocator. The resident-set sampler is what actually predicts an OOM. Report both, but size your worker off the RSS peak.

Resident memory over time for eager collect versus batched streaming Two resident-set-size traces plotted against wall-clock time for the same HTSUS release. The eager collect trace climbs steeply to a single tall peak near the top of the chart as the whole frame is materialized at once, then drops. The batched streaming trace stays low and flat, forming a repeating sawtooth of small equal humps, one per batch, each rising as a batch is read and falling as it is upserted and freed. A dashed horizontal line marks the worker memory ceiling; the eager peak crosses it and is flagged as an out-of-memory risk, while every batched hump stays well below it. The gap between the two shows that streaming trades a small increase in wall time for a large reduction in peak footprint. RSS (MiB) wall time worker ceiling eager collect() OOM risk batched iter_batches peak ≈ one batch batch 1 batch 4 batch n same release, same rows · streaming trades a little time for a much lower peak
Eager collect() forms one tall peak that can cross the worker ceiling; iter_batches streaming holds resident memory to roughly a single batch, a repeating sawtooth that stays well under the ceiling.

Verification steps

Run these checks against a real release before trusting a batch size in production:

  1. Establish the eager baseline. Run compare() once and record the eager rss_peak_mib. This is the number a naive collect() worker must survive. If it already exceeds your container limit, eager loading was never safe and the streamed path is mandatory, not an optimization.
  2. Confirm the RSS/Python gap. Check that rss_peak_mib is materially larger than py_peak_mib for every strategy. If they are nearly equal, tracemalloc is not seeing the Arrow buffers and your sampler interval is likely too coarse — drop it to 0.02s and re-run.
  3. Sweep batch size. Run load_batched at 25k, 50k, 100k, and 250k rows. Plot rss_peak_mib against batch_rows; the peak should rise roughly linearly with batch size while wall_seconds falls then flattens. Pick the largest batch whose peak leaves at least 30% headroom under the worker ceiling.
  4. Assert streamed bound independence. The streamed strategy’s rss_peak_mib should stay flat regardless of release size, because sink_parquet never holds the whole frame. If it tracks the eager peak, the streaming engine silently fell back to in-memory collection — check for an unsupported operation in the lazy plan.
  5. Row-count reconciliation. Assert sum(batch.height) across iter_release_batches equals the eager frame.height equals the streamed Parquet row count. A shortfall means a final partial batch was dropped by an off-by-one in the batched reader.
  6. Decimal fidelity spot-check. Parse general_rate from a batch into Decimal and confirm no column arrived as Float64. A single inferred float silently rounds a compound duty rate and corrupts everything downstream in the Duty Formula Calculation Frameworks.

Edge cases & gotchas

  • ru_maxrss is a high-water mark, not a live gauge. resource.getrusage reports the maximum RSS the process has ever reached, so a single earlier eager load poisons every later measurement in the same process. Profile each strategy in a fresh subprocess, or the batched run will inherit the eager peak and look far worse than it is.
  • tracemalloc under-reports Polars by design. Arrow-backed columns live in Rust-owned memory the Python allocator never sees. Never size a worker off py_peak_mib alone — it can be an order of magnitude below the real footprint and will greenlight a batch size that OOMs on release day.
  • Schema inference doubles the first-pass cost. Letting Polars scan the whole file to infer types reads it twice and can transiently spike memory above the eager steady state. Pin schema_overrides (as HTSUS_SCHEMA does) so inference is skipped entirely.
  • Float64 rate columns are a silent width and correctness trap. Beyond corrupting duty math, a wide float column inflates every batch. Keep rates as Utf8 through ingestion and parse to Decimal only at the point of computation.
  • Batch size interacts with the downstream transaction. A larger batch lowers per-row overhead but enlarges the upsert transaction and its lock footprint. The right size is a joint optimum between load-side memory here and the write-side contract in Tariff Update Ingestion Pipelines; tune them together, not in isolation.
  • Idempotency depends on stable batch boundaries. If a re-run reads the release with a different batch_size, batch N covers different rows and any per-batch checkpoint is meaningless. Fix batch_rows in config and treat it as part of the release’s ingestion identity, so a retried load resumes cleanly instead of double-applying.
  • sink_parquet needs the operation to be streamable. Some lazy operations (certain joins, sorts, and pivots) force the engine to collect fully, quietly defeating the memory bound. Keep the scan-to-sink path a straight projection and cast; do heavier reshaping after the data has landed.

Up: Tariff Update Ingestion Pipelines