feat: add bigframe.bigquery.load_data function (#2426)

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
Fixes #<issue_number_goes_here> 🦕

Author: GarrettWu

bigframes/bigquery/__init__.py

@@ -43,6 +43,7 @@
 st_regionstats,
 st_simplify,
 )
+frombigframes.bigquery._operations.ioimportload_data
 frombigframes.bigquery._operations.jsonimport (
 json_extract,
 json_extract_array,
@@ -107,6 +108,8 @@
 struct,
 # table ops
 create_external_table,
+# io ops
+load_data,
 ]
 
 _module=sys.modules[__name__]
@@ -160,6 +163,8 @@
 "struct",
 # table ops
 "create_external_table",
+# io ops
+"load_data",
 # Modules / SQL namespaces
 "ai",
 "ml",

bigframes/bigquery/_operations/io.py

@@ -0,0 +1,94 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ importannotations
+
+fromtypingimportMapping, Optional, Union
+
+importpandasaspd
+
+frombigframes.bigquery._operations.tableimport_get_table_metadata
+importbigframes.core.logging.log_adapteraslog_adapter
+importbigframes.core.sql.io
+importbigframes.session
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_io")
+defload_data(
+table_name: str,
+*,
+write_disposition: str="INTO",
+columns: Optional[Mapping[str, str]] =None,
+partition_by: Optional[list[str]] =None,
+cluster_by: Optional[list[str]] =None,
+table_options: Optional[Mapping[str, Union[str, int, float, bool, list]]] =None,
+from_files_options: Mapping[str, Union[str, int, float, bool, list]],
+with_partition_columns: Optional[Mapping[str, str]] =None,
+connection_name: Optional[str] =None,
+session: Optional[bigframes.session.Session] =None,
+) ->pd.Series:
+"""
+ Loads data into a BigQuery table.
+ See the `BigQuery LOAD DATA DDL syntax
+ <https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/load-statements>`_
+ for additional reference.
+ Args:
+ table_name (str):
+ The name of the table in BigQuery.
+ write_disposition (str, default "INTO"):
+ Whether to replace the table if it already exists ("OVERWRITE") or append to it ("INTO").
+ columns (Mapping[str, str], optional):
+ The table's schema.
+ partition_by (list[str], optional):
+ A list of partition expressions to partition the table by. See https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/load-statements#partition_expression.
+ cluster_by (list[str], optional):
+ A list of columns to cluster the table by.
+ table_options (Mapping[str, Union[str, int, float, bool, list]], optional):
+ The table options.
+ from_files_options (Mapping[str, Union[str, int, float, bool, list]]):
+ The options for loading data from files.
+ with_partition_columns (Mapping[str, str], optional):
+ The table's partition columns.
+ connection_name (str, optional):
+ The connection to use for the table.
+ session (bigframes.session.Session, optional):
+ The session to use. If not provided, the default session is used.
+ Returns:
+ pandas.Series:
+ A Series with object dtype containing the table metadata. Reference
+ the `BigQuery Table REST API reference
+ <https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table>`_
+ for available fields.
+ """
+importbigframes.pandasasbpd
+
+sql=bigframes.core.sql.io.load_data_ddl(
+table_name=table_name,
+write_disposition=write_disposition,
+columns=columns,
+partition_by=partition_by,
+cluster_by=cluster_by,
+table_options=table_options,
+from_files_options=from_files_options,
+with_partition_columns=with_partition_columns,
+connection_name=connection_name,
+ )
+
+ifsessionisNone:
+bpd.read_gbq_query(sql)
+session=bpd.get_global_session()
+else:
+session.read_gbq_query(sql)
+
+return_get_table_metadata(bqclient=session.bqclient, table_name=table_name)

bigframes/core/sql/io.py

