Add edges defined by a composite join clause (#998)

* Add composite join descriptor

* Add TODOs

* Construct general on clause

* Deal with _came_from columns

* Add errors

* Update _find_used_columns

* Add test with traversal

* Lint

* Lint

* Add todos

* Lint

* Add error messages

* Add error message

* Address coment

* Update graphql_compiler/compiler/emit_sql.py

Co-authored-by: Predrag Gruevski <2348618+obi1kenobi@users.noreply.github.com>

* Address comments

* Lint

Co-authored-by: Predrag Gruevski <2348618+obi1kenobi@users.noreply.github.com>

Author: bojanserafimov

graphql_compiler/compiler/emit_sql.py

@@ -1,7 +1,7 @@
 # Copyright 2018-present Kensho Technologies, LLC.
 """Transform a SqlNode tree into an executable SQLAlchemy query."""
 from dataclasses import dataclass
-from typing import Dict, Iterator, List, NamedTuple, Optional, Set, Tuple, Union
+from typing import AbstractSet, Dict, Iterator, List, NamedTuple, Optional, Set, Tuple, Union
 
 import six
 import sqlalchemy
@@ -21,7 +21,12 @@
 from . import blocks
 from ..global_utils import VertexPath
 from ..schema import COUNT_META_FIELD_NAME
-from ..schema.schema_info import DirectJoinDescriptor, SQLAlchemySchemaInfo
+from ..schema.schema_info import (
+    CompositeJoinDescriptor,
+    DirectJoinDescriptor,
+    JoinDescriptor,
+    SQLAlchemySchemaInfo,
+)
 from .compiler_entities import BasicBlock
 from .compiler_frontend import IrAndMetadata
 from .expressions import ContextField, Expression
@@ -143,21 +148,30 @@ def _find_used_columns(
             )
             vertex_field_name = f"{edge_direction}_{edge_name}"
             edge = sql_schema_info.join_descriptors[location_info.type.name][vertex_field_name]
-            used_columns.setdefault(get_vertex_path(location), set()).add(edge.from_column)
-            used_columns.setdefault(get_vertex_path(child_location), set()).add(edge.to_column)
+            if isinstance(edge, DirectJoinDescriptor):
+                columns_at_location = {edge.from_column}
+                columns_at_child = {edge.to_column}
+            elif isinstance(edge, CompositeJoinDescriptor):
+                columns_at_location = {column_pair[0] for column_pair in edge.column_pairs}
+                columns_at_child = {column_pair[1] for column_pair in edge.column_pairs}
+            else:
+                raise AssertionError(f"Unknown join descriptor type {edge}: {type(edge)}")
+
+            used_columns.setdefault(get_vertex_path(location), set()).update(columns_at_location)
+            used_columns.setdefault(get_vertex_path(child_location), set()).update(columns_at_child)
 
             # Check if the edge is recursive
             child_location_info = ir.query_metadata_table.get_location_info(child_location)
             if child_location_info.recursive_scopes_depth > location_info.recursive_scopes_depth:
                 # The primary key may be used if the recursive cte base semijoins to
                 # the pre-recurse cte by primary key.
                 alias = sql_schema_info.vertex_name_to_table[location_info.type.name].alias()
-                primary_key_name = _get_primary_key_name(alias, location_info.type.name, "@recurse")
-                used_columns.setdefault(get_vertex_path(location), set()).add(primary_key_name)
+                primary_keys = {column.name for column in alias.primary_key}
+                used_columns.setdefault(get_vertex_path(location), set()).update(primary_keys)
 
                 # The from_column is used at the destination as well, inside the recursive step
-                used_columns.setdefault(get_vertex_path(child_location), set()).add(
-                    edge.from_column
+                used_columns.setdefault(get_vertex_path(child_location), set()).update(
+                    columns_at_location
                 )
 
     # Find outputs used
@@ -780,7 +794,9 @@ def __init__(self, sql_schema_info: SQLAlchemySchemaInfo, ir: IrAndMetadata):
         # Move to the beginning location of the query.
         self._relocate(ir.query_metadata_table.root_location)
 
-        # Mapping aliases to the column used to join into them.
+        # Mapping aliases to one of the column used to join into them. We use this column
+        # to check for LEFT JOIN misses, since it helps us distinguish actuall NULL values
+        # from values that are NULL because of a LEFT JOIN miss.
         self._came_from: Dict[Union[Alias, ColumnRouter], Column] = {}
 
         self._recurse_needs_cte: bool = False
@@ -840,9 +856,8 @@ def _relocate(self, new_location: BaseLocation):
                     self._current_alias, self._current_location, output_fields
                 )
 
-    # TODO merge from_column and to_column into a joindescriptor
     def _join_to_parent_location(
-        self, parent_alias: Alias, from_column: str, to_column: str, optional: bool
+        self, parent_alias: Alias, join_descriptor: JoinDescriptor, optional: bool
     ):
         """Join the current location to the parent location using the column names specified."""
         if self._current_alias is None:
@@ -851,7 +866,25 @@ def _join_to_parent_location(
                 f"during fold {self}."
             )
 
-        self._came_from[self._current_alias] = self._current_alias.c[to_column]
+        # construct on clause for join
+        if isinstance(join_descriptor, DirectJoinDescriptor):
+            matching_column_pairs: AbstractSet[Tuple[str, str]] = {
+                (join_descriptor.from_column, join_descriptor.to_column),
+            }
+        elif isinstance(join_descriptor, CompositeJoinDescriptor):
+            matching_column_pairs = join_descriptor.column_pairs
+        else:
+            raise AssertionError(
+                f"Unknown join descriptor type {join_descriptor}: {type(join_descriptor)}"
+            )
+
+        if not matching_column_pairs:
+            raise AssertionError(
+                f"Invalid join descriptor {join_descriptor}, produced no matching column pairs."
+            )
+
+        _, non_null_column = sorted(matching_column_pairs)[0]
+        self._came_from[self._current_alias] = self._current_alias.c[non_null_column]
 
         if self._is_in_optional_scope() and not optional:
             # For mandatory edges in optional scope, we emit LEFT OUTER JOIN and enforce the
@@ -879,10 +912,17 @@ def _join_to_parent_location(
                 )
             )
 
