feat(bigframes): Support unstable sort_values, sort_index (#16665)

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/google-cloud-python/issues)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: TrevorBergeron

packages/bigframes/bigframes/core/array_value.py

@@ -212,11 +212,17 @@ def filter(self, predicate: ex.Expression):
             return arr.drop_columns(filter_ids)
 
     def order_by(
-        self, by: Sequence[OrderingExpression], is_total_order: bool = False
+        self,
+        by: Sequence[OrderingExpression],
+        is_total_order: bool = False,
+        stable: bool = True,
     ) -> ArrayValue:
         return ArrayValue(
             nodes.OrderByNode(
-                child=self.node, by=tuple(by), is_total_order=is_total_order
+                child=self.node,
+                by=tuple(by),
+                is_total_order=is_total_order,
+                stable=stable,
             )
         )
 

packages/bigframes/bigframes/core/blocks.py

@@ -395,9 +395,10 @@ def cols_matching_label(self, partial_label: Label) -> typing.Sequence[str]:
     def order_by(
         self,
         by: typing.Sequence[ordering.OrderingExpression],
+        stable: bool = True,
     ) -> Block:
         return Block(
-            self._expr.order_by(by),
+            self._expr.order_by(by, stable=stable),
             index_columns=self.index_columns,
             column_labels=self.column_labels,
             index_labels=self.index.names,
@@ -2412,13 +2413,13 @@ def _align_both_axes(
             rcol_indexer if (rcol_indexer is not None) else range(len(columns))
         )
 
-        left_input_lookup = (
-            lambda index: ex.deref(get_column_left[self.value_columns[index]])
+        left_input_lookup = lambda index: (
+            ex.deref(get_column_left[self.value_columns[index]])
             if index != -1
             else ex.const(None)
         )
-        righ_input_lookup = (
-            lambda index: ex.deref(get_column_right[other.value_columns[index]])
+        righ_input_lookup = lambda index: (
+            ex.deref(get_column_right[other.value_columns[index]])
             if index != -1
             else ex.const(None)
         )
@@ -2471,15 +2472,13 @@ def _align_series_block_axis_1(
             rcol_indexer if (rcol_indexer is not None) else range(len(columns))
         )
 
-        left_input_lookup = (
-            lambda index: ex.deref(get_column_left[self.value_columns[index]])
+        left_input_lookup = lambda index: (
+            ex.deref(get_column_left[self.value_columns[index]])
             if index != -1
             else ex.const(None)
         )
-        righ_input_lookup = (
-            lambda index: ex.deref(
-                get_column_right[other.transpose().value_columns[index]]
-            )
+        righ_input_lookup = lambda index: (
+            ex.deref(get_column_right[other.transpose().value_columns[index]])
             if index != -1
             else ex.const(None)
         )
@@ -2506,13 +2505,11 @@ def _align_pd_series_axis_1(
             rcol_indexer if (rcol_indexer is not None) else range(len(columns))
         )
 
-        left_input_lookup = (
-            lambda index: ex.deref(self.value_columns[index])
-            if index != -1
-            else ex.const(None)
+        left_input_lookup = lambda index: (
+            ex.deref(self.value_columns[index]) if index != -1 else ex.const(None)
         )
-        righ_input_lookup = (
-            lambda index: ex.const(other.iloc[index]) if index != -1 else ex.const(None)
+        righ_input_lookup = lambda index: (
+            ex.const(other.iloc[index]) if index != -1 else ex.const(None)
         )
 
         left_inputs = [left_input_lookup(i) for i in lcol_indexer]

packages/bigframes/bigframes/core/indexes/base.py

@@ -255,12 +255,6 @@ def query_job(self) -> bigquery.QueryJob:
             self._query_job = query_job
         return self._query_job
 
-    @property
-    def str(self) -> bigframes.operations.strings.StringMethods:
-        import bigframes.operations.strings
-
-        return bigframes.operations.strings.StringMethods(self)
-
     def get_loc(self, key) -> typing.Union[int, slice, "bigframes.series.Series"]:
         """Get integer location, slice or boolean mask for requested label.
 
@@ -436,7 +430,8 @@ def sort_values(
         *,
         inplace: bool = False,
         ascending: bool = True,
-        na_position: __builtins__.str = "last",
+        kind: str | None = None,
+        na_position: str = "last",
     ) -> Index:
         if na_position not in ["first", "last"]:
             raise ValueError("Param na_position must be one of 'first' or 'last'")
@@ -448,7 +443,8 @@ def sort_values(
             else order.descending_over(column, na_last)
             for column in index_columns
         ]
-        return Index(self._block.order_by(ordering))
+        is_stable = (kind or constants.DEFAULT_SORT_KIND) in constants.STABLE_SORT_KINDS
+        return Index(self._block.order_by(ordering, stable=is_stable))
 
     def astype(
         self,
@@ -840,6 +836,13 @@ def _apply_binary_op(
         else:
             return NotImplemented
 
+    # last so as to not shadow __builtins__.str
+    @property
+    def str(self) -> bigframes.operations.strings.StringMethods:
+        import bigframes.operations.strings
+
+        return bigframes.operations.strings.StringMethods(self)
+
 
 def _should_create_datetime_index(block: blocks.Block) -> bool:
     if len(block.index.dtypes) != 1:

packages/bigframes/bigframes/core/nodes.py

@@ -991,7 +991,8 @@ def remap_refs(
 @dataclasses.dataclass(frozen=True, eq=False)
 class OrderByNode(UnaryNode):
     by: Tuple[OrderingExpression, ...]
-    # This is an optimization, if true, can discard previous orderings.
+    stable: bool = True
+    # This is an optimization, if true, can discard previous orderings, even if doing a stable sort
     # might be a total ordering even if false
     is_total_order: bool = False
 

packages/bigframes/bigframes/core/rewrite/order.py

@@ -71,7 +71,8 @@ def pull_up_order_inner(
             child_result, child_order = pull_up_order_inner(node.child)
             return child_result, child_order.with_reverse()
         elif isinstance(node, bigframes.core.nodes.OrderByNode):
-            if node.is_total_order:
+            # unstable sorts don't care about previous order, total orders override previous order
+            if (not node.stable) or node.is_total_order:
                 new_node = remove_order(node.child)
             else:
                 new_node, child_order = pull_up_order_inner(node.child)
@@ -106,6 +107,10 @@ def pull_up_order_inner(
                         ),
                     )
                 )
+            elif not node.stable:
+                new_order = bigframes.core.ordering.RowOrdering(
+                    ordering_value_columns=tuple(new_by),
+                )
             else:
                 assert child_order
                 new_order = child_order.with_ordering_columns(new_by)

packages/bigframes/bigframes/dataframe.py

@@ -2418,6 +2418,7 @@ def sort_index(
         *,
         ascending: bool = ...,
         inplace: Literal[False] = ...,
+        kind: str | None = ...,
         na_position: Literal["first", "last"] = ...,
     ) -> DataFrame: ...
 
@@ -2427,6 +2428,7 @@ def sort_index(
         *,
         ascending: bool = ...,
         inplace: Literal[True] = ...,
+        kind: str | None = ...,
         na_position: Literal["first", "last"] = ...,
     ) -> None: ...
 
@@ -2436,6 +2438,7 @@ def sort_index(
         axis: Union[int, str] = 0,
         ascending: bool = True,
         inplace: bool = False,
+        kind: str | None = None,
         na_position: Literal["first", "last"] = "last",
     ) -> Optional[DataFrame]:
         if utils.get_axis_number(axis) == 0:
@@ -2449,7 +2452,10 @@ def sort_index(
                 else order.descending_over(column, na_last)
                 for column in index_columns
             ]
-            block = self._block.order_by(ordering)
+            is_stable = (
+                kind or constants.DEFAULT_SORT_KIND
+            ) in constants.STABLE_SORT_KINDS
+            block = self._block.order_by(ordering, stable=is_stable)
         else:  # axis=1
             _, indexer = self.columns.sort_values(
                 return_indexer=True,
@@ -2472,7 +2478,7 @@ def sort_values(
         *,
         inplace: Literal[False] = ...,
         ascending: bool | typing.Sequence[bool] = ...,
-        kind: str = ...,
+        kind: str | None = ...,
         na_position: typing.Literal["first", "last"] = ...,
     ) -> DataFrame: ...
 
@@ -2483,7 +2489,7 @@ def sort_values(
         *,
         inplace: Literal[True] = ...,
         ascending: bool | typing.Sequence[bool] = ...,
-        kind: str = ...,
+        kind: str | None = ...,
         na_position: typing.Literal["first", "last"] = ...,
     ) -> None: ...
 
@@ -2493,7 +2499,7 @@ def sort_values(
         *,
         inplace: bool = False,
         ascending: bool | typing.Sequence[bool] = True,
-        kind: str = "quicksort",
+        kind: str | None = None,
         na_position: typing.Literal["first", "last"] = "last",
     ) -> Optional[DataFrame]:
         if isinstance(by, (bigframes.series.Series, indexes.Index, DataFrame)):
@@ -2525,7 +2531,8 @@ def sort_values(
                 if is_ascending
                 else order.descending_over(column_id, na_last)
             )
-        block = self._block.order_by(ordering)
+        is_stable = (kind or constants.DEFAULT_SORT_KIND) in constants.STABLE_SORT_KINDS
+        block = self._block.order_by(ordering, stable=is_stable)
         if inplace:
             self._set_block(block)
             return None
@@ -2768,11 +2775,11 @@ def replace(
     ):
         if utils.is_dict_like(value):
             return self.apply(
-                lambda x: x.replace(
-                    to_replace=to_replace, value=value[x.name], regex=regex
+                lambda x: (
+                    x.replace(to_replace=to_replace, value=value[x.name], regex=regex)
+                    if (x.name in value)
+                    else x
                 )
-                if (x.name in value)
-                else x
             )
         return self.apply(
             lambda x: x.replace(to_replace=to_replace, value=value, regex=regex)

packages/bigframes/bigframes/series.py

@@ -1779,7 +1779,7 @@ def sort_values(
         axis=...,
         inplace: Literal[True] = ...,
         ascending: bool | typing.Sequence[bool] = ...,
-        kind: str = ...,
+        kind: str | None = ...,
         na_position: typing.Literal["first", "last"] = ...,
     ) -> None: ...
 
@@ -1790,7 +1790,7 @@ def sort_values(
         axis=...,
         inplace: Literal[False] = ...,
         ascending: bool | typing.Sequence[bool] = ...,
-        kind: str = ...,
+        kind: str | None = ...,
         na_position: typing.Literal["first", "last"] = ...,
     ) -> Series: ...
 
@@ -1800,19 +1800,21 @@ def sort_values(
         axis=0,
         inplace: bool = False,
         ascending=True,
-        kind: str = "quicksort",
+        kind: str | None = None,
         na_position: typing.Literal["first", "last"] = "last",
     ) -> Optional[Series]:
         if axis != 0 and axis != "index":
             raise ValueError(f"No axis named {axis} for object type Series")
         if na_position not in ["first", "last"]:
             raise ValueError("Param na_position must be one of 'first' or 'last'")
+        is_stable = (kind or constants.DEFAULT_SORT_KIND) in constants.STABLE_SORT_KINDS
         block = self._block.order_by(
             [
                 order.ascending_over(self._value_column, (na_position == "last"))
                 if ascending
                 else order.descending_over(self._value_column, (na_position == "last"))
             ],
+            stable=is_stable,
         )
         if inplace:
             self._set_block(block)
@@ -1822,17 +1824,35 @@ def sort_values(
 
     @typing.overload  # type: ignore[override]
     def sort_index(
-        self, *, axis=..., inplace: Literal[False] = ..., ascending=..., na_position=...
+        self,
+        *,
+        axis=...,
+        inplace: Literal[False] = ...,
+        ascending=...,
+        kind: str | None = ...,
+        na_position=...,
     ) -> Series: ...
 
     @typing.overload
     def sort_index(
-        self, *, axis=0, inplace: Literal[True] = ..., ascending=..., na_position=...
+        self,
+        *,
+        axis=0,
+        inplace: Literal[True] = ...,
+        ascending=...,
+        kind: str | None = ...,
+        na_position=...,
     ) -> None: ...
 
     @validations.requires_index
     def sort_index(
-        self, *, axis=0, inplace: bool = False, ascending=True, na_position="last"
+        self,
+        *,
+        axis=0,
+        inplace: bool = False,
+        ascending=True,
+        kind: str | None = None,
+        na_position="last",
     ) -> Optional[Series]:
         # TODO(tbergeron): Support level parameter once multi-index introduced.
         if axis != 0 and axis != "index":
@@ -1847,7 +1867,8 @@ def sort_index(
             else order.descending_over(column, na_last)
             for column in block.index_columns
         ]
-        block = block.order_by(ordering)
+        is_stable = (kind or constants.DEFAULT_SORT_KIND) in constants.STABLE_SORT_KINDS
+        block = block.order_by(ordering, stable=is_stable)
         if inplace:
             self._set_block(block)
             return None

packages/bigframes/third_party/bigframes_vendored/constants.py

@@ -55,3 +55,6 @@
     "_deferred",
 ]
 VALID_WRITE_ENGINES = typing.get_args(WriteEngineType)
+
+DEFAULT_SORT_KIND = "stable"
+STABLE_SORT_KINDS = ("stable", "mergesort")

packages/bigframes/third_party/bigframes_vendored/pandas/core/frame.py

@@ -2253,7 +2253,7 @@ def sort_values(
         *,
         inplace: bool = False,
         ascending: bool | Sequence[bool] = True,
-        kind: str = "quicksort",
+        kind: str | None = None,
         na_position: Literal["first", "last"] = "last",
     ):
         """Sort by the values along row axis.
@@ -2339,7 +2339,7 @@ def sort_values(
                 the by.
             inplace (bool, default False):
                 If True, perform operation in-place.
-            kind (str, default 'quicksort'):
+            kind (str, default None):
                 Choice of sorting algorithm. Accepts 'quicksort', 'mergesort',
                 'heapsort', 'stable'. Ignored except when determining whether to
                 sort stably. 'mergesort' or 'stable' will result in stable reorder.
@@ -2363,6 +2363,7 @@ def sort_index(
         axis: str | int = 0,
         ascending: bool = True,
         inplace: bool = False,
+        kind: str | None = None,
         na_position: Literal["first", "last"] = "last",
     ):
         """Sort object by labels (along an axis).
@@ -2375,6 +2376,10 @@ def sort_index(
                 Sort ascending vs. descending.
             inplace (bool, default False):
                 Whether to modify the DataFrame rather than creating a new one.
+            kind (str, default None):
+                Choice of sorting algorithm. Accepts 'quicksort', 'mergesort',
+                'heapsort', 'stable'. Ignored except when determining whether to
+                sort stably. 'mergesort' or 'stable' will result in stable reorder.
             na_position ({'first', 'last'}, default 'last'):
                 Puts NaNs at the beginning if `first`; `last` puts NaNs at the end.
                 Not implemented for MultiIndex.