diff --git a/experimentation/iceberg/.gitignore b/experimentation/iceberg/.gitignore new file mode 100644 index 00000000..8884117b --- /dev/null +++ b/experimentation/iceberg/.gitignore @@ -0,0 +1,9 @@ +# Isolated interpreter for this experiment (see requirements.txt). +.venv/ +__pycache__/ + +# Jetstream captures are ~21 MiB per 10 minutes and are reproducible by +# re-running capture.py. The report and results JSON are committed; the +# per-call CSV is not, since it holds one row per AWS request. +data/captures/ +data/results/*-calls.csv diff --git a/experimentation/iceberg/README.md b/experimentation/iceberg/README.md new file mode 100644 index 00000000..bbaa9602 --- /dev/null +++ b/experimentation/iceberg/README.md @@ -0,0 +1,130 @@ +# Iceberg write-amplification experiment + +Prices the four S3-facing operations in an Iceberg ingest pipeline fed by the +Bluesky Jetstream firehose: + +1. **Write records to S3** — the raw-Parquet control, no table format. +2. **Update Iceberg metadata** — the same rows through `table.append()`. +3. **Compaction** — scan, collapse each AT-URI to its latest state, drop delete + tombstones, rewrite one file per partition. +4. **Metadata cleaning** — expire snapshots, then sweep orphaned objects. + +Each runs inside a metered phase, so the output is an exact per-operation ledger +of S3 and Glue calls, latency, and cost — not an estimate. + +## How it measures + +PyIceberg's default `PyArrowFileIO` drives a C++ S3 client that Python cannot +intercept. The catalog therefore pins `py-io-impl=pyiceberg.io.fsspec.FsspecFileIO`, +routing every request through s3fs → aiobotocore → botocore, where +`s3_meter.Meter` counts it. + +`Meter.install()` wraps `botocore.session.Session.__init__`, so *every* session — +boto3's for raw writes, aiobotocore's inside s3fs, and the Glue client PyIceberg +builds for catalog commits — carries the handlers. It must run before any client +is constructed; `run_experiment.py` does this at import time. + +Two things worth knowing about the numbers: + +- **Glue commits are not S3 calls.** Swapping `metadata_location` is a Glue + `UpdateTable`, billed at $1/100k rather than $5/100k. Metering only S3 would + miss the commit path entirely, so Glue is counted as its own tier. +- **Request bytes come from the body stream, not a header.** botocore has no + `Content-Length` on the request dict at `before-call`, and large uploads switch + to `aws-chunked` encoding which carries none at all. `_body_size` measures the + `BytesIO` directly and restores its position. + +## Setup + +pyiceberg pins `rich<15` and the root project requires `rich>=15`, so this +experiment cannot share the root venv: + +```bash +uv venv --python 3.11 experimentation/iceberg/.venv +uv pip install --python experimentation/iceberg/.venv -r experimentation/iceberg/requirements.txt +``` + +## Running + +Capture and replay are separate on purpose. A 10-minute firehose capture is never +reproducible, so it is recorded once and replayed as many times as needed — that +way every write-path variant is measured against byte-identical input. + +```bash +# 1. Capture (writes data/captures/jetstream-.jsonl.gz) +./experimentation/iceberg/.venv/bin/python -m experimentation.iceberg.capture --seconds 600 + +# 2. Replay through both write paths +./experimentation/iceberg/.venv/bin/python -m experimentation.iceberg.run_experiment \ + --capture experimentation/iceberg/data/captures/jetstream-.jsonl.gz +``` + +Useful flags: `--flush-seconds` (default 60), `--max-batches`, `--skip-raw`, +`--run-id`. + +Outputs land in `data/results/-{report.md,results.json,calls.csv}`. +`calls.csv` is every individual API call, for slicing outside the report. + +## Layout + +``` +s3://lab-data-integrations-interface/experiments/iceberg// + raw/ # baseline Parquet, no table format + warehouse/ # Iceberg tables +``` + +Glue tables are `_` in the `iceberg_experiments` database, +so repeat runs never collide. + +## Data model + +Four tables — `posts`, `likes`, `reposts`, `follows` — each partitioned by +`days(created_at)`. Separate tables mean 4x the metadata commits per flush, which +is itself one of the findings. + +Bluesky `createdAt` is client-supplied. Anything more than 24h from the broker's +`time_us` falls back to ingest time, and the report splits the fallbacks into +their two very different causes: + +- **`delete` events** carry no record body, so they have no `createdAt` at all. + Structural, not a data problem — and the large majority of fallbacks. +- **Skewed timestamps** parse cleanly but sit far from the broker clock. In the + measured capture these were overwhelmingly *one* archive-import bot stamping + genuine historical dates (2011, 2013, 2016…). Those dates are arguably + correct; the rule rewrites them so a single bot cannot open a daily partition + per historical date it touches. That is a deliberate correctness-for-file-count + trade, not a data-cleaning step. + +## Duplicates vs. lifecycle collapses + +These are counted separately because they are constantly conflated: + +- A **redelivered duplicate** is the identical event twice — same `uri` *and* + same `cid`. A stable 10-minute capture contained **zero** of these. +- A **lifecycle collapse** is several distinct events about one record (create + then delete, create then update), each with its own `cid`. Collapsing these + materialises current state; it is not deduplication. + +Compaction keeps the latest row per URI and then **drops `delete` tombstones**. +Note this only cancels a create that is in the same table — a delete of a record +written before the table existed has nothing to reconcile against and is simply +discarded. A tombstone also lands in the partition of its *ingest* day, not the +partition of the record it deletes, so cross-partition deletes need equality +deletes or merge-on-read to work properly. + +## Cleanup + +```bash +aws s3 rm s3://lab-data-integrations-interface/experiments/iceberg// --recursive +python -c "from experimentation.iceberg import catalog; \ + catalog.drop_tables(catalog.build_catalog(''), '')" +``` + +## Tests + +```bash +./experimentation/iceberg/.venv/bin/python -m pytest experimentation/iceberg/tests +``` + +The suite is offline — no AWS, no network. A `conftest.py` skips collection when +the root interpreter picks it up, since that venv has no pyiceberg. diff --git a/experimentation/iceberg/capture.py b/experimentation/iceberg/capture.py new file mode 100644 index 00000000..4190fbf1 --- /dev/null +++ b/experimentation/iceberg/capture.py @@ -0,0 +1,122 @@ +"""Phase 1 -- capture the Bluesky Jetstream firehose to a local file. + +Deliberately decoupled from the S3/Iceberg write path. A 10-minute capture is +expensive to re-collect and never reproducible, so it is recorded once and then +replayed as many times as needed. That keeps every write-path variant (flush +interval, partition spec, compaction strategy) measured against byte-identical +input. + +Usage: + python -m experimentation.iceberg.capture --seconds 600 +""" + +from __future__ import annotations + +import argparse +import asyncio +import gzip +import json +import time +from collections import Counter +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from urllib.parse import urlencode + +import websockets + +from experimentation.iceberg import constants + +CAPTURE_DIR = Path(__file__).parent / "data" / "captures" + + +def build_endpoint() -> str: + """Jetstream subscribe URL filtered to the four collections we care about.""" + query = urlencode([("wantedCollections", nsid) for nsid in constants.COLLECTIONS]) + return f"{constants.JETSTREAM_ENDPOINT}?{query}" + + +def _tally(message: str, counts: Counter[str]) -> None: + """Update per-record-type counters from one raw frame.""" + counts["total"] += 1 + try: + event = json.loads(message) + except json.JSONDecodeError: + counts["unparseable"] += 1 + return + commit = event.get("commit") + if not isinstance(commit, dict): + return + record_type = constants.COLLECTIONS.get(commit.get("collection", "")) + if record_type: + counts[record_type] += 1 + + +async def _drain(socket: Any, handle: Any, counts: Counter[str], deadline: float) -> None: + """Write frames verbatim to ``handle`` until ``deadline``, tallying as we go.""" + started = time.monotonic() + while time.monotonic() < deadline: + remaining = deadline - time.monotonic() + try: + message = await asyncio.wait_for(socket.recv(), timeout=remaining) + except TimeoutError: + return + + text = message if isinstance(message, str) else message.decode("utf-8") + handle.write(text) + handle.write("\n") + _tally(text, counts) + + if counts["total"] % 20_000 == 0: + rate = counts["total"] / max(time.monotonic() - started, 1e-9) + print(f" {counts['total']:,} events {rate:,.0f}/s") + + +async def capture(seconds: int, output_path: Path) -> dict[str, int]: + """Stream Jetstream for ``seconds`` and write raw events as gzipped JSONL. + + Returns a per-record-type count. Events are written exactly as received so + the replay stage owns all parsing -- a schema change should never require + re-capturing. + """ + output_path.parent.mkdir(parents=True, exist_ok=True) + endpoint = build_endpoint() + counts: Counter[str] = Counter() + started = time.monotonic() + + print(f"connecting to {endpoint}") + print(f"capturing for {seconds}s -> {output_path}") + + # max_size=None: some posts with large embeds exceed the 1MiB default frame cap. + async with websockets.connect(endpoint, max_size=None) as socket: + with gzip.open(output_path, "wt", encoding="utf-8") as handle: + await _drain(socket, handle, counts, started + seconds) + + counts["elapsed_seconds"] = int(time.monotonic() - started) + return dict(counts) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Capture Bluesky Jetstream to a local file.") + parser.add_argument("--seconds", type=int, default=constants.DEFAULT_CAPTURE_SECONDS) + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args() + + stamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S") + output_path = args.output or CAPTURE_DIR / f"jetstream-{stamp}.jsonl.gz" + + counts = asyncio.run(capture(args.seconds, output_path)) + + size_mb = output_path.stat().st_size / 1024 / 1024 + print(f"\ncaptured {counts.get('total', 0):,} events in {counts.get('elapsed_seconds', 0)}s") + for record_type in constants.RECORD_TYPES: + print(f" {record_type:<10} {counts.get(record_type, 0):>10,}") + print(f"compressed size: {size_mb:.1f} MiB") + print(f"wrote {output_path}") + + metadata_path = output_path.with_suffix(".meta.json") + metadata_path.write_text(json.dumps(counts, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/experimentation/iceberg/catalog.py b/experimentation/iceberg/catalog.py new file mode 100644 index 00000000..0ee6acbc --- /dev/null +++ b/experimentation/iceberg/catalog.py @@ -0,0 +1,102 @@ +"""Glue catalog wiring for the experiment. + +Two things here are load-bearing: + +1. ``py-io-impl`` is pinned to ``FsspecFileIO``. PyIceberg defaults to + ``PyArrowFileIO``, whose S3 client lives in C++ and is invisible to the + meter. Fsspec routes through aiobotocore, so every request is countable. +2. Tables are named ``_`` inside a single dedicated Glue + database, so repeat runs never collide and cleanup is one prefix delete. +""" + +from __future__ import annotations + +from typing import Any + +from pyiceberg.catalog import Catalog +from pyiceberg.catalog.glue import GlueCatalog +from pyiceberg.exceptions import NamespaceAlreadyExistsError, NoSuchTableError + +from experimentation.iceberg import constants, schemas + + +def warehouse_uri(run_id: str) -> str: + return f"s3://{constants.S3_BUCKET}/{constants.S3_EXPERIMENT_PREFIX}/{run_id}/warehouse" + + +def raw_uri(run_id: str) -> str: + return f"s3://{constants.S3_BUCKET}/{constants.S3_EXPERIMENT_PREFIX}/{run_id}/raw" + + +def table_name(record_type: str, run_id: str) -> str: + return f"{record_type}_{run_id}" + + +def build_catalog(run_id: str) -> Catalog: + """Construct the Glue-backed catalog for this run.""" + return GlueCatalog( + name="iceberg_experiment", + **{ + "warehouse": warehouse_uri(run_id), + "glue.region": constants.AWS_REGION, + "s3.region": constants.AWS_REGION, + # Required for the meter to see anything -- see module docstring. + "py-io-impl": "pyiceberg.io.fsspec.FsspecFileIO", + }, + ) + + +def ensure_namespace(catalog: Catalog) -> None: + try: + catalog.create_namespace(constants.GLUE_DATABASE) + except NamespaceAlreadyExistsError: + pass + + +def create_tables(catalog: Catalog, run_id: str) -> dict[str, Any]: + """Create one partitioned table per record type. Returns record_type -> Table.""" + ensure_namespace(catalog) + tables: dict[str, Any] = {} + + for record_type in constants.RECORD_TYPES: + identifier = (constants.GLUE_DATABASE, table_name(record_type, run_id)) + tables[record_type] = catalog.create_table( + identifier=identifier, + schema=schemas.SCHEMAS[record_type], + partition_spec=schemas.PARTITION_SPEC, + location=f"{warehouse_uri(run_id)}/{record_type}", + properties={ + "format-version": "2", + "write.parquet.compression-codec": "zstd", + # Leave stale metadata.json files in place so the expiry phase + # has real work to measure. + "write.metadata.delete-after-commit.enabled": "false", + }, + ) + return tables + + +def load_tables(catalog: Catalog, run_id: str) -> dict[str, Any]: + """Load existing tables for a run, skipping any that were never created.""" + tables: dict[str, Any] = {} + for record_type in constants.RECORD_TYPES: + try: + tables[record_type] = catalog.load_table( + (constants.GLUE_DATABASE, table_name(record_type, run_id)) + ) + except NoSuchTableError: + continue + return tables + + +def drop_tables(catalog: Catalog, run_id: str) -> list[str]: + """Drop this run's Glue tables. Does not remove the S3 objects behind them.""" + dropped = [] + for record_type in constants.RECORD_TYPES: + name = table_name(record_type, run_id) + try: + catalog.drop_table((constants.GLUE_DATABASE, name)) + dropped.append(name) + except NoSuchTableError: + continue + return dropped diff --git a/experimentation/iceberg/constants.py b/experimentation/iceberg/constants.py new file mode 100644 index 00000000..94e0f2af --- /dev/null +++ b/experimentation/iceberg/constants.py @@ -0,0 +1,84 @@ +"""Tunables and fixed identifiers for the Iceberg write-amplification experiment.""" + +from __future__ import annotations + +# --- Bluesky Jetstream ------------------------------------------------------- + +JETSTREAM_ENDPOINT = "wss://jetstream2.us-east.bsky.network/subscribe" + +# Jetstream collection NSID -> the short record type we bucket it into. +COLLECTIONS: dict[str, str] = { + "app.bsky.feed.post": "posts", + "app.bsky.feed.like": "likes", + "app.bsky.feed.repost": "reposts", + "app.bsky.graph.follow": "follows", +} + +RECORD_TYPES: tuple[str, ...] = ("posts", "likes", "reposts", "follows") + +DEFAULT_CAPTURE_SECONDS = 600 # 10 minutes + +# --- S3 / Glue --------------------------------------------------------------- + +S3_BUCKET = "lab-data-integrations-interface" +S3_EXPERIMENT_PREFIX = "experiments/iceberg" +AWS_REGION = "us-east-2" + +# Dedicated Glue database so the experiment never touches `default`. +GLUE_DATABASE = "iceberg_experiments" + +# --- Write path -------------------------------------------------------------- + +DEFAULT_FLUSH_SECONDS = 60 # -> 10 commits per table across a 10-minute replay + +# Jetstream `createdAt` is client-supplied and occasionally garbage (epoch 0, +# year 2100). Anything further than this from the broker-side ingest timestamp +# falls back to ingest time so we don't spawn junk daily partitions. +MAX_CREATED_AT_SKEW_SECONDS = 86_400 + +# --- Pricing (us-east-2, USD) ------------------------------------------------ +# Used to turn measured request counts into a cost model. Verify against +# current AWS pricing before quoting these numbers externally. + +COST_PER_PUT_REQUEST = 0.005 / 1_000 # PUT, COPY, POST, LIST +COST_PER_GET_REQUEST = 0.0004 / 1_000 # GET, SELECT, and all other requests +COST_PER_DELETE_REQUEST = 0.0 # DELETE and CANCEL are free +COST_PER_GLUE_REQUEST = 1.00 / 100_000 +COST_PER_GB_MONTH_STANDARD = 0.023 + +# S3 API operations billed at the (expensive) PUT/LIST rate. Everything else +# that isn't a DELETE falls through to the GET rate. +PUT_TIER_OPERATIONS: frozenset[str] = frozenset( + { + "PutObject", + "CopyObject", + "PostObject", + "ListObjects", + "ListObjectsV2", + "ListBuckets", + "ListMultipartUploads", + "ListParts", + "CreateMultipartUpload", + "UploadPart", + "UploadPartCopy", + "CompleteMultipartUpload", + } +) + +DELETE_TIER_OPERATIONS: frozenset[str] = frozenset( + {"DeleteObject", "DeleteObjects", "AbortMultipartUpload"} +) + +# --- Phases ------------------------------------------------------------------ + +PHASE_RAW_WRITE = "raw_write" +PHASE_ICEBERG_APPEND = "iceberg_append" +PHASE_COMPACT_DEDUP = "compact_dedup" +PHASE_EXPIRE_METADATA = "expire_metadata" + +PHASES: tuple[str, ...] = ( + PHASE_RAW_WRITE, + PHASE_ICEBERG_APPEND, + PHASE_COMPACT_DEDUP, + PHASE_EXPIRE_METADATA, +) diff --git a/experimentation/iceberg/data/results/20260723_075417-report.md b/experimentation/iceberg/data/results/20260723_075417-report.md new file mode 100644 index 00000000..f8e3336d --- /dev/null +++ b/experimentation/iceberg/data/results/20260723_075417-report.md @@ -0,0 +1,97 @@ +# Iceberg write-amplification experiment -- `20260723_075417` + +- Capture: `experimentation/iceberg/data/captures/jetstream-20260723-070252.jsonl.gz` +- Warehouse: `s3://lab-data-integrations-interface/experiments/iceberg/20260723_075417/warehouse` +- Flush window: 60s -> 10 batches per table +- Rows: 131,825 posts=19,440, likes=86,742, reposts=13,666, follows=11,977 +- `createdAt` fallbacks: 7,103 rows partitioned by ingest time, of which + - 6,658 (5.05%) are `delete` events, which carry no record body and therefore have no `createdAt` at all -- structural, not a data problem + - 445 (0.34%) parse cleanly but sit more than 24h from the broker timestamp. These are mostly archive-import bots stamping genuine historical dates; the skew rule rewrites them to keep one bot from opening a daily partition per historical date it touches + +## Cost and request counts by phase + +| Phase | Wall (s) | Calls | PUT-tier | GET-tier | DELETE | Glue | Retries | Cost (USD) | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| `raw_write` | 50.7 | 50 | 50 | 0 | 0 | 0 | 0 | $0.000250 | +| `iceberg_append` | 282.2 | 572 | 210 | 282 | 0 | 80 | 0 | $0.001963 | +| `compact_dedup` | 144.9 | 516 | 30 | 470 | 0 | 16 | 0 | $0.000498 | +| `expire_metadata` | 47.6 | 88 | 12 | 52 | 4 | 20 | 0 | $0.000281 | + +## Where the requests go + +| Phase | Object class | Calls | Operations | +|---|---|---:|---| +| `raw_write` | raw | 50 | PutObject x50 | +| `iceberg_append` | metadata-json | 160 | HeadObject x80, GetObject x40, PutObject x40 | +| `iceberg_append` | other | 120 | GetTable x40, ListObjectsV2 x40, UpdateTable x40 | +| `iceberg_append` | manifest-list | 112 | PutObject x40, HeadObject x36, GetObject x36 | +| `iceberg_append` | data | 100 | PutObject x50, HeadObject x50 | +| `iceberg_append` | manifest | 80 | PutObject x40, HeadObject x40 | +| `compact_dedup` | manifest | 264 | HeadObject x132, GetObject x124, PutObject x8 | +| `compact_dedup` | data | 152 | GetObject x90, HeadObject x56, PutObject x6 | +| `compact_dedup` | manifest-list | 48 | HeadObject x20, GetObject x20, PutObject x8 | +| `compact_dedup` | metadata-json | 32 | HeadObject x16, GetObject x12, PutObject x4 | +| `compact_dedup` | other | 20 | GetTable x12, ListObjectsV2 x4, UpdateTable x4 | +| `expire_metadata` | metadata-json | 40 | HeadObject x20, GetObject x16, PutObject x4 | +| `expire_metadata` | other | 32 | GetTable x16, ListObjectsV2 x8, UpdateTable x4, DeleteObjects x4 | +| `expire_metadata` | manifest-list | 8 | HeadObject x4, GetObject x4 | +| `expire_metadata` | manifest | 8 | HeadObject x4, GetObject x4 | + +## Latency + +| Phase | p50 (ms) | p95 (ms) | p99 (ms) | max (ms) | Total AWS time (s) | +|---|---:|---:|---:|---:|---:| +| `raw_write` | 783 | 2133 | 3207 | 3207 | 49.8 | +| `iceberg_append` | 373 | 1145 | 2017 | 8757 | 315.2 | +| `compact_dedup` | 352 | 871 | 1295 | 7200 | 237.4 | +| `expire_metadata` | 351 | 1205 | 1341 | 2339 | 43.8 | + +## Iceberg vs. raw Parquet + +| Metric | Raw Parquet | Iceberg | Ratio | +|---|---:|---:|---:| +| PUT-tier requests | 50 | 210 | 4.20x | +| Total AWS calls | 50 | 572 | 11.44x | +| Bytes uploaded | 17.27 MiB | 17.92 MiB | 1.04x | +| Wall time | 50.7s | 282.2s | 5.57x | +| Cost | $0.000250 | $0.001963 | 7.85x | + +**Write amplification.** 17.27 MiB of logical Parquet produced 17.92 MiB of Iceberg uploads (1.04x) across 210 PUT-tier requests. + +## Compaction and deduplication + +| Table | Files before | Files after | Rows before | Rows after | Redelivered dupes | Lifecycle collapses | Tombstones dropped | Bytes before | Bytes after | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| posts | 19 | 2 | 19,440 | 17,123 | 0 | 212 | 2,105 (10.8%) | 3.36 MiB | 3.01 MiB | +| likes | 11 | 2 | 86,742 | 84,950 | 0 | 623 | 1,169 (1.3%) | 11.08 MiB | 9.76 MiB | +| reposts | 10 | 1 | 13,666 | 13,057 | 0 | 77 | 532 (3.9%) | 1.85 MiB | 1.56 MiB | +| follows | 10 | 1 | 11,977 | 8,984 | 0 | 143 | 2,850 (23.8%) | 0.98 MiB | 0.81 MiB | + +**Redelivered duplicates across all tables: 0.** A redelivered duplicate is the identical event twice -- same `uri` *and* same `cid`. Lifecycle collapses are different: several distinct events about one record (create then delete, create then update), each with its own `cid`. Only the first is a stream defect; the second is a record's history being materialised into current state. + +Tombstones are `delete` rows, which carry no record body. They are dropped at compaction. Note this only cancels a create that is in the same table -- a delete of a record written before this table existed has nothing to reconcile against and is simply discarded. + +## Snapshot expiry and orphan cleanup + +| Table | Snapshots before | after | Objects listed | Orphans deleted | Reclaimed | +|---|---:|---:|---:|---:|---:| +| posts | 12 | 1 | 58 | 41 | 3.46 MiB | +| likes | 12 | 1 | 50 | 33 | 11.17 MiB | +| reposts | 12 | 1 | 48 | 32 | 1.93 MiB | +| follows | 12 | 1 | 48 | 32 | 1.07 MiB | + +## Extrapolation + +Measured window covers **10.0 minutes** of firehose (131,825 rows, 220 rows/s). + +| Horizon | Rows | Ingest cost | Maintenance cost | Total | +|---|---:|---:|---:|---:| +| Measured run | 131,825 | $0.0022 | $0.0011 | $0.0033 | +| 24 hours | 18,982,800 | $0.32 | $0.16 | $0.48 | +| 30 days | 569,484,000 | $9.56 | $4.87 | $14.43 | + +_Maintenance is extrapolated at the same per-window frequency as the measured run. In production you would compact far less often than every flush, so treat this as an upper bound on the maintenance column._ + +--- + +Pricing model (us-east-2): PUT/LIST $0.005/1k, GET $0.0004/1k, DELETE free, Glue $1/100k requests. Storage is not included in the per-phase cost column. \ No newline at end of file diff --git a/experimentation/iceberg/data/results/20260723_075417-results.json b/experimentation/iceberg/data/results/20260723_075417-results.json new file mode 100644 index 00000000..2d28bda2 --- /dev/null +++ b/experimentation/iceberg/data/results/20260723_075417-results.json @@ -0,0 +1,428 @@ +{ + "run_id": "20260723_075417", + "capture": "experimentation/iceberg/data/captures/jetstream-20260723-070252.jsonl.gz", + "flush_seconds": 60, + "batches": 10, + "ingest_seconds": 333.7401512910146, + "rows_by_type": { + "posts": 19440, + "likes": 86742, + "reposts": 13666, + "follows": 11977 + }, + "total_rows": 131825, + "created_at_fallback_delete_rows": 6658, + "created_at_fallback_malformed_rows": 445, + "raw_bytes_written": 18113241, + "raw_objects_written": 50, + "pre_maintenance": { + "posts": { + "file_count": 19, + "total_bytes": 3524422, + "record_count": 19440, + "avg_file_bytes": 185495.8947368421, + "snapshots": 10 + }, + "likes": { + "file_count": 11, + "total_bytes": 11621247, + "record_count": 86742, + "avg_file_bytes": 1056477.0, + "snapshots": 10 + }, + "reposts": { + "file_count": 10, + "total_bytes": 1935969, + "record_count": 13666, + "avg_file_bytes": 193596.9, + "snapshots": 10 + }, + "follows": { + "file_count": 10, + "total_bytes": 1031079, + "record_count": 11977, + "avg_file_bytes": 103107.9, + "snapshots": 10 + } + }, + "post_maintenance": { + "posts": { + "file_count": 2, + "total_bytes": 3152810, + "record_count": 17123, + "avg_file_bytes": 1576405.0, + "snapshots": 1 + }, + "likes": { + "file_count": 2, + "total_bytes": 10237942, + "record_count": 84950, + "avg_file_bytes": 5118971.0, + "snapshots": 1 + }, + "reposts": { + "file_count": 1, + "total_bytes": 1632368, + "record_count": 13057, + "avg_file_bytes": 1632368.0, + "snapshots": 1 + }, + "follows": { + "file_count": 1, + "total_bytes": 848702, + "record_count": 8984, + "avg_file_bytes": 848702.0, + "snapshots": 1 + } + }, + "compaction": { + "posts": { + "skipped": false, + "file_count_before": 19, + "file_count_after": 2, + "bytes_before": 3524422, + "bytes_after": 3152810, + "rows_before": 19440, + "rows_after": 17123, + "tombstones_dropped": 2105, + "tombstone_pct": 10.828189300411523, + "redelivered_duplicates": 0, + "lifecycle_collapses": 212 + }, + "likes": { + "skipped": false, + "file_count_before": 11, + "file_count_after": 2, + "bytes_before": 11621247, + "bytes_after": 10237942, + "rows_before": 86742, + "rows_after": 84950, + "tombstones_dropped": 1169, + "tombstone_pct": 1.3476747135182494, + "redelivered_duplicates": 0, + "lifecycle_collapses": 623 + }, + "reposts": { + "skipped": false, + "file_count_before": 10, + "file_count_after": 1, + "bytes_before": 1935969, + "bytes_after": 1632368, + "rows_before": 13666, + "rows_after": 13057, + "tombstones_dropped": 532, + "tombstone_pct": 3.8928728230645393, + "redelivered_duplicates": 0, + "lifecycle_collapses": 77 + }, + "follows": { + "skipped": false, + "file_count_before": 10, + "file_count_after": 1, + "bytes_before": 1031079, + "bytes_after": 848702, + "rows_before": 11977, + "rows_after": 8984, + "tombstones_dropped": 2850, + "tombstone_pct": 23.79560824914419, + "redelivered_duplicates": 0, + "lifecycle_collapses": 143 + } + }, + "expiry": { + "posts": { + "snapshots_before": 12, + "snapshots_after": 1, + "expired": 11, + "objects_listed": 58, + "objects_referenced": 17, + "orphans_deleted": 41, + "orphan_bytes_reclaimed": 3625714 + }, + "likes": { + "snapshots_before": 12, + "snapshots_after": 1, + "expired": 11, + "objects_listed": 50, + "objects_referenced": 17, + "orphans_deleted": 33, + "orphan_bytes_reclaimed": 11709348 + }, + "reposts": { + "snapshots_before": 12, + "snapshots_after": 1, + "expired": 11, + "objects_listed": 48, + "objects_referenced": 16, + "orphans_deleted": 32, + "orphan_bytes_reclaimed": 2022997 + }, + "follows": { + "snapshots_before": 12, + "snapshots_after": 1, + "expired": 11, + "objects_listed": 48, + "objects_referenced": 16, + "orphans_deleted": 32, + "orphan_bytes_reclaimed": 1117046 + } + }, + "warehouse": "s3://lab-data-integrations-interface/experiments/iceberg/20260723_075417/warehouse", + "phases": { + "unattributed": { + "wall_seconds": 0.0, + "calls": 181, + "attempts": 181, + "request_bytes": 17840, + "response_bytes": 845541, + "cost_usd": 0.00034920000000000003, + "by_tier": { + "glue": 25, + "get": 148, + "put": 8, + "delete": 0 + }, + "by_operation": { + "glue:CreateDatabase": 1, + "s3:HeadObject": 76, + "s3:ListObjectsV2": 4, + "s3:PutObject": 4, + "glue:CreateTable": 4, + "glue:GetTable": 20, + "s3:GetObject": 72 + }, + "by_key_class": { + "other": 29, + "metadata-json": 48, + "manifest-list": 16, + "manifest": 88 + }, + "by_key_class_operation": { + "other": { + "CreateDatabase": 1, + "ListObjectsV2": 4, + "CreateTable": 4, + "GetTable": 20 + }, + "metadata-json": { + "HeadObject": 24, + "PutObject": 4, + "GetObject": 20 + }, + "manifest-list": { + "HeadObject": 8, + "GetObject": 8 + }, + "manifest": { + "HeadObject": 44, + "GetObject": 44 + } + }, + "latency_p50_ms": 352.1348328795284, + "latency_p95_ms": 1159.5164169557393, + "latency_p99_ms": 1713.4401248767972, + "latency_max_ms": 2389.246250037104 + }, + "raw_write": { + "wall_seconds": 50.6769603791181, + "calls": 50, + "attempts": 50, + "request_bytes": 18113241, + "response_bytes": 0, + "cost_usd": 0.00025, + "by_tier": { + "put": 50, + "get": 0, + "delete": 0, + "glue": 0 + }, + "by_operation": { + "s3:PutObject": 50 + }, + "by_key_class": { + "raw": 50 + }, + "by_key_class_operation": { + "raw": { + "PutObject": 50 + } + }, + "latency_p50_ms": 783.4664168767631, + "latency_p95_ms": 2133.2394171040505, + "latency_p99_ms": 3207.422041101381, + "latency_max_ms": 3207.422041101381 + }, + "iceberg_append": { + "wall_seconds": 282.1545720384456, + "calls": 572, + "attempts": 572, + "request_bytes": 18793945, + "response_bytes": 19055222, + "cost_usd": 0.0019628000000000002, + "by_tier": { + "put": 210, + "get": 282, + "glue": 80, + "delete": 0 + }, + "by_operation": { + "s3:PutObject": 170, + "s3:HeadObject": 206, + "glue:GetTable": 40, + "s3:GetObject": 76, + "s3:ListObjectsV2": 40, + "glue:UpdateTable": 40 + }, + "by_key_class": { + "data": 100, + "manifest": 80, + "manifest-list": 112, + "other": 120, + "metadata-json": 160 + }, + "by_key_class_operation": { + "data": { + "PutObject": 50, + "HeadObject": 50 + }, + "manifest": { + "PutObject": 40, + "HeadObject": 40 + }, + "manifest-list": { + "PutObject": 40, + "HeadObject": 36, + "GetObject": 36 + }, + "other": { + "GetTable": 40, + "ListObjectsV2": 40, + "UpdateTable": 40 + }, + "metadata-json": { + "HeadObject": 80, + "GetObject": 40, + "PutObject": 40 + } + }, + "latency_p50_ms": 372.5530421361327, + "latency_p95_ms": 1144.6151249110699, + "latency_p99_ms": 2016.860333038494, + "latency_max_ms": 8757.476707920432 + }, + "compact_dedup": { + "wall_seconds": 144.8584561671596, + "calls": 516, + "attempts": 516, + "request_bytes": 16016901, + "response_bytes": 56568988, + "cost_usd": 0.0004980000000000001, + "by_tier": { + "glue": 16, + "get": 470, + "put": 30, + "delete": 0 + }, + "by_operation": { + "glue:GetTable": 12, + "s3:HeadObject": 224, + "s3:GetObject": 246, + "s3:PutObject": 26, + "s3:ListObjectsV2": 4, + "glue:UpdateTable": 4 + }, + "by_key_class": { + "other": 20, + "metadata-json": 32, + "manifest-list": 48, + "manifest": 264, + "data": 152 + }, + "by_key_class_operation": { + "other": { + "GetTable": 12, + "ListObjectsV2": 4, + "UpdateTable": 4 + }, + "metadata-json": { + "HeadObject": 16, + "GetObject": 12, + "PutObject": 4 + }, + "manifest-list": { + "HeadObject": 20, + "GetObject": 20, + "PutObject": 8 + }, + "manifest": { + "HeadObject": 132, + "GetObject": 124, + "PutObject": 8 + }, + "data": { + "HeadObject": 56, + "GetObject": 90, + "PutObject": 6 + } + }, + "latency_p50_ms": 351.57895809970796, + "latency_p95_ms": 871.2102919816971, + "latency_p99_ms": 1294.8067081160843, + "latency_max_ms": 7199.553833110258 + }, + "expire_metadata": { + "wall_seconds": 47.621370207984, + "calls": 88, + "attempts": 88, + "request_bytes": 52589, + "response_bytes": 379172, + "cost_usd": 0.0002808, + "by_tier": { + "glue": 20, + "get": 52, + "put": 12, + "delete": 4 + }, + "by_operation": { + "glue:GetTable": 16, + "s3:HeadObject": 28, + "s3:GetObject": 24, + "s3:ListObjectsV2": 8, + "s3:PutObject": 4, + "glue:UpdateTable": 4, + "s3:DeleteObjects": 4 + }, + "by_key_class": { + "other": 32, + "metadata-json": 40, + "manifest-list": 8, + "manifest": 8 + }, + "by_key_class_operation": { + "other": { + "GetTable": 16, + "ListObjectsV2": 8, + "UpdateTable": 4, + "DeleteObjects": 4 + }, + "metadata-json": { + "HeadObject": 20, + "GetObject": 16, + "PutObject": 4 + }, + "manifest-list": { + "HeadObject": 4, + "GetObject": 4 + }, + "manifest": { + "HeadObject": 4, + "GetObject": 4 + } + }, + "latency_p50_ms": 350.61270906589925, + "latency_p95_ms": 1204.7317079268396, + "latency_p99_ms": 1340.9220839384943, + "latency_max_ms": 2339.2392091918737 + } + } +} \ No newline at end of file diff --git a/experimentation/iceberg/iceberg_writer.py b/experimentation/iceberg/iceberg_writer.py new file mode 100644 index 00000000..84f1cf31 --- /dev/null +++ b/experimentation/iceberg/iceberg_writer.py @@ -0,0 +1,53 @@ +"""Iceberg write path -- one ``append`` per flush batch per table. + +Every append is a full Iceberg commit: data files, then a manifest, then a +manifest list, then a new ``metadata.json``, then a Glue ``UpdateTable`` to swap +the pointer. Ten flushes across four tables therefore produce forty commits, +which is precisely the small-files regime this experiment exists to price. +""" + +from __future__ import annotations + +from typing import Any + +from experimentation.iceberg.raw_writer import build_arrow + + +def append_batch(table: Any, rows: list[dict[str, Any]]) -> int: + """Append one flush batch to its table. Returns the row count committed. + + Builds the Arrow table from ``table.schema()`` -- the ids the catalog + actually assigned -- not from the declared schema. Iceberg matches columns + by field id, so using the declared ids would stamp the Parquet footers with + ids the table metadata does not know about, and every such column would read + back as NULL. See the field-id note in ``schemas.py``. + """ + arrow_table = build_arrow(table.schema().as_arrow(), rows) + table.append(arrow_table) + return len(rows) + + +def table_file_stats(table: Any) -> dict[str, Any]: + """Current-snapshot file count and total size, read from table metadata. + + Uses ``inspect.files()`` rather than an S3 listing so it reflects what + Iceberg believes it owns, not what happens to be sitting in the bucket. + """ + table.refresh() + if table.current_snapshot() is None: + return {"file_count": 0, "total_bytes": 0, "record_count": 0} + + files = table.inspect.files() + sizes = files.column("file_size_in_bytes").to_pylist() if files.num_rows else [] + records = files.column("record_count").to_pylist() if files.num_rows else [] + return { + "file_count": len(sizes), + "total_bytes": sum(sizes), + "record_count": sum(records), + "avg_file_bytes": (sum(sizes) / len(sizes)) if sizes else 0, + } + + +def snapshot_count(table: Any) -> int: + table.refresh() + return len(table.snapshots()) diff --git a/experimentation/iceberg/maintenance.py b/experimentation/iceberg/maintenance.py new file mode 100644 index 00000000..3fe7fb04 --- /dev/null +++ b/experimentation/iceberg/maintenance.py @@ -0,0 +1,216 @@ +"""The two maintenance operations, and the sweeps that finish what they leave. + +``compact_table`` rewrites a table's live data into one file per partition, +collapsing each AT-URI to its latest state and dropping delete tombstones. +``expire_snapshots`` + ``sweep_orphans`` drop old snapshots and then remove the +S3 objects nothing references any more. + +PyIceberg 0.11 has no ``rewrite_data_files``, so compaction is implemented as +scan -> collapse -> ``overwrite``, which is a delete-all + append in one commit. + +Two things measured here are easy to conflate, so they are counted separately: + +- **Redelivered duplicates** -- the same event arriving twice, identified by an + identical ``(uri, cid)``. These are what "deduplication" normally means. In a + stable 10-minute Jetstream capture there were *zero* of them. +- **Lifecycle collapses** -- several distinct events about one URI (create then + delete, create then update), each with its own ``cid``. Collapsing these is a + state-materialisation decision, not deduplication. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +import boto3 +import pyarrow as pa +import pyarrow.compute as pc + +from experimentation.iceberg import constants + + +def _collapse_to_latest(table: pa.Table) -> tuple[pa.Table, dict[str, int]]: + """Reduce each ``uri`` to its most recently ingested row. + + Decomposes what was removed into redelivered duplicates (same ``uri`` *and* + ``cid``) and lifecycle collapses (same ``uri``, different ``cid``), because + the two have completely different causes and only the first is a stream + defect. + """ + if table.num_rows == 0: + return table, {"redelivered_duplicates": 0, "lifecycle_collapses": 0} + + # Sort so the surviving row per URI is the newest, then take the first + # occurrence of each URI. + sorted_table = table.sort_by([("uri", "ascending"), ("ingested_at", "descending")]) + uris = sorted_table.column("uri").to_pylist() + cids = sorted_table.column("cid").to_pylist() + + keep_indices: list[int] = [] + previous = object() + for index, uri in enumerate(uris): + if uri != previous: + keep_indices.append(index) + previous = uri + + distinct_events = len(set(zip(uris, cids, strict=True))) + stats = { + # Same URI and same content hash -> the identical event twice. + "redelivered_duplicates": table.num_rows - distinct_events, + # Distinct events about one URI -> a record's history. + "lifecycle_collapses": distinct_events - len(keep_indices), + } + return sorted_table.take(pa.array(keep_indices)), stats + + +def _drop_tombstones(table: pa.Table) -> tuple[pa.Table, int]: + """Remove ``delete`` rows, which carry no record body. + + Caveat worth knowing: this only reconciles a delete against a create that is + *in the same table*. A delete of a record written before this table existed + leaves nothing behind to cancel, and the tombstone is simply discarded. + """ + if table.num_rows == 0: + return table, 0 + kept = table.filter(pc.not_equal(table.column("operation"), "delete")) + return kept, table.num_rows - kept.num_rows + + +def compact_table(table: Any) -> dict[str, Any]: + """Compact to one file per partition, collapse each URI, drop tombstones.""" + table.refresh() + if table.current_snapshot() is None: + return {"skipped": True} + + files_before = table.inspect.files() + file_count_before = files_before.num_rows + bytes_before = sum(files_before.column("file_size_in_bytes").to_pylist()) + + scanned = table.scan().to_arrow() + rows_before = scanned.num_rows + + collapsed, collapse_stats = _collapse_to_latest(scanned) + final, tombstones = _drop_tombstones(collapsed) + + # overwrite() with no filter is delete-all + append in a single commit -- + # full-table compaction. + table.overwrite(final) + + table.refresh() + files_after = table.inspect.files() + bytes_after = sum(files_after.column("file_size_in_bytes").to_pylist()) + + return { + "skipped": False, + "file_count_before": file_count_before, + "file_count_after": files_after.num_rows, + "bytes_before": bytes_before, + "bytes_after": bytes_after, + "rows_before": rows_before, + "rows_after": final.num_rows, + "tombstones_dropped": tombstones, + "tombstone_pct": (tombstones / rows_before * 100) if rows_before else 0.0, + **collapse_stats, + } + + +def expire_snapshots(table: Any) -> dict[str, Any]: + """Expire every snapshot except the current one. + + ``older_than(now)`` rather than ``by_ids`` so the current snapshot is + protected by PyIceberg's own retention rules rather than by our bookkeeping. + """ + table.refresh() + before = len(table.snapshots()) + if before <= 1: + return {"snapshots_before": before, "snapshots_after": before, "expired": 0} + + # `older_than` protects the current snapshot itself, so passing "now" expires + # everything else without us having to track ids. + table.maintenance.expire_snapshots().older_than(datetime.now(UTC)).commit() + + table.refresh() + after = len(table.snapshots()) + return {"snapshots_before": before, "snapshots_after": after, "expired": before - after} + + +def _list_keys(client: Any, prefix: str) -> list[dict[str, Any]]: + """List every object under ``prefix``, following pagination.""" + keys: list[dict[str, Any]] = [] + paginator = client.get_paginator("list_objects_v2") + for page in paginator.paginate(Bucket=constants.S3_BUCKET, Prefix=prefix): + keys.extend(page.get("Contents", [])) + return keys + + +def _delete_keys(client: Any, keys: list[str]) -> int: + """Batch-delete keys 1000 at a time (the DeleteObjects limit).""" + deleted = 0 + for start in range(0, len(keys), 1000): + chunk = keys[start : start + 1000] + client.delete_objects( + Bucket=constants.S3_BUCKET, + Delete={"Objects": [{"Key": key} for key in chunk], "Quiet": True}, + ) + deleted += len(chunk) + return deleted + + +def _referenced_paths(table: Any) -> set[str]: + """Every S3 key the table's live metadata still points at. + + Covers data files plus the manifest and manifest-list chain of all remaining + snapshots, plus the current and logged ``metadata.json`` files. + """ + referenced: set[str] = set() + + def add(location: str | None) -> None: + if location and location.startswith("s3://"): + referenced.add(location.split("/", 3)[3]) + + table.refresh() + add(table.metadata_location) + for entry in table.metadata.metadata_log: + add(entry.metadata_file) + + io = table.io + for snapshot in table.snapshots(): + add(snapshot.manifest_list) + try: + for manifest in snapshot.manifests(io): + add(manifest.manifest_path) + for entry in manifest.fetch_manifest_entry(io, discard_deleted=False): + add(entry.data_file.file_path) + except Exception: + # A manifest list already deleted by expiry is expected here; the + # sweep just skips whatever it can no longer read. + continue + + return referenced + + +def sweep_orphans(table: Any, region: str = constants.AWS_REGION) -> dict[str, Any]: + """Delete objects under the table prefix that no live metadata references. + + This is the part ``expire_snapshots`` does not do on its own. Counted + separately so the LIST + DELETE cost of metadata hygiene is visible. + """ + client = boto3.client("s3", region_name=region) + location = table.location() + prefix = location.split("/", 3)[3].rstrip("/") + "/" + + listed = _list_keys(client, prefix) + referenced = _referenced_paths(table) + + orphans = [obj["Key"] for obj in listed if obj["Key"] not in referenced] + orphan_bytes = sum(obj["Size"] for obj in listed if obj["Key"] not in referenced) + + deleted = _delete_keys(client, orphans) if orphans else 0 + + return { + "objects_listed": len(listed), + "objects_referenced": len(referenced), + "orphans_deleted": deleted, + "orphan_bytes_reclaimed": orphan_bytes, + } diff --git a/experimentation/iceberg/raw_writer.py b/experimentation/iceberg/raw_writer.py new file mode 100644 index 00000000..023b09b2 --- /dev/null +++ b/experimentation/iceberg/raw_writer.py @@ -0,0 +1,78 @@ +"""Baseline write path: buffer -> Parquet -> S3, no table format at all. + +This is the control. Subtracting these numbers from the Iceberg append numbers +gives the metadata tax in isolation: same rows, same Parquet encoder, same +partition layout, the only difference being that nothing tracks a manifest. +""" + +from __future__ import annotations + +import io +from typing import Any + +import boto3 +import pyarrow as pa +import pyarrow.parquet as pq + +from experimentation.iceberg import constants, schemas + + +def build_arrow(schema: pa.Schema, rows: list[dict[str, Any]]) -> pa.Table: + """Project ``rows`` onto an explicit Arrow schema. + + Takes the schema as an argument rather than looking it up, because the + Iceberg path must build against the *table's* schema -- see the field-id + note in ``schemas.py``. + """ + columns = {field.name: [row.get(field.name) for row in rows] for field in schema} + return pa.Table.from_pydict(columns, schema=schema) + + +def rows_to_arrow(record_type: str, rows: list[dict[str, Any]]) -> pa.Table: + """Build an Arrow table from the declared schema for ``record_type``. + + Only for the raw baseline, which never round-trips through Iceberg metadata + and so is unaffected by field-id assignment. + """ + return build_arrow(schemas.SCHEMAS[record_type].as_arrow(), rows) + + +class RawWriter: + """Writes each flush batch as one Parquet object under a daily prefix.""" + + def __init__(self, run_id: str, region: str = constants.AWS_REGION) -> None: + self.run_id = run_id + self.client = boto3.client("s3", region_name=region) + self.bytes_written = 0 + self.objects_written = 0 + + def _key(self, record_type: str, day: str, batch_index: int) -> str: + return ( + f"{constants.S3_EXPERIMENT_PREFIX}/{self.run_id}/raw/" + f"{record_type}/created_at_day={day}/batch-{batch_index:05d}.parquet" + ) + + def write_batch(self, record_type: str, rows: list[dict[str, Any]], batch_index: int) -> int: + """Write ``rows`` as Parquet, split by day so the layout matches Iceberg's. + + Returns the number of S3 objects written. + """ + by_day: dict[str, list[dict[str, Any]]] = {} + for row in rows: + by_day.setdefault(row["created_at"].strftime("%Y-%m-%d"), []).append(row) + + for day, day_rows in by_day.items(): + table = rows_to_arrow(record_type, day_rows) + buffer = io.BytesIO() + pq.write_table(table, buffer, compression="zstd") + payload = buffer.getvalue() + + self.client.put_object( + Bucket=constants.S3_BUCKET, + Key=self._key(record_type, day, batch_index), + Body=payload, + ) + self.bytes_written += len(payload) + self.objects_written += 1 + + return len(by_day) diff --git a/experimentation/iceberg/replay.py b/experimentation/iceberg/replay.py new file mode 100644 index 00000000..b4eb323d --- /dev/null +++ b/experimentation/iceberg/replay.py @@ -0,0 +1,68 @@ +"""Turn a captured Jetstream file into deterministic flush batches. + +Batches are cut on the *broker* timestamp carried in each event, not on +wall-clock time during the replay. That makes the batch boundaries a property of +the captured data alone, so two runs on the same capture always produce the same +number of commits containing the same rows -- which is what makes the measured +request counts comparable across configurations. +""" + +from __future__ import annotations + +import gzip +import json +from collections.abc import Generator, Iterable +from pathlib import Path +from typing import Any + +from experimentation.iceberg import constants, schemas + + +def iter_events(capture_path: Path) -> Generator[dict[str, Any], None, None]: + """Yield decoded events from a gzipped JSONL capture, skipping bad lines.""" + with gzip.open(capture_path, "rt", encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if not line: + continue + try: + yield json.loads(line) + except json.JSONDecodeError: + continue + + +def iter_batches( + events: Iterable[dict[str, Any]], + flush_seconds: int = constants.DEFAULT_FLUSH_SECONDS, +) -> Generator[tuple[int, dict[str, list[dict[str, Any]]]], None, None]: + """Group parsed rows into flush windows. + + Yields ``(batch_index, {record_type: rows})``. A window closes as soon as an + event arrives whose ingest timestamp is ``flush_seconds`` past the window + start; the final partial window is always emitted. + """ + buffers: dict[str, list[dict[str, Any]]] = {rt: [] for rt in constants.RECORD_TYPES} + window_start: float | None = None + batch_index = 0 + + for event in events: + parsed = schemas.parse_event(event) + if parsed is None: + continue + record_type, row = parsed + + event_time = row["ingested_at"].timestamp() + if window_start is None: + window_start = event_time + + if event_time - window_start >= flush_seconds: + if any(buffers.values()): + yield batch_index, {rt: rows for rt, rows in buffers.items() if rows} + batch_index += 1 + buffers = {rt: [] for rt in constants.RECORD_TYPES} + window_start = event_time + + buffers[record_type].append(row) + + if any(buffers.values()): + yield batch_index, {rt: rows for rt, rows in buffers.items() if rows} diff --git a/experimentation/iceberg/report.py b/experimentation/iceberg/report.py new file mode 100644 index 00000000..ca93ae0a --- /dev/null +++ b/experimentation/iceberg/report.py @@ -0,0 +1,322 @@ +"""Render the measured call log into a markdown report and a raw CSV.""" + +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Any + +from experimentation.iceberg import constants +from experimentation.iceberg.s3_meter import CallRecord, PhaseStats + +SECONDS_PER_DAY = 86_400 +DAYS_PER_MONTH = 30 + + +def serialize_stats(stats: dict[str, PhaseStats]) -> dict[str, Any]: + """Flatten PhaseStats into JSON-safe dicts.""" + return { + phase: { + "wall_seconds": entry.wall_seconds, + "calls": entry.calls, + "attempts": entry.attempts, + "request_bytes": entry.request_bytes, + "response_bytes": entry.response_bytes, + "cost_usd": entry.cost_usd, + "by_tier": dict(entry.by_tier), + "by_operation": dict(entry.by_operation), + "by_key_class": dict(entry.by_key_class), + "by_key_class_operation": {k: dict(v) for k, v in entry.by_key_class_operation.items()}, + "latency_p50_ms": entry.percentile(50), + "latency_p95_ms": entry.percentile(95), + "latency_p99_ms": entry.percentile(99), + "latency_max_ms": max(entry.latencies_ms) if entry.latencies_ms else 0.0, + } + for phase, entry in stats.items() + } + + +def write_call_log(records: list[CallRecord], path: Path) -> None: + """Dump every individual API call, for slicing outside this report.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as handle: + writer = csv.writer(handle) + writer.writerow( + [ + "phase", + "service", + "operation", + "tier", + "key_class", + "key", + "request_bytes", + "response_bytes", + "duration_ms", + "status", + ] + ) + for record in records: + writer.writerow( + [ + record.phase, + record.service, + record.operation, + record.tier, + record.key_class, + record.key, + record.request_bytes, + record.response_bytes, + f"{record.duration_ms:.2f}", + record.status, + ] + ) + + +def _mib(value: float) -> str: + return f"{value / 1024 / 1024:.2f}" + + +def _phase_table(stats: dict[str, PhaseStats]) -> list[str]: + lines = [ + "| Phase | Wall (s) | Calls | PUT-tier | GET-tier | DELETE | Glue | Retries | Cost (USD) |", + "|---|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for phase in constants.PHASES: + entry = stats.get(phase) + if entry is None: + continue + retries = max(0, entry.attempts - entry.calls) + lines.append( + f"| `{phase}` | {entry.wall_seconds:.1f} | {entry.calls:,} | " + f"{entry.by_tier['put']:,} | {entry.by_tier['get']:,} | " + f"{entry.by_tier['delete']:,} | {entry.by_tier['glue']:,} | " + f"{retries:,} | ${entry.cost_usd:.6f} |" + ) + return lines + + +def _key_class_table(stats: dict[str, PhaseStats]) -> list[str]: + lines = [ + "| Phase | Object class | Calls | Operations |", + "|---|---|---:|---|", + ] + for phase in constants.PHASES: + entry = stats.get(phase) + if entry is None: + continue + for key_class, count in sorted(entry.by_key_class.items(), key=lambda kv: -kv[1]): + ops = entry.by_key_class_operation.get(key_class, {}) + detail = ", ".join(f"{op} x{n}" for op, n in sorted(ops.items(), key=lambda kv: -kv[1])) + lines.append(f"| `{phase}` | {key_class} | {count:,} | {detail} |") + return lines + + +def _latency_table(stats: dict[str, PhaseStats]) -> list[str]: + lines = [ + "| Phase | p50 (ms) | p95 (ms) | p99 (ms) | max (ms) | Total AWS time (s) |", + "|---|---:|---:|---:|---:|---:|", + ] + for phase in constants.PHASES: + entry = stats.get(phase) + if entry is None or not entry.latencies_ms: + continue + lines.append( + f"| `{phase}` | {entry.percentile(50):.0f} | {entry.percentile(95):.0f} | " + f"{entry.percentile(99):.0f} | {max(entry.latencies_ms):.0f} | " + f"{sum(entry.latencies_ms) / 1000:.1f} |" + ) + return lines + + +def _amplification_section(results: dict[str, Any], stats: dict[str, PhaseStats]) -> list[str]: + """Compare the Iceberg path against the raw-Parquet control.""" + raw = stats.get(constants.PHASE_RAW_WRITE) + ice = stats.get(constants.PHASE_ICEBERG_APPEND) + if raw is None or ice is None: + return ["_Baseline skipped, no comparison available._"] + + raw_puts = raw.by_tier["put"] + ice_puts = ice.by_tier["put"] + logical = results.get("raw_bytes_written", 0) + ice_bytes = ice.request_bytes + + lines = [ + "| Metric | Raw Parquet | Iceberg | Ratio |", + "|---|---:|---:|---:|", + f"| PUT-tier requests | {raw_puts:,} | {ice_puts:,} | {ice_puts / raw_puts:.2f}x |" + if raw_puts + else f"| PUT-tier requests | {raw_puts:,} | {ice_puts:,} | n/a |", + f"| Total AWS calls | {raw.calls:,} | {ice.calls:,} | " + + (f"{ice.calls / raw.calls:.2f}x |" if raw.calls else "n/a |"), + f"| Bytes uploaded | {_mib(raw.request_bytes)} MiB | {_mib(ice_bytes)} MiB | " + + (f"{ice_bytes / raw.request_bytes:.2f}x |" if raw.request_bytes else "n/a |"), + f"| Wall time | {raw.wall_seconds:.1f}s | {ice.wall_seconds:.1f}s | " + + (f"{ice.wall_seconds / raw.wall_seconds:.2f}x |" if raw.wall_seconds else "n/a |"), + f"| Cost | ${raw.cost_usd:.6f} | ${ice.cost_usd:.6f} | " + + (f"{ice.cost_usd / raw.cost_usd:.2f}x |" if raw.cost_usd else "n/a |"), + ] + + if logical: + lines += [ + "", + f"**Write amplification.** {_mib(logical)} MiB of logical Parquet produced " + f"{_mib(ice_bytes)} MiB of Iceberg uploads " + f"({ice_bytes / logical:.2f}x) across {ice_puts:,} PUT-tier requests.", + ] + return lines + + +def _extrapolation(results: dict[str, Any], stats: dict[str, PhaseStats]) -> list[str]: + """Project the measured run out to a day and a month at the same rate.""" + capture_seconds = results["batches"] * results["flush_seconds"] + if capture_seconds <= 0: + return [] + + total_cost = sum(entry.cost_usd for entry in stats.values()) + scale_day = SECONDS_PER_DAY / capture_seconds + + ingest_cost = sum( + stats[phase].cost_usd + for phase in (constants.PHASE_RAW_WRITE, constants.PHASE_ICEBERG_APPEND) + if phase in stats + ) + maintenance_cost = total_cost - ingest_cost + + rows_per_day = results["total_rows"] * scale_day + + return [ + f"Measured window covers **{capture_seconds / 60:.1f} minutes** of firehose " + f"({results['total_rows']:,} rows, {results['total_rows'] / capture_seconds:,.0f} rows/s).", + "", + "| Horizon | Rows | Ingest cost | Maintenance cost | Total |", + "|---|---:|---:|---:|---:|", + f"| Measured run | {results['total_rows']:,} | ${ingest_cost:.4f} | " + f"${maintenance_cost:.4f} | ${total_cost:.4f} |", + f"| 24 hours | {rows_per_day:,.0f} | ${ingest_cost * scale_day:.2f} | " + f"${maintenance_cost * scale_day:.2f} | ${total_cost * scale_day:.2f} |", + f"| 30 days | {rows_per_day * DAYS_PER_MONTH:,.0f} | " + f"${ingest_cost * scale_day * DAYS_PER_MONTH:.2f} | " + f"${maintenance_cost * scale_day * DAYS_PER_MONTH:.2f} | " + f"${total_cost * scale_day * DAYS_PER_MONTH:.2f} |", + "", + "_Maintenance is extrapolated at the same per-window frequency as the measured " + "run. In production you would compact far less often than every flush, so treat " + "this as an upper bound on the maintenance column._", + ] + + +def _compaction_section(results: dict[str, Any]) -> list[str]: + lines = [ + "| Table | Files before | Files after | Rows before | Rows after " + "| Redelivered dupes | Lifecycle collapses | Tombstones dropped " + "| Bytes before | Bytes after |", + "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for record_type in constants.RECORD_TYPES: + entry = results["compaction"].get(record_type, {}) + if entry.get("skipped", True): + continue + lines.append( + f"| {record_type} | {entry['file_count_before']} | {entry['file_count_after']} | " + f"{entry['rows_before']:,} | {entry['rows_after']:,} | " + f"{entry['redelivered_duplicates']:,} | " + f"{entry['lifecycle_collapses']:,} | " + f"{entry['tombstones_dropped']:,} ({entry['tombstone_pct']:.1f}%) | " + f"{_mib(entry['bytes_before'])} MiB | {_mib(entry['bytes_after'])} MiB |" + ) + + total_redelivered = sum( + e.get("redelivered_duplicates", 0) + for e in results["compaction"].values() + if not e.get("skipped", True) + ) + lines += [ + "", + f"**Redelivered duplicates across all tables: {total_redelivered:,}.** " + "A redelivered duplicate is the identical event twice -- same `uri` *and* same " + "`cid`. Lifecycle collapses are different: several distinct events about one " + "record (create then delete, create then update), each with its own `cid`. " + "Only the first is a stream defect; the second is a record's history being " + "materialised into current state.", + "", + "Tombstones are `delete` rows, which carry no record body. They are dropped at " + "compaction. Note this only cancels a create that is in the same table -- a " + "delete of a record written before this table existed has nothing to reconcile " + "against and is simply discarded.", + ] + return lines + + +def _expiry_section(results: dict[str, Any]) -> list[str]: + lines = [ + "| Table | Snapshots before | after | Objects listed | Orphans deleted | Reclaimed |", + "|---|---:|---:|---:|---:|---:|", + ] + for record_type in constants.RECORD_TYPES: + entry = results["expiry"].get(record_type, {}) + if not entry: + continue + lines.append( + f"| {record_type} | {entry['snapshots_before']} | {entry['snapshots_after']} | " + f"{entry['objects_listed']} | {entry['orphans_deleted']} | " + f"{_mib(entry['orphan_bytes_reclaimed'])} MiB |" + ) + return lines + + +def render(results: dict[str, Any], stats: dict[str, PhaseStats]) -> str: + """Assemble the full markdown report.""" + deletes = results["created_at_fallback_delete_rows"] + skewed = results["created_at_fallback_malformed_rows"] + total = results["total_rows"] or 1 + + sections: list[str] = [ + f"# Iceberg write-amplification experiment -- `{results['run_id']}`", + "", + f"- Capture: `{results['capture']}`", + f"- Warehouse: `{results['warehouse']}`", + f"- Flush window: {results['flush_seconds']}s -> {results['batches']} batches per table", + f"- Rows: {results['total_rows']:,} " + + ", ".join(f"{rt}={n:,}" for rt, n in results["rows_by_type"].items()), + f"- `createdAt` fallbacks: {deletes + skewed:,} rows partitioned by ingest time, of which", + f" - {deletes:,} ({deletes / total * 100:.2f}%) are `delete` events, which carry no " + "record body and therefore have no `createdAt` at all -- structural, not a data problem", + f" - {skewed:,} ({skewed / total * 100:.2f}%) parse cleanly but sit more than " + f"{constants.MAX_CREATED_AT_SKEW_SECONDS // 3600}h from the broker timestamp. These are " + "mostly archive-import bots stamping genuine historical dates; the skew rule rewrites them " + "to keep one bot from opening a daily partition per historical date it touches", + "", + "## Cost and request counts by phase", + "", + *_phase_table(stats), + "", + "## Where the requests go", + "", + *_key_class_table(stats), + "", + "## Latency", + "", + *_latency_table(stats), + "", + "## Iceberg vs. raw Parquet", + "", + *_amplification_section(results, stats), + "", + "## Compaction and deduplication", + "", + *_compaction_section(results), + "", + "## Snapshot expiry and orphan cleanup", + "", + *_expiry_section(results), + "", + "## Extrapolation", + "", + *_extrapolation(results, stats), + "", + "---", + "", + "Pricing model (us-east-2): PUT/LIST $0.005/1k, GET $0.0004/1k, DELETE free, " + "Glue $1/100k requests. Storage is not included in the per-phase cost column.", + ] + return "\n".join(sections) diff --git a/experimentation/iceberg/requirements.txt b/experimentation/iceberg/requirements.txt new file mode 100644 index 00000000..bdda22cf --- /dev/null +++ b/experimentation/iceberg/requirements.txt @@ -0,0 +1,15 @@ +# Isolated environment for the Iceberg experiment. +# +# This is deliberately NOT wired into the root pyproject.toml: pyiceberg caps +# rich at <15.0.0 and the main project requires rich>=15.0.0, so the two cannot +# share a resolution. Build with: +# +# uv venv --python 3.11 experimentation/iceberg/.venv +# uv pip install --python experimentation/iceberg/.venv -r experimentation/iceberg/requirements.txt +# pyiceberg-core supplies the Rust partition transforms; writing to a +# partitioned table fails without it on 0.11. +pyiceberg[glue,s3fs,sql-sqlite,pyiceberg-core]>=0.10.0 +pyarrow>=14.0.0 +boto3>=1.35.0 +websockets>=15.0.1 +pytest>=8.0.0 diff --git a/experimentation/iceberg/run_experiment.py b/experimentation/iceberg/run_experiment.py new file mode 100644 index 00000000..c67b29d9 --- /dev/null +++ b/experimentation/iceberg/run_experiment.py @@ -0,0 +1,261 @@ +"""Phase 2 -- replay a capture through both write paths and price every operation. + +Order matters at startup: the meter patches ``botocore.session.Session`` and only +sessions constructed afterwards carry the handlers, so ``METER.install()`` runs +before anything touches AWS. + +Usage: + python -m experimentation.iceberg.run_experiment --capture data/captures/.jsonl.gz +""" + +from __future__ import annotations + +import argparse +import json +import time +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from experimentation.iceberg import constants +from experimentation.iceberg.s3_meter import METER + +# Installed before any other project import can build an AWS client. +METER.install() + +from experimentation.iceberg import catalog as catalog_module # noqa: E402 +from experimentation.iceberg import iceberg_writer, maintenance, replay # noqa: E402 +from experimentation.iceberg.raw_writer import RawWriter # noqa: E402 + +RESULTS_DIR = Path(__file__).parent / "data" / "results" + + +def _new_run_id() -> str: + return datetime.now(UTC).strftime("%Y%m%d_%H%M%S") + + +@dataclass +class IngestResult: + """Totals accumulated while replaying the capture through both write paths.""" + + rows_by_type: dict[str, int] + fallback_delete_rows: int + fallback_malformed_rows: int + batches: int + seconds: float + + @property + def total_rows(self) -> int: + return sum(self.rows_by_type.values()) + + +def _count_fallbacks(rows: list[dict[str, Any]]) -> tuple[int, int]: + """Split `createdAt` fallbacks into (delete events, genuinely skewed). + + A delete carries no record body, so it structurally has no `createdAt`. + Lumping those in with bad client timestamps overstates the data-quality + problem by an order of magnitude. + """ + deletes = malformed = 0 + for row in rows: + if not row["created_at_fallback"]: + continue + if row["operation"] == "delete": + deletes += 1 + else: + malformed += 1 + return deletes, malformed + + +def _write_one( + tables: dict[str, Any], + raw_writer: RawWriter | None, + record_type: str, + rows: list[dict[str, Any]], + batch_index: int, +) -> None: + """Send one record type's slice of a batch down both write paths.""" + if raw_writer is not None: + with METER.phase(constants.PHASE_RAW_WRITE): + raw_writer.write_batch(record_type, rows, batch_index) + with METER.phase(constants.PHASE_ICEBERG_APPEND): + iceberg_writer.append_batch(tables[record_type], rows) + + +def _ingest( + capture_path: Path, + tables: dict[str, Any], + raw_writer: RawWriter | None, + flush_seconds: int, + max_batches: int | None, +) -> IngestResult: + """Replay the capture, one flush batch at a time.""" + rows_by_type = {rt: 0 for rt in constants.RECORD_TYPES} + fallback_deletes = fallback_malformed = batches = 0 + started = time.perf_counter() + + events = replay.iter_events(capture_path) + for batch_index, buffers in replay.iter_batches(events, flush_seconds=flush_seconds): + if max_batches is not None and batch_index >= max_batches: + break + batches += 1 + batch_rows = sum(len(rows) for rows in buffers.values()) + print( + f"batch {batch_index:>3} {batch_rows:>7,} rows " + + ", ".join(f"{rt}={len(rows):,}" for rt, rows in buffers.items()) + ) + + for record_type, rows in buffers.items(): + rows_by_type[record_type] += len(rows) + deletes, malformed = _count_fallbacks(rows) + fallback_deletes += deletes + fallback_malformed += malformed + _write_one(tables, raw_writer, record_type, rows, batch_index) + + return IngestResult( + rows_by_type=rows_by_type, + fallback_delete_rows=fallback_deletes, + fallback_malformed_rows=fallback_malformed, + batches=batches, + seconds=time.perf_counter() - started, + ) + + +def _table_state(tables: dict[str, Any]) -> dict[str, Any]: + return { + rt: iceberg_writer.table_file_stats(table) + | {"snapshots": iceberg_writer.snapshot_count(table)} + for rt, table in tables.items() + } + + +def _compact_all(tables: dict[str, Any]) -> dict[str, Any]: + print("compacting, collapsing lifecycles, dropping tombstones...") + results: dict[str, Any] = {} + with METER.phase(constants.PHASE_COMPACT_DEDUP): + for record_type, table in tables.items(): + result = maintenance.compact_table(table) + results[record_type] = result + if result.get("skipped"): + continue + print( + f" {record_type:<9} {result['file_count_before']:>3} -> " + f"{result['file_count_after']:>3} files, " + f"{result['redelivered_duplicates']:,} redelivered, " + f"{result['lifecycle_collapses']:,} lifecycle, " + f"{result['tombstones_dropped']:,} tombstones " + f"({result['tombstone_pct']:.1f}%)" + ) + return results + + +def _expire_all(tables: dict[str, Any]) -> dict[str, Any]: + print("\nexpiring snapshots and sweeping orphans...") + results: dict[str, Any] = {} + with METER.phase(constants.PHASE_EXPIRE_METADATA): + for record_type, table in tables.items(): + expired = maintenance.expire_snapshots(table) + swept = maintenance.sweep_orphans(table) + results[record_type] = expired | swept + print( + f" {record_type:<9} {expired['expired']:>2} snapshots expired, " + f"{swept['orphans_deleted']:>3} orphans deleted " + f"({swept['orphan_bytes_reclaimed'] / 1024 / 1024:.1f} MiB)" + ) + return results + + +def run( + capture_path: Path, + run_id: str, + flush_seconds: int, + max_batches: int | None = None, + include_raw: bool = True, +) -> dict[str, Any]: + """Execute all four measured phases against one capture file.""" + print(f"run_id {run_id}") + print(f"capture {capture_path}") + print(f"flush window {flush_seconds}s") + print(f"warehouse {catalog_module.warehouse_uri(run_id)}\n") + + catalog = catalog_module.build_catalog(run_id) + tables = catalog_module.create_tables(catalog, run_id) + print(f"created {len(tables)} Glue tables in `{constants.GLUE_DATABASE}`\n") + + raw_writer = RawWriter(run_id) if include_raw else None + ingest = _ingest(capture_path, tables, raw_writer, flush_seconds, max_batches) + print( + f"\ningested {ingest.total_rows:,} rows in {ingest.batches} batches " + f"({ingest.seconds:.1f}s)\n" + ) + + pre_maintenance = _table_state(tables) + compaction = _compact_all(tables) + expiry = _expire_all(tables) + post_maintenance = _table_state(tables) + + return { + "run_id": run_id, + "capture": str(capture_path), + "flush_seconds": flush_seconds, + "batches": ingest.batches, + "ingest_seconds": ingest.seconds, + "rows_by_type": ingest.rows_by_type, + "total_rows": ingest.total_rows, + "created_at_fallback_delete_rows": ingest.fallback_delete_rows, + "created_at_fallback_malformed_rows": ingest.fallback_malformed_rows, + "raw_bytes_written": raw_writer.bytes_written if raw_writer else 0, + "raw_objects_written": raw_writer.objects_written if raw_writer else 0, + "pre_maintenance": pre_maintenance, + "post_maintenance": post_maintenance, + "compaction": compaction, + "expiry": expiry, + "warehouse": catalog_module.warehouse_uri(run_id), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Replay a Jetstream capture into S3 + Iceberg.") + parser.add_argument("--capture", type=Path, required=True) + parser.add_argument("--run-id", type=str, default=None) + parser.add_argument("--flush-seconds", type=int, default=constants.DEFAULT_FLUSH_SECONDS) + parser.add_argument("--max-batches", type=int, default=None) + parser.add_argument("--skip-raw", action="store_true", help="Skip the non-Iceberg baseline.") + args = parser.parse_args() + + run_id = args.run_id or _new_run_id() + results = run( + capture_path=args.capture, + run_id=run_id, + flush_seconds=args.flush_seconds, + max_batches=args.max_batches, + include_raw=not args.skip_raw, + ) + + # Import here so the report module cannot pull in botocore before install(). + from experimentation.iceberg import report + + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + stats = METER.summarize() + + results_path = RESULTS_DIR / f"{run_id}-results.json" + results_path.write_text( + json.dumps(results | {"phases": report.serialize_stats(stats)}, indent=2, default=str) + ) + + report_path = RESULTS_DIR / f"{run_id}-report.md" + report_text = report.render(results, stats) + report_path.write_text(report_text) + + calls_path = RESULTS_DIR / f"{run_id}-calls.csv" + report.write_call_log(METER.records, calls_path) + + print(f"\n{report_text}") + print(f"\nwrote {results_path}") + print(f"wrote {report_path}") + print(f"wrote {calls_path}") + + +if __name__ == "__main__": + main() diff --git a/experimentation/iceberg/s3_meter.py b/experimentation/iceberg/s3_meter.py new file mode 100644 index 00000000..7dd14b99 --- /dev/null +++ b/experimentation/iceberg/s3_meter.py @@ -0,0 +1,381 @@ +"""Per-call S3 and Glue instrumentation. + +The whole experiment hangs off this module. PyIceberg's default ``PyArrowFileIO`` +drives a C++ S3 client that cannot be intercepted from Python, so the experiment +configures ``FsspecFileIO`` instead: s3fs -> aiobotocore -> botocore, which means +every request passes through botocore's event system where we can count it. + +Registration works by wrapping ``botocore.session.Session.__init__`` so that +*every* session -- boto3's for raw writes, aiobotocore's inside s3fs, and the +Glue client PyIceberg builds for catalog commits -- gets the handlers. Install +the meter before creating any client. + +Counting happens at ``before-call``/``after-call``, i.e. once per logical API +operation. ``before-send`` fires once per HTTP attempt, so ``attempts`` exceeding +``calls`` is the retry signal. +""" + +from __future__ import annotations + +import threading +import time +from collections import defaultdict +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any +from urllib.parse import unquote, urlparse + +import botocore.session + +from experimentation.iceberg import constants + +# Object classes we bucket S3 keys into. The Iceberg metadata tree is the whole +# point of the experiment, so it gets three separate buckets. +KEY_CLASS_DATA = "data" +KEY_CLASS_MANIFEST = "manifest" +KEY_CLASS_MANIFEST_LIST = "manifest-list" +KEY_CLASS_METADATA_JSON = "metadata-json" +KEY_CLASS_RAW = "raw" +KEY_CLASS_OTHER = "other" + +_UNSET_PHASE = "unattributed" + + +def classify_key(key: str) -> str: + """Bucket an S3 object key into one of the KEY_CLASS_* constants.""" + if not key: + return KEY_CLASS_OTHER + + tail = key.rsplit("/", 1)[-1] + + if "/metadata/" in key or key.startswith("metadata/"): + if tail.endswith(".metadata.json"): + return KEY_CLASS_METADATA_JSON + if tail.startswith("snap-") and tail.endswith(".avro"): + return KEY_CLASS_MANIFEST_LIST + if tail.endswith(".avro"): + return KEY_CLASS_MANIFEST + return KEY_CLASS_OTHER + + if "/raw/" in key: + return KEY_CLASS_RAW + if "/data/" in key or tail.endswith(".parquet"): + return KEY_CLASS_DATA + return KEY_CLASS_OTHER + + +def cost_tier(service: str, operation: str) -> str: + """Return the billing tier -- ``put``, ``get``, ``delete`` or ``glue``.""" + if service == "glue": + return "glue" + if operation in constants.DELETE_TIER_OPERATIONS: + return "delete" + if operation in constants.PUT_TIER_OPERATIONS: + return "put" + return "get" + + +def _extract_key(url: str, bucket: str) -> str: + """Pull the object key out of a request URL, handling both addressing styles.""" + path = unquote(urlparse(url).path).lstrip("/") + # Path-style addressing puts the bucket in front of the key. + if path == bucket: + return "" + if path.startswith(f"{bucket}/"): + return path[len(bucket) + 1 :] + return path + + +def _header_int(headers: Any, name: str) -> int: + """Best-effort Content-Length lookup across dict and HTTPHeaders shapes.""" + if not headers: + return 0 + try: + raw = headers.get(name) or headers.get(name.lower()) + except AttributeError: + return 0 + if isinstance(raw, bytes | bytearray): + raw = raw.decode("ascii", "ignore") + try: + return int(raw) + except (TypeError, ValueError): + return 0 + + +def _body_size(body: Any) -> int: + """Size of an outbound request body without consuming it. + + botocore does not set ``Content-Length`` on the request dict at + ``before-call``: PutObject bodies arrive as a ``BytesIO`` and the length is + only fixed later during signing, at which point large uploads have switched + to ``aws-chunked`` transfer encoding and carry no ``Content-Length`` at all. + Measuring the stream here -- and restoring its position -- is the one place + the true payload size is reliably available. + """ + if body is None: + return 0 + if isinstance(body, bytes | bytearray): + return len(body) + if isinstance(body, str): + return len(body.encode("utf-8")) + + seek: Any = getattr(body, "seek", None) + tell: Any = getattr(body, "tell", None) + if not (callable(seek) and callable(tell)): + return 0 + try: + original = int(tell()) + seek(0, 2) # SEEK_END + size = int(tell()) + seek(original) + return max(0, size - original) + except (OSError, ValueError, TypeError): + return 0 + + +@dataclass +class CallRecord: + """One logical AWS API operation.""" + + phase: str + service: str + operation: str + key: str + key_class: str + tier: str + request_bytes: int = 0 + response_bytes: int = 0 + duration_ms: float = 0.0 + status: int = 0 + + +@dataclass +class PhaseStats: + """Aggregates for a single phase, assembled by :meth:`Meter.summarize`.""" + + phase: str + wall_seconds: float = 0.0 + calls: int = 0 + attempts: int = 0 + request_bytes: int = 0 + response_bytes: int = 0 + by_tier: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + by_operation: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + by_key_class: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + # key_class -> {operation -> count}, the table that actually explains cost. + by_key_class_operation: dict[str, dict[str, int]] = field(default_factory=dict) + latencies_ms: list[float] = field(default_factory=list) + + @property + def cost_usd(self) -> float: + return ( + self.by_tier["put"] * constants.COST_PER_PUT_REQUEST + + self.by_tier["get"] * constants.COST_PER_GET_REQUEST + + self.by_tier["delete"] * constants.COST_PER_DELETE_REQUEST + + self.by_tier["glue"] * constants.COST_PER_GLUE_REQUEST + ) + + def percentile(self, pct: float) -> float: + """Nearest-rank percentile of call latency, in milliseconds.""" + if not self.latencies_ms: + return 0.0 + ordered = sorted(self.latencies_ms) + idx = max(0, min(len(ordered) - 1, int(round(pct / 100.0 * len(ordered))) - 1)) + return ordered[idx] + + +class Meter: + """Thread-safe collector for AWS calls, attributed to the active phase. + + PyIceberg fans manifest and data-file work out across a thread pool, so the + counters take a lock and the active phase is module-level rather than a + ``contextvar`` (context does not propagate into ``ThreadPoolExecutor`` + workers). Phases run sequentially, so a plain global is correct here. + """ + + def __init__(self, bucket: str = constants.S3_BUCKET) -> None: + self.bucket = bucket + self.records: list[CallRecord] = [] + self.attempts: dict[str, int] = defaultdict(int) + self.phase_wall: dict[str, float] = defaultdict(float) + self._phase = _UNSET_PHASE + self._lock = threading.Lock() + self._installed = False + + # -- phase control -------------------------------------------------------- + + @property + def current_phase(self) -> str: + return self._phase + + @contextmanager + def phase(self, name: str) -> Generator[None, None, None]: + """Attribute every AWS call made inside the block to ``name``.""" + previous = self._phase + self._phase = name + started = time.perf_counter() + try: + yield + finally: + elapsed = time.perf_counter() - started + with self._lock: + self.phase_wall[name] += elapsed + self._phase = previous + + # -- botocore handlers ---------------------------------------------------- + + def _on_before_call(self, params: Any = None, context: Any = None, **_: Any) -> None: + """Stash timing, phase, URL and request size while the request dict is in hand. + + The response object botocore hands to ``after-call`` carries neither the + URL nor the outbound body size, so both are captured here. + """ + if context is None: + return + context["_iceberg_meter_start"] = time.perf_counter() + context["_iceberg_meter_phase"] = self._phase + + if not isinstance(params, dict): + return + context["_iceberg_meter_url"] = params.get("url", "") + + headers = params.get("headers") + request_bytes = _header_int(headers, "Content-Length") or _header_int( + headers, "X-Amz-Decoded-Content-Length" + ) + if not request_bytes: + request_bytes = _body_size(params.get("body")) + context["_iceberg_meter_request_bytes"] = request_bytes + + def _on_before_send(self, **_: Any) -> None: + """Count HTTP attempts. ``attempts`` above ``calls`` means botocore retried.""" + phase = self._phase + with self._lock: + self.attempts[phase] += 1 + + def _on_after_call( + self, + http_response: Any = None, + model: Any = None, + context: Any = None, + **_: Any, + ) -> None: + if model is None: + return + + ctx = context if isinstance(context, dict) else {} + started = ctx.get("_iceberg_meter_start") + duration_ms = (time.perf_counter() - started) * 1000 if started else 0.0 + phase = ctx.get("_iceberg_meter_phase") or self._phase + + service = getattr(getattr(model, "service_model", None), "endpoint_prefix", "") or "" + operation = getattr(model, "name", "") or "" + + key = _extract_key(ctx.get("_iceberg_meter_url", ""), self.bucket) + request_bytes = ctx.get("_iceberg_meter_request_bytes", 0) + + response_bytes = 0 + status = 0 + if http_response is not None: + status = getattr(http_response, "status_code", 0) or 0 + response_bytes = _header_int(getattr(http_response, "headers", None), "Content-Length") + + record = CallRecord( + phase=phase, + service=service, + operation=operation, + key=key, + key_class=classify_key(key) if service == "s3" else KEY_CLASS_OTHER, + tier=cost_tier(service, operation), + request_bytes=request_bytes, + response_bytes=response_bytes, + duration_ms=duration_ms, + status=status, + ) + with self._lock: + self.records.append(record) + + # -- installation --------------------------------------------------------- + + def install(self) -> None: + """Patch botocore so all current and future sessions report to this meter. + + Idempotent. Must run before any boto3/aiobotocore client is constructed, + since sessions built earlier will not carry the handlers. + """ + if self._installed: + return + self._installed = True + + original_init = getattr( + botocore.session.Session, + "_iceberg_meter_original_init", + botocore.session.Session.__init__, + ) + + meter = self + + def patched_init(session_self: Any, *args: Any, **kwargs: Any) -> None: + original_init(session_self, *args, **kwargs) + meter._register_on(session_self) + + botocore.session.Session._iceberg_meter_original_init = original_init # type: ignore[attr-defined] + botocore.session.Session.__init__ = patched_init # type: ignore[method-assign] + + def _register_on(self, session: Any) -> None: + for service in ("s3", "glue"): + session.register( + f"before-call.{service}.*", + self._on_before_call, + unique_id=f"iceberg-meter-before-{service}", + ) + session.register( + f"after-call.{service}.*", + self._on_after_call, + unique_id=f"iceberg-meter-after-{service}", + ) + session.register( + f"before-send.{service}.*", + self._on_before_send, + unique_id=f"iceberg-meter-send-{service}", + ) + + # -- reporting ------------------------------------------------------------ + + def summarize(self) -> dict[str, PhaseStats]: + """Fold the raw call log into per-phase aggregates.""" + with self._lock: + records = list(self.records) + attempts = dict(self.attempts) + wall = dict(self.phase_wall) + + stats: dict[str, PhaseStats] = {} + for record in records: + entry = stats.setdefault(record.phase, PhaseStats(phase=record.phase)) + entry.calls += 1 + entry.request_bytes += record.request_bytes + entry.response_bytes += record.response_bytes + entry.by_tier[record.tier] += 1 + entry.by_operation[f"{record.service}:{record.operation}"] += 1 + entry.by_key_class[record.key_class] += 1 + per_class = entry.by_key_class_operation.setdefault(record.key_class, defaultdict(int)) + per_class[record.operation] += 1 + entry.latencies_ms.append(record.duration_ms) + + for phase, seconds in wall.items(): + stats.setdefault(phase, PhaseStats(phase=phase)).wall_seconds = seconds + for phase, count in attempts.items(): + stats.setdefault(phase, PhaseStats(phase=phase)).attempts = count + + return stats + + def reset(self) -> None: + with self._lock: + self.records.clear() + self.attempts.clear() + self.phase_wall.clear() + + +# Module-level singleton -- the experiment runner installs this once at startup. +METER = Meter() diff --git a/experimentation/iceberg/schemas.py b/experimentation/iceberg/schemas.py new file mode 100644 index 00000000..5013b4c2 --- /dev/null +++ b/experimentation/iceberg/schemas.py @@ -0,0 +1,201 @@ +"""Iceberg schemas for the four Bluesky record types, plus Jetstream parsing. + +Each record type gets its own table, all sharing a common header (ids 1-10) and +adding type-specific columns from id 11. Every table is partitioned by +``days(created_at)``. + +**Field ids must be contiguous from 1.** ``Catalog.create_table`` renumbers a +schema's fields sequentially, and Iceberg resolves columns by id, not by name. +Declaring a gap (say, jumping to id 20) means the table is created with the +renumbered ids while writes built from this module's schema still stamp the +declared ids into the Parquet footers -- so reads look up ids that are not +there and silently return NULL for every affected column. ``test_schemas.py`` +enforces contiguity, and ``iceberg_writer`` builds its Arrow tables from the +*table's* schema rather than this one, so the two cannot drift. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +from pyiceberg.partitioning import PartitionField, PartitionSpec +from pyiceberg.schema import Schema +from pyiceberg.transforms import DayTransform +from pyiceberg.types import ( + BooleanType, + IntegerType, + ListType, + NestedField, + StringType, + TimestamptzType, +) + +from experimentation.iceberg import constants + +# Field id of `created_at` -- the partition source for every table. +CREATED_AT_FIELD_ID = 8 + +_COMMON_FIELDS = [ + NestedField(1, "uri", StringType(), required=True), + NestedField(2, "did", StringType(), required=True), + NestedField(3, "collection", StringType(), required=True), + NestedField(4, "rkey", StringType(), required=True), + NestedField(5, "cid", StringType(), required=False), + NestedField(6, "rev", StringType(), required=False), + NestedField(7, "operation", StringType(), required=True), + NestedField(CREATED_AT_FIELD_ID, "created_at", TimestamptzType(), required=True), + NestedField(9, "ingested_at", TimestamptzType(), required=True), + # True when `createdAt` was unusable and we substituted the ingest timestamp. + NestedField(10, "created_at_fallback", BooleanType(), required=True), +] + +SCHEMAS: dict[str, Schema] = { + "posts": Schema( + *_COMMON_FIELDS, + NestedField(11, "text", StringType(), required=False), + NestedField( + 12, + "langs", + # Nested ids are assigned after every top-level field, so the list + # element follows text_length (16), not langs (12). + ListType(element_id=17, element_type=StringType(), element_required=False), + required=False, + ), + NestedField(13, "reply_root_uri", StringType(), required=False), + NestedField(14, "reply_parent_uri", StringType(), required=False), + NestedField(15, "embed_type", StringType(), required=False), + NestedField(16, "text_length", IntegerType(), required=False), + ), + "likes": Schema( + *_COMMON_FIELDS, + NestedField(11, "subject_uri", StringType(), required=False), + NestedField(12, "subject_cid", StringType(), required=False), + ), + "reposts": Schema( + *_COMMON_FIELDS, + NestedField(11, "subject_uri", StringType(), required=False), + NestedField(12, "subject_cid", StringType(), required=False), + ), + "follows": Schema( + *_COMMON_FIELDS, + NestedField(11, "subject_did", StringType(), required=False), + ), +} + +PARTITION_SPEC = PartitionSpec( + PartitionField( + source_id=CREATED_AT_FIELD_ID, + field_id=1000, + transform=DayTransform(), + name="created_at_day", + ) +) + + +def _parse_created_at(raw: Any, ingested_at: datetime) -> tuple[datetime, bool]: + """Resolve a record's ``createdAt``, falling back to ingest time when unusable. + + Bluesky ``createdAt`` is client-supplied. Malformed or wildly skewed values + would each open a junk daily partition, so anything more than + ``MAX_CREATED_AT_SKEW_SECONDS`` from the broker timestamp is rejected. + + Returns the timestamp and whether the fallback was used. + """ + if not isinstance(raw, str) or not raw: + return ingested_at, True + try: + parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return ingested_at, True + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=UTC) + parsed = parsed.astimezone(UTC) + if abs((parsed - ingested_at).total_seconds()) > constants.MAX_CREATED_AT_SKEW_SECONDS: + return ingested_at, True + return parsed, False + + +def _as_dict(value: Any) -> dict[str, Any]: + """Return ``value`` when it is a dict, otherwise an empty dict.""" + return value if isinstance(value, dict) else {} + + +def _subject_uri(record: dict[str, Any]) -> tuple[str | None, str | None]: + """Extract (uri, cid) from a like/repost subject, which is a strongref.""" + subject = _as_dict(record.get("subject")) + return subject.get("uri"), subject.get("cid") + + +def _embed_type(record: dict[str, Any]) -> str | None: + return _as_dict(record.get("embed")).get("$type") + + +def parse_event(event: dict[str, Any]) -> tuple[str, dict[str, Any]] | None: + """Turn a raw Jetstream message into ``(record_type, row)``. + + Returns ``None`` for anything that isn't a commit on a collection we track -- + identity and account events, and collections outside ``COLLECTIONS``. + """ + if event.get("kind") != "commit": + return None + commit = event.get("commit") + if not isinstance(commit, dict): + return None + + collection = commit.get("collection", "") + record_type = constants.COLLECTIONS.get(collection) + if record_type is None: + return None + + did = event.get("did", "") + rkey = commit.get("rkey", "") + if not did or not rkey: + return None + + time_us = event.get("time_us") or 0 + ingested_at = datetime.fromtimestamp(time_us / 1_000_000, tz=UTC) + + # Deletes carry no record body; they still matter for dedup realism. + record = _as_dict(commit.get("record")) + + created_at, fallback = _parse_created_at(record.get("createdAt"), ingested_at) + + row: dict[str, Any] = { + "uri": f"at://{did}/{collection}/{rkey}", + "did": did, + "collection": collection, + "rkey": rkey, + "cid": commit.get("cid"), + "rev": commit.get("rev"), + "operation": commit.get("operation", "unknown"), + "created_at": created_at, + "ingested_at": ingested_at, + "created_at_fallback": fallback, + } + + if record_type == "posts": + text = record.get("text") + langs = record.get("langs") + reply = _as_dict(record.get("reply")) + root = _as_dict(reply.get("root")) + parent = _as_dict(reply.get("parent")) + row.update( + { + "text": text, + "langs": langs if isinstance(langs, list) else None, + "reply_root_uri": root.get("uri"), + "reply_parent_uri": parent.get("uri"), + "embed_type": _embed_type(record), + "text_length": len(text) if isinstance(text, str) else None, + } + ) + elif record_type in ("likes", "reposts"): + uri, cid = _subject_uri(record) + row.update({"subject_uri": uri, "subject_cid": cid}) + elif record_type == "follows": + row["subject_did"] = ( + record.get("subject") if isinstance(record.get("subject"), str) else None + ) + + return record_type, row diff --git a/experimentation/iceberg/tests/conftest.py b/experimentation/iceberg/tests/conftest.py new file mode 100644 index 00000000..0fd50a88 --- /dev/null +++ b/experimentation/iceberg/tests/conftest.py @@ -0,0 +1,15 @@ +"""Skip this suite when collected by the root interpreter. + +These tests need pyiceberg, which cannot be installed into the root venv (it +pins rich<15, the project requires rich>=15). Run them with the experiment's own +interpreter: + + experimentation/iceberg/.venv/bin/python -m pytest experimentation/iceberg/tests +""" + +collect_ignore_glob: list[str] = [] + +try: + import pyiceberg # noqa: F401 +except ImportError: # pragma: no cover - depends on which interpreter collects + collect_ignore_glob = ["test_*.py"] diff --git a/experimentation/iceberg/tests/test_replay.py b/experimentation/iceberg/tests/test_replay.py new file mode 100644 index 00000000..3a63ed70 --- /dev/null +++ b/experimentation/iceberg/tests/test_replay.py @@ -0,0 +1,257 @@ +"""Tests for flush-window batching and the dedup rule used during compaction.""" + +from __future__ import annotations + +import gzip +import json +from datetime import UTC, datetime + +import pyarrow as pa + +from experimentation.iceberg import replay, schemas +from experimentation.iceberg.maintenance import _collapse_to_latest, _drop_tombstones + +BASE_US = 1_784_721_600_000_000 # 2026-07-22T12:00:00Z + + +def _event(offset_seconds: float, collection: str = "app.bsky.feed.post", rkey: str = "r1") -> dict: + return { + "did": "did:plc:alice", + "time_us": BASE_US + int(offset_seconds * 1_000_000), + "kind": "commit", + "commit": { + "rev": "rev1", + "operation": "create", + "collection": collection, + "rkey": rkey, + "cid": "bafy", + "record": {"createdAt": "2026-07-22T12:00:00Z", "text": "x"}, + }, + } + + +class TestBatching: + def test_events_inside_one_window_form_a_single_batch(self): + events = [_event(i, rkey=f"r{i}") for i in range(10)] + batches = list(replay.iter_batches(events, flush_seconds=60)) + assert len(batches) == 1 + assert len(batches[0][1]["posts"]) == 10 + + def test_window_closes_at_the_flush_boundary(self): + # 0..59s in window 0, 60..119s in window 1, 120s in window 2. + events = [_event(t, rkey=f"r{t}") for t in (0, 30, 59, 60, 90, 120)] + batches = list(replay.iter_batches(events, flush_seconds=60)) + assert [len(rows["posts"]) for _, rows in batches] == [3, 2, 1] + + def test_batch_indices_are_sequential(self): + events = [_event(t, rkey=f"r{t}") for t in (0, 60, 120, 180)] + assert [index for index, _ in replay.iter_batches(events, flush_seconds=60)] == [0, 1, 2, 3] + + def test_record_types_are_split_within_a_batch(self): + events = [ + _event(0, "app.bsky.feed.post", "r1"), + _event(1, "app.bsky.feed.like", "r2"), + _event(2, "app.bsky.graph.follow", "r3"), + _event(3, "app.bsky.feed.repost", "r4"), + ] + _, buffers = next(iter(replay.iter_batches(events, flush_seconds=60))) + assert set(buffers) == {"posts", "likes", "follows", "reposts"} + assert all(len(rows) == 1 for rows in buffers.values()) + + def test_empty_record_types_are_omitted(self): + _, buffers = next(iter(replay.iter_batches([_event(0)], flush_seconds=60))) + assert set(buffers) == {"posts"} + + def test_trailing_partial_window_is_emitted(self): + events = [_event(t, rkey=f"r{t}") for t in (0, 60, 61)] + batches = list(replay.iter_batches(events, flush_seconds=60)) + assert len(batches) == 2 + assert len(batches[1][1]["posts"]) == 2 + + def test_no_events_yields_no_batches(self): + assert list(replay.iter_batches([], flush_seconds=60)) == [] + + def test_untracked_events_do_not_open_a_window(self): + events = [ + {"kind": "identity", "did": "did:plc:alice", "time_us": BASE_US}, + _event(0), + ] + batches = list(replay.iter_batches(events, flush_seconds=60)) + assert len(batches) == 1 + assert len(batches[0][1]["posts"]) == 1 + + def test_batching_is_deterministic_across_runs(self): + """The whole point of replaying: identical input must give identical commits.""" + events = [_event(t * 7, rkey=f"r{t}") for t in range(40)] + first = [ + (i, {k: len(v) for k, v in b.items()}) for i, b in replay.iter_batches(list(events), 60) + ] + second = [ + (i, {k: len(v) for k, v in b.items()}) for i, b in replay.iter_batches(list(events), 60) + ] + assert first == second + + +class TestIterEvents: + def test_reads_gzipped_jsonl(self, tmp_path): + path = tmp_path / "capture.jsonl.gz" + with gzip.open(path, "wt", encoding="utf-8") as handle: + for i in range(3): + handle.write(json.dumps(_event(i, rkey=f"r{i}")) + "\n") + assert len(list(replay.iter_events(path))) == 3 + + def test_malformed_and_blank_lines_are_skipped(self, tmp_path): + path = tmp_path / "capture.jsonl.gz" + with gzip.open(path, "wt", encoding="utf-8") as handle: + handle.write(json.dumps(_event(0)) + "\n") + handle.write("{not json\n") + handle.write("\n") + handle.write(json.dumps(_event(1, rkey="r2")) + "\n") + assert len(list(replay.iter_events(path))) == 2 + + +def _arrow_rows(rows: list[dict]) -> pa.Table: + schema = schemas.SCHEMAS["posts"].as_arrow() + return pa.Table.from_pydict( + {f.name: [r.get(f.name) for r in rows] for f in schema}, schema=schema + ) + + +def _row( + uri: str, + ingested_offset: int, + text: str, + operation: str = "create", + cid: str | None = "c", +) -> dict: + ingested = datetime.fromtimestamp(BASE_US / 1e6 + ingested_offset, tz=UTC) + return { + "uri": uri, + "did": "did:plc:alice", + "collection": "app.bsky.feed.post", + "rkey": uri.rsplit("/", 1)[-1], + "cid": cid, + "rev": "r", + "operation": operation, + "created_at": ingested, + "ingested_at": ingested, + "created_at_fallback": False, + "text": text, + "langs": None, + "reply_root_uri": None, + "reply_parent_uri": None, + "embed_type": None, + "text_length": len(text), + } + + +class TestCollapseToLatest: + def test_same_uri_and_cid_is_a_redelivered_duplicate(self): + """Identical event twice -- the only thing that is really a duplicate.""" + table = _arrow_rows( + [_row("at://a/1", 0, "same", cid="c1"), _row("at://a/1", 10, "same", cid="c1")] + ) + collapsed, stats = _collapse_to_latest(table) + assert stats["redelivered_duplicates"] == 1 + assert stats["lifecycle_collapses"] == 0 + assert collapsed.num_rows == 1 + + def test_same_uri_different_cid_is_a_lifecycle_collapse(self): + """Create then edit -- two real events about one record, not a duplicate.""" + table = _arrow_rows( + [_row("at://a/1", 0, "v1", cid="c1"), _row("at://a/1", 10, "v2", cid="c2")] + ) + collapsed, stats = _collapse_to_latest(table) + assert stats["redelivered_duplicates"] == 0 + assert stats["lifecycle_collapses"] == 1 + assert collapsed.column("text").to_pylist() == ["v2"] + + def test_create_then_delete_counts_as_lifecycle_not_duplicate(self): + table = _arrow_rows( + [ + _row("at://a/1", 0, "hello", cid="c1"), + _row("at://a/1", 10, "", operation="delete", cid=None), + ] + ) + collapsed, stats = _collapse_to_latest(table) + assert stats["redelivered_duplicates"] == 0 + assert stats["lifecycle_collapses"] == 1 + assert collapsed.column("operation").to_pylist() == ["delete"] + + def test_distinct_uris_are_all_kept(self): + table = _arrow_rows([_row(f"at://a/{i}", i, "t") for i in range(5)]) + collapsed, stats = _collapse_to_latest(table) + assert stats == {"redelivered_duplicates": 0, "lifecycle_collapses": 0} + assert collapsed.num_rows == 5 + + def test_three_events_keep_only_the_newest(self): + table = _arrow_rows( + [ + _row("at://a/1", 0, "v1", cid="c1"), + _row("at://a/1", 5, "v2", cid="c2"), + _row("at://a/1", 9, "v3", cid="c3"), + ] + ) + collapsed, stats = _collapse_to_latest(table) + assert stats["lifecycle_collapses"] == 2 + assert collapsed.column("text").to_pylist() == ["v3"] + + def test_empty_table(self): + collapsed, stats = _collapse_to_latest(_arrow_rows([])) + assert stats == {"redelivered_duplicates": 0, "lifecycle_collapses": 0} + assert collapsed.num_rows == 0 + + def test_all_columns_survive(self): + table = _arrow_rows( + [_row("at://a/1", 0, "old", cid="c1"), _row("at://a/1", 10, "new", cid="c2")] + ) + collapsed, _ = _collapse_to_latest(table) + assert collapsed.schema == table.schema + + +class TestDropTombstones: + def test_delete_rows_are_removed(self): + table = _arrow_rows( + [ + _row("at://a/1", 0, "kept"), + _row("at://a/2", 1, "", operation="delete", cid=None), + _row("at://a/3", 2, "kept too"), + ] + ) + kept, dropped = _drop_tombstones(table) + assert dropped == 1 + assert kept.column("uri").to_pylist() == ["at://a/1", "at://a/3"] + + def test_updates_are_not_tombstones(self): + table = _arrow_rows([_row("at://a/1", 0, "edited", operation="update")]) + kept, dropped = _drop_tombstones(table) + assert dropped == 0 + assert kept.num_rows == 1 + + def test_all_deletes_yields_empty_table(self): + table = _arrow_rows( + [_row(f"at://a/{i}", i, "", operation="delete", cid=None) for i in range(3)] + ) + kept, dropped = _drop_tombstones(table) + assert dropped == 3 + assert kept.num_rows == 0 + assert kept.schema == table.schema + + def test_empty_table(self): + kept, dropped = _drop_tombstones(_arrow_rows([])) + assert dropped == 0 + assert kept.num_rows == 0 + + def test_create_delete_pair_vanishes_end_to_end(self): + """The full compaction rule: collapse to latest, then drop the tombstone.""" + table = _arrow_rows( + [ + _row("at://a/1", 0, "posted", cid="c1"), + _row("at://a/1", 10, "", operation="delete", cid=None), + _row("at://a/2", 5, "survivor", cid="c3"), + ] + ) + collapsed, _ = _collapse_to_latest(table) + final, dropped = _drop_tombstones(collapsed) + assert dropped == 1 + assert final.column("uri").to_pylist() == ["at://a/2"] diff --git a/experimentation/iceberg/tests/test_s3_meter.py b/experimentation/iceberg/tests/test_s3_meter.py new file mode 100644 index 00000000..19038075 --- /dev/null +++ b/experimentation/iceberg/tests/test_s3_meter.py @@ -0,0 +1,235 @@ +"""Tests for the metering layer -- key classification, billing tiers, attribution.""" + +from __future__ import annotations + +import io +import threading + +from experimentation.iceberg import constants +from experimentation.iceberg.s3_meter import ( + KEY_CLASS_DATA, + KEY_CLASS_MANIFEST, + KEY_CLASS_MANIFEST_LIST, + KEY_CLASS_METADATA_JSON, + KEY_CLASS_OTHER, + KEY_CLASS_RAW, + Meter, + _body_size, + _extract_key, + classify_key, + cost_tier, +) + +BUCKET = "lab-data-integrations-interface" + + +class TestClassifyKey: + def test_metadata_json(self): + key = "experiments/iceberg/run1/warehouse/posts/metadata/00003-abc.metadata.json" + assert classify_key(key) == KEY_CLASS_METADATA_JSON + + def test_manifest_list_beats_generic_avro(self): + key = "experiments/iceberg/run1/warehouse/posts/metadata/snap-123-1-abc.avro" + assert classify_key(key) == KEY_CLASS_MANIFEST_LIST + + def test_manifest(self): + key = "experiments/iceberg/run1/warehouse/posts/metadata/abc-m0.avro" + assert classify_key(key) == KEY_CLASS_MANIFEST + + def test_data_file(self): + key = "experiments/iceberg/run1/warehouse/posts/data/created_at_day=2026-07-22/x.parquet" + assert classify_key(key) == KEY_CLASS_DATA + + def test_raw_baseline_is_not_counted_as_iceberg_data(self): + key = "experiments/iceberg/run1/raw/posts/created_at_day=2026-07-22/batch-00001.parquet" + assert classify_key(key) == KEY_CLASS_RAW + + def test_empty_key(self): + assert classify_key("") == KEY_CLASS_OTHER + + +class TestCostTier: + def test_put_operations(self): + assert cost_tier("s3", "PutObject") == "put" + assert cost_tier("s3", "ListObjectsV2") == "put" + assert cost_tier("s3", "CompleteMultipartUpload") == "put" + + def test_get_is_the_fallback(self): + assert cost_tier("s3", "GetObject") == "get" + assert cost_tier("s3", "HeadObject") == "get" + + def test_delete_is_free(self): + assert cost_tier("s3", "DeleteObject") == "delete" + assert cost_tier("s3", "DeleteObjects") == "delete" + + def test_glue_is_its_own_tier(self): + assert cost_tier("glue", "UpdateTable") == "glue" + + +class TestExtractKey: + def test_virtual_hosted_style(self): + url = f"https://{BUCKET}.s3.us-east-2.amazonaws.com/experiments/iceberg/a/b.parquet" + assert _extract_key(url, BUCKET) == "experiments/iceberg/a/b.parquet" + + def test_path_style_strips_bucket(self): + url = f"https://s3.us-east-2.amazonaws.com/{BUCKET}/experiments/iceberg/a/b.parquet" + assert _extract_key(url, BUCKET) == "experiments/iceberg/a/b.parquet" + + def test_bucket_root(self): + assert _extract_key(f"https://s3.us-east-2.amazonaws.com/{BUCKET}", BUCKET) == "" + + def test_percent_encoded_key_is_decoded(self): + url = f"https://{BUCKET}.s3.us-east-2.amazonaws.com/a/created_at_day%3D2026-07-22/x.parquet" + assert _extract_key(url, BUCKET) == "a/created_at_day=2026-07-22/x.parquet" + + +class TestBodySize: + """botocore hands us a BytesIO with no Content-Length; sizing must still work.""" + + def test_bytes_body(self): + assert _body_size(b"A" * 4096) == 4096 + + def test_str_body_is_measured_in_utf8_bytes(self): + assert _body_size("é" * 10) == 20 + + def test_bytesio_body(self): + assert _body_size(io.BytesIO(b"A" * 4096)) == 4096 + + def test_bytesio_position_is_restored(self): + """Consuming the stream here would corrupt the upload.""" + body = io.BytesIO(b"A" * 4096) + body.seek(100) + assert _body_size(body) == 3996 + assert body.tell() == 100 + + def test_none_and_unmeasurable_bodies(self): + assert _body_size(None) == 0 + assert _body_size(object()) == 0 + + def test_chunked_upload_falls_back_to_decoded_length_header(self): + """Large uploads switch to aws-chunked and drop Content-Length entirely.""" + meter = Meter(bucket=BUCKET) + context: dict = {} + params = { + "url": f"https://{BUCKET}.s3.us-east-2.amazonaws.com/w/posts/data/x.parquet", + "headers": { + "X-Amz-Decoded-Content-Length": b"8388608", + "Transfer-Encoding": b"chunked", + }, + "body": None, + } + meter._on_before_call(params=params, context=context) + assert context["_iceberg_meter_request_bytes"] == 8_388_608 + + +def _record_call(meter: Meter, operation: str, key: str, service: str = "s3") -> None: + """Drive the meter's handlers the way botocore would.""" + context: dict = {} + model = type( + "Model", + (), + {"name": operation, "service_model": type("SM", (), {"endpoint_prefix": service})()}, + )() + params = { + "url": f"https://{BUCKET}.s3.us-east-2.amazonaws.com/{key}", + "headers": {"Content-Length": "100"}, + } + meter._on_before_call(params=params, context=context) + response = type("Resp", (), {"status_code": 200, "headers": {"Content-Length": "250"}})() + meter._on_after_call(http_response=response, model=model, context=context) + + +class TestPhaseAttribution: + def test_calls_land_in_the_active_phase(self): + meter = Meter(bucket=BUCKET) + with meter.phase(constants.PHASE_ICEBERG_APPEND): + _record_call(meter, "PutObject", "w/posts/data/x.parquet") + with meter.phase(constants.PHASE_COMPACT_DEDUP): + _record_call(meter, "GetObject", "w/posts/data/x.parquet") + + stats = meter.summarize() + assert stats[constants.PHASE_ICEBERG_APPEND].by_tier["put"] == 1 + assert stats[constants.PHASE_COMPACT_DEDUP].by_tier["get"] == 1 + + def test_nested_phases_restore_the_outer_phase(self): + meter = Meter(bucket=BUCKET) + with meter.phase("outer"): + with meter.phase("inner"): + pass + assert meter.current_phase == "outer" + assert meter.current_phase == "unattributed" + + def test_bytes_and_key_class_are_recorded(self): + meter = Meter(bucket=BUCKET) + with meter.phase(constants.PHASE_ICEBERG_APPEND): + _record_call(meter, "PutObject", "w/posts/metadata/00001-a.metadata.json") + + stats = meter.summarize()[constants.PHASE_ICEBERG_APPEND] + assert stats.request_bytes == 100 + assert stats.response_bytes == 250 + assert stats.by_key_class[KEY_CLASS_METADATA_JSON] == 1 + + def test_concurrent_calls_are_all_counted(self): + """PyIceberg writes manifests from a thread pool; the counters must hold.""" + meter = Meter(bucket=BUCKET) + with meter.phase(constants.PHASE_ICEBERG_APPEND): + threads = [ + threading.Thread( + target=_record_call, args=(meter, "PutObject", f"w/posts/data/{i}.parquet") + ) + for i in range(50) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert meter.summarize()[constants.PHASE_ICEBERG_APPEND].calls == 50 + + def test_retries_show_up_as_attempts_above_calls(self): + meter = Meter(bucket=BUCKET) + with meter.phase(constants.PHASE_ICEBERG_APPEND): + _record_call(meter, "PutObject", "w/posts/data/x.parquet") + meter._on_before_send() + meter._on_before_send() + + stats = meter.summarize()[constants.PHASE_ICEBERG_APPEND] + assert stats.calls == 1 + assert stats.attempts == 2 + + +class TestCostModel: + def test_cost_uses_the_right_rate_per_tier(self): + meter = Meter(bucket=BUCKET) + with meter.phase(constants.PHASE_ICEBERG_APPEND): + for i in range(1000): + _record_call(meter, "PutObject", f"w/posts/data/{i}.parquet") + for i in range(1000): + _record_call(meter, "GetObject", f"w/posts/data/{i}.parquet") + _record_call(meter, "UpdateTable", "", service="glue") + + stats = meter.summarize()[constants.PHASE_ICEBERG_APPEND] + expected = 0.005 + 0.0004 + constants.COST_PER_GLUE_REQUEST + assert abs(stats.cost_usd - expected) < 1e-9 + + def test_deletes_are_free(self): + meter = Meter(bucket=BUCKET) + with meter.phase(constants.PHASE_EXPIRE_METADATA): + for i in range(500): + _record_call(meter, "DeleteObject", f"w/posts/data/{i}.parquet") + + assert meter.summarize()[constants.PHASE_EXPIRE_METADATA].cost_usd == 0.0 + + def test_percentiles(self): + meter = Meter(bucket=BUCKET) + with meter.phase("p"): + _record_call(meter, "PutObject", "a.parquet") + stats = meter.summarize()["p"] + stats.latencies_ms = [float(i) for i in range(1, 101)] + assert stats.percentile(50) == 50.0 + assert stats.percentile(95) == 95.0 + + def test_percentile_of_empty_is_zero(self): + from experimentation.iceberg.s3_meter import PhaseStats + + assert PhaseStats(phase="x").percentile(95) == 0.0 diff --git a/experimentation/iceberg/tests/test_schemas.py b/experimentation/iceberg/tests/test_schemas.py new file mode 100644 index 00000000..c7d379c5 --- /dev/null +++ b/experimentation/iceberg/tests/test_schemas.py @@ -0,0 +1,248 @@ +"""Tests for Jetstream event parsing and the createdAt fallback rule.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +from experimentation.iceberg import constants, schemas + +# 2026-07-22T12:00:00Z +INGEST_US = 1_784_721_600_000_000 +INGEST_DT = datetime.fromtimestamp(INGEST_US / 1_000_000, tz=UTC) + + +def _event(collection: str, record: dict | None, operation: str = "create") -> dict: + commit = { + "rev": "3l3qo2vutsw2b", + "operation": operation, + "collection": collection, + "rkey": "3l3qo2vuowo2b", + "cid": "bafyreiabc", + } + if record is not None: + commit["record"] = record + return {"did": "did:plc:alice", "time_us": INGEST_US, "kind": "commit", "commit": commit} + + +class TestRouting: + def test_post(self): + parsed = schemas.parse_event( + _event("app.bsky.feed.post", {"text": "hello", "createdAt": "2026-07-22T12:00:00.000Z"}) + ) + assert parsed is not None + record_type, row = parsed + assert record_type == "posts" + assert row["text"] == "hello" + assert row["text_length"] == 5 + assert row["uri"] == "at://did:plc:alice/app.bsky.feed.post/3l3qo2vuowo2b" + + def test_like_extracts_subject_strongref(self): + record = { + "createdAt": "2026-07-22T12:00:00Z", + "subject": {"uri": "at://did:plc:bob/app.bsky.feed.post/xyz", "cid": "bafy1"}, + } + parsed = schemas.parse_event(_event("app.bsky.feed.like", record)) + assert parsed is not None + record_type, row = parsed + assert record_type == "likes" + assert row["subject_uri"] == "at://did:plc:bob/app.bsky.feed.post/xyz" + assert row["subject_cid"] == "bafy1" + + def test_follow_subject_is_a_plain_did(self): + parsed = schemas.parse_event( + _event( + "app.bsky.graph.follow", + {"createdAt": "2026-07-22T12:00:00Z", "subject": "did:plc:bob"}, + ) + ) + assert parsed is not None + record_type, row = parsed + assert record_type == "follows" + assert row["subject_did"] == "did:plc:bob" + + def test_repost(self): + record = {"createdAt": "2026-07-22T12:00:00Z", "subject": {"uri": "at://x", "cid": "c"}} + parsed = schemas.parse_event(_event("app.bsky.feed.repost", record)) + assert parsed is not None + assert parsed[0] == "reposts" + + def test_untracked_collection_is_dropped(self): + assert ( + schemas.parse_event( + _event("app.bsky.actor.profile", {"createdAt": "2026-07-22T12:00:00Z"}) + ) + is None + ) + + def test_non_commit_kinds_are_dropped(self): + assert schemas.parse_event({"kind": "identity", "did": "did:plc:alice"}) is None + + def test_missing_commit_is_dropped(self): + assert schemas.parse_event({"kind": "commit", "did": "did:plc:alice"}) is None + + def test_missing_rkey_is_dropped(self): + event = _event("app.bsky.feed.post", {"createdAt": "2026-07-22T12:00:00Z"}) + event["commit"]["rkey"] = "" + assert schemas.parse_event(event) is None + + +class TestCreatedAtFallback: + def test_valid_timestamp_is_kept(self): + parsed = schemas.parse_event( + _event("app.bsky.feed.post", {"createdAt": "2026-07-22T11:59:00Z"}) + ) + assert parsed is not None + _, row = parsed + assert row["created_at_fallback"] is False + assert row["created_at"] == datetime(2026, 7, 22, 11, 59, tzinfo=UTC) + + def test_malformed_timestamp_falls_back(self): + parsed = schemas.parse_event(_event("app.bsky.feed.post", {"createdAt": "not-a-date"})) + assert parsed is not None + _, row = parsed + assert row["created_at_fallback"] is True + assert row["created_at"] == INGEST_DT + + def test_far_future_timestamp_falls_back(self): + """A year-2100 createdAt would otherwise open a junk daily partition.""" + parsed = schemas.parse_event( + _event("app.bsky.feed.post", {"createdAt": "2100-01-01T00:00:00Z"}) + ) + assert parsed is not None + _, row = parsed + assert row["created_at_fallback"] is True + assert row["created_at"] == INGEST_DT + + def test_epoch_zero_falls_back(self): + parsed = schemas.parse_event( + _event("app.bsky.feed.post", {"createdAt": "1970-01-01T00:00:00Z"}) + ) + assert parsed is not None + assert parsed[1]["created_at_fallback"] is True + + def test_skew_inside_the_window_is_kept(self): + """Genuinely backdated-but-plausible records must not be rewritten.""" + parsed = schemas.parse_event( + _event("app.bsky.feed.post", {"createdAt": "2026-07-22T00:00:01Z"}) + ) + assert parsed is not None + assert parsed[1]["created_at_fallback"] is False + + def test_missing_created_at_falls_back(self): + parsed = schemas.parse_event(_event("app.bsky.feed.post", {"text": "no timestamp"})) + assert parsed is not None + assert parsed[1]["created_at_fallback"] is True + + def test_delete_has_no_record_body(self): + parsed = schemas.parse_event(_event("app.bsky.feed.post", None, operation="delete")) + assert parsed is not None + _, row = parsed + assert row["operation"] == "delete" + assert row["created_at_fallback"] is True + assert row["text"] is None + + def test_naive_timestamp_is_treated_as_utc(self): + parsed = schemas.parse_event( + _event("app.bsky.feed.post", {"createdAt": "2026-07-22T11:59:00"}) + ) + assert parsed is not None + assert parsed[1]["created_at_fallback"] is False + + +class TestNestedFields: + def test_reply_refs(self): + record = { + "createdAt": "2026-07-22T12:00:00Z", + "text": "reply", + "reply": {"root": {"uri": "at://root"}, "parent": {"uri": "at://parent"}}, + } + parsed = schemas.parse_event(_event("app.bsky.feed.post", record)) + assert parsed is not None + _, row = parsed + assert row["reply_root_uri"] == "at://root" + assert row["reply_parent_uri"] == "at://parent" + + def test_malformed_nested_values_do_not_raise(self): + record = { + "createdAt": "2026-07-22T12:00:00Z", + "reply": "garbage", + "embed": 42, + "langs": "en", + } + parsed = schemas.parse_event(_event("app.bsky.feed.post", record)) + assert parsed is not None + _, row = parsed + assert row["reply_root_uri"] is None + assert row["embed_type"] is None + assert row["langs"] is None + + def test_embed_type(self): + record = {"createdAt": "2026-07-22T12:00:00Z", "embed": {"$type": "app.bsky.embed.images"}} + parsed = schemas.parse_event(_event("app.bsky.feed.post", record)) + assert parsed is not None + assert parsed[1]["embed_type"] == "app.bsky.embed.images" + + +class TestFieldIds: + """Guards against a silent-NULL bug that already happened once. + + ``Catalog.create_table`` renumbers schema fields contiguously from 1, and + Iceberg resolves columns by field id rather than name. Declaring + non-contiguous ids means writes stamp the declared ids into Parquet while + the table metadata records the renumbered ones -- reads then find nothing + and return NULL for every affected column, with no error anywhere. + """ + + def test_top_level_ids_are_contiguous_from_one(self): + for record_type, schema in schemas.SCHEMAS.items(): + ids = [field.field_id for field in schema.fields] + assert ids == list(range(1, len(ids) + 1)), ( + f"{record_type} declares {ids}; create_table would assign " + f"{list(range(1, len(ids) + 1))}, and the mismatch reads back as NULL" + ) + + def test_nested_element_ids_follow_every_top_level_field(self): + """Nested ids are assigned after all top-level fields, not next to their parent.""" + posts = schemas.SCHEMAS["posts"] + langs = posts.find_field("langs") + top_level_count = len(posts.fields) + assert langs.field_type.element_id == top_level_count + 1 + + def test_common_header_is_identical_across_tables(self): + header = [(f.field_id, f.name) for f in schemas.SCHEMAS["posts"].fields[:10]] + for record_type, schema in schemas.SCHEMAS.items(): + assert [(f.field_id, f.name) for f in schema.fields[:10]] == header, record_type + + +class TestSchemaShape: + def test_every_record_type_has_a_schema(self): + assert set(schemas.SCHEMAS) == set(constants.RECORD_TYPES) + + def test_partition_source_is_created_at(self): + for record_type, schema in schemas.SCHEMAS.items(): + field = schema.find_field(schemas.CREATED_AT_FIELD_ID) + assert field.name == "created_at", record_type + + def test_parsed_rows_only_use_declared_columns(self): + """A key not in the schema would be silently dropped at Arrow conversion.""" + samples = { + "app.bsky.feed.post": {"createdAt": "2026-07-22T12:00:00Z", "text": "x"}, + "app.bsky.feed.like": { + "createdAt": "2026-07-22T12:00:00Z", + "subject": {"uri": "u", "cid": "c"}, + }, + "app.bsky.feed.repost": { + "createdAt": "2026-07-22T12:00:00Z", + "subject": {"uri": "u", "cid": "c"}, + }, + "app.bsky.graph.follow": { + "createdAt": "2026-07-22T12:00:00Z", + "subject": "did:plc:bob", + }, + } + for collection, record in samples.items(): + parsed = schemas.parse_event(_event(collection, record)) + assert parsed is not None + record_type, row = parsed + declared = {field.name for field in schemas.SCHEMAS[record_type].fields} + assert set(row) <= declared, f"{record_type}: {set(row) - declared}" diff --git a/pyproject.toml b/pyproject.toml index 2841cf93..4ddffebb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,7 +83,10 @@ pythonVersion = "3.11" typeCheckingMode = "standard" venvPath = "." venv = ".venv" -exclude = ["tests", "experiments"] +# experimentation/iceberg resolves against its own venv (pyiceberg pins +# rich<15, the root project requires rich>=15), so the root interpreter cannot +# type-check it. +exclude = ["tests", "experiments", "experimentation/iceberg"] [tool.vulture] exclude = ["*venv*", "temp.py"]