feat: enhance batch processing with strict result count checks and null handling

Author: srimon12

docs/reference.md

@@ -171,7 +171,7 @@ Both sync and async connections support:
 
 ## Project Structure
 
-```
+```text
 qql/
 ├── pyproject.toml          # Package config; installs the `qql` CLI command
 ├── src/

src/qql/async_connection.py

@@ -190,22 +190,36 @@ async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseExc
         if not self._queries:
             return
         results = await self.connection.run_queries_batch(self._queries)
-        for proxy, res in zip(self._proxies, results):
+        for proxy, res in zip(self._proxies, results, strict=False):
             proxy._resolve(res)
+        if len(results) != len(self._proxies):
+            error = RuntimeError(
+                "Batch result count mismatch: "
+                f"expected {len(self._proxies)}, got {len(results)}"
+            )
+            for proxy in self._proxies[len(results):]:
+                proxy._reject(error)
+            raise error
 
 
 class AsyncOperationProxy:
     """Proxy handle that resolves to an ExecutionResult after QQLAsyncBatch exits."""
 
     def __init__(self) -> None:
         self._result: ExecutionResult | None = None
+        self._exception: RuntimeError | None = None
 
     def _resolve(self, result: ExecutionResult) -> None:
         self._result = result
 
+    def _reject(self, exception: RuntimeError) -> None:
+        self._exception = exception
+
     @property
     def result(self) -> ExecutionResult:
         """The resolved ExecutionResult."""
+        if self._exception is not None:
+            raise self._exception
         if self._result is None:
             raise RuntimeError("AsyncBatch has not been executed yet.")
         return self._result

src/qql/async_executor.py

@@ -138,12 +138,16 @@ async def _ensure_collection(
         vector_size: int,
         topology: CollectionTopology,
         explicit_vector: str | None,
-    ) -> None:
+    ) -> CollectionTopology:
         if topology.exists:
             info = await self._client.get_collection(name)
             vectors = info.config.params.vectors  # type: ignore[union-attr]
+            sparse_vectors = info.config.params.sparse_vectors or {}
+            current_topology = CollectionTopology(
+                **collection_topology_kwargs(vectors, sparse_vectors)
+            )
             if isinstance(vectors, dict):
-                vector_name = topology.dense_using(explicit_vector)
+                vector_name = current_topology.dense_using(explicit_vector)
                 if vector_name is None:
                     raise QQLRuntimeError("Collection has no dense vector")
                 vector_config = vectors[vector_name]
@@ -164,12 +168,12 @@ async def _ensure_collection(
                     )
             else:
                 raise QQLRuntimeError("Collection has no dense vector")
+            return current_topology
         else:
             async with self._creation_lock:
                 current_topology = await self._resolve_topology(name)
                 if current_topology.exists:
-                    await self._ensure_collection(name, vector_size, current_topology, explicit_vector)
-                    return
+                    return await self._ensure_collection(name, vector_size, current_topology, explicit_vector)
 
                 await self._create_collection_and_wait(
                     collection_name=name,
@@ -179,6 +183,7 @@ async def _ensure_collection(
                         )
                     },
                 )
+                return await self._resolve_topology(name)
 
     async def _create_collection_and_wait(self, **kwargs: Any) -> None:
         collection_name = kwargs["collection_name"]
@@ -290,7 +295,7 @@ async def _execute_insert(self, node: InsertStmt) -> ExecutionResult:
         embedder = Embedder(model_name)
         vector = embedder.embed(node.values["text"])
 
