diff --git a/gigl/common/beam/sharded_read.py b/gigl/common/beam/sharded_read.py index affff00fa..e14608f43 100644 --- a/gigl/common/beam/sharded_read.py +++ b/gigl/common/beam/sharded_read.py @@ -4,9 +4,9 @@ from apache_beam.io.gcp.bigquery import BigQueryQueryPriority from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference from apache_beam.pvalue import PBegin -from google.cloud import bigquery from gigl.common.logger import Logger +from gigl.src.common.utils.bq import BqUtils logger = Logger() @@ -31,11 +31,11 @@ def _assert_shard_key_in_table(table_name: str, shard_key: str) -> None: """ Validate that the shard key is a valid column in the BigQuery table. """ - client = bigquery.Client() - table_ref = bigquery.TableReference.from_string(table_name) - - table = client.get_table(table_ref) - column_names = [field.name for field in table.schema] + # BqUtils.fetch_bq_table_schema resolves the table via BigQuery's own path + # parsing, so it accepts the same `project.dataset.table` strings (and, as a + # superset, `project:dataset.table`) as the previous TableReference lookup. + bq_utils = BqUtils() + column_names = list(bq_utils.fetch_bq_table_schema(bq_table=table_name).keys()) if shard_key not in column_names: raise ValueError( diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index e4ebcca2a..874be4b99 100644 --- a/gigl/common/data/export.py +++ b/gigl/common/data/export.py @@ -25,6 +25,7 @@ from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger from gigl.common.utils.retry import retry +from gigl.src.common.utils.bq import BqUtils from gigl.src.common.utils.file_loader import FileLoader logger = Logger() @@ -340,7 +341,6 @@ def add_prediction( self.add_record(batched_records) -# TODO(kmonte): We should migrate this over to `BqUtils.load_files_to_bq` once that is implemented. def _load_records_to_bigquery( gcs_folder: GcsUri, project_id: str, @@ -360,20 +360,18 @@ def _load_records_to_bigquery( schema (Sequence[bigquery.SchemaField]): The BigQuery schema for the records. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. + The BigQuery quota / billing project may be overridden via the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable, which is resolved by + :class:`~gigl.src.common.utils.bq.BqUtils`. + Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows user to monitor and retrieve details about the job status and result. The returned job will be done if `should_run_async=False` and will be returned immediately after creation (not necessarily complete) if - `should_run_asnyc=True`. + `should_run_async=True`. """ - start = time.perf_counter() logger.info(f"Loading records from {gcs_folder} to BigQuery.") - # Initialize the BigQuery client - bigquery_client = bigquery.Client(project=project_id) - - # Construct dataset and table references - dataset_ref = bigquery_client.dataset(dataset_id) - table_ref = dataset_ref.table(table_id) + bq_utils = BqUtils(project=project_id) # Configure the load job job_config = bigquery.LoadJobConfig( @@ -382,24 +380,13 @@ def _load_records_to_bigquery( schema=schema, ) - load_job = bigquery_client.load_table_from_uri( + return bq_utils.load_files_to_bq( source_uris=os.path.join(gcs_folder.uri, "*.avro"), - destination=table_ref, + bq_path=bq_utils.join_path(project_id, dataset_id, table_id), job_config=job_config, + should_run_async=should_run_async, ) - if should_run_async: - logger.info( - f"Started loading process for {dataset_id}:{table_id} with job id {load_job.job_id}, running asynchronously" - ) - else: - load_job.result() # Wait for the job to complete. - logger.info( - f"Loading {load_job.output_rows:,} rows into {dataset_id}:{table_id} in {time.perf_counter() - start:.2f} seconds." - ) - - return load_job - def load_embeddings_to_bigquery( gcs_folder: GcsUri, @@ -426,11 +413,15 @@ def load_embeddings_to_bigquery( table_id (str): The BigQuery table ID. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. + The BigQuery quota / billing project may be overridden via the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable, which is resolved by + :class:`~gigl.src.common.utils.bq.BqUtils`. + Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows user to monitor and retrieve details about the job status and result. The returned job will be done if `should_run_async=False` and will be returned immediately after creation (not necessarily complete) if - `should_run_asnyc=True`. + `should_run_async=True`. """ return _load_records_to_bigquery( gcs_folder, @@ -467,11 +458,15 @@ def load_predictions_to_bigquery( table_id (str): The BigQuery table ID. should_run_async (bool): Whether loading to bigquery step should happen asynchronously. Defaults to False. + The BigQuery quota / billing project may be overridden via the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable, which is resolved by + :class:`~gigl.src.common.utils.bq.BqUtils`. + Returns: LoadJob: A BigQuery LoadJob object representing the load operation, which allows user to monitor and retrieve details about the job status and result. The returned job will be done if `should_run_async=False` and will be returned immediately after creation (not necessarily complete) if - `should_run_asnyc=True`. + `should_run_async=True`. """ return _load_records_to_bigquery( gcs_folder, diff --git a/gigl/env/constants.py b/gigl/env/constants.py index 78a8c5505..7d84ea0de 100644 --- a/gigl/env/constants.py +++ b/gigl/env/constants.py @@ -26,6 +26,13 @@ GIGL_CUDA_DOCKER_URI_ENV_KEY: Final[str] = "GIGL_CUDA_DOCKER_URI" GIGL_COMPONENT_ENV_KEY: Final[str] = "GIGL_COMPONENT" +# Optional override for the quota / billing project used by GiGL's BigQuery +# clients. When set to a non-empty value, every ``BqUtils`` BigQuery client +# bills quota to this project. The override is scoped to BigQuery clients only, +# leaving other Google clients (e.g. Cloud Storage) on their default quota +# project. Unset (or empty) means no override. +GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: Final[str] = "GIGL_BIGQUERY_QUOTA_PROJECT" + _LEGACY_RESOURCE_CONFIG_ENV_KEY: Final[str] = "RESOURCE_CONFIG_PATH" _legacy_resource_config_env_warned: bool = False _logger = Logger() diff --git a/gigl/src/common/utils/bq.py b/gigl/src/common/utils/bq.py index dfdec91b9..ab24d8c36 100644 --- a/gigl/src/common/utils/bq.py +++ b/gigl/src/common/utils/bq.py @@ -1,18 +1,22 @@ import datetime import itertools +import os import re -from typing import Iterable, Optional, Tuple, Union +import time +from typing import Iterable, Optional, Sequence, Tuple, Union import google.api_core.retry import google.cloud.bigquery as bigquery +from google.api_core.client_options import ClientOptions from google.api_core.exceptions import NotFound from google.cloud.bigquery._helpers import _record_field_to_json -from google.cloud.bigquery.job import _AsyncJob +from google.cloud.bigquery.job import LoadJob, _AsyncJob from google.cloud.bigquery.table import RowIterator from gigl.common import GcsUri, LocalUri, Uri from gigl.common.logger import Logger from gigl.common.utils.retry import retry +from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY from gigl.src.common.constants.time import DEFAULT_DATE_FORMAT from gigl.src.common.utils.time import convert_days_to_ms, current_datetime @@ -58,9 +62,44 @@ def _load_file_to_bq_with_retry( class BqUtils: - def __init__(self, project: Optional[str] = None) -> None: - logger.info(f"BqUtils initialized with project: {project}") - self.__bq_client = bigquery.Client(project=project) + def __init__( + self, + project: Optional[str] = None, + quota_project_id: Optional[str] = None, + ) -> None: + """Initializes a BqUtils instance wrapping a single BigQuery client. + + The quota / billing project for the underlying client is resolved as: + the explicit ``quota_project_id`` argument if provided, else the + ``GIGL_BIGQUERY_QUOTA_PROJECT`` environment variable if set to a + non-empty value, else no override (google-auth's default resolution). + + The override scopes only this BigQuery client; other Google clients in + the process are unaffected. + + Args: + project (Optional[str]): The GCP project the client acts on + (default project for datasets, tables, and jobs). + quota_project_id (Optional[str]): The GCP project billed for + BigQuery quota / usage for this client. + """ + # An empty string (from either the argument or the env var) is treated + # as "unset" so callers can opt out with an empty value. + resolved_quota_project_id = quota_project_id or os.environ.get( + GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY + ) + client_options = ( + ClientOptions(quota_project_id=resolved_quota_project_id) + if resolved_quota_project_id + else None + ) + logger.info( + f"BqUtils initialized with project: {project}, " + f"quota project: {resolved_quota_project_id}" + ) + self.__bq_client = bigquery.Client( + project=project, client_options=client_options + ) def create_bq_dataset(self, dataset_id, exists_ok=True) -> None: dataset = bigquery.Dataset(dataset_id) @@ -412,6 +451,53 @@ def fetch_bq_table_schema(self, bq_table: str) -> dict[str, bigquery.SchemaField schema_dict = {field.name: field for field in bq_schema} return schema_dict + def load_files_to_bq( + self, + source_uris: Union[str, Sequence[str]], + bq_path: str, + job_config: bigquery.LoadJobConfig, + should_run_async: bool = False, + ) -> LoadJob: + """Loads one or more GCS files into a BigQuery table with a single load job. + + ``source_uris`` may contain a trailing wildcard (e.g. + ``gs://bucket/folder/*.avro``) to load every matching file. + + Args: + source_uris (Union[str, Sequence[str]]): GCS URI(s), optionally with + a ``*`` wildcard, of the file(s) to load. + bq_path (str): Destination table path in ``project.dataset.table`` + format. + job_config (bigquery.LoadJobConfig): The load job configuration + (source format, write disposition, schema, ...). + should_run_async (bool): If True, return immediately after the load + job is created without waiting for completion. Defaults to False. + + Returns: + LoadJob: The BigQuery load job. Completed if + ``should_run_async=False``; possibly still running if + ``should_run_async=True``. + """ + # Timing is measured around the load job here (rather than at the call + # site) so all callers get consistent load-duration logging. + start = time.perf_counter() + load_job = self.__bq_client.load_table_from_uri( + source_uris=source_uris, + destination=bq_path, + job_config=job_config, + ) + if should_run_async: + logger.info( + f"Started load job {load_job.job_id} for {bq_path}, running asynchronously." + ) + else: + load_job.result() # Waits for the job to complete. + logger.info( + f"Loaded {load_job.output_rows:,} rows into {bq_path} in " + f"{time.perf_counter() - start:.2f} seconds." + ) + return load_job + def load_file_to_bq( self, source_path: Uri, diff --git a/tests/integration/common/beam/__init__.py b/tests/integration/common/beam/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/common/beam/sharded_read_test.py b/tests/integration/common/beam/sharded_read_test.py new file mode 100644 index 000000000..a51069d0b --- /dev/null +++ b/tests/integration/common/beam/sharded_read_test.py @@ -0,0 +1,121 @@ +import uuid + +import apache_beam as beam +from absl.testing import absltest +from apache_beam.options.pipeline_options import GoogleCloudOptions, PipelineOptions +from apache_beam.testing.util import assert_that, equal_to +from google.api_core.exceptions import NotFound + +from gigl.common.beam.sharded_read import ( + BigQueryShardedReadConfig, + ShardedExportRead, +) +from gigl.env.pipelines_config import get_resource_config +from gigl.src.common.utils.bq import BqUtils +from tests.test_assets.test_case import TestCase + +NUM_ROWS = 20 + + +class ShardedReadIntegrationTest(TestCase): + """Integration tests for ShardedExportRead against real BigQuery.""" + + def setUp(self): + resource_config = get_resource_config() + test_unique_name = f"gigl_sharded_read_test_{uuid.uuid4().hex}" + + self.bq_utils = BqUtils(project=resource_config.project) + self.bq_project = resource_config.project + self.bq_dataset = resource_config.temp_assets_bq_dataset_name + self.temp_location = ( + f"{resource_config.temp_assets_regional_bucket_path}/" + f"sharded_read_test/{test_unique_name}" + ) + self.table_path = self.bq_utils.join_path( + self.bq_project, + self.bq_dataset, + test_unique_name, + ) + + table_path = self.bq_utils.format_bq_path(self.table_path) + create_query = f""" + CREATE OR REPLACE TABLE `{table_path}` AS + SELECT + id as user_id, + CONCAT('test_name_', CAST(id AS STRING)) as name + FROM UNNEST(GENERATE_ARRAY(0, {NUM_ROWS - 1})) as id + """ + self.bq_utils.run_query(query=create_query, labels={}) + + def tearDown(self): + self.bq_utils.delete_bq_table_if_exist( + bq_table_path=self.table_path, not_found_ok=True + ) + + def test_reads_all_rows_across_shards(self): + sharded_read_config = BigQueryShardedReadConfig( + shard_key="user_id", + project_id=self.bq_project, + temp_dataset_name=self.bq_dataset, + num_shards=2, + ) + + # Construction validates the shard key against the real table schema. + sharded_read = ShardedExportRead( + table_name=self.table_path, + sharded_read_info=sharded_read_config, + ) + + # ReadFromBigQuery's EXPORT method stages query results in GCS, so the + # DirectRunner pipeline needs a project and temp_location to run against. + options = PipelineOptions() + google_cloud_options = options.view_as(GoogleCloudOptions) + google_cloud_options.project = self.bq_project + google_cloud_options.temp_location = self.temp_location + + # ShardedExportRead flattens every shard, so the read yields one dict per + # row with the table's column names as keys. We assert that all NUM_ROWS + # rows are recovered across the shards. + expected_rows = [ + {"user_id": i, "name": f"test_name_{i}"} for i in range(NUM_ROWS) + ] + with beam.Pipeline(options=options) as pipeline: + rows = pipeline | sharded_read + assert_that(rows, equal_to(expected_rows)) + + def test_raises_when_shard_key_is_not_a_column(self): + sharded_read_config = BigQueryShardedReadConfig( + shard_key="not_a_column", + project_id=self.bq_project, + temp_dataset_name=self.bq_dataset, + num_shards=2, + ) + + with self.assertRaises(ValueError): + ShardedExportRead( + table_name=self.table_path, + sharded_read_info=sharded_read_config, + ) + + def test_raises_when_table_does_not_exist(self): + missing_table_path = self.bq_utils.join_path( + self.bq_project, + self.bq_dataset, + f"gigl_sharded_read_missing_{uuid.uuid4().hex}", + ) + sharded_read_config = BigQueryShardedReadConfig( + shard_key="user_id", + project_id=self.bq_project, + temp_dataset_name=self.bq_dataset, + num_shards=2, + ) + + with self.assertRaises(NotFound): + ShardedExportRead( + table_name=missing_table_path, + sharded_read_info=sharded_read_config, + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/integration/common/data/export_test.py b/tests/integration/common/data/export_test.py index cb54344b5..f83b3fa9e 100644 --- a/tests/integration/common/data/export_test.py +++ b/tests/integration/common/data/export_test.py @@ -5,7 +5,12 @@ from parameterized import param, parameterized from gigl.common import GcsUri -from gigl.common.data.export import EmbeddingExporter, load_embeddings_to_bigquery +from gigl.common.data.export import ( + EmbeddingExporter, + PredictionExporter, + load_embeddings_to_bigquery, + load_predictions_to_bigquery, +) from gigl.common.logger import Logger from gigl.common.utils.gcs import GcsUtils from gigl.env.pipelines_config import get_resource_config @@ -105,3 +110,58 @@ def test_embedding_export(self, _, should_run_async: bool): bq_client.count_number_of_rows_in_bq_table(bq_export_table_path), num_nodes * 2, ) + + +class PredictionExportIntegrationTest(TestCase): + def setUp(self): + resource_config = get_resource_config() + test_unique_name = f"GiGL-Integration-Prediction-Exporter-{uuid.uuid4().hex}" + self.prediction_output_dir = GcsUri.join( + resource_config.temp_assets_regional_bucket_path, + test_unique_name, + "predictions", + ) + self.prediction_output_bq_project = resource_config.project + self.prediction_output_bq_dataset = resource_config.temp_assets_bq_dataset_name + self.prediction_output_bq_table = test_unique_name + + def tearDown(self): + gcs_utils = GcsUtils() + gcs_utils.delete_files_in_bucket_dir(self.prediction_output_dir) + bq_client = BqUtils() + bq_export_table_path = bq_client.join_path( + self.prediction_output_bq_project, + self.prediction_output_bq_dataset, + self.prediction_output_bq_table, + ) + bq_client.delete_bq_table_if_exist( + bq_table_path=bq_export_table_path, + ) + + def test_prediction_export(self): + num_nodes = 100 + with PredictionExporter(export_dir=self.prediction_output_dir) as exporter: + for i in torch.arange(num_nodes): + exporter.add_prediction(torch.tensor([i]), torch.ones(1) * i, "node") + + bq_client = BqUtils() + bq_export_table_path = bq_client.join_path( + self.prediction_output_bq_project, + self.prediction_output_bq_dataset, + self.prediction_output_bq_table, + ) + logger.info( + f"Will try exporting {self.prediction_output_dir} to BQ: {bq_export_table_path}" + ) + load_job = load_predictions_to_bigquery( + gcs_folder=self.prediction_output_dir, + project_id=self.prediction_output_bq_project, + dataset_id=self.prediction_output_bq_dataset, + table_id=self.prediction_output_bq_table, + ) + + self.assertEqual(load_job.output_rows, num_nodes) + self.assertEqual( + bq_client.count_number_of_rows_in_bq_table(bq_export_table_path), + num_nodes, + ) diff --git a/tests/unit/common/data/export_test.py b/tests/unit/common/data/export_test.py index 8f1f300b9..e17441746 100644 --- a/tests/unit/common/data/export_test.py +++ b/tests/unit/common/data/export_test.py @@ -2,7 +2,7 @@ import tempfile from pathlib import Path from typing import Optional -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import fastavro @@ -20,8 +20,6 @@ _PREDICTION_KEY, EmbeddingExporter, PredictionExporter, - load_embeddings_to_bigquery, - load_predictions_to_bigquery, ) from gigl.common.utils.retry import RetriesFailedException from tests.test_assets.test_case import TestCase @@ -382,51 +380,6 @@ def test_skips_flush_if_empty(self, mock_gcs_utils_class): exporter = EmbeddingExporter(export_dir=gcs_base_uri) exporter.flush_records() - @parameterized.expand( - [ - param( - "Test if we can load embeddings synchronously", - should_run_async=False, - ), - param( - "Test if we can load embeddings asynchronously", - should_run_async=True, - ), - ] - ) - @patch("gigl.common.data.export.bigquery.Client") - def test_load_embedding_to_bigquery( - self, _, mock_bigquery_client, should_run_async: bool - ): - # Mock inputs - gcs_folder = GcsUri("gs://test-bucket/test-folder") - project_id = "test-project" - dataset_id = "test-dataset" - table_id = "test-table" - - # Mock BigQuery client and load job - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client - - # Call the function - load_job = load_embeddings_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, - ) - - # Assertions - mock_bigquery_client.assert_called_once_with(project=project_id) - mock_client.load_table_from_uri.assert_called_once_with( - source_uris=f"{gcs_folder.uri}/*.avro", - destination=mock_client.dataset.return_value.table.return_value, - job_config=ANY, - ) - self.assertEqual(load_job.output_rows, 1000) - class TestPredictionsExporter(TestCase): def setUp(self): @@ -734,51 +687,6 @@ def test_skips_flush_if_empty(self, mock_gcs_utils_class): exporter = PredictionExporter(export_dir=gcs_base_uri) exporter.flush_records() - @parameterized.expand( - [ - param( - "Test if we can load predictions synchronously", - should_run_async=False, - ), - param( - "Test if we can load predictions asynchronously", - should_run_async=True, - ), - ] - ) - @patch("gigl.common.data.export.bigquery.Client") - def test_load_prediction_to_bigquery( - self, _, mock_bigquery_client, should_run_async: bool - ): - # Mock inputs - gcs_folder = GcsUri("gs://test-bucket/test-folder") - project_id = "test-project" - dataset_id = "test-dataset" - table_id = "test-table" - - # Mock BigQuery client and load job - mock_client = MagicMock() - mock_client.load_table_from_uri.return_value.output_rows = 1000 - mock_bigquery_client.return_value = mock_client - - # Call the function - load_job = load_predictions_to_bigquery( - gcs_folder, - project_id, - dataset_id, - table_id, - should_run_async=should_run_async, - ) - - # Assertions - mock_bigquery_client.assert_called_once_with(project=project_id) - mock_client.load_table_from_uri.assert_called_once_with( - source_uris=f"{gcs_folder.uri}/*.avro", - destination=mock_client.dataset.return_value.table.return_value, - job_config=ANY, - ) - self.assertEqual(load_job.output_rows, 1000) - if __name__ == "__main__": absltest.main() diff --git a/tests/unit/src/common/utils/bq_test.py b/tests/unit/src/common/utils/bq_test.py index efa02e342..f4c6df008 100644 --- a/tests/unit/src/common/utils/bq_test.py +++ b/tests/unit/src/common/utils/bq_test.py @@ -1,8 +1,10 @@ +import os from typing import Sequence from unittest.mock import MagicMock, patch from parameterized import param, parameterized +from gigl.env.constants import GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY from gigl.src.common.utils.bq import BqUtils from tests.test_assets.test_case import TestCase @@ -141,3 +143,43 @@ def test_returns_latest_table_with_hourly_suffix( bq_table_path_prefix=self.PREFIX, table_partition_suffix="YYYYMMDDHH" ) self.assertEqual(result, "myproject.mydataset.events_2025010112") + + +@patch("gigl.src.common.utils.bq.bigquery.Client") +class BqClientConstructionTest(TestCase): + def test_no_quota_project_when_env_unset_and_no_param( + self, mock_client_cls: MagicMock + ) -> None: + with patch.dict(os.environ, {}, clear=True): + BqUtils(project="my-project") + _, call_kwargs = mock_client_cls.call_args + self.assertEqual(call_kwargs["project"], "my-project") + self.assertIsNone(call_kwargs["client_options"]) + + def test_env_var_sets_quota_project(self, mock_client_cls: MagicMock) -> None: + with patch.dict( + os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: "env-quota-project"} + ): + BqUtils(project="my-project") + _, call_kwargs = mock_client_cls.call_args + self.assertEqual( + call_kwargs["client_options"].quota_project_id, "env-quota-project" + ) + + def test_explicit_param_overrides_env_var(self, mock_client_cls: MagicMock) -> None: + with patch.dict( + os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: "env-quota-project"} + ): + BqUtils(project="my-project", quota_project_id="param-quota-project") + _, call_kwargs = mock_client_cls.call_args + self.assertEqual( + call_kwargs["client_options"].quota_project_id, "param-quota-project" + ) + + def test_empty_env_var_is_treated_as_unset( + self, mock_client_cls: MagicMock + ) -> None: + with patch.dict(os.environ, {GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: ""}): + BqUtils(project="my-project") + _, call_kwargs = mock_client_cls.call_args + self.assertIsNone(call_kwargs["client_options"])