From 391e558703e2d19f5fbbde6592f7bb1171f78811 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Link?= Date: Tue, 30 Jun 2026 21:21:54 -0300 Subject: [PATCH 1/3] Add compact-parquet pipeline and CLI command Compacts small hive-partitioned Parquet files on GCS into larger files using a staging-based swap to keep the original path intact. Adds pyarrow and gcsfs dependencies. Co-Authored-By: Claude Sonnet 4.6 --- .../scratch/nmea-decoded-1-day.yaml | 10 + pyproject.toml | 2 + requirements.txt | 63 +++- src/gfw/ops/cli/commands/compact_parquet.py | 61 ++++ src/gfw/ops/cli/main.py | 2 + .../ops/pipelines/compact_parquet/__init__.py | 3 + src/gfw/ops/pipelines/compact_parquet/main.py | 295 ++++++++++++++++++ tests/pipelines/compact_parquet/__init__.py | 0 tests/pipelines/compact_parquet/test_main.py | 223 +++++++++++++ 9 files changed, 657 insertions(+), 2 deletions(-) create mode 100644 config/compact-parquet/scratch/nmea-decoded-1-day.yaml create mode 100644 src/gfw/ops/cli/commands/compact_parquet.py create mode 100644 src/gfw/ops/pipelines/compact_parquet/__init__.py create mode 100644 src/gfw/ops/pipelines/compact_parquet/main.py create mode 100644 tests/pipelines/compact_parquet/__init__.py create mode 100644 tests/pipelines/compact_parquet/test_main.py diff --git a/config/compact-parquet/scratch/nmea-decoded-1-day.yaml b/config/compact-parquet/scratch/nmea-decoded-1-day.yaml new file mode 100644 index 0000000..108bb4e --- /dev/null +++ b/config/compact-parquet/scratch/nmea-decoded-1-day.yaml @@ -0,0 +1,10 @@ +gcs_output_path: gs://engineering-development-ttl30/tomas.link/parquet-migration/messages/ +gcs_staging_path: gs://engineering-development-ttl30/tomas.link/parquet-migration/_compacted_messages/ +event_source: wf827-pipe-nmea-parsed +start_date: "2023-01-01" +end_date: "2023-01-02" +project: world-fishing-827 +partition_prefix: event_ +target_file_size_mb: 512 +# staging_path defaults to gs://engineering-development-ttl30/tomas.link/parquet-migration/_compact_messages_staging +dry_run: false diff --git a/pyproject.toml b/pyproject.toml index ce37c1e..906a968 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,8 @@ classifiers = [ requires-python = ">= 3.12" dependencies = [ "cloudpathlib[gs]~=0.24", + "duckdb~=1.3", + "gcsfs~=2025.5", "gfw-common[bq]~=0.14", "google-cloud-bigquery~=3.0", "google-cloud-storage~=2.0", diff --git a/requirements.txt b/requirements.txt index 13e7919..5f260bb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,13 @@ # This file was autogenerated by uv via the following command: # uv pip compile --prerelease=allow -o requirements.txt pyproject.toml +aiohappyeyeballs==2.7.1 + # via aiohttp +aiohttp==3.14.1 + # via gcsfs +aiosignal==1.4.0 + # via aiohttp +attrs==26.1.0 + # via aiohttp certifi==2026.6.17 # via requests cffi==2.0.0 @@ -12,6 +20,18 @@ cryptography==47.0.0 # via # google-auth # pyopenssl +decorator==5.3.1 + # via gcsfs +duckdb==1.6.0.dev214 + # via gfw-ops (pyproject.toml) +frozenlist==1.8.0 + # via + # aiohttp + # aiosignal +fsspec==2025.12.0 + # via gcsfs +gcsfs==2025.12.0 + # via gfw-ops (pyproject.toml) gfw-common==0.14.0 # via gfw-ops (pyproject.toml) google-api-core==2.31.0 @@ -19,12 +39,18 @@ google-api-core==2.31.0 # google-cloud-bigquery # google-cloud-core # google-cloud-storage + # google-cloud-storage-control google-auth==2.55.0 # via + # gcsfs # google-api-core + # google-auth-oauthlib # google-cloud-bigquery # google-cloud-core # google-cloud-storage + # google-cloud-storage-control +google-auth-oauthlib==1.4.0 + # via gcsfs google-cloud-bigquery==3.42.0 # via # gfw-ops (pyproject.toml) @@ -37,6 +63,9 @@ google-cloud-storage==2.19.0 # via # gfw-ops (pyproject.toml) # cloudpathlib + # gcsfs +google-cloud-storage-control==1.12.0 + # via gcsfs google-crc32c==1.8.0 # via # google-cloud-storage @@ -48,15 +77,23 @@ google-resumable-media==2.10.0 googleapis-common-protos==1.75.0 # via # google-api-core + # grpc-google-iam-v1 # grpcio-status +grpc-google-iam-v1==0.14.4 + # via google-cloud-storage-control grpcio==1.81.1 # via # google-api-core + # google-cloud-storage-control + # googleapis-common-protos + # grpc-google-iam-v1 # grpcio-status grpcio-status==1.81.1 # via google-api-core idna==3.18 - # via requests + # via + # requests + # yarl jinja2==3.1.6 # via gfw-common markdown-it-py==4.2.0 @@ -65,14 +102,28 @@ markupsafe==3.0.3 # via jinja2 mdurl==0.1.2 # via markdown-it-py +multidict==6.7.1 + # via + # aiohttp + # yarl +oauthlib==3.3.1 + # via requests-oauthlib packaging==26.2 # via google-cloud-bigquery +propcache==0.5.2 + # via + # aiohttp + # yarl proto-plus==1.28.0 - # via google-api-core + # via + # google-api-core + # google-cloud-storage-control protobuf==6.33.6 # via # google-api-core + # google-cloud-storage-control # googleapis-common-protos + # grpc-google-iam-v1 # grpcio-status # proto-plus pyarrow==18.1.0 @@ -97,9 +148,13 @@ pyyaml==6.0.3 # gfw-common requests==2.34.2 # via + # gcsfs # google-api-core # google-cloud-bigquery # google-cloud-storage + # requests-oauthlib +requests-oauthlib==2.0.0 + # via google-auth-oauthlib rich==14.3.4 # via # gfw-ops (pyproject.toml) @@ -113,7 +168,11 @@ sqlparse==0.5.0 # via gfw-common typing-extensions==4.15.0 # via + # aiohttp + # aiosignal # grpcio # pyopenssl urllib3==2.7.0 # via requests +yarl==1.24.2 + # via aiohttp diff --git a/src/gfw/ops/cli/commands/compact_parquet.py b/src/gfw/ops/cli/commands/compact_parquet.py new file mode 100644 index 0000000..c860906 --- /dev/null +++ b/src/gfw/ops/cli/commands/compact_parquet.py @@ -0,0 +1,61 @@ +"""CLI command for compacting hive-partitioned Parquet files on GCS.""" +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +from gfw.common.cli import Command, Option + +from gfw.ops.pipelines import compact_parquet as pipeline + + +_DESCRIPTION = "Compacts small hive-partitioned Parquet files on GCS in-place into larger files." + +HELP_PROJECT = "GCP project for billing." +HELP_GCS_OUTPUT_PATH = "GCS path prefix (gs://bucket/path) of the hive-partitioned files." +HELP_EVENT_SOURCE = "Value for the {prefix}source hive partition key." +HELP_START_DATE = "Start date to compact, inclusive (YYYY-MM-DD)." +HELP_END_DATE = "End date to compact, exclusive (YYYY-MM-DD)." +HELP_PART_PREFIX = "Prefix applied to partition key names in the hive path." +HELP_TARGET_SIZE = "Target output file size in MB." +HELP_MEMORY_LIMIT = "DuckDB memory cap (e.g. '8GB'). DuckDB spills to disk beyond this limit." +HELP_STAGING = ( + "When omitted, operates in swap mode: compacted files are staged to an auto-generated " + "sibling path, then swapped in-place. When set, operates in copy mode: compacted files " + "are written to this path and source files are left untouched." +) +HELP_DRY_RUN = "Log planned compaction and exit without modifying files." + + +class CompactParquet(Command): + """Compact small hive-partitioned Parquet files on GCS in-place into larger files.""" + + @property + def name(self) -> str: + """Command name.""" + return "compact-parquet" + + @property + def description(self) -> str: + """Command description.""" + return _DESCRIPTION + + @property + def options(self) -> list[Option]: + """Command options.""" + return [ + Option("--project", type=str, required=True, help=HELP_PROJECT), + Option("--gcs-output-path", type=str, required=True, help=HELP_GCS_OUTPUT_PATH), + Option("--event-source", type=str, required=True, help=HELP_EVENT_SOURCE), + Option("--start-date", type=str, required=True, help=HELP_START_DATE), + Option("--end-date", type=str, required=True, help=HELP_END_DATE), + Option("--partition-prefix", type=str, default="event_", help=HELP_PART_PREFIX), + Option("--target-file-size-mb", type=int, default=512, help=HELP_TARGET_SIZE), + Option("--memory-limit", type=str, default="8GB", help=HELP_MEMORY_LIMIT), + Option("--gcs-staging-path", type=str, default=None, help=HELP_STAGING), + Option("--dry-run", type=bool, default=False, help=HELP_DRY_RUN), + ] + + def run(self, config: SimpleNamespace, **kwargs: Any) -> None: + """Run the compaction.""" + pipeline.run(**vars(config), **kwargs) diff --git a/src/gfw/ops/cli/main.py b/src/gfw/ops/cli/main.py index 6f7e906..e94bd26 100644 --- a/src/gfw/ops/cli/main.py +++ b/src/gfw/ops/cli/main.py @@ -8,6 +8,7 @@ from gfw.common.logging import LoggerConfig from gfw.ops.cli.commands.bq_to_parquet import BqToParquet +from gfw.ops.cli.commands.compact_parquet import CompactParquet from gfw.ops.cli.commands.create_external_table import CreateExternalTable from gfw.ops.cli.commands.sharded_to_partitioned import ShardedToPartitioned from gfw.ops.version import __version__ @@ -21,6 +22,7 @@ def run(args: list[str]) -> None: subcommands=[ ShardedToPartitioned, BqToParquet, + CompactParquet, CreateExternalTable, ], version=__version__, diff --git a/src/gfw/ops/pipelines/compact_parquet/__init__.py b/src/gfw/ops/pipelines/compact_parquet/__init__.py new file mode 100644 index 0000000..9c4f0db --- /dev/null +++ b/src/gfw/ops/pipelines/compact_parquet/__init__.py @@ -0,0 +1,3 @@ +from .main import run + +__all__ = ["run"] diff --git a/src/gfw/ops/pipelines/compact_parquet/main.py b/src/gfw/ops/pipelines/compact_parquet/main.py new file mode 100644 index 0000000..aaf28e8 --- /dev/null +++ b/src/gfw/ops/pipelines/compact_parquet/main.py @@ -0,0 +1,295 @@ +"""Compact small hive-partitioned Parquet files on GCS into larger files.""" +from __future__ import annotations + +import datetime +import logging +import time +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + +import gcsfs +import duckdb +from cloudpathlib import GSPath +from google.cloud import storage + + +logger = logging.getLogger(__name__) + +_MB = 1024 * 1024 + + +def _date_range(start_date: str, end_date: str) -> list[datetime.date]: + start = datetime.date.fromisoformat(start_date) + end = datetime.date.fromisoformat(end_date) + dates: list[datetime.date] = [] + current = start + while current < end: + dates.append(current) + current += datetime.timedelta(days=1) + return dates + + +def _duckdb_conn(memory_limit: str = "8GB", threads: int = 4) -> duckdb.DuckDBPyConnection: + """Return a DuckDB connection authenticated against GCS via application default credentials.""" + conn = duckdb.connect() + conn.execute(f"SET memory_limit='{memory_limit}'") + conn.execute(f"SET threads={threads}") + conn.execute("SET preserve_insertion_order=false") + conn.register_filesystem(gcsfs.GCSFileSystem()) + return conn + + +@dataclass +class Compactor: + """Compacts hive-partitioned Parquet files for a given source and date range. + + When ``gcs_staging_path`` is ``None`` (the default), operates in *swap* mode: DuckDB + writes compacted output to an auto-generated staging sibling path, the originals are + deleted, and the staged files are moved back so the external table path never changes. + If the process is interrupted between the delete and copy steps, the next run detects + the staging files and resumes automatically. + + When ``gcs_staging_path`` is set, operates in *copy* mode: compacted files are written + directly to ``gcs_staging_path`` without touching the source. Use this when you want + to keep the original uncompacted files alongside the compacted output (e.g. for + benchmarking both versions via separate external tables). + """ + + gcs_client: storage.Client + gcs_output_path: GSPath + event_source: str + partition_prefix: str + target_file_size_mb: int + gcs_staging_path: Optional[GSPath] = None + conn_factory: Callable[[], duckdb.DuckDBPyConnection] = _duckdb_conn + swap: bool = field(init=False) + + def __post_init__(self): + self.swap = False + if self.gcs_staging_path is None: + self.gcs_staging_path = self.default_staging_path + self.swap = True + + @property + def default_staging_path(self) -> GSPath: + return self.gcs_output_path.parent / f"_compact_{self.gcs_output_path.name}_staging" + + def run(self, dates: list[datetime.date]) -> None: + """Compact each date partition, resuming any interrupted swaps first.""" + total = len(dates) + logger.info(f"Starting compaction: {total} date(s) to process") + for i, date in enumerate(dates, 1): + logger.info(f"[{i}/{total}] {date}") + self._compact(date) + + logger.info(f"Finished compaction: {total} date(s) processed") + + def _partition_path(self, base: GSPath, date: datetime.date) -> GSPath: + p = self.partition_prefix + return base / f"{p}source={self.event_source}" / f"{p}date={date.isoformat()}" + + def _list_parquet_blobs(self, base: GSPath, date: datetime.date) -> list[storage.Blob]: + part = self._partition_path(base, date) + blobs = self.gcs_client.list_blobs(part.bucket, prefix=f"{part.blob}/") + return [b for b in blobs if b.name.endswith(".parquet")] + + def _compact(self, date: datetime.date) -> None: + source_blobs = self._list_parquet_blobs(self.gcs_output_path, date) + dest_blobs = self._list_parquet_blobs(self.gcs_staging_path, date) + + if self.swap and dest_blobs and not source_blobs: + # Interrupted between delete-originals and copy-from-staging: resume the swap. + logger.warning( + f"Resuming interrupted swap for {date}: " + f"{len(dest_blobs)} staged file(s) found, originals already deleted" + ) + self._copy_to_partition(date, dest_blobs, self.gcs_output_path) + self._delete_blobs(dest_blobs) + return + + if not source_blobs: + logger.info(f"No source files for {date}, skipping") + return + + if self.swap and len(source_blobs) == 1: + logger.info(f"Skipping {date}: already a single file") + return + + if dest_blobs: + logger.info(f"Removing {len(dest_blobs)} existing file(s) from dest for {date}") + self._delete_blobs(dest_blobs) + + source_part = self._partition_path(self.gcs_output_path, date) + source_uris = [f"gs://{source_part.bucket}/{b.name}" for b in source_blobs] + total_mb = sum(b.size for b in source_blobs) / _MB + + logger.info( + f"Compacting {len(source_uris)} file(s) for {date} ({total_mb:.0f} MB compressed)" + ) + + t0 = time.monotonic() + written = self._write_compacted(date, source_uris, self.gcs_staging_path) + logger.info(f"DuckDB done in {time.monotonic() - t0:.1f}s — {len(written)} file(s)") + + if self.swap: + # Originals are deleted before the copy so we never have duplicate data visible + # to the external table. The staging files serve as the recovery point if the + # copy step is interrupted. + t0 = time.monotonic() + self._delete_blobs(source_blobs) + self._copy_to_partition(date, written, self.gcs_output_path) + self._delete_blobs(written) + logger.info(f"Swap done in {time.monotonic() - t0:.1f}s") + + def _write_compacted( + self, date: datetime.date, source_uris: list[str], dest_path: GSPath + ) -> list[storage.Blob]: + """Compact source files into dest_path using DuckDB.""" + dest_part = self._partition_path(dest_path, date) + dest_dir = f"gs://{dest_part.bucket}/{dest_part.blob}" + files_sql = "[" + ", ".join(f"'{u}'" for u in source_uris) + "]" + target_bytes = self.target_file_size_mb * _MB + + conn = self.conn_factory() + logger.info(f"Writing compacted output to {dest_dir}") + conn.execute( + f""" + COPY (SELECT * FROM read_parquet({files_sql})) + TO '{dest_dir}' + ( + FORMAT PARQUET, + COMPRESSION SNAPPY, + FILE_SIZE_BYTES {target_bytes}, + PER_THREAD_OUTPUT false + ) + """ + ) + conn.close() + + written = self._list_parquet_blobs(dest_path, date) + logger.info(f"Wrote {len(written)} file(s) to {dest_dir}") + return written + + def _copy_to_partition( + self, date: datetime.date, blobs: list[storage.Blob], dest_path: GSPath + ) -> None: + dest_part = self._partition_path(dest_path, date) + bucket = self.gcs_client.bucket(dest_part.bucket) + for blob in blobs: + filename = blob.name.rsplit("/", 1)[-1] + dest_name = f"{dest_part.blob}/{filename}" + bucket.copy_blob(blob, bucket, new_name=dest_name) + logger.info(f"Copied to gs://{dest_part.bucket}/{dest_name}") + + def _delete_blobs(self, blobs: list[storage.Blob]) -> None: + for blob in blobs: + blob.delete() + logger.info(f"Deleted gs://{blob.bucket.name}/{blob.name}") + + +def run( + project: str, + gcs_output_path: str, + event_source: str, + start_date: str, + end_date: str, + partition_prefix: str = "event_", + target_file_size_mb: int = 512, + memory_limit: str = "8GB", + gcs_staging_path: str | None = None, + dry_run: bool = False, + gcs_client_factory: Callable[[str], storage.Client] = storage.Client, + conn_factory: Callable[[], duckdb.DuckDBPyConnection] | None = None, + unknown_unparsed_args: tuple = (), + unknown_parsed_args: dict | None = None, + **kwargs: Any, +) -> None: + """Compact small hive-partitioned Parquet files on GCS into larger files. + + For each date partition, DuckDB reads all source files in parallel and writes + compacted output to a staging path. The staged files are then swapped into the + original partition path so the external table path never changes. If the process + is interrupted after deleting originals, the next run detects the staged files + and resumes the copy step automatically. + + Path structure:: + + {gcs_output_path}/{prefix}source={event_source}/{prefix}date=YYYY-MM-DD/*.parquet + + Staging path defaults to ``gs://{bucket}/{parent}/_compact_{table}_staging`` + (sibling of the table directory, same bucket so copies are server-side). + + Args: + project: + GCP project used for billing. + + gcs_output_path: + GCS path prefix (gs://bucket/path) of the hive-partitioned files. + + event_source: + Value of the ``{prefix}source`` hive partition key + (e.g. ``"wf827-pipe-nmea-parsed"``). + + start_date: + Start date, inclusive (YYYY-MM-DD). + + end_date: + End date, exclusive (YYYY-MM-DD). + + partition_prefix: + Prefix applied to partition key names. Defaults to ``"event_"``. + + target_file_size_mb: + Target size for each output file in MB. Defaults to 512. + + gcs_staging_path: + When ``None`` (default), operates in swap mode: compacted files are written + to an auto-generated staging sibling path, the originals are deleted, and the + staged files are moved back in-place. When set, operates in copy mode: compacted + files are written to this path and the source files are left untouched, giving + you both uncompacted and compacted paths for separate external tables. + + dry_run: + Log planned compaction and exit without modifying files. + + gcs_client_factory: + Injectable factory for :class:`~google.cloud.storage.Client`. + + memory_limit: + DuckDB memory cap (e.g. ``"8GB"``). DuckDB spills to disk beyond this limit. + Defaults to ``"8GB"``. Set lower on machines with less RAM, higher on GKE pods. + + conn_factory: + Injectable factory for a configured DuckDB connection. Defaults to + :func:`_duckdb_conn` which authenticates via application default credentials. + Override in tests to avoid real GCS calls. + + unknown_unparsed_args: + Extra unparsed CLI args (ignored). + + unknown_parsed_args: + Extra parsed args (ignored). + """ + source = GSPath(gcs_output_path.rstrip("/")) + dates = _date_range(start_date, end_date) + + logger.info(f"Compacting {gcs_output_path} for [{start_date}, {end_date}) ({len(dates)} days)") + + if dry_run: + for date in dates: + p = partition_prefix + logger.info(f"[dry-run] {gcs_output_path}/{p}source={event_source}/{p}date={date}") + return + + resolved_conn_factory = conn_factory or (lambda: _duckdb_conn(memory_limit=memory_limit)) + gcs_client = gcs_client_factory(project=project) + compactor = Compactor( + gcs_client=gcs_client, + gcs_output_path=source, + event_source=event_source, + partition_prefix=partition_prefix, + target_file_size_mb=target_file_size_mb, + gcs_staging_path=GSPath(gcs_staging_path.rstrip("/")) if gcs_staging_path else None, + conn_factory=resolved_conn_factory, + ) + compactor.run(dates) diff --git a/tests/pipelines/compact_parquet/__init__.py b/tests/pipelines/compact_parquet/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/pipelines/compact_parquet/test_main.py b/tests/pipelines/compact_parquet/test_main.py new file mode 100644 index 0000000..2459c85 --- /dev/null +++ b/tests/pipelines/compact_parquet/test_main.py @@ -0,0 +1,223 @@ +import datetime +from unittest.mock import MagicMock, patch + +from gfw.ops.pipelines.compact_parquet.main import Compactor, _date_range, run +from cloudpathlib import GSPath + + +DATE = datetime.date(2024, 1, 1) +GCS_PATH = GSPath("gs://bucket/path/messages") + + +def _make_blob(name: str, size: int = 10 * 1024 * 1024) -> MagicMock: + blob = MagicMock() + blob.name = name + blob.size = size + blob.bucket.name = "bucket" + return blob + + +def _source_blobs(n: int = 3) -> list[MagicMock]: + return [ + _make_blob(f"path/messages/event_source=src/event_date=2024-01-01/part_{i}.parquet") + for i in range(n) + ] + + +def _staging_blobs(n: int = 1) -> list[MagicMock]: + return [ + _make_blob( + f"path/_compact_messages_staging/event_source=src/" + f"event_date=2024-01-01/part_{i}.parquet" + ) + for i in range(n) + ] + + +def _make_compactor( + source_blobs=None, + staging_blobs=None, + conn_factory=None, +) -> Compactor: + """Return a swap-mode Compactor with mocked GCS client and DuckDB connection.""" + source_blobs = source_blobs if source_blobs is not None else _source_blobs() + staging_blobs = staging_blobs if staging_blobs is not None else [] + + mock_gcs = MagicMock() + + def list_blobs(bucket, prefix): + if "_staging" in prefix: + return iter(staging_blobs) + return iter(source_blobs) + + mock_gcs.list_blobs.side_effect = list_blobs + + if conn_factory is None: + def conn_factory(): + return MagicMock() + + return Compactor( + gcs_client=mock_gcs, + gcs_output_path=GCS_PATH, + event_source="src", + partition_prefix="event_", + target_file_size_mb=512, + conn_factory=conn_factory, + ) + + +# --- _date_range --- + + +def test_date_range_basic(): + dates = _date_range("2024-01-01", "2024-01-04") + assert dates == [ + datetime.date(2024, 1, 1), + datetime.date(2024, 1, 2), + datetime.date(2024, 1, 3), + ] + + +def test_date_range_empty(): + assert _date_range("2024-01-01", "2024-01-01") == [] + + +# --- skip logic --- + + +def test_compact_skips_when_no_source_files(): + compactor = _make_compactor(source_blobs=[], staging_blobs=[]) + with patch.object(compactor, "_write_compacted") as mock_write: + compactor._compact(DATE) + mock_write.assert_not_called() + + +def test_compact_skips_when_already_single_file(): + compactor = _make_compactor(source_blobs=_source_blobs(1)) + with patch.object(compactor, "_write_compacted") as mock_write: + compactor._compact(DATE) + mock_write.assert_not_called() + + +# --- normal compaction flow --- + + +def test_compact_normal_flow_order(): + """write staging → delete source → copy staging → delete staging.""" + staged = _staging_blobs(1) + compactor = _make_compactor() + calls = [] + + with patch.object( + compactor, "_write_compacted", return_value=staged + ) as mock_write: + with patch.object( + compactor, "_delete_blobs", side_effect=lambda b: calls.append(("delete", b)) + ): + with patch.object( + compactor, "_copy_to_partition", side_effect=lambda d, b, p: calls.append(("copy", b)) + ): + compactor._compact(DATE) + + mock_write.assert_called_once() + assert calls[0][0] == "delete" # source deleted first + assert calls[1][0] == "copy" # then staging copied + assert calls[2][0] == "delete" # then staging deleted + + +def test_compact_deletes_source_before_copy(): + """Source blobs must be deleted before the copy step.""" + source = _source_blobs(3) + staged = _staging_blobs(1) + compactor = _make_compactor(source_blobs=source) + deleted = [] + + with patch.object(compactor, "_write_compacted", return_value=staged): + with patch.object(compactor, "_delete_blobs", side_effect=lambda b: deleted.append(b)): + with patch.object(compactor, "_copy_to_partition"): + compactor._compact(DATE) + + assert deleted[0] == source + + +# --- resume interrupted swap --- + + +def test_compact_resumes_interrupted_swap(): + """If staging exists but source is gone, copy staging and clean up.""" + staged = _staging_blobs(1) + compactor = _make_compactor(source_blobs=[], staging_blobs=staged) + + with patch.object(compactor, "_copy_to_partition") as mock_copy: + with patch.object(compactor, "_delete_blobs") as mock_delete: + compactor._compact(DATE) + + mock_copy.assert_called_once_with(DATE, staged, GCS_PATH) + mock_delete.assert_called_once_with(staged) + + +# --- leftover staging cleanup --- + + +def test_compact_cleans_up_leftover_staging_before_writing(): + """Leftover staging files from a prior run are deleted before writing new ones.""" + leftover = _staging_blobs(2) + staged = _staging_blobs(1) + compactor = _make_compactor(staging_blobs=leftover) + deleted = [] + + with patch.object(compactor, "_write_compacted", return_value=staged): + with patch.object(compactor, "_delete_blobs", side_effect=lambda b: deleted.append(b)): + with patch.object(compactor, "_copy_to_partition"): + compactor._compact(DATE) + + assert deleted[0] == leftover # leftover deleted first + + +# --- DuckDB SQL generation --- + + +def test_write_compacted_to_staging_executes_copy_sql(): + mock_conn = MagicMock() + compactor = _make_compactor(conn_factory=lambda: mock_conn) + + source_uris = [ + "gs://bucket/path/messages/event_source=src/event_date=2024-01-01/part_0.parquet" + ] + + with patch.object(compactor, "_list_parquet_blobs", return_value=_staging_blobs(1)): + compactor._write_compacted(DATE, source_uris, compactor.gcs_staging_path) + + executed_sql = mock_conn.execute.call_args[0][0] + assert "COPY" in executed_sql + assert "FORMAT PARQUET" in executed_sql + assert "COMPRESSION SNAPPY" in executed_sql + assert source_uris[0] in executed_sql + assert str(512 * 1024 * 1024) in executed_sql + + +def test_write_compacted_to_staging_closes_connection(): + mock_conn = MagicMock() + compactor = _make_compactor(conn_factory=lambda: mock_conn) + + with patch.object(compactor, "_list_parquet_blobs", return_value=_staging_blobs(1)): + compactor._write_compacted(DATE, ["gs://bucket/x/part.parquet"], compactor.gcs_staging_path) + + mock_conn.close.assert_called_once() + + +# --- dry run --- + + +def test_run_dry_run_makes_no_gcs_calls(): + mock_gcs = MagicMock() + run( + project="proj", + gcs_output_path="gs://bucket/messages", + event_source="src", + start_date="2024-01-01", + end_date="2024-01-03", + dry_run=True, + gcs_client_factory=lambda project: mock_gcs, + ) + mock_gcs.list_blobs.assert_not_called() From b11e11e2e6b49994d2376d22847e8f190a7007b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Link?= Date: Tue, 30 Jun 2026 21:21:54 -0300 Subject: [PATCH 2/3] Add benchmark-parquet pipeline and CLI command Co-Authored-By: Claude Sonnet 4.6 --- Dockerfile | 4 +- Makefile | 4 +- .../scratch/nmea-parsed-1-day.yaml | 22 ++ .../scratch/nmea-parsed-1-month.yaml | 30 ++ pyproject.toml | 7 + requirements.txt | 37 +-- src/gfw/ops/cli/commands/benchmark_parquet.py | 70 ++++ src/gfw/ops/cli/commands/compact_parquet.py | 12 +- src/gfw/ops/cli/main.py | 2 + .../pipelines/benchmark_parquet/__init__.py | 5 + .../ops/pipelines/benchmark_parquet/main.py | 302 ++++++++++++++++++ src/gfw/ops/pipelines/compact_parquet/main.py | 150 ++++++--- tests/cli/commands/test_compact_parquet.py | 28 ++ tests/pipelines/compact_parquet/test_main.py | 269 ++++++++++++---- 14 files changed, 801 insertions(+), 141 deletions(-) create mode 100644 config/benchmark-parquet/scratch/nmea-parsed-1-day.yaml create mode 100644 config/benchmark-parquet/scratch/nmea-parsed-1-month.yaml create mode 100644 src/gfw/ops/cli/commands/benchmark_parquet.py create mode 100644 src/gfw/ops/pipelines/benchmark_parquet/__init__.py create mode 100644 src/gfw/ops/pipelines/benchmark_parquet/main.py create mode 100644 tests/cli/commands/test_compact_parquet.py diff --git a/Dockerfile b/Dockerfile index 123304a..7e0838e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,7 +15,7 @@ ENV UV_COMPILE_BYTECODE=1 COPY pyproject.toml requirements.txt README.md MANIFEST.in ./ RUN uv pip install --system --upgrade pip && \ uv pip install --system build && \ - uv pip install --system --prefix=/install --prerelease=allow -r requirements.txt + uv pip install --system --prefix=/install -r requirements.txt COPY src ./src RUN uv pip install --system --prefix=/install --no-deps . @@ -48,7 +48,7 @@ RUN apt-get update && \ WORKDIR /opt/project COPY . . -RUN uv pip install --system --prerelease=allow -e .[lint,dev,build] && \ +RUN uv pip install --system -e .[lint,dev,build] && \ uv pip install --system -r requirements-test.txt # --------------------------------------------------------------------------------------- diff --git a/Makefile b/Makefile index cb4d228..bbc8a5a 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,7 @@ UV_VERSION := 0.10.9 VENV:=uv venv PIP:=uv pip -PIP_COMPILE:=uv pip compile --prerelease=allow +PIP_COMPILE:=uv pip compile @@ -81,7 +81,7 @@ install-test: upgrade-pip .PHONY: install ## Install the package in editable mode & all dependencies for local development. install: upgrade-pip - ${PIP} install --prerelease=allow -e .[lint,dev,build] + ${PIP} install -e .[lint,dev,build] make install-test .PHONY: test ## Run all unit tests exporting coverage.xml report. diff --git a/config/benchmark-parquet/scratch/nmea-parsed-1-day.yaml b/config/benchmark-parquet/scratch/nmea-parsed-1-day.yaml new file mode 100644 index 0000000..bb22f8d --- /dev/null +++ b/config/benchmark-parquet/scratch/nmea-parsed-1-day.yaml @@ -0,0 +1,22 @@ +project: world-fishing-827 +start_date: "2023-02-07" +end_date: "2023-02-08" +queries: + - aggregation +agg_column: tagblock_timestamp +runs: 3 +location: US +max_gb_billed: 100.0 +tables: + native: + table: world-fishing-827.scratch_tomas_ttl30d.pipe_nmea_parsed_2023_feb + column: tagblock_timestamp + type: native + external_uncompacted: + table: gfw-int-ais-staging-datalake.vessel_transmissions_v1.messages__PR_2627__feature_bq_to_parquet_dag + column: event_date + type: external + external_compacted: + table: gfw-int-ais-staging-datalake.vessel_transmissions_v1._messages_compacted__PR_2627__feature_bq_to_parquet_dag + column: event_date + type: external diff --git a/config/benchmark-parquet/scratch/nmea-parsed-1-month.yaml b/config/benchmark-parquet/scratch/nmea-parsed-1-month.yaml new file mode 100644 index 0000000..73d9b0a --- /dev/null +++ b/config/benchmark-parquet/scratch/nmea-parsed-1-month.yaml @@ -0,0 +1,30 @@ +project: world-fishing-827 +start_date: "2023-02-01" +end_date: "2023-03-01" +queries: + - aggregation +agg_column: tagblock_timestamp +runs: 3 +location: US +max_gb_billed: 2000.0 +tables: + native: + table: world-fishing-827.scratch_tomas_ttl30d.pipe_nmea_parsed_2023_feb + column: tagblock_timestamp + type: native + external_uncompacted: + table: gfw-int-ais-staging-datalake.vessel_transmissions_v1.messages__PR_2627__feature_bq_to_parquet_dag + column: event_date + type: external + external_compacted: + table: gfw-int-ais-staging-datalake.vessel_transmissions_v1._messages_compacted__PR_2627__feature_bq_to_parquet_dag + column: event_date + type: external + external_uncompacted_cached: + table: gfw-int-ais-staging-datalake.vessel_transmissions_v1.messages__PR_2627__feature_bq_to_parquet_dag_cached + column: event_date + type: external + external_compacted_cached: + table: gfw-int-ais-staging-datalake.vessel_transmissions_v1._messages_compacted__PR_2627__feature_bq_to_parquet_dag_cached + column: event_date + type: external \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 906a968..4d73460 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "pyyaml~=6.0", "rich~=14.0", "rich-argparse~=1.6", + "tenacity~=9.0", ] [project.urls] @@ -150,6 +151,12 @@ addopts = "-v --cov=src --cov=init_project --cov-report=term-missing" branch = true parallel = true context = "${CONTEXT}" +omit = [ + "src/gfw/ops/pipelines/benchmark_parquet/main.py", # Analysis, not used in production. + "src/gfw/ops/cli/commands/benchmark_parquet.py", # Analysis, not used in production. + "**/__init__.py", + +] [tool.coverage.report] precision = 0 diff --git a/requirements.txt b/requirements.txt index 5f260bb..d7b5a22 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile --prerelease=allow -o requirements.txt pyproject.toml +# uv pip compile -o requirements.txt pyproject.toml aiohappyeyeballs==2.7.1 # via aiohttp aiohttp==3.14.1 @@ -10,19 +10,17 @@ attrs==26.1.0 # via aiohttp certifi==2026.6.17 # via requests -cffi==2.0.0 +cffi==2.1.0 # via cryptography -charset-normalizer==3.4.7 +charset-normalizer==3.4.9 # via requests cloudpathlib==0.24.0 # via gfw-ops (pyproject.toml) -cryptography==47.0.0 - # via - # google-auth - # pyopenssl +cryptography==49.0.0 + # via google-auth decorator==5.3.1 # via gcsfs -duckdb==1.6.0.dev214 +duckdb==1.5.4 # via gfw-ops (pyproject.toml) frozenlist==1.8.0 # via @@ -32,7 +30,7 @@ fsspec==2025.12.0 # via gcsfs gcsfs==2025.12.0 # via gfw-ops (pyproject.toml) -gfw-common==0.14.0 +gfw-common==0.14.1 # via gfw-ops (pyproject.toml) google-api-core==2.31.0 # via @@ -40,7 +38,7 @@ google-api-core==2.31.0 # google-cloud-core # google-cloud-storage # google-cloud-storage-control -google-auth==2.55.0 +google-auth==2.55.2 # via # gcsfs # google-api-core @@ -51,7 +49,7 @@ google-auth==2.55.0 # google-cloud-storage-control google-auth-oauthlib==1.4.0 # via gcsfs -google-cloud-bigquery==3.42.0 +google-cloud-bigquery==3.42.2 # via # gfw-ops (pyproject.toml) # gfw-common @@ -81,14 +79,14 @@ googleapis-common-protos==1.75.0 # grpcio-status grpc-google-iam-v1==0.14.4 # via google-cloud-storage-control -grpcio==1.81.1 +grpcio==1.82.1 # via # google-api-core # google-cloud-storage-control # googleapis-common-protos # grpc-google-iam-v1 # grpcio-status -grpcio-status==1.81.1 +grpcio-status==1.82.1 # via google-api-core idna==3.18 # via @@ -114,11 +112,11 @@ propcache==0.5.2 # via # aiohttp # yarl -proto-plus==1.28.0 +proto-plus==1.28.1 # via # google-api-core # google-cloud-storage-control -protobuf==6.33.6 +protobuf==7.35.1 # via # google-api-core # google-cloud-storage-control @@ -128,7 +126,7 @@ protobuf==6.33.6 # proto-plus pyarrow==18.1.0 # via gfw-common -pyasn1==0.6.3 +pyasn1==0.6.4 # via pyasn1-modules pyasn1-modules==0.4.2 # via google-auth @@ -136,8 +134,6 @@ pycparser==3.0 # via cffi pygments==2.20.0 # via rich -pyopenssl==26.2.0 - # via google-auth python-dateutil==2.9.0.post0 # via # gfw-common @@ -166,12 +162,13 @@ six==1.17.0 # via python-dateutil sqlparse==0.5.0 # via gfw-common -typing-extensions==4.15.0 +tenacity==9.1.4 + # via gfw-ops (pyproject.toml) +typing-extensions==4.16.0 # via # aiohttp # aiosignal # grpcio - # pyopenssl urllib3==2.7.0 # via requests yarl==1.24.2 diff --git a/src/gfw/ops/cli/commands/benchmark_parquet.py b/src/gfw/ops/cli/commands/benchmark_parquet.py new file mode 100644 index 0000000..330f295 --- /dev/null +++ b/src/gfw/ops/cli/commands/benchmark_parquet.py @@ -0,0 +1,70 @@ +"""CLI command for benchmarking BigQuery table query performance.""" +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +from gfw.common.cli import Command, Option +from gfw.common.cli.actions import NestedKeyValueAction + +from gfw.ops.pipelines import benchmark_parquet as pipeline + + +_DESCRIPTION = "Benchmark SELECT * query performance across one or more BigQuery tables." + +HELP_PROJECT = "GCP project for billing." +HELP_START_DATE = "Start of date range, inclusive (YYYY-MM-DD)." +HELP_END_DATE = "End of date range, exclusive (YYYY-MM-DD)." +HELP_TABLES = ( + "Table configurations as dotted key=value pairs. " + "Each table is keyed by its display label. " + "Required fields: table, column. " + "Optional: type (native|external, default: external). " + "Example: --tables native.table=project.dataset.t native.column=ts native.type=native " + "external.table=project.dataset.ext external.column=event_date" +) +HELP_QUERIES = ( + "Which queries to run. Valid values: select_star, aggregation. " + "Defaults to all. Example: --queries aggregation" +) +HELP_AGG_COLUMN = ( + "Timestamp column name to aggregate by hour, applied to all tables. " + "Example: --agg-column tagblock_timestamp. " + "If omitted, the aggregation query is skipped." +) +HELP_RUNS = "Number of times each query is repeated. Results are averaged." +HELP_LOCATION = "BigQuery job location." +HELP_MAX_GB = "Safety cap on bytes billed per query in GB. Pass 0 to disable." + + +class BenchmarkParquet(Command): + """Benchmark SELECT * query performance across one or more BigQuery tables.""" + + @property + def name(self) -> str: + """Command name.""" + return "benchmark-parquet" + + @property + def description(self) -> str: + """Command description.""" + return _DESCRIPTION + + @property + def options(self) -> list[Option]: + """Command options.""" + return [ + Option("--project", type=str, required=True, help=HELP_PROJECT), + Option("--start-date", type=str, required=True, help=HELP_START_DATE), + Option("--end-date", type=str, required=True, help=HELP_END_DATE), + Option("--tables", type=str, nargs="+", action=NestedKeyValueAction, help=HELP_TABLES), + Option("--queries", type=str, nargs="+", default=None, help=HELP_QUERIES), + Option("--agg-column", type=str, default=None, help=HELP_AGG_COLUMN), + Option("--runs", type=int, default=1, help=HELP_RUNS), + Option("--location", type=str, default="US", help=HELP_LOCATION), + Option("--max-gb-billed", type=float, default=100.0, help=HELP_MAX_GB), + ] + + def run(self, config: SimpleNamespace, **kwargs: Any) -> None: + """Run the benchmark.""" + pipeline.run(**vars(config), **kwargs) diff --git a/src/gfw/ops/cli/commands/compact_parquet.py b/src/gfw/ops/cli/commands/compact_parquet.py index c860906..3d2fe78 100644 --- a/src/gfw/ops/cli/commands/compact_parquet.py +++ b/src/gfw/ops/cli/commands/compact_parquet.py @@ -18,13 +18,15 @@ HELP_END_DATE = "End date to compact, exclusive (YYYY-MM-DD)." HELP_PART_PREFIX = "Prefix applied to partition key names in the hive path." HELP_TARGET_SIZE = "Target output file size in MB." -HELP_MEMORY_LIMIT = "DuckDB memory cap (e.g. '8GB'). DuckDB spills to disk beyond this limit." -HELP_STAGING = ( +HELP_THREADS = "Number of DuckDB threads for parallel reads/writes." +HELP_MEMORY_LIMIT = "DuckDB memory cap in GB. DuckDB spills to disk beyond this limit." +HELP_STAGING_PATH = ( "When omitted, operates in swap mode: compacted files are staged to an auto-generated " "sibling path, then swapped in-place. When set, operates in copy mode: compacted files " "are written to this path and source files are left untouched." ) HELP_DRY_RUN = "Log planned compaction and exit without modifying files." +HELP_MAX_RETRIES = "Number of retries on transient DuckDB I/O errors before failing the task." class CompactParquet(Command): @@ -51,9 +53,11 @@ def options(self) -> list[Option]: Option("--end-date", type=str, required=True, help=HELP_END_DATE), Option("--partition-prefix", type=str, default="event_", help=HELP_PART_PREFIX), Option("--target-file-size-mb", type=int, default=512, help=HELP_TARGET_SIZE), - Option("--memory-limit", type=str, default="8GB", help=HELP_MEMORY_LIMIT), - Option("--gcs-staging-path", type=str, default=None, help=HELP_STAGING), + Option("--memory-limit-gb", type=int, default=12, help=HELP_MEMORY_LIMIT), + Option("--threads", type=int, default=4, help=HELP_THREADS), + Option("--gcs-staging-path", type=str, default=None, help=HELP_STAGING_PATH), Option("--dry-run", type=bool, default=False, help=HELP_DRY_RUN), + Option("--max-retries", type=int, default=3, help=HELP_MAX_RETRIES), ] def run(self, config: SimpleNamespace, **kwargs: Any) -> None: diff --git a/src/gfw/ops/cli/main.py b/src/gfw/ops/cli/main.py index e94bd26..cc663d1 100644 --- a/src/gfw/ops/cli/main.py +++ b/src/gfw/ops/cli/main.py @@ -7,6 +7,7 @@ from gfw.common.cli import CLI from gfw.common.logging import LoggerConfig +from gfw.ops.cli.commands.benchmark_parquet import BenchmarkParquet from gfw.ops.cli.commands.bq_to_parquet import BqToParquet from gfw.ops.cli.commands.compact_parquet import CompactParquet from gfw.ops.cli.commands.create_external_table import CreateExternalTable @@ -24,6 +25,7 @@ def run(args: list[str]) -> None: BqToParquet, CompactParquet, CreateExternalTable, + BenchmarkParquet, ], version=__version__, allow_unknown=True, diff --git a/src/gfw/ops/pipelines/benchmark_parquet/__init__.py b/src/gfw/ops/pipelines/benchmark_parquet/__init__.py new file mode 100644 index 0000000..3b22766 --- /dev/null +++ b/src/gfw/ops/pipelines/benchmark_parquet/__init__.py @@ -0,0 +1,5 @@ +"""Benchmark query performance of a native BigQuery table vs a Parquet external table.""" +from gfw.ops.pipelines.benchmark_parquet import main + + +run = main.run diff --git a/src/gfw/ops/pipelines/benchmark_parquet/main.py b/src/gfw/ops/pipelines/benchmark_parquet/main.py new file mode 100644 index 0000000..405a6f1 --- /dev/null +++ b/src/gfw/ops/pipelines/benchmark_parquet/main.py @@ -0,0 +1,302 @@ +"""Benchmark query performance of a native BigQuery table vs a Parquet external table.""" +from __future__ import annotations + +import logging +import time +import uuid +from dataclasses import dataclass, field +from typing import Any, Callable + +from google.cloud import bigquery +from rich.console import Console +from rich.table import Table + + +logger = logging.getLogger(__name__) + +_DATE_FILTERS = { + "native": "DATE({column}) >= '{start_date}' AND DATE({column}) < '{end_date}'", + "external": "{column} >= '{start_date}' AND {column} < '{end_date}'", +} + +_QUERIES = { + "select_star": "SELECT * FROM `{table}` WHERE {date_filter}", + "aggregation": ( + "SELECT EXTRACT(HOUR FROM {agg_column}) AS hour, COUNT(*) AS cnt" + " FROM `{table}` WHERE {date_filter}" + " GROUP BY 1 ORDER BY 1" + ), +} + + +@dataclass +class QueryStats: + """Stats captured from a single BigQuery job execution.""" + + bytes_processed: int + slot_ms: int + elapsed_s: float + + +@dataclass +class BenchmarkResult: + """Aggregated stats for one query × one table across N runs.""" + + label: str + query: str + runs: list[QueryStats] = field(default_factory=list) + + @property + def avg_bytes(self) -> float: + """Average bytes processed across runs.""" + return sum(r.bytes_processed for r in self.runs) / len(self.runs) + + @property + def avg_slot_ms(self) -> float: + """Average slot milliseconds across runs.""" + return sum(r.slot_ms for r in self.runs) / len(self.runs) + + @property + def avg_elapsed_s(self) -> float: + """Average wall-clock elapsed time across runs.""" + return sum(r.elapsed_s for r in self.runs) / len(self.runs) + + @property + def min_elapsed_s(self) -> float: + """Minimum wall-clock elapsed time across runs.""" + return min(r.elapsed_s for r in self.runs) + + +_SCRATCH_DATASET = "gfw_benchmark_scratch" + + +def _ensure_scratch_dataset(client: bigquery.Client, location: str) -> None: + dataset_ref = bigquery.DatasetReference(client.project, _SCRATCH_DATASET) + dataset = bigquery.Dataset(dataset_ref) + dataset.location = location + client.create_dataset(dataset, exists_ok=True) + + +def _run_query( + client: bigquery.Client, + sql: str, + location: str, + max_bytes_billed: int | None, +) -> QueryStats: + dest = f"{client.project}.{_SCRATCH_DATASET}.run_{uuid.uuid4().hex[:12]}" + job_config = bigquery.QueryJobConfig( + use_query_cache=False, + maximum_bytes_billed=max_bytes_billed, + allow_large_results=True, + destination=dest, + write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE, + ) + job = client.query(sql, job_config=job_config, location=location) + + while job.state != "DONE": + time.sleep(1) + job = client.get_job(job.job_id, location=location) + + client.delete_table(dest, not_found_ok=True) + + if job.errors: + raise RuntimeError(f"Query job failed: {job.errors}") + + elapsed_s = (job.ended - job.started).total_seconds() + + return QueryStats( + bytes_processed=job.total_bytes_processed or 0, + slot_ms=job.slot_millis or 0, + elapsed_s=elapsed_s, + ) + + +def _fmt_bytes(n: float) -> str: + if n < 1024: + return f"{n:.0f} B" + + if n < 1024**2: + return f"{n/1024:.1f} KB" + + if n < 1024**3: + return f"{n/1024**2:.1f} MB" + + return f"{n/1024**3:.2f} GB" + + +def _ratio(a: float, b: float) -> str: + if b == 0: + return "—" + + return f"{a/b:.1f}×" + + +def _build_queries( + queries_to_include: set[str], + table_ref: str, + date_filter: str, + agg_column: str | None, +) -> list[tuple[str, str]]: + result = [] + for name in queries_to_include: + if name not in _QUERIES: + raise ValueError(f"Unknown query {name!r}. Valid values: {set(_QUERIES)}.") + + result.append( + ( + name, + _QUERIES[name].format( + table=table_ref, date_filter=date_filter, agg_column=agg_column + ), + ) + ) + + return result + + +def _print_results(results: list[BenchmarkResult]) -> None: + console = Console() + query_names = list(dict.fromkeys(r.query for r in results)) + + for qname in query_names: + rows = [r for r in results if r.query == qname] + t = Table(title=f"Query: {qname}", show_lines=True) + t.add_column("Table") + t.add_column("Bytes Processed", justify="right") + t.add_column("Slot ms (avg)", justify="right") + t.add_column("Elapsed avg", justify="right") + t.add_column("Elapsed min", justify="right") + + for r in rows: + t.add_row( + r.label, + _fmt_bytes(r.avg_bytes), + f"{r.avg_slot_ms:,.0f}", + f"{r.avg_elapsed_s:.2f}s", + f"{r.min_elapsed_s:.2f}s", + ) + + if len(rows) == 2: + ref, cmp = rows[0], rows[1] + t.add_row( + f"[bold]ratio ({ref.label}/{cmp.label})[/bold]", + _ratio(ref.avg_bytes, cmp.avg_bytes), + _ratio(ref.avg_slot_ms, cmp.avg_slot_ms), + _ratio(ref.avg_elapsed_s, cmp.avg_elapsed_s), + "—", + style="bold", + ) + + console.print(t) + + +def run( + project: str, + start_date: str, + end_date: str, + tables: dict[str, dict] | None = None, + queries: list[str] | None = None, + agg_column: str | None = None, + runs: int = 1, + location: str = "US", + max_gb_billed: float = 100.0, + bq_client_factory: Callable[..., bigquery.Client] = bigquery.Client, + unknown_unparsed_args: tuple = (), + unknown_parsed_args: dict | None = None, + **kwargs: Any, +) -> None: + """Benchmark query performance across one or more BigQuery tables. + + Runs ``SELECT *`` and (optionally) an aggregation query against each table, + filtered to a single date partition. Prints a comparison of bytes processed, + slot milliseconds, and wall-clock elapsed time from BQ job stats. Query + cache is disabled so each run reflects real scan cost. + + Tables are specified as a dict of dicts — the key is the display label, and + each value must contain: + + - ``table``: fully-qualified BigQuery table (``project.dataset.table``). + - ``column``: column used for date filtering. + - ``type``: ``"native"`` wraps the column in ``DATE()`` (for TIMESTAMP + columns); ``"external"`` compares directly (for hive partition DATE keys). + Defaults to ``"external"``. + + Args: + project: + GCP project used for billing. + + start_date: + Start of date range, inclusive (``YYYY-MM-DD``). + + end_date: + End of date range, exclusive (``YYYY-MM-DD``). + + tables: + Dict of table configurations keyed by display label. See above. + + queries: + List of query names to run. Valid values: ``select_star``, + ``aggregation``. Defaults to all queries. + + agg_column: + Timestamp column name to aggregate by hour, applied to all tables + (e.g. ``tagblock_timestamp``). If omitted, the aggregation query is + skipped. + + runs: + Number of times each query is repeated. Results are averaged. + + location: + BigQuery job location. Defaults to ``US``. + + max_gb_billed: + Safety cap on bytes billed per query (in GB). Defaults to 100 GB. + Pass 0 to disable. + """ + if not tables: + raise ValueError("--tables is required: specify at least one table configuration.") + + queries_to_include = set(queries) if queries else set(_QUERIES) + + if "aggregation" in queries_to_include and not agg_column: + raise ValueError("--agg-column is required when running the aggregation query.") + + max_bytes_billed = int(max_gb_billed * 1024**3) if max_gb_billed else None + client = bq_client_factory(project=project) + _ensure_scratch_dataset(client, location) + + results: list[BenchmarkResult] = [] + + for label, cfg in tables.items(): + table_ref = cfg["table"] + column = cfg["column"] + table_type = cfg.get("type", "external") + + if table_type not in _DATE_FILTERS: + raise ValueError( + f"Unknown table type {table_type!r} for {label!r}. Use 'native' or 'external'." + ) + + date_filter = _DATE_FILTERS[table_type].format( + column=column, start_date=start_date, end_date=end_date + ) + queries_to_run = _build_queries(queries_to_include, table_ref, date_filter, agg_column) + + for query_name, sql in queries_to_run: + result = BenchmarkResult(label=label, query=query_name) + logger.info(f"[{label}:{query_name}] {sql}") + + for i in range(runs): + logger.info(f" run {i + 1}/{runs}") + stats = _run_query( + client, sql, location=location, max_bytes_billed=max_bytes_billed + ) + result.runs.append(stats) + logger.info( + f" → {_fmt_bytes(stats.bytes_processed)}, " + f"{stats.slot_ms:,} slot ms, {stats.elapsed_s:.2f}s" + ) + + results.append(result) + + _print_results(results) diff --git a/src/gfw/ops/pipelines/compact_parquet/main.py b/src/gfw/ops/pipelines/compact_parquet/main.py index aaf28e8..8ed9f49 100644 --- a/src/gfw/ops/pipelines/compact_parquet/main.py +++ b/src/gfw/ops/pipelines/compact_parquet/main.py @@ -5,12 +5,14 @@ import logging import time from dataclasses import dataclass, field + from typing import Any, Callable, Optional -import gcsfs import duckdb +import gcsfs from cloudpathlib import GSPath from google.cloud import storage +from tenacity import Retrying, retry_if_exception_type, stop_after_attempt logger = logging.getLogger(__name__) @@ -29,16 +31,46 @@ def _date_range(start_date: str, end_date: str) -> list[datetime.date]: return dates -def _duckdb_conn(memory_limit: str = "8GB", threads: int = 4) -> duckdb.DuckDBPyConnection: - """Return a DuckDB connection authenticated against GCS via application default credentials.""" - conn = duckdb.connect() - conn.execute(f"SET memory_limit='{memory_limit}'") - conn.execute(f"SET threads={threads}") - conn.execute("SET preserve_insertion_order=false") +def _duckdb_conn(memory_limit: int = 8, threads: int = 4) -> duckdb.DuckDBPyConnection: + conn = duckdb.connect( + config={ + "memory_limit": f"{memory_limit}GB", + "threads": threads, + "preserve_insertion_order": False, + } + ) conn.register_filesystem(gcsfs.GCSFileSystem()) return conn +@dataclass +class CompactionQuery: + """Renders the DuckDB COPY statement for a single compaction job.""" + + source_uris: list[str] + dest_dir: str + target_file_size_mb: int + + _TEMPLATE = ( + "COPY (SELECT * FROM read_parquet({files_sql}))\n" + "TO '{dest_dir}'\n" + "(\n" + " FORMAT PARQUET,\n" + " COMPRESSION SNAPPY,\n" + " FILE_SIZE_BYTES {target_bytes},\n" + " PER_THREAD_OUTPUT false\n" + ")" + ) + + def render(self) -> str: + files_sql = "[" + ", ".join(f"'{u}'" for u in self.source_uris) + "]" + return self._TEMPLATE.format( + files_sql=files_sql, + dest_dir=self.dest_dir, + target_bytes=self.target_file_size_mb * _MB, + ) + + @dataclass class Compactor: """Compacts hive-partitioned Parquet files for a given source and date range. @@ -60,16 +92,36 @@ class Compactor: event_source: str partition_prefix: str target_file_size_mb: int - gcs_staging_path: Optional[GSPath] = None - conn_factory: Callable[[], duckdb.DuckDBPyConnection] = _duckdb_conn + gcs_staging_path: GSPath | None = None + memory_limit_gb: int = 8 + threads: int = 4 + conn_factory: Callable[..., duckdb.DuckDBPyConnection] = _duckdb_conn + max_retries: int = 3 swap: bool = field(init=False) + _connection: Optional[duckdb.DuckDBPyConnection] = field(default=None, init=False, repr=False) + def __post_init__(self): self.swap = False if self.gcs_staging_path is None: self.gcs_staging_path = self.default_staging_path self.swap = True + @property + def connection(self) -> duckdb.DuckDBPyConnection: + if self._connection is None: + self._connection = self.conn_factory( + memory_limit=self.memory_limit_gb, threads=self.threads + ) + return self._connection + + def close_connection(self) -> None: + if self._connection is None: + return + + self._connection.close() + self._connection = None + @property def default_staging_path(self) -> GSPath: return self.gcs_output_path.parent / f"_compact_{self.gcs_output_path.name}_staging" @@ -80,10 +132,30 @@ def run(self, dates: list[datetime.date]) -> None: logger.info(f"Starting compaction: {total} date(s) to process") for i, date in enumerate(dates, 1): logger.info(f"[{i}/{total}] {date}") - self._compact(date) - + self._compact_with_retry(date) + self.close_connection() logger.info(f"Finished compaction: {total} date(s) processed") + def _compact_with_retry(self, date: datetime.date) -> None: + def before_sleep(retry_state: Retrying) -> None: + logger.warning( + f"Attempt {retry_state.attempt_number}/{self.max_retries} failed for {date}: " + f"{retry_state.outcome.exception()}. Retrying..." + ) + self.close_connection() + + for attempt in Retrying( + # InvalidInputException covers Snappy decompression failures from partial GCS reads, + # which DuckDB surfaces as a parse/input error rather than an I/O error. Confirmed + # transient: retrying the same date succeeds. + retry=retry_if_exception_type((duckdb.IOException, duckdb.InvalidInputException)), + stop=stop_after_attempt(self.max_retries), + before_sleep=before_sleep, + reraise=True, + ): + with attempt: + self._compact(date) + def _partition_path(self, base: GSPath, date: datetime.date) -> GSPath: p = self.partition_prefix return base / f"{p}source={self.event_source}" / f"{p}date={date.isoformat()}" @@ -147,24 +219,13 @@ def _write_compacted( """Compact source files into dest_path using DuckDB.""" dest_part = self._partition_path(dest_path, date) dest_dir = f"gs://{dest_part.bucket}/{dest_part.blob}" - files_sql = "[" + ", ".join(f"'{u}'" for u in source_uris) + "]" - target_bytes = self.target_file_size_mb * _MB - - conn = self.conn_factory() - logger.info(f"Writing compacted output to {dest_dir}") - conn.execute( - f""" - COPY (SELECT * FROM read_parquet({files_sql})) - TO '{dest_dir}' - ( - FORMAT PARQUET, - COMPRESSION SNAPPY, - FILE_SIZE_BYTES {target_bytes}, - PER_THREAD_OUTPUT false - ) - """ + query = CompactionQuery( + source_uris=source_uris, + dest_dir=dest_dir, + target_file_size_mb=self.target_file_size_mb, ) - conn.close() + logger.info(f"Writing compacted output to {dest_dir}") + self.connection.execute(query.render()) written = self._list_parquet_blobs(dest_path, date) logger.info(f"Wrote {len(written)} file(s) to {dest_dir}") @@ -195,7 +256,9 @@ def run( end_date: str, partition_prefix: str = "event_", target_file_size_mb: int = 512, - memory_limit: str = "8GB", + memory_limit_gb: int = 8, + threads: int = 4, + max_retries: int = 3, gcs_staging_path: str | None = None, dry_run: bool = False, gcs_client_factory: Callable[[str], storage.Client] = storage.Client, @@ -207,16 +270,16 @@ def run( """Compact small hive-partitioned Parquet files on GCS into larger files. For each date partition, DuckDB reads all source files in parallel and writes - compacted output to a staging path. The staged files are then swapped into the - original partition path so the external table path never changes. If the process - is interrupted after deleting originals, the next run detects the staged files - and resumes the copy step automatically. + compacted output to an auto-generated staging sibling path. The staged files are + then swapped into the original partition path so the external table path never + changes. If the process is interrupted after deleting originals, the next run + detects the staged files and resumes the copy step automatically. Path structure:: {gcs_output_path}/{prefix}source={event_source}/{prefix}date=YYYY-MM-DD/*.parquet - Staging path defaults to ``gs://{bucket}/{parent}/_compact_{table}_staging`` + Staging path: ``gs://{bucket}/{parent}/_compact_{table}_staging`` (sibling of the table directory, same bucket so copies are server-side). Args: @@ -242,6 +305,10 @@ def run( target_file_size_mb: Target size for each output file in MB. Defaults to 512. + max_retries: + Number of retries on transient DuckDB I/O errors before failing the task. + Defaults to 3. + gcs_staging_path: When ``None`` (default), operates in swap mode: compacted files are written to an auto-generated staging sibling path, the originals are deleted, and the @@ -255,15 +322,6 @@ def run( gcs_client_factory: Injectable factory for :class:`~google.cloud.storage.Client`. - memory_limit: - DuckDB memory cap (e.g. ``"8GB"``). DuckDB spills to disk beyond this limit. - Defaults to ``"8GB"``. Set lower on machines with less RAM, higher on GKE pods. - - conn_factory: - Injectable factory for a configured DuckDB connection. Defaults to - :func:`_duckdb_conn` which authenticates via application default credentials. - Override in tests to avoid real GCS calls. - unknown_unparsed_args: Extra unparsed CLI args (ignored). @@ -281,7 +339,6 @@ def run( logger.info(f"[dry-run] {gcs_output_path}/{p}source={event_source}/{p}date={date}") return - resolved_conn_factory = conn_factory or (lambda: _duckdb_conn(memory_limit=memory_limit)) gcs_client = gcs_client_factory(project=project) compactor = Compactor( gcs_client=gcs_client, @@ -290,6 +347,9 @@ def run( partition_prefix=partition_prefix, target_file_size_mb=target_file_size_mb, gcs_staging_path=GSPath(gcs_staging_path.rstrip("/")) if gcs_staging_path else None, - conn_factory=resolved_conn_factory, + memory_limit_gb=memory_limit_gb, + threads=threads, + max_retries=max_retries, + conn_factory=conn_factory or _duckdb_conn, ) compactor.run(dates) diff --git a/tests/cli/commands/test_compact_parquet.py b/tests/cli/commands/test_compact_parquet.py new file mode 100644 index 0000000..5fb1adb --- /dev/null +++ b/tests/cli/commands/test_compact_parquet.py @@ -0,0 +1,28 @@ +import pytest +from unittest.mock import MagicMock + +from gfw.common.cli import CLI + +from gfw.ops.cli.commands.compact_parquet import CompactParquet + + +def test_dry_run(): + CLI(subcommands=[CompactParquet]).execute( + [ + "compact-parquet", + "--project", "proj", + "--gcs-output-path", "gs://bucket/messages", + "--event-source", "src", + "--start-date", "2024-01-01", + "--end-date", "2024-01-02", + "--dry-run", + ], + gcs_client_factory=lambda project: MagicMock(), + conn_factory=MagicMock, + ) + + +def test_help(): + with pytest.raises(SystemExit) as exc: + CLI(subcommands=[CompactParquet]).execute(["compact-parquet", "--help"]) + assert exc.value.code == 0 diff --git a/tests/pipelines/compact_parquet/test_main.py b/tests/pipelines/compact_parquet/test_main.py index 2459c85..4b0d4bf 100644 --- a/tests/pipelines/compact_parquet/test_main.py +++ b/tests/pipelines/compact_parquet/test_main.py @@ -1,12 +1,16 @@ import datetime from unittest.mock import MagicMock, patch -from gfw.ops.pipelines.compact_parquet.main import Compactor, _date_range, run +import duckdb +import pytest from cloudpathlib import GSPath +from gfw.ops.pipelines.compact_parquet.main import Compactor, _date_range, _duckdb_conn, run + DATE = datetime.date(2024, 1, 1) GCS_PATH = GSPath("gs://bucket/path/messages") +COPY_STAGING_PATH = GSPath("gs://bucket/path/compacted_messages") def _make_blob(name: str, size: int = 10 * 1024 * 1024) -> MagicMock: @@ -36,34 +40,59 @@ def _staging_blobs(n: int = 1) -> list[MagicMock]: def _make_compactor( source_blobs=None, - staging_blobs=None, + staging_blobs_sequence=None, conn_factory=None, + gcs_staging_path=None, ) -> Compactor: - """Return a swap-mode Compactor with mocked GCS client and DuckDB connection.""" + """Return a Compactor with mocked GCS client and DuckDB connection factory. + + staging_blobs_sequence is a list of blob lists consumed in order for each + staging-path list_blobs call. Defaults to [] (always returns empty). + Defaults to swap mode. Pass gcs_staging_path to get copy mode. + """ source_blobs = source_blobs if source_blobs is not None else _source_blobs() - staging_blobs = staging_blobs if staging_blobs is not None else [] + staging_iter = iter(staging_blobs_sequence or []) mock_gcs = MagicMock() def list_blobs(bucket, prefix): - if "_staging" in prefix: - return iter(staging_blobs) - return iter(source_blobs) + if prefix.startswith(GCS_PATH.blob): + return iter(source_blobs) + return iter(next(staging_iter, [])) mock_gcs.list_blobs.side_effect = list_blobs - if conn_factory is None: - def conn_factory(): - return MagicMock() - return Compactor( gcs_client=mock_gcs, gcs_output_path=GCS_PATH, event_source="src", partition_prefix="event_", target_file_size_mb=512, - conn_factory=conn_factory, + gcs_staging_path=gcs_staging_path, + conn_factory=conn_factory or MagicMock, + ) + + +# --- _duckdb_conn --- + + +def test_duckdb_conn_configures_connection(): + mock_conn = MagicMock() + with patch( + "gfw.ops.pipelines.compact_parquet.main.duckdb.connect", return_value=mock_conn + ) as mock_connect: + with patch("gfw.ops.pipelines.compact_parquet.main.gcsfs.GCSFileSystem") as mock_fs: + conn = _duckdb_conn(memory_limit=4, threads=2) + + assert conn is mock_conn + mock_connect.assert_called_once_with( + config={ + "memory_limit": "4GB", + "threads": 2, + "preserve_insertion_order": False, + } ) + mock_conn.register_filesystem.assert_called_once_with(mock_fs.return_value) # --- _date_range --- @@ -86,58 +115,47 @@ def test_date_range_empty(): def test_compact_skips_when_no_source_files(): - compactor = _make_compactor(source_blobs=[], staging_blobs=[]) - with patch.object(compactor, "_write_compacted") as mock_write: - compactor._compact(DATE) - mock_write.assert_not_called() + compactor = _make_compactor(source_blobs=[]) + compactor._compact(DATE) # no error = no write attempted def test_compact_skips_when_already_single_file(): compactor = _make_compactor(source_blobs=_source_blobs(1)) - with patch.object(compactor, "_write_compacted") as mock_write: - compactor._compact(DATE) - mock_write.assert_not_called() + compactor._compact(DATE) # no error = no write attempted # --- normal compaction flow --- -def test_compact_normal_flow_order(): +def test_compact_swap_full_flow(): """write staging → delete source → copy staging → delete staging.""" + source = _source_blobs(3) staged = _staging_blobs(1) - compactor = _make_compactor() - calls = [] + # 1st staging call (existing-check): empty; 2nd (post-write list): staged + compactor = _make_compactor(source_blobs=source, staging_blobs_sequence=[[], staged]) - with patch.object( - compactor, "_write_compacted", return_value=staged - ) as mock_write: - with patch.object( - compactor, "_delete_blobs", side_effect=lambda b: calls.append(("delete", b)) - ): - with patch.object( - compactor, "_copy_to_partition", side_effect=lambda d, b, p: calls.append(("copy", b)) - ): - compactor._compact(DATE) + compactor._compact(DATE) - mock_write.assert_called_once() - assert calls[0][0] == "delete" # source deleted first - assert calls[1][0] == "copy" # then staging copied - assert calls[2][0] == "delete" # then staging deleted + for blob in source: + blob.delete.assert_called_once() + staged[0].delete.assert_called_once() + compactor.gcs_client.bucket.return_value.copy_blob.assert_called_once() def test_compact_deletes_source_before_copy(): """Source blobs must be deleted before the copy step.""" source = _source_blobs(3) staged = _staging_blobs(1) - compactor = _make_compactor(source_blobs=source) - deleted = [] + compactor = _make_compactor(source_blobs=source, staging_blobs_sequence=[[], staged]) + ops = [] + source[0].delete.side_effect = lambda: ops.append("delete") + compactor.gcs_client.bucket.return_value.copy_blob.side_effect = lambda *a, **k: ops.append( + "copy" + ) - with patch.object(compactor, "_write_compacted", return_value=staged): - with patch.object(compactor, "_delete_blobs", side_effect=lambda b: deleted.append(b)): - with patch.object(compactor, "_copy_to_partition"): - compactor._compact(DATE) + compactor._compact(DATE) - assert deleted[0] == source + assert ops.index("delete") < ops.index("copy") # --- resume interrupted swap --- @@ -146,14 +164,12 @@ def test_compact_deletes_source_before_copy(): def test_compact_resumes_interrupted_swap(): """If staging exists but source is gone, copy staging and clean up.""" staged = _staging_blobs(1) - compactor = _make_compactor(source_blobs=[], staging_blobs=staged) + compactor = _make_compactor(source_blobs=[], staging_blobs_sequence=[staged]) - with patch.object(compactor, "_copy_to_partition") as mock_copy: - with patch.object(compactor, "_delete_blobs") as mock_delete: - compactor._compact(DATE) + compactor._compact(DATE) - mock_copy.assert_called_once_with(DATE, staged, GCS_PATH) - mock_delete.assert_called_once_with(staged) + compactor.gcs_client.bucket.return_value.copy_blob.assert_called_once() + staged[0].delete.assert_called_once() # --- leftover staging cleanup --- @@ -162,33 +178,29 @@ def test_compact_resumes_interrupted_swap(): def test_compact_cleans_up_leftover_staging_before_writing(): """Leftover staging files from a prior run are deleted before writing new ones.""" leftover = _staging_blobs(2) - staged = _staging_blobs(1) - compactor = _make_compactor(staging_blobs=leftover) - deleted = [] + # 1st staging call: leftover exists; 2nd (post-write): fresh empty result + compactor = _make_compactor(staging_blobs_sequence=[leftover, []]) - with patch.object(compactor, "_write_compacted", return_value=staged): - with patch.object(compactor, "_delete_blobs", side_effect=lambda b: deleted.append(b)): - with patch.object(compactor, "_copy_to_partition"): - compactor._compact(DATE) + compactor._compact(DATE) - assert deleted[0] == leftover # leftover deleted first + for blob in leftover: + blob.delete.assert_called_once() # --- DuckDB SQL generation --- -def test_write_compacted_to_staging_executes_copy_sql(): +def test_write_compacted_executes_copy_sql(): mock_conn = MagicMock() - compactor = _make_compactor(conn_factory=lambda: mock_conn) + compactor = _make_compactor(conn_factory=MagicMock(return_value=mock_conn)) source_uris = [ "gs://bucket/path/messages/event_source=src/event_date=2024-01-01/part_0.parquet" ] - with patch.object(compactor, "_list_parquet_blobs", return_value=_staging_blobs(1)): - compactor._write_compacted(DATE, source_uris, compactor.gcs_staging_path) + compactor._write_compacted(DATE, source_uris, compactor.gcs_staging_path) - executed_sql = mock_conn.execute.call_args[0][0] + executed_sql = compactor.connection.execute.call_args[0][0] assert "COPY" in executed_sql assert "FORMAT PARQUET" in executed_sql assert "COMPRESSION SNAPPY" in executed_sql @@ -196,14 +208,115 @@ def test_write_compacted_to_staging_executes_copy_sql(): assert str(512 * 1024 * 1024) in executed_sql -def test_write_compacted_to_staging_closes_connection(): +def test_run_closes_connection(): mock_conn = MagicMock() - compactor = _make_compactor(conn_factory=lambda: mock_conn) + compactor = _make_compactor(conn_factory=MagicMock(return_value=mock_conn)) + compactor.run([DATE]) + mock_conn.close.assert_called_once() - with patch.object(compactor, "_list_parquet_blobs", return_value=_staging_blobs(1)): - compactor._write_compacted(DATE, ["gs://bucket/x/part.parquet"], compactor.gcs_staging_path) - mock_conn.close.assert_called_once() +def test_close_connection_noop_if_never_opened(): + compactor = _make_compactor() + compactor.close_connection() # no error, no connection created + + +# --- copy mode --- + + +def test_compactor_copy_mode_does_not_set_swap(): + compactor = _make_compactor(gcs_staging_path=COPY_STAGING_PATH) + assert compactor.swap is False + assert compactor.gcs_staging_path == COPY_STAGING_PATH + + +def test_compact_copy_mode_does_not_touch_source(): + """In copy mode, source files are left untouched after compaction.""" + source = _source_blobs(3) + compactor = _make_compactor(source_blobs=source, gcs_staging_path=COPY_STAGING_PATH) + + compactor._compact(DATE) + + for blob in source: + blob.delete.assert_not_called() + compactor.gcs_client.bucket.return_value.copy_blob.assert_not_called() + + +# --- retry --- + + +def test_compact_retries_on_io_error(): + """A transient IOException triggers a retry; connection is reset between attempts.""" + calls = 0 + + def flaky_conn(**kwargs): + conn = MagicMock() + nonlocal calls + calls += 1 + if calls == 1: + conn.execute.side_effect = duckdb.IOException("transient") + return conn + + source = _source_blobs(3) + staged = _staging_blobs(1) + compactor = _make_compactor( + source_blobs=source, + staging_blobs_sequence=[[], [], staged], + conn_factory=flaky_conn, + ) + + compactor._compact_with_retry(DATE) + + assert calls == 2 + + +def test_compact_reraises_after_max_retries(): + """Exhausting all retries re-raises the last exception.""" + conn = MagicMock() + conn.execute.side_effect = duckdb.IOException("persistent") + compactor = _make_compactor( + staging_blobs_sequence=[[], [], [], []], + conn_factory=lambda **kwargs: conn, + ) + compactor.max_retries = 2 + + with pytest.raises(duckdb.IOException): + compactor._compact_with_retry(DATE) + + +# --- Compactor.run() --- + + +def test_run_processes_all_dates(): + compactor = _make_compactor() + compactor.run([datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)]) + + prefixes = [call.kwargs["prefix"] for call in compactor.gcs_client.list_blobs.call_args_list] + assert any("2024-01-01" in p for p in prefixes) + assert any("2024-01-02" in p for p in prefixes) + + +# --- _copy_to_partition and _delete_blobs --- + + +def test_copy_to_partition_calls_gcs_copy(): + compactor = _make_compactor() + blob = _make_blob("path/messages/event_source=src/event_date=2024-01-01/part_0.parquet") + + compactor._copy_to_partition(DATE, [blob], GCS_PATH) + + bucket = compactor.gcs_client.bucket.return_value + expected_dest = "path/messages/event_source=src/event_date=2024-01-01/part_0.parquet" + bucket.copy_blob.assert_called_once_with(blob, bucket, new_name=expected_dest) + + +def test_delete_blobs_deletes_each_blob(): + compactor = _make_compactor() + blobs = [_make_blob(f"path/part_{i}.parquet") for i in range(3)] + + compactor._delete_blobs(blobs) + + for blob in blobs: + blob.delete.assert_called_once() # --- dry run --- @@ -216,8 +329,28 @@ def test_run_dry_run_makes_no_gcs_calls(): gcs_output_path="gs://bucket/messages", event_source="src", start_date="2024-01-01", - end_date="2024-01-03", + end_date="2024-01-02", dry_run=True, gcs_client_factory=lambda project: mock_gcs, ) mock_gcs.list_blobs.assert_not_called() + + +# --- run() non-dry-run --- + + +def test_run_function_runs_compactor(): + mock_gcs = MagicMock() + mock_gcs.list_blobs.return_value = iter([]) + + run( + project="proj", + gcs_output_path="gs://bucket/messages", + event_source="src", + start_date="2024-01-01", + end_date="2024-01-02", + gcs_client_factory=lambda project: mock_gcs, + conn_factory=MagicMock, + ) + + mock_gcs.list_blobs.assert_called() From fa8b88837a13e79675585715c917dd34696d43a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Link?= Date: Mon, 13 Jul 2026 22:12:08 -0300 Subject: [PATCH 3/3] Fix silent data loss in compact-parquet swap on partial delete failure A crash partway through deleting original blobs left the live partition with a partial remnant that the old resume check ("staging present, source empty") couldn't distinguish from a fresh partition, causing the correct staged replacement to be wiped and only the survivors recompacted. Record the exact list of originals to delete in a manifest before starting the delete, so a resumed run can finish from that authoritative list regardless of how far the previous attempt got. Raise instead of guessing when staging is non-empty, source is empty, and no manifest exists, since that state is unverifiable and should not occur post-migration. Co-Authored-By: Claude Sonnet 5 --- src/gfw/ops/pipelines/compact_parquet/main.py | 111 +++++++++++++++--- tests/pipelines/compact_parquet/test_main.py | 110 +++++++++++++++-- 2 files changed, 195 insertions(+), 26 deletions(-) diff --git a/src/gfw/ops/pipelines/compact_parquet/main.py b/src/gfw/ops/pipelines/compact_parquet/main.py index 8ed9f49..1c671c9 100644 --- a/src/gfw/ops/pipelines/compact_parquet/main.py +++ b/src/gfw/ops/pipelines/compact_parquet/main.py @@ -2,15 +2,18 @@ from __future__ import annotations import datetime +import json import logging import time -from dataclasses import dataclass, field +from dataclasses import dataclass, field from typing import Any, Callable, Optional import duckdb import gcsfs + from cloudpathlib import GSPath +from google.api_core.exceptions import NotFound from google.cloud import storage from tenacity import Retrying, retry_if_exception_type, stop_after_attempt @@ -18,6 +21,7 @@ logger = logging.getLogger(__name__) _MB = 1024 * 1024 +_MANIFEST_NAME = "_manifest.json" def _date_range(start_date: str, end_date: str) -> list[datetime.date]: @@ -78,8 +82,11 @@ class Compactor: When ``gcs_staging_path`` is ``None`` (the default), operates in *swap* mode: DuckDB writes compacted output to an auto-generated staging sibling path, the originals are deleted, and the staged files are moved back so the external table path never changes. - If the process is interrupted between the delete and copy steps, the next run detects - the staging files and resumes automatically. + Before deleting any originals, the exact list of blobs to delete is committed to a + manifest file in staging. If the process is interrupted at any point during the delete + (even partway through), the next run finds the manifest and knows staging already holds + a complete, verified replacement — so it finishes the delete from the recorded list and + copies back, rather than re-deriving state from what's currently on disk. When ``gcs_staging_path`` is set, operates in *copy* mode: compacted files are written directly to ``gcs_staging_path`` without touching the source. Use this when you want @@ -166,18 +173,46 @@ def _list_parquet_blobs(self, base: GSPath, date: datetime.date) -> list[storage return [b for b in blobs if b.name.endswith(".parquet")] def _compact(self, date: datetime.date) -> None: + manifest = self._read_manifest(date) if self.swap else None source_blobs = self._list_parquet_blobs(self.gcs_output_path, date) dest_blobs = self._list_parquet_blobs(self.gcs_staging_path, date) + if manifest is not None: + # A manifest exists only once compaction fully succeeded and originals were + # about to be deleted, so staging is guaranteed to hold the complete replacement + # regardless of how far the delete-of-originals got before the interruption. + logger.warning( + f"Resuming committed swap for {date}: finishing delete of " + f"{len(manifest)} recorded original(s) and copy-back from staging" + ) + self._finish_swap(date, manifest, dest_blobs) + return + if self.swap and dest_blobs and not source_blobs: - # Interrupted between delete-originals and copy-from-staging: resume the swap. + # Unreachable under this code's own logic: the manifest is written before any + # delete starts and removed only after the swap fully completes, so a crash + # anywhere in between always leaves a manifest behind. Reaching this state means + # either leftover data from a pre-manifest run, or something outside this code's + # control touched the paths — either way, there's no way to verify dest is the + # complete, correct replacement. Stop and surface it rather than guessing. + raise RuntimeError( + f"Ambiguous state for {date}: {len(dest_blobs)} staged file(s) in " + f"{self.gcs_staging_path}, no source files in {self.gcs_output_path}, and no " + "manifest recording a committed swap. Refusing to auto-recover — verify " + "manually whether the staged files are the complete replacement before " + "deleting or restoring anything." + ) + + if self.swap and dest_blobs and source_blobs: + # Originals are still present and no manifest was recorded, so the + # delete-of-originals step never started. Any staging output here is leftover + # from an interrupted compaction attempt — safe to discard and recompact. logger.warning( - f"Resuming interrupted swap for {date}: " - f"{len(dest_blobs)} staged file(s) found, originals already deleted" + f"Clearing {len(dest_blobs)} stale staged file(s) for {date} " + "(interrupted compaction attempt)" ) - self._copy_to_partition(date, dest_blobs, self.gcs_output_path) self._delete_blobs(dest_blobs) - return + dest_blobs = [] if not source_blobs: logger.info(f"No source files for {date}, skipping") @@ -187,7 +222,7 @@ def _compact(self, date: datetime.date) -> None: logger.info(f"Skipping {date}: already a single file") return - if dest_blobs: + if not self.swap and dest_blobs: logger.info(f"Removing {len(dest_blobs)} existing file(s) from dest for {date}") self._delete_blobs(dest_blobs) @@ -204,13 +239,13 @@ def _compact(self, date: datetime.date) -> None: logger.info(f"DuckDB done in {time.monotonic() - t0:.1f}s — {len(written)} file(s)") if self.swap: - # Originals are deleted before the copy so we never have duplicate data visible - # to the external table. The staging files serve as the recovery point if the - # copy step is interrupted. + # The manifest is written before anything is deleted, so its existence is what + # tells a resumed run that staging already holds a verified, complete replacement + # — independent of how much of the delete loop below actually completed. t0 = time.monotonic() - self._delete_blobs(source_blobs) - self._copy_to_partition(date, written, self.gcs_output_path) - self._delete_blobs(written) + names = [b.name for b in source_blobs] + self._write_manifest(date, names) + self._finish_swap(date, names, written) logger.info(f"Swap done in {time.monotonic() - t0:.1f}s") def _write_compacted( @@ -247,6 +282,52 @@ def _delete_blobs(self, blobs: list[storage.Blob]) -> None: blob.delete() logger.info(f"Deleted gs://{blob.bucket.name}/{blob.name}") + def _delete_named_blobs(self, base: GSPath, names: list[str]) -> None: + bucket = self.gcs_client.bucket(base.bucket) + for name in names: + try: + bucket.blob(name).delete() + logger.info(f"Deleted gs://{base.bucket}/{name}") + except NotFound: + logger.info(f"Already deleted gs://{base.bucket}/{name}") + + def _finish_swap( + self, date: datetime.date, original_names: list[str], staged_blobs: list[storage.Blob] + ) -> None: + """Delete originals, copy staged output into place, and clear the manifest. + + Shared by both a fresh compaction and a manifest-driven resume: every step here + is idempotent (deleting an already-gone original, re-copying an already-copied + file, deleting an already-gone manifest all no-op), so re-running this from the + top after any partial failure is always safe. + """ + self._delete_named_blobs(self.gcs_output_path, original_names) + self._copy_to_partition(date, staged_blobs, self.gcs_output_path) + self._delete_blobs(staged_blobs) + self._delete_manifest(date) + + def _manifest_blob(self, date: datetime.date) -> storage.Blob: + part = self._partition_path(self.gcs_staging_path, date) + bucket = self.gcs_client.bucket(part.bucket) + return bucket.blob(f"{part.blob}/{_MANIFEST_NAME}") + + def _read_manifest(self, date: datetime.date) -> Optional[list[str]]: + blob = self._manifest_blob(date) + if not blob.exists(): + return None + return list(json.loads(blob.download_as_text())) + + def _write_manifest(self, date: datetime.date, names: list[str]) -> None: + blob = self._manifest_blob(date) + blob.upload_from_string(json.dumps(names), content_type="application/json") + + def _delete_manifest(self, date: datetime.date) -> None: + blob = self._manifest_blob(date) + try: + blob.delete() + except NotFound: + pass + def run( project: str, diff --git a/tests/pipelines/compact_parquet/test_main.py b/tests/pipelines/compact_parquet/test_main.py index 4b0d4bf..a678b0c 100644 --- a/tests/pipelines/compact_parquet/test_main.py +++ b/tests/pipelines/compact_parquet/test_main.py @@ -1,8 +1,11 @@ import datetime +import json + from unittest.mock import MagicMock, patch import duckdb import pytest + from cloudpathlib import GSPath from gfw.ops.pipelines.compact_parquet.main import Compactor, _date_range, _duckdb_conn, run @@ -43,12 +46,14 @@ def _make_compactor( staging_blobs_sequence=None, conn_factory=None, gcs_staging_path=None, + manifest=None, ) -> Compactor: """Return a Compactor with mocked GCS client and DuckDB connection factory. staging_blobs_sequence is a list of blob lists consumed in order for each staging-path list_blobs call. Defaults to [] (always returns empty). Defaults to swap mode. Pass gcs_staging_path to get copy mode. + Pass manifest (a list of blob names) to simulate a pre-existing manifest file. """ source_blobs = source_blobs if source_blobs is not None else _source_blobs() staging_iter = iter(staging_blobs_sequence or []) @@ -62,6 +67,11 @@ def list_blobs(bucket, prefix): mock_gcs.list_blobs.side_effect = list_blobs + manifest_blob = mock_gcs.bucket.return_value.blob.return_value + manifest_blob.exists.return_value = manifest is not None + if manifest is not None: + manifest_blob.download_as_text.return_value = json.dumps(manifest) + return Compactor( gcs_client=mock_gcs, gcs_output_path=GCS_PATH, @@ -128,7 +138,7 @@ def test_compact_skips_when_already_single_file(): def test_compact_swap_full_flow(): - """write staging → delete source → copy staging → delete staging.""" + """Write staging → delete source → copy staging → delete staging.""" source = _source_blobs(3) staged = _staging_blobs(1) # 1st staging call (existing-check): empty; 2nd (post-write list): staged @@ -136,10 +146,13 @@ def test_compact_swap_full_flow(): compactor._compact(DATE) - for blob in source: - blob.delete.assert_called_once() + bucket = compactor.gcs_client.bucket.return_value + deleted_names = [ + call.args[0] for call in bucket.blob.call_args_list if not call.args[0].endswith(".json") + ] + assert set(deleted_names) == {b.name for b in source} staged[0].delete.assert_called_once() - compactor.gcs_client.bucket.return_value.copy_blob.assert_called_once() + bucket.copy_blob.assert_called_once() def test_compact_deletes_source_before_copy(): @@ -148,10 +161,9 @@ def test_compact_deletes_source_before_copy(): staged = _staging_blobs(1) compactor = _make_compactor(source_blobs=source, staging_blobs_sequence=[[], staged]) ops = [] - source[0].delete.side_effect = lambda: ops.append("delete") - compactor.gcs_client.bucket.return_value.copy_blob.side_effect = lambda *a, **k: ops.append( - "copy" - ) + bucket = compactor.gcs_client.bucket.return_value + bucket.blob.return_value.delete.side_effect = lambda: ops.append("delete") + bucket.copy_blob.side_effect = lambda *a, **k: ops.append("copy") compactor._compact(DATE) @@ -161,15 +173,90 @@ def test_compact_deletes_source_before_copy(): # --- resume interrupted swap --- -def test_compact_resumes_interrupted_swap(): - """If staging exists but source is gone, copy staging and clean up.""" +def test_compact_raises_on_ambiguous_state_without_manifest(): + """Staging non-empty, source empty, no manifest: unverifiable state — must not + auto-recover. Refuse and surface the problem instead of guessing. + """ staged = _staging_blobs(1) compactor = _make_compactor(source_blobs=[], staging_blobs_sequence=[staged]) + with pytest.raises(RuntimeError, match="Ambiguous state"): + compactor._compact(DATE) + + compactor.gcs_client.bucket.return_value.copy_blob.assert_not_called() + staged[0].delete.assert_not_called() + + +def test_compact_resumes_committed_swap_via_manifest(): + """A manifest takes priority over directory state: finish deleting exactly the + recorded originals and copy back, even though the live partition currently shows + only a partial remnant (2 of 3 originals already deleted before a crash). + """ + all_originals = _source_blobs(3) + manifest_names = [b.name for b in all_originals] + surviving_remnant = all_originals[:1] + staged = _staging_blobs(1) + + compactor = _make_compactor( + source_blobs=surviving_remnant, + staging_blobs_sequence=[staged], + manifest=manifest_names, + ) + compactor._compact(DATE) - compactor.gcs_client.bucket.return_value.copy_blob.assert_called_once() + bucket = compactor.gcs_client.bucket.return_value + deleted_names = [ + call.args[0] for call in bucket.blob.call_args_list if call.args[0] in manifest_names + ] + assert set(deleted_names) == set(manifest_names) + bucket.copy_blob.assert_called_once() staged[0].delete.assert_called_once() + assert compactor._connection is None # never recompacted the surviving remnant + + +def test_write_and_read_manifest_round_trip(): + compactor = _make_compactor() + + compactor._write_manifest(DATE, ["a.parquet", "b.parquet"]) + + bucket = compactor.gcs_client.bucket.return_value + written_json = bucket.blob.return_value.upload_from_string.call_args[0][0] + assert json.loads(written_json) == ["a.parquet", "b.parquet"] + + +def test_read_manifest_returns_none_when_absent(): + compactor = _make_compactor() + assert compactor._read_manifest(DATE) is None + + +def test_read_manifest_returns_recorded_names(): + names = ["a.parquet", "b.parquet"] + compactor = _make_compactor(manifest=names) + assert compactor._read_manifest(DATE) == names + + +def test_delete_named_blobs_tolerates_already_deleted(): + """A blob already removed by a partial prior delete (404 on retry) is a no-op.""" + from google.api_core.exceptions import NotFound + + compactor = _make_compactor() + bucket = compactor.gcs_client.bucket.return_value + bucket.blob.return_value.delete.side_effect = [None, NotFound("gone"), None] + + compactor._delete_named_blobs(GCS_PATH, ["a.parquet", "b.parquet", "c.parquet"]) + + assert bucket.blob.return_value.delete.call_count == 3 + + +def test_delete_manifest_tolerates_already_deleted(): + from google.api_core.exceptions import NotFound + + compactor = _make_compactor() + bucket = compactor.gcs_client.bucket.return_value + bucket.blob.return_value.delete.side_effect = NotFound("gone") + + compactor._delete_manifest(DATE) # no error raised # --- leftover staging cleanup --- @@ -342,6 +429,7 @@ def test_run_dry_run_makes_no_gcs_calls(): def test_run_function_runs_compactor(): mock_gcs = MagicMock() mock_gcs.list_blobs.return_value = iter([]) + mock_gcs.bucket.return_value.blob.return_value.exists.return_value = False run( project="proj",