SniperGold_ML/engine/chunkmap.py

"""Line-safe chunk map (spec 9). Newline-scan only; pure function of bytes.

A chunk's nominal end is ``byte_start + chunk_bytes_nominal``; the actual
``byte_end`` is the first line terminator at or after the nominal end,
INCLUSIVE of terminator bytes. ``byte_end`` is therefore always a line
boundary and no chunk ever splits a CSV record. A file smaller than one
chunk yields exactly one chunk.
"""

import os

from .util import canonical_json, canonical_json_sha256, sha256_file

SCAN_BLOCK = 8 * 1024 * 1024


def iter_line_spans(fh, block=SCAN_BLOCK):
    """Yield (start, end_exclusive_of_terminator, terminator_len) for every
    raw line, streaming. End-inclusive byte position is start+end+term_len? No:
    span_end_exclusive = start + (line length) + terminator_len.
    terminator_len is 0 for a final unterminated line.
    """
    carry = b""
    base = 0

    def scan(data, dbase, at_eof):
        i = 0
        n = len(data)
        start = 0
        while i < n:
            b = data[i]
            if b == 0x0A:                       # \n
                yield (dbase + start, dbase + i, 1)
                i += 1
                start = i
            elif b == 0x0D:                     # \r
                if i + 1 < n:
                    if data[i + 1] == 0x0A:     # \r\n
                        yield (dbase + start, dbase + i, 2)
                        i += 2
                    else:
                        yield (dbase + start, dbase + i, 1)
                        i += 1
                    start = i
                elif at_eof:
                    yield (dbase + start, dbase + i, 1)
                    i += 1
                    start = i
                else:
                    break                      # hold trailing \r for next block
            else:
                i += 1
        if start < n:
            return data[start:], start
        return b"", n

    while True:
        chunk = fh.read(block)
        if not chunk:
            break
        data = carry + chunk
        dbase = base
        gen = scan(data, dbase, False)
        while True:
            try:
                item = next(gen)
            except StopIteration as done:
                carry, consumed = done.value
                break
            yield item
        base = dbase + consumed
    if carry:
        for item in scan(carry, base, True):
            yield item


def build_chunk_map(source_path, chunk_bytes_nominal, workload_bytes_nominal):
    """Build the deterministic chunk map for a source file (read-only scan)."""
    st = os.stat(source_path)
    total_bytes = st.st_size
    source_id = canonical_json_sha256({
        "path": os.path.abspath(source_path),
        "size": total_bytes,
        "mtime_ns": st.st_mtime_ns,
    })

    chunks = []
    with open(source_path, "rb") as fh:
        chunk_start = 0
        line_count = 0
        idx = 0
        for (lstart, lend, term) in iter_line_spans(fh):
            line_count += 1
            end_inclusive = lend + term
            if end_inclusive - chunk_start >= chunk_bytes_nominal:
                chunks.append({"index": idx, "byte_start": chunk_start,
                               "byte_end": end_inclusive, "line_count": line_count})
                idx += 1
                chunk_start = end_inclusive
                line_count = 0
        if chunk_start < total_bytes:
            # final chunk ends at EOF (inclusive of any final terminator)
            chunks.append({"index": idx, "byte_start": chunk_start,
                           "byte_end": total_bytes, "line_count": line_count})
        elif not chunks:
            # empty file -> single empty chunk (0 bytes)
            chunks.append({"index": 0, "byte_start": 0, "byte_end": 0,
                           "line_count": 0})

    # Workload grouping: accumulate whole chunks; final workload is partial.
    workloads = []
    wl_bytes = 0
    wl_chunks = []
    wl_total = 0
    for c in chunks:
        size = c["byte_end"] - c["byte_start"]
        wl_chunks.append(c["index"])
        wl_bytes += size
        wl_total += size
        if wl_bytes >= workload_bytes_nominal:
            workloads.append({"index": len(workloads), "chunks": list(wl_chunks),
                              "total_bytes": wl_total, "partial": False})
            wl_chunks = []
            wl_bytes = 0
            wl_total = 0
    if wl_chunks:
        workloads.append({"index": len(workloads), "chunks": list(wl_chunks),
                          "total_bytes": wl_total, "partial": True})
    if not workloads:
        workloads.append({"index": 0, "chunks": [], "total_bytes": 0, "partial": True})

    wl_of_chunk = {}
    for w in workloads:
        for ci in w["chunks"]:
            wl_of_chunk[ci] = w["index"]
    for c in chunks:
        c["workload_index"] = wl_of_chunk.get(c["index"], len(workloads) - 1)

    payload = {
        "source_id": source_id,
        "chunk_bytes_nominal": chunk_bytes_nominal,
        "workload_bytes_nominal": workload_bytes_nominal,
        "chunks": chunks,
        "workloads": workloads,
        "total_bytes": total_bytes,
        "total_chunks": len(chunks),
        "total_workloads": len(workloads),
    }
    payload["chunkmap_sha256"] = canonical_json_sha256(
        {k: v for k, v in payload.items() if k != "chunkmap_sha256"})
    return payload


def chunkmap_digest(chunkmap):
    return chunkmap["chunkmap_sha256"]