feat: Add MultimodalDataset.toBigframes() method to convert dataset to a Bigframes Dataframe object and inspect the dataset in the notebook.

PiperOrigin-RevId: 758308530

Author: Frances Hubis Thoma

google/cloud/aiplatform/preview/datasets.py

@@ -1093,6 +1093,18 @@ def from_gemini_request_jsonl(
  create_request_timeout=create_request_timeout,
  )
 
+ def to_bigframes(self) -> "bigframes.pandas.DataFrame": # type: ignore # noqa: F821
+ """Converts a multimodal dataset to a BigFrames dataframe.
+
+ This is the preferred method to inspect the multimodal dataset in a
+ notebook.
+
+ Returns:
+ A BigFrames dataframe.
+ """
+ bigframes = _try_import_bigframes()
+ return bigframes.pandas.read_gbq_table(self.bigquery_table().lstrip("bq://"))
+
  @classmethod
  @base.optional_sync()
  def _create_from_bigquery(

tests/system/aiplatform/test_multimodal_dataset.py

@@ -171,6 +171,25 @@ def test_create_from_bigframes(self, shared_state):
  bigquery_client.delete_table(table_id, not_found_ok=True)
  ds.delete()
 
+ def test_export_to_bigframes(self, shared_state):
+ assert shared_state["bigquery_client"]
+ bigquery_client = shared_state["bigquery_client"]
+
+ bpd.options.bigquery.project = _TEST_PROJECT
+ bpd.options.bigquery.location = _TEST_LOCATION
+
+ try:
+ bf_df_source = bpd.DataFrame(_TEST_DATASET)
+ ds = datasets.MultimodalDataset.from_bigframes(dataframe=bf_df_source)
+ bf_df_exported = ds.to_bigframes()
+ table_id = _uri_to_table_id(ds.bigquery_table)
+
+ assert len(bf_df_exported) == len(_TEST_DATASET)
+ assert set(bf_df_exported.columns) == {"Question", "Answer"}
+ finally:
+ bigquery_client.delete_table(table_id, not_found_ok=True)
+ ds.delete()
+
  def test_assemble_dataset(self, shared_state):
  assert shared_state["bigquery_client"]
  assert shared_state["bigquery_test_table"]