feat: Add the possibility to create multimodal datasets without explicitly specifying a bigquery dataset/table.

PiperOrigin-RevId: 744018606

Author: fthoele

google/cloud/aiplatform/preview/datasets.py

@@ -17,6 +17,7 @@
 
 import dataclasses
 from typing import Dict, List, Optional, Tuple
+import uuid
 
 from google.auth import credentials as auth_credentials
 from google.cloud import storage
@@ -41,7 +42,8 @@
 _MULTIMODAL_METADATA_SCHEMA_URI = (
     "gs://google-cloud-aiplatform/schema/dataset/metadata/multimodal_1.0.0.yaml"
 )
-
+_DEFAULT_BQ_DATASET_PREFIX = "vertex_datasets"
+_DEFAULT_BQ_TABLE_PREFIX = "multimodal_dataset"
 _INPUT_CONFIG_FIELD = "inputConfig"
 _BIGQUERY_SOURCE_FIELD = "bigquerySource"
 _URI_FIELD = "uri"
@@ -147,6 +149,37 @@ def _normalize_and_validate_table_id(
     return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}"
 
 
+def _create_default_bigquery_dataset_if_not_exists(
+    *,
+    project: Optional[str] = None,
+    location: Optional[str] = None,
+    credentials: Optional[auth_credentials.Credentials] = None,
+) -> str:
+    # Loading bigquery lazily to avoid auto-loading it when importing vertexai
+    from google.cloud import bigquery  # pylint: disable=g-import-not-at-top
+
+    if not project:
+        project = initializer.global_config.project
+    if not location:
+        location = initializer.global_config.location
+    if not credentials:
+        credentials = initializer.global_config.credentials
+
+    bigquery_client = bigquery.Client(project=project, credentials=credentials)
+    location_str = location.lower().replace("-", "_")
+    dataset_id = bigquery.DatasetReference(
+        project, f"{_DEFAULT_BQ_DATASET_PREFIX}_{location_str}"
+    )
+    dataset = bigquery.Dataset(dataset_ref=dataset_id)
+    dataset.location = location
+    bigquery_client.create_dataset(dataset, exists_ok=True)
+    return f"{dataset_id.project}.{dataset_id.dataset_id}"
+
+
+def _generate_target_table_id(dataset_id: str):
+    return f"{dataset_id}.{_DEFAULT_BQ_TABLE_PREFIX}_{str(uuid.uuid4())}"
+
+
 class GeminiExample:
     """A class representing a Gemini example."""
 
