diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 32a2908a42..5c0d9b78e1 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -46,7 +46,6 @@ import bigframes.core.guid import bigframes.core.indexers as indexers import bigframes.core.indexes as indexes -import bigframes.core.io import bigframes.core.ordering as order import bigframes.core.utils as utils import bigframes.core.window @@ -56,6 +55,7 @@ import bigframes.operations.aggregations as agg_ops import bigframes.series import bigframes.series as bf_series +import bigframes.session._io.bigquery import third_party.bigframes_vendored.pandas.core.frame as vendored_pandas_frame import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing @@ -2201,9 +2201,9 @@ def to_csv( raise NotImplementedError(ERROR_IO_REQUIRES_WILDCARD) result_table = self._run_io_query( - index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID + index=index, ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID ) - export_data_statement = bigframes.core.io.create_export_csv_statement( + export_data_statement = bigframes.session._io.bigquery.create_export_csv_statement( f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path_or_buf, field_delimiter=sep, @@ -2243,9 +2243,9 @@ def to_json( ) result_table = self._run_io_query( - index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID + index=index, ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID ) - export_data_statement = bigframes.core.io.create_export_data_statement( + export_data_statement = bigframes.session._io.bigquery.create_export_data_statement( f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path_or_buf, format="JSON", @@ -2319,9 +2319,9 @@ def to_parquet( export_options["compression"] = compression.upper() result_table = self._run_io_query( - index=index, ordering_id=bigframes.core.io.IO_ORDERING_ID + index=index, ordering_id=bigframes.session._io.bigquery.IO_ORDERING_ID ) - export_data_statement = bigframes.core.io.create_export_data_statement( + export_data_statement = bigframes.session._io.bigquery.create_export_data_statement( f"{result_table.project}.{result_table.dataset_id}.{result_table.table_id}", uri=path, format="PARQUET", diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 5ec3da1a5a..db9c5a353c 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -67,13 +67,13 @@ import bigframes.core as core import bigframes.core.blocks as blocks import bigframes.core.guid as guid -import bigframes.core.io as bigframes_io from bigframes.core.ordering import IntegerEncoding, OrderingColumnReference import bigframes.core.utils as utils import bigframes.dataframe as dataframe import bigframes.formatting_helpers as formatting_helpers from bigframes.remote_function import read_gbq_function as bigframes_rgf from bigframes.remote_function import remote_function as bigframes_rf +import bigframes.session._io.bigquery as bigframes_io import bigframes.session.clients import bigframes.version diff --git a/bigframes/session/_io/__init__.py b/bigframes/session/_io/__init__.py new file mode 100644 index 0000000000..1dc90d1848 --- /dev/null +++ b/bigframes/session/_io/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2023 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. diff --git a/bigframes/core/io.py b/bigframes/session/_io/bigquery.py similarity index 100% rename from bigframes/core/io.py rename to bigframes/session/_io/bigquery.py diff --git a/tests/unit/core/test_io.py b/tests/unit/session/test_io_bigquery.py similarity index 93% rename from tests/unit/core/test_io.py rename to tests/unit/session/test_io_bigquery.py index afb38a5f75..d2255d5edf 100644 --- a/tests/unit/core/test_io.py +++ b/tests/unit/session/test_io_bigquery.py @@ -18,7 +18,7 @@ import google.cloud.bigquery as bigquery import pytest -import bigframes.core.io +import bigframes.session._io.bigquery def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): @@ -26,7 +26,7 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): "my-test-project._e8166e0cdb.anonbb92cd" ) - sql = bigframes.core.io.create_snapshot_sql( + sql = bigframes.session._io.bigquery.create_snapshot_sql( table_ref, datetime.datetime.now(datetime.timezone.utc) ) @@ -40,7 +40,7 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets(): def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg") - sql = bigframes.core.io.create_snapshot_sql( + sql = bigframes.session._io.bigquery.create_snapshot_sql( table_ref, datetime.datetime.now(datetime.timezone.utc) ) @@ -101,4 +101,5 @@ def test_create_snapshot_sql_doesnt_timetravel_session_datasets(): ), ) def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str): - pass + sql = bigframes.session._io.bigquery.bq_schema_to_sql(schema) + assert sql == expected