#!/usr/bin/env python3
"""MapReduce materialization vs. a dataflow pipeline, DDIA 2e Ch. 11.

The SAME three-stage batch workflow -- map -> group-by -> aggregate -- run over
ONE dataset TWO ways, so you can measure the cost of materializing intermediate
state between stages.

  MapReduce-style: every stage writes its FULL output to a "distributed
  filesystem" before the next stage may read it. We model the DFS honestly:
  the bytes are written to a real temp file, fsync'd, and counted REPLICATION=3
  times (one durable copy per replica). Because each stage's file must be
  complete before the next stage opens it, the stages cannot overlap and the
  first FINAL record cannot appear until the entire input has flowed through
  all three stages -- "its file-based I/O prevents job pipelining."

  Dataflow / pipelined-style: the three stages are chained generators. A record
  is pulled through map -> group -> aggregate on demand, nothing is written
  between stages, and the first FINAL record is emitted as soon as the first
  group is complete -- after consuming only a handful of input records.

Two deterministic measurements:
  (1) INTERMEDIATE BYTES written to the DFS between stages. MapReduce writes
      ~3x the intermediate volume at every stage boundary; the pipeline writes
      zero (intermediate state lives in memory).
  (2) TIME-TO-FIRST-OUTPUT, as a records-consumed proxy (not wall-clock): how
      many source records are pulled before the first FINAL record is emitted.
      MapReduce must consume the ENTIRE input first; the pipeline needs only
      the first group.

Pure standard library, deterministic -> reproducible.
"""
import os
import tempfile

# --- The dataset -------------------------------------------------------------
# 100,000 records, laid out already grouped by key (1,000 keys x 100 records).
# Grouped/sorted input is what lets a dataflow engine stream the group-by and
# emit a group's aggregate the moment the key changes -- see the Go-deeper note
# on why an UNsorted group-by would re-introduce a barrier even in a pipeline.
NUM_KEYS = 1000
GROUP_SIZE = 100
TOTAL = NUM_KEYS * GROUP_SIZE        # 100,000
REPLICATION = 3                       # DFS replication factor: 3 durable copies


class Source:
    """The input. Yields raw record strings and counts every record pulled,
    so we can measure how many the consumer touched before its first output."""

    def __init__(self):
        self.pulled = 0

    def __iter__(self):
        for k in range(NUM_KEYS):
            for i in range(GROUP_SIZE):
                self.pulled += 1
                # A raw log line the map stage will parse.
                yield f"evt key={k} val={(k * 13 + i * i) % 100}"


def parse(raw):
    """Stage 1 (map): turn a raw log line into a (key, value) pair."""
    _, key_part, val_part = raw.split(" ")
    return int(key_part[4:]), int(val_part[4:])


# --- The DFS model -----------------------------------------------------------

def materialize(lines):
    """Write an intermediate stage's full output to the 'distributed filesystem'.

    We do it for real: serialize to bytes, write to a temp file, fsync so the
    bytes are durable, then account for REPLICATION copies. Returns the number
    of bytes charged to the DFS (data size x replication factor)."""
    data = "".join(lines).encode()
    fd, path = tempfile.mkstemp(prefix="dfs_stage_")
    try:
        os.write(fd, data)
        os.fsync(fd)                  # durability: the stage's output survives a crash
    finally:
        os.close(fd)
        os.unlink(path)
    return len(data) * REPLICATION    # 3 replicas each hold the full copy


# --- MapReduce-style: materialize every stage to the DFS ---------------------

def run_mapreduce(source):
    """map -> group -> aggregate, materializing each stage to the DFS first.

    Returns (final_results, intermediate_bytes, source_pulls_at_first_output)."""
    intermediate_bytes = 0

    # Stage 1 (map): consume the ENTIRE input, then write the whole map output.
    mapped = []
    for raw in source:                        # pulls all TOTAL records
        mapped.append(parse(raw))
    intermediate_bytes += materialize(f"{k}\t{v}\n" for k, v in mapped)

    # Stage 2 (group-by): read the full map output, then write the whole grouping.
    groups = {}
    for k, v in mapped:
        groups.setdefault(k, []).append(v)
    grouped = sorted(groups.items())
    intermediate_bytes += materialize(
        f"{k}\t{','.join(map(str, vs))}\n" for k, vs in grouped
    )

    # Stage 3 (aggregate): read the full grouping, produce the final result.
    # The first final record only exists NOW -- after all TOTAL records have
    # already flowed through stages 1 and 2. Capture the source counter here.
    pulls_at_first_output = source.pulled
    final = [(k, sum(vs)) for k, vs in grouped]

    return final, intermediate_bytes, pulls_at_first_output


# --- Dataflow-style: chain the stages as generators, materialize nothing ----

def streaming_group(pairs):
    """Stage 2 as a generator: emit (key, values) the instant the key changes.
    Works because the input is grouped by key -- a group is complete as soon as
    a different key arrives, so we never have to buffer the whole dataset."""
    cur_key = None
    acc = []
    for k, v in pairs:
        if cur_key is None:
            cur_key = k
        if k != cur_key:
            yield cur_key, acc
            cur_key, acc = k, []
        acc.append(v)
    if cur_key is not None:
        yield cur_key, acc


def run_pipeline(source):
    """map -> group -> aggregate as fused generators. Nothing is written between
    stages. Returns (final_results, intermediate_bytes, source_pulls_at_first)."""
    mapped = (parse(raw) for raw in source)                 # stage 1
    grouped = streaming_group(mapped)                        # stage 2
    aggregated = ((k, sum(vs)) for k, vs in grouped)         # stage 3

    final = []
    pulls_at_first_output = None
    for record in aggregated:
        if pulls_at_first_output is None:
            # First FINAL record just came out -- how many source records did
            # it take? Only the first group plus the one record that closed it.
            pulls_at_first_output = source.pulled
        final.append(record)

    intermediate_bytes = 0            # generators keep intermediate state in memory
    return final, intermediate_bytes, pulls_at_first_output


def human(n):
    return f"{n:,} bytes ({n / 1024:.1f} KiB)"


def main():
    print(f"one workflow (map -> group-by -> aggregate) over {TOTAL:,} records, two engines")
    print(f"    {NUM_KEYS:,} keys x {GROUP_SIZE} records each, DFS replication factor = {REPLICATION}\n")

    mr_final, mr_bytes, mr_first = run_mapreduce(Source())
    pl_final, pl_bytes, pl_first = run_pipeline(Source())

    # Same workflow, same data -> the two engines must agree on the answer.
    assert mr_final == pl_final, "engines disagree!"
    print(f"both engines produce the identical {len(mr_final):,} aggregates "
          f"(e.g. key 0 -> {mr_final[0][1]}, key 999 -> {mr_final[-1][1]}).\n")

    print("(1) intermediate bytes written to the DFS between stages")
    print(f"    MapReduce (materialize map + group, each x{REPLICATION}): {human(mr_bytes)}")
    print(f"    dataflow pipeline (generators, nothing written):          {human(pl_bytes)}")
    print(f"    ratio: MapReduce writes {mr_bytes:,} bytes of intermediate; the pipeline writes {pl_bytes}.\n")

    print("(2) time-to-first-output  (source records pulled before the FIRST final record)")
    print(f"    MapReduce:         {mr_first:,} of {TOTAL:,} records  (the ENTIRE input, at every stage barrier)")
    print(f"    dataflow pipeline: {pl_first:,} of {TOTAL:,} records  (just the first group + the record that closed it)")
    print(f"    ratio: MapReduce consumes {mr_first // pl_first}x more input before it can emit anything.")


if __name__ == "__main__":
    main()