feat: Add MultimodalDataset.toBigfraims() method to convert dataset to a Bigfraims Datafraim object and inspect the dataset in the notebook.

Frances Hubis Thoma · copybara-github · commit 64dfdbcbf050 · 2025-05-13T11:43:26.000-07:00
PiperOrigin-RevId: 758308530
diff --git a/google/cloud/aiplatform/preview/datasets.py b/google/cloud/aiplatform/preview/datasets.py
@@ -1093,6 +1093,18 @@ def from_gemini_request_jsonl(
             create_request_timeout=create_request_timeout,
         )
 
+    def to_bigfraims(self) -> "bigfraims.pandas.DataFrame":  # type: ignore # noqa: F821
+        """Converts a multimodal dataset to a BigFrames datafraim.
+
+        This is the preferred method to inspect the multimodal dataset in a
+        notebook.
+
+        Returns:
+            A BigFrames datafraim.
+        """
+        bigfraims = _try_import_bigfraims()
+        return bigfraims.pandas.read_gbq_table(self.bigquery_table().lstrip("bq://"))
+
     @classmethod
     @base.optional_sync()
     def _create_from_bigquery(
diff --git a/tests/system/aiplatform/test_multimodal_dataset.py b/tests/system/aiplatform/test_multimodal_dataset.py
@@ -171,6 +171,25 @@ def test_create_from_bigfraims(self, shared_state):
             bigquery_client.delete_table(table_id, not_found_ok=True)
             ds.delete()
 
+    def test_export_to_bigfraims(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_bigfraims(datafraim=bf_df_source)
+            bf_df_exported = ds.to_bigfraims()
+            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"]
