Refactors Query Language for Thicket

hatchet/graphframe.py

@@ -17,7 +17,13 @@
 from .node import Node
 from .graph import Graph
 from .frame import Frame
-from .query import AbstractQuery, QueryMatcher, parse_cypher_query
+from .query import (
+    is_hatchet_query,
+    ObjectQuery,
+    parse_string_dialect,
+    QueryEngine,
+    AbstractQuery,
+)
 from .external.console import ConsoleRenderer
 from .util.dot import trees_to_dot
 from .util.deprecated import deprecated_params
@@ -86,6 +92,7 @@ def __init__(
         self.inc_metrics = [] if inc_metrics is None else inc_metrics
         self.default_metric = default_metric
         self.metadata = metadata
+        self.query_engine = QueryEngine()
 
     @staticmethod
     def from_hpctoolkit(dirname):
@@ -466,16 +473,16 @@ def filter(
                 filtered_rows = dataframe_copy.apply(filter_obj, axis=1)
                 filtered_df = dataframe_copy[filtered_rows]
 
-        elif isinstance(filter_obj, (list, str)) or issubclass(
-            type(filter_obj), AbstractQuery
-        ):
+        elif isinstance(filter_obj, (list, str)) or is_hatchet_query(filter_obj):
             # use a callpath query to apply the filter
             query = filter_obj
             if isinstance(filter_obj, list):
-                query = QueryMatcher(filter_obj)
+                query = ObjectQuery(filter_obj)
             elif isinstance(filter_obj, str):
-                query = parse_cypher_query(filter_obj)
-            query_matches = query.apply(self)
+                query = parse_string_dialect(filter_obj)
+            elif issubclass(type(filter_obj), AbstractQuery):
+                query = filter_obj._get_new_query()
+            query_matches = self.query_engine.apply(query, self.graph, self.dataframe)
             # match_set = list(set().union(*query_matches))
             # filtered_df = dataframe_copy.loc[dataframe_copy["node"].isin(match_set)]
             filtered_df = dataframe_copy.loc[dataframe_copy["node"].isin(query_matches)]

hatchet/query/__init__.py

@@ -0,0 +1,96 @@
+# Copyright 2017-2023 Lawrence Livermore National Security, LLC and other
+# Hatchet Project Developers. See the top-level LICENSE file for details.
+#
+# SPDX-License-Identifier: MIT
+
+# Make flake8 ignore unused names in this file
+# flake8: noqa: F401
+
+from .query import Query
+from .compound import (
+    CompoundQuery,
+    ConjunctionQuery,
+    DisjunctionQuery,
+    ExclusiveDisjunctionQuery,
+    NegationQuery,
+)
+from .object_dialect import ObjectQuery
+from .string_dialect import StringQuery, parse_string_dialect
+from .engine import QueryEngine
+from .errors import (
+    InvalidQueryPath,
+    InvalidQueryFilter,
+    RedundantQueryFilterWarning,
+    BadNumberNaryQueryArgs,
+)
+
+from .compat import (
+    AbstractQuery,
+    NaryQuery,
+    AndQuery,
+    IntersectionQuery,
+    OrQuery,
+    UnionQuery,
+    XorQuery,
+    SymDifferenceQuery,
+    NotQuery,
+    QueryMatcher,
+    CypherQuery,
+    parse_cypher_query,
+)
+
+
+def combine_via_conjunction(query0, query1):
+    return ConjunctionQuery(query0, query1)
+
+
+def combine_via_disjunction(query0, query1):
+    return DisjunctionQuery(query0, query1)
+
+
+def combine_via_exclusive_disjunction(query0, query1):
+    return ExclusiveDisjunctionQuery(query0, query1)
+
+
+def negate_query(query):
+    return NegationQuery(query)
+
+
+Query.__and__ = combine_via_conjunction
+Query.__or__ = combine_via_disjunction
+Query.__xor__ = combine_via_exclusive_disjunction
+Query.__not__ = negate_query
+
+
+CompoundQuery.__and__ = combine_via_conjunction
+CompoundQuery.__or__ = combine_via_disjunction
+CompoundQuery.__xor__ = combine_via_exclusive_disjunction
+CompoundQuery.__not__ = negate_query
+
+
+def is_hatchet_query(query_obj):
+    return (
+        issubclass(type(query_obj), Query)
+        or issubclass(type(query_obj), CompoundQuery)
+        or issubclass(type(query_obj), AbstractQuery)
+    )
+
+
+# Uncomment when accessing old-style queries using
+# 'from hatchet.query import *' is removed
+#
+# __all__ = [
+#     "Query",
+#     "CompoundQuery",
+#     "ConjunctionQuery",
+#     "DisjunctionQuery",
+#     "ExclusiveDisjunctionQuery",
+#     "NegationQuery",
+#     "ObjectQuery",
+#     "StringQuery",
+#     "parse_string_dialect",
+#     "InvalidQueryFilter",
+#     "InvalidQueryPath",
+#     "RedundantQueryFilterWarning",
+#     "BadNumberNaryQueryArgs",
+# ]