Partial partition rewrite for in-place inventory modifications & treatments

services/api/api/resources/inventories/cache.py

@@ -1,84 +1,71 @@
 """
 api/v2/resources/inventories/cache.py
 
-Cached async parquet metadata access for inventory data streaming.
+Cached async access to an inventory's partitioned-Parquet metadata and data.
+
+Counts come from the dask DataFrame, which reads each partition's own footer (and
+ignores any aggregated ``_metadata``). This stays correct after standgen rewrites
+only the partitions a modification changed — an in-place overwrite leaves the
+``_metadata`` aggregate stale, but ``len(ddf)`` does not trust it.
 """
 
 import asyncio
 from dataclasses import dataclass
 
+import dask.dataframe as dd
 import pandas as pd
-import pyarrow.parquet as pq
 from ring import lru
 
 from lib.config import INVENTORIES_BUCKET
 
 
-@dataclass
-class PartitionMeta:
-    index: int
-    num_rows: int
-    path: str
-
-
 @dataclass
 class InventoryMeta:
     num_partitions: int
     total_rows: int
     columns: list[str]
-    partitions: list[PartitionMeta]
-
-
-def _read_metadata_sync(inventory_id: str) -> InventoryMeta:
-    metadata_path = f"gs://{INVENTORIES_BUCKET}/{inventory_id}/_metadata"
-    pf = pq.ParquetFile(metadata_path)
-    return _parse_metadata(pf.metadata)
 
 
-def _parse_metadata(metadata: pq.FileMetaData) -> InventoryMeta:
-    partitions_by_path: dict[str, int] = {}
-    for i in range(metadata.num_row_groups):
-        rg = metadata.row_group(i)
-        path = rg.column(0).file_path
-        partitions_by_path[path] = partitions_by_path.get(path, 0) + rg.num_rows
+def _inventory_path(inventory_id: str) -> str:
+    return f"gs://{INVENTORIES_BUCKET}/{inventory_id}"
 
-    columns = metadata.schema.to_arrow_schema().names
-
-    sorted_paths = sorted(partitions_by_path.keys())
-    partitions = [
-        PartitionMeta(index=idx, num_rows=partitions_by_path[path], path=path)
-        for idx, path in enumerate(sorted_paths)
-    ]
 
+def _read_metadata_sync(inventory_id: str) -> InventoryMeta:
+    ddf = dd.read_parquet(_inventory_path(inventory_id))
     return InventoryMeta(
-        num_partitions=len(partitions),
-        total_rows=sum(p.num_rows for p in partitions),
-        columns=columns,
-        partitions=partitions,
+        num_partitions=ddf.npartitions,
+        total_rows=len(ddf),  # footer-level row count; ignores stale _metadata
+        columns=list(ddf.columns),
     )
 
 
 @lru(maxsize=128, force_asyncio=True)
 async def get_inventory_metadata(
     inventory_id: str, checksum: str | None
 ) -> InventoryMeta:
-    """Read an inventory's partitioned-Parquet metadata, cached per content.
+    """Read an inventory's partition metadata, cached per content.
 
     ``checksum`` participates only in the cache key (the body ignores it). An
-    in-place modification (POST .../modifications) rewrites the data under the
-    same ``inventory_id`` and re-assigns the inventory's ``checksum``, so a new
-    checksum bypasses the stale entry. Without it, this LRU would keep serving
-    pre-modification partition paths and row counts. ``None`` (legacy
-    inventories without a checksum) is a valid, stable key — such inventories
-    predate in-place edits, and any in-place edit assigns a fresh checksum.
+    in-place modification rewrites the changed partitions under the same
+    ``inventory_id`` and re-assigns the inventory's ``checksum``, so a new
+    checksum bypasses the stale entry. ``None`` (legacy inventories without a
+    checksum) is a valid, stable key.
     """
     return await asyncio.to_thread(_read_metadata_sync, inventory_id)
 
 
+def _read_partition_sync(
+    inventory_id: str, partition_index: int, columns: list[str] | None
+) -> pd.DataFrame:
+    ddf = dd.read_parquet(_inventory_path(inventory_id), columns=columns)
+    return ddf.partitions[partition_index].compute().reset_index(drop=True)
+
+
 async def read_partition(
     inventory_id: str,
-    partition_path: str,
+    partition_index: int,
     columns: list[str] | None = None,
 ) -> pd.DataFrame:
-    path = f"gs://{INVENTORIES_BUCKET}/{inventory_id}/{partition_path}"
-    return await asyncio.to_thread(pd.read_parquet, path, columns=columns)
+    return await asyncio.to_thread(
+        _read_partition_sync, inventory_id, partition_index, columns
+    )

services/api/api/resources/inventories/router.py

@@ -45,7 +45,6 @@
     InventoryDataMetadata,
     InventoryDataResponse,
     InventoryJsonOrientation,
-    InventoryPartitionInfo,
     InventorySortField,
     InventoryType,
     ListInventoriesResponse,
@@ -599,9 +598,9 @@ async def get_inventory_data_metadata(
     """
     # Get Inventory Data Metadata
 
-    Returns partition count, total rows, per-partition row counts, and column
-    names for a completed inventory. Reads only the `_metadata` file from GCS
-    (cached after first access).
+    Returns the partition count, total row count, and column names for a
+    completed inventory (cached after first access). To get a single partition's
+    rows — and its row count — fetch it from the partition data endpoint below.
 
     ## Path Parameters
 
@@ -611,10 +610,10 @@ async def get_inventory_data_metadata(
     ## Response
 
     - **inventory_id**: The inventory ID.
-    - **num_partitions**: Number of Parquet partitions.
+    - **num_partitions**: Number of Parquet partitions (addressable as `0` to
+      `num_partitions - 1` via the partition data endpoint).
     - **total_rows**: Total row count across all partitions.
     - **columns**: List of column names.
-    - **partitions**: Per-partition index and row count.
 
     ## Error Responses
 
@@ -647,10 +646,6 @@ async def get_inventory_data_metadata(
         num_partitions=meta.num_partitions,
         total_rows=meta.total_rows,
         columns=meta.columns,
-        partitions=[
-            InventoryPartitionInfo(index=p.index, num_rows=p.num_rows)
-            for p in meta.partitions
-        ],
     )
 
 
@@ -748,8 +743,7 @@ async def get_inventory_data(
                 detail=(f"Columns not found: {missing}. Available: {meta.columns}"),
             )
 
-    partition = meta.partitions[partition_index]
-    df = await read_partition(inventory_id, partition.path, columns=requested_columns)
+    df = await read_partition(inventory_id, partition_index, columns=requested_columns)
 
     col_names = list(df.columns)
 

services/api/api/resources/inventories/schema.py

@@ -211,17 +211,11 @@ class InventoryJsonOrientation(StrEnum):
     records = "records"
 
 
-class InventoryPartitionInfo(BaseModel):
-    index: int
-    num_rows: int
-
-
 class InventoryDataMetadata(BaseModel):
     inventory_id: str
     num_partitions: int
     total_rows: int
     columns: list[str]
-    partitions: list[InventoryPartitionInfo]
 
 
 class InventoryDataResponse(BaseModel):

services/api/benchmarks/.gitignore

@@ -0,0 +1,2 @@
+# Benchmark result files (timestamped, carry run-specific resource IDs).
+bench_inplace_*.json