Merge pull request #41 from microsoft/order_by_expressions_support

Order by expressions support

Author: niclasko

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

@@ -1,15 +1,19 @@
 """Represents an ORDER BY operation that sorts results."""
 
-from typing import Any, Dict, List
+import functools
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from .operation import Operation
 
+if TYPE_CHECKING:
+    from ..expressions.expression import Expression
+
 
 class SortField:
-    """A single sort specification: field name and direction."""
+    """A single sort specification: expression and direction."""
 
-    def __init__(self, field: str, direction: str = "asc"):
-        self.field = field
+    def __init__(self, expression: 'Expression', direction: str = "asc"):
+        self.expression = expression
         self.direction = direction
 
 
@@ -19,27 +23,63 @@ class OrderBy(Operation):
     Can be attached to a RETURN operation (sorting its results),
     or used as a standalone accumulating operation after a non-aggregate WITH.
 
-    Example:
+    Supports both simple field references and arbitrary expressions:
+
+    Example::
+
         RETURN x ORDER BY x DESC
+        RETURN x ORDER BY toLower(x.name) ASC
+        RETURN x ORDER BY string_distance(toLower(x.name), toLower('Thomas')) ASC
     """
 
     def __init__(self, fields: List[SortField]):
         super().__init__()
         self._fields = fields
         self._results: List[Dict[str, Any]] = []
+        self._sort_keys: List[List[Any]] = []
 
     @property
     def fields(self) -> List[SortField]:
         return self._fields
 
-    def sort(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """Sorts an array of records according to the sort fields."""
-        import functools
+    def capture_sort_keys(self) -> None:
+        """Evaluate every sort-field expression against the current runtime
+        context and store the resulting values.  Must be called once per
+        accumulated row (from ``Return.run()``)."""
+        self._sort_keys.append([f.expression.value() for f in self._fields])
 
-        def compare(a: Dict[str, Any], b: Dict[str, Any]) -> int:
-            for sf in self._fields:
-                a_val = a.get(sf.field)
-                b_val = b.get(sf.field)
+    def sort(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Sort records using pre-computed sort keys captured during
+        accumulation.  When no keys have been captured (e.g. aggregated
+        returns), falls back to looking up simple reference identifiers
+        in each record."""
+        from ..expressions.reference import Reference
+
+        use_keys = len(self._sort_keys) == len(records)
+        keys = self._sort_keys
+
+        # Pre-compute fallback field names for when sort keys aren't
+        # available (aggregated returns).
+        fallback_fields: List[Optional[str]] = []
+        for f in self._fields:
+            root = f.expression.first_child()
+            if isinstance(root, Reference) and f.expression.child_count() == 1:
+                fallback_fields.append(root.identifier)
+            else:
+                fallback_fields.append(None)
+
+        indices = list(range(len(records)))
+
+        def compare(ai: int, bi: int) -> int:
+            for f_idx, sf in enumerate(self._fields):
+                if use_keys:
+                    a_val = keys[ai][f_idx]
+                    b_val = keys[bi][f_idx]
+                elif fallback_fields[f_idx] is not None:
+                    a_val = records[ai].get(fallback_fields[f_idx])  # type: ignore[arg-type]
+                    b_val = records[bi].get(fallback_fields[f_idx])  # type: ignore[arg-type]
+                else:
+                    continue
                 cmp = 0
                 if a_val is None and b_val is None:
                     cmp = 0
@@ -55,7 +95,8 @@ def compare(a: Dict[str, Any], b: Dict[str, Any]) -> int:
                     return -cmp if sf.direction == "desc" else cmp
             return 0
 
-        return sorted(records, key=functools.cmp_to_key(compare))
+        indices.sort(key=functools.cmp_to_key(compare))
+        return [records[i] for i in indices]
 
     async def run(self) -> None:
         """When used as a standalone operation, passes through to next."""
@@ -64,6 +105,7 @@ async def run(self) -> None:
 
     async def initialize(self) -> None:
         self._results = []
+        self._sort_keys = []
         if self.next:
             await self.next.initialize()
 

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

@@ -68,6 +68,9 @@ async def run(self) -> None:
             # Deep copy objects to preserve their state
             value = copy.deepcopy(raw) if isinstance(raw, (dict, list)) else raw
             record[alias] = value
