feat: ummd.MultimodalDataset.from_bigquery() now also accepts a table id (not just a BQ table URI).

PiperOrigin-RevId: 781393598

Author: Frances Hubis Thoma

google/cloud/aiplatform/preview/datasets.py

@@ -690,7 +690,7 @@ def bigquery_table(self) -> str:
  def from_bigquery(
  cls,
  *,
- bigquery_uri: str,
+ bigquery_source: str,
  display_name: Optional[str] = None,
  project: Optional[str] = None,
  location: Optional[str] = None,
@@ -702,10 +702,10 @@ def from_bigquery(
  """Creates a multimodal dataset from a BigQuery table.
 
  Args:
- bigquery_uri (str):
- Required. The BigQuery table URI to be used for the created
- dataset. The table uri can be in the format of
- "bq://dataset.table" or "bq://project.dataset.table".
+ bigquery_source (str):
+ Required. The BigQuery table URI or ID to be used for the created
+ dataset, which can be in the format of "bq://dataset.table",
+ "bq://project.dataset.table" or "project.dataset.table".
  display_name (str):
  Optional. The user-defined name of the dataset. The name can be
  up to 128 characters long and can consist of any UTF-8
@@ -741,9 +741,10 @@ def from_bigquery(
  dataset (MultimodalDataset):
  The created multimodal dataset.
  """
+ if not bigquery_source.startswith("bq://"):
+ bigquery_source = f"bq://{bigquery_source}"
  return cls._create_from_bigquery(
- bigquery_uri=bigquery_uri,
- metadata=_get_metadata_for_bq(bq_uri=bigquery_uri),
+ metadata=_get_metadata_for_bq(bq_uri=bigquery_source),
  display_name=display_name,
  project=project,
  location=location,
@@ -856,7 +857,6 @@ def from_pandas(
 
  bigquery_uri = f"bq://{target_table_id}"
  return cls._create_from_bigquery(
- bigquery_uri=bigquery_uri,
  metadata=_get_metadata_for_bq(bq_uri=bigquery_uri),
  display_name=display_name,
  project=project,
@@ -958,7 +958,6 @@ def from_bigframes(
 
  bigquery_uri = f"bq://{target_table_id}"
  return cls._create_from_bigquery(
- bigquery_uri=bigquery_uri,
  metadata=_get_metadata_for_bq(bq_uri=bigquery_uri),
  display_name=display_name,
  project=project,
@@ -1095,7 +1094,6 @@ def from_gemini_request_jsonl(
 
  bigquery_uri = f"bq://{target_table_id}"
  return cls._create_from_bigquery(
- bigquery_uri=bigquery_uri,
  metadata=_get_metadata_for_bq(
  bq_uri=bigquery_uri, request_column_name=request_column_name
  ),
@@ -1125,7 +1123,6 @@ def to_bigframes(self) -> "bigframes.pandas.DataFrame": # type: ignore # noqa:
  def _create_from_bigquery(
  cls,
  *,
- bigquery_uri: str,
  metadata: struct_pb2.Value,
  display_name: Optional[str] = None,
  project: Optional[str] = None,

tests/system/aiplatform/test_multimodal_dataset.py

@@ -105,14 +105,30 @@ def copy_sample_data(self, shared_state):
  # to clean up the table.
  yield
 
- def test_create_new_dataset(self, shared_state):
+ def test_create_new_dataset_from_bq_table_uri(self, shared_state):
  assert shared_state["bigquery_test_table"]
  bigquery_table = f"bq://{shared_state['bigquery_test_table']}"
  display_name = "test dataset"
  labels = {"label1": "value1", "label2": "value2"}
  try:
  ds = datasets.MultimodalDataset.from_bigquery(
- bigquery_uri=bigquery_table, display_name=display_name, labels=labels
+ bigquery_source=bigquery_table, display_name=display_name, labels=labels
+ )
+ assert ds.display_name == display_name
+ assert ds.bigquery_table == bigquery_table
+ assert ds.labels == labels
+ assert ds.location == _TEST_LOCATION
+ finally:
+ ds.delete()
+
+ def test_create_new_dataset_from_bq_table_id(self, shared_state):
+ assert shared_state["bigquery_test_table"]
+ bigquery_table = f"{shared_state['bigquery_test_table']}"
+ display_name = "test dataset"
+ labels = {"label1": "value1", "label2": "value2"}
+ try:
+ ds = datasets.MultimodalDataset.from_bigquery(
+ bigquery_source=bigquery_table, display_name=display_name, labels=labels
  )
  assert ds.display_name == display_name
  assert ds.bigquery_table == bigquery_table

tests/unit/aiplatform/test_multimodal_datasets.py

@@ -438,7 +438,7 @@ def test_dataset_bigquery_table(self):
  def test_create_dataset_from_bigquery(self, create_dataset_mock, sync):
  aiplatform.init(project=_TEST_PROJECT)
  new_dataset = ummd.MultimodalDataset.from_bigquery(
- bigquery_uri=_TEST_SOURCE_URI_BQ,
+ bigquery_source=_TEST_SOURCE_URI_BQ,
  display_name=_TEST_DISPLAY_NAME,
  sync=sync,
  )