Skip to content

Commit fa7e365

Browse files
committed
docs: add code snippet for storing dataframes to a CSV file
1 parent 03ac660 commit fa7e365

File tree

3 files changed

+33
-9
lines changed

3 files changed

+33
-9
lines changed

samples/snippets/conftest.py

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,12 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from typing import Iterator
15+
from typing import Iterator, Generator
1616

1717
from google.cloud import bigquery
1818
import pytest
1919
import test_utils.prefixer
20+
from google.cloud import storage
2021

2122
import bigframes.pandas as bpd
2223

@@ -42,11 +43,27 @@ def bigquery_client() -> bigquery.Client:
4243
return bigquery_client
4344

4445

46+
@pytest.fixture(scope="session")
47+
def storage_client(project_id) -> storage.Client:
48+
return storage.Client(project=project_id)
49+
50+
4551
@pytest.fixture(scope="session")
4652
def project_id(bigquery_client: bigquery.Client) -> str:
4753
return bigquery_client.project
4854

4955

56+
@pytest.fixture(scope="session")
57+
def gcs_bucket(storage_client) -> Generator[str, None, None]:
58+
bucket_name = "bigframes-gcs-test"
59+
60+
yield bucket_name
61+
62+
bucket = storage_client.get_bucket(bucket_name)
63+
for blob in bucket.list_blobs():
64+
blob.delete()
65+
66+
5067
@pytest.fixture(autouse=True)
5168
def reset_session() -> None:
5269
"""An autouse fixture ensuring each sample runs in a fresh session.
@@ -78,11 +95,6 @@ def dataset_id_eu(bigquery_client: bigquery.Client, project_id: str) -> Iterator
7895
bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True)
7996

8097

81-
@pytest.fixture(scope="session")
82-
def gcs_dst_bucket() -> str:
83-
return "gs://bigframes_blob_test"
84-
85-
8698
@pytest.fixture
8799
def random_model_id(
88100
bigquery_client: bigquery.Client, project_id: str, dataset_id: str

samples/snippets/multimodal_test.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@
1313
# limitations under the License.
1414

1515

16-
def test_multimodal_dataframe(gcs_dst_bucket: str) -> None:
16+
def test_multimodal_dataframe(gcs_bucket: str) -> None:
1717
# destination folder must be in a GCS bucket that the BQ connection service account (default or user provided) has write access to.
18-
dst_bucket = gcs_dst_bucket
18+
dst_bucket = f"gs://{gcs_bucket}"
1919
# [START bigquery_dataframes_multimodal_dataframe_create]
2020
import bigframes
2121

samples/snippets/sessions_and_io_test.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,13 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from google.cloud import storage
1516

16-
def test_sessions_and_io(project_id: str, dataset_id: str) -> None:
17+
18+
def test_sessions_and_io(project_id: str, dataset_id: str, gcs_bucket: str) -> None:
1719
YOUR_PROJECT_ID = project_id
1820
YOUR_LOCATION = "us"
21+
YOUR_BUCKET = gcs_bucket
1922

2023
# [START bigquery_dataframes_create_and_use_session_instance]
2124
import bigframes
@@ -138,6 +141,15 @@ def test_sessions_and_io(project_id: str, dataset_id: str) -> None:
138141
# [END bigquery_dataframes_read_data_from_csv]
139142
assert df is not None
140143

144+
# [START bigquery_dataframes_write_data_to_csv]
145+
import bigframes.pandas as bpd
146+
147+
df = bpd.DataFrame({"my_col": [1, 2, 3]})
148+
# Write a dataframe to a CSV file in GCS
149+
df.to_csv(f"gs://{YOUR_BUCKET}/myfile*.csv")
150+
# [END bigquery_dataframes_write_data_to_csv]
151+
assert df is not None
152+
141153
# [START bigquery_dataframes_read_data_from_bigquery_table]
142154
import bigframes.pandas as bpd
143155

0 commit comments

Comments
 (0)