Add support for custom SQL functions with DuckDB, closes #1047

Author: davidmezzetti

setup.py

@@ -55,7 +55,7 @@
 
 extras["console"] = ["rich>=12.0.1"]
 
-extras["database"] = ["duckdb>=0.7.1", "pillow>=7.1.2", "sqlalchemy>=2.0.20"]
+extras["database"] = ["duckdb>=0.8.0", "pillow>=7.1.2", "sqlalchemy>=2.0.20"]
 
 extras["graph"] = ["grand-cypher>=0.6.0", "grand-graph>=0.6.0", "networkx>=2.7.1", "sqlalchemy>=2.0.20"]
 

src/python/txtai/database/base.py

@@ -263,11 +263,11 @@ def registerfunctions(self, config):
         if inputs:
             functions = []
             for fn in inputs:
-                name, argcount = None, -1
+                name, argcount, deterministic = None, -1, None
 
                 # Optional function configuration
                 if isinstance(fn, dict):
-                    name, argcount, fn = fn.get("name"), fn.get("argcount", -1), fn["function"]
+                    name, argcount, fn, deterministic = (fn.get("name"), fn.get("argcount", -1), fn["function"], fn.get("deterministic"))
 
                 # Determine if this is a callable object or a function
                 if not isinstance(fn, types.FunctionType) and hasattr(fn, "__call__"):
@@ -277,7 +277,7 @@ def registerfunctions(self, config):
                     name = name if name else fn.__name__.lower()
 
                 # Store function details
-                functions.append((name, argcount, fn))
+                functions.append((name, argcount, fn, deterministic))
 
             # pylint: disable=W0201
             self.functions = functions

src/python/txtai/database/duckdb.py

@@ -5,6 +5,7 @@
 import os
 import re
 
+from typing import get_type_hints
 from tempfile import TemporaryDirectory
 
 # Conditional import
@@ -80,8 +81,15 @@ def rows(self):
             rows = self.cursor.fetchmany(batch)
 
     def addfunctions(self):
-        # DuckDB doesn't currently support scalar functions
-        return
+        if self.connection and self.functions:
+            for name, _, fn, deterministic in self.functions:
+                # Get function type hints
+                hints = get_type_hints(fn)
+
+                # Create database functions
+                self.connection.create_function(
+                    name, fn, return_type=hints.get("return", str), side_effects=not deterministic if deterministic is not None else False
+                )
 
     def copy(self, path):
         # Delete existing file, if necessary

src/python/txtai/database/sqlite.py

@@ -34,8 +34,9 @@ def addfunctions(self):
             # Enable callback tracebacks to show user-defined function errors
             sqlite3.enable_callback_tracebacks(True)
 
-            for name, argcount, fn in self.functions:
-                self.connection.create_function(name, argcount, fn)
+            # Create database functions
+            for name, argcount, fn, deterministic in self.functions:
+                self.connection.create_function(name, argcount, fn, deterministic=deterministic if deterministic is not None else False)
 
     def copy(self, path):
         # Delete existing file, if necessary

test/python/testdatabase/testduckdb.py

@@ -53,3 +53,32 @@ def testArchive(self):
         """
 
         super().testArchive()
+
+    def testFunction(self):
+        """
+        Test custom functions
+        """
+
+        embeddings = Embeddings(
+            {
+                "path": "sentence-transformers/nli-mpnet-base-v2",
+                "content": self.backend,
+                "functions": [{"name": "length", "function": "testdatabase.testduckdb.length"}],
+            }
+        )
+
+        # Create an index for the list of text
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        # Search for best match
+        result = embeddings.search("select length(text) length from txtai where id = 0", 1)[0]
+
+        self.assertEqual(result["length"], 39)
+
+
+def length(text):
+    """
+    Custom SQL function.
+    """
+
+    return len(text)