fix: Ensure binops with pandas objects returns bigquery dataframes (#1404)

Author: TrevorBergeron

bigframes/core/indexes/base.py

@@ -48,6 +48,8 @@ class Index(vendored_pandas_index.Index):
  _linked_frame: Union[
  bigframes.dataframe.DataFrame, bigframes.series.Series, None
  ] = None
+ # Must be above 5000 for pandas to delegate to bigframes for binops
+ __pandas_priority__ = 12000
 
  # Overrided on __new__ to create subclasses like pandas does
  def __new__(

bigframes/dataframe.py

@@ -103,6 +103,8 @@ class DataFrame(vendored_pandas_frame.DataFrame):
  __doc__ = vendored_pandas_frame.DataFrame.__doc__
  # internal flag to disable cache at all
  _disable_cache_override: bool = False
+ # Must be above 5000 for pandas to delegate to bigframes for binops
+ __pandas_priority__ = 15000
 
  def __init__(
  self,

bigframes/series.py

@@ -78,6 +78,9 @@
 
 @log_adapter.class_logger
 class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Series):
+ # Must be above 5000 for pandas to delegate to bigframes for binops
+ __pandas_priority__ = 13000
+
  def __init__(self, *args, **kwargs):
  self._query_job: Optional[bigquery.QueryJob] = None
  super().__init__(*args, **kwargs)

tests/system/small/test_dataframe.py

@@ -2471,6 +2471,20 @@ def test_listlike_binop_axis_1_in_memory_data(scalars_dfs, input):
  assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
 
 
+@skip_legacy_pandas
+def test_df_reverse_binop_pandas(scalars_dfs):
+ scalars_df, scalars_pandas_df = scalars_dfs
+
+ pd_series = pd.Series([100, 200, 300])
+
+ df_columns = ["int64_col", "float64_col", "int64_too"]
+
+ bf_result = pd_series + scalars_df[df_columns].to_pandas()
+ pd_result = pd_series + scalars_pandas_df[df_columns]
+
+ assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
+
+
 def test_listlike_binop_axis_1_bf_index(scalars_dfs):
  scalars_df, scalars_pandas_df = scalars_dfs
 

tests/system/small/test_series.py

@@ -1656,6 +1656,27 @@ def test_series_binop_w_other_types(scalars_dfs, other):
  )
 
 
+@pytest.mark.parametrize(
+ ("other",),
+ [
+ ([-1.4, 2.3, None],),
+ (pd.Index([-1.4, 2.3, None]),),
+ (pd.Series([-1.4, 2.3, None], index=[44, 2, 1]),),
+ ],
+)
+@skip_legacy_pandas
+def test_series_reverse_binop_w_other_types(scalars_dfs, other):
+ scalars_df, scalars_pandas_df = scalars_dfs
+
+ bf_result = (other + scalars_df["int64_col"].head(3)).to_pandas()
+ pd_result = other + scalars_pandas_df["int64_col"].head(3)
+
+ assert_series_equal(
+ bf_result,
+ pd_result,
+ )
+
+
 @skip_legacy_pandas
 def test_series_combine_first(scalars_dfs):
  scalars_df, scalars_pandas_df = scalars_dfs