From 502569b1d01f4b051dee75ee4b5feb0a99faa256 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Link?= Date: Wed, 5 Aug 2026 15:09:17 -0300 Subject: [PATCH] Avoid writing empty placeholder files for date partitions with no data PartitionedSource addresses each date via BigQuery's partition decorator (table$YYYYMMDD), which returns zero rows rather than raising NotFound when a partition has no data -- unlike ShardedSource, where a missing shard table does raise NotFound. This meant non-sharded sources silently wrote a schema-only placeholder file for every empty date (e.g. before a source's real start date, or any gap day), cluttering the output with tiny files that hurt query performance on the external table. Source stays a pure addressing/naming type: dates_with_data_query() returns the check as data (a query + parameters), not by running it, so it never needs a BQ client. Exporter (which already owns the BQ client) runs the query and filters dates accordingly. There's no safe generic default for this on the base class, same as separator: a source that can't support the check must explicitly return None (a deliberate choice, not an inherited omission), which Exporter treats as "assume every date has data". ShardedSource gets the same proactive check via INFORMATION_SCHEMA.TABLES, instead of relying on catching NotFound from a submitted extract job. NotFound is a generic "resource missing" error -- it fires the same way for a missing shard, a typo'd dataset, or a missing project, so treating it as proof "this date has no data" was unreliable: a misconfigured source could make every date look like a benign gap instead of a hard failure. A truly wrong project/dataset now fails immediately anyway, since Exporter.dates_with_data() runs its check with no try/except, before any job is submitted. Since dates are now confirmed to exist before a job is ever submitted, a NotFound at job-completion time is an anomaly (e.g. a genuine race: the table dropped between the check and the job running), not a routine outcome. It's handled by the same generic except Exception path as any other unexpected failure -- no special case needed, since the outcome (fail hard, surface via run()'s RuntimeError, get retried) is identical either way. ExportStatus. NOT_FOUND and ExportJobResults.skipped are removed as a result: nothing should intentionally produce that status anymore. --- .../ops/pipelines/bq_to_parquet/__init__.py | 1 + .../pipelines/bq_to_parquet/destination.py | 1 + src/gfw/ops/pipelines/bq_to_parquet/job.py | 13 +- src/gfw/ops/pipelines/bq_to_parquet/main.py | 37 +++- src/gfw/ops/pipelines/bq_to_parquet/source.py | 80 ++++++- tests/pipelines/bq_to_parquet/test_main.py | 207 ++++++++++++++++-- 6 files changed, 302 insertions(+), 37 deletions(-) diff --git a/src/gfw/ops/pipelines/bq_to_parquet/__init__.py b/src/gfw/ops/pipelines/bq_to_parquet/__init__.py index 9c4f0db..fb95ddc 100644 --- a/src/gfw/ops/pipelines/bq_to_parquet/__init__.py +++ b/src/gfw/ops/pipelines/bq_to_parquet/__init__.py @@ -1,3 +1,4 @@ from .main import run + __all__ = ["run"] diff --git a/src/gfw/ops/pipelines/bq_to_parquet/destination.py b/src/gfw/ops/pipelines/bq_to_parquet/destination.py index 3485056..b5f83b0 100644 --- a/src/gfw/ops/pipelines/bq_to_parquet/destination.py +++ b/src/gfw/ops/pipelines/bq_to_parquet/destination.py @@ -2,6 +2,7 @@ from __future__ import annotations import datetime + from dataclasses import dataclass, field from cloudpathlib import GSPath diff --git a/src/gfw/ops/pipelines/bq_to_parquet/job.py b/src/gfw/ops/pipelines/bq_to_parquet/job.py index 36e6985..e963355 100644 --- a/src/gfw/ops/pipelines/bq_to_parquet/job.py +++ b/src/gfw/ops/pipelines/bq_to_parquet/job.py @@ -3,10 +3,10 @@ import datetime import logging + from dataclasses import dataclass, field -from enum import auto, StrEnum +from enum import StrEnum, auto -from google.api_core.exceptions import NotFound from google.cloud import bigquery @@ -17,7 +17,6 @@ class ExportStatus(StrEnum): """Outcome of a single BQ extract job.""" SUCCESS = auto() - NOT_FOUND = auto() FAILED = auto() @@ -35,9 +34,6 @@ def wait(self) -> ExportJob: self.job.result() logger.info(f"Completed {self.job.job_id}") self.status = ExportStatus.SUCCESS - except NotFound: - logger.warning(f"Skipped {self.date}: not found") - self.status = ExportStatus.NOT_FOUND except Exception: logger.exception(f"Failed {self.job.job_id} for {self.date}") self.status = ExportStatus.FAILED @@ -61,11 +57,6 @@ def succeeded(self) -> list[ExportJob]: """Jobs that completed successfully.""" return self.filter(ExportStatus.SUCCESS) - @property - def skipped(self) -> list[ExportJob]: - """Jobs skipped because the source table or shard was not found.""" - return self.filter(ExportStatus.NOT_FOUND) - def filter(self, status: ExportStatus) -> list[ExportJob]: """Return jobs matching the given status.""" return [job for job in self.jobs if job.status == status] diff --git a/src/gfw/ops/pipelines/bq_to_parquet/main.py b/src/gfw/ops/pipelines/bq_to_parquet/main.py index 74fc8a3..78c4d78 100644 --- a/src/gfw/ops/pipelines/bq_to_parquet/main.py +++ b/src/gfw/ops/pipelines/bq_to_parquet/main.py @@ -3,6 +3,7 @@ import datetime import logging + from dataclasses import dataclass from typing import Any, Callable @@ -58,7 +59,34 @@ def submit_job(self, date: datetime.date) -> ExportJob: def remaining_dates(self, dates: list[datetime.date]) -> list[datetime.date]: exported = self.destination.existing_dates(dates) - return [d for d in dates if d not in exported] + if exported: + logger.info(f"{len(exported)} date(s) already exported, skipping") + logger.debug(f"Already exported: {sorted(exported)}") + + candidates = [d for d in dates if d not in exported] + if not candidates: + logger.info("No remaining dates to export") + return [] + + with_data = self.dates_with_data(candidates) + for date in candidates: + if date not in with_data: + logger.info(f"Skipping {date}: no data in source") + + return [d for d in candidates if d in with_data] + + def dates_with_data(self, dates: list[datetime.date]) -> set[datetime.date]: + """Return the subset of dates that have data in the source. + + Runs the query the source describes using this exporter's own BQ + client; the source itself performs no I/O. Assumes dates is non-empty. + """ + query = self.source.dates_with_data_query(dates) + job_config = bigquery.QueryJobConfig(query_parameters=query.parameters) + rows = self.bq_client.query(query.sql, job_config=job_config).result() + ids = {row.date_id for row in rows if row.date_id is not None} + + return {d for d in dates if d.strftime("%Y%m%d") in ids} def run( @@ -105,7 +133,8 @@ def run( sharded: Set to ``True`` for date-sharded tables (``table_YYYYMMDD``). Each shard is - addressed directly. Missing shards are skipped with a warning. + addressed directly. Dates with no matching shard are skipped before + submitting any job (see Exporter.dates_with_data()). partition_prefix: Prefix applied to partition key names in the hive output path. @@ -162,6 +191,4 @@ def run( "Retrying this task will skip already-exported dates and resume from the failures." ) - logger.info( - f"Done: {len(results.succeeded)} exported, {len(results.skipped)} skipped (not found)" - ) + logger.info(f"Done: {len(results.succeeded)} exported") diff --git a/src/gfw/ops/pipelines/bq_to_parquet/source.py b/src/gfw/ops/pipelines/bq_to_parquet/source.py index 9fb0ac2..cd5479a 100644 --- a/src/gfw/ops/pipelines/bq_to_parquet/source.py +++ b/src/gfw/ops/pipelines/bq_to_parquet/source.py @@ -2,8 +2,23 @@ from __future__ import annotations import datetime + from dataclasses import dataclass -from typing import ClassVar +from typing import ClassVar, NamedTuple + +from google.cloud import bigquery + + +class DatesWithDataQuery(NamedTuple): + """A metadata-only query whose rows identify dates that have data. + + Executing ``sql`` (with ``parameters`` bound) is expected to return one row + per date that has data, each exposing that date as ``YYYYMMDD`` under a + ``date_id`` column. + """ + + sql: str + parameters: list[bigquery.ScalarQueryParameter] @dataclass(frozen=True) @@ -25,6 +40,19 @@ def ref(self, date: datetime.date) -> str: """Return the BQ table reference for the given date.""" return f"{self.table}{self.separator}{date.strftime('%Y%m%d')}" + def dates_with_data_query(self, dates: list[datetime.date]) -> DatesWithDataQuery: + """Return a query to find which of ``dates`` have data. + + This describes the check as data (SQL + parameters) without touching + BigQuery -- callers (e.g. :class:`~.main.Exporter`, which owns the BQ + client) are responsible for running it. + + Subclasses must implement this, same as ``separator``: there's no + safe generic default, since silently skipping the check would mean + every date is assumed to have data. + """ + raise NotImplementedError + @dataclass(frozen=True) class PartitionedSource(Source): @@ -32,9 +60,59 @@ class PartitionedSource(Source): separator: ClassVar[str] = "$" + def dates_with_data_query(self, dates: list[datetime.date]) -> DatesWithDataQuery: + """Return a query for dates whose partition has at least one row. + + Extracting an empty or non-existent partition via the ``table$YYYYMMDD`` + decorator does not raise an error: BigQuery returns zero rows and still + writes a schema-only output file. Checking partition metadata up front + avoids submitting jobs -- and writing placeholder files -- for dates + with no data. + """ + project, dataset, table_name = self.table.split(".") + sql = f""" + SELECT partition_id AS date_id + FROM `{project}.{dataset}.INFORMATION_SCHEMA.PARTITIONS` + WHERE table_name = @table_name + AND partition_id BETWEEN @min_id AND @max_id + AND total_rows > 0 + """ + parameters = [ + bigquery.ScalarQueryParameter("table_name", "STRING", table_name), + bigquery.ScalarQueryParameter("min_id", "STRING", min(dates).strftime("%Y%m%d")), + bigquery.ScalarQueryParameter("max_id", "STRING", max(dates).strftime("%Y%m%d")), + ] + return DatesWithDataQuery(sql, parameters) + @dataclass(frozen=True) class ShardedSource(Source): """A date-sharded BigQuery table. Addresses each shard as ``table_YYYYMMDD``.""" separator: ClassVar[str] = "_" + + def dates_with_data_query(self, dates: list[datetime.date]) -> DatesWithDataQuery: + """Return a query for dates whose shard table exists. + + Checks table existence directly via ``INFORMATION_SCHEMA.TABLES``, + rather than submitting an extract job and inferring "the shard doesn't + exist" from a ``NotFound`` error: that exception is raised for *any* + missing resource (table, dataset, project), so catching it alone can't + tell a genuinely missing shard apart from, say, a typo'd dataset name + that would make every date look like a benign gap. + """ + project, dataset, table_name = self.table.split(".") + sql = f""" + SELECT SUBSTR(table_name, -8) AS date_id + FROM `{project}.{dataset}.INFORMATION_SCHEMA.TABLES` + WHERE STARTS_WITH(table_name, @table_prefix) + AND SUBSTR(table_name, -8) BETWEEN @min_id AND @max_id + """ + parameters = [ + bigquery.ScalarQueryParameter( + "table_prefix", "STRING", f"{table_name}{self.separator}" + ), + bigquery.ScalarQueryParameter("min_id", "STRING", min(dates).strftime("%Y%m%d")), + bigquery.ScalarQueryParameter("max_id", "STRING", max(dates).strftime("%Y%m%d")), + ] + return DatesWithDataQuery(sql, parameters) diff --git a/tests/pipelines/bq_to_parquet/test_main.py b/tests/pipelines/bq_to_parquet/test_main.py index fb5cfe8..9d77cb8 100644 --- a/tests/pipelines/bq_to_parquet/test_main.py +++ b/tests/pipelines/bq_to_parquet/test_main.py @@ -1,9 +1,11 @@ import datetime -import pytest + from unittest.mock import MagicMock -from google.api_core.exceptions import NotFound +import pytest + from cloudpathlib import GSPath +from google.api_core.exceptions import NotFound from gfw.ops.pipelines.bq_to_parquet.destination import HiveDestination from gfw.ops.pipelines.bq_to_parquet.main import Exporter, _date_range, run @@ -40,6 +42,46 @@ def list_blobs(bucket, prefix, delimiter): return mock_client +def _make_bq_client(extract_job=None, dates_without_data=None): + """Return a mock BQ client. + + extract_job: + Value returned by extract_table(). Defaults to a MagicMock representing + a successful job. + + dates_without_data: + Dates that a source's dates_with_data_query() check should report as + having no data. Defaults to none -- every date in the requested range + is reported as having data, matching pre-existing test expectations + that don't exercise the empty/missing-date check. Applies to whichever + query the source under test builds (Exporter.dates_with_data() reads + min_id/max_id from the query's parameters, shared by both source + types), so this works for both PartitionedSource and ShardedSource. + """ + dates_without_data = dates_without_data or set() + mock_client = MagicMock() + mock_client.extract_table.return_value = extract_job or MagicMock() + + def query(sql, job_config=None): + params = {p.name: p.value for p in job_config.query_parameters} + start = datetime.datetime.strptime(params["min_id"], "%Y%m%d").date() + end = datetime.datetime.strptime(params["max_id"], "%Y%m%d").date() + + rows = [] + current = start + while current <= end: + if current not in dates_without_data: + rows.append(MagicMock(date_id=current.strftime("%Y%m%d"))) + current += datetime.timedelta(days=1) + + mock_query_job = MagicMock() + mock_query_job.result.return_value = rows + return mock_query_job + + mock_client.query.side_effect = query + return mock_client + + def _make_exporter( mock_bq, mock_gcs, bq_in="proj.ds.table", sharded=False, event_source="wf827-table" ): @@ -77,8 +119,94 @@ def test_sharded_source_ref(): assert ShardedSource("proj.ds.table").ref(date) == "proj.ds.table_20240115" -def test_dry_run_does_not_submit_jobs(): +def test_source_dates_with_data_query_is_not_implemented_by_default(): + dates = [datetime.date(2024, 1, 1)] + with pytest.raises(NotImplementedError): + Source("proj.ds.table").dates_with_data_query(dates) + + +def test_partitioned_source_dates_with_data_query(): + dates = [datetime.date(2024, 1, 1), datetime.date(2024, 1, 3)] + + query = PartitionedSource("proj.ds.table").dates_with_data_query(dates) + + assert "INFORMATION_SCHEMA.PARTITIONS" in query.sql + assert "total_rows > 0" in query.sql + params = {p.name: p.value for p in query.parameters} + assert params == { + "table_name": "table", + "min_id": "20240101", + "max_id": "20240103", + } + + +def test_partitioned_source_dates_with_data_query_empty_input(): + """Exporter never calls this with an empty list -- it short-circuits first. + A direct call violating that precondition should fail loudly, not return + an ambiguous None (which means something else: "this source can't check"). + """ + with pytest.raises(ValueError): + PartitionedSource("proj.ds.table").dates_with_data_query([]) + + +def test_sharded_source_dates_with_data_query(): + dates = [datetime.date(2024, 1, 1), datetime.date(2024, 1, 3)] + + query = ShardedSource("proj.ds.table").dates_with_data_query(dates) + + assert "INFORMATION_SCHEMA.TABLES" in query.sql + params = {p.name: p.value for p in query.parameters} + assert params == { + "table_prefix": "table_", + "min_id": "20240101", + "max_id": "20240103", + } + + +def test_sharded_source_dates_with_data_query_empty_input(): + with pytest.raises(ValueError): + ShardedSource("proj.ds.table").dates_with_data_query([]) + + +def test_exporter_dates_with_data_empty_input(): + """remaining_dates() never calls this with an empty list -- it short-circuits + first. A direct call violating that precondition fails loudly instead of + returning an empty set silently. + """ + mock_bq = MagicMock() + exporter = _make_exporter(mock_bq, _make_gcs_client()) + + with pytest.raises(ValueError): + exporter.dates_with_data([]) + mock_bq.query.assert_not_called() + + +def test_run_skips_bq_check_when_all_dates_already_exported(): mock_bq = MagicMock() + already_exported = {datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)} + + results = _make_exporter(mock_bq, _make_gcs_client(already_exported)).run([ + datetime.date(2024, 1, 1), + datetime.date(2024, 1, 2), + ]) + + assert results.succeeded == [] + mock_bq.query.assert_not_called() + mock_bq.extract_table.assert_not_called() + + +def test_exporter_dates_with_data_filters_using_source_query(): + mock_bq = _make_bq_client(dates_without_data={datetime.date(2024, 1, 2)}) + exporter = _make_exporter(mock_bq, _make_gcs_client()) + dates = [datetime.date(2024, 1, 1), datetime.date(2024, 1, 2), datetime.date(2024, 1, 3)] + + assert exporter.dates_with_data(dates) == { + datetime.date(2024, 1, 1), datetime.date(2024, 1, 3) + } + + +def test_dry_run_does_not_submit_jobs(): + mock_bq = _make_bq_client() run( bq_in="proj.ds.table", gcs_out="gs://bucket/path", @@ -95,8 +223,7 @@ def test_dry_run_does_not_submit_jobs(): def test_run_submits_one_job_per_day(): mock_job = MagicMock() - mock_bq = MagicMock() - mock_bq.extract_table.return_value = mock_job + mock_bq = _make_bq_client(extract_job=mock_job) results = _make_exporter(mock_bq, _make_gcs_client()).run([ datetime.date(2024, 1, 1), @@ -110,8 +237,7 @@ def test_run_submits_one_job_per_day(): def test_run_hive_path_includes_source_and_date(): mock_job = MagicMock() - mock_bq = MagicMock() - mock_bq.extract_table.return_value = mock_job + mock_bq = _make_bq_client(extract_job=mock_job) run( bq_in="proj.ds.table", @@ -131,8 +257,7 @@ def test_run_hive_path_includes_source_and_date(): def test_run_skips_already_exported_dates(): mock_job = MagicMock() - mock_bq = MagicMock() - mock_bq.extract_table.return_value = mock_job + mock_bq = _make_bq_client(extract_job=mock_job) already_exported = {datetime.date(2024, 1, 1)} results = _make_exporter(mock_bq, _make_gcs_client(already_exported)).run([ @@ -144,11 +269,43 @@ def test_run_skips_already_exported_dates(): assert len(results.succeeded) == 1 +def test_run_skips_dates_with_no_data_in_source(): + mock_job = MagicMock() + mock_bq = _make_bq_client( + extract_job=mock_job, dates_without_data={datetime.date(2024, 1, 1)} + ) + + results = _make_exporter(mock_bq, _make_gcs_client()).run([ + datetime.date(2024, 1, 1), + datetime.date(2024, 1, 2), + ]) + + assert mock_bq.extract_table.call_count == 1 + assert mock_bq.extract_table.call_args[0][0] == "proj.ds.table$20240102" + assert len(results.succeeded) == 1 + + +def test_run_sharded_skips_dates_with_missing_shard(): + """A missing shard is caught by the proactive check, before any job is submitted.""" + mock_job = MagicMock() + mock_bq = _make_bq_client( + extract_job=mock_job, dates_without_data={datetime.date(2024, 1, 1)} + ) + + results = _make_exporter(mock_bq, _make_gcs_client(), sharded=True).run([ + datetime.date(2024, 1, 1), + datetime.date(2024, 1, 2), + ]) + + assert mock_bq.extract_table.call_count == 1 + assert mock_bq.extract_table.call_args[0][0] == "proj.ds.table_20240102" + assert len(results.succeeded) == 1 + + def test_run_raises_on_job_failure(): mock_job = MagicMock() mock_job.result.side_effect = Exception("BQ backend error") - mock_bq = MagicMock() - mock_bq.extract_table.return_value = mock_job + mock_bq = _make_bq_client(extract_job=mock_job) with pytest.raises(RuntimeError, match="Export failed for 1 date"): run( @@ -163,17 +320,27 @@ def test_run_raises_on_job_failure(): ) -def test_run_sharded_skips_missing(): +def test_run_raises_on_unexpected_not_found(): + """Even though missing shards are normally caught up front, a NotFound that + slips through (e.g. a genuine race: the shard is dropped between the check + and the job running) is treated as a hard failure, not silently skipped -- + it's no longer the expected way to detect a missing date. + """ mock_job = MagicMock() mock_job.result.side_effect = [None, NotFound("not found")] - mock_bq = MagicMock() - mock_bq.extract_table.return_value = mock_job + mock_bq = _make_bq_client(extract_job=mock_job) - results = _make_exporter(mock_bq, _make_gcs_client(), sharded=True).run([ - datetime.date(2024, 1, 1), - datetime.date(2024, 1, 2), - ]) + with pytest.raises(RuntimeError, match="Export failed for 1 date"): + run( + bq_in="proj.ds.table", + gcs_out="gs://bucket/path", + project="proj", + event_source="wf827-table", + start_date="2024-01-01", + end_date="2024-01-03", + sharded=True, + bq_client_factory=_make_bq_factory(mock_bq), + gcs_client_factory=_make_gcs_factory(_make_gcs_client()), + ) assert mock_bq.extract_table.call_count == 2 - assert len(results.succeeded) == 1 - assert len(results.skipped) == 1