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

fix: read_pandas inline respects location #412

Merged
gcf-merge-on-green merged 14 commits into from
Mar 12, 2024
Merged

fix: read_pandas inline respects location #412

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
8f2ee8d
fix: read_pandas inline respects location
GarrettWu Mar 5, 2024
147c3be
fix tests
GarrettWu Mar 6, 2024
6bca066
fix tests
GarrettWu Mar 6, 2024
da6a97e
fix tests
GarrettWu Mar 6, 2024
3a3c8e6
retrigger tests
GarrettWu Mar 7, 2024
cb4589d
add construction test
GarrettWu Mar 7, 2024
82c7bde
Merge remote-tracking branch 'github/main' into garrettwu-fix
GarrettWu Mar 7, 2024
8ec4ccc
fix test
GarrettWu Mar 7, 2024
801bcf5
Merge branch 'main' into garrettwu-fix
gcf-merge-on-green[bot] Mar 8, 2024
eba7a27
fix tests
GarrettWu Mar 8, 2024
60c57bc
Merge remote-tracking branch 'github/garrettwu-fix' into garrettwu-fix
GarrettWu Mar 8, 2024
2cb5e0c
fix tests
GarrettWu Mar 8, 2024
054338d
mark ReadLocalNode session optional
GarrettWu Mar 11, 2024
7aa8d76
Merge branch 'main' into garrettwu-fix
GarrettWu Mar 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions bigframes/core/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def from_ibis(
return cls(node)

@classmethod
def from_pandas(cls, pd_df: pandas.DataFrame):
def from_pandas(cls, pd_df: pandas.DataFrame, session: bigframes.Session):
iobytes = io.BytesIO()
# Use alphanumeric identifiers, to avoid downstream problems with escaping.
as_ids = [
Expand All @@ -78,7 +78,7 @@ def from_pandas(cls, pd_df: pandas.DataFrame):
]
unique_ids = tuple(bigframes.core.utils.disambiguate_ids(as_ids))
pd_df.reset_index(drop=True).set_axis(unique_ids, axis=1).to_feather(iobytes)
node = nodes.ReadLocalNode(iobytes.getvalue())
node = nodes.ReadLocalNode(feather_bytes=iobytes.getvalue(), session=session)
return cls(node)

@property
Expand Down
4 changes: 2 additions & 2 deletions bigframes/core/blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def __init__(
self._stats_cache[" ".join(self.index_columns)] = {}

@classmethod
def from_local(cls, data) -> Block:
def from_local(cls, data, session: bigframes.Session) -> Block:
pd_data = pd.DataFrame(data)
columns = pd_data.columns

Expand All @@ -162,7 +162,7 @@ def from_local(cls, data) -> Block:
)
index_ids = pd_data.columns[: len(index_labels)]

keys_expr = core.ArrayValue.from_pandas(pd_data)
keys_expr = core.ArrayValue.from_pandas(pd_data, session)
return cls(
keys_expr,
column_labels=columns,
Expand Down
1 change: 1 addition & 0 deletions bigframes/core/nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ def __hash__(self):
@dataclass(frozen=True)
class ReadLocalNode(BigFrameNode):
feather_bytes: bytes
session: typing.Optional[bigframes.session.Session] = None

Copy link
Copy Markdown
Contributor

@TrevorBergeron TrevorBergeron Mar 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Local data is session-independent, we don't want to add a session constrain to the node. Don't worry about a dataframe/block not having session, that just means you can execute it anywhere, as all the data sources are local.

Copy link
Copy Markdown
Contributor Author

@GarrettWu GarrettWu Mar 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The data itself is independent of session, yes. But when reading a local data, a specific session will be used. And when executing the query, we'd call that particular session. Do we have other options than keep a record here?

Copy link
Copy Markdown
Contributor

@TrevorBergeron TrevorBergeron Mar 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A session will be used to execute the tree, yes, but the choice of session need not be constrained by the tree itself. You can check a tree to see if it has a required session, and otherwise, just use the default session to execute trees that don't depend on a specific session.

Copy link
Copy Markdown
Contributor

@TrevorBergeron TrevorBergeron Mar 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed offline, lets just make the session an optional field that we set when users have a specific session they used to read the local data with session.read_gbq.

Copy link
Copy Markdown
Contributor Author

@GarrettWu GarrettWu Mar 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated.

def __hash__(self):
return self._node_hash
Expand Down
2 changes: 1 addition & 1 deletion bigframes/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1646,7 +1646,7 @@ def _reindex_rows(
raise NotImplementedError(
"Cannot reindex with index with different nlevels"
)
new_indexer = DataFrame(index=index)[[]]
new_indexer = DataFrame(index=index, session=self._session)[[]]
# multiindex join is senstive to index names, so we will set all these
result = new_indexer.rename_axis(range(new_indexer.index.nlevels)).join(
self.rename_axis(range(self.index.nlevels)),
Expand Down
2 changes: 1 addition & 1 deletion bigframes/session/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -916,7 +916,7 @@ def _read_pandas(
def _read_pandas_inline(
self, pandas_dataframe: pandas.DataFrame
) -> dataframe.DataFrame:
return dataframe.DataFrame(blocks.Block.from_local(pandas_dataframe))
return dataframe.DataFrame(blocks.Block.from_local(pandas_dataframe, self))

def _read_pandas_load_job(
self, pandas_dataframe: pandas.DataFrame, api_name: str
Expand Down
2 changes: 1 addition & 1 deletion notebooks/location/regionalized.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2791,7 +2791,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.9"
},
"orig_nbformat": 4
},
Expand Down
17 changes: 17 additions & 0 deletions tests/system/small/test_dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,23 @@ def test_df_construct_from_dict():
)


def test_df_construct_inline_respects_location():
import bigframes.pandas as bpd

bpd.close_session()
bpd.options.bigquery.location = "europe-west1"

df = bpd.DataFrame([[1, 2, 3], [4, 5, 6]])
repr(df)

table = bpd.get_global_session().bqclient.get_table(df.query_job.destination)
assert table.location == "europe-west1"

# Reset global session
bpd.close_session()
bpd.options.bigquery.location = "us"


def test_get_column(scalars_dfs):
scalars_df, scalars_pandas_df = scalars_dfs
col_name = "int64_col"
Expand Down
11 changes: 11 additions & 0 deletions tests/system/small/test_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,6 +369,17 @@ def test_read_pandas(session, scalars_dfs):
pd.testing.assert_frame_equal(result, expected)


def test_read_pandas_inline_respects_location():
options = bigframes.BigQueryOptions(location="europe-west1")
session = bigframes.Session(options)

df = session.read_pandas(pd.DataFrame([[1, 2, 3], [4, 5, 6]]))
Copy link
Copy Markdown
Contributor

@shobsi shobsi Mar 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we add a test for inline data too, and verify that it creates the result tables in the intended location? bpd.Dataframe([[1, 2, 3], [4, 5, 6]])

Copy link
Copy Markdown
Contributor Author

@GarrettWu GarrettWu Mar 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added

repr(df)

table = session.bqclient.get_table(df.query_job.destination)
assert table.location == "europe-west1"


def test_read_pandas_col_label_w_space(session: bigframes.Session):
expected = pd.DataFrame(
{
Expand Down
9 changes: 8 additions & 1 deletion tests/unit/core/test_blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from unittest import mock

import pandas
import pandas.testing
import pytest

import bigframes
import bigframes.core.blocks as blocks


Expand Down Expand Up @@ -74,8 +77,12 @@
)
def test_block_from_local(data):
expected = pandas.DataFrame(data)
mock_session = mock.create_autospec(spec=bigframes.Session)

# hard-coded the returned dimension of the session for that each of the test case contains 3 rows.
mock_session._execute.return_value = (iter([[3]]), None)

block = blocks.Block.from_local(data)
block = blocks.Block.from_local(data, mock_session)

pandas.testing.assert_index_equal(block.column_labels, expected.columns)
assert tuple(block.index.names) == tuple(expected.index.names)
Expand Down