diff-stream 0.1.0.dev0__py3-none-any.whl

diff_stream/_core/partition.py

@@ -0,0 +1,158 @@
+"""Bounded whole-group hash partitions, independent of panel/DiD semantics.
+
+Consumers implement _leaf and optionally _oversized_group. Rows of a group
+always remain together; the decoded leaf budget can include projected columns.
+Intended for eventual hdfe_stream.core.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+import pyarrow.parquet as pq
+
+
+def preprocess_budget_mb(preprocess_mb, budget):
+    """The working memory for partitioning and sorting, in decimal MB.
+
+    An explicit value is validated and used as given. None (the default) uses
+    a quarter of the job's current memory headroom, between 16 MB and 32 GB,
+    so a panel that fits in memory is sorted in one piece instead of being
+    hash-partitioned through disk, and a smaller ``max_memory_gb`` gives
+    smaller partitions automatically. A quarter, not more: headroom shrinks
+    while the data are staged, and each leaf is checked against it again.
+    """
+    if preprocess_mb is None:
+        return float(min(32000.0, max(16.0, budget.headroom_gb() * 250)))
+    if isinstance(preprocess_mb, bool) or not np.isfinite(preprocess_mb) or preprocess_mb < 1:
+        raise ValueError("preprocess_mb must be finite and at least 1 MB, or None")
+    return float(preprocess_mb)
+
+
+def staging_bytes(preprocess_mb, store):
+    """Expected memory of staging: the preprocessing allowance, or less when the
+    staged rows are small. Sorting a panel in one piece peaked at about six
+    times its rows' size in the measurements (more on small panels, where the
+    runtime's own use dominates), so eight times plus a fixed 0.4 GB covers
+    them; the allowance stays the upper bound."""
+    rows = getattr(store, "n_rows_used", None)
+    if not rows:
+        return preprocess_mb * 1e6
+    design = getattr(store, "design", None)
+    columns = 4 + (len(design.names) if design is not None else 0)
+    columns += int(getattr(store, "weighted", False))
+    return min(preprocess_mb * 1e6, 8.0 * rows * 8 * columns + 0.4e9)
+
+
+@dataclass
+class Part:
+    path: Path
+    rows: int
+    size: int
+
+
+def whole_group_batches(path, key, max_groups, read_rows=32768):
+    """Read sorted Parquet without splitting a group across yielded frames.
+
+    Memory is bounded by one reader batch, one unfinished group, and the
+    consumer's batch. Keys need only be unique within this file.
+    """
+    pending = None
+    with pq.ParquetFile(path, pre_buffer=False) as pf:
+        for batch in pf.iter_batches(batch_size=read_rows):
+            frame = pl.from_arrow(batch)
+            if pending is not None:
+                frame = pl.concat([pending, frame])
+            values = frame[key].to_numpy()
+            starts = np.r_[0, np.flatnonzero(values[1:] != values[:-1]) + 1]
+            end = int(starts[-1])
+            for j in range(0, len(starts) - 1, max_groups):
+                stop = int(starts[min(j + max_groups, len(starts) - 1)])
+                yield frame.slice(int(starts[j]), stop - int(starts[j]))
+            pending = frame.slice(end)
+        if pending is not None:
+            yield pending
+
+
+class GroupPartitions:
+    def __init__(self, root, budget, preprocess_mb, *, key="id"):
+        self.root = Path(root)
+        self.budget = budget
+        self.key = key
+        self.leaf_bytes = int(preprocess_mb * 1e6 / 4)
+        self.read_rows = max(256, min(262144, self.leaf_bytes // 128))
+        self.extra_row_bytes = 0
+        self.max_leaf_bytes = 0
+        self.max_depth = 0
+        self.partition_count = 0
+        self._serial = 0
+
+    def _oversized_group(self, part):
+        raise MemoryError("a single group exceeds preprocess_mb; increase the preprocessing budget")
+
+    def _guard(self, size):
+        if not self.budget.fits(size / 1e9):
+            threads = pl.thread_pool_size()
+            hint = (
+                f" (Polars is using {threads} threads, and the memory it holds while reading grows "
+                "with them, about 130 MB each for row groups of a million rows; set "
+                "STREAM_LIMITS_MAX_THREADS or POLARS_MAX_THREADS before Python starts to lower it)"
+                if threads > 4
+                else ""
+            )
+            raise MemoryError(
+                "preprocessing buffers exceed job headroom; reduce preprocess_mb or request more memory"
+                + hint
+            )
+
+    def _name(self, prefix):
+        self._serial += 1
+        return self.root / f"{prefix}-{self._serial}.parquet"
+
+    def _prepare(self, part, depth):
+        effective = part.size + part.rows * self.extra_row_bytes
+        if effective <= self.leaf_bytes:
+            self._leaf(part)
+            return
+        if depth >= 16:
+            raise MemoryError("unit/hash skew exceeds the partition budget; increase preprocess_mb")
+        self.max_depth = max(self.max_depth, depth + 1)
+        fanout = min(16, max(2, int(np.ceil(2 * effective / self.leaf_bytes))))
+        children = {}
+        writers = {}
+        first_id, same_id = None, True
+        try:
+            with pq.ParquetFile(part.path, pre_buffer=False) as pf:
+                for b in pf.iter_batches(batch_size=self.read_rows):
+                    self._guard(32e6)
+                    d = pl.from_arrow(b)
+                    if first_id is None:
+                        first_id = d[self.key][0]
+                    same_id = same_id and (d[self.key] == first_id).all()
+                    d = d.with_columns(
+                        (pl.col(self.key).hash(seed=72931 + depth) % fanout).alias("bucket")
+                    )
+                    for key, frame in d.partition_by("bucket", as_dict=True).items():
+                        k = key[0]
+                        table = frame.drop("bucket").to_arrow()
+                        if k not in writers:
+                            p = self._name("hash")
+                            children[k] = Part(p, 0, 0)
+                            writers[k] = pq.ParquetWriter(
+                                p, table.schema, compression="zstd", use_dictionary=False
+                            )
+                        writers[k].write_table(table, row_group_size=self.read_rows)
+                        children[k].rows += table.num_rows
+                        children[k].size += table.nbytes
+        finally:
+            for w in writers.values():
+                w.close()
+        part.path.unlink()
+        if same_id:
+            self._oversized_group(part)
+        # Do not retain the parent's last decoded batch and writer metadata
+        # while recursively processing its children.
+        del writers, d, frame, table, b, w
+        for child in children.values():
+            self._prepare(child, depth + 1)

diff_stream/_core/process.py

@@ -0,0 +1,61 @@
+"""Small standard-library process batch, intended for eventual shared core.
+
+Threads only wait on child processes; numerical work belongs in the children.
+On failure or interruption, stop and reap children before consumers delete
+their scratch. Commands are argument vectors, never shell strings.
+"""
+
+import subprocess
+from concurrent.futures import CancelledError, ThreadPoolExecutor
+from threading import Lock
+
+
+class ProcessBatch:
+    def __init__(self, workers):
+        self._executor = ThreadPoolExecutor(max_workers=workers)
+        self._lock = Lock()
+        self._children = set()
+        self._closing = False
+
+    def __enter__(self):
+        return self
+
+    def submit(self, command, log):
+        return self._executor.submit(self._run, command, log)
+
+    def _run(self, command, log):
+        with open(log, "wb") as output:
+            with self._lock:
+                if self._closing:
+                    raise CancelledError()
+                child = subprocess.Popen(
+                    command, stdin=subprocess.DEVNULL, stdout=output, stderr=subprocess.STDOUT
+                )
+                self._children.add(child)
+            try:
+                return child.wait()
+            finally:
+                with self._lock:
+                    self._children.discard(child)
+
+    def __exit__(self, exc_type, exc, tb):
+        with self._lock:
+            self._closing = True
+            children = list(self._children)
+        # On normal exit all submitted futures have already completed. A
+        # consumer exception/interrupt must not leave numerical fits running.
+        for child in children:
+            try:
+                child.terminate()
+            except ProcessLookupError:
+                pass
+        for child in children:
+            try:
+                child.wait(timeout=2)
+            except subprocess.TimeoutExpired:
+                try:
+                    child.kill()
+                except ProcessLookupError:
+                    pass  # Exited between the timeout and kill.
+                child.wait()
+        self._executor.shutdown(wait=True, cancel_futures=True)

diff_stream/_core/resample.py

@@ -0,0 +1,171 @@
+"""Bounded, disk-backed ordinary whole-group resampling.
+
+Generic adapter intended for eventual hdfe_stream.core: consumers own sample
+selection and model semantics. Uniform draws use N independent integer indices;
+copied groups receive distinct draw IDs, rather than frequency weights.
+"""
+
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+import pyarrow.parquet as pq
+
+from .partition import GroupPartitions, Part
+
+
+@dataclass
+class ResampleLeaf:
+    path: Path
+    start: int
+    stop: int
+
+
+class DiskGroupResampler(GroupPartitions):
+    def __init__(
+        self,
+        snapshot,
+        root,
+        budget,
+        preprocess_mb,
+        *,
+        key,
+        sort_columns,
+        select,
+        index_column,
+        draw_column,
+    ):
+        super().__init__(root, budget, preprocess_mb, key=key)
+        self.sort_columns, self.select = sort_columns, select
+        self.index_column, self.draw_column = index_column, draw_column
+        self.leaves = []
+        self.n_units = self.n_rows = 0
+        self.identity = hashlib.sha256()
+        size = 0
+        with pq.ParquetFile(snapshot, pre_buffer=False) as pf:
+            rows = pf.metadata.num_rows
+            for batch in pf.iter_batches(batch_size=self.read_rows):
+                self._guard(32e6 + batch.nbytes * 3)
+                size += batch.nbytes
+        self._prepare(Part(Path(snapshot), rows, size), 0)
+        self.ends = np.array([leaf.stop for leaf in self.leaves], dtype=np.int64)
+
+    def _leaf(self, part):
+        self._guard(part.size * 4 + 32e6)
+        with pq.ParquetFile(part.path, pre_buffer=False) as pf:
+            data = pl.from_arrow(pf.read()).sort(self.sort_columns)
+        data = self.select(data)
+        if len(data):
+            ids = data[self.key].to_numpy()
+            starts = np.r_[0, np.flatnonzero(ids[1:] != ids[:-1]) + 1]
+            lengths = np.diff(np.r_[starts, len(data)])
+            for value in ids[starts]:
+                encoded = repr(value.item() if isinstance(value, np.generic) else value).encode()
+                self.identity.update(len(encoded).to_bytes(8, "little"))
+                self.identity.update(encoded)
+            data = data.with_columns(
+                pl.Series(
+                    self.index_column,
+                    np.repeat(np.arange(self.n_units, self.n_units + len(starts)), lengths),
+                )
+            )
+            path = self._name("resample-source")
+            data.write_parquet(path, row_group_size=32768)
+            self.leaves.append(ResampleLeaf(path, self.n_units, self.n_units + len(starts)))
+            self.n_units += len(starts)
+            self.n_rows += len(data)
+            self.partition_count += 1
+            self.max_leaf_bytes = max(self.max_leaf_bytes, part.size)
+        part.path.unlink()
+
+    def materialize(self, rng, root, *, time_column, chunk_rows=32768):
+        """Write one exact N-group resample. Return path, dates, rows, plan hash.
+
+        Integer plans are routed to leaf-specific binary files, never collected
+        into an N-row in-memory manifest. Output buffers and join multiplicity
+        are bounded separately from the number of groups in the population.
+        """
+        root = Path(root)
+        root.mkdir()
+        digest = hashlib.sha256()
+        for offset in range(0, self.n_units, chunk_rows):
+            count = min(chunk_rows, self.n_units - offset)
+            indices = rng.integers(0, self.n_units, size=count, dtype=np.int64)
+            digest.update(indices.astype("<i8", copy=False).tobytes())
+            buckets = np.searchsorted(self.ends, indices, side="right")
+            order = np.argsort(buckets, kind="stable")
+            ordered = buckets[order]
+            starts = np.r_[0, np.flatnonzero(ordered[1:] != ordered[:-1]) + 1, count]
+            for a, b in zip(starts[:-1], starts[1:], strict=True):
+                chosen = order[a:b]
+                pairs = np.column_stack([indices[chosen], offset + chosen]).astype("<i8")
+                with (root / f"plan-{ordered[a]}.bin").open("ab") as f:
+                    pairs.tofile(f)
+        output = root / "resample.parquet"
+        writer, pending, pending_rows, pending_bytes = None, [], 0, 0
+        dates, rows = set(), 0
+        try:
+            for j, leaf in enumerate(self.leaves):
+                plan = root / f"plan-{j}.bin"
+                if not plan.exists():
+                    continue
+                with pq.ParquetFile(leaf.path, pre_buffer=False) as pf:
+                    data = pl.from_arrow(pf.read())
+                ids = data[self.index_column].to_numpy()
+                starts = np.r_[0, np.flatnonzero(ids[1:] != ids[:-1]) + 1]
+                lengths = np.diff(np.r_[starts, len(data)])
+                leaf_bytes = data.estimated_size()
+
+                def blocks(pairs, leaf=leaf, lengths=lengths, leaf_bytes=leaf_bytes, data=data):
+                    counts = np.bincount(pairs[:, 0] - leaf.start, minlength=leaf.stop - leaf.start)
+                    nrows = int(counts @ lengths)
+                    upper = leaf_bytes * int(counts.max()) + 16 * nrows
+                    if upper > self.leaf_bytes * 2 and len(pairs) > 1:
+                        mid = len(pairs) // 2
+                        yield from blocks(pairs[:mid])
+                        yield from blocks(pairs[mid:])
+                        return
+                    self._guard(4 * upper + 2 * leaf_bytes + 32e6)
+                    selection = pl.DataFrame(
+                        {self.index_column: pairs[:, 0], self.draw_column: pairs[:, 1]}
+                    )
+                    yield data.join(selection, on=self.index_column, how="inner").drop(
+                        self.index_column
+                    )
+
+                with plan.open("rb") as f:
+                    while True:
+                        pairs = np.fromfile(f, dtype="<i8", count=512).reshape(-1, 2)
+                        if not len(pairs):
+                            break
+                        for frame in blocks(pairs):
+                            dates.update(frame[time_column].unique().to_list())
+                            rows += len(frame)
+                            if writer is None:
+                                writer = pq.ParquetWriter(
+                                    output,
+                                    frame.to_arrow().schema,
+                                    compression="zstd",
+                                    use_dictionary=False,
+                                )
+                            pending.append(frame)
+                            pending_rows += len(frame)
+                            pending_bytes += frame.estimated_size()
+                            if (
+                                pending_rows >= 32768
+                                or pending_bytes >= self.leaf_bytes * 2
+                                or len(pending) >= 128
+                            ):
+                                writer.write_table(
+                                    pl.concat(pending).to_arrow(), row_group_size=32768
+                                )
+                                pending, pending_rows, pending_bytes = [], 0, 0
+                plan.unlink()
+            if pending:
+                writer.write_table(pl.concat(pending).to_arrow(), row_group_size=32768)
+        finally:
+            if writer is not None:
+                writer.close()
+        return output, sorted(dates), rows, digest.hexdigest()

diff_stream/_core/staging_kernels.py

@@ -0,0 +1,90 @@
+"""One compiled pass over a sorted unit-period table for the staging checks.
+
+Staging validates the panel and sums outcomes and weights by cohort and period.
+In NumPy that is about fifteen passes over every row, each single-threaded.
+The kernel here does the checks and sums in one parallel pass over chunks of
+units; the chunks' partial sums are added in order, so the result does not
+depend on the number of threads. ``PanelStore._leaf`` keeps the NumPy code for
+the cases the kernel does not cover (treatment columns, spaced periods,
+complete-case selection, original cohorts), and the tests compare the two.
+"""
+
+import numpy as np
+from numba import njit, prange
+
+ENABLED = True  # tests switch it off to compare with the NumPy path
+CHUNKS = 256  # at most; fixed so results do not depend on the number of threads
+
+DUPLICATE, COHORT, WEIGHT, BAD_WEIGHT, PERIOD = range(5)
+
+
+@njit(parallel=True, cache=True)
+def leaf_pass(starts, ends, t, g, y, w, has_w, scale, code, t0, per_chunk, flags, sums, mass, counts, mean_w):  # fmt: skip
+    """Checks and sums of one sorted table; ``starts``/``ends`` delimit the units
+    (``code`` is each unit's cohort index), ``flags[chunk]`` marks the failed
+    checks and ``sums``, ``mass``, ``counts`` are per chunk by cohort and period."""
+    n_units = starts.shape[0]
+    T = sums.shape[2]
+    chunks = flags.shape[0]
+    for c in prange(chunks):
+        lo = c * per_chunk
+        hi = min(n_units, lo + per_chunk)
+        for u in range(lo, hi):
+            s = starts[u]
+            e = ends[u]
+            k = code[u]
+            wsum = 0.0
+            for r in range(s, e):
+                if r > s and t[r] == t[r - 1]:
+                    flags[c, DUPLICATE] = 1
+                if g[r] != g[s]:
+                    flags[c, COHORT] = 1
+                ti = t[r] - t0
+                if ti < 0 or ti >= T:
+                    flags[c, PERIOD] = 1
+                    continue
+                wr = 1.0
+                if has_w:
+                    if w[r] != w[s]:
+                        flags[c, WEIGHT] = 1
+                    wr = w[r] / scale
+                    if not (np.isfinite(wr) and wr > 0.0):
+                        flags[c, BAD_WEIGHT] = 1
+                wsum += wr
+                sums[c, k, ti] += y[r] * wr
+                mass[c, k, ti] += wr
+                counts[c, k, ti] += 1
+            mean_w[u] = wsum / (e - s)
+
+
+def summarize(starts, lengths, t, g, y, w, scale, code, n_cohorts, t0, periods):
+    """Flags and sums (``(cohorts, periods)`` arrays) of a sorted table; ``w`` is
+    None when unweighted."""
+    n_units = len(starts)
+    chunks = min(CHUNKS, max(1, n_units))
+    per_chunk = -(-n_units // chunks)
+    chunks = -(-n_units // per_chunk)
+    flags = np.zeros((chunks, 5), dtype=np.int8)
+    shape = (chunks, n_cohorts, periods)
+    sums, mass = np.zeros(shape), np.zeros(shape)
+    counts = np.zeros(shape, dtype=np.int64)
+    mean_w = np.zeros(n_units)
+    leaf_pass(
+        starts.astype(np.int64),
+        (starts + lengths).astype(np.int64),
+        np.ascontiguousarray(t, dtype=np.int64),
+        np.ascontiguousarray(g, dtype=np.int64),
+        np.ascontiguousarray(y, dtype=np.float64),
+        np.ascontiguousarray(w, dtype=np.float64) if w is not None else np.zeros(1),
+        w is not None,
+        float(scale),
+        np.ascontiguousarray(code, dtype=np.int64),
+        int(t0),
+        per_chunk,
+        flags,
+        sums,
+        mass,
+        counts,
+        mean_w,
+    )  # fmt: skip
+    return flags.max(axis=0), sums.sum(axis=0), mass.sum(axis=0), counts.sum(axis=0), mean_w