Delete redundant joins by substitution equi-join keys for their mirror to render one side pruneable

pydough/conversion/hybrid_tree.py

@@ -858,7 +858,7 @@ def is_singular(self) -> bool:
             case _:
                 return False
         # The current level is fine, so check any levels above it next.
-        return True if self.parent is None else self.parent.always_exists()
+        return True if self.parent is None else self.parent.is_singular()
 
     def equals_ignoring_successors(self, other: "HybridTree") -> bool:
         """

pydough/conversion/join_key_substitution.py

@@ -0,0 +1,108 @@
+"""
+Logic for switching references to join keys from one side of a join to the other
+when certain conditions are met, thus allowing the join to be removed by the
+column pruner. The conditions are:
+- The join is an inner join.
+- The join has equi-join keys.
+- The cardinality in either direction is singular-access.
+- The only columns used from one side of the join (the one being referenced in
+  a singular-access manner) are the join keys (or a subset thereof).
+"""
+
+from pydough.relational import (
+    ColumnReference,
+    ColumnReferenceFinder,
+    Join,
+    JoinCardinality,
+    JoinType,
+    RelationalExpression,
+    RelationalNode,
+    RelationalShuttle,
+)
+from pydough.relational.rel_util import (
+    apply_substitution,
+    extract_equijoin_keys,
+)
+
+
+class JoinKeySubstitutionShuttle(RelationalShuttle):
+    """
+    The relational shuttle that performs join key substitution optimization.
+    """
+
+    def visit_join(self, join: Join) -> RelationalNode:
+        # Build up a mapping of join key substitutions mapping input columns
+        # from one input to another when the optimization case is detected:
+        # requires an inner join with equi-join keys.
+        join_substitution: dict[RelationalExpression, RelationalExpression] = {}
+        if join.join_type == JoinType.INNER:
+            lhs_keys_list, rhs_keys_list = extract_equijoin_keys(join)
+            if len(lhs_keys_list) > 0 and len(rhs_keys_list) > 0:
+                # Identify which columns are used by the join columns that come
+                # from the left and right inputs.
+                lhs_keys: set[ColumnReference] = set(lhs_keys_list)
+                rhs_keys: set[ColumnReference] = set(rhs_keys_list)
+                col_finder = ColumnReferenceFinder()
+                for value in join.columns.values():
+                    value.accept(col_finder)
+                col_refs: set[ColumnReference] = col_finder.get_column_references()
+                lhs_refs = {
+                    ref
+                    for ref in col_refs
+                    if ref.input_name == join.default_input_aliases[0]
+                }
+                rhs_refs = col_refs - lhs_refs
+                # If each row on the left side (LHS) matches exactly one row on the right side (RHS)
+                # (i.e., singular access)
+                # and the query only references columns from the RHS that are join keys,
+                # then we can substitute the RHS join keys with the corresponding LHS join keys.
+                # This allows the join to potentially be removed later since it adds no new data.
+                if (
+                    join.cardinality == JoinCardinality.SINGULAR_ACCESS
+                    and rhs_refs <= rhs_keys
+                ):
+                    for lhs_key, rhs_key in zip(lhs_keys_list, rhs_keys_list):
+                        join_substitution[rhs_key] = lhs_key
+
+                # If the right side is singular access, and all the columns used
+                # from the left side are just the join keys, then we can
+                # substitute the left join keys with the right join keys.
+                elif (
+                    join.reverse_cardinality == JoinCardinality.SINGULAR_ACCESS
+                    and lhs_refs <= lhs_keys
+                ):
+                    for lhs_key, rhs_key in zip(lhs_keys_list, rhs_keys_list):
+                        join_substitution[lhs_key] = rhs_key
+
+        # If any substitutions were identified, create a new Join node
+        # with the substitutions applied to its columns.
+        if len(join_substitution) > 0:
+            join = Join(
+                join.inputs,
+                join.condition,
+                join.join_type,
+                {
+                    name: apply_substitution(expr, join_substitution, {})
+                    for name, expr in join.columns.items()
+                },
+                join.cardinality,
+                join.reverse_cardinality,
+                join.correl_name,
+            )
+
+        # Recursively visit the inputs to the join to transform them as well.
+        return super().visit_join(join)
+
+
+def join_key_substitution(root: RelationalNode) -> RelationalNode:
+    """
+    The main entry point for join key substitution optimization.
+
+    Args:
+        `root`: The root of the relational tree being optimized.
+
+    Returns:
+        The optimized relational tree.
+    """
+    shuttle: JoinKeySubstitutionShuttle = JoinKeySubstitutionShuttle()
+    return root.accept_shuttle(shuttle)

pydough/conversion/relational_converter.py

@@ -88,6 +88,7 @@
 from .hybrid_translator import HybridTranslator
 from .hybrid_tree import HybridTree
 from .join_aggregate_transpose import pull_aggregates_above_joins
+from .join_key_substitution import join_key_substitution
 from .masking_shuttles import MaskLiteralComparisonShuttle
 from .merge_projects import merge_projects
 from .projection_pullup import pullup_projections
@@ -1584,6 +1585,12 @@ def optimize_relational_tree(
     pruner: ColumnPruner = ColumnPruner()
     root = pruner.prune_unused_columns(root)
 
+    # Run a pass that substitutes join keys when the only columns used by one
+    # side of the join are the join keys. This will make some joins redundant
+    # and allow them to be deleted later. Then, re-run column pruning.
+    root = confirm_root(join_key_substitution(root))
+    root = pruner.prune_unused_columns(root)
+
     # Bubble up names from the leaf nodes to further encourage simpler naming
     # without aliases, and also to delete duplicate columns where possible.
     # This is done early to maximize the chances that a nicer name will be used
@@ -1631,7 +1638,8 @@ def optimize_relational_tree(
     #   D: join-aggregate transpose
     #   E: projection pullup again
     #   F: redundant aggregation removal
-    #   G: column pruning
+    #   G: join key substitution
+    #   H: column pruning
     # This is done because pullup will create more opportunities for expression
     # simplification, which will allow more filters to be pushed further down,
     # and the combination of those together will create more opportunities for
@@ -1644,6 +1652,7 @@ def optimize_relational_tree(
         root = confirm_root(pull_aggregates_above_joins(root))
         root = confirm_root(pullup_projections(root))
         root = remove_redundant_aggs(root)
+        root = confirm_root(join_key_substitution(root))
         root = pruner.prune_unused_columns(root)
 
     # Re-run projection merging, without pushing into joins. This will allow

pydough/conversion/relational_simplification.py

@@ -1283,8 +1283,14 @@ def simplify_function_call(
                         )
 
                     case _:
-                        # All other cases remain non-simplified.
-                        pass
+                        # Simplify comparing an expression to itself as
+                        # True/False. All other cases remain non-simplified.
+                        if expr.inputs[0] == expr.inputs[1]:
+                            is_eq: bool = expr.op in (pydop.EQU, pydop.LEQ, pydop.GEQ)
+                            output_expr = LiteralExpression(is_eq, expr.data_type)
+                            output_predicates |= PredicateSet(
+                                not_null=True, not_negative=True, positive=is_eq
+                            )
 
                 output_predicates.not_negative = True
 

tests/test_metadata/sample_graphs.json

@@ -673,7 +673,7 @@
         "parent collection": "parts",
         "child collection": "supply_records",
         "singular": false,
-        "always matches": false,
+        "always matches": true,
         "keys": {"key": ["part_key"]},
         "description": "The records indicating which companies supply the part",
         "synonyms": ["producers", "vendors", "suppliers of part"]

tests/test_metadata/snowflake_sample_graphs.json

@@ -673,7 +673,7 @@
         "parent collection": "parts",
         "child collection": "supply_records",
         "singular": false,
-        "always matches": false,
+        "always matches": true,
         "keys": {"key": ["part_key"]},
         "description": "The records indicating which companies supply the part",
         "synonyms": ["producers", "vendors", "suppliers of part"]