fix more tests

Author: TrevorBergeron

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1030,7 +1030,7 @@ def to_timedelta_op_impl(x: ibis_types.Value, op: ops.ToTimedeltaOp):
 
 @scalar_op_compiler.register_unary_op(ops.timedelta_floor_op)
 def timedelta_floor_op_impl(x: ibis_types.NumericValue):
-    return ibis_api.case().when(x > 0, x.floor()).else_(x.ceil()).end()
+    return ibis_api.case().when(x > ibis.literal(0), x.floor()).else_(x.ceil()).end()
 
 
 @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True)

bigframes/core/utils.py

@@ -113,7 +113,7 @@ def get_standardized_ids(
     """
     col_ids = [
         UNNAMED_COLUMN_ID
-        if col_label is None
+        if pd.isna(col_label)  # type: ignore
         else label_to_identifier(col_label, strict=strict)
         for col_label in col_labels
     ]

bigframes/dataframe.py

@@ -4398,7 +4398,7 @@ def to_excel(
         **kwargs,
     ) -> None:
         return self.to_pandas(allow_large_results=allow_large_results).to_excel(
-            excel_writer, sheet_name, **kwargs
+            excel_writer, sheet_name=sheet_name, **kwargs
         )
 
     def to_latex(

bigframes/series.py

@@ -2292,7 +2292,7 @@ def to_excel(
         self, excel_writer, sheet_name="Sheet1", *, allow_large_results=None, **kwargs
     ) -> None:
         return self.to_pandas(allow_large_results=allow_large_results).to_excel(
-            excel_writer, sheet_name, **kwargs
+            excel_writer, sheet_name=sheet_name, **kwargs
         )
 
     def to_json(

bigframes/testing/utils.py

@@ -91,8 +91,7 @@ def assert_series_equivalent(pd_series: pd.Series, bf_series: bpd.Series, **kwar
 
 
 def _normalize_all_nulls(col: pd.Series) -> pd.Series:
-    # This over-normalizes probably, make more conservative later
-    if col.hasnans and (pd_types.is_float_dtype(col.dtype)):
+    if pd_types.is_float_dtype(col.dtype):
         col = col.astype("float64").astype("Float64")
     return col
 
@@ -172,8 +171,8 @@ def assert_series_equal(
         right.index = right.index.astype("Int64")
 
     if nulls_are_nan:
-        left = _normalize_all_nulls(left)
-        right = _normalize_all_nulls(right)
+        left = _normalize_all_nulls(left.infer_objects())
+        right = _normalize_all_nulls(right.infer_objects())
         left.index = _normalize_index_nulls(left.index)
         right.index = _normalize_index_nulls(right.index)
         left.name = pd.NA if pd.isna(left.name) else left.name  # type: ignore

tests/system/small/bigquery/test_datetime.py

@@ -19,6 +19,7 @@
 import pytest
 
 from bigframes import bigquery
+import bigframes.testing
 
 _TIMESTAMP_DTYPE = pd.ArrowDtype(pa.timestamp("us", tz="UTC"))
 
@@ -40,7 +41,7 @@ def test_unix_seconds(scalars_dfs):
         .apply(lambda ts: _to_unix_epoch(ts, "s"))
         .astype("Int64")
     )
-    pd.testing.assert_series_equal(actual_res, expected_res)
+    bigframes.testing.assert_series_equal(actual_res, expected_res)
 
 
 def test_unix_seconds_after_type_casting(int_series):
@@ -53,7 +54,9 @@ def test_unix_seconds_after_type_casting(int_series):
         .apply(lambda ts: _to_unix_epoch(ts, "s"))
         .astype("Int64")
     )
-    pd.testing.assert_series_equal(actual_res, expected_res, check_index_type=False)
+    bigframes.testing.assert_series_equal(
+        actual_res, expected_res, check_index_type=False
+    )
 
 
 def test_unix_seconds_incorrect_input_type_raise_error(scalars_dfs):
@@ -73,7 +76,7 @@ def test_unix_millis(scalars_dfs):
         .apply(lambda ts: _to_unix_epoch(ts, "ms"))
         .astype("Int64")
     )
-    pd.testing.assert_series_equal(actual_res, expected_res)
+    bigframes.testing.assert_series_equal(actual_res, expected_res)
 
 
 def test_unix_millis_after_type_casting(int_series):
@@ -86,7 +89,9 @@ def test_unix_millis_after_type_casting(int_series):
         .apply(lambda ts: _to_unix_epoch(ts, "ms"))
         .astype("Int64")
     )
-    pd.testing.assert_series_equal(actual_res, expected_res, check_index_type=False)
+    bigframes.testing.assert_series_equal(
+        actual_res, expected_res, check_index_type=False
+    )
 
 
 def test_unix_millis_incorrect_input_type_raise_error(scalars_dfs):
@@ -106,7 +111,7 @@ def test_unix_micros(scalars_dfs):
         .apply(lambda ts: _to_unix_epoch(ts, "us"))
         .astype("Int64")
     )
-    pd.testing.assert_series_equal(actual_res, expected_res)
+    bigframes.testing.assert_series_equal(actual_res, expected_res)
 
 
 def test_unix_micros_after_type_casting(int_series):
@@ -119,7 +124,9 @@ def test_unix_micros_after_type_casting(int_series):
         .apply(lambda ts: _to_unix_epoch(ts, "us"))
         .astype("Int64")
     )
-    pd.testing.assert_series_equal(actual_res, expected_res, check_index_type=False)
+    bigframes.testing.assert_series_equal(
+        actual_res, expected_res, check_index_type=False
+    )
 
 
 def test_unix_micros_incorrect_input_type_raise_error(scalars_dfs):

tests/system/small/bigquery/test_geo.py

@@ -32,6 +32,7 @@
 import bigframes.bigquery as bbq
 import bigframes.geopandas
 import bigframes.session
+import bigframes.testing
 
 
 def test_geo_st_area(session: bigframes.session.Session):
@@ -56,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
 
-    pd.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         geopd_s_result,
         check_dtype=False,
@@ -109,7 +110,7 @@ def test_st_length_various_geometries(session):
 
     # Test default use_spheroid
     result_default = st_length(geoseries).to_pandas()
-    pd.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         result_default,
         expected_lengths,
         rtol=1e-3,
@@ -118,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()
-    pd.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         result_explicit_false,
         expected_lengths,
         rtol=1e-3,
@@ -152,7 +153,7 @@ def test_geo_st_difference_with_geometry_objects(session: bigframes.session.Sess
         index=[0, 1, 2],
         dtype=geopandas.array.GeometryDtype(),
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -191,7 +192,7 @@ def test_geo_st_difference_with_single_geometry_object(
         index=[0, 1, 2],
         dtype=geopandas.array.GeometryDtype(),
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -217,7 +218,7 @@ def test_geo_st_difference_with_similar_geometry_objects(
         index=[0, 1, 2],
         dtype=geopandas.array.GeometryDtype(),
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -273,7 +274,7 @@ def test_geo_st_distance_with_geometry_objects(session: bigframes.session.Sessio
         index=[0, 1, 2, 3],
         dtype="Float64",
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -320,7 +321,7 @@ def test_geo_st_distance_with_single_geometry_object(
         ],
         dtype="Float64",
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -355,7 +356,7 @@ def test_geo_st_intersection_with_geometry_objects(session: bigframes.session.Se
         index=[0, 1, 2],
         dtype=geopandas.array.GeometryDtype(),
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -394,7 +395,7 @@ def test_geo_st_intersection_with_single_geometry_object(
         index=[0, 1, 2],
         dtype=geopandas.array.GeometryDtype(),
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -424,7 +425,7 @@ def test_geo_st_intersection_with_similar_geometry_objects(
         index=[0, 1, 2],
         dtype=geopandas.array.GeometryDtype(),
     )
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         geobf_s_result,
         expected,
         check_index_type=False,
@@ -465,7 +466,7 @@ def test_geo_st_isclosed(session: bigframes.session.Session):
     ]
     expected_series = pd.Series(data=expected_data, dtype="boolean")
 
-    pd.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         bf_result,
         expected_series,
         # We default to Int64 (nullable) dtype, but pandas defaults to int64 index.

tests/system/small/operations/test_dates.py

@@ -17,10 +17,10 @@
 
 from packaging import version
 import pandas as pd
-import pandas.testing
 import pytest
 
 from bigframes import dtypes
+import bigframes.testing
 
 
 def test_date_diff_between_series(session):
@@ -35,7 +35,7 @@ def test_date_diff_between_series(session):
     actual_result = (bf_df["col_1"] - bf_df["col_2"]).to_pandas()
 
     expected_result = (pd_df["col_1"] - pd_df["col_2"]).astype(dtypes.TIMEDELTA_DTYPE)
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
 
@@ -47,7 +47,7 @@ def test_date_diff_literal_sub_series(scalars_dfs):
     actual_result = (literal - bf_df["date_col"]).to_pandas()
 
     expected_result = (literal - pd_df["date_col"]).astype(dtypes.TIMEDELTA_DTYPE)
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
 
@@ -59,7 +59,7 @@ def test_date_diff_series_sub_literal(scalars_dfs):
     actual_result = (bf_df["date_col"] - literal).to_pandas()
 
     expected_result = (pd_df["date_col"] - literal).astype(dtypes.TIMEDELTA_DTYPE)
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
 
@@ -70,7 +70,7 @@ def test_date_series_diff_agg(scalars_dfs):
     actual_result = bf_df["date_col"].diff().to_pandas()
 
     expected_result = pd_df["date_col"].diff().astype(dtypes.TIMEDELTA_DTYPE)
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         actual_result, expected_result, check_index_type=False
     )
 
@@ -86,6 +86,6 @@ def test_date_can_cast_after_accessor(scalars_dfs):
         pd.to_datetime(pd_df["date_col"]).dt.isocalendar().week.astype("Int64")
     )
 
-    pandas.testing.assert_series_equal(
+    bigframes.testing.assert_series_equal(
         actual_result, expected_result, check_dtype=False, check_index_type=False
     )

tests/system/small/operations/test_datetimes.py

@@ -576,9 +576,12 @@ def test_timestamp_diff_series_sub_literal(scalars_dfs, column, value):
     bf_series = bf_df[column]
     pd_series = pd_df[column]
 
-    actual_result = (bf_series - value).to_pandas()
+    # Pandas doesn't handle nulls properly here so we ffill
+    # overflows for no good reason
+    # related? https://github.com/apache/arrow/issues/43031
+    actual_result = (bf_series.ffill() - value).to_pandas()
 
-    expected_result = pd_series - value
+    expected_result = pd_series.ffill() - value
     assert_series_equal(actual_result, expected_result)
 
 
@@ -594,9 +597,12 @@ def test_timestamp_diff_literal_sub_series(scalars_dfs, column, value):
     bf_series = bf_df[column]
     pd_series = pd_df[column]
 
-    actual_result = (value - bf_series).to_pandas()
+    # Pandas doesn't handle nulls properly here so we ffill
+    # overflows for no good reason
+    # related? https://github.com/apache/arrow/issues/43031
+    actual_result = (value - bf_series.ffill()).to_pandas()
 
-    expected_result = value - pd_series
+    expected_result = value - pd_series.ffill()
     assert_series_equal(actual_result, expected_result)
 
 

tests/system/small/test_dataframe.py

@@ -778,6 +778,8 @@ def test_drop_bigframes_index(scalars_dfs):
 
 
 def test_drop_bigframes_index_with_na(scalars_dfs):
+    if pd.__version__.startswith("3"):
+        pytest.skip("Pandas 3.0 doesn't doesn't support drop with pd.NA values")
     scalars_df, scalars_pandas_df = scalars_dfs
     scalars_df = scalars_df.copy()
     scalars_pandas_df = scalars_pandas_df.copy()