Skip to content
This repository was archived by the owner on May 7, 2026. It is now read-only.

Commit cdfa759

Browse files
GarrettWuGarrettWu
committed
feat: add bigframe.bigquery.load_data function
1 parent cfa1c8e commit cdfa759

6 files changed

Lines changed: 373 additions & 0 deletions

File tree

‎bigframes/bigquery/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
st_regionstats,
4444
st_simplify,
4545
)
46+
from bigframes.bigquery._operations.io import load_data
4647
from bigframes.bigquery._operations.json import (
4748
json_extract,
4849
json_extract_array,
@@ -107,6 +108,8 @@
107108
struct,
108109
# table ops
109110
create_external_table,
111+
# io ops
112+
load_data,
110113
]
111114

112115
_module = sys.modules[__name__]
@@ -160,6 +163,8 @@
160163
"struct",
161164
# table ops
162165
"create_external_table",
166+
# io ops
167+
"load_data",
163168
# Modules / SQL namespaces
164169
"ai",
165170
"ml",

‎bigframes/bigquery/_operations/io.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from typing import Mapping, Optional, Union
18+
19+
import bigframes_vendored.constants
20+
import google.cloud.bigquery
21+
import pandas as pd
22+
23+
import bigframes.core.logging.log_adapter as log_adapter
24+
import bigframes.core.sql.io
25+
import bigframes.session
26+
27+
28+
def _get_table_metadata(
29+
*,
30+
bqclient: google.cloud.bigquery.Client,
31+
table_name: str,
32+
) -> pd.Series:
33+
table_metadata = bqclient.get_table(table_name)
34+
table_dict = table_metadata.to_api_repr()
35+
return pd.Series(table_dict)
36+
37+
38+
@log_adapter.method_logger(custom_base_name="bigquery_io")
39+
def load_data(
40+
table_name: str,
41+
*,
42+
replace: bool = False,
43+
columns: Optional[Mapping[str, str]] = None,
44+
partition_by: Optional[list[str]] = None,
45+
cluster_by: Optional[list[str]] = None,
46+
table_options: Optional[Mapping[str, Union[str, int, float, bool, list]]] = None,
47+
from_files_options: Mapping[str, Union[str, int, float, bool, list]],
48+
with_partition_columns: Optional[Mapping[str, str]] = None,
49+
connection_name: Optional[str] = None,
50+
session: Optional[bigframes.session.Session] = None,
51+
) -> pd.Series:
52+
"""
53+
Loads data into a BigQuery table.
54+
See the `BigQuery LOAD DATA DDL syntax
55+
<https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/load-statements>`_
56+
for additional reference.
57+
Args:
58+
table_name (str):
59+
The name of the table in BigQuery.
60+
replace (bool, default False):
61+
Whether to replace the table if it already exists.
62+
columns (Mapping[str, str], optional):
63+
The table's schema.
64+
partition_by (list[str], optional):
65+
A list of columns to partition the table by.
66+
cluster_by (list[str], optional):
67+
A list of columns to cluster the table by.
68+
table_options (Mapping[str, Union[str, int, float, bool, list]], optional):
69+
The table options.
70+
from_files_options (Mapping[str, Union[str, int, float, bool, list]]):
71+
The options for loading data from files.
72+
with_partition_columns (Mapping[str, str], optional):
73+
The table's partition columns.
74+
connection_name (str, optional):
75+
The connection to use for the table.
76+
session (bigframes.session.Session, optional):
77+
The session to use. If not provided, the default session is used.
78+
Returns:
79+
pandas.Series:
80+
A Series with object dtype containing the table metadata. Reference
81+
the `BigQuery Table REST API reference
82+
<https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table>`_
83+
for available fields.
84+
"""
85+
import bigframes.pandas as bpd
86+
87+
sql = bigframes.core.sql.io.load_data_ddl(
88+
table_name=table_name,
89+
replace=replace,
90+
columns=columns,
91+
partition_by=partition_by,
92+
cluster_by=cluster_by,
93+
table_options=table_options,
94+
from_files_options=from_files_options,
95+
with_partition_columns=with_partition_columns,
96+
connection_name=connection_name,
97+
)
98+
99+
if session is None:
100+
bpd.read_gbq_query(sql)
101+
session = bpd.get_global_session()
102+
assert (
103+
session is not None
104+
), f"Missing connection to BigQuery. Please report how you encountered this error at {bigframes_vendored.constants.FEEDBACK_LINK}."
105+
else:
106+
session.read_gbq_query(sql)
107+
108+
return _get_table_metadata(bqclient=session.bqclient, table_name=table_name)

