fix transpose test

Author: TrevorBergeron

bigframes/testing/utils.py

@@ -73,7 +73,7 @@
 def assert_dfs_equivalent(pd_df: pd.DataFrame, bf_df: bpd.DataFrame, **kwargs):
     bf_df_local = bf_df.to_pandas()
     ignore_order = not bf_df._session._strictly_ordered
-    assert_pandas_df_equal(bf_df_local, pd_df, ignore_order=ignore_order, **kwargs)
+    assert_frame_equal(bf_df_local, pd_df, ignore_order=ignore_order, **kwargs)
 
 
 def assert_series_equivalent(pd_series: pd.Series, bf_series: bpd.Series, **kwargs):
@@ -82,21 +82,40 @@ def assert_series_equivalent(pd_series: pd.Series, bf_series: bpd.Series, **kwar
     assert_series_equal(bf_df_local, pd_series, ignore_order=ignore_order, **kwargs)
 
 
-def assert_pandas_df_equal(df0, df1, ignore_order: bool = False, **kwargs):
+def _normalize_all_nulls(col: pd.Series) -> pd.Series:
+    if col.dtype == bigframes.dtypes.FLOAT_DTYPE:
+        col = col.astype("float64")
+    if pd_types.is_object_dtype(col):
+        col = col.fillna(float("nan"))
+    return col
+
+
+def assert_frame_equal(
+    left: pd.DataFrame,
+    right: pd.DataFrame,
+    *,
+    ignore_order: bool = False,
+    nulls_are_nan: bool = True,
+    **kwargs,
+):
     if ignore_order:
         # Sort by a column to get consistent results.
-        if df0.index.name != "rowindex":
-            df0 = df0.sort_values(
-                list(df0.columns.drop("geography_col", errors="ignore"))
+        if left.index.name != "rowindex":
+            left = left.sort_values(
+                list(left.columns.drop("geography_col", errors="ignore"))
             ).reset_index(drop=True)
-            df1 = df1.sort_values(
-                list(df1.columns.drop("geography_col", errors="ignore"))
+            right = right.sort_values(
+                list(right.columns.drop("geography_col", errors="ignore"))
             ).reset_index(drop=True)
         else:
-            df0 = df0.sort_index()
-            df1 = df1.sort_index()
+            left = left.sort_index()
+            right = right.sort_index()
+
+    if nulls_are_nan:
+        left = left.apply(_normalize_all_nulls)
+        right = right.apply(_normalize_all_nulls)
 
-    pd.testing.assert_frame_equal(df0, df1, **kwargs)
+    pd.testing.assert_frame_equal(left, right, **kwargs)
 
 
 def assert_series_equal(
@@ -116,14 +135,8 @@ def assert_series_equal(
             right = right.sort_index()
 
     if nulls_are_nan:
-        if left.dtype == bigframes.dtypes.FLOAT_DTYPE:
-            left = left.astype("float64")
-        if right.dtype == bigframes.dtypes.FLOAT_DTYPE:
-            right = right.astype("float64")
-        if pd_types.is_object_dtype(left):
-            left = left.fillna(float("nan"))
-        if pd_types.is_object_dtype(right):
-            right = right.fillna(float("nan"))
+        left = _normalize_all_nulls(left)
+        right = _normalize_all_nulls(right)
 
     pd.testing.assert_series_equal(left, right, **kwargs)
 

tests/system/large/functions/test_remote_function.py

@@ -36,7 +36,7 @@
 import bigframes.pandas as bpd
 import bigframes.series
 from bigframes.testing.utils import (
-    assert_pandas_df_equal,
+    assert_frame_equal,
     cleanup_function_assets,
     delete_cloud_function,
     get_cloud_functions,
@@ -214,7 +214,7 @@ def square(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(square, session.bqclient, session.cloudfunctionsclient)
@@ -261,7 +261,7 @@ def add_one(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -349,7 +349,7 @@ def square(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(square, session.bqclient, session.cloudfunctionsclient)
@@ -403,7 +403,7 @@ def sign(num):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -453,7 +453,7 @@ def circumference(radius):
         pd_result_col = pd_result_col.astype(pandas.Float64Dtype())
         pd_result = pd_float64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -503,7 +503,7 @@ def find_team(num):
         pd_result_col = pd_result_col.astype(pandas.StringDtype(storage="pyarrow"))
         pd_result = pd_float64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -591,7 +591,7 @@ def inner_test():
             pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
             pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-            assert_pandas_df_equal(bf_result, pd_result)
+            assert_frame_equal(bf_result, pd_result)
 
         # Test that the remote function works as expected
         inner_test()
@@ -683,7 +683,7 @@ def is_odd(num):
         pd_result_col = pd_int64_col.mask(is_odd)
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -727,7 +727,7 @@ def is_odd(num):
         pd_result_col = pd_int64_col[pd_int64_col.notnull()].mask(is_odd, -1)
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -770,7 +770,7 @@ def test_remote_udf_lambda(session, scalars_dfs, dataset_id, bq_cf_connection):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -829,7 +829,7 @@ def square(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -884,7 +884,7 @@ def pd_np_foo(x) -> None:
         # comparing for the purpose of this test
         pd_result.result = pd_result.result.astype(pandas.Float64Dtype())
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -928,7 +928,7 @@ def test_internal(rf, udf):
             pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
             pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-            assert_pandas_df_equal(bf_result, pd_result)
+            assert_frame_equal(bf_result, pd_result)
 
         # Create an explicit name for the remote function
         prefixer = test_utils.prefixer.Prefixer("foo", "")
@@ -1109,7 +1109,7 @@ def square(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(square, session.bqclient, session.cloudfunctionsclient)
@@ -1150,7 +1150,7 @@ def square(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(square, session.bqclient, session.cloudfunctionsclient)
@@ -1225,7 +1225,7 @@ def square(x):
         pd_result_col = pd_result_col.astype(pandas.Int64Dtype())
         pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result)
+        assert_frame_equal(bf_result, pd_result)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(square, session.bqclient, session.cloudfunctionsclient)
@@ -1283,7 +1283,7 @@ def square_num(x):
         pd_result_col = pd_int64_col.apply(lambda x: x if x is None else x * x)
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
+        assert_frame_equal(bf_result, pd_result, check_dtype=False)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -1357,7 +1357,7 @@ def square_num(x):
         pd_result_col = pd_int64_col.apply(lambda x: x if x is None else x * x)
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
+        assert_frame_equal(bf_result, pd_result, check_dtype=False)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(
@@ -1416,7 +1416,7 @@ def square_num(x):
         pd_result_col = df["num"].apply(lambda x: x if x is None else x * x)
         pd_result = df.assign(result=pd_result_col)
 
-        assert_pandas_df_equal(
+        assert_frame_equal(
             bf_result, pd_result, check_dtype=False, check_index_type=False
         )
 
@@ -1504,7 +1504,7 @@ def square_num(x):
         pd_result_col = pd_int64_col.apply(square_num)
         pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
 
-        assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
+        assert_frame_equal(bf_result, pd_result, check_dtype=False)
     finally:
         # clean up the gcp assets created for the remote function
         cleanup_function_assets(

tests/system/small/bigquery/test_vector_search.py

@@ -23,7 +23,7 @@
 
 import bigframes.bigquery as bbq
 import bigframes.pandas as bpd
-from bigframes.testing.utils import assert_pandas_df_equal
+from bigframes.testing.utils import assert_frame_equal
 
 # Need at least 5,000 rows to create a vector index.
 VECTOR_DF = pd.DataFrame(
@@ -154,7 +154,7 @@ def test_vector_search_basic_params_with_df():
         },
         index=pd.Index([1, 0, 0, 1], dtype="Int64"),
     )
-    assert_pandas_df_equal(
+    assert_frame_equal(
         expected.sort_values("id"),
         vector_search_result.sort_values("id"),
         check_dtype=False,

tests/system/small/functions/test_remote_function.py

@@ -34,7 +34,7 @@
 from bigframes.functions import _utils as bff_utils
 from bigframes.functions import function as bff
 import bigframes.session._io.bigquery
-from bigframes.testing.utils import assert_pandas_df_equal, get_function_name
+from bigframes.testing.utils import assert_frame_equal, get_function_name
 
 _prefixer = test_utils.prefixer.Prefixer("bigframes", "")
 
@@ -159,7 +159,7 @@ def square(x):
     pd_result_col = pd_result_col.astype(pd.Int64Dtype())
     pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -208,7 +208,7 @@ def square(x):
     pd_result_col = pd_result_col.astype(pd.Int64Dtype())
     pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -300,7 +300,7 @@ def square(x):
     pd_result_col = pd_result_col.astype(pd.Int64Dtype())
     pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -388,7 +388,7 @@ def square(x):
     pd_result_col = pd_result_col.astype(pd.Int64Dtype())
     pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -437,7 +437,7 @@ def square(x):
     pd_result_col = pd_result_col.astype(pd.Int64Dtype())
     pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -482,7 +482,7 @@ def square(x):
     pd_result_col = pd_result_col.astype(pd.Int64Dtype())
     pd_result = pd_int64_col_filtered.to_frame().assign(result=pd_result_col)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -517,7 +517,7 @@ def add_one(x):
     for col in pd_result:
         pd_result[col] = pd_result[col].astype(pd_int64_df_filtered[col].dtype)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -552,7 +552,7 @@ def add_one(x):
     for col in pd_result:
         pd_result[col] = pd_result[col].astype(pd_int64_df_filtered[col].dtype)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -585,7 +585,7 @@ def add_one(x):
     for col in pd_result:
         pd_result[col] = pd_result[col].astype(pd_int64_df[col].dtype)
 
-    assert_pandas_df_equal(bf_result, pd_result)
+    assert_frame_equal(bf_result, pd_result)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
@@ -738,7 +738,7 @@ def square1(x):
     s2_result_col = int64_col_filtered.apply(square2)
     s2_result = int64_col_filtered.to_frame().assign(result=s2_result_col)
 
-    assert_pandas_df_equal(s1_result.to_pandas(), s2_result.to_pandas())
+    assert_frame_equal(s1_result.to_pandas(), s2_result.to_pandas())
 
 
 def test_read_gbq_function_runs_existing_udf(session):
@@ -937,7 +937,7 @@ def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id):
         indirect_df = indirect_df.assign(y=indirect_df.x.apply(square))
         converted_indirect_df = indirect_df.to_pandas()
 
-        assert_pandas_df_equal(
+        assert_frame_equal(
             direct_df, converted_indirect_df, ignore_order=True, check_index_type=False
         )