Centralize BigQuery Avro loads in BqUtils

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,

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,

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,
+        )

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()