+        on_clause = sqlalchemy.and_(
+            *(
+                parent_alias.c[from_column] == self._current_alias.c[to_column]
+                for from_column, to_column in sorted(matching_column_pairs)
+            )
+        )
+
         # Join to where we came from.
         self._from_clause = self._from_clause.join(
             self._current_alias,
-            onclause=(parent_alias.c[from_column] == self._current_alias.c[to_column]),
+            onclause=on_clause,
             isouter=self._is_in_optional_scope(),
         )
 
@@ -932,11 +972,14 @@ def traverse(self, vertex_field: str, optional: bool) -> None:
                     "Attempting to traverse inside a fold while the _current_location was not a "
                     f"FoldScopeLocation. _current_location was set to {self._current_location}."
                 )
+            if not isinstance(edge, DirectJoinDescriptor):
+                raise NotImplementedError(
+                    f"Edge {vertex_field} is backed by a CompositeJoinDescriptor, "
+                    "so it can't be used inside a @fold scope."
+                )
             self._current_fold.add_traversal(edge, previous_alias, self._current_alias)
         else:
-            self._join_to_parent_location(
-                previous_alias, edge.from_column, edge.to_column, optional
-            )
+            self._join_to_parent_location(previous_alias, edge, optional)
 
     def _wrap_into_cte(self) -> None:
         """Wrap the current query into a cte."""
@@ -1017,6 +1060,11 @@ def recurse(self, vertex_field: str, depth: int) -> None:
             )
 
         edge = self._sql_schema_info.join_descriptors[self._current_classname][vertex_field]
+        if not isinstance(edge, DirectJoinDescriptor):
+            raise NotImplementedError(
+                f"Edge {vertex_field} requires a JOIN across a composite key, which is currently "
+                f"not supported for use with @recurse."
+            )
         primary_key = self._get_current_primary_key_name("@recurse")
 
         # Wrap the query so far into a CTE if it would speed up the recursive query.
@@ -1074,8 +1122,8 @@ def recurse(self, vertex_field: str, depth: int) -> None:
             .where(base.c[CTE_DEPTH_NAME] < literal_depth)
         )
 
-        # TODO(bojanserafimov): This creates an unused alias if there's no tags or outputs so far
-        self._join_to_parent_location(previous_alias, primary_key, CTE_KEY_NAME, False)
+        join_descriptor = DirectJoinDescriptor(primary_key, CTE_KEY_NAME)
+        self._join_to_parent_location(previous_alias, join_descriptor, False)
 
     def start_global_operations(self) -> None:
         """Execute a GlobalOperationsStart block."""
@@ -1131,6 +1179,11 @@ def fold(self, fold_scope_location: FoldScopeLocation) -> None:
         join_descriptor = self._sql_schema_info.join_descriptors[self._current_classname][
             full_edge_name
         ]
