Aggregate partial query results when clustering, closes #184

src/python/txtai/api/cluster.py

@@ -8,6 +8,8 @@ import zlib
 
 import aiohttp
 
+from ..database.sql import Aggregate
+
 
 class Cluster:
     """
@@ -31,6 +33,9 @@ class Cluster:
         if "shards" in self.config:
             self.shards = self.config["shards"]
 
+        # Query aggregator
+        self.aggregate = Aggregate()
+
     def search(self, query, limit=None):
         """
         Finds documents in the embeddings cluster most similar to the input query. Returns
@@ -55,9 +60,8 @@ class Cluster:
         for result in self.execute("get", action):
             results.extend(result)
 
-        # Sort results - there will be no score when there is no associated similarity query
-        if results and "score" in results[0]:
-            results = sorted(results, key=lambda x: x["score"], reverse=True)
+        # Combine aggregate functions and sort
+        results = self.aggregate(query, results)
 
         # Limit results
         return results[: (limit if limit else 10)]
@@ -86,13 +90,13 @@ class Cluster:
 
         # Combine results per query
         results = []
-        for x in range(len(queries)):
+        for x, query in enumerate(queries):
             result = []
             for section in batch:
                 result.extend(section[x])
 
-            # Sort and limit results
-            results.append(sorted(result, key=lambda x: x["score"], reverse=True)[: (limit if limit else 10)])
+            # Aggregate, sort and limit results
+            results.append(self.aggregate(query, result)[: (limit if limit else 10)])
 
         return results
 
@@ -209,7 +213,7 @@ class Cluster:
             json results if any
         """
 
-        async with aiohttp.ClientSession() as session:
+        async with aiohttp.ClientSession(raise_for_status=True) as session:
             tasks = []
 
             for x, url in enumerate(urls):

src/python/txtai/database/sql/__init__.py

@@ -2,5 +2,6 @@
 SQL imports
 """
 
+from .aggregate import Aggregate
 from .base import SQL, SQLException
 from .expression import Expression

src/python/txtai/database/sql/aggregate.py

@@ -0,0 +1,187 @@
+"""
+Aggregate module
+"""
+
+import itertools
+import operator
+
+from .base import SQL
+
+
+class Aggregate(SQL):
+    """
+    Aggregates partial results from queries. Partial results come from queries when working with sharded indexes.
+    """
+
+    def __init__(self, database=None):
+        # Always return token lists as this method requires them
+        super().__init__(database, True)
+
+    def __call__(self, query, results):
+        """
+        Analyzes query results, combines aggregate function results and applies ordering.
+
+        Args:
+            query: input query
+            results: query results
+
+        Returns:
+            aggregated query results
+        """
+
+        # Parse query
+        query = super().__call__(query)
+
+        # Detect if this is a SQL query
+        if "select" in query:
+            # Get list of unique and aggregate columns. If no aggregate columns or order by found, skip
+            columns = query["select"]
+            aggcolumns = self.aggcolumns(columns)
+            if aggcolumns or query["orderby"]:
+                # Query formatted columns
+                qcolumns = list(results[0].keys())
+
+                # Combine internal column names with initial result values for consistent processing
+                results = [dict(zip(columns, result.values())) for result in results]
+
+                # Merge aggregate columns
+                if aggcolumns:
+                    results = self.aggregate(query, results, columns, aggcolumns)
+
+                # Sort results
+                results = self.orderby(query, results) if query["orderby"] else self.defaultsort(results)
+
+                # Combine query formatted columns with result values for final results
+                return [dict(zip(qcolumns, result.values())) for result in results]
+
+        # Otherwise, run default sort
+        return self.defaultsort(results)
+
+    def aggcolumns(self, columns):
+        """
+        Filters columns for columns that have an aggregate function call.
+
+        Args:
+            columns: list of columns
+
+        Returns:
+            list of aggregate columns
+        """
+
+        aggregates = {}
+        for column in columns:
+            column = column.lower()
+            if column.startswith(("count(", "sum(", "total(")):
+                aggregates[column] = sum
+            elif column.startswith("max("):
+                aggregates[column] = max
+            elif column.startswith("min("):
+                aggregates[column] = min
+            elif column.startswith("avg("):
+                aggregates[column] = lambda x: sum(x) / len(x)
+
+        return aggregates
+
+    def aggregate(self, query, results, columns, aggcolumns):
+        """
+        Merges aggregate columns in results.
+
+        Args:
+            query: input query
+            results: query results
+            columns: list of select columns
+            aggcolumns: list of aggregate columns
+
+        Returns:
+            results with aggregates merged
+        """
+
+        # Group data, if necessary
+        if query["groupby"]:
+            results = self.groupby(query, results, columns)
+        else:
+            results = [results]
+
+        # Compute column values
+        rows = []
+        for result in results:
+            # Calculate/copy column values
+            row = {}
+            for column in columns:
+                if column in aggcolumns:
+                    # Calculate aggregate value
+                    function = aggcolumns[column]
+                    row[column] = function([r[column] for r in result])
+                else:
+                    # Non aggregate column value repeat, use first value
+                    row[column] = result[0][column]
+
+            # Add row using original query columns
+            rows.append(row)
+
+        return rows
+
+    def groupby(self, query, results, columns):
+        """
+        Groups results using query group by clause.
+
+        Args:
+            query: input query
+            results: query results
+            columns: list of select colum
+
+        Returns:
+            results grouped using group by clause
+        """
+
+        groupby = [column for column in columns if column.lower() in query["groupby"]]
+        if groupby:
+            results = sorted(results, key=operator.itemgetter(*groupby))
+            return [list(value) for _, value in itertools.groupby(results, operator.itemgetter(*groupby))]
+
+        return [results]
+
+    def orderby(self, query, results):
+        """
+        Applies an order by clause to results.
+
+        Args:
+            query: input query
+            results: query results
+
+        Returns:
+            results ordered using order by clause
+        """
+
+        # Sort in reverse order
+        for clause in query["orderby"][::-1]:
+            # Order by columns must be selected
+            reverse = False
+            if clause.lower().endswith(" asc"):
+                clause = clause.rsplit(" ")[0]
+            elif clause.lower().endswith(" desc"):
+                clause = clause.rsplit(" ")[0]
+                reverse = True
+
+            # Order by columns must be in select clause
+            if clause in query["select"]:
+                results = sorted(results, key=operator.itemgetter(clause), reverse=reverse)
+
+        return results
+
+    def defaultsort(self, results):
+        """
+        Default sorting algorithm for results. Sorts by score descending, if available.
+
+        Args:
+            results: query results
+
+        Returns:
+            results ordered by score descending
+        """
+
+        # Sort standard query using score column, if present
+        if results and "score" in results[0]:
+            return sorted(results, key=lambda x: x["score"], reverse=True)
+
+        return results

