diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index e4ebcca2a..ba9af221b 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, @@ -364,16 +364,10 @@ def _load_records_to_bigquery( 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 +376,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, @@ -430,7 +413,7 @@ def load_embeddings_to_bigquery( 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, @@ -471,7 +454,7 @@ def load_predictions_to_bigquery( 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/src/common/utils/bq.py b/gigl/src/common/utils/bq.py index dfdec91b9..c81ae007d 100644 --- a/gigl/src/common/utils/bq.py +++ b/gigl/src/common/utils/bq.py @@ -1,13 +1,14 @@ import datetime import itertools 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.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 @@ -412,6 +413,44 @@ 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: + """Load one or more GCS files into a BigQuery table. + + Args: + source_uris (Union[str, Sequence[str]]): GCS URI or URIs to load. + Wildcards are supported. + bq_path (str): Destination table in ``project.dataset.table`` format. + job_config (bigquery.LoadJobConfig): BigQuery load configuration. + should_run_async (bool): Whether to return before the load finishes. + Defaults to False. + + Returns: + LoadJob: The created BigQuery load job. + """ + start_time = 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() + logger.info( + f"Loaded {load_job.output_rows:,} rows into {bq_path} in " + f"{time.perf_counter() - start_time:.2f} seconds." + ) + return load_job + def load_file_to_bq( self, source_path: Uri, 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()