|
| 1 | +# Copyright 2025 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 | + |
| 16 | +def test_sessions_and_io(project_id: str, dataset_id: str) -> None: |
| 17 | + YOUR_PROJECT_ID = project_id |
| 18 | + YOUR_LOCATION = "us" |
| 19 | + |
| 20 | + # [START bigquery_dataframes_create_and_use_session_instance] |
| 21 | + import bigframes |
| 22 | + import bigframes.pandas as bpd |
| 23 | + |
| 24 | + # Create session object |
| 25 | + context = bigframes.BigQueryOptions( |
| 26 | + project=YOUR_PROJECT_ID, |
| 27 | + location=YOUR_LOCATION, |
| 28 | + ) |
| 29 | + session = bigframes.Session(context) |
| 30 | + |
| 31 | + # Load a BigQuery table into a dataframe |
| 32 | + df1 = session.read_gbq("bigquery-public-data.ml_datasets.penguins") |
| 33 | + |
| 34 | + # Create a dataframe with local data: |
| 35 | + df2 = bpd.DataFrame({"my_col": [1, 2, 3]}, session=session) |
| 36 | + # [END bigquery_dataframes_create_and_use_session_instance] |
| 37 | + assert df1 is not None |
| 38 | + assert df2 is not None |
| 39 | + |
| 40 | + # [START bigquery_dataframes_combine_data_from_multiple_sessions_raise_error] |
| 41 | + import bigframes |
| 42 | + import bigframes.pandas as bpd |
| 43 | + |
| 44 | + context = bigframes.BigQueryOptions(location=YOUR_LOCATION, project=YOUR_PROJECT_ID) |
| 45 | + |
| 46 | + session1 = bigframes.Session(context) |
| 47 | + session2 = bigframes.Session(context) |
| 48 | + |
| 49 | + series1 = bpd.Series([1, 2, 3, 4, 5], session=session1) |
| 50 | + series2 = bpd.Series([1, 2, 3, 4, 5], session=session2) |
| 51 | + |
| 52 | + try: |
| 53 | + series1 + series2 |
| 54 | + except ValueError as e: |
| 55 | + print(e) # Error message: Cannot use combine sources from multiple sessions |
| 56 | + # [END bigquery_dataframes_combine_data_from_multiple_sessions_raise_error] |
| 57 | + |
| 58 | + # [START bigquery_dataframes_set_options_for_global_session] |
| 59 | + import bigframes.pandas as bpd |
| 60 | + |
| 61 | + # Set project ID for the global session |
| 62 | + bpd.options.bigquery.project = YOUR_PROJECT_ID |
| 63 | + # Update the global default session location |
| 64 | + bpd.options.bigquery.location = YOUR_LOCATION |
| 65 | + # [END bigquery_dataframes_set_options_for_global_session] |
| 66 | + |
| 67 | + # [START bigquery_dataframes_global_session_is_the_default_session] |
| 68 | + # The following two statements are essentiall the same |
| 69 | + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") |
| 70 | + df = bpd.get_global_session().read_gbq("bigquery-public-data.ml_datasets.penguins") |
| 71 | + # [END bigquery_dataframes_global_session_is_the_default_session] |
| 72 | + assert df is not None |
| 73 | + |
| 74 | + # [START bigquery_dataframes_create_dataframe_from_py_and_np] |
| 75 | + import numpy as np |
| 76 | + |
| 77 | + import bigframes.pandas as bpd |
| 78 | + |
| 79 | + s = bpd.Series([1, 2, 3]) |
| 80 | + |
| 81 | + # Create a dataframe with Python dict |
| 82 | + df = bpd.DataFrame( |
| 83 | + { |
| 84 | + "col_1": [1, 2, 3], |
| 85 | + "col_2": [4, 5, 6], |
| 86 | + } |
| 87 | + ) |
| 88 | + |
| 89 | + # Create a series with Numpy |
| 90 | + s = bpd.Series(np.arange(10)) |
| 91 | + # [END bigquery_dataframes_create_dataframe_from_py_and_np] |
| 92 | + assert s is not None |
| 93 | + |
| 94 | + # [START bigquery_dataframes_create_dataframe_from_pandas] |
| 95 | + import numpy as np |
| 96 | + import pandas as pd |
| 97 | + |
| 98 | + import bigframes.pandas as bpd |
| 99 | + |
| 100 | + pd_df = pd.DataFrame(np.random.randn(4, 2)) |
| 101 | + |
| 102 | + # Convert Pandas dataframe to BigQuery DataFrame with read_pandas() |
| 103 | + df_1 = bpd.read_pandas(pd_df) |
| 104 | + # Convert Pandas dataframe to BigQuery DataFrame with the dataframe constructor |
| 105 | + df_2 = bpd.DataFrame(pd_df) |
| 106 | + # [END bigquery_dataframes_create_dataframe_from_pandas] |
| 107 | + assert df_1 is not None |
| 108 | + assert df_2 is not None |
| 109 | + |
| 110 | + # [START bigquery_dataframes_convert_bq_dataframe_to_pandas] |
| 111 | + import bigframes.pandas as bpd |
| 112 | + |
| 113 | + bf_df = bpd.DataFrame({"my_col": [1, 2, 3]}) |
| 114 | + # Returns a Pandas Dataframe |
| 115 | + bf_df.to_pandas() |
| 116 | + |
| 117 | + bf_s = bpd.Series([1, 2, 3]) |
| 118 | + # Returns a Pandas Series |
| 119 | + bf_s.to_pandas() |
| 120 | + # [END bigquery_dataframes_convert_bq_dataframe_to_pandas] |
| 121 | + assert bf_s.to_pandas() is not None |
| 122 | + |
| 123 | + # [START bigquery_dataframes_to_pandas_dry_run] |
| 124 | + import bigframes.pandas as bpd |
| 125 | + |
| 126 | + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") |
| 127 | + |
| 128 | + # Returns a Pandas series with dry run stats |
| 129 | + df.to_pandas(dry_run=True) |
| 130 | + # [END bigquery_dataframes_to_pandas_dry_run] |
| 131 | + assert df.to_pandas(dry_run=True) is not None |
| 132 | + |
| 133 | + # [START bigquery_dataframes_read_data_from_csv] |
| 134 | + import bigframes.pandas as bpd |
| 135 | + |
| 136 | + # Read a CSV file from GCS |
| 137 | + df = bpd.read_csv("gs://cloud-samples-data/bigquery/us-states/us-states.csv") |
| 138 | + # [END bigquery_dataframes_read_data_from_csv] |
| 139 | + assert df is not None |
| 140 | + |
| 141 | + # [START bigquery_dataframes_read_data_from_bigquery_table] |
| 142 | + import bigframes.pandas as bpd |
| 143 | + |
| 144 | + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") |
| 145 | + # [END bigquery_dataframes_read_data_from_bigquery_table] |
| 146 | + assert df is not None |
| 147 | + |
| 148 | + # [START bigquery_dataframes_read_from_sql_query] |
| 149 | + import bigframes.pandas as bpd |
| 150 | + |
| 151 | + sql = """ |
| 152 | + SELECT species, island, body_mass_g |
| 153 | + FROM bigquery-public-data.ml_datasets.penguins |
| 154 | + WHERE sex = 'MALE' |
| 155 | + """ |
| 156 | + |
| 157 | + df = bpd.read_gbq(sql) |
| 158 | + # [END bigquery_dataframes_read_from_sql_query] |
| 159 | + assert df is not None |
| 160 | + |
| 161 | + table_name = "snippets-session-and-io-test" |
| 162 | + |
| 163 | + # [START bigquery_dataframes_dataframe_to_bigquery_table] |
| 164 | + import bigframes.pandas as bpd |
| 165 | + |
| 166 | + df = bpd.DataFrame({"my_col": [1, 2, 3]}) |
| 167 | + |
| 168 | + df.to_gbq(f"{project_id}.{dataset_id}.{table_name}") |
| 169 | + # [END bigquery_dataframes_dataframe_to_bigquery_table] |
0 commit comments