DM-55338: Implement staging of update records

python/lsst/dax/ppdb/bigquery/chunk_uploader.py

@@ -230,8 +230,9 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None:
         if not upload_file_list and not manifest.is_empty_chunk:
             raise ChunkUploadError(chunk_id, f"No files found to upload in {chunk_dir} for non-empty chunk")
 
+        gcs_uri = posixpath.join(self.config.bucket_name, gcs_prefix)
         try:
-            # 1) Upload the parquet files to GCS for non-empty chunks.
+            # 1) Upload the parquet files to GCS if there are any.
             if upload_file_list:
                 gcs_names = {
                     file_path: posixpath.join(gcs_prefix, file_path.name) for file_path in upload_file_list
@@ -251,8 +252,7 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None:
                 except* UploadError as eg:
                     raise ChunkUploadError(chunk_id, f"{len(eg.exceptions)} upload(s) failed") from eg
 
-            # 2) Upload manifest, even for empty chunks.
-            gcs_uri = posixpath.join(self.config.bucket_name, gcs_prefix)
+            # 2) Upload the manifest JSON file, even for empty chunks.
             try:
                 self._storage.upload_from_string(
                     posixpath.join(gcs_prefix, manifest_path.name),
@@ -261,33 +261,30 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None:
             except UploadError as e:
                 raise ChunkUploadError(chunk_id, "Manifest upload failed") from e
 
-            # Next two steps are inapplicable to empty chunks.
             if not manifest.is_empty_chunk:
-                # 3) Update the status and GCS URI in the database, marking
-                # chunks with no table data as "staged" since they don't need
-                # to go through the staging process in Dataflow.
-                status = ChunkStatus.UPLOADED if manifest.has_table_data else ChunkStatus.STAGED
-                updated_replica_chunk = replica_chunk.with_new_status(status).with_new_gcs_uri(
+                # 3) Update the database record for the chunk, setting status
+                # to 'staged' and populating the GCS URI.
+                new_status = ChunkStatus.UPLOADED
+                updated_replica_chunk = replica_chunk.with_new_status(new_status).with_new_gcs_uri(
                     f"gs://{gcs_uri}"
                 )
                 try:
                     self._ppdb.update_chunks([updated_replica_chunk], fields={"status", "gcs_uri"})
                     _LOG.info(
                         "Updated replica chunk %d in database with status '%s' and GCS URI: %s",
                         chunk_id,
-                        status,
+                        new_status,
                         gcs_uri,
                     )
                 except Exception as e:
                     raise ChunkUploadError(chunk_id, "Failed to update replica chunk in database") from e
 
                 # 4) Publish Pub/Sub event to trigger the Dataflow staging
-                # job for chunks with table data (skipped for updates-only).
-                if manifest.has_table_data:
-                    try:
-                        self._post_to_stage_chunk_topic(gcs_uri, chunk_id)
-                    except Exception as e:
-                        raise ChunkUploadError(chunk_id, "Failed to publish staging message") from e
+                # for chunks with table data.
+                try:
+                    self._post_to_stage_chunk_topic(gcs_uri, chunk_id)
+                except Exception as e:
+                    raise ChunkUploadError(chunk_id, "Failed to publish staging message") from e
 
         except ChunkUploadError as err:
             try:

python/lsst/dax/ppdb/bigquery/ppdb_bigquery.py

@@ -781,7 +781,7 @@ def _handle_updates(
             The directory containing the chunk's data.
         """
         update_records = UpdateRecords(
-            records=list(apdb_update_records),
+            records=[(replica_chunk.id, record) for record in apdb_update_records],
         )
         parquet_path = chunk_dir / UpdateRecords.PARQUET_FILE_NAME
         update_records.write_parquet_file(parquet_path)

python/lsst/dax/ppdb/bigquery/schema/dataset_builder.py

@@ -217,6 +217,18 @@ def build_tables(self) -> list[bigquery.Table]:
             )
             _update_schema_fields(table, apdb_replica_chunk_field)
 
+        # Add the table which will hold the raw update records.
+        updates_table = bigquery.Table(
+            self._config.fqn_for(self.dataset_type, "updates"),
+            schema=[
+                bigquery.SchemaField("update_time_ns", "INT64", mode="REQUIRED"),
+                bigquery.SchemaField("update_order", "INT64", mode="REQUIRED"),
+                bigquery.SchemaField("json_payload", "STRING", mode="REQUIRED"),
+                bigquery.SchemaField("apdb_replica_chunk", "INT64", mode="REQUIRED"),
+            ],
+        )
+        tables.append(updates_table)
+
         return tables
 
 

python/lsst/dax/ppdb/bigquery/updates/__init__.py

@@ -22,5 +22,4 @@
 from .expanded_update_record import *
 from .updates_merger import *
 from .update_records import *
-from .update_record_expander import *
-from .updates_table import *
+from .expanded_updates_table import *

python/lsst/dax/ppdb/bigquery/updates/expanded_update_record.py

@@ -27,6 +27,8 @@
 
 from pydantic import BaseModel, Field
 
+from lsst.dax.apdb.apdbUpdateRecord import ApdbUpdateRecord
+
 
 class ExpandedUpdateRecord(BaseModel):
     """A single normalized (expanded) update row.
@@ -62,7 +64,7 @@ class ExpandedUpdateRecord(BaseModel):
         description=("New value for the field."),
     )
 
-    replica_chunk_id: int = Field(
+    apdb_replica_chunk: int = Field(
         ...,
         ge=0,
         description=("Source replica chunk identifier associated with this update."),
@@ -76,3 +78,36 @@ class ExpandedUpdateRecord(BaseModel):
     update_order: int = Field(
         description=("Ordering value within the replica chunk or update batch."),
     )
+
+    @classmethod
+    def from_update_record(
+        cls, update_record: ApdbUpdateRecord, apdb_replica_chunk: int
+    ) -> list[ExpandedUpdateRecord]:
+        """Expand a single APDB update record into field-level update rows.
+
+        Parameters
+        ----------
+        update_record
+            The APDB update record to expand.
+        apdb_replica_chunk
+            The replica chunk ID associated with this update record.
+
+        Returns
+        -------
+        `list` [ `ExpandedUpdateRecord` ]
+            One expanded record per field being updated.
+        """
+        table_name = update_record.apdb_table.name
+        record_id_values = tuple(value for _, value in update_record.record_id())
+        return [
+            cls(
+                table_name=table_name,
+                record_id=record_id_values,
+                field_name=field_name,
+                field_value=field_value,
+                apdb_replica_chunk=apdb_replica_chunk,
+                update_order=update_record.update_order,
+                update_time_ns=update_record.update_time_ns,
+            )
+            for field_name, field_value in update_record.record_payload()
+        ]

python/lsst/dax/ppdb/bigquery/updates/updates_table.py → python/lsst/dax/ppdb/bigquery/updates/expanded_updates_table.py

@@ -21,7 +21,7 @@
 
 from __future__ import annotations
 
-__all__ = ["UpdatesTable"]
+__all__ = ["ExpandedUpdatesTable"]
 
 from collections.abc import Iterable
 from typing import Any
@@ -31,9 +31,13 @@
 from .expanded_update_record import ExpandedUpdateRecord
 
 
-class UpdatesTable:
-    """Manage the BigQuery table used for inserting expanded update records
-    which contain one update per row.
+class ExpandedUpdatesTable:
+    """Manage the BigQuery tables holding expanded update records.
+
+    This manages two related tables in the promotion dataset: the
+    ``expanded_updates`` table, which contains one row per updated field, and
+    the ``latest_only`` table, which is derived from it and contains only the
+    latest update for each field.
 
     Parameters
     ----------
@@ -44,35 +48,38 @@ class UpdatesTable:
     dataset_id
         BigQuery dataset ID.
     table_name
-        Name of the updates table. Defaults to ``"updates"``.
+        Name of the expanded updates table. Defaults to
+        ``"expanded_updates"``.
     latest_only_table_name
-        Name of the latest-only updates table. Defaults to
-        ``"updates_latest_only"``.
+        Name of the latest-only updates table. Defaults to ``"latest_only"``.
     """
 
-    _DEFAULT_TABLE_NAME: str = "updates"
-    _DEFAULT_LATEST_ONLY_TABLE_NAME: str = "updates_latest_only"
+    EXPANDED_UPDATES_NAME: str = "expanded_updates"
+    LATEST_ONLY_NAME: str = "latest_only"
 
     def __init__(
         self,
         client: bigquery.Client,
         project_id: str,
         dataset_id: str,
-        table_name: str | None = None,
-        latest_only_table_name: str | None = None,
+        expanded_updates_name: str | None = None,
+        latest_only_name: str | None = None,
     ) -> None:
         self._client: bigquery.Client = client
-        table_name = table_name or self._DEFAULT_TABLE_NAME
-        latest_only_table_name = latest_only_table_name or self._DEFAULT_LATEST_ONLY_TABLE_NAME
-        self._table_fqn = f"{project_id}.{dataset_id}.{table_name}"
-        self._latest_only_table_fqn = f"{project_id}.{dataset_id}.{latest_only_table_name}"
+        expanded_updates_name = expanded_updates_name or self.EXPANDED_UPDATES_NAME
+        latest_only_name = latest_only_name or self.LATEST_ONLY_NAME
+        self._expanded_updates_fqn = f"{project_id}.{dataset_id}.{expanded_updates_name}"
+        self._latest_only_fqn = f"{project_id}.{dataset_id}.{latest_only_name}"
 
     @property
-    def latest_only_table_fqn(self) -> str:
-        """Fully-qualified BigQuery latest-only table name in the form
-        ``"project.dataset.table"`` (`str`, read-only).
-        """
-        return self._latest_only_table_fqn
+    def expanded_updates_fqn(self) -> str:
+        """Fully-qualified expanded updates table name (`str`, read-only)."""
+        return self._expanded_updates_fqn
+
+    @property
+    def latest_only_fqn(self) -> str:
+        """Fully-qualified latest-only table name (`str`, read-only)."""
+        return self._latest_only_fqn
 
     @staticmethod
     def _make_record_key(record_id: Iterable[int]) -> str:
@@ -90,15 +97,13 @@ def _make_record_key(record_id: Iterable[int]) -> str:
         """
         return "-".join(str(x) for x in record_id)
 
-    @property
-    def table_fqn(self) -> str:
-        """Fully-qualified BigQuery table name in the form
-        ``"project.dataset.table"`` (`str`, read-only).
-        """
-        return self._table_fqn
+    def create(self, drop_if_exists: bool = False) -> bigquery.Table:
+        """Create the expanded updates table.
 
-    def create(self) -> bigquery.Table:
-        """Create the updates table.
+        Parameters
+        ----------
+        drop_if_exists
+            If True, drop the table if it already exists before creating it.
 
         Returns
         -------
@@ -109,50 +114,29 @@ def create(self) -> bigquery.Table:
         ------
         google.api_core.exceptions.Conflict
             Raised if the table already exists.
-
-        Notes
-        -----
-        Schema:
-
-        - table_name: STRING (REQUIRED)
-        - record_id: ARRAY<INT64> (REQUIRED)
-        - record_key: STRING (REQUIRED)
-        - field_name: STRING (REQUIRED)
-        - value_json: JSON (REQUIRED)
-        - replica_chunk_id: INT64 (REQUIRED)
-        - update_order: INT64 (REQUIRED)
-        - update_time_ns: INT64 (REQUIRED)
         """
+        if drop_if_exists:
+            self._client.delete_table(self._expanded_updates_fqn, not_found_ok=True)
         schema: list[bigquery.SchemaField] = [
             bigquery.SchemaField("table_name", "STRING", mode="REQUIRED"),
             bigquery.SchemaField("record_id", "INT64", mode="REPEATED"),
             bigquery.SchemaField("record_key", "STRING", mode="REQUIRED"),
             bigquery.SchemaField("field_name", "STRING", mode="REQUIRED"),
             bigquery.SchemaField("value_json", "JSON", mode="REQUIRED"),
-            bigquery.SchemaField("replica_chunk_id", "INT64", mode="REQUIRED"),
+            bigquery.SchemaField("apdb_replica_chunk", "INT64", mode="REQUIRED"),
             bigquery.SchemaField("update_order", "INT64", mode="REQUIRED"),
             bigquery.SchemaField("update_time_ns", "INT64", mode="REQUIRED"),
         ]
-
-        table = bigquery.Table(self._table_fqn, schema=schema)
+        table = bigquery.Table(self.expanded_updates_fqn, schema=schema)
         return self._client.create_table(table)
 
-    def drop(self) -> None:
-        """Drop the table if it exists."""
-        self._client.delete_table(self._table_fqn, not_found_ok=True)
-
-    def recreate(self) -> None:
-        """Drop the table if it exists and then create it."""
-        self.drop()
-        self.create()
-
     def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob:
-        """Insert `ExpandedUpdateRecord` rows into the updates table.
+        """Insert `ExpandedUpdateRecord` rows into the expanded updates table.
 
         Parameters
         ----------
         records
-            Iterable of update records to insert.
+            Iterable of expanded update records to insert.
 
         Returns
         -------
@@ -176,7 +160,7 @@ def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob:
                 "record_key": self._make_record_key(r.record_id),
                 "field_name": r.field_name,
                 "value_json": r.field_value,
-                "replica_chunk_id": r.replica_chunk_id,
+                "apdb_replica_chunk": r.apdb_replica_chunk,
                 "update_order": r.update_order,
                 "update_time_ns": r.update_time_ns,
             }
@@ -185,7 +169,7 @@ def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob:
 
         job = self._client.load_table_from_json(
             rows,
-            self._table_fqn,
+            self._expanded_updates_fqn,
             job_config=bigquery.LoadJobConfig(
                 write_disposition=bigquery.WriteDisposition.WRITE_APPEND,
             ),
@@ -199,30 +183,30 @@ def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob:
 
     def create_latest_only(self) -> None:
         """Select only the latest update for each unique
-        ``(table_name, record_id, field_name)`` combination and write them to a
-        new table.
+        ``(table_name, record_key, field_name)`` combination and write them to
+        the latest-only table.
 
         Notes
         -----
         This keeps only the latest record with an update on an identical
-        ``(table_name, record_id, field_name)``, based on the descending
-        ordering of ``replica_chunk_id``, ``update_time_ns``, and
+        ``(table_name, record_key, field_name)``, based on the descending
+        ordering of ``apdb_replica_chunk``, ``update_time_ns``, and
         ``update_order``.
         """
         query = f"""
-        CREATE OR REPLACE TABLE `{self._latest_only_table_fqn}`
+        CREATE OR REPLACE TABLE `{self._latest_only_fqn}`
         AS
         SELECT * EXCEPT(row_num)
         FROM (
             SELECT *,
                 ROW_NUMBER() OVER (
                     PARTITION BY table_name, record_key, field_name
                     ORDER BY
-                        replica_chunk_id DESC,
+                        apdb_replica_chunk DESC,
                         update_time_ns DESC,
                         update_order DESC
                 ) as row_num
-            FROM `{self._table_fqn}`
+            FROM `{self._expanded_updates_fqn}`
         )
         WHERE row_num = 1
         """