‎bigframes/core/sql/io.py‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from typing import Mapping, Optional, Union
18+
19+
20+
def load_data_ddl(
21+
table_name: str,
22+
*,
23+
replace: bool = False,
24+
columns: Optional[Mapping[str, str]] = None,
25+
partition_by: Optional[list[str]] = None,
26+
cluster_by: Optional[list[str]] = None,
27+
table_options: Optional[Mapping[str, Union[str, int, float, bool, list]]] = None,
28+
from_files_options: Mapping[str, Union[str, int, float, bool, list]],
29+
with_partition_columns: Optional[Mapping[str, str]] = None,
30+
connection_name: Optional[str] = None,
31+
) -> str:
32+
"""Generates the LOAD DATA DDL statement."""
33+
statement = ["LOAD DATA"]
34+
if replace:
35+
statement.append("OVERWRITE")
36+
else:
37+
statement.append("INTO")
38+
statement.append(table_name)
39+
40+
if columns:
41+
column_defs = ", ".join([f"{name} {typ}" for name, typ in columns.items()])
42+
statement.append(f"({column_defs})")
43+
44+
if partition_by:
45+
statement.append(f"PARTITION BY {', '.join(partition_by)}")
46+
47+
if cluster_by:
48+
statement.append(f"CLUSTER BY {', '.join(cluster_by)}")
49+
50+
if table_options:
51+
opts = []
52+
for key, value in table_options.items():
53+
if isinstance(value, str):
54+
value_sql = repr(value)
55+
opts.append(f"{key} = {value_sql}")
56+
elif isinstance(value, bool):
57+
opts.append(f"{key} = {str(value).upper()}")
58+
elif isinstance(value, list):
59+
list_str = ", ".join([repr(v) for v in value])
60+
opts.append(f"{key} = [{list_str}]")
61+
else:
62+
opts.append(f"{key} = {value}")
63+
options_str = ", ".join(opts)
64+
statement.append(f"OPTIONS ({options_str})")
65+
66+
opts = []
67+
for key, value in from_files_options.items():
68+
if isinstance(value, str):
69+
value_sql = repr(value)
70+
opts.append(f"{key} = {value_sql}")
71+
elif isinstance(value, bool):
72+
opts.append(f"{key} = {str(value).upper()}")
73+
elif isinstance(value, list):
74+
list_str = ", ".join([repr(v) for v in value])
75+
opts.append(f"{key} = [{list_str}]")
76+
else:
77+
opts.append(f"{key} = {value}")
78+
options_str = ", ".join(opts)
79+
statement.append(f"FROM FILES ({options_str})")
80+
81+
if with_partition_columns:
82+
part_defs = ", ".join(
83+
[f"{name} {typ}" for name, typ in with_partition_columns.items()]
84+
)
85+
statement.append(f"WITH PARTITION COLUMNS ({part_defs})")
86+
87+
if connection_name:
88+
statement.append(f"WITH CONNECTION `{connection_name}`")
89+
90+
return " ".join(statement)

‎tests/system/large/bigquery/test_io.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import bigframes.bigquery as bbq
16+
17+
18+
def test_load_data(session, dataset_id):
19+
table_name = f"{dataset_id}.test_load_data"
20+
uri = "gs://cloud-samples-data/bigquery/us-states/us-states.csv"
21+
22+
# Create the external table
23+
table = bbq.load_data(
24+
table_name,
25+
columns={
26+
"name": "STRING",
27+
"post_abbr": "STRING",
28+
},
29+
from_files_options={"format": "CSV", "uris": [uri], "skip_leading_rows": 1},
30+
session=session,
31+
)
32+
assert table is not None
33+
34+
# Read the table to verify
35+
import bigframes.pandas as bpd
36+
37+
bf_df = bpd.read_gbq(table_name)
38+
pd_df = bf_df.to_pandas()
39+
assert len(pd_df) > 0

