Connected sets

Connected sets shows these at work.

hdfe_stream.connected_components(data, fe, *, min_degree=1, leave_one_out=False, rank_by='rows', batch_rows=2000000, verbose=True, logger=None, log_level=20)

Connected components of the graph of two fixed-effect dimensions, and so the largest connected set.

The graph has one vertex per level of each dimension, and each row joins its two levels. Its components are found by union-find over the rows as they stream.

Parameters:

• data (str, Path, polars.LazyFrame or polars.DataFrame) – Parquet path or glob, Polars LazyFrame, or DataFrame. It is read in batches and never held in memory.

• fe (str or (str, str)) – The two dimensions, as a formula’s fixed-effect part ("worker_id + firm_id") or a pair of names (("worker_id", "firm_id")). Either may be an interaction ("firm_id^year"). Rows with a null or NaN in any of their columns are in no component.

• min_degree (int or (int, int), default 1) – Remove every level with fewer distinct partners in the other dimension (distinct firms for a worker, distinct workers for a firm) than this, repeating until none is left, before finding the components. An int applies to both dimensions; a pair is (first, second). (2, 1) keeps only movers: workers seen at two or more firms, and the firms they were seen at. 2 also drops firms with a single worker, and so on until every worker and firm left has two partners. (1, k) keeps firms with at least k distinct workers. Partners are distinct levels, so a worker with many rows at one firm has one; with firm_id^year, a worker who stays at a firm for two years has two. The default, 1, removes nothing. min_degree=2 is not the leave-one-out connected set: see below.

• leave_one_out (bool, default False) –

  Prune to the leave-one-out connected set that Kline-Saggio-Solvsten estimation needs (leave_out_kss applies it), in which every effect stays estimable when any one worker is dropped. The first dimension plays the workers. Following LeaveOutTwoWay’s pruning_unbal_v3.m and VarianceComponentsHDFE.jl:

  1. take the largest connected component;

  2. drop workers observed only once (their leverage is exactly 1);

  3. remove every worker whose removal would split the network (an articulation point), and take the largest piece left;

  4. repeat 3 until no such worker is left.

  Step 3 repeats because removing one such worker can make another. The result is one component, 0; every other level’s component is null. Unlike min_degree=2 it keeps stayers seen at least twice, and it removes movers that min_degree=2 keeps, so neither set contains the other, although every mover and firm it keeps min_degree=2 keeps too. With min_degree, the pruning by degree comes first. Expect it to remove a lot: in KSS’s own application, about half the firms, mostly ones attached to a single mover.

• rank_by (str, default "rows") – What “largest” means: "rows", or the name of either dimension to count its levels. Ties go to the component whose smallest level of the first dimension sorts first.

• batch_rows (int, default 2_000_000) – Rows per batch as the data streams.

• verbose – Progress messages, as for feols_stream.

• logger – Progress messages, as for feols_stream.

• log_level – Progress messages, as for feols_stream.

Returns:

ConnectedComponents – Small level tables, not rows. Its ConnectedComponents.largest and ConnectedComponents.label apply them to the data lazily; it keeps a reference to data for that (a DataFrame is then kept alive with it), and either takes other data with the same levels instead.

Notes

Pruning comes first, then the components: the components are those of what is left, so the largest one is connected and every level in it has at least min_degree partners in it. (Taking the largest component first and pruning it can disconnect it, once min_degree exceeds 2.)

Memory. Without pruning: the level tables themselves, about 20 bytes more per level, and a batch of rows. Pruning also holds the distinct pairs, at 12 bytes each, and briefly about twice that while they are merged; leave_one_out adds the network’s adjacency, about 8 bytes per pair more, and 50 per level. In practice the peak is set by Polars finding the distinct levels and reading the Parquet: on 150 million rows with 25 million 9-character worker ids, about 1.3 GB above what a bare scan of the two columns takes.

Three or more dimensions. Connectedness of one pair is necessary for the effects to be identified but not sufficient: a further dimension can make levels of the pair redundant that no graph of the pair shows (year effects, say, when every move between firms coincides with a change of year). The components here are those of the pair only; they are not combined with any other dimension, which would be wrong rather than conservative, since two firms sharing years are not comparable unless a worker moves between them.

In a fit. The same set is what feols_stream(..., connected="largest") keeps, without pruning, of the first two fixed effects listed in the formula, but found among the rows the model can use (rows with a missing outcome or covariate do not connect anything there). Filter those out before calling this to get the same rows.

Examples

>>> cc = connected_components("data/*.parquet", "worker_id + firm_id")
>>> cc.sizes                       # rows and levels per component
>>> akm = cc.largest()             # the largest connected set, lazily
>>> rows = cc.label()              # every row, with its component

class hdfe_stream.ConnectedComponents

The connected components of the graph of two fixed-effect dimensions, as returned by connected_components: small level tables, not rows.

Attributes:

• fe ((str, str)) – The two dimensions, e.g. ("worker_id", "firm_id^year").

• levels (dict of str to polars.DataFrame) – {dimension: DataFrame}: every level of the dimension found in a complete row (its column, or columns for an interaction), and its component. Components are numbered by size, so component == 0 is the largest; levels removed by min_degree or leave_one_out have a null component.

• sizes (polars.DataFrame) – One row per component, in order: component, rows, and the number of levels of each dimension (a column named after it).

• min_degree ((int, int)) – (first dimension, second dimension), as used.

• leave_one_out (bool) – Whether this is the leave-one-out connected set.

• diagnostics (dict) – Counts and timings (see summary).

property n_components

The number of components.

largest(data=None)

The rows of the data in the largest component, lazily and in their original order: the largest connected set.

By default the data is the data the components were found in, with all its columns. Other data must be a subset of it, or at least have the same levels: each dimension is filtered with whichever is smaller, its levels in the largest component (a semi-join) or the rest (an anti-join), so a level that was never seen may be let through. Without min_degree, one dimension’s table settles it, since both levels of a row are in the same component.

label(data=None)

The data (by default the data the components were found in, as for largest), lazily and in its original order, with two columns added: component (UInt32), the row’s component, null when the row is in none (a null or NaN level, or a level removed by min_degree or leave_one_out); and in_largest, whether the row is in the largest component.

It joins the full level table of the dimension with fewer levels (and, after pruning, the smaller list of the other’s removed or surviving levels), so it holds that table, not the rows.

summary()

A short report: the components, the largest, and what pruning removed.

__init__(fe, levels, sizes, min_degree=(1, 1), leave_one_out=False, diagnostics=<factory>, _after_min_degree=None, _source=None)

classmethod __new__(*args, **kwargs)