Wide designs¶
The in-memory libraries build the whole design, rows × covariates, before they start. hdfe_stream never does: it holds one bucket or one batch of rows at a time, so its memory grows with the number of covariates times the batch, not times the data. On the 8.5-million-row benchmark panel, 167 age indicators ran pyfixest and xhdfe out of memory on a 30.8 GB machine, while hdfe_stream finished in 5.3 GB (see many covariates).
Three things matter for a wide design.
1. Let hdfe_stream build the design¶
A formula term such as i(age_bin) is carried through the partitioning as the
one column it is computed from, and the indicators are built one bucket at a
time. That works for any expression whose value on a row depends on that row
alone, which formula terms always do. With the low-level
StreamingHDFE, an expression that needs the whole
column (pl.col("x") - pl.col("x").mean(), a window, a shift) is detected and
the design is built before partitioning instead, with a warning: the estimates
are the same, the memory is not. Define such a column in the input LazyFrame to
avoid it. result.diagnostics["design_evaluated"] says which happened:
"per bucket" is the one you want.
2. Scale the batches down with the width¶
At its defaults hdfe_stream holds a whole bucket of up to 20 million rows, with
every covariate, while it sorts. That is fine for a handful of covariates and
not for hundreds. Every step that holds rows holds all the covariates, so scale
rows_per_bucket, batch_rows and row_group_size down with the width. The
benchmarks aim at about 1 GB of design per bucket and 250 MB per batch and row
group, never above the defaults, so a narrow design is unaffected:
def sized_options(k, bucket_bytes=1e9, batch_bytes=250e6):
"""Bucket, batch and row-group sizes for a design of k covariates."""
per_row = 8 * (k + 4) # k covariates, plus outcome and ids
batch = int(min(2_000_000, batch_bytes // per_row))
return {"rows_per_bucket": int(min(20_000_000, bucket_bytes // per_row)),
"batch_rows": batch, "row_group_size": min(500_000, batch)}
fit = feols_stream(f"log_earn ~ i(age_bin, ref={first}) | worker_id + firm_id",
data, workdir="scratch", vcov={"CRV1": "worker_id"},
save_resid=False, **sized_options(k))
3. Expect time to grow with the square of the width¶
The cross-products cost k × k per row. They are formed chunk by chunk as BLAS matrix products, so they are fast, but 503 indicators take several times as long as 167. Disk grows with the width too, since the working copy of the rows carries every covariate.
The example builds age in months at an interview month that varies from year to year (were it fixed within a worker, fine age bins would be collinear with the worker effects), in bins of twelve, three and one months. It is small so that it runs in seconds; the point is the pattern:
bins covariates seconds disk GB design evaluated
12 mo 41 0.9 0.013 per bucket
3 mo 167 2.9 0.039 per bucket
1 mo 503 9.6 0.111 per bucket
The full script: examples/wide_designs.py
"""Many covariates: a categorical expanded into hundreds of indicators.
The in-memory libraries build the whole design -- rows x covariates -- before
they start. hdfe_stream never does: it holds one bucket or one batch of rows at
a time, so its memory grows with the number of covariates times the batch, not
times the data. On the 8.5-million-row benchmark panel, 167 age indicators ran
pyfixest and xhdfe out of memory on a 30.8 GB machine, while hdfe_stream finished
in 5.3 GB (see
benchmarks/covariates_benchmark.py and the README).
Three things matter for a wide design:
1. Let the design be built by hdfe_stream. A formula term like i(age_bin) is
carried through the partition as the one column it comes from, and the
indicators are built bucket by bucket. (With the low-level interface the
same happens for any row-wise Polars expression; one that needs the whole
column, like x - x.mean(), makes the fit build the design up front instead,
with a warning. `diagnostics["design_evaluated"]` says which happened.)
2. Scale the batch sizes down with the width. Every step that holds rows holds
all the covariates, so rows_per_bucket, batch_rows and row_group_size are
the knobs; `sized_options` below is the rule the benchmark uses.
3. Expect time to grow with the square of the number of covariates: the
cross-products are k x k per row.
This example is small so it runs in seconds; the point is the pattern.
python examples/wide_designs.py
"""
import polars as pl
from hdfe_stream import feols_stream
from simulated_data import OUTPUT, panel_path, workdir
# Age in months at an interview month that varies from year to year (were it
# fixed within a worker, fine age bins would be collinear with the worker
# effects), in bins of one to twelve months.
wide = OUTPUT / "wide.parquet"
if not wide.exists():
month = (pl.struct("worker_id", "year").hash(seed=1) % 12).cast(pl.Int32)
age_months = (pl.col("age") * 12).cast(pl.Int32) + month
(pl.scan_parquet(panel_path())
.with_columns(**{f"age_bin_{w}": age_months // w for w in (12, 3, 1)})
.sink_parquet(wide))
def sized_options(k, bucket_bytes=1e9, batch_bytes=250e6):
"""Bucket, batch and row-group sizes for a design of k covariates, so that
each holds roughly the given bytes of design whatever k is. Never above
the defaults, so a narrow design is unaffected."""
per_row = 8 * (k + 4)
batch = int(min(2_000_000, batch_bytes // per_row))
return {"rows_per_bucket": int(min(20_000_000, bucket_bytes // per_row)),
"batch_rows": batch, "row_group_size": min(500_000, batch)}
print(f"{'bins':>9} {'covariates':>10} {'seconds':>8} {'disk GB':>8} design evaluated")
for width in (12, 3, 1):
column = f"age_bin_{width}"
levels = pl.scan_parquet(wide).select(pl.col(column).unique().sort()).collect()[column]
k = len(levels) - 1
fit = feols_stream(f"log_earn ~ i({column}, ref={levels[0]}) | worker_id + firm_id",
str(wide), workdir=workdir("wide"), vcov={"CRV1": "worker_id"},
verbose=False, save_resid=False, **sized_options(k))
d = fit.diagnostics
print(f"{width:>6} mo {len(fit.coefnames):>10} {d['seconds_total']:>8.1f} "
f"{d['disk_peak_gb']:>8.3f} {d['design_evaluated']}")
fit.cleanup()