feat: allow table targets in multi-region datasets when creating multimodal datasets

PiperOrigin-RevId: 744717895

Author: cleop-google

google/cloud/aiplatform/preview/datasets.py

@@ -51,6 +51,7 @@
 _GEMINI_TEMPLATE_CONFIG_FIELD = "geminiTemplateConfig"
 _PROMPT_URI_FIELD = "promptUri"
 _REQUEST_COLUMN_NAME_FIELD = "requestColumnName"
+_BQ_MULTIREGIONS = {"us", "eu"}
 
 _LOGGER = base.Logger(__name__)
 
@@ -107,6 +108,16 @@ def _get_metadata_for_bq(
     return json_format.ParseDict(input_config, struct_pb2.Value())
 
 
+def _bq_dataset_location_allowed(
+    vertex_location: str, bq_dataset_location: str
+) -> bool:
+    if bq_dataset_location == vertex_location:
+        return True
+    if bq_dataset_location in _BQ_MULTIREGIONS:
+        return vertex_location.startswith(bq_dataset_location)
+    return False
+
+
 def _normalize_and_validate_table_id(
     *,
     table_id: str,
@@ -138,7 +149,7 @@ def _normalize_and_validate_table_id(
     )
     client = bigquery.Client(project=project, credentials=credentials)
     bq_dataset = client.get_dataset(dataset_ref=dataset_ref)
-    if bq_dataset.location != vertex_location:
+    if not _bq_dataset_location_allowed(vertex_location, bq_dataset.location):
         raise ValueError(
             f"The BigQuery dataset"
             f" `{dataset_ref.project}.{dataset_ref.dataset_id}` must be in the"

tests/unit/aiplatform/test_multimodal_datasets.py

@@ -198,6 +198,9 @@ def bigframes_import_mock():
     bigframes_module.bigquery = bbq_module
     sys.modules["bigframes"] = bigframes_module
 
+    bigframes_module.BigQueryOptions = mock.MagicMock()
+    bigframes_module.connect = mock.MagicMock()
+
     yield bigframes_module, bpd_module, bbq_module
 
     del sys.modules["bigframes"]
@@ -461,8 +464,6 @@ def test_create_dataset_from_gemini_request_jsonl(
         bpd_module.read_pandas = mock.MagicMock()
         bbq_module.parse_json = lambda x: x
 
-        bf_module.BigQueryOptions = mock.MagicMock()
-        bf_module.connect = mock.MagicMock()
         session_mock = mock.MagicMock()
         bf_module.connect.return_value.__enter__.return_value = session_mock
 
@@ -744,6 +745,53 @@ def test_assemble_request_column_name(self, assemble_data_mock):
         )
         assert result_table_id == _TEST_ASSEMBLE_DATA_BIGQUERY_DESTINATION[5:]
 
+    @pytest.mark.usefixtures("get_dataset_mock")
+    def test_create_dataset_from_pandas_multiregion_target_table_allowed(
+        self, create_dataset_mock, bigframes_import_mock, bq_client_mock
+    ):
+        bq_client_mock.return_value.get_dataset.return_value.location = "us"
+
+        _, bpd_module, _ = bigframes_import_mock
+
+        bpd_module.read_pandas = lambda x: mock.Mock()
+        aiplatform.init(project=_TEST_PROJECT)
+        dataframe = pandas.DataFrame(
+            {
+                "question": ["question"],
+                "answer": ["answer"],
+            }
+        )
+        ummd.MultimodalDataset.from_pandas(
+            dataframe=dataframe,
+            target_table_id=_TEST_TARGET_BQ_TABLE,
+            display_name=_TEST_DISPLAY_NAME,
+            location="us-central1",
+        )
+        create_dataset_mock.assert_called_once()
+
+    def test_create_dataset_from_pandas_multiregion_target_table_location_mismatch_throws_error(
+        self, bigframes_import_mock, bq_client_mock
+    ):
+        bq_client_mock.return_value.get_dataset.return_value.location = "eu"
+
+        _, bpd_module, _ = bigframes_import_mock
+
+        bpd_module.read_pandas = lambda x: mock.Mock()
+        aiplatform.init(project=_TEST_PROJECT)
+        dataframe = pandas.DataFrame(
+            {
+                "question": ["question"],
+                "answer": ["answer"],
+            }
+        )
+        with pytest.raises(ValueError):
+            ummd.MultimodalDataset.from_pandas(
+                dataframe=dataframe,
+                target_table_id=_TEST_TARGET_BQ_TABLE,
+                display_name=_TEST_DISPLAY_NAME,
+                location="us-central1",
+            )
+
 
 class TestGeminiExample:
     """Tests for the GeminiExample class."""