X Tutup
Skip to content

chore: fix google3 import colabsqlviz tests #2499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
GarrettWu merged 2 commits into from
Mar 10, 2026
Merged

chore: fix google3 import colabsqlviz tests #2499

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
513343b
chore: fix google3 import colabsqlviz tests
GarrettWu Mar 10, 2026
ca3ca55
fix: mypy
GarrettWu Mar 10, 2026
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
7 changes: 1 addition & 6 deletions bigframes/testing/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,5 @@
These modules are provided for testing the BigQuery DataFrames package. The
interface is not considered stable.
"""
from bigframes.testing.utils import (
assert_frame_equal,
assert_index_equal,
assert_series_equal,
)

__all__ = ["assert_frame_equal", "assert_series_equal", "assert_index_equal"]
# Do not import modules contains pytest. (b/490160312)
14 changes: 7 additions & 7 deletions tests/system/small/bigquery/test_datetime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import pytest

from bigframes import bigquery
import bigframes.testing
import bigframes.testing.utils

_TIMESTAMP_DTYPE = pd.ArrowDtype(pa.timestamp("us", tz="UTC"))

Expand All @@ -41,7 +41,7 @@ def test_unix_seconds(scalars_dfs):
.apply(lambda ts: _to_unix_epoch(ts, "s"))
.astype("Int64")
)
bigframes.testing.assert_series_equal(actual_res, expected_res)
bigframes.testing.utils.assert_series_equal(actual_res, expected_res)


def test_unix_seconds_after_type_casting(int_series):
Expand All @@ -54,7 +54,7 @@ def test_unix_seconds_after_type_casting(int_series):
.apply(lambda ts: _to_unix_epoch(ts, "s"))
.astype("Int64")
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
actual_res, expected_res, check_index_type=False
)

Expand All @@ -76,7 +76,7 @@ def test_unix_millis(scalars_dfs):
.apply(lambda ts: _to_unix_epoch(ts, "ms"))
.astype("Int64")
)
bigframes.testing.assert_series_equal(actual_res, expected_res)
bigframes.testing.utils.assert_series_equal(actual_res, expected_res)


def test_unix_millis_after_type_casting(int_series):
Expand All @@ -89,7 +89,7 @@ def test_unix_millis_after_type_casting(int_series):
.apply(lambda ts: _to_unix_epoch(ts, "ms"))
.astype("Int64")
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
actual_res, expected_res, check_index_type=False
)

Expand All @@ -111,7 +111,7 @@ def test_unix_micros(scalars_dfs):
.apply(lambda ts: _to_unix_epoch(ts, "us"))
.astype("Int64")
)
bigframes.testing.assert_series_equal(actual_res, expected_res)
bigframes.testing.utils.assert_series_equal(actual_res, expected_res)


def test_unix_micros_after_type_casting(int_series):
Expand All @@ -124,7 +124,7 @@ def test_unix_micros_after_type_casting(int_series):
.apply(lambda ts: _to_unix_epoch(ts, "us"))
.astype("Int64")
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
actual_res, expected_res, check_index_type=False
)

Expand Down
26 changes: 13 additions & 13 deletions tests/system/small/bigquery/test_geo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
import bigframes.bigquery as bbq
import bigframes.geopandas
import bigframes.session
import bigframes.testing
import bigframes.testing.utils


def test_geo_st_area(session: bigframes.session.Session):
Expand All @@ -57,7 +57,7 @@ def test_geo_st_area(session: bigframes.session.Session):
geobf_s_result = bbq.st_area(geobf_s).to_pandas().round(-3)
assert geobf_s_result.iloc[0] >= 1000

bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
geopd_s_result,
check_dtype=False,
Expand Down Expand Up @@ -110,7 +110,7 @@ def test_st_length_various_geometries(session):

# Test default use_spheroid
result_default = st_length(geoseries).to_pandas()
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
result_default,
expected_lengths,
rtol=1e-3,
Expand All @@ -119,7 +119,7 @@ def test_st_length_various_geometries(session):

# Test explicit use_spheroid=False
result_explicit_false = st_length(geoseries, use_spheroid=False).to_pandas()
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
result_explicit_false,
expected_lengths,
rtol=1e-3,
Expand Down Expand Up @@ -153,7 +153,7 @@ def test_geo_st_difference_with_geometry_objects(session: bigframes.session.Sess
index=[0, 1, 2],
dtype=geopandas.array.GeometryDtype(),
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -192,7 +192,7 @@ def test_geo_st_difference_with_single_geometry_object(
index=[0, 1, 2],
dtype=geopandas.array.GeometryDtype(),
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand All @@ -218,7 +218,7 @@ def test_geo_st_difference_with_similar_geometry_objects(
index=[0, 1, 2],
dtype=geopandas.array.GeometryDtype(),
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -274,7 +274,7 @@ def test_geo_st_distance_with_geometry_objects(session: bigframes.session.Sessio
index=[0, 1, 2, 3],
dtype="Float64",
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -321,7 +321,7 @@ def test_geo_st_distance_with_single_geometry_object(
],
dtype="Float64",
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -356,7 +356,7 @@ def test_geo_st_intersection_with_geometry_objects(session: bigframes.session.Se
index=[0, 1, 2],
dtype=geopandas.array.GeometryDtype(),
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -395,7 +395,7 @@ def test_geo_st_intersection_with_single_geometry_object(
index=[0, 1, 2],
dtype=geopandas.array.GeometryDtype(),
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -425,7 +425,7 @@ def test_geo_st_intersection_with_similar_geometry_objects(
index=[0, 1, 2],
dtype=geopandas.array.GeometryDtype(),
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
geobf_s_result,
expected,
check_index_type=False,
Expand Down Expand Up @@ -466,7 +466,7 @@ def test_geo_st_isclosed(session: bigframes.session.Session):
]
expected_series = pd.Series(data=expected_data, dtype="boolean")

bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
bf_result,
expected_series,
# We default to Int64 (nullable) dtype, but pandas defaults to int64 index.
Expand Down
34 changes: 25 additions & 9 deletions tests/system/small/bigquery/test_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
import bigframes.bigquery as bbq
import bigframes.dtypes as dtypes
import bigframes.pandas as bpd
import bigframes.testing
import bigframes.testing.utils


def test_sql_scalar_for_all_scalar_types(scalars_df_null_index):
Expand Down Expand Up @@ -60,7 +60,9 @@ def test_sql_scalar_for_bool_series(scalars_df_index):
result = bbq.sql_scalar("CAST({0} AS INT64)", [series])
expected = series.astype(dtypes.INT_DTYPE)
expected.name = result.name
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


@pytest.mark.parametrize(
Expand All @@ -84,7 +86,9 @@ def test_sql_scalar_outputs_all_scalar_types(scalars_df_index, column_name):
result = bbq.sql_scalar("{0}", [series])
expected = series
expected.name = result.name
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


def test_sql_scalar_for_array_series(repeated_df):
Expand Down Expand Up @@ -114,14 +118,18 @@ def test_sql_scalar_for_array_series(repeated_df):
+ repeated_df["numeric_list_col"].list.len()
+ repeated_df["string_list_col"].list.len()
)
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


def test_sql_scalar_outputs_array_series(repeated_df):
result = bbq.sql_scalar("{0}", [repeated_df["int_list_col"]])
expected = repeated_df["int_list_col"]
expected.name = result.name
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


def test_sql_scalar_for_struct_series(nested_structs_df):
Expand All @@ -132,14 +140,18 @@ def test_sql_scalar_for_struct_series(nested_structs_df):
expected = nested_structs_df["person"].struct.field(
"name"
).str.len() + nested_structs_df["person"].struct.field("age")
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


def test_sql_scalar_outputs_struct_series(nested_structs_df):
result = bbq.sql_scalar("{0}", [nested_structs_df["person"]])
expected = nested_structs_df["person"]
expected.name = result.name
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


def test_sql_scalar_for_json_series(json_df):
Expand All @@ -151,11 +163,15 @@ def test_sql_scalar_for_json_series(json_df):
)
expected = bbq.json_value(json_df["json_col"], "$.int_value")
expected.name = result.name
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)


def test_sql_scalar_outputs_json_series(json_df):
result = bbq.sql_scalar("{0}", [json_df["json_col"]])
expected = json_df["json_col"]
expected.name = result.name
bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
bigframes.testing.utils.assert_series_equal(
result.to_pandas(), expected.to_pandas()
)
4 changes: 2 additions & 2 deletions tests/system/small/bigquery/test_struct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

import bigframes.bigquery as bbq
import bigframes.series as series
import bigframes.testing
import bigframes.testing.utils


@pytest.mark.parametrize(
Expand Down Expand Up @@ -53,7 +53,7 @@ def test_struct_from_dataframe(columns_arg):
srs = series.Series(
columns_arg,
)
bigframes.testing.assert_series_equal(
bigframes.testing.utils.assert_series_equal(
srs.to_pandas(),
bbq.struct(srs.struct.explode()).to_pandas(),
check_index_type=False,
Expand Down
4 changes: 2 additions & 2 deletions tests/system/small/core/test_reshape.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

from bigframes import session
from bigframes.core.reshape import merge
import bigframes.testing
import bigframes.testing.utils


@pytest.mark.parametrize(
Expand Down Expand Up @@ -56,7 +56,7 @@ def test_join_with_index(
how=how,
)

bigframes.testing.assert_frame_equal(
bigframes.testing.utils.assert_frame_equal(
bf_result, pd_result, check_dtype=False, check_index_type=False
)

Expand Down
Loading
Loading
X Tutup