feat: Add DataFrame.mask method (#1302)

Author: TrevorBergeron

bigframes/dataframe.py

@@ -2365,6 +2365,9 @@ def where(self, cond, other=None):
  result.columns.names = self.columns.names
  return result
 
+ def mask(self, cond, other=None):
+ return self.where(~cond, other=other)
+
  def dropna(
  self,
  *,

tests/system/small/test_dataframe.py

@@ -331,6 +331,17 @@ def test_where_series_cond(scalars_df_index, scalars_pandas_df_index):
  pandas.testing.assert_frame_equal(bf_result, pd_result)
 
 
+def test_mask_series_cond(scalars_df_index, scalars_pandas_df_index):
+ cond_bf = scalars_df_index["int64_col"] > 0
+ cond_pd = scalars_pandas_df_index["int64_col"] > 0
+
+ bf_df = scalars_df_index[["int64_too", "int64_col", "float64_col"]]
+ pd_df = scalars_pandas_df_index[["int64_too", "int64_col", "float64_col"]]
+ bf_result = bf_df.mask(cond_bf, bf_df + 1).to_pandas()
+ pd_result = pd_df.mask(cond_pd, pd_df + 1)
+ pandas.testing.assert_frame_equal(bf_result, pd_result)
+
+
 def test_where_series_multi_index(scalars_df_index, scalars_pandas_df_index):
  # Test when a dataframe has multi-index or multi-columns.
  columns = ["int64_col", "float64_col"]

third_party/bigframes_vendored/pandas/core/frame.py

@@ -2048,6 +2048,98 @@ def where(self, cond, other):
  """
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+ def mask(self, cond, other):
+ """Replace values where the condition is False.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> bpd.options.display.progress_bar = None
+
+ >>> df = bpd.DataFrame({'a': [20, 10, 0], 'b': [0, 10, 20]})
+ >>> df
+ a b
+ 0 20 0
+ 1 10 10
+ 2 0 20
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ You can filter the values in the dataframe based on a condition. The
+ values matching the condition would be kept, and not matching would be
+ replaced. The default replacement value is ``NA``. For example, when the
+ condition is a dataframe:
+
+ >>> df.mask(df > 0)
+ a b
+ 0 <NA> 0
+ 1 <NA> <NA>
+ 2 0 <NA>
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ You can specify a custom replacement value for non-matching values.
+
+ >>> df.mask(df > 0, -1)
+ a b
+ 0 -1 0
+ 1 -1 -1
+ 2 0 -1
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ Besides dataframe, the condition can be a series too. For example:
+
+ >>> df.mask(df['a'] > 10, -1)
+ a b
+ 0 -1 -1
+ 1 10 10
+ 2 0 20
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ As for the replacement, it can be a dataframe too. For example:
+
+ >>> df.mask(df > 10, -df)
+ a b
+ 0 -20 0
+ 1 10 10
+ 2 0 -20
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ >>> df.mask(df['a'] > 10, -df)
+ a b
+ 0 -20 0
+ 1 10 10
+ 2 0 20
+ <BLANKLINE>
+ [3 rows x 2 columns]
+
+ Please note, replacement doesn't support Series for now. In pandas, when
+ specifying a Series as replacement, the axis value should be specified
+ at the same time, which is not supported in bigframes DataFrame.
+
+ Args:
+ cond (bool Series/DataFrame, array-like, or callable):
+ Where cond is False, keep the original value. Where True, replace
+ with corresponding value from other. If cond is callable, it is
+ computed on the Series/DataFrame and returns boolean
+ Series/DataFrame or array. The callable must not change input
+ Series/DataFrame (though pandas doesn’t check it).
+ other (scalar, DataFrame, or callable):
+ Entries where cond is True are replaced with corresponding value
+ from other. If other is callable, it is computed on the
+ DataFrame and returns scalar or DataFrame. The callable must not
+ change input DataFrame (though pandas doesn’t check it). If not
+ specified, entries will be filled with the corresponding NULL
+ value (np.nan for numpy dtypes, pd.NA for extension dtypes).
+
+ Returns:
+ DataFrame: DataFrame after the replacement.
+ """
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
  # ----------------------------------------------------------------------
  # Sorting