‎tests/unit/bigquery/_operations/test_io.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from unittest import mock
16+
17+
import pytest
18+
19+
import bigframes.bigquery._operations.io
20+
import bigframes.core.sql.io
21+
import bigframes.session
22+
23+
24+
@pytest.fixture
25+
def mock_session():
26+
return mock.create_autospec(spec=bigframes.session.Session)
27+
28+
29+
@mock.patch("bigframes.bigquery._operations.io._get_table_metadata")
30+
def test_load_data(get_table_metadata_mock, mock_session):
31+
bigframes.bigquery._operations.io.load_data(
32+
"my-project.my_dataset.my_table",
33+
columns={"col1": "INT64", "col2": "STRING"},
34+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
35+
session=mock_session,
36+
)
37+
mock_session.read_gbq_query.assert_called_once()
38+
generated_sql = mock_session.read_gbq_query.call_args[0][0]
39+
expected = "LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
40+
assert generated_sql == expected
41+
get_table_metadata_mock.assert_called_once()

‎tests/unit/core/sql/test_io.py‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
14+
import bigframes.core.sql.io
15+
16+
17+
def test_load_data_ddl():
18+
sql = bigframes.core.sql.io.load_data_ddl(
19+
"my-project.my_dataset.my_table",
20+
columns={"col1": "INT64", "col2": "STRING"},
21+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
22+
)
23+
expected = "LOAD DATA INTO my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
24+
assert sql == expected
25+
26+
27+
def test_load_data_ddl_replace():
28+
sql = bigframes.core.sql.io.load_data_ddl(
29+
"my-project.my_dataset.my_table",
30+
replace=True,
31+
columns={"col1": "INT64", "col2": "STRING"},
32+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
33+
)
34+
expected = "LOAD DATA OVERWRITE my-project.my_dataset.my_table (col1 INT64, col2 STRING) FROM FILES (format = 'CSV', uris = ['gs://bucket/path*'])"
35+
assert sql == expected
36+
37+
38+
def test_load_data_ddl_with_partition_columns():
39+
sql = bigframes.core.sql.io.load_data_ddl(
40+
"my-project.my_dataset.my_table",
41+
columns={"col1": "INT64", "col2": "STRING"},
42+
with_partition_columns={"part1": "DATE", "part2": "STRING"},
43+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
44+
)
45+
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)"
46+
assert sql == expected
47+
48+
49+
def test_load_data_ddl_connection():
50+
sql = bigframes.core.sql.io.load_data_ddl(
51+
"my-project.my_dataset.my_table",
52+
columns={"col1": "INT64", "col2": "STRING"},
53+
connection_name="my-connection",
54+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
55+
)
56+
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`"
57+
assert sql == expected
58+
59+
60+
def test_load_data_ddl_partition_by():
61+
sql = bigframes.core.sql.io.load_data_ddl(
62+
"my-project.my_dataset.my_table",
63+
columns={"col1": "INT64", "col2": "STRING"},
64+
partition_by=["date_col"],
65+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
66+
)
67+
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*'])"
68+
assert sql == expected
69+
70+
71+
def test_load_data_ddl_cluster_by():
72+
sql = bigframes.core.sql.io.load_data_ddl(
73+
"my-project.my_dataset.my_table",
74+
columns={"col1": "INT64", "col2": "STRING"},
75+
cluster_by=["cluster_col"],
76+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
77+
)
78+
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*'])"
79+
assert sql == expected
80+
81+
82+
def test_load_data_ddl_table_options():
83+
sql = bigframes.core.sql.io.load_data_ddl(
84+
"my-project.my_dataset.my_table",
85+
columns={"col1": "INT64", "col2": "STRING"},
86+
table_options={"description": "my table"},
87+
from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]},
88+
)
89+
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*'])"
90+
assert sql == expected

0 commit comments

Comments
 (0)