postgresqlPackages.pgvectorscale: init at 0.7.0

Closes #345457 Can be installed with { services.postgresql.extensions = ps: [ ps.pgvector ps.pgvectorscale ]; }
2025-03-27 12:49:59 +01:00
parent b9f0161238
commit 8aedff852a
2 changed files with 2901 additions and 0 deletions
@@ -0,0 +1,89 @@
+{
+  buildPgrxExtension,
+  postgresql,
+  fetchFromGitHub,
+  lib,
+  postgresqlTestExtension,
+}:
+
+let
+  finalPackage = buildPgrxExtension rec {
+    pname = "pgvectorscale";
+    version = "0.7.0";
+
+    src = fetchFromGitHub {
+      owner = "timescale";
+      repo = "pgvectorscale";
+      tag = version;
+      hash = "sha256-dy481k2SvyYXwwcsyLZSl3XlhSk9C5+4LfEfciB1DK4=";
+    };
+
+    doCheck = false;
+
+    useFetchCargoVendor = true;
+    cargoHash = "sha256-CeRyDn9VhxfjWFJ1/Z/XvOUQOSnDoHHZAqgfYTeKU0o=";
+    cargoPatches = [
+      ./add-Cargo.lock.patch
+    ];
+
+    cargoPgrxFlags = [
+      "-p"
+      "vectorscale"
+    ];
+
+    inherit postgresql;
+
+    passthru.tests.extension = postgresqlTestExtension {
+      inherit finalPackage;
+      withPackages = [ "pgvector" ];
+      sql =
+        let
+          genCheck =
+            id: compare: expected:
+            let
+              vecStr = "[${lib.concatMapStringsSep "," toString compare}]";
+            in
+            ''
+              ASSERT (
+                SELECT id
+                FROM document_embedding
+                WHERE ${toString expected} = (embedding <-> '${vecStr}')
+              ) = ${toString id},
+              'Expected vector of row with ID=${toString id} to have a euclidean distance from ${vecStr} of ${toString expected}';
+            '';
+        in
+        ''
+          CREATE EXTENSION vectorscale CASCADE;
+          CREATE TABLE IF NOT EXISTS document_embedding  (
+              id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
+              embedding VECTOR(3)
+          );
+
+          INSERT INTO document_embedding (id, embedding) VALUES
+          (1, '[1,2,4]'),
+          (2, '[1,2,5]');
+
+          CREATE INDEX document_embedding_idx ON document_embedding
+          USING diskann (embedding vector_cosine_ops);
+
+          DO $$
+            BEGIN
+            ${genCheck 1 [ 1 2 3 ] 1}
+            ${genCheck 2 [ 1 2 3 ] 2}
+            END;
+          $$
+          LANGUAGE PLPGSQL;
+        '';
+    };
+
+    meta = {
+      homepage = "https://github.com/timescale/pgvectorscale";
+      maintainers = lib.teams.flyingcircus.members;
+      description = "Complement to pgvector for high performance, cost efficient vector search on large workloads";
+      license = lib.licenses.postgresql;
+      platforms = postgresql.meta.platforms;
+      changelog = "https://github.com/timescale/pgvectorscale/releases/tag/${version}";
+    };
+  };
+in
+finalPackage