src/python/txtai/database/sql/base.py

@@ -16,16 +16,17 @@ class SQL:
     # List of clauses to parse
     CLAUSES = ["select", "from", "where", "group", "having", "order", "limit"]
 
-    def __init__(self, database):
+    def __init__(self, database=None, tolist=False):
         """
         Creates a new SQL query parser.
 
         Args:
-            database: database instance that provides resolver callback
+            database: database instance that provides resolver callback, if any
+            tolist: outputs expression lists if True, expression text otherwise, defaults to False
         """
 
         # Expression parser
-        self.expression = Expression(database.resolve)
+        self.expression = Expression(database.resolve if database else self.defaultresolve, tolist)
 
     def __call__(self, query):
         """
@@ -67,6 +68,22 @@ class SQL:
         # Return clauses, default to full query if this is not a SQL query
         return clauses if clauses else {"similar": [[query]]}
 
+    # pylint: disable=W0613
+    def defaultresolve(self, name, alias=False, compound=False):
+        """
+        Default resolve function. Performs no processing, only returns name.
+
+        Args:
+            name: query column name
+            alias: True if an alias clause should be added, defaults to False
+            compound: True if this column is part of a compound expression, defaults to False
+
+        Returns:
+            name if alias is False, otherwise empty string
+        """
+
+        return name if not alias else ""
+
     def issql(self, query):
         """
         Detects if this is a SQL query.

src/python/txtai/database/sql/expression.py

@@ -20,15 +20,17 @@ class Expression:
     # Default list of sort order operators
     SORT_ORDER = ["asc", "desc"]
 
-    def __init__(self, resolver):
+    def __init__(self, resolver, tolist):
         """
         Creates a new expression parser.
 
         Args:
             resolver: function to call to resolve query column names with database column names
+            tolist: outputs expression lists if True, text if False
         """
 
         self.resolver = resolver
+        self.tolist = tolist
 
     def __call__(self, tokens, alias=False, similar=None):
         """
@@ -48,7 +50,7 @@ class Expression:
         tokens = self.resolve(list(tokens), alias, similar)
 
         # Re-write query and return
-        return self.build(tokens)
+        return self.buildlist(tokens) if self.tolist is True else self.buildtext(tokens)
 
     def resolve(self, tokens, alias, similar):
         """
@@ -88,7 +90,7 @@ class Expression:
         # Remove replaced tokens
         return [token for token in tokens if token]
 