+        if not isinstance(join_descriptor, DirectJoinDescriptor):
+            raise NotImplementedError(
+                f"Edge {full_edge_name} requires a JOIN across a composite key, which is currently "
+                "not supported for use with @fold."
+            )
 
         # 3. Initialize fold object.
         self._current_fold = FoldSubqueryBuilder(

graphql_compiler/compiler/ir_lowering_sql/__init__.py

@@ -3,6 +3,7 @@
 
 from .. import blocks, expressions
 from ...compiler.compiler_frontend import IrAndMetadata
+from ...schema.schema_info import CompositeJoinDescriptor, DirectJoinDescriptor
 from ..helpers import FoldScopeLocation, get_edge_direction_and_name
 from ..ir_lowering_common import common
 
@@ -40,9 +41,14 @@ def _find_non_null_columns(schema_info, query_metadata_table):
             vertex_field_name = "{}_{}".format(edge_direction, edge_name)
             edge = schema_info.join_descriptors[location_info.type.name][vertex_field_name]
 
-            # The value of the column used to join to this table is an indicator of whether
+            # The value of any column used to join to this table is an indicator of whether
             # the left join was a hit or a miss.
-            non_null_column[child_location.query_path] = edge.to_column
+            if isinstance(edge, DirectJoinDescriptor):
+                non_null_column[child_location.query_path] = edge.to_column
+            elif isinstance(edge, CompositeJoinDescriptor):
+                non_null_column[child_location.query_path] = sorted(edge.column_pairs)[0][1]
+            else:
+                raise AssertionError(f"Unknown join descriptor type {edge}: {type(edge)}")
 
     return non_null_column
 

graphql_compiler/schema/schema_info.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 from enum import Enum, Flag, auto, unique
 from functools import partial
-from typing import Dict, Mapping, Optional, Sequence
+from typing import AbstractSet, Dict, Mapping, Optional, Sequence, Tuple, Union
 
 from graphql.type import GraphQLSchema
 from graphql.type.definition import GraphQLInterfaceType, GraphQLObjectType
@@ -20,20 +20,46 @@
 from ..schema_generation.schema_graph import SchemaGraph
 
 
-# Describes the intent to join two tables using the specified columns.
-#
-# The resulting join expression could be something like:
-# JOIN origin_table.from_column = destination_table.to_column
-#
-# The type of join (inner vs left, etc.) is not specified.
-# The tables are not specified.
-DirectJoinDescriptor = namedtuple(
-    "DirectJoinDescriptor",
-    (
-        "from_column",  # The column in the source table we intend to join on.
-        "to_column",  # The column in the destination table we intend to join on.
-    ),
-)
+@dataclass
+class DirectJoinDescriptor:
+    """Describes the ability to join two tables using the specified columns.
+
+    The resulting join expression could be something like:
+    JOIN origin_table.from_column = destination_table.to_column
+
+    The type of join (inner vs left, etc.) is not specified.
+    The tables are not specified.
+    """
+
+    from_column: str  # The column in the source table we intend to join on.
+    to_column: str  # The column in the destination table we intend to join on.
+
+
+@dataclass
+class CompositeJoinDescriptor:
+    """Describes the ability to join two tables with a composite relationship.
+
+    The resulting join expression could be something like:
+    JOIN
+        origin_table.from_column_1 == destination_table.to_column_1 AND
+        origin_table.from_column_2 == destination_table.to_column_2 AND
+        origin_table.from_column_3 == destination_table.to_column_3
+
+    The type of join (inner vs left, etc.) is not specified.
+    The tables are not specified.
+    """
+
+    # (from_column, to_column) pairs, where from_column is on the origin table
+    # and to_column is on the destination table of the join.
+    column_pairs: AbstractSet[Tuple[str, str]]
+
+    def __post_init__(self) -> None:
+        """Validate fields."""
+        if not self.column_pairs:
+            raise AssertionError("The column_pairs field is expected to be non-empty.")
+
+
+JoinDescriptor = Union[DirectJoinDescriptor, CompositeJoinDescriptor]
 
 
 @dataclass
@@ -143,17 +169,17 @@ class SQLSchemaInfo(BackendSpecificSchemaInfo):
 
     vertex_name_to_table: Dict[str, sqlalchemy.Table]
     # dict mapping every GraphQL object type or interface type name in the schema to
-    # dict mapping every vertex field name at that type to a DirectJoinDescriptor.
+    # dict mapping every vertex field name at that type to a JoinDescriptor.
     # The tables the join is to be performed on are not specified.
     # They are inferred from the schema and the tables dictionary.
-    join_descriptors: Dict[str, Dict[str, DirectJoinDescriptor]]
+    join_descriptors: Dict[str, Dict[str, JoinDescriptor]]
 
 
 def _create_sql_schema_info(
     dialect: Dialect,
     schema: GraphQLSchema,
     vertex_name_to_table: Dict[str, sqlalchemy.Table],
-    join_descriptors: Dict[str, Dict[str, DirectJoinDescriptor]],
+    join_descriptors: Dict[str, Dict[str, JoinDescriptor]],
     type_equivalence_hints: Optional[Dict[str, str]] = None,
 ) -> SQLSchemaInfo:
     """Create a SQLSchemaInfo object for a database using a flavor of SQL."""
@@ -239,7 +265,7 @@ def _create_sql_schema_info(
         # All tables are expected to have primary keys.
         "vertex_name_to_table",
         # dict mapping every graphql object type or interface type name in the schema to:
-        #    dict mapping every vertex field name at that type to a DirectJoinDescriptor. The
+        #    dict mapping every vertex field name at that type to a JoinDescriptor. The
         #    tables the join is to be performed on are not specified. They are inferred from
         #    the schema and the tables dictionary.
         "join_descriptors",
@@ -252,7 +278,7 @@ def make_sqlalchemy_schema_info(
     type_equivalence_hints: TypeEquivalenceHintsType,
     dialect: Dialect,
     vertex_name_to_table: Dict[str, sqlalchemy.Table],
-    join_descriptors: Dict[str, Dict[str, DirectJoinDescriptor]],
+    join_descriptors: Dict[str, Dict[str, JoinDescriptor]],
     validate: bool = True,
 ) -> SQLAlchemySchemaInfo:
     """Make a SQLAlchemySchemaInfo if the input provided is valid.
@@ -282,7 +308,7 @@ def make_sqlalchemy_schema_info(
                               schema to a SQLAlchemy table
         join_descriptors: dict mapping GraphQL object and interface type names in the schema to:
                           dict mapping every vertex field name at that type to a
-                          DirectJoinDescriptor. The tables on which the join is to be performed
+                          JoinDescriptor. The tables on which the join is to be performed
                           are not specified. They are inferred from the schema and the tables
                           dictionary.
         validate: whether to validate that the given inputs are valid for creation of

graphql_compiler/schema_generation/sqlalchemy/edge_descriptors.py

@@ -115,6 +115,7 @@ def generate_direct_edge_descriptors_from_foreign_keys(
             else:
                 number_of_composite_foreign_keys += 1
 
+    # TODO(bojanserafimov): Infer CompositeJoinDescriptor objects
     if number_of_composite_foreign_keys:
         warnings.warn(
             "Ignored {} edges implied by composite foreign keys. We currently do not "

graphql_compiler/tests/test_compiler.py

@@ -11154,3 +11154,41 @@ def test_recursive_field_type_is_subtype_of_parent_field(self) -> None:
             expected_cypher,
             expected_sql,
         )
+
+    def test_animal_born_at_traversal(self) -> None:
+        """Ensure that sql composite key traversals work."""
+        test_data = test_input_data.animal_born_at_traversal()
+
+        expected_match = SKIP_TEST
+        expected_gremlin = SKIP_TEST
+        expected_mssql = """
+            SELECT
+                [Animal_1].name AS animal_name,
+                [BirthEvent_1].name AS birth_event_name
+            FROM
+                db_1.schema_1.[Animal] AS [Animal_1]
+                JOIN db_1.schema_1.[BirthEvent] AS [BirthEvent_1]
+                    ON [Animal_1].birth_date = [BirthEvent_1].event_date
+                    AND [Animal_1].birth_uuid = [BirthEvent_1].uuid
+        """
+        expected_postgresql = """
+            SELECT
+                "Animal_1".name AS animal_name,
+                "BirthEvent_1".name AS birth_event_name
+            FROM
+                schema_1."Animal" AS "Animal_1"
+                JOIN schema_1."BirthEvent" AS "BirthEvent_1"
+                    ON "Animal_1".birth_date = "BirthEvent_1".event_date
+                    AND "Animal_1".birth_uuid = "BirthEvent_1".uuid
+        """
+        expected_cypher = SKIP_TEST
+
+        check_test_data(
+            self,
+            test_data,
+            expected_match,
+            expected_gremlin,
+            expected_mssql,
+            expected_cypher,
+            expected_postgresql,
+        )