From e0f228f41ab27cda27c21736ff09a28667b08276 Mon Sep 17 00:00:00 2001
From: Neukz <kacperneumann1@gmail.com>
Date: Mon, 15 Jun 2026 13:02:19 +0200
Subject: [PATCH] feat: add extraction-metadata column and entity partitioning
 to writer

---
 src/repolytics/ingestion/writer.py | 38 ++++++++++++++++++++++++------
 tests/unit/test_writer.py          | 24 +++++++++++++++++++
 2 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/src/repolytics/ingestion/writer.py b/src/repolytics/ingestion/writer.py
index d29514f..74bb2f4 100644
--- a/src/repolytics/ingestion/writer.py
+++ b/src/repolytics/ingestion/writer.py
@@ -1,6 +1,8 @@
 """Parquet writer for raw API responses.
 
 Lands each record as a single JSON `data` column plus a `_loaded_at` timestamp.
+Optional extraction metadata (e.g. the source repo) is stamped as extra
+`_`-prefixed string columns, kept separate from the pristine `data` blob.
 """
 
 import json
@@ -9,37 +11,59 @@
 
 import polars as pl
 
-# Raw-landing schema.
+# Raw-landing schema; metadata columns are appended per call.
 _SCHEMA = {"data": pl.Utf8, "_loaded_at": pl.Datetime(time_unit="us", time_zone="UTC")}
 
 
-def write_parquet(records: list[dict], path: str | Path) -> Path:
+def write_parquet(
+    records: list[dict],
+    path: str | Path,
+    *,
+    metadata: dict[str, str] | None = None,
+) -> Path:
     """Write `records` to `path` as Parquet (JSON `data` column + `_loaded_at`).
 
     Each record is serialized to a JSON string and stamped with a single batch
-    timestamp. An empty `records` list still writes a zero-row file with the
-    correct schema, so the partition stays present and schema-stable.
+    timestamp. `metadata` adds one constant string column per key to every row
+    (extraction provenance such as the source repo), distinct from the payload.
+    An empty `records` list still writes a zero-row file with the full schema,
+    so the partition stays present and schema-stable.
     """
+    metadata = metadata or {}
     out = Path(path)
     out.parent.mkdir(parents=True, exist_ok=True)
     loaded_at = datetime.now(UTC)
+    schema = {**_SCHEMA, **dict.fromkeys(metadata, pl.Utf8)}
     frame = pl.DataFrame(
         {
             "data": [json.dumps(record, ensure_ascii=False) for record in records],
             "_loaded_at": [loaded_at] * len(records),
+            **{key: [value] * len(records) for key, value in metadata.items()},
         },
-        schema=_SCHEMA,
+        schema=schema,
     )
     frame.write_parquet(out)
     return out
 
 
 def partition_path(
-    root: str | Path, source: str, table: str, date: datetime | str
+    root: str | Path,
+    source: str,
+    table: str,
+    date: datetime | str,
+    *,
+    entity: str | None = None,
 ) -> Path:
     """Build the date-partitioned landing path for a source/table on a date.
 
+    When `entity` is given (e.g. a repo or package), it is inserted before
+    the date and its `/` slugified to `__`, so each entity lands in its own
+    file without collisions.
+
     Returns `{root}/{source}/{table}/{YYYY-MM-DD}.parquet`.
     """
     day = date.strftime("%Y-%m-%d") if isinstance(date, datetime) else date
-    return Path(root) / source / table / f"{day}.parquet"
+    base = Path(root) / source / table
+    if entity is not None:
+        base = base / entity.replace("/", "__")
+    return base / f"{day}.parquet"
diff --git a/tests/unit/test_writer.py b/tests/unit/test_writer.py
index 47c75a3..4207e8d 100644
--- a/tests/unit/test_writer.py
+++ b/tests/unit/test_writer.py
@@ -46,6 +46,23 @@ def test_write_parquet_empty_writes_schema_only(tmp_path: Path) -> None:
     assert frame.height == 0
 
 
+def test_write_parquet_metadata_columns(tmp_path: Path) -> None:
+    records = [{"sha": "a"}, {"sha": "b"}]
+    out = write_parquet(
+        records, tmp_path / "commits.parquet", metadata={"_repo": "o/r"}
+    )
+    frame = pl.read_parquet(out)
+    assert frame.columns == ["data", "_loaded_at", "_repo"]
+    assert frame["_repo"].to_list() == ["o/r", "o/r"]
+
+
+def test_write_parquet_empty_keeps_metadata_columns(tmp_path: Path) -> None:
+    out = write_parquet([], tmp_path / "empty.parquet", metadata={"_repo": "o/r"})
+    frame = pl.read_parquet(out)
+    assert frame.columns == ["data", "_loaded_at", "_repo"]
+    assert frame.height == 0
+
+
 def test_write_parquet_creates_parent_dirs(tmp_path: Path) -> None:
     out = write_parquet([{"id": 1}], tmp_path / "a" / "b" / "c.parquet")
     assert out.exists()
@@ -58,3 +75,10 @@ def test_partition_path_formats_date() -> None:
 
     from_string = partition_path("data/raw", "github", "commits", "2024-01-02")
     assert from_string == Path("data/raw/github/commits/2024-01-02.parquet")
+
+
+def test_partition_path_with_entity_slugifies_slash() -> None:
+    path = partition_path(
+        "data/raw", "github", "commits", "2024-01-02", entity="encode/httpx"
+    )
+    assert path == Path("data/raw/github/commits/encode__httpx/2024-01-02.parquet")