pyrox.dev
1{
2 buildPgrxExtension,
3 cargo-pgrx_0_14_1,
4 clang,
5 fetchFromGitHub,
6 lib,
7 nix-update-script,
8 postgresql,
9 postgresqlTestExtension,
10 replaceVars,
11 rust-jemalloc-sys,
12 stdenv,
13}:
14let
15 # Follow upstream and use rust-jemalloc-sys on linux aarch64 and x86_64
16 # Additionally, disable init exec TLS, since it causes issues with postgres.
17 # https://github.com/tensorchord/VectorChord/blob/0.4.2/Cargo.toml#L43-L44
18 useSystemJemalloc =
19 stdenv.hostPlatform.isLinux && (stdenv.hostPlatform.isAarch64 || stdenv.hostPlatform.isx86_64);
20 rust-jemalloc-sys' = (
21 rust-jemalloc-sys.override (old: {
22 jemalloc = old.jemalloc.override { disableInitExecTls = true; };
23 })
24 );
25in
26buildPgrxExtension (finalAttrs: {
27 inherit postgresql;
28 cargo-pgrx = cargo-pgrx_0_14_1;
29
30 pname = "vectorchord";
31 version = "0.4.2";
32
33 src = fetchFromGitHub {
34 owner = "tensorchord";
35 repo = "vectorchord";
36 tag = finalAttrs.version;
37 hash = "sha256-EdMuSNcWwCBsAY0e3d0WVug1KBWYWldvKStF6cf/uRs=";
38 };
39
40 patches = [
41 # Tell the `simd` crate to use the flags from the rust bindgen hook
42 (replaceVars ./0001-read-clang-flags-from-environment.diff {
43 clang = lib.getExe clang;
44 })
45 # Add feature flags needed for features not yet stabilised in rustc stable
46 ./0002-add-feature-flags.diff
47 ];
48
49 buildInputs = lib.optionals (useSystemJemalloc) [
50 rust-jemalloc-sys'
51 ];
52
53 cargoHash = "sha256-8NwfsJn5dnvog3fexzLmO3v7/3+L7xtv+PHWfCCWoHY=";
54
55 # Include upgrade scripts in the final package
56 # https://github.com/tensorchord/VectorChord/blob/0.4.2/crates/make/src/main.rs#L224
57 postInstall = ''
58 cp sql/upgrade/* $out/share/postgresql/extension/
59 '';
60
61 env = {
62 # Bypass rust nightly features not being available on rust stable
63 RUSTC_BOOTSTRAP = 1;
64 };
65
66 # This crate does not have the "pg_test" feature
67 usePgTestCheckFeature = false;
68
69 passthru = {
70 updateScript = nix-update-script { };
71
72 tests.extension = postgresqlTestExtension {
73 inherit (finalAttrs) finalPackage;
74 withPackages = [ "pgvector" ]; # vectorchord depends on pgvector at runtime
75 postgresqlExtraSettings = ''
76 shared_preload_libraries = 'vchord'
77 '';
78
79 sql = ''
80 CREATE EXTENSION vchord CASCADE;
81
82 CREATE TABLE items (id bigint PRIMARY KEY, embedding vector(3));
83 INSERT INTO items (id, embedding) VALUES
84 (1, '[1,2,4]'),
85 (2, '[1,2,5]'),
86 (3, '[0,0,3]'),
87 (4, '[0,0,2]'),
88 (5, '[0,0,1]');
89
90 CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops) WITH (options = $$
91 residual_quantization = true
92 [build.internal]
93 lists = [4096]
94 spherical_centroids = false
95 $$);
96
97 SET vchordrq.probes = 1;
98 '';
99
100 asserts = [
101 {
102 query = "SELECT extversion FROM pg_extension WHERE extname = 'vchord'";
103 expected = "'${finalAttrs.version}'";
104 description = "Expected installed version to match the derivation's version";
105 }
106 {
107 query = "SELECT id FROM items WHERE embedding <-> '[1,2,3]' = 1";
108 expected = "1";
109 description = "Expected vector of row with ID=1 to have an euclidean distance from [1,2,3] of 1.";
110 }
111 {
112 query = "SELECT id FROM items WHERE embedding <-> '[1,2,3]' = 2";
113 expected = "2";
114 description = "Expected vector of row with ID=2 to have an euclidean distance from [1,2,3] of 2.";
115 }
116 {
117 query = "SELECT id FROM items ORDER BY embedding <-> '[2,3,7]' LIMIT 1";
118 expected = "2";
119 description = "Expected vector of row with ID=2 to be the closest to [2,3,7].";
120 }
121 ];
122 };
123 };
124
125 meta = {
126 # PostgreSQL 18 is not yet supported
127 # Will be supported in the next release (likely 0.5.0), as it's already supported in the main branch
128 # Check after next package update.
129 broken = lib.warnIf (
130 finalAttrs.version != "0.4.2"
131 ) "Is postgresql18Packages.vectorchord still broken?" (lib.versionAtLeast postgresql.version "18");
132 changelog = "https://github.com/tensorchord/VectorChord/releases/tag/${finalAttrs.version}";
133 description = "Scalable, fast, and disk-friendly vector search in Postgres, the successor of pgvecto.rs";
134 homepage = "https://github.com/tensorchord/VectorChord";
135 license = lib.licenses.agpl3Only; # dual licensed with Elastic License v2 (ELv2)
136 maintainers = with lib.maintainers; [
137 diogotcorreia
138 ];
139 platforms = postgresql.meta.platforms;
140 };
141})