From 1894fc92a6b90cb39579e7b059e1d9c77425e23d Mon Sep 17 00:00:00 2001
From: Christian Leopoldseder <cleop@google.com>
Date: Mon, 13 Apr 2026 07:53:52 -0700
Subject: [PATCH] feat: Add Vertex Dataset input and output options for batch
 jobs

PiperOrigin-RevId: 898998803
---
 vertexai/_genai/types/common.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py
index 97ecc9c3eb..d01c803a81 100644
--- a/vertexai/_genai/types/common.py
+++ b/vertexai/_genai/types/common.py
@@ -13325,6 +13325,12 @@ def to_bigframes(
             raise ValueError("Multimodal dataset bigquery source uri is not set.")
         return bigframes.pandas.read_gbq_table(self.bigquery_uri.removeprefix("bq://"))
 
+    def to_batch_job_source(self) -> BatchJobSource:
+        """Converts the dataset to a BatchJobSource."""
+        return genai_types.BatchJobSource(
+            vertex_dataset_name=self.name,
+        )
+
 
 class MultimodalDatasetDict(TypedDict, total=False):
     """Represents a multimodal dataset."""