@@ -610,7 +643,7 @@ def from_pandas(
         cls,
         *,
         dataframe: pandas.DataFrame,
-        target_table_id: str,
+        target_table_id: Optional[str] = None,
         display_name: Optional[str] = None,
         project: Optional[str] = None,
         location: Optional[str] = None,
@@ -625,12 +658,14 @@ def from_pandas(
             dataframe (pandas.DataFrame):
                 The pandas dataframe to be used for the created dataset.
             target_table_id (str):
-                The BigQuery table id where the dataframe will be uploaded. The
-                table id can be in the format of "dataset.table" or
-                "project.dataset.table". If a table already exists with the
+                Optional. The BigQuery table id where the dataframe will be
+                uploaded. The table id can be in the format of "dataset.table"
+                or "project.dataset.table". If a table already exists with the
                 given table id, it will be overwritten. Note that the BigQuery
                 dataset must already exist and be in the same location as the
-                multimodal dataset.
+                multimodal dataset. If not provided, a generated table id will
+                be created in the `vertex_datasets` dataset (e.g.
+                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
             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
@@ -667,21 +702,43 @@ def from_pandas(
                 The created multimodal dataset.
         """
         bigframes = _try_import_bigframes()
-        # TODO(b/400355374): `table_id` should be optional, and if not provided,
-        # we generate a random table id. Also, check if we can use a default
-        # dataset that's created from the SDK.
-        target_table_id = _normalize_and_validate_table_id(
-            table_id=target_table_id,
-            project=project,
-            vertex_location=location,
+        from google.cloud import bigquery  # pylint: disable=g-import-not-at-top
+
+        if not project:
+            project = initializer.global_config.project
+        if not location:
+            location = initializer.global_config.location
+        if not credentials:
+            credentials = initializer.global_config.credentials
+
+        if target_table_id:
+            target_table_id = _normalize_and_validate_table_id(
+                table_id=target_table_id,
+                project=project,
+                vertex_location=location,
+                credentials=credentials,
+            )
+        else:
+            dataset_id = _create_default_bigquery_dataset_if_not_exists(
+                project=project, location=location, credentials=credentials
+            )
+            target_table_id = _generate_target_table_id(dataset_id)
+
+        session_options = bigframes.BigQueryOptions(
             credentials=credentials,
+            project=project,
+            location=location,
         )
-
-        temp_bigframes_df = bigframes.pandas.read_pandas(dataframe)
-        temp_bigframes_df.to_gbq(
-            destination_table=target_table_id,
-            if_exists="replace",
+        with bigframes.connect(session_options) as session:
+            temp_bigframes_df = session.read_pandas(dataframe)
+            temp_table_id = temp_bigframes_df.to_gbq()
+        client = bigquery.Client(project=project, credentials=credentials)
+        copy_job = client.copy_table(
+            sources=temp_table_id,
+            destination=target_table_id,
         )
+        copy_job.result()
+
         bigquery_uri = f"bq://{target_table_id}"
         return cls._create_from_bigquery(
             bigquery_uri=bigquery_uri,
@@ -700,7 +757,7 @@ def from_bigframes(
         cls,
         *,
         dataframe: "bigframes.pandas.DataFrame",  # type: ignore # noqa: F821
-        target_table_id: str,
+        target_table_id: Optional[str] = None,
         display_name: Optional[str] = None,
         project: Optional[str] = None,
         location: Optional[str] = None,
@@ -716,12 +773,14 @@ def from_bigframes(
                 The BigFrames dataframe that will be used for the created
                 dataset.
             target_table_id (str):
-                The BigQuery table id where the dataframe will be uploaded. The
-                table id can be in the format of "dataset.table" or
-                "project.dataset.table". If a table already exists with the
+                Optional. The BigQuery table id where the dataframe will be
+                uploaded. The table id can be in the format of "dataset.table"
+                or "project.dataset.table". If a table already exists with the
                 given table id, it will be overwritten. Note that the BigQuery
                 dataset must already exist and be in the same location as the
-                multimodal dataset.
+                multimodal dataset. If not provided, a generated table id will
+                be created in the `vertex_datasets` dataset (e.g.
+                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
             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
@@ -756,19 +815,32 @@ def from_bigframes(
         Returns:
             The created multimodal dataset.
         """
-        # TODO(b/400355374): `table_id` should be optional, and if not provided,
-        # we generate a random table id. Also, check if we can use a default
-        # dataset that's created from the SDK.
-        target_table_id = _normalize_and_validate_table_id(
-            table_id=target_table_id,
-            project=project,
-            vertex_location=location,
-            credentials=credentials,
-        )
-        dataframe.to_gbq(
-            destination_table=target_table_id,
-            if_exists="replace",
+        from google.cloud import bigquery  # pylint: disable=g-import-not-at-top
+
+        if target_table_id:
+            target_table_id = _normalize_and_validate_table_id(
+                table_id=target_table_id,
+                project=project,
+                vertex_location=location,
+                credentials=credentials,
+            )
+        else:
+            dataset_id = _create_default_bigquery_dataset_if_not_exists(
+                project=project, location=location, credentials=credentials
+            )
+            target_table_id = _generate_target_table_id(dataset_id)
+
+        if not project:
+            project = initializer.global_config.project
+
+        temp_table_id = dataframe.to_gbq()
+        client = bigquery.Client(project=project, credentials=credentials)
+        copy_job = client.copy_table(
+            sources=temp_table_id,
+            destination=target_table_id,
         )
+        copy_job.result()
+
         bigquery_uri = f"bq://{target_table_id}"
         return cls._create_from_bigquery(
             bigquery_uri=bigquery_uri,
@@ -787,7 +859,7 @@ def from_gemini_request_jsonl(
         cls,
         *,
         gcs_uri: str,
-        target_table_id: str,
+        target_table_id: Optional[str] = None,
         display_name: Optional[str] = None,
         project: Optional[str] = None,
         location: Optional[str] = None,
@@ -808,11 +880,14 @@ def from_gemini_request_jsonl(
                 The Google Cloud Storage URI of the JSONL file to import.
                 For example, 'gs://my-bucket/path/to/data.jsonl'
             target_table_id (str):
-                The BigQuery table id where the dataframe will be uploaded. The
-                table id can be in the format of "dataset.table" or
-                "project.dataset.table". If a table already exists with the
+                Optional. The BigQuery table id where the dataframe will be
+                uploaded. The table id can be in the format of "dataset.table"
+                or "project.dataset.table". If a table already exists with the
                 given table id, it will be overwritten. Note that the BigQuery
-                dataset must already exist.
+                dataset must already exist and be in the same location as the
+                multimodal dataset. If not provided, a generated table id will
+                be created in the `vertex_datasets` dataset (e.g.
+                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
             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
@@ -848,14 +923,23 @@ def from_gemini_request_jsonl(
             The created multimodal dataset.
         """
         bigframes = _try_import_bigframes()
+        from google.cloud import bigquery  # pylint: disable=g-import-not-at-top
+
         if not project:
             project = initializer.global_config.project
-        # TODO(b/400355374): `table_id` should be optional, and if not provided,
-        # we generate a random table id. Also, check if we can use a default
-        # dataset that's created from the SDK.
-        target_table_id = _normalize_and_validate_table_id(
-            table_id=target_table_id, project=project
-        )
+
+        if target_table_id:
+            target_table_id = _normalize_and_validate_table_id(
+                table_id=target_table_id,
+                project=project,
+                vertex_location=location,
+                credentials=credentials,
+            )
+        else:
+            dataset_id = _create_default_bigquery_dataset_if_not_exists(
+                project=project, location=location, credentials=credentials
+            )
+            target_table_id = _generate_target_table_id(dataset_id)
 
         gcs_uri_prefix = "gs://"
         if gcs_uri.startswith(gcs_uri_prefix):
@@ -877,13 +961,21 @@ def from_gemini_request_jsonl(
         lines = [line.strip() for line in jsonl_string.splitlines() if line.strip()]
         df = pandas.DataFrame(lines, columns=[request_column_name])
 
-        temp_bigframes_df = bigframes.pandas.read_pandas(df)
-        temp_bigframes_df[request_column_name] = bigframes.bigquery.parse_json(
-            temp_bigframes_df[request_column_name]
+        session_options = bigframes.BigQueryOptions(
+            credentials=credentials,
+            project=project,
+            location=location,
         )
-        temp_bigframes_df.to_gbq(
-            destination_table=target_table_id,
-            if_exists="replace",
+        with bigframes.connect(session_options) as session:
+            temp_bigframes_df = session.read_pandas(df)
+            temp_bigframes_df[request_column_name] = bigframes.bigquery.parse_json(
+                temp_bigframes_df[request_column_name]
+            )
+            temp_table_id = temp_bigframes_df.to_gbq()
+        client = bigquery.Client(project=project, credentials=credentials)
+        client.copy_table(
+            sources=temp_table_id,
+            destination=target_table_id,
         )
 
         bigquery_uri = f"bq://{target_table_id}"