mseep-txtai 9.1.1__py3-none-any.whl

txtai/database/schema/statement.py

@@ -0,0 +1,98 @@
+"""
+Statement module
+"""
+
+
+class Statement:
+    """
+    Standard database schema SQL statements.
+    """
+
+    # Temporary table for working with id batches
+    CREATE_BATCH = """
+        CREATE TEMP TABLE IF NOT EXISTS batch (
+            indexid INTEGER,
+            id TEXT,
+            batch INTEGER
+        )
+    """
+
+    DELETE_BATCH = "DELETE FROM batch"
+    INSERT_BATCH_INDEXID = "INSERT INTO batch (indexid, batch) VALUES (?, ?)"
+    INSERT_BATCH_ID = "INSERT INTO batch (id, batch) VALUES (?, ?)"
+
+    # Temporary table for joining similarity scores
+    CREATE_SCORES = """
+        CREATE TEMP TABLE IF NOT EXISTS scores (
+            indexid INTEGER PRIMARY KEY,
+            score REAL
+        )
+    """
+
+    DELETE_SCORES = "DELETE FROM scores"
+    INSERT_SCORE = "INSERT INTO scores VALUES (?, ?)"
+
+    # Documents - stores full content
+    CREATE_DOCUMENTS = """
+        CREATE TABLE IF NOT EXISTS documents (
+            id TEXT PRIMARY KEY,
+            data JSON,
+            tags TEXT,
+            entry DATETIME
+        )
+    """
+
+    INSERT_DOCUMENT = "INSERT OR REPLACE INTO documents VALUES (?, ?, ?, ?)"
+    DELETE_DOCUMENTS = "DELETE FROM documents WHERE id IN (SELECT id FROM batch)"
+
+    # Objects - stores binary content
+    CREATE_OBJECTS = """
+        CREATE TABLE IF NOT EXISTS objects (
+            id TEXT PRIMARY KEY,
+            object BLOB,
+            tags TEXT,
+            entry DATETIME
+        )
+    """
+
+    INSERT_OBJECT = "INSERT OR REPLACE INTO objects VALUES (?, ?, ?, ?)"
+    DELETE_OBJECTS = "DELETE FROM objects WHERE id IN (SELECT id FROM batch)"
+
+    # Sections - stores section text
+    CREATE_SECTIONS = """
+        CREATE TABLE IF NOT EXISTS %s (
+            indexid INTEGER PRIMARY KEY,
+            id TEXT,
+            text TEXT,
+            tags TEXT,
+            entry DATETIME
+        )
+    """
+
+    CREATE_SECTIONS_INDEX = "CREATE INDEX section_id ON sections(id)"
+    INSERT_SECTION = "INSERT INTO sections VALUES (?, ?, ?, ?, ?)"
+    DELETE_SECTIONS = "DELETE FROM sections WHERE id IN (SELECT id FROM batch)"
+    COPY_SECTIONS = (
+        "INSERT INTO %s SELECT (select count(*) - 1 from sections s1 where s.indexid >= s1.indexid) indexid, "
+        + "s.id, %s AS text, s.tags, s.entry FROM sections s LEFT JOIN documents d ON s.id = d.id ORDER BY indexid"
+    )
+    STREAM_SECTIONS = (
+        "SELECT s.id, s.text, data, object, s.tags FROM %s s "
+        + "LEFT JOIN documents d ON s.id = d.id "
+        + "LEFT JOIN objects o ON s.id = o.id ORDER BY indexid"
+    )
+    DROP_SECTIONS = "DROP TABLE sections"
+    RENAME_SECTIONS = "ALTER TABLE %s RENAME TO sections"
+
+    # Queries
+    SELECT_IDS = "SELECT indexid, id FROM sections WHERE id in (SELECT id FROM batch)"
+    COUNT_IDS = "SELECT count(indexid) FROM sections"
+
+    # Partial sql clauses
+    TABLE_CLAUSE = (
+        "SELECT %s FROM sections s "
+        + "LEFT JOIN documents d ON s.id = d.id "
+        + "LEFT JOIN objects o ON s.id = o.id "
+        + "LEFT JOIN scores sc ON s.indexid = sc.indexid"
+    )
+    IDS_CLAUSE = "s.indexid in (SELECT indexid from batch WHERE batch=%s)"

txtai/database/sql/__init__.py

@@ -0,0 +1,8 @@
+"""
+SQL imports
+"""
+
+from .aggregate import Aggregate
+from .base import SQL, SQLError
+from .expression import Expression
+from .token import Token

txtai/database/sql/aggregate.py

