Skip to content

graflo.hq.bulk.fold

Fold staged runs by identity in DuckDB: one row per vertex and per edge.

The rules are the upsert's, applied once to the whole load (the file backend's read-side fold applies the same ones):

  • per field, the value of the latest row (largest seq) that carries it; a field a later row lacks keeps its earlier value;
  • merge: union fields take the members of every row in seq order, each once, nulls dropped;
  • an edge's endpoint documents fold the same way as its properties.

Rows whose key occurs once pass through; only repeated keys are exploded into one row per field and reduced. Edges whose endpoint is not a folded vertex are rejected with their reason, as are edges owned through an owner_field whose owner is not a folded vertex. DuckDB spills to the staging directory under the memory limit it is given.

Output, in the staging directory: folded/<slot>.jsonl.zst sorted by key (an edge run split into cells is a directory instead, see :func:folded_cell), rejects/, and counts.json.

Attributes

COUNTS = 'counts.json' module-attribute

FOLDED = 'folded' module-attribute

Classes

FoldCounts

Bases: BaseModel

Every run of a fold, vertices first.

Source code in graflo/hq/bulk/fold.py
class FoldCounts(BaseModel):
    """Every run of a fold, vertices first."""

    runs: list[RunCounts]

    def vertex(self, name: str) -> RunCounts | None:
        return next(
            (r for r in self.runs if r.type.kind == "vertex" and r.type.name == name),
            None,
        )

    @classmethod
    def load(cls, staging_dir: Path) -> FoldCounts:
        return cls.model_validate_json((Path(staging_dir) / COUNTS).read_text())

Attributes

runs instance-attribute

Methods:

load(staging_dir) classmethod
Source code in graflo/hq/bulk/fold.py
@classmethod
def load(cls, staging_dir: Path) -> FoldCounts:
    return cls.model_validate_json((Path(staging_dir) / COUNTS).read_text())
vertex(name)
Source code in graflo/hq/bulk/fold.py
def vertex(self, name: str) -> RunCounts | None:
    return next(
        (r for r in self.runs if r.type.kind == "vertex" and r.type.name == name),
        None,
    )

RunCounts

Bases: BaseModel

One run's rows: staged, after the fold, rejected, and its dense-id range.

Source code in graflo/hq/bulk/fold.py
class RunCounts(BaseModel):
    """One run's rows: staged, after the fold, rejected, and its dense-id range."""

    slot: str
    type: RunType
    staged: int
    folded: int
    rejected: int
    dense: tuple[int, int] | None = None
    buckets: int | None = None
    """Buckets per endpoint of an edge run split into cells."""

Attributes

buckets = None class-attribute instance-attribute

Buckets per endpoint of an edge run split into cells.

dense = None class-attribute instance-attribute
folded instance-attribute
rejected instance-attribute
slot instance-attribute
staged instance-attribute
type instance-attribute

Functions:

counts_json(staging_dir)

counts.json as data.

Source code in graflo/hq/bulk/fold.py
def counts_json(staging_dir: Path) -> dict[str, Any]:
    """``counts.json`` as data."""
    return json.loads((Path(staging_dir) / COUNTS).read_text())

fold(staging_dir, *, dense_ids=False, edge_buckets=None, memory_limit=None, threads=None)

Fold every staged run of staging_dir; write the folded runs and counts.

Parameters:

Name Type Description Default
staging_dir Path

A directory a :class:~graflo.hq.bulk.stage.Stager wrote.

required
dense_ids bool

Number every folded vertex 0..n-1 across types, in type then key order, and carry the numbers onto edge endpoints.

False
edge_buckets int | None

Split each edge run into cells by a hash of its source and target keys into this many buckets each.

None
memory_limit str | None

DuckDB's memory limit ("4GB"); spills beyond it.

None
threads int | None

DuckDB worker threads; its default when omitted.

None

Raises:

Type Description
RuntimeError

The bulk extra is not installed.

Source code in graflo/hq/bulk/fold.py
def fold(
    staging_dir: Path,
    *,
    dense_ids: bool = False,
    edge_buckets: int | None = None,
    memory_limit: str | None = None,
    threads: int | None = None,
) -> FoldCounts:
    """Fold every staged run of *staging_dir*; write the folded runs and counts.

    Args:
        staging_dir: A directory a :class:`~graflo.hq.bulk.stage.Stager` wrote.
        dense_ids: Number every folded vertex ``0..n-1`` across types, in type
            then key order, and carry the numbers onto edge endpoints.
        edge_buckets: Split each edge run into cells by a hash of its source
            and target keys into this many buckets each.
        memory_limit: DuckDB's memory limit (``"4GB"``); spills beyond it.
        threads: DuckDB worker threads; its default when omitted.

    Raises:
        RuntimeError: The ``bulk`` extra is not installed.
    """
    duckdb = _deps.duckdb()
    root = Path(staging_dir)
    (root / FOLDED).mkdir(parents=True, exist_ok=True)
    (root / REJECTS).mkdir(parents=True, exist_ok=True)
    con = duckdb.connect()
    try:
        con.execute(f"SET temp_directory = {_lit(str(root / 'duckdb.tmp'))}")
        con.execute("SET preserve_insertion_order = false")
        if memory_limit:
            con.execute(f"SET memory_limit = {_lit(memory_limit)}")
        if threads:
            con.execute(f"SET threads = {int(threads)}")
        runs = list(staged_runs(root))
        vertex_runs = [(s, t) for s, t in runs if t.kind == "vertex"]
        edge_runs = [(s, t) for s, t in runs if t.kind == "edge"]
        counts: list[RunCounts] = []
        key_tables: dict[str, str] = {}
        offset = 0
        for index, (slot, run) in enumerate(vertex_runs):
            table = f"vk{index}"
            result = _fold_vertices(
                con, root, slot, run, table, offset if dense_ids else None
            )
            if dense_ids:
                offset += result.folded
            key_tables[str(run.name)] = table
            counts.append(result)
        for slot, run in edge_runs:
            counts.append(
                _fold_edges(con, root, slot, run, key_tables, dense_ids, edge_buckets)
            )
    finally:
        con.close()
    out = FoldCounts(runs=counts)
    (root / COUNTS).write_text(out.model_dump_json(indent=2))
    return out

folded_cell(staging_dir, slot, cell)

The files of one cell (source bucket, target bucket) of an edge run.

Source code in graflo/hq/bulk/fold.py
def folded_cell(staging_dir: Path, slot: str, cell: tuple[int, int]) -> list[Path]:
    """The files of one cell ``(source bucket, target bucket)`` of an edge run."""
    a, b = cell
    directory = Path(staging_dir) / FOLDED / slot / f"sb={a}" / f"tb={b}"
    return sorted(directory.glob("*.json*"))

folded_path(staging_dir, slot)

Source code in graflo/hq/bulk/fold.py
def folded_path(staging_dir: Path, slot: str) -> Path:
    return Path(staging_dir) / FOLDED / f"{slot}.jsonl.zst"