feat: Add DataFrame.round method (#1742)

* feat: Add DataFrame.round method * docstring full typename * restrict round test to pandas 2.x * add type hint

Author: TrevorBergeron

bigframes/core/compile/scalar_op_compiler.py

@@ -1826,6 +1826,13 @@ def fillna_op(
 
 @scalar_op_compiler.register_binary_op(ops.round_op)
 def round_op(x: ibis_types.Value, y: ibis_types.Value):
+ if x.type().is_integer():
+ # bq produces float64, but pandas returns int
+ return (
+ typing.cast(ibis_types.NumericValue, x)
+ .round(digits=typing.cast(ibis_types.IntegerValue, y))
+ .cast(ibis_dtypes.int64)
+ )
  return typing.cast(ibis_types.NumericValue, x).round(
  digits=typing.cast(ibis_types.IntegerValue, y)
  )

bigframes/dataframe.py

@@ -27,6 +27,7 @@
 from typing import (
  Callable,
  Dict,
+ Hashable,
  Iterable,
  List,
  Literal,
@@ -3608,6 +3609,47 @@ def _groupby_series(
  def abs(self) -> DataFrame:
  return self._apply_unary_op(ops.abs_op)
 
+ def round(self, decimals: Union[int, dict[Hashable, int]] = 0) -> DataFrame:
+ is_mapping = utils.is_dict_like(decimals)
+ if not (is_mapping or isinstance(decimals, int)):
+ raise TypeError("'decimals' must be either a dict-like or integer.")
+ block = self._block
+ exprs = []
+ for label, col_id, dtype in zip(
+ block.column_labels, block.value_columns, block.dtypes
+ ):
+ if dtype in set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE) - {
+ bigframes.dtypes.BOOL_DTYPE
+ }:
+ if is_mapping:
+ if label in decimals: # type: ignore
+ exprs.append(
+ ops.round_op.as_expr(
+ col_id,
+ ex.const(
+ decimals[label], dtype=bigframes.dtypes.INT_DTYPE # type: ignore
+ ),
+ )
+ )
+ else:
+ exprs.append(ex.deref(col_id))
+ else:
+ exprs.append(
+ ops.round_op.as_expr(
+ col_id,
+ ex.const(
+ typing.cast(int, decimals),
+ dtype=bigframes.dtypes.INT_DTYPE,
+ ),
+ )
+ )
+ else:
+ exprs.append(ex.deref(col_id))
+
+ return DataFrame(
+ block.project_exprs(exprs, labels=block.column_labels, drop=True)
+ )
+
  def isna(self) -> DataFrame:
  return self._apply_unary_op(ops.isnull_op)
 

bigframes/operations/numeric_ops.py

@@ -289,7 +289,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 )
 
 round_op = base_ops.create_binary_op(
- name="round", type_signature=op_typing.BINARY_REAL_NUMERIC
+ name="round", type_signature=op_typing.BINARY_NUMERIC
 )
 
 unsafe_pow_op = base_ops.create_binary_op(

tests/system/small/test_dataframe.py

@@ -1635,6 +1635,25 @@ def test_merge_left_on_right_on(scalars_dfs, merge_how):
  )
 
 
+@pytest.mark.parametrize(
+ ("decimals",),
+ [
+ (2,),
+ ({"float64_col": 0, "bool_col": 1, "int64_too": -3},),
+ ({},),
+ ],
+)
+def test_dataframe_round(scalars_dfs, decimals):
+ if pd.__version__.startswith("1."):
+ pytest.skip("Rounding doesn't work as expected in pandas 1.x")
+ scalars_df, scalars_pandas_df = scalars_dfs
+
+ bf_result = scalars_df.round(decimals).to_pandas()
+ pd_result = scalars_pandas_df.round(decimals)
+
+ assert_pandas_df_equal(bf_result, pd_result)
+
+
 def test_get_dtypes(scalars_df_default_index):
  dtypes = scalars_df_default_index.dtypes
  dtypes_dict: Dict[str, bigframes.dtypes.Dtype] = {

third_party/bigframes_vendored/pandas/core/frame.py

@@ -4788,6 +4788,83 @@ def merge(
  """
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+ def round(self, decimals):
+ """
+ Round a DataFrame to a variable number of decimal places.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> bpd.options.display.progress_bar = None
+ >>> df = bpd.DataFrame([(.21, .32), (.01, .67), (.66, .03), (.21, .18)],
+ ... columns=['dogs', 'cats'])
+ >>> df
+ dogs cats
+ 0 0.21 0.32
+ 1 0.01 0.67
+ 2 0.66 0.03
+ 3 0.21 0.18
+ <BLANKLINE>
+ [4 rows x 2 columns]
+
+ By providing an integer each column is rounded to the same number
+ of decimal places
+
+ >>> df.round(1)
+ dogs cats
+ 0 0.2 0.3
+ 1 0.0 0.7
+ 2 0.7 0.0
+ 3 0.2 0.2
+ <BLANKLINE>
+ [4 rows x 2 columns]
+
+ With a dict, the number of places for specific columns can be
+ specified with the column names as key and the number of decimal
+ places as value
+
+ >>> df.round({'dogs': 1, 'cats': 0})
+ dogs cats
+ 0 0.2 0.0
+ 1 0.0 1.0
+ 2 0.7 0.0
+ 3 0.2 0.0
+ <BLANKLINE>
+ [4 rows x 2 columns]
+
+ Using a Series, the number of places for specific columns can be
+ specified with the column names as index and the number of
+ decimal places as value
+
+ >>> decimals = pd.Series([0, 1], index=['cats', 'dogs'])
+ >>> df.round(decimals)
+ dogs cats
+ 0 0.2 0.0
+ 1 0.0 1.0
+ 2 0.7 0.0
+ 3 0.2 0.0
+ <BLANKLINE>
+ [4 rows x 2 columns]
+
+ Args:
+ decimals (int, dict, Series):
+ Number of decimal places to round each column to. If an int is
+ given, round each column to the same number of places.
+ Otherwise dict and Series round to variable numbers of places.
+ Column names should be in the keys if `decimals` is a
+ dict-like, or in the index if `decimals` is a Series. Any
+ columns not included in `decimals` will be left as is. Elements
+ of `decimals` which are not columns of the input will be
+ ignored.
+
+ Returns:
+ bigframes.pandas.DataFrame:
+ A DataFrame with the affected columns rounded to the specified
+ number of decimal places.
+
+ """
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
  def apply(self, func, *, axis=0, args=(), **kwargs):
  """Apply a function along an axis of the DataFrame.