Include location type in ContextField and related expressions. (kensho-technologies#229)

* Include location type in ContextField and related types.

* Fix the type check in the equality comparison of compiler entities.

Author: obi1kenobi

graphql_compiler/compiler/compiler_entities.py

@@ -3,6 +3,7 @@
 
 from abc import ABCMeta, abstractmethod
 
+from graphql import is_type
 import six
 
 
@@ -43,9 +44,23 @@ def __repr__(self):
     # pylint: disable=protected-access
     def __eq__(self, other):
         """Return True if the CompilerEntity objects are equal, and False otherwise."""
-        return (type(self) == type(other) and
-                self._print_args == other._print_args and
-                self._print_kwargs == other._print_kwargs)
+        if type(self) != type(other):
+            return False
+
+        if len(self._print_args) != len(other._print_args):
+            return False
+
+        # The args sometimes contain GraphQL type objects, which unfortunately do not define "==".
+        # We have to split them out and compare them using "is_same_type()" instead.
+        for self_arg, other_arg in six.moves.zip(self._print_args, other._print_args):
+            if is_type(self_arg):
+                if not self_arg.is_same_type(other_arg):
+                    return False
+            else:
+                if self_arg != other_arg:
+                    return False
+
+        return self._print_kwargs == other._print_kwargs
     # pylint: enable=protected-access
 
     def __ne__(self, other):

graphql_compiler/compiler/compiler_frontend.py

@@ -688,7 +688,8 @@ def _compile_ast_node_to_ir(schema, current_schema_type, ast, location, context)
 
             visitor_fn = expressions.make_type_replacement_visitor(
                 expressions.ContextField,
-                lambda context_field: expressions.GlobalContextField(context_field.location))
+                lambda context_field: expressions.GlobalContextField(
+                    context_field.location, context_field.field_type))
             filter_block = filter_block.visit_and_update_expressions(visitor_fn)
 
             set_fold_count_filter(context)

graphql_compiler/compiler/expressions.py

@@ -270,9 +270,9 @@ def to_gremlin(self):
 class GlobalContextField(Expression):
     """A field drawn from the global context, for use in a global operations WHERE statement."""
 
-    __slots__ = ('location',)
+    __slots__ = ('location', 'field_type')
 
-    def __init__(self, location):
+    def __init__(self, location, field_type):
         """Construct a new GlobalContextField object that references a field at a given location.
 
         Args:
@@ -281,8 +281,9 @@ def __init__(self, location):
         Returns:
             new GlobalContextField object
         """
-        super(GlobalContextField, self).__init__(location)
+        super(GlobalContextField, self).__init__(location, field_type)
         self.location = location
+        self.field_type = field_type
         self.validate()
 
     def validate(self):
@@ -295,6 +296,9 @@ def validate(self):
             raise AssertionError(u'Received Location without a field: {}'
                                  .format(self.location))
 
+        if not is_graphql_type(self.field_type):
+            raise ValueError(u'Invalid value of "field_type": {}'.format(self.field_type))
+
     def to_match(self):
         """Return a unicode object with the MATCH representation of this GlobalContextField."""
         self.validate()
@@ -314,22 +318,24 @@ def to_gremlin(self):
 class ContextField(Expression):
     """A field drawn from the global context, e.g. if selected earlier in the query."""
 
-    __slots__ = ('location',)
+    __slots__ = ('location', 'field_type')
 
-    def __init__(self, location):
+    def __init__(self, location, field_type):
         """Construct a new ContextField object that references a field from the global context.
 
         Args:
             location: Location, specifying where the field was declared. If the Location points
                       to a vertex, the field refers to the data captured at the location vertex.
                       Otherwise, if the Location points to a property, the field refers to
                       the particular value of that property.
+            field_type: GraphQLType object, specifying the type of the field being output
 
         Returns:
             new ContextField object
         """
-        super(ContextField, self).__init__(location)
+        super(ContextField, self).__init__(location, field_type)
         self.location = location
+        self.field_type = field_type
         self.validate()
 
     def validate(self):
@@ -338,6 +344,9 @@ def validate(self):
             raise TypeError(u'Expected Location location, got: {} {}'.format(
                 type(self.location).__name__, self.location))
 
+        if not is_graphql_type(self.field_type):
+            raise ValueError(u'Invalid value of "field_type": {}'.format(self.field_type))
+
     def to_match(self):
         """Return a unicode object with the MATCH representation of this ContextField."""
         self.validate()

graphql_compiler/compiler/filters.py

@@ -148,7 +148,7 @@ def _represent_argument(directive_location, context, argument, inferred_type):
         if field_is_local:
             representation = expressions.LocalField(argument_name)
         else:
-            representation = expressions.ContextField(location)
+            representation = expressions.ContextField(location, tag_inferred_type)
 
         return (representation, non_existence_expression)
     else:

graphql_compiler/compiler/ir_lowering_common.py

@@ -53,30 +53,34 @@ def to_match(self):
         return mark_name
 
 
-def lower_context_field_existence(ir_blocks):
+def lower_context_field_existence(ir_blocks, query_metadata_table):
     """Lower ContextFieldExistence expressions into lower-level expressions."""
     def regular_visitor_fn(expression):
         """Expression visitor function that rewrites ContextFieldExistence expressions."""
         if not isinstance(expression, ContextFieldExistence):
             return expression
 
+        location_type = query_metadata_table.get_location_info(expression.location).type
+
         # Since this function is only used in blocks that aren't ConstructResult,
         # the location check is performed using a regular ContextField expression.
         return BinaryComposition(
             u'!=',
-            ContextField(expression.location),
+            ContextField(expression.location, location_type),
             NullLiteral)
 
     def construct_result_visitor_fn(expression):
         """Expression visitor function that rewrites ContextFieldExistence expressions."""
         if not isinstance(expression, ContextFieldExistence):
             return expression
 
+        location_type = query_metadata_table.get_location_info(expression.location).type
+
         # Since this function is only used in ConstructResult blocks,
         # the location check is performed using the special OutputContextVertex expression.
         return BinaryComposition(
             u'!=',
-            OutputContextVertex(expression.location),
+            OutputContextVertex(expression.location, location_type),
             NullLiteral)
 
     new_ir_blocks = []

graphql_compiler/compiler/ir_lowering_gremlin/__init__.py

@@ -39,7 +39,7 @@ def lower_ir(ir_blocks, query_metadata_table, type_equivalence_hints=None):
     """
     sanity_check_ir_blocks_from_frontend(ir_blocks, query_metadata_table)
 
-    ir_blocks = lower_context_field_existence(ir_blocks)
+    ir_blocks = lower_context_field_existence(ir_blocks, query_metadata_table)
     ir_blocks = optimize_boolean_expression_comparisons(ir_blocks)
 
     if type_equivalence_hints:

graphql_compiler/compiler/ir_lowering_match/__init__.py

@@ -80,17 +80,18 @@ def lower_ir(ir_blocks, query_metadata_table, type_equivalence_hints=None):
     # Append global operation block(s) to filter out incorrect results
     # from simple optional match traverses (using a WHERE statement)
     if len(simple_optional_root_info) > 0:
-        where_filter_predicate = construct_where_filter_predicate(simple_optional_root_info)
+        where_filter_predicate = construct_where_filter_predicate(
+            query_metadata_table, simple_optional_root_info)
         ir_blocks.insert(-1, GlobalOperationsStart())
         ir_blocks.insert(-1, Filter(where_filter_predicate))
 
     # These lowering / optimization passes work on IR blocks.
-    ir_blocks = lower_context_field_existence(ir_blocks)
+    ir_blocks = lower_context_field_existence(ir_blocks, query_metadata_table)
     ir_blocks = optimize_boolean_expression_comparisons(ir_blocks)
     ir_blocks = rewrite_binary_composition_inside_ternary_conditional(ir_blocks)
     ir_blocks = merge_consecutive_filter_clauses(ir_blocks)
     ir_blocks = lower_has_substring_binary_compositions(ir_blocks)
-    ir_blocks = orientdb_eval_scheduling.workaround_lowering_pass(ir_blocks)
+    ir_blocks = orientdb_eval_scheduling.workaround_lowering_pass(ir_blocks, query_metadata_table)
 
     # Here, we lower from raw IR blocks into a MatchQuery object.
     # From this point on, the lowering / optimization passes work on the MatchQuery representation.

graphql_compiler/compiler/ir_lowering_match/ir_lowering.py

@@ -13,7 +13,7 @@
 from ..blocks import Backtrack, CoerceType, MarkLocation, QueryRoot
 from ..expressions import (
     BinaryComposition, ContextField, ContextFieldExistence, FalseLiteral, FoldedContextField,
-    Literal, TernaryConditional, TrueLiteral
+    GlobalContextField, Literal, TernaryConditional, TrueLiteral
 )
 from ..helpers import FoldScopeLocation
 from .utils import convert_coerce_type_to_instanceof_filter
@@ -254,17 +254,24 @@ def _translate_equivalent_locations(match_query, location_translations):
 
     def visitor_fn(expression):
         """Expression visitor function used to rewrite expressions with updated Location data."""
-        if isinstance(expression, (ContextField, ContextFieldExistence)):
-            old_location = expression.location
+        if isinstance(expression, (ContextField, GlobalContextField)):
+            old_location = expression.location.at_vertex()
             new_location = location_translations.get(old_location, old_location)
+            if expression.location.field is not None:
+                new_location = new_location.navigate_to_field(expression.location.field)
 
             # The Expression could be one of many types, including:
             #   - ContextField
-            #   - ContextFieldExistence
+            #   - GlobalContextField
             # We determine its exact class to make sure we return an object of the same class
-            # as the replacement expression.
+            # as the expression being replaced.
             expression_cls = type(expression)
-            return expression_cls(new_location)
+            return expression_cls(new_location, expression.field_type)
+        elif isinstance(expression, ContextFieldExistence):
+            old_location = expression.location
+            new_location = location_translations.get(old_location, old_location)
+
+            return ContextFieldExistence(new_location)
         elif isinstance(expression, FoldedContextField):
             # Update the Location within FoldedContextField
             old_location = expression.fold_scope_location.base_location
@@ -322,8 +329,13 @@ def visitor_fn(expression):
     # Rewrite the Locations in the ConstructResult output block.
     new_output_block = match_query.output_block.visit_and_update_expressions(visitor_fn)
 
+    # Rewrite the Locations in the global where block.
+    new_where_block = None
+    if match_query.where_block is not None:
+        new_where_block = match_query.where_block.visit_and_update_expressions(visitor_fn)
+
     return match_query._replace(match_traversals=new_match_traversals, folds=new_folds,
-                                output_block=new_output_block)
+                                output_block=new_output_block, where_block=new_where_block)
 
 
 def lower_folded_coerce_types_into_filter_blocks(folded_ir_blocks):

graphql_compiler/compiler/ir_lowering_match/utils.py

@@ -9,7 +9,7 @@
     BinaryComposition, Expression, GlobalContextField, Literal, LocalField, NullLiteral,
     TrueLiteral, UnaryTransformation, ZeroLiteral
 )
-from ..helpers import Location, get_only_element_from_collection, is_vertex_field_name
+from ..helpers import get_only_element_from_collection, is_vertex_field_name
 
 
 def convert_coerce_type_to_instanceof_filter(coerce_type_block):
@@ -134,7 +134,8 @@ def filter_edge_field_non_existence(edge_expression):
     return BinaryComposition(u'||', field_null_check, field_size_check)
 
 
-def _filter_orientdb_simple_optional_edge(optional_edge_location, inner_location_name):
+def _filter_orientdb_simple_optional_edge(
+        query_metadata_table, optional_edge_location, inner_location_name):
     """Return an Expression that is False for rows that don't follow the @optional specification.
 
     OrientDB does not filter correctly within optionals. Namely, a result where the optional edge
