Merge pull request #63 from microsoft/graph_pattern_traversal_improvements

Graph pattern traversal improvements

Author: niclasko

flowquery-py/src/graph/node.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Optional
 
 from ..parsing.ast_node import ASTNode
 from ..parsing.expressions.expression import Expression
@@ -29,7 +29,6 @@ def __init__(
         self._incoming: Optional['Relationship'] = None
         self._outgoing: Optional['Relationship'] = None
         self._data: Optional['NodeData'] = None
-        self._todo_next: Optional[Callable[[], Union[None, Awaitable[None]]]] = None
 
     @property
     def identifier(self) -> Optional[str]:
@@ -112,7 +111,7 @@ def incoming(self, relationship: Optional['Relationship']) -> None:
     def set_data(self, data: Optional['NodeData']) -> None:
         self._data = data
 
-    async def next(self) -> None:
+    async def next(self) -> AsyncIterator[None]:
         if self._data:
             self._data.reset()
             while self._data.next():
@@ -122,10 +121,12 @@ async def next(self) -> None:
                     if not self._matches_properties():
                         continue
                     if self._outgoing and self._value:
-                        await self._outgoing.find(self._value['id'])
-                    await self.run_todo_next()
+                        async for _ in self._outgoing.find(self._value['id']):
+                            yield
+                    else:
+                        yield
 
-    async def find(self, id_: str, hop: int = 0) -> None:
+    async def find(self, id_: str, hop: int = 0) -> AsyncIterator[None]:
         if self._data:
             self._data.reset()
             while self._data.find(id_, hop):
@@ -137,19 +138,7 @@ async def find(self, id_: str, hop: int = 0) -> None:
                     if self._incoming:
                         self._incoming.set_end_node(self)
                     if self._outgoing and self._value:
-                        await self._outgoing.find(self._value['id'], hop)
-                    await self.run_todo_next()
-
-    @property
-    def todo_next(self) -> Optional[Callable[[], Union[None, Awaitable[None]]]]:
-        return self._todo_next
-
-    @todo_next.setter
-    def todo_next(self, func: Optional[Callable[[], Union[None, Awaitable[None]]]]) -> None:
-        self._todo_next = func
-
-    async def run_todo_next(self) -> None:
-        if self._todo_next:
-            result = self._todo_next()
-            if result is not None:
-                await result
+                        async for _ in self._outgoing.find(self._value['id'], hop):
+                            yield
+                    else:
+                        yield

flowquery-py/src/graph/node_reference.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any, AsyncIterator, Optional
 
 from ..parsing.ast_node import ASTNode
 from .node import Node
@@ -30,22 +30,26 @@ def referred(self) -> ASTNode:
     def value(self) -> Optional[Any]:
         return self._reference.value() if self._reference else None
 
-    async def next(self) -> None:
+    async def next(self) -> AsyncIterator[None]:
         """Process next using the referenced node's value."""
         ref_value = self._reference.value()
         if ref_value is None:
             return
         self.set_value(dict(ref_value))
         if self._outgoing and self._value:
-            await self._outgoing.find(self._value['id'])
-        await self.run_todo_next()
+            async for _ in self._outgoing.find(self._value['id']):
+                yield
+        else:
+            yield
 
-    async def find(self, id_: str, hop: int = 0) -> None:
+    async def find(self, id_: str, hop: int = 0) -> AsyncIterator[None]:
         """Find by ID, only matching if it equals the referenced node's ID."""
         referenced = self._reference.value()
         if referenced is None or id_ != referenced.get('id'):
             return
         self.set_value(dict(referenced))
         if self._outgoing and self._value:
-            await self._outgoing.find(self._value['id'], hop)
-        await self.run_todo_next()
+            async for _ in self._outgoing.find(self._value['id'], hop):
+                yield
+        else:
+            yield

flowquery-py/src/graph/pattern.py

@@ -112,4 +112,5 @@ async def initialize(self) -> None:
     async def traverse(self) -> None:
         first = self.first_node()
         if first and isinstance(first, Node):
-            await first.next()
+            async for _ in first.next():
+                pass

flowquery-py/src/graph/pattern_expression.py

@@ -51,13 +51,9 @@ async def evaluate(self) -> None:
         Sets _evaluation to True if the pattern is matched, False otherwise.
         """
         self._evaluation = False
-
-        async def set_evaluation_true() -> None:
+        async for _ in self.start_node.next():
             self._evaluation = True
 
-        self.end_node.todo_next = set_evaluation_true
-        await self.start_node.next()
-
     def value(self) -> Any:
         """Returns the result of the pattern evaluation."""
         return self._evaluation

flowquery-py/src/graph/patterns.py

@@ -1,6 +1,6 @@
 """Collection of graph patterns for FlowQuery."""
 
-from typing import Awaitable, Callable, List, Optional
+from typing import AsyncIterator, List, Optional
 
 from .pattern import Pattern
 
@@ -10,33 +10,25 @@ class Patterns:
 
     def __init__(self, patterns: Optional[List[Pattern]] = None) -> None:
         self._patterns = patterns or []
-        self._to_do_next: Optional[Callable[[], Awaitable[None]]] = None
 
     @property
     def patterns(self) -> List[Pattern]:
         return self._patterns
 
-    @property
-    def to_do_next(self) -> Optional[Callable[[], Awaitable[None]]]:
-        return self._to_do_next
-
-    @to_do_next.setter
-    def to_do_next(self, func: Optional[Callable[[], Awaitable[None]]]) -> None:
-        self._to_do_next = func
-        if self._patterns:
-            self._patterns[-1].end_node.todo_next = func
-
     async def initialize(self) -> None:
-        previous: Optional[Pattern] = None
         for pattern in self._patterns:
-            await pattern.fetch_data()  # Ensure data is loaded
-            if previous is not None:
-                # Chain the patterns together
-                async def next_pattern_start(p: Pattern = pattern) -> None:
-                    await p.start_node.next()
-                previous.end_node.todo_next = next_pattern_start
-            previous = pattern
-
-    async def traverse(self) -> None:
-        if self._patterns:
-            await self._patterns[0].start_node.next()
+            await pattern.fetch_data()
+
+    async def traverse(self) -> AsyncIterator[None]:
+        if not self._patterns:
+            return
+        async for _ in self._chain_patterns(0):
+            yield
+
+    async def _chain_patterns(self, index: int) -> AsyncIterator[None]:
+        async for _ in self._patterns[index].start_node.next():
+            if index + 1 < len(self._patterns):
+                async for _ in self._chain_patterns(index + 1):
+                    yield
+            else:
+                yield

flowquery-py/src/graph/relationship.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Union
 
 from ..parsing.ast_node import ASTNode
 from .hops import Hops
@@ -151,44 +151,49 @@ def set_end_node(self, node: 'Node') -> None:
     def _left_id_or_right_id(self) -> str:
         return "left_id" if self._direction == "left" else "right_id"
 
-    async def find(self, left_id: str, hop: int = 0) -> None:
+    async def find(self, left_id: str, hop: int = 0) -> AsyncIterator[None]:
         """Find relationships starting from the given node ID."""
         # Save original source node
         original = self._source
         if hop > 0:
             # For hops greater than 0, the source becomes the target of the previous hop
             self._source = self._target
-        if hop == 0:
-            if self._data:
-                self._data.reset()
-
-            # Handle zero-hop case: when min is 0 on a variable-length relationship,
-            # match source node as target (no traversal)
-            if self._hops and self._hops.multi() and self._hops.min == 0 and self._target:
-                # For zero-hop, target finds the same node as source (left_id)
-                # No relationship match is pushed since no edge is traversed
-                await self._target.find(left_id, hop)
-
-        while self._data and self._data.find(left_id, hop, self._direction):
-            data = self._data.current(hop)
-            if data is None:
-                continue
-            id = data[self._left_id_or_right_id()]
-            if hop + 1 >= self._hops.min:
-                self.set_value(self, left_id)
-                if not self._matches_properties(hop):
+        try:
+            if hop == 0:
+                if self._data:
+                    self._data.reset()
+
+                # Handle zero-hop case: when min is 0 on a variable-length relationship,
+                # match source node as target (no traversal)
+                if self._hops and self._hops.multi() and self._hops.min == 0 and self._target:
+                    # For zero-hop, target finds the same node as source (left_id)
+                    # No relationship match is pushed since no edge is traversed
+                    async for _ in self._target.find(left_id, hop):
+                        yield
+
+            while self._data and self._data.find(left_id, hop, self._direction):
+                data = self._data.current(hop)
+                if data is None:
                     continue
-                if self._target:
-                    await self._target.find(id, hop)
-                if hop + 1 < self._hops.max:
-                    if self._matches.is_circular(id):
-                        self._matches.pop()
+                id = data[self._left_id_or_right_id()]
+                if hop + 1 >= self._hops.min:
+                    self.set_value(self, left_id)
+                    if not self._matches_properties(hop):
                         continue
-                    await self.find(id, hop + 1)
-                self._matches.pop()
-            else:
-                # Below minimum hops: traverse the edge without yielding a match
-                await self.find(id, hop + 1)
-
-        # Restore original source node
-        self._source = original
+                    if self._target:
+                        async for _ in self._target.find(id, hop):
+                            yield
+                    if hop + 1 < self._hops.max:
+                        if self._matches.is_circular(id):
+                            self._matches.pop()
+                            continue
+                        async for _ in self.find(id, hop + 1):
+                            yield
+                    self._matches.pop()
+                else:
+                    # Below minimum hops: traverse the edge without yielding a match
+                    async for _ in self.find(id, hop + 1):
+                        yield
+        finally:
+            # Restore original source node
+            self._source = original

flowquery-py/src/parsing/data_structures/lookup.py

@@ -44,8 +44,13 @@ def value(self) -> Any:
         # Try dict-like access first, then fall back to attribute access for objects
         try:
             return obj[key]
-        except (TypeError, KeyError):
+        except KeyError:
             # For objects with attributes (like dataclasses), use getattr
+            if hasattr(obj, key):
+                return getattr(obj, key)
+            # Return None for missing keys, matching JavaScript obj[key] behavior
+            return None
+        except TypeError:
             if hasattr(obj, key):
                 return getattr(obj, key)
             raise

flowquery-py/src/parsing/operations/match.py

@@ -41,15 +41,11 @@ async def run(self) -> None:
         await self._patterns.initialize()
         matched = False
 
-        async def to_do_next() -> None:
-            nonlocal matched
+        async for _ in self._patterns.traverse():
             matched = True
             if self.next:
                 await self.next.run()
 
-        self._patterns.to_do_next = to_do_next
-        await self._patterns.traverse()
-
         # For OPTIONAL MATCH: if nothing matched, continue with None values
         if not matched and self._optional:
             for pattern in self._patterns.patterns:

flowquery-py/tests/compute/test_runner.py

@@ -1113,6 +1113,42 @@ async def test_return_with_expression_alias_which_starts_with_keyword(self):
         assert len(results) == 1
         assert results[0] == {"return1": 1, "notes": ["hello", "world"]}
 
+    @pytest.mark.asyncio
+    async def test_lookup_missing_property_returns_null(self):
+        """Test that accessing a missing property returns null instead of raising KeyError."""
+        runner = Runner('RETURN {a: 1}.b as result')
+        await runner.run()
+        results = runner.results
+        assert len(results) == 1
+        assert results[0] == {"result": None}
+
+    @pytest.mark.asyncio
+    async def test_lookup_missing_property_bracket_notation_returns_null(self):
+        """Test that bracket notation on a missing property returns null."""
+        runner = Runner('RETURN {a: 1}["b"] as result')
+        await runner.run()
+        results = runner.results
+        assert len(results) == 1
+        assert results[0] == {"result": None}
+
+    @pytest.mark.asyncio
+    async def test_lookup_missing_property_with_coalesce(self):
+        """Test coalesce with a missing property lookup."""
+        runner = Runner('RETURN coalesce({a: 1}.b, "default") as result')
+        await runner.run()
+        results = runner.results
+        assert len(results) == 1
+        assert results[0] == {"result": "default"}
+
+    @pytest.mark.asyncio
+    async def test_lookup_on_null_returns_null(self):
+        """Test that lookup on null returns null."""
+        runner = Runner('WITH null as obj RETURN obj.x as result')
+        await runner.run()
+        results = runner.results
+        assert len(results) == 1
+        assert results[0] == {"result": None}
+
     @pytest.mark.asyncio
     async def test_return_with_where_clause(self):
         """Test return with where clause."""