Merge pull request #34 from microsoft/bug_fixes

Bug fixes

Author: niclasko

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

@@ -11,6 +11,13 @@ def __init__(self, limit: int):
         self._count = 0
         self._limit = limit
 
+    @property
+    def is_limit_reached(self) -> bool:
+        return self._count >= self._limit
+
+    def increment(self) -> None:
+        self._count += 1
+
     async def run(self) -> None:
         if self._count >= self._limit:
             return

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

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from ..ast_node import ASTNode
+from .limit import Limit
 from .projection import Projection
 
 if TYPE_CHECKING:
@@ -24,6 +25,7 @@ def __init__(self, expressions: List[ASTNode]) -> None:
         super().__init__(expressions)
         self._where: Optional['Where'] = None
         self._results: List[Dict[str, Any]] = []
+        self._limit: Optional[Limit] = None
 
     @property
     def where(self) -> Any:
@@ -35,16 +37,28 @@ def where(self) -> Any:
     def where(self, where: 'Where') -> None:
         self._where = where
 
+    @property
+    def limit(self) -> Optional[Limit]:
+        return self._limit
+
+    @limit.setter
+    def limit(self, limit: Limit) -> None:
+        self._limit = limit
+
     async def run(self) -> None:
         if not self.where:
             return
+        if self._limit is not None and self._limit.is_limit_reached:
+            return
         record: Dict[str, Any] = {}
         for expression, alias in self.expressions():
             raw = expression.value()
             # Deep copy objects to preserve their state
             value = copy.deepcopy(raw) if isinstance(raw, (dict, list)) else raw
             record[alias] = value
         self._results.append(record)
+        if self._limit is not None:
+            self._limit.increment()
 
     async def initialize(self) -> None:
         self._results = []

flowquery-py/src/parsing/parser.py

@@ -116,6 +116,9 @@ def _parse_tokenized(self, is_sub_query: bool = False) -> ASTNode:
             if self.token.is_union():
                 break
 
+            if self.token.is_eof():
+                break
+
             operation = self._parse_operation()
             if operation is None and not is_sub_query:
                 raise ValueError("Expected one of WITH, UNWIND, RETURN, LOAD, OR CALL")
@@ -145,8 +148,11 @@ def _parse_tokenized(self, is_sub_query: bool = False) -> ASTNode:
 
             limit = self._parse_limit()
             if limit is not None:
-                operation.add_sibling(limit)
-                operation = limit
+                if isinstance(operation, Return):
+                    operation.limit = limit
+                else:
+                    operation.add_sibling(limit)
+                    operation = limit
 
             previous = operation
 
@@ -539,13 +545,11 @@ def _parse_node(self) -> Optional[Node]:
         node.properties = dict(self._parse_properties())
         if identifier is not None and identifier in self._state.variables:
             reference = self._state.variables.get(identifier)
-            # Resolve through Expression -> Reference -> Node (e.g., after WITH)
-            ref_child = reference.first_child() if isinstance(reference, Expression) else None
-            if isinstance(ref_child, Reference):
-                inner = ref_child.referred
-                if isinstance(inner, Node):
-                    reference = inner
-            if reference is None or (not isinstance(reference, Node) and not isinstance(reference, Unwind)):
+            if reference is None or (
+                not isinstance(reference, Node)
+                and not isinstance(reference, Unwind)
+                and not isinstance(reference, Expression)
+            ):
                 raise ValueError(f"Undefined node reference: {identifier}")
             node = NodeReference(node, reference)
         elif identifier is not None:

flowquery-py/tests/compute/test_runner.py

@@ -925,6 +925,20 @@ async def test_limit(self):
         results = runner.results
         assert len(results) == 50
 
+    @pytest.mark.asyncio
+    async def test_limit_as_last_operation(self):
+        """Test limit as the last operation after return."""
+        runner = Runner(
+            """
+            unwind range(1, 10) as i
+            return i
+            limit 5
+            """
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 5
+
     @pytest.mark.asyncio
     async def test_range_lookup(self):
         """Test range lookup."""
@@ -1457,6 +1471,71 @@ async def test_match_with_referenced_to_previous_variable(self):
         assert results[0] == {"name1": "Person 1", "name2": "Person 2", "name3": "Person 3"}
         assert results[1] == {"name1": "Person 2", "name2": "Person 3", "name3": "Person 4"}
 
+    @pytest.mark.asyncio
+    async def test_match_with_aggregated_with_and_subsequent_match(self):
+        """Test match with aggregated WITH followed by another match using the same node reference."""
+        await Runner(
+            """
+            CREATE VIRTUAL (:AggUser) AS {
+                unwind [
+                    {id: 1, name: 'Alice'},
+                    {id: 2, name: 'Bob'},
+                    {id: 3, name: 'Carol'}
+                ] as record
+                RETURN record.id as id, record.name as name
+            }
+            """
+        ).run()
+        await Runner(
+            """
+            CREATE VIRTUAL (:AggUser)-[:KNOWS]-(:AggUser) AS {
+                unwind [
+                    {left_id: 1, right_id: 2},
+                    {left_id: 1, right_id: 3}
+                ] as record
+                RETURN record.left_id as left_id, record.right_id as right_id
+            }
+            """
+        ).run()
+        await Runner(
+            """
+            CREATE VIRTUAL (:AggProject) AS {
+                unwind [
+                    {id: 1, name: 'Project A'},
+                    {id: 2, name: 'Project B'}
+                ] as record
+                RETURN record.id as id, record.name as name
+            }
+            """
+        ).run()
+        await Runner(
+            """
+            CREATE VIRTUAL (:AggUser)-[:WORKS_ON]-(:AggProject) AS {
+                unwind [
+                    {left_id: 1, right_id: 1},
+                    {left_id: 1, right_id: 2}
+                ] as record
+                RETURN record.left_id as left_id, record.right_id as right_id
+            }
+            """
+        ).run()
+        match = Runner(
+            """
+            MATCH (u:AggUser)-[:KNOWS]->(s:AggUser)
+            WITH u, count(s) as acquaintances
+            MATCH (u)-[:WORKS_ON]->(p:AggProject)
+            RETURN u.name as name, acquaintances, collect(p.name) as projects
+            """
+        )
+        await match.run()
+        results = match.results
+        assert len(results) == 1
+        assert results[0] == {
+            "name": "Alice",
+            "acquaintances": 2,
+            "projects": ["Project A", "Project B"],
+        }
+
     @pytest.mark.asyncio
     async def test_match_and_return_full_node(self):
         """Test match and return full node."""

src/parsing/operations/group_by.ts

@@ -52,10 +52,12 @@ class GroupBy extends Projection {
         let node: Node = this.current;
         for (const mapper of this.mappers) {
             const value: any = mapper.value();
-            let child: Node | undefined = node.children.get(value);
+            const key: string =
+                typeof value === "object" && value !== null ? JSON.stringify(value) : String(value);
+            let child: Node | undefined = node.children.get(key);
             if (child === undefined) {
                 child = new Node(value);
-                node.children.set(value, child);
+                node.children.set(key, child);
             }
             node = child;
         }

src/parsing/operations/limit.ts

@@ -7,6 +7,12 @@ class Limit extends Operation {
         super();
         this.limit = limit;
     }
+    public get isLimitReached(): boolean {
+        return this.count >= this.limit;
+    }
+    public increment(): void {
+        this.count++;
+    }
     public async run(): Promise<void> {
         if (this.count >= this.limit) {
             return;
@@ -19,4 +25,4 @@ class Limit extends Operation {
     }
 }
 
-export default Limit;
+export default Limit;

src/parsing/operations/return.ts

@@ -1,3 +1,4 @@
+import Limit from "./limit";
 import Projection from "./projection";
 import Where from "./where";
 
@@ -15,6 +16,7 @@ import Where from "./where";
 class Return extends Projection {
     protected _where: Where | null = null;
     protected _results: Record<string, any>[] = [];
+    private _limit: Limit | null = null;
     public set where(where: Where) {
         this._where = where;
     }
@@ -24,17 +26,26 @@ class Return extends Projection {
         }
         return this._where.value();
     }
+    public set limit(limit: Limit) {
+        this._limit = limit;
+    }
     public async run(): Promise<void> {
         if (!this.where) {
             return;
         }
+        if (this._limit !== null && this._limit.isLimitReached) {
+            return;
+        }
         const record: Map<string, any> = new Map();
         for (const [expression, alias] of this.expressions()) {
             const raw = expression.value();
             const value: any = typeof raw === "object" && raw !== null ? structuredClone(raw) : raw;
             record.set(alias, value);
         }
         this._results.push(Object.fromEntries(record));
+        if (this._limit !== null) {
+            this._limit.increment();
+        }
     }
     public async initialize(): Promise<void> {
         this._results = [];

src/parsing/parser.ts

@@ -112,6 +112,9 @@ class Parser extends BaseParser {
             if (this.token.isUnion()) {
                 break;
             }
+            if (this.token.isEOF()) {
+                break;
+            }
             operation = this.parseOperation();
             if (operation === null && !isSubQuery) {
                 throw new Error("Expected one of WITH, UNWIND, RETURN, LOAD, OR CALL");
@@ -142,8 +145,12 @@ class Parser extends BaseParser {
             }
             const limit = this.parseLimit();
             if (limit !== null) {
-                operation!.addSibling(limit);
-                operation = limit;
+                if (operation instanceof Return) {
+                    (operation as Return).limit = limit;
+                } else {
+                    operation!.addSibling(limit);
+                    operation = limit;
+                }
             }
             previous = operation;
         }
@@ -494,16 +501,11 @@ class Parser extends BaseParser {
         node.label = label!;
         if (identifier !== null && this._state.variables.has(identifier)) {
             let reference = this._state.variables.get(identifier);
-            // Resolve through Expression -> Reference -> Node (e.g., after WITH)
-            if (reference instanceof Expression && reference.firstChild() instanceof Reference) {
-                const inner = (reference.firstChild() as Reference).referred;
-                if (inner instanceof Node) {
-                    reference = inner;
-                }
-            }
             if (
                 reference === undefined ||
-                (!(reference instanceof Node) && !(reference instanceof Unwind))
+                (!(reference instanceof Node) &&
+                    !(reference instanceof Unwind) &&
+                    !(reference instanceof Expression))
             ) {
                 throw new Error(`Undefined node reference: ${identifier}`);
             }

tests/compute/runner.test.ts

@@ -842,6 +842,19 @@ test("Test limit", async () => {
     expect(results.length).toBe(50);
 });
 
+test("Test limit as last operation", async () => {
+    const runner = new Runner(
+        `
+        unwind range(1, 10) as i
+        return i
+        limit 5
+        `
+    );
+    await runner.run();
+    const results = runner.results;
+    expect(results.length).toBe(5);
+});
+
 test("Test range lookup", async () => {
     const runner = new Runner(
         `
@@ -1327,6 +1340,60 @@ test("Test match with referenced to previous variable", async () => {
     expect(results[1]).toEqual({ name1: "Person 2", name2: "Person 3", name3: "Person 4" });
 });
 
+test("Test match with aggregated with and subsequent match", async () => {
+    await new Runner(`
+        CREATE VIRTUAL (:User) AS {
+            unwind [
+                {id: 1, name: 'Alice'},
+                {id: 2, name: 'Bob'},
+                {id: 3, name: 'Carol'}
+            ] as record
+            RETURN record.id as id, record.name as name
+        }
+    `).run();
+    await new Runner(`
+        CREATE VIRTUAL (:User)-[:KNOWS]-(:User) AS {
+            unwind [
+                {left_id: 1, right_id: 2},
+                {left_id: 1, right_id: 3}
+            ] as record
+            RETURN record.left_id as left_id, record.right_id as right_id
+        }
+    `).run();
+    await new Runner(`
+        CREATE VIRTUAL (:Project) AS {
+            unwind [
+                {id: 1, name: 'Project A'},
+                {id: 2, name: 'Project B'}
+            ] as record
+            RETURN record.id as id, record.name as name
+        }
+    `).run();
+    await new Runner(`
+        CREATE VIRTUAL (:User)-[:WORKS_ON]-(:Project) AS {
+            unwind [
+                {left_id: 1, right_id: 1},
+                {left_id: 1, right_id: 2}
+            ] as record
+            RETURN record.left_id as left_id, record.right_id as right_id
+        }
+    `).run();
+    const match = new Runner(`
+        MATCH (u:User)-[:KNOWS]->(s:User)
+        WITH u, count(s) as acquaintances
+        MATCH (u)-[:WORKS_ON]->(p:Project)
+        RETURN u.name as name, acquaintances, collect(p.name) as projects
+    `);
+    await match.run();
+    const results = match.results;
+    expect(results.length).toBe(1);
+    expect(results[0]).toEqual({
+        name: "Alice",
+        acquaintances: 2,
+        projects: ["Project A", "Project B"],
+    });
+});
+
 test("Test match and return full node", async () => {
     await new Runner(`
         CREATE VIRTUAL (:Person) AS {