refactor: Make join order propogation explicit in node (#1434)

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -452,6 +452,7 @@ def relational_join(
         other: ArrayValue,
         conditions: typing.Tuple[typing.Tuple[str, str], ...] = (),
         type: typing.Literal["inner", "outer", "left", "right", "cross"] = "inner",
+        propogate_order: Optional[bool] = None,
     ) -> typing.Tuple[ArrayValue, typing.Tuple[dict[str, str], dict[str, str]]]:
         l_mapping = {  # Identity mapping, only rename right side
             lcol.name: lcol.name for lcol in self.node.ids
@@ -465,6 +466,7 @@ def relational_join(
                 for l_col, r_col in conditions
             ),
             type=type,
+            propogate_order=propogate_order or self.session._strictly_ordered,
         )
         return ArrayValue(join_node), (l_mapping, r_mapping)
 

bigframes/core/compile/compiler.py

@@ -68,9 +68,7 @@ def compile_sql(
         node = self._replace_unsupported_ops(node)
         # prune before pulling up order to avoid unnnecessary row_number() ops
         node = rewrites.column_pruning(node)
-        node, ordering = rewrites.pull_up_order(
-            node, order_root=ordered, ordered_joins=self.strict
-        )
+        node, ordering = rewrites.pull_up_order(node, order_root=ordered)
         # final pruning to cleanup up any leftovers unused values
         node = rewrites.column_pruning(node)
         return self.compile_node(node).to_sql(
@@ -87,9 +85,7 @@ def compile_raw(
     ]:
         node = self._replace_unsupported_ops(node)
         node = rewrites.column_pruning(node)
-        node, ordering = rewrites.pull_up_order(
-            node, order_root=True, ordered_joins=self.strict
-        )
+        node, ordering = rewrites.pull_up_order(node, order_root=True)
         node = rewrites.column_pruning(node)
         sql = self.compile_node(node).to_sql()
         return sql, node.schema.to_bigquery(), ordering

bigframes/core/nodes.py

@@ -287,6 +287,7 @@ class JoinNode(BigFrameNode):
     right_child: BigFrameNode
     conditions: typing.Tuple[typing.Tuple[ex.DerefOp, ex.DerefOp], ...]
     type: typing.Literal["inner", "outer", "left", "right", "cross"]
+    propogate_order: bool
 
     def _validate(self):
         assert not (
@@ -311,8 +312,7 @@ def order_ambiguous(self) -> bool:
 
     @property
     def explicitly_ordered(self) -> bool:
-        # Do not consider user pre-join ordering intent - they need to re-order post-join in unordered mode.
-        return False
+        return self.propogate_order
 
     @property
     def fields(self) -> Iterable[Field]:

bigframes/core/rewrite/order.py

@@ -29,7 +29,6 @@ def pull_up_order(
     root: bigframes.core.nodes.BigFrameNode,
     *,
     order_root: bool = True,
-    ordered_joins: bool = True,
 ) -> Tuple[bigframes.core.nodes.BigFrameNode, bigframes.core.ordering.RowOrdering]:
     """
     Pull the ordering up, putting full order definition into window ops.
@@ -92,7 +91,7 @@ def pull_up_order_inner(
             child_result, child_order = pull_up_order_inner(node.child)
             return node.replace_child(child_result), child_order
         elif isinstance(node, bigframes.core.nodes.JoinNode):
-            if ordered_joins:
+            if node.propogate_order:
                 return pull_order_join(node)
             else:
                 return (