+        # Capture sort-key values while expression bindings are still live.
+        if self._order_by is not None:
+            self._order_by.capture_sort_keys()
         self._results.append(record)
         if self._order_by is None and self._limit is not None:
             self._limit.increment()

flowquery-py/src/parsing/parser.py

@@ -767,10 +767,9 @@ def _parse_order_by(self) -> Optional[OrderBy]:
         self._expect_and_skip_whitespace_and_comments()
         fields: list[SortField] = []
         while True:
-            if not self.token.is_identifier_or_keyword():
-                raise ValueError("Expected field name in ORDER BY")
-            field = self.token.value
-            self.set_next_token()
+            expression = self._parse_expression()
+            if expression is None:
+                raise ValueError("Expected expression in ORDER BY")
             self._skip_whitespace_and_comments()
             direction = "asc"
             if self.token.is_asc():
@@ -781,7 +780,7 @@ def _parse_order_by(self) -> Optional[OrderBy]:
                 direction = "desc"
                 self.set_next_token()
                 self._skip_whitespace_and_comments()
-            fields.append(SortField(field, direction))
+            fields.append(SortField(expression, direction))
             if self.token.is_comma():
                 self.set_next_token()
                 self._skip_whitespace_and_comments()

flowquery-py/tests/compute/test_runner.py

@@ -4408,6 +4408,133 @@ async def test_order_by_with_where(self):
         assert results[3] == {"x": 4}
         assert results[4] == {"x": 3}
 
