Merge pull request #35 from microsoft/bug_fixes_and_improvements

Bug fixes and improvements

Author: niclasko

flowquery-py/src/parsing/expressions/operator.py

@@ -163,23 +163,23 @@ def value(self) -> int:
 
 class Is(Operator):
     def __init__(self) -> None:
-        super().__init__(-1, True)
+        super().__init__(0, True)
 
     def value(self) -> int:
         return 1 if self.lhs.value() == self.rhs.value() else 0
 
 
 class IsNot(Operator):
     def __init__(self) -> None:
-        super().__init__(-1, True)
+        super().__init__(0, True)
 
     def value(self) -> int:
         return 1 if self.lhs.value() != self.rhs.value() else 0
 
 
 class In(Operator):
     def __init__(self) -> None:
-        super().__init__(-1, True)
+        super().__init__(0, True)
 
     def value(self) -> int:
         lst = self.rhs.value()
@@ -190,7 +190,7 @@ def value(self) -> int:
 
 class NotIn(Operator):
     def __init__(self) -> None:
-        super().__init__(-1, True)
+        super().__init__(0, True)
 
     def value(self) -> int:
         lst = self.rhs.value()

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

@@ -10,6 +10,7 @@
 from .load import Load
 from .match import Match
 from .operation import Operation
+from .order_by import OrderBy, SortField
 from .projection import Projection
 from .return_op import Return
 from .union import Union
@@ -36,4 +37,6 @@
     "CreateRelationship",
     "Union",
     "UnionAll",
+    "OrderBy",
+    "SortField",
 ]

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

@@ -19,4 +19,7 @@ async def run(self) -> None:
     def results(self) -> List[Dict[str, Any]]:
         if self._where is not None:
             self._group_by.where = self._where
-        return list(self._group_by.generate_results())
+        results = list(self._group_by.generate_results())
+        if self._order_by is not None:
+            results = self._order_by.sort(results)
+        return results

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

@@ -15,6 +15,10 @@ def __init__(self, limit: int):
     def is_limit_reached(self) -> bool:
         return self._count >= self._limit
 
+    @property
+    def limit_value(self) -> int:
+        return self._limit
+
     def increment(self) -> None:
         self._count += 1
 

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

@@ -0,0 +1,72 @@
+"""Represents an ORDER BY operation that sorts results."""
+
+from typing import Any, Dict, List
+
+from .operation import Operation
+
+
+class SortField:
+    """A single sort specification: field name and direction."""
+
+    def __init__(self, field: str, direction: str = "asc"):
+        self.field = field
+        self.direction = direction
+
+
+class OrderBy(Operation):
+    """Represents an ORDER BY operation that sorts results.
+
+    Can be attached to a RETURN operation (sorting its results),
+    or used as a standalone accumulating operation after a non-aggregate WITH.
+
+    Example:
+        RETURN x ORDER BY x DESC
+    """
+
+    def __init__(self, fields: List[SortField]):
+        super().__init__()
+        self._fields = fields
+        self._results: List[Dict[str, 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 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)
+                cmp = 0
+                if a_val is None and b_val is None:
+                    cmp = 0
+                elif a_val is None:
+                    cmp = -1
+                elif b_val is None:
+                    cmp = 1
+                elif a_val < b_val:
+                    cmp = -1
+                elif a_val > b_val:
+                    cmp = 1
+                if cmp != 0:
+                    return -cmp if sf.direction == "desc" else cmp
+            return 0
+
+        return sorted(records, key=functools.cmp_to_key(compare))
+
+    async def run(self) -> None:
+        """When used as a standalone operation, passes through to next."""
+        if self.next:
+            await self.next.run()
+
+    async def initialize(self) -> None:
+        self._results = []
+        if self.next:
+            await self.next.initialize()
+
+    @property
+    def results(self) -> List[Dict[str, Any]]:
+        return self._results

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

@@ -5,6 +5,7 @@
 
 from ..ast_node import ASTNode
 from .limit import Limit
+from .order_by import OrderBy
 from .projection import Projection
 
 if TYPE_CHECKING:
@@ -26,6 +27,7 @@ def __init__(self, expressions: List[ASTNode]) -> None:
         self._where: Optional['Where'] = None
         self._results: List[Dict[str, Any]] = []
         self._limit: Optional[Limit] = None
+        self._order_by: Optional[OrderBy] = None
 
     @property
     def where(self) -> Any:
@@ -45,10 +47,20 @@ def limit(self) -> Optional[Limit]:
     def limit(self, limit: Limit) -> None:
         self._limit = limit
 
+    @property
+    def order_by(self) -> Optional[OrderBy]:
+        return self._order_by
+
+    @order_by.setter
+    def order_by(self, order_by: OrderBy) -> None:
+        self._order_by = order_by
+
     async def run(self) -> None:
         if not self.where:
             return
-        if self._limit is not None and self._limit.is_limit_reached:
+        # When ORDER BY is present, skip limit during accumulation;
+        # limit will be applied after sorting in results property
+        if self._order_by is None and self._limit is not None and self._limit.is_limit_reached:
             return
         record: Dict[str, Any] = {}
         for expression, alias in self.expressions():
@@ -57,12 +69,17 @@ async def run(self) -> None:
             value = copy.deepcopy(raw) if isinstance(raw, (dict, list)) else raw
             record[alias] = value
         self._results.append(record)
-        if self._limit is not None:
+        if self._order_by is None and self._limit is not None:
             self._limit.increment()
 
     async def initialize(self) -> None:
         self._results = []
 
     @property
     def results(self) -> List[Dict[str, Any]]:
-        return self._results
+        result = self._results
+        if self._order_by is not None:
+            result = self._order_by.sort(result)
+        if self._order_by is not None and self._limit is not None:
+            result = result[:self._limit.limit_value]
+        return result

flowquery-py/src/parsing/parser.py

@@ -61,6 +61,7 @@
 from .operations.load import Load
 from .operations.match import Match
 from .operations.operation import Operation
+from .operations.order_by import OrderBy, SortField
 from .operations.return_op import Return
 from .operations.union import Union
 from .operations.union_all import UnionAll
@@ -146,6 +147,14 @@ def _parse_tokenized(self, is_sub_query: bool = False) -> ASTNode:
                     operation.add_sibling(where)
                     operation = where
 
+            order_by = self._parse_order_by()
+            if order_by is not None:
+                if isinstance(operation, Return):
+                    operation.order_by = order_by
+                else:
+                    operation.add_sibling(order_by)
+                    operation = order_by
+
             limit = self._parse_limit()
             if limit is not None:
                 if isinstance(operation, Return):
@@ -694,6 +703,41 @@ def _parse_limit(self) -> Optional[Limit]:
         self.set_next_token()
         return limit
 
+    def _parse_order_by(self) -> Optional[OrderBy]:
+        self._skip_whitespace_and_comments()
+        if not self.token.is_order():
+            return None
+        self._expect_previous_token_to_be_whitespace_or_comment()
+        self.set_next_token()
+        self._expect_and_skip_whitespace_and_comments()
+        if not self.token.is_by():
+            raise ValueError("Expected BY after ORDER")
+        self.set_next_token()
+        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()
+            self._skip_whitespace_and_comments()
+            direction = "asc"
+            if self.token.is_asc():
+                direction = "asc"
+                self.set_next_token()
+                self._skip_whitespace_and_comments()
+            elif self.token.is_desc():
+                direction = "desc"
+                self.set_next_token()
+                self._skip_whitespace_and_comments()
+            fields.append(SortField(field, direction))
+            if self.token.is_comma():
+                self.set_next_token()
+                self._skip_whitespace_and_comments()
+            else:
+                break
+        return OrderBy(fields)
+
     def _parse_expressions(
         self, alias_option: AliasOption = AliasOption.NOT_ALLOWED
     ) -> Iterator[Expression]:

flowquery-py/src/tokenization/token.py

@@ -630,6 +630,34 @@ def ALL() -> Token:
     def is_all(self) -> bool:
         return self._type == TokenType.KEYWORD and self._value == Keyword.ALL.value
 
+    @staticmethod
+    def ORDER() -> Token:
+        return Token(TokenType.KEYWORD, Keyword.ORDER.value)
+
+    def is_order(self) -> bool:
+        return self._type == TokenType.KEYWORD and self._value == Keyword.ORDER.value
+
+    @staticmethod
+    def BY() -> Token:
+        return Token(TokenType.KEYWORD, Keyword.BY.value)
+
+    def is_by(self) -> bool:
+        return self._type == TokenType.KEYWORD and self._value == Keyword.BY.value
+
+    @staticmethod
+    def ASC() -> Token:
+        return Token(TokenType.KEYWORD, Keyword.ASC.value)
+
+    def is_asc(self) -> bool:
+        return self._type == TokenType.KEYWORD and self._value == Keyword.ASC.value
+
+    @staticmethod
+    def DESC() -> Token:
+        return Token(TokenType.KEYWORD, Keyword.DESC.value)
+
+    def is_desc(self) -> bool:
+        return self._type == TokenType.KEYWORD and self._value == Keyword.DESC.value
+
     # End of file token
 
     @staticmethod