refactor: convert ordering definitions to use flexible scalar express… (#490)

…ions

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://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) 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> 🦕

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -30,7 +30,7 @@
 import bigframes.core.join_def as join_def
 import bigframes.core.local_data as local_data
 import bigframes.core.nodes as nodes
-from bigframes.core.ordering import OrderingColumnReference
+from bigframes.core.ordering import OrderingExpression
 import bigframes.core.ordering as orderings
 import bigframes.core.rewrite
 import bigframes.core.schema as schemata
@@ -162,7 +162,7 @@ def filter_by_id(self, predicate_id: str, keep_null: bool = False) -> ArrayValue
     def filter(self, predicate: ex.Expression):
         return ArrayValue(nodes.FilterNode(child=self.node, predicate=predicate))
 
-    def order_by(self, by: Sequence[OrderingColumnReference]) -> ArrayValue:
+    def order_by(self, by: Sequence[OrderingExpression]) -> ArrayValue:
         return ArrayValue(nodes.OrderByNode(child=self.node, by=tuple(by)))
 
     def reversed(self) -> ArrayValue:

bigframes/core/block_transforms.py

@@ -176,7 +176,7 @@ def _interpolate_column(
 ) -> typing.Tuple[blocks.Block, str]:
     if interpolate_method not in ["linear", "nearest", "ffill"]:
         raise ValueError("interpolate method not supported")
-    window_ordering = (ordering.OrderingColumnReference(x_values),)
+    window_ordering = (ordering.OrderingExpression(ex.free_var(x_values)),)
     backwards_window = windows.WindowSpec(following=0, ordering=window_ordering)
     forwards_window = windows.WindowSpec(preceding=0, ordering=window_ordering)
 
@@ -338,8 +338,8 @@ def value_counts(
     if sort:
         block = block.order_by(
             [
-                ordering.OrderingColumnReference(
-                    count_id,
+                ordering.OrderingExpression(
+                    ex.free_var(count_id),
                     direction=ordering.OrderingDirection.ASC
                     if ascending
                     else ordering.OrderingDirection.DESC,
@@ -398,8 +398,8 @@ def rank(
         window = windows.WindowSpec(
             # BigQuery has syntax to reorder nulls with "NULLS FIRST/LAST", but that is unavailable through ibis presently, so must order on a separate nullity expression first.
             ordering=(
-                ordering.OrderingColumnReference(
-                    col,
+                ordering.OrderingExpression(
+                    ex.free_var(col),
                     ordering.OrderingDirection.ASC
                     if ascending
                     else ordering.OrderingDirection.DESC,
@@ -481,8 +481,8 @@ def nsmallest(
     if keep == "last":
         block = block.reversed()
     order_refs = [
-        ordering.OrderingColumnReference(
-            col_id, direction=ordering.OrderingDirection.ASC
+        ordering.OrderingExpression(
+            ex.free_var(col_id), direction=ordering.OrderingDirection.ASC
         )
         for col_id in column_ids
     ]
@@ -511,8 +511,8 @@ def nlargest(
     if keep == "last":
         block = block.reversed()
     order_refs = [
-        ordering.OrderingColumnReference(
-            col_id, direction=ordering.OrderingDirection.DESC
+        ordering.OrderingExpression(
+            ex.free_var(col_id), direction=ordering.OrderingDirection.DESC
         )
         for col_id in column_ids
     ]
@@ -804,9 +804,9 @@ def _idx_extrema(
         )
         # Have to find the min for each
         order_refs = [
-            ordering.OrderingColumnReference(value_col, direction),
+            ordering.OrderingExpression(ex.free_var(value_col), direction),
             *[
-                ordering.OrderingColumnReference(idx_col)
+                ordering.OrderingExpression(ex.free_var(idx_col))
                 for idx_col in original_block.index_columns
             ],
         ]

bigframes/core/blocks.py

@@ -270,7 +270,7 @@ def cols_matching_label(self, partial_label: Label) -> typing.Sequence[str]:
 
     def order_by(
         self,
-        by: typing.Sequence[ordering.OrderingColumnReference],
+        by: typing.Sequence[ordering.OrderingExpression],
     ) -> Block:
         return Block(
             self._expr.order_by(by),
@@ -610,7 +610,9 @@ def _split(
             string_ordering_col, random_state_col, ops.strconcat_op
         )
         block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op)
-        block = block.order_by([ordering.OrderingColumnReference(hash_string_sum_col)])
+        block = block.order_by(
+            [ordering.OrderingExpression(ex.free_var(hash_string_sum_col))]
+        )
 
         intervals = []
         cur = 0
@@ -628,15 +630,17 @@ def _split(
             sliced_blocks = [
                 sliced_block.order_by(
                     [
-                        ordering.OrderingColumnReference(idx_col)
+                        ordering.OrderingExpression(ex.free_var(idx_col))
                         for idx_col in sliced_block.index_columns
                     ]
                 )
                 for sliced_block in sliced_blocks
             ]
         elif sort is False:
             sliced_blocks = [
-                sliced_block.order_by([ordering.OrderingColumnReference(ordering_col)])
+                sliced_block.order_by(
+                    [ordering.OrderingExpression(ex.free_var(ordering_col))]
+                )
                 for sliced_block in sliced_blocks
             ]
 
@@ -1706,7 +1710,10 @@ def merge(
         if sort:
             # sort uses coalesced join keys always
             joined_expr = joined_expr.order_by(
-                [ordering.OrderingColumnReference(col_id) for col_id in coalesced_ids],
+                [
+                    ordering.OrderingExpression(ex.free_var(col_id))
+                    for col_id in coalesced_ids
+                ],
             )
 
         joined_expr = joined_expr.select_columns(result_columns)
@@ -2028,7 +2035,10 @@ def join_mono_indexed(
     )
     if sort:
         combined_expr = combined_expr.order_by(
-            [ordering.OrderingColumnReference(col_id) for col_id in coalesced_join_cols]
+            [
+                ordering.OrderingExpression(ex.free_var(col_id))
+                for col_id in coalesced_join_cols
+            ]
         )
     block = Block(
         combined_expr,
@@ -2117,7 +2127,10 @@ def join_multi_indexed(
     )
     if sort:
         combined_expr = combined_expr.order_by(
-            [ordering.OrderingColumnReference(col_id) for col_id in coalesced_join_cols]
+            [
+                ordering.OrderingExpression(ex.free_var(col_id))
+                for col_id in coalesced_join_cols
+            ]
         )
 
     if left.index.nlevels == 1: