test: Fix prerelease and pandas 3.0 test compat

Author: TrevorBergeron

bigframes/testing/__init__.py

@@ -17,3 +17,9 @@
 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_series_equal
+
+__all__ = [
+    "assert_frame_equal",
+    "assert_series_equal",
+]

noxfile.py

@@ -20,7 +20,6 @@
 import multiprocessing
 import os
 import pathlib
-import re
 import shutil
 import time
 from typing import Dict, List
@@ -588,99 +587,36 @@ def prerelease(session: nox.sessions.Session, tests_path, extra_pytest_options=(
     constraints_path = str(
         CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt"
     )
-
-    # Ignore officially released versions of certain packages specified in
-    # testing/constraints-*.txt and install a more recent, pre-release versions
-    # directly
-    already_installed = set()
+    session.install(
+        *set(UNIT_TEST_STANDARD_DEPENDENCIES + SYSTEM_TEST_STANDARD_DEPENDENCIES),
+        "-c",
+        constraints_path,
+        "-e",
+        ".",
+    )
 
     # PyArrow prerelease packages are published to an alternative PyPI host.
     # https://arrow.apache.org/docs/python/install.html#installing-nightly-packages
     session.install(
+        "--no-deps",
+        "--upgrade",
         "--extra-index-url",
         "https://pypi.fury.io/arrow-nightlies/",
-        "--prefer-binary",
-        "--pre",
-        "--upgrade",
         "pyarrow",
-    )
-    already_installed.add("pyarrow")
-
-    session.install(
-        "--prefer-binary",
-        "--pre",
-        "--upgrade",
         # We exclude each version individually so that we can continue to test
         # some prerelease packages. See:
         # https://github.com/googleapis/python-bigquery-dataframes/pull/268#discussion_r1423205172
         # "pandas!=2.1.4, !=2.2.0rc0, !=2.2.0, !=2.2.1",
         "pandas",
-    )
-    already_installed.add("pandas")
-
-    # Try to avoid a cap on our SQLGlot so that bigframes
-    # can be integrated with SQLMesh. See:
-    # https://github.com/googleapis/python-bigquery-dataframes/issues/942
-    # If SQLGlot introduces something that breaks us, lets file an issue
-    # upstream and/or make sure we fix bigframes to work with it.
-    session.install(
-        "--upgrade",
-        "git+https://github.com/tobymao/sqlglot.git#egg=sqlglot",
-    )
-    already_installed.add("sqlglot")
-
-    # Workaround https://github.com/googleapis/python-db-dtypes-pandas/issues/178
-    session.install("--no-deps", "db-dtypes")
-    already_installed.add("db-dtypes")
-
-    # Ensure we catch breaking changes in the client libraries early.
-    session.install(
-        "--upgrade",
+        # Workaround https://github.com/googleapis/python-db-dtypes-pandas/issues/178
+        "db-dtypes",
+        # Ensure we catch breaking changes in the client libraries early.
         "git+https://github.com/googleapis/python-bigquery.git#egg=google-cloud-bigquery",
-    )
-    already_installed.add("google-cloud-bigquery")
-    session.install(
         "--upgrade",
         "-e",
         "git+https://github.com/googleapis/google-cloud-python.git#egg=google-cloud-bigquery-storage&subdirectory=packages/google-cloud-bigquery-storage",
-    )
-    already_installed.add("google-cloud-bigquery-storage")
-    session.install(
-        "--upgrade",
         "git+https://github.com/googleapis/python-bigquery-pandas.git#egg=pandas-gbq",
     )
-    already_installed.add("pandas-gbq")
-
-    session.install(
-        *set(UNIT_TEST_STANDARD_DEPENDENCIES + SYSTEM_TEST_STANDARD_DEPENDENCIES),
-        "-c",
-        constraints_path,
-    )
-
-    # Because we test minimum dependency versions on the minimum Python
-    # version, the first version we test with in the unit tests sessions has a
-    # constraints file containing all dependencies and extras.
-    with open(
-        CURRENT_DIRECTORY / "testing" / f"constraints-{DEFAULT_PYTHON_VERSION}.txt",
-        encoding="utf-8",
-    ) as constraints_file:
-        constraints_text = constraints_file.read()
-
-    # Ignore leading whitespace and comment lines.
-    deps = [
-        match.group(1)
-        for match in re.finditer(
-            r"^\s*(\S+)(?===\S+)", constraints_text, flags=re.MULTILINE
-        )
-        if match.group(1) not in already_installed
-    ]
-
-    print(already_installed)
-
-    # We use --no-deps to ensure that pre-release versions aren't overwritten
-    # by the version ranges in setup.py.
-    session.install(*deps)
-    session.install("--no-deps", "-e", ".")
 
     # Print out prerelease package versions.
     session.run("python", "-m", "pip", "freeze")

tests/system/small/bigquery/test_array.py

@@ -67,7 +67,11 @@
 )
 def test_array_length(input_data, expected):
     series = bpd.Series(input_data)
-    expected = pd.Series(expected, dtype=bigframes.dtypes.INT_DTYPE)
+    expected = pd.Series(
+        expected,
+        index=pd.Index(range(len(input_data)), dtype="Int64"),
+        dtype=bigframes.dtypes.INT_DTYPE,
+    )
     pd.testing.assert_series_equal(
         bbq.array_length(series).to_pandas(),
         expected,

tests/system/small/bigquery/test_sql.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pandas as pd
 import pytest
 
 import bigframes.bigquery as bbq
 import bigframes.dtypes as dtypes
 import bigframes.pandas as bpd
+import bigframes.testing
 
 
 def test_sql_scalar_for_all_scalar_types(scalars_df_null_index):
@@ -59,8 +59,8 @@ def test_sql_scalar_for_bool_series(scalars_df_index):
     series: bpd.Series = scalars_df_index["bool_col"]
     result = bbq.sql_scalar("CAST({0} AS INT64)", [series])
     expected = series.astype(dtypes.INT_DTYPE)
-    expected.name = None
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    expected.name = result.name
+    bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
 
 
 @pytest.mark.parametrize(
@@ -83,8 +83,8 @@ def test_sql_scalar_outputs_all_scalar_types(scalars_df_index, column_name):
     series: bpd.Series = scalars_df_index[column_name]
     result = bbq.sql_scalar("{0}", [series])
     expected = series
-    expected.name = None
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    expected.name = result.name
+    bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
 
 
 def test_sql_scalar_for_array_series(repeated_df):
@@ -114,14 +114,14 @@ def test_sql_scalar_for_array_series(repeated_df):
         + repeated_df["numeric_list_col"].list.len()
         + repeated_df["string_list_col"].list.len()
     )
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    bigframes.testing.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 = None
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    expected.name = result.name
+    bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
 
 
 def test_sql_scalar_for_struct_series(nested_structs_df):
@@ -132,14 +132,14 @@ 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")
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    bigframes.testing.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 = None
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    expected.name = result.name
+    bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
 
 
 def test_sql_scalar_for_json_series(json_df):
@@ -150,12 +150,12 @@ def test_sql_scalar_for_json_series(json_df):
         ],
     )
     expected = bbq.json_value(json_df["json_col"], "$.int_value")
-    expected.name = None
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    expected.name = result.name
+    bigframes.testing.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 = None
-    pd.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())
+    expected.name = result.name
+    bigframes.testing.assert_series_equal(result.to_pandas(), expected.to_pandas())