@@ -0,0 +1,87 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ importannotations
+
+fromtypingimportMapping, Optional, Union
+
+
+defload_data_ddl(
+table_name: str,
+*,
+write_disposition: str="INTO",
+columns: Optional[Mapping[str, str]] =None,
+partition_by: Optional[list[str]] =None,
+cluster_by: Optional[list[str]] =None,
+table_options: Optional[Mapping[str, Union[str, int, float, bool, list]]] =None,
+from_files_options: Mapping[str, Union[str, int, float, bool, list]],
+with_partition_columns: Optional[Mapping[str, str]] =None,
+connection_name: Optional[str] =None,
+) ->str:
+"""Generates the LOAD DATA DDL statement."""
+statement= ["LOAD DATA"]
+statement.append(write_disposition)
+statement.append(table_name)
+
+ifcolumns:
+column_defs=", ".join([f"{name}{typ}"forname, typincolumns.items()])
+statement.append(f"({column_defs})")
+
+ifpartition_by:
+statement.append(f"PARTITION BY {', '.join(partition_by)}")
+
+ifcluster_by:
+statement.append(f"CLUSTER BY {', '.join(cluster_by)}")
+
+iftable_options:
+opts= []
+forkey, valueintable_options.items():
+ifisinstance(value, str):
+value_sql=repr(value)
+opts.append(f"{key} = {value_sql}")
+elifisinstance(value, bool):
+opts.append(f"{key} = {str(value).upper()}")
+elifisinstance(value, list):
+list_str=", ".join([repr(v) forvinvalue])
+opts.append(f"{key} = [{list_str}]")
+else:
+opts.append(f"{key} = {value}")
+options_str=", ".join(opts)
+statement.append(f"OPTIONS ({options_str})")
+
+opts= []
+forkey, valueinfrom_files_options.items():
+ifisinstance(value, str):
+value_sql=repr(value)
+opts.append(f"{key} = {value_sql}")
+elifisinstance(value, bool):
+opts.append(f"{key} = {str(value).upper()}")
+elifisinstance(value, list):
+list_str=", ".join([repr(v) forvinvalue])
+opts.append(f"{key} = [{list_str}]")
+else:
+opts.append(f"{key} = {value}")
+options_str=", ".join(opts)
+statement.append(f"FROM FILES ({options_str})")
+
+ifwith_partition_columns:
+part_defs=", ".join(
+ [f"{name}{typ}"forname, typinwith_partition_columns.items()]
+ )
+statement.append(f"WITH PARTITION COLUMNS ({part_defs})")
+
+ifconnection_name:
+statement.append(f"WITH CONNECTION `{connection_name}`")
+
+return" ".join(statement)

tests/system/large/bigquery/test_io.py

@@ -0,0 +1,39 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for for the specific language governing permissions and
+# limitations under the License.
+
+importbigframes.bigqueryasbbq
+
+
+deftest_load_data(session, dataset_id):
+table_name=f"{dataset_id}.test_load_data"
+uri="gs://cloud-samples-data/bigquery/us-states/us-states.csv"
+
+# Create the external table
+table=bbq.load_data(
+table_name,
+columns={
+"name": "STRING",
+"post_abbr": "STRING",
+ },
+from_files_options={"format": "CSV", "uris": [uri], "skip_leading_rows": 1},
+session=session,
+ )
+asserttableisnotNone
+
+# Read the table to verify
+importbigframes.pandasasbpd
+
+bf_df=bpd.read_gbq(table_name)
+pd_df=bf_df.to_pandas()
+assertlen(pd_df) >0

tests/unit/bigquery/_operations/test_io.py

@@ -0,0 +1,41 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+fromunittestimportmock
+
+importpytest
+
+importbigframes.bigquery._operations.io
+importbigframes.core.sql.io
+importbigframes.session
+
+
+@pytest.fixture
+defmock_session():
+returnmock.create_autospec(spec=bigframes.session.Session)
+
+
+@mock.patch("bigframes.bigquery._operations.io._get_table_metadata")
+deftest_load_data(get_table_metadata_mock, mock_session):
+bigframes.bigquery._operations.io.load_data(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+session=mock_session,
+ )
+mock_session.read_gbq_query.assert_called_once()
+generated_sql=mock_session.read_gbq_query.call_args[0][0]
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
+assertgenerated_sql==expected
+get_table_metadata_mock.assert_called_once()

tests/unit/core/sql/test_io.py

@@ -0,0 +1,90 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+importbigframes.core.sql.io
+
+
+deftest_load_data_ddl():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
+assertsql==expected
+
+
+deftest_load_data_ddl_overwrite():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+write_disposition="OVERWRITE",
+columns={"col1": "INT64", "col2": "STRING"},
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA OVERWRITE my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
+assertsql==expected
+
+
+deftest_load_data_ddl_with_partition_columns():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+with_partition_columns={"part1": "DATE", "part2": "STRING"},
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*']) WITH PARTITION COLUMNS (part1 DATE, part2 STRING)"
+assertsql==expected
+
+
+deftest_load_data_ddl_connection():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+connection_name="my-connection",
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*']) WITH CONNECTION `my-connection`"
+assertsql==expected
+
+
+deftest_load_data_ddl_partition_by():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+partition_by=["date_col"],
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) PARTITION BY date_col FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
+assertsql==expected
+
+
+deftest_load_data_ddl_cluster_by():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+cluster_by=["cluster_col"],
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) CLUSTER BY cluster_col FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
+assertsql==expected
+
+
+deftest_load_data_ddl_table_options():
+sql=bigframes.core.sql.io.load_data_ddl(
+"my-project.my_dataset.my_table",
+columns={"col1": "INT64", "col2": "STRING"},
+table_options={"description": "my table"},
+from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
+ )
+expected="LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) OPTIONS (description = 'my table') FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
+assertsql==expected