@@ -162,6 +163,9 @@ def _filter_orientdb_simple_optional_edge(optional_edge_location, inner_location
         Here, the `optional_edge_location` is `Animal___1.out_Animal_ParentOf`.
 
     Args:
+        query_metadata_table: QueryMetadataTable object containing all metadata collected during
+                              query processing, including location metadata (e.g. which locations
+                              are folded or optional).
         optional_edge_location: Location object representing the optional edge field
         inner_location_name: string representing location within the corresponding optional traverse
 
@@ -171,13 +175,21 @@ def _filter_orientdb_simple_optional_edge(optional_edge_location, inner_location
     inner_local_field = LocalField(inner_location_name)
     inner_location_existence = BinaryComposition(u'!=', inner_local_field, NullLiteral)
 
-    edge_context_field = GlobalContextField(optional_edge_location)
+    # The optional_edge_location here is actually referring to the edge field itself.
+    # This is definitely non-standard, but required to get the proper semantics.
+    # To get its type, we construct the location of the vertex field on the other side of the edge.
+    vertex_location = (
+        optional_edge_location.at_vertex().navigate_to_subpath(optional_edge_location.field)
+    )
+    location_type = query_metadata_table.get_location_info(vertex_location).type
+
+    edge_context_field = GlobalContextField(optional_edge_location, location_type)
     edge_field_non_existence = filter_edge_field_non_existence(edge_context_field)
 
     return BinaryComposition(u'||', edge_field_non_existence, inner_location_existence)
 
 
-def construct_where_filter_predicate(simple_optional_root_info):
+def construct_where_filter_predicate(query_metadata_table, simple_optional_root_info):
     """Return an Expression that is True if and only if each simple optional filter is True.
 
     Construct filters for each simple optional, that are True if and only if `edge_field` does
@@ -186,6 +198,9 @@ def construct_where_filter_predicate(simple_optional_root_info):
     evaluate to True (conjunction).
 
     Args:
+        query_metadata_table: QueryMetadataTable object containing all metadata collected during
+                              query processing, including location metadata (e.g. which locations
+                              are folded or optional).
         simple_optional_root_info: dict mapping from simple_optional_root_location -> dict
                                    containing keys
                                    - 'inner_location_name': Location object correspoding to the
@@ -204,9 +219,9 @@ def construct_where_filter_predicate(simple_optional_root_info):
         inner_location_name = root_info_dict['inner_location_name']
         edge_field = root_info_dict['edge_field']
 
-        optional_edge_location = Location(root_location.query_path, field=edge_field)
+        optional_edge_location = root_location.navigate_to_field(edge_field)
         optional_edge_where_filter = _filter_orientdb_simple_optional_edge(
-            optional_edge_location, inner_location_name)
+            query_metadata_table, optional_edge_location, inner_location_name)
         inner_location_name_to_where_filter[inner_location_name] = optional_edge_where_filter
 
     # Sort expressions by inner_location_name to obtain deterministic order

graphql_compiler/compiler/workarounds/orientdb_eval_scheduling.py

@@ -16,21 +16,21 @@
 )
 
 
-def workaround_lowering_pass(ir_blocks):
+def workaround_lowering_pass(ir_blocks, query_metadata_table):
     """Extract locations from TernaryConditionals and rewrite their Filter blocks as necessary."""
     new_ir_blocks = []
 
     for block in ir_blocks:
         if isinstance(block, Filter):
-            new_block = _process_filter_block(block)
+            new_block = _process_filter_block(query_metadata_table, block)
         else:
             new_block = block
         new_ir_blocks.append(new_block)
 
     return new_ir_blocks
 
 
-def _process_filter_block(block):
+def _process_filter_block(query_metadata_table, block):
     """Rewrite the provided Filter block if necessary."""
     # For a given Filter block with BinaryComposition predicate expression X,
     # let L be the set of all Locations referenced in any TernaryConditional
@@ -85,7 +85,7 @@ def extract_locations_visitor(expression):
                                  u'{} {}'.format(ternary, return_value))
 
     tautologies = [
-        _create_tautological_expression_for_location(location)
+        _create_tautological_expression_for_location(query_metadata_table, location)
         for location in problematic_locations
     ]
 
@@ -98,8 +98,12 @@ def extract_locations_visitor(expression):
     return Filter(final_predicate)
 
 
-def _create_tautological_expression_for_location(location):
+def _create_tautological_expression_for_location(query_metadata_table, location):
     """For a given location, create a BinaryComposition that always evaluates to 'true'."""
-    location_exists = BinaryComposition(u'!=', ContextField(location), NullLiteral)
-    location_does_not_exist = BinaryComposition(u'=', ContextField(location), NullLiteral)
+    location_type = query_metadata_table.get_location_info(location).type
+
+    location_exists = BinaryComposition(
+        u'!=', ContextField(location, location_type), NullLiteral)
+    location_does_not_exist = BinaryComposition(
+        u'=', ContextField(location, location_type), NullLiteral)
     return BinaryComposition(u'||', location_exists, location_does_not_exist)