Out of core

Nearest-neighbor matching on a propensity score for a million units in four Parquet files, never loaded at once. The score model is fitted by feglm_stream out of core, the strata are matched a batch at a time (memory="by_stratum", memory bounded by the largest stratum), and the pairs are streamed back to Parquet.

import tempfile
from pathlib import Path

import polars as pl

from match_stream import match_nn, simulate_matching

pl.Config.set_tbl_hide_dataframe_shape(True)

tmp = Path(tempfile.mkdtemp())
data, out = tmp / "extract", tmp / "output"
data.mkdir()
out.mkdir()
for part in range(4):                                     # four files, as an extract would be
    simulate_matching(250_000, seed=part).with_columns(
        (pl.col("id") + part * 250_000).alias("id")).write_parquet(data / f"part{part}.parquet")

# effect() and balance() read the source again (lazily): keep outputs out of
# the glob, and do not change the files between matching and those calls
m = match_nn("d ~ x1 + x2 + x3", str(data / "*.parquet"), exact=["r", "i"],
             replace=False, caliper=0.2, id="id", memory="by_stratum", workdir=tmp,
             verbose=False)
print("memory path:", m.memory)
m.summary()
m.matches().sink_parquet(out / "pairs.parquet")
print(f"{pl.scan_parquet(out / 'pairs.parquet').select(pl.len()).collect().item():,} pairs")
m.effect("y").summary()

Output:

memory path: by_stratum
match-stream: nn, 1:1 without replacement, estimand ATT, order largest
  distance: ps (score model: d ~ x1 + x2 + x3, logit)
  exact: r, i (20 strata)
  calipers: ps 0.2 SD
               available     matched   unmatched   discarded
  treated        253,312     252,034       1,278           0
  control        746,688     252,034     494,654           0
252,034 pairs
### y: nearest-neighbor matching on ps
Estimand: ATT   Matched focal units: 252,034
Standard errors: influence function, weights fixed, clustered by subclass (252,034)
┌─────────────┬──────────┬────────────┬────────────┬──────────┐
│ Coefficient ┆ Estimate ┆ Std. Error ┆ t value    ┆ Pr(>|t|) │
│ ---         ┆ ---      ┆ ---        ┆ ---        ┆ ---      │
│ str         ┆ f64      ┆ f64        ┆ f64        ┆ f64      │
╞═════════════╪══════════╪════════════╪════════════╪══════════╡
│ ATT:y       ┆ 1.009677 ┆ 0.003894   ┆ 259.263754 ┆ 0.0      │
└─────────────┴──────────┴────────────┴────────────┴──────────┘

Notes.

  • Memory paths. memory="auto" (the default) picks the in-memory path when the narrow columns fit the memory headroom (stream-limits), else the by-stratum path. Both give identical matches.

  • The source is read again. effect() and balance() re-read the source for the outcome and covariates. Keep it unchanged until then, and keep outputs out of its glob.

  • Scratch space. workdir= is where the score model’s files and the by-stratum copy go.