+    @pytest.mark.asyncio
+    async def test_order_by_with_property_access_expression(self):
+        """Test ORDER BY with property access expression."""
+        runner = Runner(
+            "unwind [{name: 'Charlie', age: 30}, {name: 'Alice', age: 25}, {name: 'Bob', age: 35}] as person "
+            "return person.name as name, person.age as age "
+            "order by person.name asc"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"name": "Alice", "age": 25}
+        assert results[1] == {"name": "Bob", "age": 35}
+        assert results[2] == {"name": "Charlie", "age": 30}
+
+    @pytest.mark.asyncio
+    async def test_order_by_with_function_expression(self):
+        """Test ORDER BY with function expression."""
+        runner = Runner(
+            "unwind ['BANANA', 'apple', 'Cherry'] as fruit "
+            "return fruit "
+            "order by toLower(fruit)"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"fruit": "apple"}
+        assert results[1] == {"fruit": "BANANA"}
+        assert results[2] == {"fruit": "Cherry"}
+
+    @pytest.mark.asyncio
+    async def test_order_by_with_function_expression_descending(self):
+        """Test ORDER BY with function expression descending."""
+        runner = Runner(
+            "unwind ['BANANA', 'apple', 'Cherry'] as fruit "
+            "return fruit "
+            "order by toLower(fruit) desc"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"fruit": "Cherry"}
+        assert results[1] == {"fruit": "BANANA"}
+        assert results[2] == {"fruit": "apple"}
+
+    @pytest.mark.asyncio
+    async def test_order_by_with_nested_function_expression(self):
+        """Test ORDER BY with nested function expression."""
+        runner = Runner(
+            "unwind ['Alice', 'Bob', 'ALICE', 'bob'] as name "
+            "return name "
+            "order by string_distance(toLower(name), toLower('alice')) asc"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 4
+        # 'Alice' and 'ALICE' have distance 0 from 'alice', should come first
+        assert results[0]["name"] == "Alice"
+        assert results[1]["name"] == "ALICE"
+        # 'Bob' and 'bob' have higher distance from 'alice'
+        assert results[2]["name"] == "Bob"
+        assert results[3]["name"] == "bob"
+
+    @pytest.mark.asyncio
+    async def test_order_by_with_arithmetic_expression(self):
+        """Test ORDER BY with arithmetic expression."""
+        runner = Runner(
+            "unwind [{a: 3, b: 1}, {a: 1, b: 5}, {a: 2, b: 2}] as item "
+            "return item.a as a, item.b as b "
+            "order by item.a + item.b asc"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"a": 3, "b": 1}  # sum = 4
+        assert results[1] == {"a": 2, "b": 2}  # sum = 4
+        assert results[2] == {"a": 1, "b": 5}  # sum = 6
+
+    @pytest.mark.asyncio
+    async def test_order_by_expression_does_not_leak_synthetic_keys(self):
+        """Test ORDER BY expression does not leak synthetic keys."""
+        runner = Runner(
+            "unwind ['B', 'a', 'C'] as x "
+            "return x "
+            "order by toLower(x) asc"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        # Results should only contain 'x', no extra keys
+        for r in results:
+            assert list(r.keys()) == ["x"]
+        assert results[0] == {"x": "a"}
+        assert results[1] == {"x": "B"}
+        assert results[2] == {"x": "C"}
+
+    @pytest.mark.asyncio
+    async def test_order_by_with_expression_and_limit(self):
+        """Test ORDER BY with expression and limit."""
+        runner = Runner(
+            "unwind ['BANANA', 'apple', 'Cherry', 'date', 'ELDERBERRY'] as fruit "
+            "return fruit "
+            "order by toLower(fruit) asc "
+            "limit 3"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"fruit": "apple"}
+        assert results[1] == {"fruit": "BANANA"}
+        assert results[2] == {"fruit": "Cherry"}
+
+    @pytest.mark.asyncio
+    async def test_order_by_with_mixed_simple_and_expression_fields(self):
+        """Test ORDER BY with mixed simple and expression fields."""
+        runner = Runner(
+            "unwind [{name: 'Alice', score: 3}, {name: 'Alice', score: 1}, {name: 'Bob', score: 2}] as item "
+            "return item.name as name, item.score as score "
+            "order by name asc, item.score desc"
+        )
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"name": "Alice", "score": 3}  # Alice, score 3 desc
+        assert results[1] == {"name": "Alice", "score": 1}  # Alice, score 1 desc
+        assert results[2] == {"name": "Bob", "score": 2}    # Bob
+
     @pytest.mark.asyncio
     async def test_delete_virtual_node_operation(self):
         """Test delete virtual node operation."""

flowquery-py/tests/parsing/test_parser.py

@@ -1172,3 +1172,66 @@ def test_optional_without_match_throws_error(self):
         parser = Parser()
         with pytest.raises(Exception, match="Expected MATCH after OPTIONAL"):
             parser.parse("OPTIONAL RETURN 1")
+
+    # ORDER BY expression tests
+
+    def test_order_by_simple_identifier(self):
+        """Test ORDER BY with a simple identifier parses correctly."""
+        parser = Parser()
+        ast = parser.parse("unwind [1, 2] as x return x order by x")
+        assert ast is not None
+
+    def test_order_by_property_access(self):
+        """Test ORDER BY with property access parses correctly."""
+        parser = Parser()
+        ast = parser.parse(
+            "unwind [{name: 'Bob'}, {name: 'Alice'}] as person "
+            "return person.name as name order by person.name asc"
+        )
+        assert ast is not None
+
+    def test_order_by_function_call(self):
+        """Test ORDER BY with function call parses correctly."""
+        parser = Parser()
+        ast = parser.parse(
+            "unwind ['HELLO', 'WORLD'] as word "
+            "return word order by toLower(word) asc"
+        )
+        assert ast is not None
+
+    def test_order_by_nested_function_calls(self):
+        """Test ORDER BY with nested function calls parses correctly."""
+        parser = Parser()
+        ast = parser.parse(
+            "unwind ['Alice', 'Bob'] as name "
+            "return name order by string_distance(toLower(name), toLower('alice')) asc"
+        )
+        assert ast is not None
+
+    def test_order_by_arithmetic_expression(self):
+        """Test ORDER BY with arithmetic expression parses correctly."""
+        parser = Parser()
+        ast = parser.parse(
+            "unwind [{a: 3, b: 1}, {a: 1, b: 5}] as item "
+            "return item.a as a, item.b as b order by item.a + item.b desc"
+        )
+        assert ast is not None
+
+    def test_order_by_multiple_expression_fields(self):
+        """Test ORDER BY with multiple expression fields parses correctly."""
+        parser = Parser()
+        ast = parser.parse(
+            "unwind [{a: 1, b: 2}] as item "
+            "return item.a as a, item.b as b "
+            "order by toLower(item.a) asc, item.b desc"
+        )
+        assert ast is not None
+
+    def test_order_by_expression_with_limit(self):
+        """Test ORDER BY with expression and LIMIT parses correctly."""
+        parser = Parser()
+        ast = parser.parse(
+            "unwind ['c', 'a', 'b'] as x "
+            "return x order by toLower(x) asc limit 2"
+        )
+        assert ast is not None