feat: Add groupby.rank() (#1433)

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -26,6 +26,7 @@
 import bigframes.core.expression as ex
 import bigframes.core.ordering as ordering
 import bigframes.core.window_spec as windows
+import bigframes.dtypes
 import bigframes.dtypes as dtypes
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
@@ -409,6 +410,8 @@ def rank(
  method: str = "average",
  na_option: str = "keep",
  ascending: bool = True,
+ grouping_cols: tuple[str, ...] = (),
+ columns: tuple[str, ...] = (),
 ):
  if method not in ["average", "min", "max", "first", "dense"]:
  raise ValueError(
@@ -417,8 +420,8 @@ def rank(
  if na_option not in ["keep", "top", "bottom"]:
  raise ValueError("na_option must be one of 'keep', 'top', or 'bottom'")
 
- columns = block.value_columns
- labels = block.column_labels
+ columns = columns or tuple(col for col in block.value_columns)
+ labels = [block.col_id_to_label[id] for id in columns]
  # Step 1: Calculate row numbers for each row
  # Identify null values to be treated according to na_option param
  rownum_col_ids = []
@@ -442,9 +445,13 @@ def rank(
  block, rownum_id = block.apply_window_op(
  col if na_option == "keep" else nullity_col_id,
  agg_ops.dense_rank_op if method == "dense" else agg_ops.count_op,
- window_spec=windows.unbound(ordering=window_ordering)
+ window_spec=windows.unbound(
+ grouping_keys=grouping_cols, ordering=window_ordering
+ )
  if method == "dense"
- else windows.rows(following=0, ordering=window_ordering),
+ else windows.rows(
+ following=0, ordering=window_ordering, grouping_keys=grouping_cols
+ ),
  skip_reproject_unsafe=(col != columns[-1]),
  )
  rownum_col_ids.append(rownum_id)
@@ -462,12 +469,32 @@ def rank(
  block, result_id = block.apply_window_op(
  rownum_col_ids[i],
  agg_op,
- window_spec=windows.unbound(grouping_keys=(columns[i],)),
+ window_spec=windows.unbound(grouping_keys=(columns[i], *grouping_cols)),
  skip_reproject_unsafe=(i < (len(columns) - 1)),
  )
  post_agg_rownum_col_ids.append(result_id)
  rownum_col_ids = post_agg_rownum_col_ids
 
+ # Pandas masks all values where any grouping column is null
+ # Note: we use pd.NA instead of float('nan')
+ if grouping_cols:
+ predicate = functools.reduce(
+ ops.and_op.as_expr,
+ [ops.notnull_op.as_expr(column_id) for column_id in grouping_cols],
+ )
+ block = block.project_exprs(
+ [
+ ops.where_op.as_expr(
+ ex.deref(col),
+ predicate,
+ ex.const(None),
+ )
+ for col in rownum_col_ids
+ ],
+ labels=labels,
+ )
+ rownum_col_ids = list(block.value_columns[-len(rownum_col_ids) :])
+
  # Step 3: post processing: mask null values and cast to float
  if method in ["min", "max", "first", "dense"]:
  # Pandas rank always produces Float64, so must cast for aggregation types that produce ints

bigframes/core/compile/ibis_types.py

@@ -397,6 +397,7 @@ def literal_to_ibis_scalar(
  )
  # "correct" way would be to use ibis.array, but this produces invalid BQ SQL syntax
  return tuple(literal)
+
  if not pd.api.types.is_list_like(literal) and pd.isna(literal):
  if ibis_dtype:
  return bigframes_vendored.ibis.null().cast(ibis_dtype)

bigframes/core/groupby/__init__.py

@@ -174,6 +174,20 @@ def median(self, numeric_only: bool = False, *, exact: bool = True) -> df.DataFr
  return self.quantile(0.5)
  return self._aggregate_all(agg_ops.median_op, numeric_only=True)
 
+ def rank(
+ self, method="average", ascending: bool = True, na_option: str = "keep"
+ ) -> df.DataFrame:
+ return df.DataFrame(
+ block_ops.rank(
+ self._block,
+ method,
+ na_option,
+ ascending,
+ grouping_cols=tuple(self._by_col_ids),
+ columns=tuple(self._selected_cols),
+ )
+ )
+
  def quantile(
  self, q: Union[float, Sequence[float]] = 0.5, *, numeric_only: bool = False
  ) -> df.DataFrame:
@@ -574,6 +588,20 @@ def sum(self, *args) -> series.Series:
  def mean(self, *args) -> series.Series:
  return self._aggregate(agg_ops.mean_op)
 
+ def rank(
+ self, method="average", ascending: bool = True, na_option: str = "keep"
+ ) -> series.Series:
+ return series.Series(
+ block_ops.rank(
+ self._block,
+ method,
+ na_option,
+ ascending,
+ grouping_cols=tuple(self._by_col_ids),
+ columns=(self._value_column,),
+ )
+ )
+
  def median(
  self,
  *args,

tests/system/small/test_groupby.py

@@ -16,7 +16,7 @@
 import pytest
 
 import bigframes.pandas as bpd
-from tests.system.utils import assert_pandas_df_equal
+from tests.system.utils import assert_pandas_df_equal, skip_legacy_pandas
 
 # =================
 # DataFrame.groupby
@@ -94,6 +94,72 @@ def test_dataframe_groupby_quantile(scalars_df_index, scalars_pandas_df_index, q
  )
 
 
+@skip_legacy_pandas
+@pytest.mark.parametrize(
+ ("na_option", "method", "ascending"),
+ [
+ (
+ "keep",
+ "average",
+ True,
+ ),
+ (
+ "top",
+ "min",
+ False,
+ ),
+ (
+ "bottom",
+ "max",
+ False,
+ ),
+ (
+ "top",
+ "first",
+ False,
+ ),
+ (
+ "bottom",
+ "dense",
+ False,
+ ),
+ ],
+)
+def test_dataframe_groupby_rank(
+ scalars_df_index,
+ scalars_pandas_df_index,
+ na_option,
+ method,
+ ascending,
+):
+ col_names = ["int64_too", "float64_col", "int64_col", "string_col"]
+ bf_result = (
+ scalars_df_index[col_names]
+ .groupby("string_col")
+ .rank(
+ na_option=na_option,
+ method=method,
+ ascending=ascending,
+ )
+ ).to_pandas()
+ pd_result = (
+ (
+ scalars_pandas_df_index[col_names]
+ .groupby("string_col")
+ .rank(
+ na_option=na_option,
+ method=method,
+ ascending=ascending,
+ )
+ )
+ .astype("float64")
+ .astype("Float64")
+ )
+ pd.testing.assert_frame_equal(
+ pd_result, bf_result, check_dtype=False, check_index_type=False
+ )
+
+
 @pytest.mark.parametrize(
  ("operator"),
  [
@@ -534,6 +600,72 @@ def test_series_groupby_agg_list(scalars_df_index, scalars_pandas_df_index):
  )
 
 
+@skip_legacy_pandas
+@pytest.mark.parametrize(
+ ("na_option", "method", "ascending"),
+ [
+ (
+ "keep",
+ "average",
+ True,
+ ),
+ (
+ "top",
+ "min",
+ False,
+ ),
+ (
+ "bottom",
+ "max",
+ False,
+ ),
+ (
+ "top",
+ "first",
+ False,
+ ),
+ (
+ "bottom",
+ "dense",
+ False,
+ ),
+ ],
+)
+def test_series_groupby_rank(
+ scalars_df_index,
+ scalars_pandas_df_index,
+ na_option,
+ method,
+ ascending,
+):
+ col_names = ["int64_col", "string_col"]
+ bf_result = (
+ scalars_df_index[col_names]
+ .groupby("string_col")["int64_col"]
+ .rank(
+ na_option=na_option,
+ method=method,
+ ascending=ascending,
+ )
+ ).to_pandas()
+ pd_result = (
+ (
+ scalars_pandas_df_index[col_names]
+ .groupby("string_col")["int64_col"]
+ .rank(
+ na_option=na_option,
+ method=method,
+ ascending=ascending,
+ )
+ )
+ .astype("float64")
+ .astype("Float64")
+ )
+ pd.testing.assert_series_equal(
+ pd_result, bf_result, check_dtype=False, check_index_type=False
+ )
+
+
 @pytest.mark.parametrize("dropna", [True, False])
 def test_series_groupby_head(scalars_df_index, scalars_pandas_df_index, dropna):
  bf_result = (

third_party/bigframes_vendored/pandas/core/groupby/__init__.py

@@ -363,6 +363,77 @@ def var(
  """
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+ def rank(
+ self,
+ method: str = "average",
+ ascending: bool = True,
+ na_option: str = "keep",
+ ):
+ """
+ Provide the rank of values within each group.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> import numpy as np
+ >>> bpd.options.display.progress_bar = None
+
+ >>> df = bpd.DataFrame(
+ ... {
+ ... "group": ["a", "a", "a", "a", "a", "b", "b", "b", "b", "b"],
+ ... "value": [2, 4, 2, 3, 5, 1, 2, 4, 1, 5],
+ ... }
+ ... )
+ >>> df
+ group value
+ 0 a 2
+ 1 a 4
+ 2 a 2
+ 3 a 3
+ 4 a 5
+ 5 b 1
+ 6 b 2
+ 7 b 4
+ 8 b 1
+ 9 b 5
+ <BLANKLINE>
+ [10 rows x 2 columns]
+ >>> for method in ['average', 'min', 'max', 'dense', 'first']:
+ ... df[f'{method}_rank'] = df.groupby('group')['value'].rank(method)
+ >>> df
+ group value average_rank min_rank max_rank dense_rank first_rank
+ 0 a 2 1.5 1.0 2.0 1.0 1.0
+ 1 a 4 4.0 4.0 4.0 3.0 4.0
+ 2 a 2 1.5 1.0 2.0 1.0 2.0
+ 3 a 3 3.0 3.0 3.0 2.0 3.0
+ 4 a 5 5.0 5.0 5.0 4.0 5.0
+ 5 b 1 1.5 1.0 2.0 1.0 1.0
+ 6 b 2 3.0 3.0 3.0 2.0 3.0
+ 7 b 4 4.0 4.0 4.0 3.0 4.0
+ 8 b 1 1.5 1.0 2.0 1.0 2.0
+ 9 b 5 5.0 5.0 5.0 4.0 5.0
+ <BLANKLINE>
+ [10 rows x 7 columns]
+
+ Args:
+ method ({'average', 'min', 'max', 'first', 'dense'}, default 'average'):
+ * average: average rank of group.
+ * min: lowest rank in group.
+ * max: highest rank in group.
+ * first: ranks assigned in order they appear in the array.
+ * dense: like 'min', but rank always increases by 1 between groups.
+ ascending (bool, default True):
+ False for ranks by high (1) to low (N).
+ na_option ({'keep', 'top', 'bottom'}, default 'keep'):
+ * keep: leave NA values where they are.
+ * top: smallest rank if ascending.
+ * bottom: smallest rank if descending.
+
+ Returns:
+ DataFrame with ranking of values within each group
+ """
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
  def skew(
  self,
  *,