-        await self._ensure_collection(
+        topology = await self._ensure_collection(
             node.collection, len(vector), topology, node.dense_vector
         )
         point_vector = build_dense_point_vector(
@@ -351,22 +356,6 @@ async def _execute_insert_bulk(self, node: InsertBulkStmt) -> ExecutionResult:
             sparse_objs = [sparse_embedder.embed(vals["text"]) for vals in node.values_list]
 
             first_dense_vector = dense_vectors[0] if dense_vectors else None
-            points: list[PointStruct] = []
-            for idx, vals in enumerate(node.values_list):
-                point_id, payload = extract_point_id_and_payload(vals)
-                dense_vector = dense_vectors[idx]
-                sparse_obj = sparse_objs[idx]
-                sparse_vector = SparseVector(
-                    indices=sparse_obj["indices"], values=sparse_obj["values"]
-                )
-                points.append(
-                    PointStruct(
-                        id=point_id,
-                        vector={dense_name: dense_vector, sparse_name: sparse_vector},
-                        payload=payload,
-                    )
-                )
-
             if not topology.exists:
                 assert first_dense_vector is not None
                 async with self._creation_lock:
@@ -385,6 +374,22 @@ async def _execute_insert_bulk(self, node: InsertBulkStmt) -> ExecutionResult:
                         dense_name = current_topology.dense_using(node.dense_vector) or dense_name
                         sparse_name = current_topology.sparse_using(node.sparse_vector)
 
+            points: list[PointStruct] = []
+            for idx, vals in enumerate(node.values_list):
+                point_id, payload = extract_point_id_and_payload(vals)
+                dense_vector = dense_vectors[idx]
+                sparse_obj = sparse_objs[idx]
+                sparse_vector = SparseVector(
+                    indices=sparse_obj["indices"], values=sparse_obj["values"]
+                )
+                points.append(
+                    PointStruct(
+                        id=point_id,
+                        vector={dense_name: dense_vector, sparse_name: sparse_vector},
+                        payload=payload,
+                    )
+                )
+
             try:
                 await self._client.upsert(
                     collection_name=node.collection,
@@ -406,6 +411,10 @@ async def _execute_insert_bulk(self, node: InsertBulkStmt) -> ExecutionResult:
         vectors = [embedder.embed(vals["text"]) for vals in node.values_list]
 
         first_vector = vectors[0] if vectors else None
+        assert first_vector is not None
+        topology = await self._ensure_collection(
+            node.collection, len(first_vector), topology, node.dense_vector
+        )
         points = []
         for idx, vals in enumerate(node.values_list):
             vector = vectors[idx]
@@ -420,11 +429,6 @@ async def _execute_insert_bulk(self, node: InsertBulkStmt) -> ExecutionResult:
                 PointStruct(id=point_id, vector=point_vector, payload=payload)
             )
 
-        assert first_vector is not None
-        await self._ensure_collection(
-            node.collection, len(first_vector), topology, node.dense_vector
-        )
-
         try:
             await self._client.upsert(
                 collection_name=node.collection,

src/qql/connection.py

@@ -198,22 +198,36 @@ def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException
         if not self._queries:
             return
         results = self.connection.run_queries_batch(self._queries)
-        for proxy, res in zip(self._proxies, results):
+        for proxy, res in zip(self._proxies, results, strict=False):
             proxy._resolve(res)
+        if len(results) != len(self._proxies):
+            error = RuntimeError(
+                "Batch result count mismatch: "
+                f"expected {len(self._proxies)}, got {len(results)}"
+            )
+            for proxy in self._proxies[len(results):]:
+                proxy._reject(error)
+            raise error
 
 
 class OperationProxy:
     """Proxy handle that resolves to an ExecutionResult after QQLBatch exits."""
 
     def __init__(self) -> None:
         self._result: ExecutionResult | None = None
+        self._exception: RuntimeError | None = None
 
     def _resolve(self, result: ExecutionResult) -> None:
         self._result = result
 
+    def _reject(self, exception: RuntimeError) -> None:
+        self._exception = exception
+
     @property
     def result(self) -> ExecutionResult:
         """The resolved ExecutionResult."""
+        if self._exception is not None:
+            raise self._exception
         if self._result is None:
             raise RuntimeError("Batch has not been executed yet.")
         return self._result

src/qql/parser.py

@@ -72,6 +72,8 @@ def __init__(self, tokens: list[Token]) -> None:
 
     def parse(self) -> ASTNode:
         node = self._parse_single_statement()
+        while self._peek().kind == TokenKind.SEMICOLON:
+            self._advance()
         self._expect(TokenKind.EOF)
         return node
 
@@ -1078,12 +1080,15 @@ def _parse_field_path(self) -> str:
             f"Expected a field name, got '{tok.value}'", tok.pos
         )
 
-    def _parse_literal(self) -> str | int | float | bool:
-        """STRING | INTEGER | FLOAT | boolean"""
+    def _parse_literal(self) -> str | int | float | bool | None:
+        """STRING | INTEGER | FLOAT | boolean | NULL"""
         tok = self._peek()
         if tok.kind == TokenKind.STRING:
             self._advance()
             return tok.value
+        if tok.kind == TokenKind.NULL:
+            self._advance()
+            return None
         if tok.kind == TokenKind.INTEGER:
             self._advance()
             return int(tok.value)
@@ -1099,7 +1104,7 @@ def _parse_literal(self) -> str | int | float | bool:
                 self._advance()
                 return False
         raise QQLSyntaxError(
-            f"Expected a literal value (string, integer, float, or boolean), got '{tok.value}'",
+            f"Expected a literal value (string, integer, float, boolean, or null), got '{tok.value}'",
             tok.pos,
         )
 
@@ -1116,10 +1121,10 @@ def _parse_number(self) -> int | float:
             f"Expected a number, got '{tok.value}'", tok.pos
         )
 
-    def _parse_literal_list(self) -> list[str | int | float | bool]:
+    def _parse_literal_list(self) -> list[str | int | float | bool | None]:
         """'(' literal { ',' literal } [','] ')'  — used by IN / NOT IN."""
         self._expect(TokenKind.LPAREN)
-        items: list[str | int | float | bool] = []
+        items: list[str | int | float | bool | None] = []
         if self._peek().kind == TokenKind.RPAREN:
             self._advance()
             return items

src/qql/utils.py

@@ -78,19 +78,65 @@ class SearchGroupByOptions:
 
 
 def render_parameterized_query(template: str, params: dict[str, Any]) -> str:
-    query_str = template
-    for k in sorted(params.keys(), key=len, reverse=True):
-        val = params[k]
-        placeholder = f":{k}"
-        if isinstance(val, str):
-            escaped_val = val.replace("'", "\\'")
-            repr_val = f"'{escaped_val}'"
-        elif isinstance(val, bool):
-            repr_val = "true" if val else "false"
-        else:
-            repr_val = str(val)
-        query_str = query_str.replace(placeholder, repr_val)
-    return query_str
+    rendered = []
+    in_string = False
+    quote_char = ""
+    i = 0
+    while i < len(template):
+        ch = template[i]
+        if in_string:
+            rendered.append(ch)
+            if ch == "\\" and i + 1 < len(template):
+                rendered.append(template[i + 1])
+                i += 2
+                continue
+            if ch == quote_char:
+                in_string = False
+                quote_char = ""
+            i += 1
+            continue
+
+        if ch in ("'", '"'):
+            in_string = True
+            quote_char = ch
+            rendered.append(ch)
+            i += 1
+            continue
+
+        if ch == ":":
+            name_start = i + 1
+            name_end = name_start
+            while name_end < len(template) and (
+                template[name_end].isalnum() or template[name_end] == "_"
+            ):
+                name_end += 1
+            name = template[name_start:name_end]
+            if name in params:
+                rendered.append(_qql_literal(params[name]))
+                i = name_end
+                continue
+
+        rendered.append(ch)
+        i += 1
+
+    return "".join(rendered)
+
+
+def _qql_literal(value: Any) -> str:
+    if value is None:
+        return "null"
+    if isinstance(value, str):
+        escaped = (
+            value.replace("\\", "\\\\")
+            .replace("'", "\\'")
+            .replace("\n", "\\n")
+            .replace("\t", "\\t")
+            .replace("\r", "\\r")
+        )
+        return f"'{escaped}'"
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    return str(value)
 
 
 def collection_topology_kwargs(vectors: Any, sparse_vectors: Any) -> dict[str, Any]:
@@ -158,6 +204,14 @@ def group_batch_statements(statements: tuple[ASTNode, ...]) -> list[BatchGroup]:
     current_group: list[ASTNode] = []
 
     for stmt in statements:
+        if isinstance(stmt, SearchStmt) and stmt.group_by is not None:
+            _append_batch_group(groups, current_type, current_collection, current_group)
+            groups.append(BatchGroup("other", stmt.collection, [stmt]))
+            current_type = None
+            current_collection = None
+            current_group = []
+            continue
+
         if isinstance(stmt, (SearchStmt, RecommendStmt)):
             coll = stmt.collection
             if current_type == "query" and current_collection == coll:
@@ -247,6 +301,12 @@ def build_qdrant_filter(expr: FilterExpr) -> Any:
     if isinstance(expr, NotExpr):
         return Filter(must_not=[build_qdrant_filter(expr.operand)])
     if isinstance(expr, CompareExpr):
+        if expr.value is None:
+            null_condition = IsNullCondition(is_null=PayloadField(key=expr.field))
+            if expr.op == "=":
+                return null_condition
+            if expr.op == "!=":
+                return Filter(must_not=[null_condition])
         if expr.op == "=":
             return FieldCondition(key=expr.field, match=MatchValue(value=expr.value))
         if expr.op == "!=":