@@ -0,0 +1,178 @@
+"""
+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)
+
+        # Check 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 = list(results[0].keys())
+            aggcolumns = self.aggcolumns(columns)
+            if aggcolumns or query["orderby"]:
+                # Merge aggregate columns
+                if aggcolumns:
+                    results = self.aggregate(query, results, columns, aggcolumns)
+
+                # Sort results and return
+                return self.orderby(query, results) if query["orderby"] else self.defaultsort(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 columns
+
+        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

txtai/database/sql/base.py

@@ -0,0 +1,189 @@
+"""
+SQL module
+"""
+
+from io import StringIO
+from shlex import shlex
+
+from .expression import Expression
+
+
+class SQL:
+    """
+    Translates txtai SQL statements into database native queries.
+    """
+
+    # List of clauses to parse
+    CLAUSES = ["select", "from", "where", "group", "having", "order", "limit", "offset"]
+
+    def __init__(self, database=None, tolist=False):
+        """
+        Creates a new SQL query parser.
+
+        Args:
+            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 if database else self.defaultresolve, tolist)
+
+    def __call__(self, query):
+        """
+        Parses an input SQL query and normalizes column names in the query clauses. This method will also embed
+        similarity search placeholders into the query.
+
+        Args:
+            query: input query
+
+        Returns:
+            {clause name: clause text}
+        """
+
+        clauses = None
+        if self.issql(query):
+            # Ignore multiple statements
+            query = query.split(";")[0]
+
+            # Tokenize query
+            tokens, positions = self.tokenize(query)
+
+            # Alias clauses and similar queries
+            aliases, similar = {}, []
+
+            # Parse SQL clauses
+            clauses = {
+                "select": self.parse(tokens, positions, "select", alias=True, aliases=aliases),
+                "where": self.parse(tokens, positions, "where", aliases=aliases, similar=similar),
+                "groupby": self.parse(tokens, positions, "group", offset=2, aliases=aliases),
+                "having": self.parse(tokens, positions, "having", aliases=aliases),
+                "orderby": self.parse(tokens, positions, "order", offset=2, aliases=aliases),
+                "limit": self.parse(tokens, positions, "limit", aliases=aliases),
+                "offset": self.parse(tokens, positions, "offset", aliases=aliases),
+            }
+
+            # Add parsed similar queries, if any
+            if similar:
+                clauses["similar"] = similar
+
+        # 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=None):
+        """
+        Default resolve function. Performs no processing, only returns name.
+
+        Args:
+            name: query column name
+            alias: alias name, defaults to None
+
+        Returns:
+            name
+        """
+
+        return name
+
+    def issql(self, query):
+        """
+        Detects if this is a SQL query.
+
+        Args:
+            query: input query
+
+        Returns:
+            True if this is a valid SQL query, False otherwise
+        """
+
+        if isinstance(query, str):
+            # Reduce query to a lower-cased single line stripped of leading/trailing whitespace
+            query = query.lower().strip(";").replace("\n", " ").replace("\t", " ").strip()
+
+            # Detect if this is a valid txtai SQL statement
+            return query.startswith("select ") and (" from txtai " in query or query.endswith(" from txtai"))
+
+        return False
+
+    def snippet(self, text):
+        """
+        Parses a partial SQL snippet.
+
+        Args:
+            text: SQL snippet
+
+        Returns:
+            parsed snippet
+        """
+
+        tokens, _ = self.tokenize(text)
+        return self.expression(tokens)
+
+    def tokenize(self, query):
+        """
+        Tokenizes SQL query into tokens.
+
+        Args:
+            query: input query
+
+        Returns:
+            (tokenized query, token positions)
+        """
+
+        # Build a simple SQL lexer
+        #   - Punctuation chars are parsed as standalone tokens which helps identify operators
+        #   - Add additional wordchars to prevent splitting on those values
+        #   - Disable comments
+        tokens = shlex(StringIO(query), punctuation_chars="=!<>+-*/%|")
+        tokens.wordchars += ":@#"
+        tokens.commenters = ""
+        tokens = list(tokens)
+
+        # Identify sql clause token positions
+        positions = {}
+
+        # Get position of clause keywords. For multi-term clauses, validate next token matches as well
+        for x, token in enumerate(tokens):
+            t = token.lower()
+            if t not in positions and t in SQL.CLAUSES and (t not in ["group", "order"] or (x + 1 < len(tokens) and tokens[x + 1].lower() == "by")):
+                positions[t] = x
+
+        return (tokens, positions)
+
+    def parse(self, tokens, positions, name, offset=1, alias=False, aliases=None, similar=None):
+        """
+        Runs query column name to database column name mappings for clauses. This method will also
+        parse SIMILAR() function calls, extract parameters for those calls and leave a placeholder
+        to be filled in with similarity results.
+
+        Args:
+            tokens: query tokens
+            positions: token positions - used to locate the start of sql clauses
+            name: current query clause name
+            offset: how many tokens are in the clause name
+            alias: True if terms in the clause should be aliased (i.e. column as alias)
+            aliases: dict of generated aliases, if present these tokens should NOT be resolved
+            similar: list where parsed similar clauses should be stored
+
+        Returns:
+            formatted clause
+        """
+
+        clause = None
+        if name in positions:
+            # Find the next clause token
+            end = [positions.get(x, len(tokens)) for x in SQL.CLAUSES[SQL.CLAUSES.index(name) + 1 :]]
+            end = min(end) if end else len(tokens)
+
+            # Start after current clause token and end before next clause or end of string
+            clause = tokens[positions[name] + offset : end]
+
+            # Parse and resolve parameters
+            clause = self.expression(clause, alias, aliases, similar)
+
+        return clause
+
+
+class SQLError(Exception):
+    """
+    Raised for errors generated by user SQL queries
+    """