mypy

GarrettWu · GarrettWu · commit eba021235afa · 2025-08-27T00:03:48.000Z
diff --git a/bigframes/ml/utils.py b/bigframes/ml/utils.py
@@ -81,7 +81,7 @@ def batch_convert_to_series(
 
 def batch_convert_to_bf_array_type(
     *input: ArrayType, session: Optional[Session] = None
-) -> Generator[bpd.Series, None, None]:
+) -> Generator[Union[bpd.DataFrame, bpd.Series], None, None]:
     """Converts the input to BigFrames DataFrame or Series.
 
     Args:
