[core][feat] Rewrite queries to make them more efficient (#2093)

Author: aquamatthias

fixcore/fixcore/db/arango_query.py

@@ -11,6 +11,7 @@
 
 from fixcore.constants import less_greater_then_operations as lgt_ops, arangodb_matches_null_ops
 from fixcore.db import EstimatedSearchCost, EstimatedQueryCostRating as Rating
+from fixcore.db.arango_query_rewrite import rewrite_query
 from fixcore.db.arangodb_functions import as_arangodb_function
 from fixcore.db.model import QueryModel
 from fixcore.model.graph_access import Section, Direction
@@ -110,7 +111,7 @@ def to_query(
     id_column: str = "_key",
 ) -> Tuple[str, Json]:
     ctx = ArangoQueryContext()
-    query = query_model.query
+    query = rewrite_query(query_model)
     start = from_collection or f"`{db.graph_vertex_name()}`"
     cursor, query_str = query_string(db, query, query_model, start, with_edges, ctx, id_column=id_column)
     last_limit = f" LIMIT {ll.offset}, {ll.length}" if (ll := query.current_part.limit) else ""

fixcore/fixcore/db/arango_query_rewrite.py

@@ -0,0 +1,86 @@
+from typing import List, cast
+
+from attr import evolve
+
+from fixcore.db.model import QueryModel
+from fixcore.model.model import predefined_kinds_by_name
+from fixcore.query.model import (
+    Query,
+    Predicate,
+    FulltextTerm,
+    Term,
+    ContextTerm,
+    MergeTerm,
+    CombinedTerm,
+    IsTerm,
+    NotTerm,
+)
+
+
+def add_is_term(query_model: QueryModel) -> Query:
+    model = query_model.model
+
+    def and_combined(term: Term) -> bool:
+        if isinstance(term, NotTerm):
+            return True
+        return isinstance(term, CombinedTerm) and term.op == "and"
+
+    def no_context(term: Term) -> bool:
+        return not isinstance(term, (MergeTerm, ContextTerm))
+
+    def combine_term_if_possible(term: Term, predicates: List[Predicate]) -> Term:
+        kinds = set()
+        for pred in predicates:
+            for res in query_model.owners(pred.name):
+                if res.fqn not in predefined_kinds_by_name:
+                    kinds.add(res.fqn)
+        kinds.discard("resource")  # all resources have this base kind - ignore it
+        return IsTerm(kinds=sorted(kinds)).and_term(term) if kinds else term
+
+    def change_term(term: Term) -> Term:
+        if isinstance(term, CombinedTerm) and term.op == "or":
+            left = change_term(term.left)
+            right = change_term(term.right)
+            return evolve(term, left=left, right=right)
+        elif isinstance(term, CombinedTerm) and term.op == "and":
+            li = term.left.find_term(lambda t: isinstance(t, IsTerm), and_combined)
+            ri = term.right.find_term(lambda t: isinstance(t, IsTerm), and_combined)
+            if li is None and ri is None:
+                predicates = cast(List[Predicate], term.find_terms(lambda t: isinstance(t, Predicate), no_context))
+                contexts = cast(List[ContextTerm], term.find_terms(lambda t: isinstance(t, ContextTerm), no_context))
+                all_preds = predicates + [pred for ctx in contexts for pred in ctx.visible_predicates()]
+                return combine_term_if_possible(term, all_preds)
+        elif isinstance(term, NotTerm):
+            return NotTerm(change_term(term.term))
+        elif isinstance(term, Predicate):
+            return combine_term_if_possible(term, [term])
+        elif isinstance(term, ContextTerm):
+            return combine_term_if_possible(term, term.visible_predicates())
+        elif isinstance(term, MergeTerm):
+            pre = change_term(term.pre_filter)
+            post = change_term(term.post_filter) if term.post_filter else None
+            queries = [evolve(mq, query=add_is_term(QueryModel(mq.query, model))) for mq in term.merge]
+            return MergeTerm(pre_filter=pre, post_filter=post, merge=queries)
+        return term
+
+    part = query_model.query.first_part
+    part = evolve(part, term=change_term(part.term))
+    return evolve(query_model.query, parts=query_model.query.parts[:-1] + [part])
+
+
+def rewrite_query(
+    query_model: QueryModel,
+) -> Query:
+    q = query_model.query
+    p = q.first_part
+
+    #  check for single tags predicate. use fulltext index: tags.foo==bar --> "bar" and tags.foo==bar
+    if isinstance(p.term, Predicate) and p.term.name.startswith("reported.tags.") and p.term.op in ("==", "in"):
+        value = " ".join(p.term.value) if isinstance(p.term.value, list) else str(p.term.value)
+        part = evolve(p, term=FulltextTerm(value).and_term(p.term))
+        return evolve(q, parts=q.parts[:-1] + [part])
+
+    # try to add an IsTerm if not already provided
+    q = add_is_term(query_model)
+
+    return q

fixcore/fixcore/db/model.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import Dict, Any, Optional, Tuple
+from typing import Dict, Any, Optional, Tuple, List
 
 from attr import define
 
 from fixcore.model.graph_access import Section
-from fixcore.model.model import Model, ResolvedPropertyPath
+from fixcore.model.model import Model, ResolvedPropertyPath, ComplexKind
 from fixcore.model.resolve_in_graph import GraphResolver
 from fixcore.query.model import Query
 from fixcore.util import first
@@ -30,15 +30,23 @@ def is_set(self, name: str) -> bool:
                 return value.lower() in ["1", "true", "yes", "y"]
         return False
 
-    def prop_kind(self, path: str) -> Tuple[ResolvedPropertyPath, Optional[str]]:  # prop, merge_name
+    def __prop(self, path: str) -> Tuple[str, Optional[str]]:
         merge_name = first(lambda name: path.startswith(name + "."), self.query.merge_names) or first(
             lambda name: path.startswith(name + "."), ancestor_merges
         )
         # remove merge_name and section part (if existent) from the local_path
         lookup = Section.without_section(path[len(merge_name) + 1 :] if merge_name else path)  # noqa: E203
+        return lookup, merge_name
+
+    def prop_kind(self, path: str) -> Tuple[ResolvedPropertyPath, Optional[str]]:  # prop, merge_name
+        lookup, merge_name = self.__prop(path)
         resolved = self.model.property_by_path(lookup)
         return resolved, merge_name
 
+    def owners(self, path: str) -> List[ComplexKind]:
+        lookup, _ = self.__prop(path)
+        return self.model.owners_by_path(lookup)
+
 
 @define(repr=True, eq=True)
 class GraphUpdate(ABC):

fixcore/fixcore/model/model.py

@@ -6,6 +6,7 @@
 import sys
 import textwrap
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from datetime import datetime, timezone, date
 from functools import lru_cache
 from json import JSONDecodeError
@@ -920,6 +921,7 @@ def __init__(
         self.__resolved_props: Dict[str, Tuple[Property, Kind]] = {}
         self.__resolved_direct_props: Dict[str, Tuple[Property, Kind]] = {}
         self.__resolved_bases: Dict[str, ComplexKind] = {}
+        self.__owner_lookup: Dict[PropertyPath, ComplexKind] = {}
         self.__all_props: List[Property] = list(self.properties)
         self.__resolved_hierarchy: Set[str] = {fqn}
         self.__property_by_path: List[ResolvedPropertyPath] = []
@@ -988,7 +990,7 @@ def resolve(self, model: Dict[str, Kind]) -> None:
                         self.__resolved_hierarchy.update(base.__resolved_hierarchy)
 
             # property path -> kind
-            self.__property_by_path = ComplexKind.resolve_properties(self, model)
+            self.__property_by_path, self.__owner_lookup = ComplexKind.resolve_properties(self, model)
             self.__synthetic_props = [p for p in self.__property_by_path if p.prop.synthetic]
 
             # resolve predecessor kinds
@@ -1051,6 +1053,9 @@ def is_a(self, kind: str) -> bool:
     def resolved_property_paths(self) -> List[ResolvedPropertyPath]:
         return self.__property_by_path
 
+    def owned_paths(self) -> Dict[PropertyPath, ComplexKind]:
+        return self.__owner_lookup
+
     def resolved_bases(self) -> Dict[str, ComplexKind]:
         return self.__resolved_bases
 
@@ -1244,16 +1249,19 @@ def walk_element(
 
     @staticmethod
     def resolve_properties(
-        complex_kind: ComplexKind,
-        model: Dict[str, Kind],
-        from_path: PropertyPath = EmptyPath,
-        maybe_visited: Optional[Dict[str, PropertyPath]] = None,
-    ) -> List[ResolvedPropertyPath]:
-        visited = maybe_visited or {}
+        complex_kind: ComplexKind, model: Dict[str, Kind]
+    ) -> Tuple[List[ResolvedPropertyPath], Dict[PropertyPath, ComplexKind]]:
+        visited: Dict[str, PropertyPath] = {}
         result: List[ResolvedPropertyPath] = []
+        owner_lookup: Dict[PropertyPath, ComplexKind] = {}
 
         def path_for(
-            prop: Property, kind: Kind, path: PropertyPath, array: bool = False, add_prop_to_path: bool = True
+            owner: ComplexKind,
+            prop: Property,
+            kind: Kind,
+            path: PropertyPath,
+            array: bool = False,
+            add_prop_to_path: bool = True,
         ) -> None:
             prop_name = f"{prop.name}[]" if array else prop.name
             # Detect object cycles: remember the path when we have visited this property.
@@ -1267,27 +1275,35 @@ def path_for(
             kind.resolve(model)
             if isinstance(kind, SimpleKind):
                 result.append(ResolvedPropertyPath(relative, prop, kind))
+                owner_lookup[relative] = owner
             elif isinstance(kind, ArrayKind):
                 if name := relative.last_part:
                     result.append(ResolvedPropertyPath(relative, Property(name, kind.fqn), kind))
-                path_for(prop, kind.inner, path, True)
+                    owner_lookup[relative] = owner
+                path_for(owner, prop, kind.inner, path, True)
             elif isinstance(kind, DictionaryKind):
                 child = relative.child(None)
                 if name := relative.last_part:
                     result.append(ResolvedPropertyPath(relative, Property(name, kind.fqn), kind))
+                    owner_lookup[relative] = owner
                     # Any child path accessing this dictionary will get a property of value kind.
                     value = kind.value_kind
                     result.append(ResolvedPropertyPath(child, Property("any", value.fqn), value))
-                path_for(prop, kind.value_kind, child, add_prop_to_path=False)
+                    owner_lookup[child] = owner
+                path_for(owner, prop, kind.value_kind, child, add_prop_to_path=False)
             elif isinstance(kind, ComplexKind):
                 if name := relative.last_part:
                     result.append(ResolvedPropertyPath(relative, Property(name, kind.fqn), kind))
-                result.extend(ComplexKind.resolve_properties(kind, model, relative, visited))
+                    owner_lookup[relative] = owner
+                for_complex_kind(owner, kind, relative)
 
-        for x in complex_kind.all_props():
-            path_for(x, complex_kind.__resolved_props[x.name][1], from_path)
+        def for_complex_kind(owner: ComplexKind, current: ComplexKind, relative: PropertyPath) -> None:
+            for cpx in list(current.resolved_bases().values()) + [current]:
+                for prop in cpx.properties:
+                    path_for(owner, prop, cpx.__resolved_props[prop.name][1], relative)
 
-        return result
+        for_complex_kind(complex_kind, complex_kind, PropertyPath([], ""))
+        return result, owner_lookup
 
 
 string_kind = StringKind("string")
@@ -1362,29 +1378,39 @@ def path_for(
 class Model:
     @staticmethod
     def empty() -> Model:
-        return Model({}, [])
+        return Model({}, [], {})
 
     @staticmethod
     def from_kinds(kinds: List[Kind]) -> Model:
         all_kinds = kinds + predefined_kinds
         kind_dict = {kind.fqn: kind for kind in all_kinds}
         for kind in all_kinds:
             kind.resolve(kind_dict)
-        resolved = list(
-            # several complex kinds might have the same property
-            # reduce the list by hash over the path.
-            {
-                r.path: r
-                for c in all_kinds
-                if isinstance(c, ComplexKind) and c.aggregate_root
-                for r in c.resolved_property_paths()
-            }.values()
-        )
-        return Model(kind_dict, resolved)
+        # several complex kinds might have the same property
+        # reduce the list by hash over the path.
+        prop_kinds_by_path = {}
+        # lookup map to get the aggregate root that defined a specific property path
+        # Example: instance_cores
+        complex_by_path_distinct: Dict[PropertyPath, Dict[str, ComplexKind]] = defaultdict(dict)
+        for c in all_kinds:
+            if isinstance(c, ComplexKind) and c.aggregate_root:
+                for r in c.resolved_property_paths():
+                    prop_kinds_by_path[r.path] = r
+                for path, cpl in c.owned_paths().items():
+                    complex_by_path_distinct[path][cpl.fqn] = cpl
+
+        complex_by_path = {k: list(v.values()) for k, v in complex_by_path_distinct.items()}
+        return Model(kind_dict, list(prop_kinds_by_path.values()), complex_by_path)
 
-    def __init__(self, kinds: Dict[str, Kind], property_kind_by_path: List[ResolvedPropertyPath]):
+    def __init__(
+        self,
+        kinds: Dict[str, Kind],
+        property_kind_by_path: List[ResolvedPropertyPath],
+        complex_kinds_by_path: Dict[PropertyPath, List[ComplexKind]],
+    ):
         self.kinds = kinds
         self.__property_kind_by_path: List[ResolvedPropertyPath] = property_kind_by_path
+        self.__complex_kinds_by_path: Dict[PropertyPath, List[ComplexKind]] = complex_kinds_by_path
 
     def __contains__(self, name_or_object: Union[str, Json]) -> bool:
         if isinstance(name_or_object, str):
@@ -1434,6 +1460,10 @@ def property_by_path(self, path_: Union[str, List[str]]) -> ResolvedPropertyPath
     def kind_by_path(self, path: Union[str, List[str]]) -> Kind:
         return self.property_by_path(path).kind
 
+    def owners_by_path(self, path_: Union[str, List[str]]) -> List[ComplexKind]:
+        path = PropertyPath.from_string(path_) if isinstance(path_, str) else PropertyPath.from_list(path_)
+        return self.__complex_kinds_by_path.get(path, [])
+
     def coerce(self, js: Json) -> Json:
         try:
             kind: Kind = self[js["kind"]]
@@ -1619,7 +1649,7 @@ def all_predecessor_kinds(kind: ComplexKind) -> Dict[EdgeType, List[str]]:
                 )
             else:
                 result[kind.fqn] = kind
-        return Model(result, self.__property_kind_by_path)
+        return Model(result, self.__property_kind_by_path, self.__complex_kinds_by_path)
 
     def filter_complex(
         self, filter_fn: Callable[[ComplexKind], bool], with_bases: bool = True, with_prop_types: bool = True
@@ -1649,7 +1679,7 @@ def add_kind(cpl: ComplexKind) -> None:
             if isinstance(kind, ComplexKind) and filter_fn(kind):
                 add_kind(kind)
 
-        return Model(kinds, self.__property_kind_by_path)
+        return Model(kinds, self.__property_kind_by_path, self.__complex_kinds_by_path)
 
     def complete_path(
         self,