-    def build(self, tokens):
+    def buildtext(self, tokens):
         """
         Builds a new expression from tokens. This method applies a set of rules to generate whitespace between tokens.
 
@@ -108,6 +110,42 @@ class Expression:
         # Remove any leading/trailing whitespace and return
         return text.strip()
 
+    def buildlist(self, tokens):
+        """
+        Builds a new expression from tokens. This method returns a list of expression components. These components can be joined together
+        on commas to form a text expression.
+
+        Args:
+            tokens: input expression
+
+        Returns:
+            expression list
+        """
+
+        parts, current, parens, brackets = [], [], 0, 0
+
+        for token in tokens:
+            if token == "," and not parens and not brackets:
+                parts.append("".join(current))
+                current = []
+            else:
+                if token == "(":
+                    parens += 1
+                elif token == ")":
+                    parens -= 1
+                elif token == "[":
+                    brackets += 1
+                elif token == "]":
+                    brackets -= 1
+                elif self.issortorder(token):
+                    token = f" {token}"
+                current.append(token)
+
+        if current:
+            parts.append("".join(current))
+
+        return parts
+
     def issimilar(self, tokens, x, similar):
         """
         Checks if tokens[x] is a similar() function.

test/python/testapi/testcluster.py

@@ -6,6 +6,7 @@ import json
 import os
 import tempfile
 import unittest
+import urllib.parse
 
 from http.server import HTTPServer, BaseHTTPRequestHandler
 from threading import Thread
@@ -35,7 +36,15 @@ class RequestHandler(BaseHTTPRequestHandler):
         """
 
         if self.path == "/count":
-            response = 5
+            response = 26
+        elif self.path.startswith("/search?query=select"):
+            if "group" in self.path:
+                response = [{"count(*)": 12, "text": "This is a test"}, {"count(*)": 14, "text": "And another test"}]
+            else:
+                if self.server.server_port == 8002:
+                    response = [{"count(*)": 12, "min(indexid)": 0, "max(indexid)": 11, "avg(indexid)": 6.3}]
+                else:
+                    response = [{"count(*)": 16, "min(indexid)": 2, "max(indexid)": 14, "avg(indexid)": 6.7}]
         elif self.path.startswith("/search"):
             response = [{"id": 4, "score": 0.40}]
         else:
@@ -136,7 +145,7 @@ class TestCluster(unittest.TestCase):
         Test cluster count
         """
 
-        self.assertEqual(self.client.get("count").json(), 10)
+        self.assertEqual(self.client.get("count").json(), 52)
 
     def testDelete(self):
         """
@@ -157,7 +166,8 @@ class TestCluster(unittest.TestCase):
         Test cluster search
         """
 
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
         self.assertEqual(uid, 4)
 
     def testSearchBatch(self):
@@ -170,6 +180,31 @@ class TestCluster(unittest.TestCase):
         uids = [result[0]["id"] for result in results]
         self.assertEqual(uids, [4, 1])
 
+    def testSQL(self):
+        """
+        Test cluster SQL statement
+        """
+
+        query = urllib.parse.quote("select count(*), min(indexid), max(indexid), avg(indexid) from txtai where text=[This is a test]")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(), [{"count(*)": 28, "min(indexid)": 0, "max(indexid)": 14, "avg(indexid)": 6.5}]
+        )
+
+        query = urllib.parse.quote("select count(*), text from txtai group by text order by count(*) desc")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(),
+            [{"count(*)": 28, "text": "And another test"}, {"count(*)": 24, "text": "This is a test"}],
+        )
+
+        query = urllib.parse.quote("select count(*), text from txtai group by text order by count(*) asc")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(),
+            [{"count(*)": 24, "text": "This is a test"}, {"count(*)": 28, "text": "And another test"}],
+        )
+
+        query = urllib.parse.quote("select count(*) from txtai group by id order by count(*)")
+        self.assertEqual(self.client.get(f"search?query={query}").json(), [{"count(*)": 52}])
+
     def testUpsert(self):
         """
         Test cluster upsert
@@ -180,6 +215,7 @@ class TestCluster(unittest.TestCase):
         self.client.get("upsert")
 
         # Search for best match
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
+        query = "feel good story"
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
 
         self.assertEqual(uid, 4)

test/python/testapi/testembeddings.py

@@ -5,6 +5,7 @@ Embeddings API module tests
 import os
 import tempfile
 import unittest
+import urllib.parse
 
 from unittest.mock import patch
 
@@ -103,7 +104,8 @@ class TestEmbeddings(unittest.TestCase):
         self.assertEqual(ids, [4])
 
         # Search for best match
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
 
         self.assertEqual(self.client.get("count").json(), 5)
         self.assertEqual(uid, 5)
@@ -173,7 +175,8 @@ class TestEmbeddings(unittest.TestCase):
         Test search via API
         """
 
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
         self.assertEqual(uid, 4)
 
     def testSearchBatch(self):
@@ -232,7 +235,8 @@ class TestEmbeddings(unittest.TestCase):
         self.client.get("upsert")
 
         # Search for best match
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
 
         self.assertEqual(uid, 0)
 
@@ -249,7 +253,8 @@ class TestEmbeddings(unittest.TestCase):
         self.client = TestEmbeddings.start(False)
 
         # Test search
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
         self.assertEqual(uid, 4)
 
         # Test similarity