Optimize Jmespath with key lookup

Author: zprobst

nodestream/pipeline/value_providers/jmespath_value_provider.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from typing import Any, Iterable, Type
 
 import jmespath
@@ -8,6 +9,51 @@
 from .value_provider import ValueProvider, ValueProviderException
 
 
+# `QueryStrategy` is here to provide the seam for different optimizations
+# for executing jmespath queries. We can either execute a "fully fledged"
+# jmespath query or we can implement some simple access patterns that
+# are faster to execute. For example, if the expression is a simple key
+# lookup, we can just use the key directly instead of compiling the
+# jmespath expression and then executing it with all the weight and
+# overhead that comes with it.
+
+
+class QueryStrategy(ABC):
+    @classmethod
+    def from_string_expression(cls, expression: str):
+        if expression.isalpha():
+            return KeyLookup(expression)
+
+        compiled_query = jmespath.compile(expression)
+        return ExecuteJmespath(compiled_query)
+
+    @abstractmethod
+    def search(self, context: ProviderContext):
+        pass
+
+
+class ExecuteJmespath(QueryStrategy):
+    def __init__(self, compiled_query: ParsedResult) -> None:
+        self.compiled_query = compiled_query
+
+    def search(self, context: ProviderContext):
+        return self.compiled_query.search(context.document)
+
+    def __str__(self) -> str:
+        return str(self.compiled_query.expression)
+
+
+class KeyLookup(QueryStrategy):
+    def __init__(self, key: str) -> None:
+        self.key = key
+
+    def search(self, context: ProviderContext):
+        return context.document.get(self.key, None)
+
+    def __str__(self) -> str:
+        return self.key
+
+
 class JmespathValueProvider(ValueProvider):
     """A `ValueProvider` that uses JMESPath to extract values from a document."""
 
@@ -24,13 +70,13 @@ def install_yaml_tag(cls, loader: Type[SafeLoader]):
 
     @classmethod
     def from_string_expression(cls, expression: str):
-        return cls(jmespath.compile(expression))
+        return cls(QueryStrategy.from_string_expression(expression))
 
-    def __init__(self, compiled_query: ParsedResult) -> None:
-        self.compiled_query = compiled_query
+    def __init__(self, strategy: QueryStrategy) -> None:
+        self.strategy = strategy
 
     def search(self, context: ProviderContext):
-        raw_search = self.compiled_query.search(context.document)
+        raw_search = self.strategy.search(context)
         if raw_search is None:
             return
         if isinstance(raw_search, list):
@@ -51,14 +97,12 @@ def many_values(self, context: ProviderContext) -> Iterable[Any]:
             raise ValueProviderException(str(context.document), self) from e
 
     def __str__(self) -> str:
-        return (
-            f"JmespathValueProvider: { {'expression': self.compiled_query.expression} }"
-        )
+        return f"JmespathValueProvider: { {'expression': str(self.strategy)} }"
 
 
 SafeDumper.add_representer(
     JmespathValueProvider,
     lambda dumper, jmespath: dumper.represent_scalar(
-        "!jmespath", jmespath.compiled_query.expression
+        "!jmespath", str(jmespath.strategy)
     ),
 )

tests/unit/pipeline/value_providers/test_jmespath_value_provider.py

@@ -7,37 +7,37 @@
 
 
 def test_single_value_present(blank_context_with_document):
-    subject = JmespathValueProvider(jmespath.compile("team.name"))
+    subject = JmespathValueProvider.from_string_expression("team.name")
     assert_that(
         subject.single_value(blank_context_with_document), equal_to("nodestream")
     )
 
 
 def test_single_value_missing(blank_context_with_document):
-    subject = JmespathValueProvider(jmespath.compile("team.description"))
+    subject = JmespathValueProvider.from_string_expression("team.description")
     assert_that(subject.single_value(blank_context_with_document), none())
 
 
 def test_single_value_is_list(blank_context_with_document):
-    subject = JmespathValueProvider(jmespath.compile("project.tags"))
+    subject = JmespathValueProvider.from_string_expression("project.tags")
     result = subject.single_value(blank_context_with_document)
     assert_that(result, equal_to("graphdb"))
 
 
 def test_multiple_values_missing(blank_context_with_document):
-    subject = JmespathValueProvider(jmespath.compile("project.labels"))
+    subject = JmespathValueProvider.from_string_expression("team.description")
     assert_that(list(subject.many_values(blank_context_with_document)), has_length(0))
 
 
 def test_multiple_values_returns_one_value(blank_context_with_document):
-    subject = JmespathValueProvider(jmespath.compile("team.name"))
+    subject = JmespathValueProvider.from_string_expression("team.name")
     result = list(subject.many_values(blank_context_with_document))
     assert_that(result, has_length(1))
     assert_that(result[0], equal_to("nodestream"))
 
 
 def test_multiple_values_hit(blank_context_with_document):
-    subject = JmespathValueProvider(jmespath.compile("project.tags"))
+    subject = JmespathValueProvider.from_string_expression("project.tags")
     result = subject.many_values(blank_context_with_document)
     assert_that(list(result), equal_to(["graphdb", "python"]))
 
@@ -46,7 +46,7 @@ def test_single_value_error(blank_context_with_document):
     some_text_from_document = blank_context_with_document.document["team"]["name"]
     # this will error because team2 does not exist causing the join to throw an error
     expression_with_error = "join('/', [team.name || '', team2.name])"
-    subject = JmespathValueProvider(jmespath.compile(expression_with_error))
+    subject = JmespathValueProvider.from_string_expression(expression_with_error)
 
     with pytest.raises(ValueProviderException) as e_info:
         subject.single_value(blank_context_with_document)
@@ -59,7 +59,7 @@ def test_single_value_error(blank_context_with_document):
 def test_multiple_values_error(blank_context_with_document):
     # this will error because team2 does not exist causing the join to throw an error
     expression_with_error = "join('/', [team.name || '', team2.name])"
-    subject = JmespathValueProvider(jmespath.compile(expression_with_error))
+    subject = JmespathValueProvider.from_string_expression(expression_with_error)
 
     with pytest.raises(Exception) as e_info:
         generator = subject.many_values(blank_context_with_document)