feat: add write_engine parameter to read_FORMATNAME methods to control how data is written to BigQuery

bigframes/dtypes.py

@@ -406,12 +406,19 @@ def arrow_dtype_to_bigframes_dtype(arrow_dtype: pa.DataType) -> Dtype:
  return pd.ArrowDtype(arrow_dtype)
  if pa.types.is_struct(arrow_dtype):
  return pd.ArrowDtype(arrow_dtype)
+
+ # BigFrames doesn't distinguish between string and large_string because the
+ # largest string (2 GB) is already larger than the largest BigQuery row.
+ if pa.types.is_string(arrow_dtype) or pa.types.is_large_string(arrow_dtype):
+ return STRING_DTYPE
+
  if arrow_dtype == pa.null():
  return DEFAULT_DTYPE
- else:
- raise ValueError(
- f"Unexpected Arrow data type {arrow_dtype}. {constants.FEEDBACK_LINK}"
- )
+
+ # No other types matched.
+ raise ValueError(
+ f"Unexpected Arrow data type {arrow_dtype}. {constants.FEEDBACK_LINK}"
+ )
 
 
 _BIGFRAMES_TO_ARROW = {

bigframes/pandas/io/api.py

@@ -30,6 +30,7 @@
  Union,
 )
 
+import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.io.gbq as vendored_pandas_gbq
 from google.cloud import bigquery
 import numpy
@@ -103,6 +104,7 @@ def read_csv(
  Literal["c", "python", "pyarrow", "python-fwf", "bigquery"]
  ] = None,
  encoding: Optional[str] = None,
+ write_engine: constants.WriteEngineType = "default",
  **kwargs,
 ) -> bigframes.dataframe.DataFrame:
  return global_session.with_default_session(
@@ -116,6 +118,7 @@ def read_csv(
  dtype=dtype,
  engine=engine,
  encoding=encoding,
+ write_engine=write_engine,
  **kwargs,
  )
 
@@ -133,6 +136,7 @@ def read_json(
  encoding: Optional[str] = None,
  lines: bool = False,
  engine: Literal["ujson", "pyarrow", "bigquery"] = "ujson",
+ write_engine: constants.WriteEngineType = "default",
  **kwargs,
 ) -> bigframes.dataframe.DataFrame:
  return global_session.with_default_session(
@@ -143,6 +147,7 @@ def read_json(
  encoding=encoding,
  lines=lines,
  engine=engine,
+ write_engine=write_engine,
  **kwargs,
  )
 
@@ -243,24 +248,41 @@ def read_gbq_table(
 
 
 @typing.overload
-def read_pandas(pandas_dataframe: pandas.DataFrame) -> bigframes.dataframe.DataFrame:
+def read_pandas(
+ pandas_dataframe: pandas.DataFrame,
+ *,
+ write_engine: constants.WriteEngineType = "default",
+) -> bigframes.dataframe.DataFrame:
  ...
 
 
 @typing.overload
-def read_pandas(pandas_dataframe: pandas.Series) -> bigframes.series.Series:
+def read_pandas(
+ pandas_dataframe: pandas.Series,
+ *,
+ write_engine: constants.WriteEngineType = "default",
+) -> bigframes.series.Series:
  ...
 
 
 @typing.overload
-def read_pandas(pandas_dataframe: pandas.Index) -> bigframes.core.indexes.Index:
+def read_pandas(
+ pandas_dataframe: pandas.Index,
+ *,
+ write_engine: constants.WriteEngineType = "default",
+) -> bigframes.core.indexes.Index:
  ...
 
 
-def read_pandas(pandas_dataframe: Union[pandas.DataFrame, pandas.Series, pandas.Index]):
+def read_pandas(
+ pandas_dataframe: Union[pandas.DataFrame, pandas.Series, pandas.Index],
+ *,
+ write_engine: constants.WriteEngineType = "default",
+):
  return global_session.with_default_session(
  bigframes.session.Session.read_pandas,
  pandas_dataframe,
+ write_engine=write_engine,
  )
 
 
@@ -271,25 +293,32 @@ def read_pickle(
  filepath_or_buffer: FilePath | ReadPickleBuffer,
  compression: CompressionOptions = "infer",
  storage_options: StorageOptions = None,
+ *,
+ write_engine: constants.WriteEngineType = "default",
 ):
  return global_session.with_default_session(
  bigframes.session.Session.read_pickle,
  filepath_or_buffer=filepath_or_buffer,
  compression=compression,
  storage_options=storage_options,
+ write_engine=write_engine,
  )
 
 
 read_pickle.__doc__ = inspect.getdoc(bigframes.session.Session.read_pickle)
 
 
 def read_parquet(
- path: str | IO["bytes"], *, engine: str = "auto"
+ path: str | IO["bytes"],
+ *,
+ engine: str = "auto",
+ write_engine: constants.WriteEngineType = "default",
 ) -> bigframes.dataframe.DataFrame:
  return global_session.with_default_session(
  bigframes.session.Session.read_parquet,
  path,
  engine=engine,
+ write_engine=write_engine,
  )