graflo.hq.bulk.run

Run a bulk load: open the staging directory, fold, write, verify.

Casting into the staging directory is the caster's (it hands each batch to a :class:~graflo.hq.bulk.stage.Stager); this module does the rest. Every phase is recorded in the run ledger, so a run stopped at any point resumes from the last thing that completed: staging is redone, a fold is kept, written chunks are skipped and the one in doubt is replayed.

Attributes

logger = logging.getLogger(__name__) module-attribute

Classes

BulkVerifyError

Bases: RuntimeError

The target does not hold what the fold produced.

Source code in graflo/hq/bulk/run.py

class BulkVerifyError(RuntimeError):
    """The target does not hold what the fold produced."""

Functions:

dense_ids_for(conn_conf)

Whether the target's bulk writer takes dense integer ids.

Raises:

Type	Description
ValueError	The target has no bulk writer.

Source code in graflo/hq/bulk/run.py

def dense_ids_for(conn_conf: DBConfig) -> bool:
    """Whether the target's bulk writer takes dense integer ids.

    Raises:
        ValueError: The target has no bulk writer.
    """
    conn_cls = ConnectionManager.require(conn_conf, ConnectionCapability.BULK_LOAD)
    return bool(conn_cls.bulk_dense_ids)

edge_buckets_for(conn_conf, cfg)

Buckets per endpoint to split edge runs into, when the writer asks for it.

Raises:

Type	Description
ValueError	The target has no bulk writer.

Source code in graflo/hq/bulk/run.py

def edge_buckets_for(conn_conf: DBConfig, cfg: BulkConfig) -> int | None:
    """Buckets per endpoint to split edge runs into, when the writer asks for it.

    Raises:
        ValueError: The target has no bulk writer.
    """
    conn_cls = ConnectionManager.require(conn_conf, ConnectionCapability.BULK_LOAD)
    return cfg.parallelism if conn_cls.bulk_partition_edges else None

fold_staged(staging_dir, ledger, cfg, *, dense_ids, edge_buckets=None)

Fold the staged runs once; a resumed run reuses the fold.

Source code in graflo/hq/bulk/run.py

def fold_staged(
    staging_dir: Path,
    ledger: RunLedger,
    cfg: BulkConfig,
    *,
    dense_ids: bool,
    edge_buckets: int | None = None,
) -> FoldCounts:
    """Fold the staged runs once; a resumed run reuses the fold."""
    if ledger.has(Phase.FOLDED):
        return FoldCounts.load(staging_dir)
    counts = fold(
        staging_dir,
        dense_ids=dense_ids,
        edge_buckets=edge_buckets,
        memory_limit=cfg.memory_limit,
        threads=cfg.fold_threads,
    )
    ledger.mark(Phase.FOLDED)
    for run in counts.runs:
        logger.info(
            "Folded %s: %s staged, %s folded, %s rejected",
            run.type.label,
            run.staged,
            run.folded,
            run.rejected,
        )
    return counts

open_staging(staging_dir, mode='bulk')

The ledger of a new or resumable mode run (bulk or sync) in staging_dir.

A run stopped while staging is staged again from the start; one stopped later resumes.

Raises:

Type	Description
RuntimeError	The bulk extra is not installed.
ValueError	staging_dir holds files but no run, or a run of the other mode.

Source code in graflo/hq/bulk/run.py

def open_staging(staging_dir: Path, mode: str = "bulk") -> RunLedger:
    """The ledger of a new or resumable *mode* run (``bulk`` or ``sync``) in *staging_dir*.

    A run stopped while staging is staged again from the start; one stopped
    later resumes.

    Raises:
        RuntimeError: The ``bulk`` extra is not installed.
        ValueError: *staging_dir* holds files but no run, or a run of the other mode.
    """
    _deps.require()
    root = Path(staging_dir)
    root.mkdir(parents=True, exist_ok=True)
    ledger = RunLedger(root)
    if not ledger.exists:
        if any(root.iterdir()):
            raise ValueError(
                f"staging directory {root} is not empty and holds no bulk or sync "
                "run; give an empty or new directory"
            )
        ledger.set_mode(mode)
        ledger.mark(Phase.OPENED)
    elif (ledger.mode or "bulk") != mode:
        raise ValueError(
            f"staging directory {root} holds a {ledger.mode or 'bulk'} run; a "
            f"{mode} run needs its own directory"
        )
    elif not ledger.has(Phase.STAGED):
        for name in _OWN:
            path = root / name
            if path.is_dir():
                shutil.rmtree(path)
            elif path.exists():
                path.unlink()
        logger.info("Bulk run in %s was stopped while staging; staging again", root)
    return ledger

verify(expected, actual)

Refuse a load whose per-type counts differ from the fold's.

Raises:

Type	Description
BulkVerifyError	Listing every type that differs.

Source code in graflo/hq/bulk/run.py

def verify(expected: BulkCounts, actual: BulkCounts) -> None:
    """Refuse a load whose per-type counts differ from the fold's.

    Raises:
        BulkVerifyError: Listing every type that differs.
    """
    diffs = [
        f"vertex {name}: expected {expected.vertices.get(name, 0)}, "
        f"target holds {actual.vertices.get(name, 0)}"
        for name in sorted(set(expected.vertices) | set(actual.vertices))
        if expected.vertices.get(name, 0) != actual.vertices.get(name, 0)
    ] + [
        f"edge {edge}: expected {expected.edges.get(edge, 0)}, "
        f"target holds {actual.edges.get(edge, 0)}"
        for edge in sorted(set(expected.edges) | set(actual.edges), key=str)
        if expected.edges.get(edge, 0) != actual.edges.get(edge, 0)
    ]
    if diffs:
        raise BulkVerifyError("bulk load verification failed: " + "; ".join(diffs))

write_folded(schema, conn_conf, staging_dir, ledger, cfg, *, recreate=False)

Write the folded runs to the target and check its counts against the fold.

Parameters:

Name	Type	Description	Default
schema	Schema	The schema, logical names, initialized for the target flavor.	required
conn_conf	DBConfig	The target; it must declare supports_bulk_load.	required
staging_dir	Path	A folded staging directory.	required
ledger	RunLedger	Its ledger.	required
cfg	BulkConfig	Chunk size and parallelism.	required
recreate	bool	Drop the target's schema first instead of refusing an existing one.	False

Raises:

Type	Description
ValueError	The target has no bulk writer, or (from the writer's first-write guard) already holds a schema.
BulkVerifyError	The target's counts differ from the fold's.

Source code in graflo/hq/bulk/run.py

def write_folded(
    schema: Schema,
    conn_conf: DBConfig,
    staging_dir: Path,
    ledger: RunLedger,
    cfg: BulkConfig,
    *,
    recreate: bool = False,
) -> BulkCounts:
    """Write the folded runs to the target and check its counts against the fold.

    Args:
        schema: The schema, logical names, initialized for the target flavor.
        conn_conf: The target; it must declare ``supports_bulk_load``.
        staging_dir: A folded staging directory.
        ledger: Its ledger.
        cfg: Chunk size and parallelism.
        recreate: Drop the target's schema first instead of refusing an
            existing one.

    Raises:
        ValueError: The target has no bulk writer, or (from the writer's
            first-write guard) already holds a schema.
        BulkVerifyError: The target's counts differ from the fold's.
    """
    ConnectionManager.require(conn_conf, ConnectionCapability.BULK_LOAD)
    counts = FoldCounts.load(staging_dir)
    ctx = WriteContext.for_flavor(schema, conn_conf.connection_type)
    target_schema = physical_schema(schema, conn_conf.connection_type)
    if ledger.has(Phase.FINISHED):
        logger.info("Bulk run in %s is already finished", staging_dir)
    resume = ledger.has(Phase.BEGINNING)
    with ConnectionManager(connection_config=conn_conf) as db:
        ledger.mark(Phase.BEGINNING)
        plan = db.bulk_begin(target_schema, cfg, recreate=recreate, resume=resume)
        if resume and plan.restart:
            logger.warning(
                "The target lost the chunks written before; writing all again"
            )
            ledger.restart()
        ledger.mark(Phase.BEGUN)
    parallel = min(cfg.parallelism, plan.max_parallel or cfg.parallelism)

    vertex_runs = [r for r in counts.runs if r.type.kind == "vertex"]
    edge_runs = [r for r in counts.runs if r.type.kind == "edge" and not r.buckets]
    celled_runs = [r for r in counts.runs if r.type.kind == "edge" and r.buckets]
    with _Connections(conn_conf) as conns:
        _write_chunks(
            (
                c
                for run in vertex_runs
                for c in _vertex_chunks(ctx, staging_dir, run, cfg)
            ),
            lambda chunk: conns.get().bulk_write_vertices(target_schema, chunk),
            ledger,
            parallel,
        )
        if not ledger.has(Phase.EDGES_PREPARED):
            conns.get().bulk_prepare_edges(target_schema)
            ledger.mark(Phase.EDGES_PREPARED)
        _write_chunks(
            (
                c
                for run in edge_runs
                for c in _edge_chunks(
                    schema,
                    ctx,
                    run,
                    [folded_path(staging_dir, run.slot)],
                    run.slot,
                    cfg,
                )
            ),
            lambda chunk: conns.get().bulk_write_edges(target_schema, chunk),
            ledger,
            parallel,
        )
        # One run at a time: two runs may share an endpoint type.
        for run in celled_runs:
            _write_rounds(
                _edge_rounds(schema, ctx, staging_dir, run, cfg),
                lambda chunk: conns.get().bulk_write_edges(target_schema, chunk),
                ledger,
                parallel,
            )
        actual = conns.get().bulk_finish(target_schema)
    verify(_expected(schema, ctx, counts), actual)
    ledger.mark(Phase.FINISHED)
    return actual