nixpkgs/pkgs/servers/sql/postgresql/ext/vectorchord/package.nix

{
  buildPgrxExtension,
  cargo-pgrx_0_17_0,
  fetchFromGitHub,
  lib,
  nix-update-script,
  postgresql,
  postgresqlTestExtension,
  rustc,
  stdenv,
}:
buildPgrxExtension (finalAttrs: {
  inherit postgresql;
  cargo-pgrx = cargo-pgrx_0_17_0;

  pname = "vectorchord";
  version = "1.1.1";

  src = fetchFromGitHub {
    owner = "supervc-stack";
    repo = "VectorChord";
    tag = finalAttrs.version;
    hash = "sha256-QL9XGSQFOcrpww03Y5F0JuDbpo0v8oidUqucLxggkqE=";
  };

  cargoHash = "sha256-IXOCzKJArNOcb/2TcJbLz1XdCquUpyF/cLHYU5vmlko=";

  patches =
    lib.optional (lib.versionOlder rustc.llvm.version "22.0.0" && stdenv.hostPlatform.isx86_64)
      [
        # Due to a bug in LLVM 21, build fails on x86_64 with:
        # `rustc-LLVM ERROR: Cannot select: intrinsic %llvm.x86.avx512.vpdpbusd.512`.
        # This has been fixed in Rust's build of LLVM and LLVM 22 with https://github.com/rust-lang/llvm-project/commit/94e2c19f86a699d7a19ff0f4130b696699189c8d,
        # but this is not reflected in nixpkgs' packaging of rustc.
        # For this reason, we temporarily disable avx512vnni support until this is fixed in nixpkgs.
        # This might cause a performance penalty, but should not affect correctness.
        # See https://github.com/NixOS/nixpkgs/pull/537113#issuecomment-4846239887
        ./0001-disable-avx512vnni.patch
      ];

  # Include upgrade scripts in the final package
  # https://github.com/supervc-stack/VectorChord/blob/0.5.0/crates/make/src/main.rs#L366
  postInstall = ''
    cp sql/upgrade/* $out/share/postgresql/extension/
  '';

  # This crate does not have the "pg_test" feature
  usePgTestCheckFeature = false;

  passthru = {
    updateScript = nix-update-script { };

    tests.extension = postgresqlTestExtension {
      inherit (finalAttrs) finalPackage;
      withPackages = [ "pgvector" ]; # vectorchord depends on pgvector at runtime
      postgresqlExtraSettings = ''
        shared_preload_libraries = 'vchord'
      '';

      sql = ''
        CREATE EXTENSION vchord CASCADE;

        CREATE TABLE items (id bigint PRIMARY KEY, embedding vector(3));
        INSERT INTO items (id, embedding) VALUES
        (1, '[1,2,4]'),
        (2, '[1,2,5]'),
        (3, '[0,0,3]'),
        (4, '[0,0,2]'),
        (5, '[0,0,1]');

        CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops) WITH (options = $$
        residual_quantization = true
        [build.internal]
        lists = [4096]
        spherical_centroids = false
        $$);

        SET vchordrq.probes = 1;
      '';

      asserts = [
        {
          query = "SELECT extversion FROM pg_extension WHERE extname = 'vchord'";
          expected = "'${finalAttrs.version}'";
          description = "Expected installed version to match the derivation's version";
        }
        {
          query = "SELECT id FROM items WHERE embedding <-> '[1,2,3]' = 1";
          expected = "1";
          description = "Expected vector of row with ID=1 to have an euclidean distance from [1,2,3] of 1.";
        }
        {
          query = "SELECT id FROM items WHERE embedding <-> '[1,2,3]' = 2";
          expected = "2";
          description = "Expected vector of row with ID=2 to have an euclidean distance from [1,2,3] of 2.";
        }
        {
          query = "SELECT id FROM items ORDER BY embedding <-> '[2,3,7]' LIMIT 1";
          expected = "2";
          description = "Expected vector of row with ID=2 to be the closest to [2,3,7].";
        }
      ];
    };
  };

  meta = {
    # PostgreSQL 19 is not yet supported
    # See https://github.com/supervc-stack/VectorChord/issues/464
    # Check after next package update.
    broken = lib.warnIf (
      finalAttrs.version != "1.1.1"
    ) "Is postgresql19Packages.vectorchord still broken?" (lib.versionAtLeast postgresql.version "19");
    changelog = "https://github.com/supervc-stack/VectorChord/releases/tag/${finalAttrs.version}";
    description = "Scalable, fast, and disk-friendly vector search in Postgres, the successor of pgvecto.rs";
    homepage = "https://github.com/supervc-stack/VectorChord";
    license = lib.licenses.OR [
      lib.licenses.agpl3Only
      lib.licenses.elastic20
    ];
    maintainers = with lib.maintainers; [
      diogotcorreia
    ];
    platforms = postgresql.meta.platforms;
  };
})