├── .github └── workflows │ └── build.yml ├── .gitignore ├── CHANGELOG.md ├── Examples ├── Cohere │ ├── Package.swift │ └── Sources │ │ └── main.swift ├── Hybrid │ ├── Package.swift │ └── Sources │ │ └── main.swift ├── OpenAI │ ├── Package.swift │ └── Sources │ │ └── main.swift └── Sparse │ ├── Package.swift │ └── Sources │ └── main.swift ├── LICENSE.txt ├── Package.swift ├── README.md ├── Sources ├── Pgvector │ ├── HalfVector.swift │ ├── SparseVector.swift │ └── Vector.swift ├── PgvectorClientKit │ ├── HalfVector.swift │ ├── SparseVector.swift │ └── Vector.swift └── PgvectorNIO │ ├── HalfVector.swift │ ├── PgvectorNIO.swift │ ├── SparseVector.swift │ └── Vector.swift └── Tests └── PgvectorTests ├── HalfVectorTests.swift ├── PostgresClientKitTests.swift ├── PostgresNIOTests.swift ├── SparseVectorTests.swift └── VectorTests.swift /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | on: [push, pull_request] 3 | jobs: 4 | build: 5 | runs-on: ubuntu-latest 6 | steps: 7 | - uses: actions/checkout@v4 8 | - uses: ankane/setup-postgres@v1 9 | with: 10 | database: pgvector_swift_test 11 | dev-files: true 12 | - run: | 13 | cd /tmp 14 | git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git 15 | cd pgvector 16 | make 17 | sudo make install 18 | - run: psql -d pgvector_swift_test -c "CREATE EXTENSION vector" 19 | - run: swift test 20 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | .build 3 | /Packages 4 | /*.xcodeproj 5 | xcuserdata/ 6 | DerivedData/ 7 | .swiftpm/config/registries.json 8 | .swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata 9 | .netrc 10 | Package.resolved 11 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## 0.1.0 (2025-03-05) 2 | 3 | - First release 4 | -------------------------------------------------------------------------------- /Examples/Cohere/Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 6.0 2 | 3 | import PackageDescription 4 | 5 | let package = Package( 6 | name: "Example", 7 | platforms: [ 8 | .macOS(.v13) 9 | ], 10 | dependencies: [ 11 | .package(url: "https://github.com/vapor/postgres-nio", from: "1.24.0") 12 | ], 13 | targets: [ 14 | .executableTarget( 15 | name: "Example", 16 | dependencies: [ 17 | .product(name: "PostgresNIO", package: "postgres-nio") 18 | ]) 19 | ] 20 | ) 21 | -------------------------------------------------------------------------------- /Examples/Cohere/Sources/main.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | import PostgresNIO 3 | 4 | guard let apiKey = ProcessInfo.processInfo.environment["CO_API_KEY"] else { 5 | print("Set CO_API_KEY") 6 | exit(1) 7 | } 8 | 9 | let config = PostgresClient.Configuration( 10 | host: "localhost", 11 | port: 5432, 12 | username: ProcessInfo.processInfo.environment["USER"]!, 13 | password: nil, 14 | database: "pgvector_example", 15 | tls: .disable 16 | ) 17 | 18 | let client = PostgresClient(configuration: config) 19 | 20 | struct ApiData: Encodable { 21 | var texts: [String] 22 | var model: String 23 | var inputType: String 24 | var embeddingTypes: [String] 25 | } 26 | 27 | struct EmbedResponse: Decodable { 28 | var embeddings: EmbeddingsObject 29 | } 30 | 31 | struct EmbeddingsObject: Decodable { 32 | var ubinary: [[UInt8]] 33 | } 34 | 35 | func embed(texts: [String], inputType: String, apiKey: String) async throws -> [String] { 36 | let url = URL(string: "https://api.cohere.com/v2/embed")! 37 | let data = ApiData( 38 | texts: texts, 39 | model: "embed-v4.0", 40 | inputType: inputType, 41 | embeddingTypes: ["ubinary"] 42 | ) 43 | 44 | var request = URLRequest(url: url) 45 | request.httpMethod = "POST" 46 | request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") 47 | request.setValue("application/json", forHTTPHeaderField: "Content-Type") 48 | 49 | let encoder = JSONEncoder() 50 | encoder.keyEncodingStrategy = .convertToSnakeCase 51 | request.httpBody = try encoder.encode(data) 52 | 53 | let (body, _) = try await URLSession.shared.data(for: request) 54 | let response = try JSONDecoder().decode(EmbedResponse.self, from: body) 55 | 56 | return response.embeddings.ubinary.map { 57 | $0.map { String(repeating: "0", count: 8 - String($0, radix: 2).count) + String($0, radix: 2) }.joined() 58 | } 59 | } 60 | 61 | try await withThrowingTaskGroup(of: Void.self) { taskGroup in 62 | taskGroup.addTask { 63 | await client.run() 64 | } 65 | 66 | try await client.query("CREATE EXTENSION IF NOT EXISTS vector") 67 | try await client.query("DROP TABLE IF EXISTS documents") 68 | try await client.query("CREATE TABLE documents (id serial PRIMARY KEY, content text, embedding bit(1536))") 69 | 70 | let input = [ 71 | "The dog is barking", 72 | "The cat is purring", 73 | "The bear is growling", 74 | ] 75 | let embeddings = try await embed(texts: input, inputType: "search_document", apiKey: apiKey) 76 | for (content, embedding) in zip(input, embeddings) { 77 | try await client.query("INSERT INTO documents (content, embedding) VALUES (\(content), \(embedding)::bit(1536))") 78 | } 79 | 80 | let query = "forest" 81 | let queryEmbedding = (try await embed(texts: [query], inputType: "search_query", apiKey: apiKey))[0] 82 | let rows = try await client.query("SELECT content FROM documents ORDER BY embedding <~> \(queryEmbedding)::bit(1536) LIMIT 5") 83 | for try await row in rows { 84 | print(row) 85 | } 86 | 87 | taskGroup.cancelAll() 88 | } 89 | -------------------------------------------------------------------------------- /Examples/Hybrid/Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 6.0 2 | 3 | import PackageDescription 4 | 5 | let package = Package( 6 | name: "Example", 7 | platforms: [ 8 | .macOS(.v13) 9 | ], 10 | dependencies: [ 11 | .package(url: "https://github.com/vapor/postgres-nio", from: "1.24.0"), 12 | .package(url: "https://github.com/pgvector/pgvector-swift", from: "0.1.0"), 13 | ], 14 | targets: [ 15 | .executableTarget( 16 | name: "Example", 17 | dependencies: [ 18 | .product(name: "PostgresNIO", package: "postgres-nio"), 19 | .product(name: "Pgvector", package: "pgvector-swift"), 20 | .product(name: "PgvectorNIO", package: "pgvector-swift"), 21 | ]) 22 | ] 23 | ) 24 | -------------------------------------------------------------------------------- /Examples/Hybrid/Sources/main.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | import Pgvector 3 | import PgvectorNIO 4 | import PostgresNIO 5 | 6 | let config = PostgresClient.Configuration( 7 | host: "localhost", 8 | port: 5432, 9 | username: ProcessInfo.processInfo.environment["USER"]!, 10 | password: nil, 11 | database: "pgvector_example", 12 | tls: .disable 13 | ) 14 | 15 | let client = PostgresClient(configuration: config) 16 | 17 | struct ApiData: Encodable { 18 | var model: String 19 | var input: [String] 20 | } 21 | 22 | struct ApiResponse: Decodable { 23 | var embeddings: [[Float]] 24 | } 25 | 26 | func embed(input: [String], taskType: String) async throws -> [[Float]] { 27 | // nomic-embed-text uses a task prefix 28 | // https://huggingface.co/nomic-ai/nomic-embed-text-v1.5 29 | let input = input.map { taskType + ": " + $0 } 30 | 31 | let url = URL(string: "http://localhost:11434/api/embed")! 32 | let data = ApiData( 33 | model: "nomic-embed-text", 34 | input: input 35 | ) 36 | 37 | var request = URLRequest(url: url) 38 | request.httpMethod = "POST" 39 | request.httpBody = try JSONEncoder().encode(data) 40 | let (body, _) = try await URLSession.shared.data(for: request) 41 | let response = try JSONDecoder().decode(ApiResponse.self, from: body) 42 | 43 | return response.embeddings 44 | } 45 | 46 | try await withThrowingTaskGroup(of: Void.self) { taskGroup in 47 | taskGroup.addTask { 48 | await client.run() 49 | } 50 | 51 | try await client.query("CREATE EXTENSION IF NOT EXISTS vector") 52 | try await PgvectorNIO.registerTypes(client) 53 | 54 | try await client.query("DROP TABLE IF EXISTS documents") 55 | try await client.query("CREATE TABLE documents (id serial PRIMARY KEY, content text, embedding vector(768))") 56 | 57 | let input = [ 58 | "The dog is barking", 59 | "The cat is purring", 60 | "The bear is growling", 61 | ] 62 | let embeddings = try await embed(input: input, taskType: "search_document") 63 | for (content, embedding) in zip(input, embeddings) { 64 | let embedding = Vector(embedding) 65 | try await client.query("INSERT INTO documents (content, embedding) VALUES (\(content), \(embedding))") 66 | } 67 | 68 | let query = "growling bear" 69 | let queryEmbedding = Vector((try await embed(input: [query], taskType: "search_query"))[0]) 70 | let k = 60 71 | let sql: PostgresQuery = """ 72 | WITH semantic_search AS ( 73 | SELECT id, RANK () OVER (ORDER BY embedding <=> \(queryEmbedding)) AS rank 74 | FROM documents 75 | ORDER BY embedding <=> \(queryEmbedding) 76 | LIMIT 20 77 | ), 78 | keyword_search AS ( 79 | SELECT id, RANK () OVER (ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC) 80 | FROM documents, plainto_tsquery('english', \(query)) query 81 | WHERE to_tsvector('english', content) @@ query 82 | ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC 83 | LIMIT 20 84 | ) 85 | SELECT 86 | COALESCE(semantic_search.id, keyword_search.id) AS id, 87 | COALESCE(1.0 / (\(k) + semantic_search.rank), 0.0) + 88 | COALESCE(1.0 / (\(k) + keyword_search.rank), 0.0) AS score 89 | FROM semantic_search 90 | FULL OUTER JOIN keyword_search ON semantic_search.id = keyword_search.id 91 | ORDER BY score DESC 92 | LIMIT 5 93 | """ 94 | let rows = try await client.query(sql) 95 | for try await row in rows { 96 | print(row) 97 | } 98 | 99 | taskGroup.cancelAll() 100 | } 101 | -------------------------------------------------------------------------------- /Examples/OpenAI/Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 6.0 2 | 3 | import PackageDescription 4 | 5 | let package = Package( 6 | name: "Example", 7 | platforms: [ 8 | .macOS(.v13) 9 | ], 10 | dependencies: [ 11 | .package(url: "https://github.com/vapor/postgres-nio", from: "1.24.0"), 12 | .package(url: "https://github.com/pgvector/pgvector-swift", from: "0.1.0"), 13 | ], 14 | targets: [ 15 | .executableTarget( 16 | name: "Example", 17 | dependencies: [ 18 | .product(name: "PostgresNIO", package: "postgres-nio"), 19 | .product(name: "Pgvector", package: "pgvector-swift"), 20 | .product(name: "PgvectorNIO", package: "pgvector-swift"), 21 | ]) 22 | ] 23 | ) 24 | -------------------------------------------------------------------------------- /Examples/OpenAI/Sources/main.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | import Pgvector 3 | import PgvectorNIO 4 | import PostgresNIO 5 | 6 | guard let apiKey = ProcessInfo.processInfo.environment["OPENAI_API_KEY"] else { 7 | print("Set OPENAI_API_KEY") 8 | exit(1) 9 | } 10 | 11 | let config = PostgresClient.Configuration( 12 | host: "localhost", 13 | port: 5432, 14 | username: ProcessInfo.processInfo.environment["USER"]!, 15 | password: nil, 16 | database: "pgvector_example", 17 | tls: .disable 18 | ) 19 | 20 | let client = PostgresClient(configuration: config) 21 | 22 | struct ApiData: Encodable { 23 | var model: String 24 | var input: [String] 25 | } 26 | 27 | struct ApiResponse: Decodable { 28 | var data: [ApiObject] 29 | } 30 | 31 | struct ApiObject: Decodable { 32 | var embedding: [Float] 33 | } 34 | 35 | func embed(input: [String], apiKey: String) async throws -> [[Float]] { 36 | let url = URL(string: "https://api.openai.com/v1/embeddings")! 37 | let data = ApiData( 38 | model: "text-embedding-3-small", 39 | input: input 40 | ) 41 | 42 | var request = URLRequest(url: url) 43 | request.httpMethod = "POST" 44 | request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") 45 | request.setValue("application/json", forHTTPHeaderField: "Content-Type") 46 | request.httpBody = try JSONEncoder().encode(data) 47 | let (body, _) = try await URLSession.shared.data(for: request) 48 | let response = try JSONDecoder().decode(ApiResponse.self, from: body) 49 | 50 | return response.data.map { $0.embedding } 51 | } 52 | 53 | try await withThrowingTaskGroup(of: Void.self) { taskGroup in 54 | taskGroup.addTask { 55 | await client.run() 56 | } 57 | 58 | try await client.query("CREATE EXTENSION IF NOT EXISTS vector") 59 | try await PgvectorNIO.registerTypes(client) 60 | 61 | try await client.query("DROP TABLE IF EXISTS documents") 62 | try await client.query("CREATE TABLE documents (id serial PRIMARY KEY, content text, embedding vector(1536))") 63 | 64 | let input = [ 65 | "The dog is barking", 66 | "The cat is purring", 67 | "The bear is growling", 68 | ] 69 | let embeddings = try await embed(input: input, apiKey: apiKey) 70 | for (content, embedding) in zip(input, embeddings) { 71 | let embedding = Vector(embedding) 72 | try await client.query("INSERT INTO documents (content, embedding) VALUES (\(content), \(embedding))") 73 | } 74 | 75 | let query = "forest" 76 | let queryEmbedding = Vector((try await embed(input: [query], apiKey: apiKey))[0]) 77 | let rows = try await client.query("SELECT content FROM documents ORDER BY embedding <=> \(queryEmbedding) LIMIT 5") 78 | for try await row in rows { 79 | print(row) 80 | } 81 | 82 | taskGroup.cancelAll() 83 | } 84 | -------------------------------------------------------------------------------- /Examples/Sparse/Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 6.0 2 | 3 | import PackageDescription 4 | 5 | let package = Package( 6 | name: "Example", 7 | platforms: [ 8 | .macOS(.v13) 9 | ], 10 | dependencies: [ 11 | .package(url: "https://github.com/vapor/postgres-nio", from: "1.24.0"), 12 | .package(url: "https://github.com/pgvector/pgvector-swift", from: "0.1.0"), 13 | ], 14 | targets: [ 15 | .executableTarget( 16 | name: "Example", 17 | dependencies: [ 18 | .product(name: "PostgresNIO", package: "postgres-nio"), 19 | .product(name: "Pgvector", package: "pgvector-swift"), 20 | .product(name: "PgvectorNIO", package: "pgvector-swift"), 21 | ]) 22 | ] 23 | ) 24 | -------------------------------------------------------------------------------- /Examples/Sparse/Sources/main.swift: -------------------------------------------------------------------------------- 1 | // good resources 2 | // https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/ 3 | // https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1 4 | // 5 | // run with 6 | // text-embeddings-router --model-id opensearch-project/opensearch-neural-sparse-encoding-v1 --pooling splade 7 | 8 | import Foundation 9 | import Pgvector 10 | import PgvectorNIO 11 | import PostgresNIO 12 | 13 | let config = PostgresClient.Configuration( 14 | host: "localhost", 15 | port: 5432, 16 | username: ProcessInfo.processInfo.environment["USER"]!, 17 | password: nil, 18 | database: "pgvector_example", 19 | tls: .disable 20 | ) 21 | 22 | let client = PostgresClient(configuration: config) 23 | 24 | struct ApiElement: Decodable { 25 | var index: Int 26 | var value: Float 27 | } 28 | 29 | func embed(input: [String]) async throws -> [[Int: Float]] { 30 | let url = URL(string: "http://localhost:3000/embed_sparse")! 31 | let data = [ 32 | "inputs": input 33 | ] 34 | 35 | var request = URLRequest(url: url) 36 | request.httpMethod = "POST" 37 | request.setValue("application/json", forHTTPHeaderField: "Content-Type") 38 | request.httpBody = try JSONEncoder().encode(data) 39 | let (body, _) = try await URLSession.shared.data(for: request) 40 | let response = try JSONDecoder().decode([[ApiElement]].self, from: body) 41 | 42 | return response.map { Dictionary(uniqueKeysWithValues: $0.map { ($0.index, $0.value) }) } 43 | } 44 | 45 | try await withThrowingTaskGroup(of: Void.self) { taskGroup in 46 | taskGroup.addTask { 47 | await client.run() 48 | } 49 | 50 | try await client.query("CREATE EXTENSION IF NOT EXISTS vector") 51 | try await PgvectorNIO.registerTypes(client) 52 | 53 | try await client.query("DROP TABLE IF EXISTS documents") 54 | try await client.query("CREATE TABLE documents (id serial PRIMARY KEY, content text, embedding sparsevec(30522))") 55 | 56 | let input = [ 57 | "The dog is barking", 58 | "The cat is purring", 59 | "The bear is growling", 60 | ] 61 | let embeddings = try await embed(input: input) 62 | for (content, embedding) in zip(input, embeddings) { 63 | let embedding = SparseVector(embedding, dim: 30522) 64 | try await client.query("INSERT INTO documents (content, embedding) VALUES (\(content), \(embedding))") 65 | } 66 | 67 | let query = "forest" 68 | let embedding = SparseVector((try await embed(input: [query]))[0], dim: 30522) 69 | let rows = try await client.query("SELECT content FROM documents ORDER BY embedding <#> \(embedding) LIMIT 5") 70 | for try await row in rows { 71 | print(row) 72 | } 73 | 74 | taskGroup.cancelAll() 75 | } 76 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2023-2025 Andrew Kane 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | THE SOFTWARE. 22 | -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 5.7 2 | 3 | import PackageDescription 4 | 5 | let package = Package( 6 | name: "Pgvector", 7 | platforms: [ 8 | .macOS(.v13) 9 | ], 10 | products: [ 11 | .library( 12 | name: "Pgvector", 13 | targets: ["Pgvector"]), 14 | .library( 15 | name: "PgvectorClientKit", 16 | targets: ["PgvectorClientKit"]), 17 | .library( 18 | name: "PgvectorNIO", 19 | targets: ["PgvectorNIO"]), 20 | ], 21 | dependencies: [ 22 | .package(url: "https://github.com/codewinsdotcom/PostgresClientKit", from: "1.0.0"), 23 | .package(url: "https://github.com/vapor/postgres-nio", from: "1.0.0"), 24 | ], 25 | targets: [ 26 | .target( 27 | name: "Pgvector", 28 | dependencies: []), 29 | .target( 30 | name: "PgvectorClientKit", 31 | dependencies: [ 32 | "Pgvector", 33 | "PostgresClientKit", 34 | ]), 35 | .target( 36 | name: "PgvectorNIO", 37 | dependencies: [ 38 | "Pgvector", 39 | .product(name: "PostgresNIO", package: "postgres-nio"), 40 | ]), 41 | .testTarget( 42 | name: "PgvectorTests", 43 | dependencies: [ 44 | "Pgvector", 45 | "PgvectorClientKit", 46 | "PgvectorNIO", 47 | "PostgresClientKit", 48 | .product(name: "PostgresNIO", package: "postgres-nio"), 49 | ]), 50 | ] 51 | ) 52 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pgvector-swift 2 | 3 | [pgvector](https://github.com/pgvector/pgvector) support for Swift 4 | 5 | Supports [PostgresNIO](https://github.com/vapor/postgres-nio) and [PostgresClientKit](https://github.com/codewinsdotcom/PostgresClientKit) 6 | 7 | [![Build Status](https://github.com/pgvector/pgvector-swift/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-swift/actions) 8 | 9 | ## Getting Started 10 | 11 | Follow the instructions for your database library: 12 | 13 | - [PostgresNIO](#postgresnio) 14 | - [PostgresClientKit](#postgresclientkit) 15 | 16 | Or check out an example: 17 | 18 | - [Embeddings](Examples/OpenAI/Sources/main.swift) with OpenAI 19 | - [Binary embeddings](Examples/Cohere/Sources/main.swift) with Cohere 20 | - [Hybrid search](Examples/Hybrid/Sources/main.swift) with Ollama 21 | - [Sparse search](Examples/Sparse/Sources/main.swift) with Text Embeddings Inference 22 | 23 | ## PostgresNIO 24 | 25 | Add to your application’s `Package.swift` 26 | 27 | ```diff 28 | dependencies: [ 29 | + .package(url: "https://github.com/pgvector/pgvector-swift", from: "0.1.0") 30 | ], 31 | targets: [ 32 | .executableTarget(name: "App", dependencies: [ 33 | + .product(name: "Pgvector", package: "pgvector-swift"), 34 | + .product(name: "PgvectorNIO", package: "pgvector-swift") 35 | ]) 36 | ] 37 | ``` 38 | 39 | Import the packages 40 | 41 | ```swift 42 | import Pgvector 43 | import PgvectorNIO 44 | ``` 45 | 46 | Enable the extension 47 | 48 | ```swift 49 | try await client.query("CREATE EXTENSION IF NOT EXISTS vector") 50 | ``` 51 | 52 | Register the types 53 | 54 | ```swift 55 | try await PgvectorNIO.registerTypes(client) 56 | ``` 57 | 58 | Create a table 59 | 60 | ```swift 61 | try await client.query("CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))") 62 | ``` 63 | 64 | Insert vectors 65 | 66 | ```swift 67 | let embedding1 = Vector([1, 1, 1]) 68 | let embedding2 = Vector([2, 2, 2]) 69 | let embedding3 = Vector([1, 1, 2]) 70 | try await client.query("INSERT INTO items (embedding) VALUES (\(embedding1)), (\(embedding2)), (\(embedding3))") 71 | ``` 72 | 73 | Get the nearest neighbors 74 | 75 | ```swift 76 | let embedding = Vector([1, 1, 1]) 77 | let rows = try await client.query("SELECT id, embedding::text FROM items ORDER BY embedding <-> \(embedding) LIMIT 5") 78 | for try await row in rows { 79 | print(row) 80 | } 81 | ``` 82 | 83 | Add an approximate index 84 | 85 | ```swift 86 | try await client.query("CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)") 87 | // or 88 | try await client.query("CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)") 89 | ``` 90 | 91 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 92 | 93 | See a [full example](Tests/PgvectorTests/PostgresNIOTests.swift) 94 | 95 | ## PostgresClientKit 96 | 97 | Add to your application’s `Package.swift` 98 | 99 | ```diff 100 | dependencies: [ 101 | + .package(url: "https://github.com/pgvector/pgvector-swift", from: "0.1.0") 102 | ], 103 | targets: [ 104 | .executableTarget(name: "App", dependencies: [ 105 | + .product(name: "Pgvector", package: "pgvector-swift"), 106 | + .product(name: "PgvectorClientKit", package: "pgvector-swift") 107 | ]) 108 | ] 109 | ``` 110 | 111 | Import the packages 112 | 113 | ```swift 114 | import Pgvector 115 | import PgvectorClientKit 116 | ``` 117 | 118 | Enable the extension 119 | 120 | ```swift 121 | let text = "CREATE EXTENSION IF NOT EXISTS vector" 122 | let statement = try connection.prepareStatement(text: text) 123 | try statement.execute() 124 | ``` 125 | 126 | Create a table 127 | 128 | ```swift 129 | let text = "CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))" 130 | let statement = try connection.prepareStatement(text: text) 131 | try statement.execute() 132 | ``` 133 | 134 | Insert vectors 135 | 136 | ```swift 137 | let text = "INSERT INTO items (embedding) VALUES ($1), ($2), ($3)" 138 | let statement = try connection.prepareStatement(text: text) 139 | try statement.execute(parameterValues: [Vector([1, 1, 1]), Vector([2, 2, 2]), Vector([1, 1, 2])]) 140 | ``` 141 | 142 | Get the nearest neighbors 143 | 144 | ```swift 145 | let text = "SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5" 146 | let statement = try connection.prepareStatement(text: text) 147 | let cursor = try statement.execute(parameterValues: [Vector([1, 1, 1])]) 148 | 149 | for row in cursor { 150 | let columns = try row.get().columns 151 | let id = try columns[0].int() 152 | let embedding = try columns[1].vector() 153 | print(id, embedding) 154 | } 155 | ``` 156 | 157 | Add an approximate index 158 | 159 | ```swift 160 | let text = "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)" 161 | // or 162 | let text = "CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)" 163 | let statement = try connection.prepareStatement(text: text) 164 | try statement.execute() 165 | ``` 166 | 167 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 168 | 169 | See a [full example](Tests/PgvectorTests/PostgresClientKitTests.swift) 170 | 171 | ## Reference 172 | 173 | ### Vectors 174 | 175 | Create a vector from an array 176 | 177 | ```swift 178 | let vec = Vector([1, 2, 3]) 179 | ``` 180 | 181 | ### Half Vectors 182 | 183 | Create a half vector from an array 184 | 185 | ```swift 186 | let vec = HalfVector([1, 2, 3]) 187 | ``` 188 | 189 | ### Sparse Vectors 190 | 191 | Create a sparse vector from an array 192 | 193 | ```swift 194 | let vec = SparseVector([1, 0, 2, 0, 3, 0]) 195 | ``` 196 | 197 | Or a dictionary of non-zero elements 198 | 199 | ```swift 200 | let vec = SparseVector([0: 1, 2: 2, 4: 3], dim: 6)! 201 | ``` 202 | 203 | Note: Indices start at 0 204 | 205 | Get the number of dimensions 206 | 207 | ```swift 208 | let dim = vec.dim 209 | ``` 210 | 211 | Get the indices and values of non-zero elements 212 | 213 | ```swift 214 | let indices = vec.indices 215 | let values = vec.values 216 | ``` 217 | 218 | ## History 219 | 220 | View the [changelog](https://github.com/pgvector/pgvector-swift/blob/master/CHANGELOG.md) 221 | 222 | ## Contributing 223 | 224 | Everyone is encouraged to help improve this project. Here are a few ways you can help: 225 | 226 | - [Report bugs](https://github.com/pgvector/pgvector-swift/issues) 227 | - Fix bugs and [submit pull requests](https://github.com/pgvector/pgvector-swift/pulls) 228 | - Write, clarify, or fix documentation 229 | - Suggest or add new features 230 | 231 | To get started with development: 232 | 233 | ```sh 234 | git clone https://github.com/pgvector/pgvector-swift.git 235 | cd pgvector-swift 236 | createdb pgvector_swift_test 237 | swift test 238 | ``` 239 | 240 | To run an example: 241 | 242 | ```sh 243 | cd Examples/Ollama 244 | createdb pgvector_example 245 | swift run 246 | ``` 247 | -------------------------------------------------------------------------------- /Sources/Pgvector/HalfVector.swift: -------------------------------------------------------------------------------- 1 | public struct HalfVector: Equatable { 2 | public private(set) var value: [Float16] 3 | 4 | public init(_ value: [Float16]) { 5 | self.value = value 6 | } 7 | 8 | public init?(_ string: String) { 9 | guard string.count >= 2, string.first == "[", string.last == "]" else { 10 | return nil 11 | } 12 | let start = string.index(string.startIndex, offsetBy: 1) 13 | let end = string.index(string.endIndex, offsetBy: -1) 14 | let parts = string[start.. String { 23 | return "[" + value.map { String($0) }.joined(separator: ",") + "]" 24 | } 25 | 26 | public static func == (lhs: HalfVector, rhs: HalfVector) -> Bool { 27 | return lhs.value == rhs.value 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /Sources/Pgvector/SparseVector.swift: -------------------------------------------------------------------------------- 1 | public struct SparseVector: Equatable { 2 | public private(set) var dim: Int 3 | public private(set) var indices: [Int] 4 | public private(set) var values: [Float] 5 | 6 | public init(_ value: [Float]) { 7 | var indices: [Int] = [] 8 | var values: [Float] = [] 9 | for (i, v) in value.enumerated() { 10 | if v != 0 { 11 | indices.append(i) 12 | values.append(v) 13 | } 14 | } 15 | 16 | self.dim = value.count 17 | self.indices = indices 18 | self.values = values 19 | } 20 | 21 | public init?(_ dictionary: [Int: Float], dim: Int) { 22 | guard dim >= 0 else { 23 | return nil 24 | } 25 | 26 | guard dictionary.allSatisfy({ $0.0 >= 0 && $0.0 < dim }) else { 27 | return nil 28 | } 29 | 30 | var elements = dictionary.filter { $1 != 0 }.map { ($0, $1) } 31 | elements.sort { $0.0 < $1.0 } 32 | 33 | self.dim = dim 34 | self.indices = elements.map { $0.0 } 35 | self.values = elements.map { $0.1 } 36 | } 37 | 38 | public init?(_ string: String) { 39 | let parts = string.split(separator: "/", maxSplits: 2) 40 | guard parts.count == 2 else { 41 | return nil 42 | } 43 | 44 | let elements = parts[0] 45 | guard elements.first == "{", elements.last == "}" else { 46 | return nil 47 | } 48 | 49 | guard let dim = Int(parts[1]) else { 50 | return nil 51 | } 52 | 53 | var indices: [Int] = [] 54 | var values: [Float] = [] 55 | 56 | if elements.count > 2 { 57 | let start = elements.index(elements.startIndex, offsetBy: 1) 58 | let end = elements.index(elements.endIndex, offsetBy: -1) 59 | for e in elements[start.. String { 76 | let elements = zip(indices, values).map { String($0 + 1) + ":" + String($1) } 77 | return "{" + elements.joined(separator: ",") + "}/" + String(dim) 78 | } 79 | 80 | public static func == (lhs: SparseVector, rhs: SparseVector) -> Bool { 81 | return lhs.dim == rhs.dim && lhs.indices == rhs.indices && lhs.values == rhs.values 82 | } 83 | } 84 | -------------------------------------------------------------------------------- /Sources/Pgvector/Vector.swift: -------------------------------------------------------------------------------- 1 | public struct Vector: Equatable { 2 | public private(set) var value: [Float] 3 | 4 | public init(_ value: [Float]) { 5 | self.value = value 6 | } 7 | 8 | public init?(_ string: String) { 9 | guard string.count >= 2, string.first == "[", string.last == "]" else { 10 | return nil 11 | } 12 | let start = string.index(string.startIndex, offsetBy: 1) 13 | let end = string.index(string.endIndex, offsetBy: -1) 14 | let parts = string[start.. String { 23 | return "[" + value.map { String($0) }.joined(separator: ",") + "]" 24 | } 25 | 26 | public static func == (lhs: Vector, rhs: Vector) -> Bool { 27 | return lhs.value == rhs.value 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /Sources/PgvectorClientKit/HalfVector.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresClientKit 3 | 4 | extension HalfVector: @retroactive PostgresValueConvertible { 5 | public var postgresValue: PostgresValue { 6 | return PostgresValue(text()) 7 | } 8 | } 9 | 10 | extension PostgresValue { 11 | public func halfVector() throws -> HalfVector { 12 | if isNull { 13 | throw PostgresError.valueIsNil 14 | } 15 | return try optionalHalfVector()! 16 | } 17 | 18 | public func optionalHalfVector() throws -> HalfVector? { 19 | guard let rawValue = rawValue else { return nil } 20 | 21 | guard let vector = HalfVector(rawValue) else { 22 | throw PostgresError.valueConversionError(value: self, type: HalfVector.self) 23 | } 24 | 25 | return vector 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /Sources/PgvectorClientKit/SparseVector.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresClientKit 3 | 4 | extension SparseVector: @retroactive PostgresValueConvertible { 5 | public var postgresValue: PostgresValue { 6 | return PostgresValue(text()) 7 | } 8 | } 9 | 10 | extension PostgresValue { 11 | public func sparseVector() throws -> SparseVector { 12 | if isNull { 13 | throw PostgresError.valueIsNil 14 | } 15 | return try optionalSparseVector()! 16 | } 17 | 18 | public func optionalSparseVector() throws -> SparseVector? { 19 | guard let rawValue = rawValue else { return nil } 20 | 21 | guard let vector = SparseVector(rawValue) else { 22 | throw PostgresError.valueConversionError(value: self, type: SparseVector.self) 23 | } 24 | 25 | return vector 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /Sources/PgvectorClientKit/Vector.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresClientKit 3 | 4 | extension Vector: @retroactive PostgresValueConvertible { 5 | public var postgresValue: PostgresValue { 6 | return PostgresValue(text()) 7 | } 8 | } 9 | 10 | extension PostgresValue { 11 | public func vector() throws -> Vector { 12 | if isNull { 13 | throw PostgresError.valueIsNil 14 | } 15 | return try optionalVector()! 16 | } 17 | 18 | public func optionalVector() throws -> Vector? { 19 | guard let rawValue = rawValue else { return nil } 20 | 21 | guard let vector = Vector(rawValue) else { 22 | throw PostgresError.valueConversionError(value: self, type: Vector.self) 23 | } 24 | 25 | return vector 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /Sources/PgvectorNIO/HalfVector.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresNIO 3 | 4 | extension HalfVector: @retroactive PostgresDynamicTypeEncodable { 5 | public static var psqlType: PostgresDataType? 6 | 7 | public var psqlType: PostgresDataType { 8 | HalfVector.psqlType! 9 | } 10 | 11 | public var psqlFormat: PostgresFormat { 12 | .binary 13 | } 14 | 15 | public func encode( 16 | into byteBuffer: inout ByteBuffer, 17 | context: PostgresEncodingContext 18 | ) { 19 | byteBuffer.writeInteger(Int16(value.count), as: Int16.self) 20 | byteBuffer.writeInteger(0, as: Int16.self) 21 | for v in value { 22 | byteBuffer.writeInteger(v.bitPattern, as: UInt16.self) 23 | } 24 | } 25 | } 26 | 27 | extension HalfVector: @retroactive PostgresDecodable { 28 | public init( 29 | from buffer: inout ByteBuffer, 30 | type: PostgresDataType, 31 | format: PostgresFormat, 32 | context: PostgresDecodingContext 33 | ) throws { 34 | guard type == HalfVector.psqlType else { 35 | throw PostgresDecodingError.Code.typeMismatch 36 | } 37 | 38 | guard format == .binary else { 39 | throw PostgresDecodingError.Code.failure 40 | } 41 | 42 | guard buffer.readableBytes >= 2, let dim = buffer.readInteger(as: Int16.self) else { 43 | throw PostgresDecodingError.Code.failure 44 | } 45 | 46 | guard buffer.readableBytes >= 2, let unused = buffer.readInteger(as: Int16.self), unused == 0 else { 47 | throw PostgresDecodingError.Code.failure 48 | } 49 | 50 | var value: [Float16] = [] 51 | for _ in 0..= 2, let v = buffer.readInteger(as: UInt16.self) else { 53 | throw PostgresDecodingError.Code.failure 54 | } 55 | value.append(Float16(bitPattern: v)) 56 | } 57 | self.init(value) 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /Sources/PgvectorNIO/PgvectorNIO.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresNIO 3 | 4 | public enum PgvectorError: Error { 5 | case string(String) 6 | } 7 | 8 | public struct PgvectorNIO { 9 | // note: global OID for each type is not ideal since it will be different for each database 10 | public static func registerTypes(_ client: PostgresClient) async throws { 11 | let rows = try await client.query("SELECT regtype('vector')::oid::integer, regtype('halfvec')::oid::integer, regtype('sparsevec')::oid::integer") 12 | 13 | var iterator = rows.makeAsyncIterator() 14 | guard let row = try await iterator.next() else { 15 | throw PgvectorError.string("unreachable") 16 | } 17 | let (vectorOid, halfvecOid, sparsevecOid) = try row.decode((Int?, Int?, Int?).self) 18 | 19 | if let oid = vectorOid { 20 | Vector.psqlType = PostgresDataType(UInt32(oid)) 21 | } else { 22 | throw PgvectorError.string("vector type not found in the database") 23 | } 24 | 25 | if let oid = halfvecOid { 26 | HalfVector.psqlType = PostgresDataType(UInt32(oid)) 27 | } 28 | 29 | if let oid = sparsevecOid { 30 | SparseVector.psqlType = PostgresDataType(UInt32(oid)) 31 | } 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /Sources/PgvectorNIO/SparseVector.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresNIO 3 | 4 | extension SparseVector: @retroactive PostgresDynamicTypeEncodable { 5 | public static var psqlType: PostgresDataType? 6 | 7 | public var psqlType: PostgresDataType { 8 | SparseVector.psqlType! 9 | } 10 | 11 | public var psqlFormat: PostgresFormat { 12 | .binary 13 | } 14 | 15 | public func encode( 16 | into byteBuffer: inout ByteBuffer, 17 | context: PostgresEncodingContext 18 | ) { 19 | byteBuffer.writeInteger(Int32(dim), as: Int32.self) 20 | byteBuffer.writeInteger(Int32(indices.count), as: Int32.self) 21 | byteBuffer.writeInteger(0, as: Int32.self) 22 | for v in indices { 23 | byteBuffer.writeInteger(Int32(v), as: Int32.self) 24 | } 25 | for v in values { 26 | byteBuffer.writeInteger(v.bitPattern, as: UInt32.self) 27 | } 28 | } 29 | } 30 | 31 | extension SparseVector: @retroactive PostgresDecodable { 32 | public init( 33 | from buffer: inout ByteBuffer, 34 | type: PostgresDataType, 35 | format: PostgresFormat, 36 | context: PostgresDecodingContext 37 | ) throws { 38 | guard type == SparseVector.psqlType else { 39 | throw PostgresDecodingError.Code.typeMismatch 40 | } 41 | 42 | guard format == .binary else { 43 | throw PostgresDecodingError.Code.failure 44 | } 45 | 46 | guard buffer.readableBytes >= 4, let dim = buffer.readInteger(as: Int32.self) else { 47 | throw PostgresDecodingError.Code.failure 48 | } 49 | 50 | guard buffer.readableBytes >= 4, let nnz = buffer.readInteger(as: Int32.self) else { 51 | throw PostgresDecodingError.Code.failure 52 | } 53 | 54 | guard buffer.readableBytes >= 4, let unused = buffer.readInteger(as: Int32.self), unused == 0 else { 55 | throw PostgresDecodingError.Code.failure 56 | } 57 | 58 | var indices: [Int] = [] 59 | for _ in 0..= 4, let v = buffer.readInteger(as: Int32.self) else { 61 | throw PostgresDecodingError.Code.failure 62 | } 63 | indices.append(Int(v)) 64 | } 65 | 66 | var values: [Float] = [] 67 | for _ in 0..= 4, let v = buffer.readInteger(as: UInt32.self) else { 69 | throw PostgresDecodingError.Code.failure 70 | } 71 | values.append(Float(bitPattern: v)) 72 | } 73 | 74 | self.init(Dictionary(uniqueKeysWithValues: zip(indices, values)), dim: Int(dim))! 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /Sources/PgvectorNIO/Vector.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import PostgresNIO 3 | 4 | extension Vector: @retroactive PostgresDynamicTypeEncodable { 5 | public static var psqlType: PostgresDataType? 6 | 7 | public var psqlType: PostgresDataType { 8 | Vector.psqlType! 9 | } 10 | 11 | public var psqlFormat: PostgresFormat { 12 | .binary 13 | } 14 | 15 | public func encode( 16 | into byteBuffer: inout ByteBuffer, 17 | context: PostgresEncodingContext 18 | ) { 19 | byteBuffer.writeInteger(Int16(value.count), as: Int16.self) 20 | byteBuffer.writeInteger(0, as: Int16.self) 21 | for v in value { 22 | byteBuffer.writeInteger(v.bitPattern) 23 | } 24 | } 25 | } 26 | 27 | extension Vector: @retroactive PostgresDecodable { 28 | public init( 29 | from buffer: inout ByteBuffer, 30 | type: PostgresDataType, 31 | format: PostgresFormat, 32 | context: PostgresDecodingContext 33 | ) throws { 34 | guard type == Vector.psqlType else { 35 | throw PostgresDecodingError.Code.typeMismatch 36 | } 37 | 38 | guard format == .binary else { 39 | throw PostgresDecodingError.Code.failure 40 | } 41 | 42 | guard buffer.readableBytes >= 2, let dim = buffer.readInteger(as: Int16.self) else { 43 | throw PostgresDecodingError.Code.failure 44 | } 45 | 46 | guard buffer.readableBytes >= 2, let unused = buffer.readInteger(as: Int16.self), unused == 0 else { 47 | throw PostgresDecodingError.Code.failure 48 | } 49 | 50 | var value: [Float] = [] 51 | for _ in 0..= 4, let v = buffer.readInteger(as: UInt32.self) else { 53 | throw PostgresDecodingError.Code.failure 54 | } 55 | value.append(Float(bitPattern: v)) 56 | } 57 | self.init(value) 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /Tests/PgvectorTests/HalfVectorTests.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import Testing 3 | 4 | final class HalfVectorTests { 5 | @Test func equatable() { 6 | let a = HalfVector([1, 2, 3]) 7 | let b = HalfVector([1, 2, 3]) 8 | let c = HalfVector([1, 2, 4]) 9 | #expect(a == a) 10 | #expect(a == b) 11 | #expect(a != c) 12 | } 13 | } 14 | -------------------------------------------------------------------------------- /Tests/PgvectorTests/PostgresClientKitTests.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | import Pgvector 3 | import PgvectorClientKit 4 | import PostgresClientKit 5 | import Testing 6 | 7 | final class PostgresClientKitTests { 8 | @Test func example() throws { 9 | var configuration = PostgresClientKit.ConnectionConfiguration() 10 | configuration.database = "pgvector_swift_test" 11 | configuration.ssl = false 12 | configuration.user = ProcessInfo.processInfo.environment["USER"]! 13 | 14 | let connection = try PostgresClientKit.Connection(configuration: configuration) 15 | defer { connection.close() } 16 | 17 | var text = "CREATE EXTENSION IF NOT EXISTS vector" 18 | var statement = try connection.prepareStatement(text: text) 19 | try statement.execute() 20 | 21 | text = "DROP TABLE IF EXISTS items" 22 | statement = try connection.prepareStatement(text: text) 23 | try statement.execute() 24 | 25 | text = "CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))" 26 | statement = try connection.prepareStatement(text: text) 27 | try statement.execute() 28 | 29 | text = "INSERT INTO items (embedding) VALUES ($1), ($2), ($3)" 30 | statement = try connection.prepareStatement(text: text) 31 | try statement.execute(parameterValues: [Vector([1, 1, 1]), Vector([2, 2, 2]), Vector([1, 1, 2])]) 32 | 33 | text = "SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5" 34 | statement = try connection.prepareStatement(text: text) 35 | let cursor = try statement.execute(parameterValues: [Vector([1, 1, 1])]) 36 | 37 | for row in cursor { 38 | let columns = try row.get().columns 39 | let id = try columns[0].int() 40 | let embedding = try columns[1].vector() 41 | print(id, embedding) 42 | } 43 | 44 | text = "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)" 45 | statement = try connection.prepareStatement(text: text) 46 | try statement.execute() 47 | 48 | text = "SELECT $1::vector, $2::halfvec, $3::sparsevec" 49 | statement = try connection.prepareStatement(text: text) 50 | let typesCursor = try statement.execute(parameterValues: [Vector([1, 2, 3]), HalfVector([1, 2, 3]), SparseVector([1, 0, 2, 0, 3, 0])]) 51 | 52 | for row in typesCursor { 53 | let columns = try row.get().columns 54 | let embedding = try columns[0].vector() 55 | let halfEmbedding = try columns[1].halfVector() 56 | let sparseEmbedding = try columns[2].sparseVector() 57 | print(embedding, halfEmbedding, sparseEmbedding) 58 | } 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /Tests/PgvectorTests/PostgresNIOTests.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | import Pgvector 3 | import PgvectorNIO 4 | import PostgresNIO 5 | import Testing 6 | 7 | final class PostgresNIOTests { 8 | @Test func example() async throws { 9 | let config = PostgresClient.Configuration( 10 | host: "localhost", 11 | port: 5432, 12 | username: ProcessInfo.processInfo.environment["USER"]!, 13 | password: nil, 14 | database: "pgvector_swift_test", 15 | tls: .disable 16 | ) 17 | 18 | let client = PostgresClient(configuration: config) 19 | 20 | try await withThrowingTaskGroup(of: Void.self) { taskGroup in 21 | taskGroup.addTask { 22 | await client.run() 23 | } 24 | 25 | try await client.query("CREATE EXTENSION IF NOT EXISTS vector") 26 | try await PgvectorNIO.registerTypes(client) 27 | 28 | try await client.query("DROP TABLE IF EXISTS nio_items") 29 | try await client.query("CREATE TABLE nio_items (id bigserial PRIMARY KEY, embedding vector(3))") 30 | 31 | let embedding1 = Vector([1, 1, 1]) 32 | let embedding2 = Vector([2, 2, 2]) 33 | let embedding3 = Vector([1, 1, 2]) 34 | try await client.query("INSERT INTO nio_items (embedding) VALUES (\(embedding1)), (\(embedding2)), (\(embedding3))") 35 | 36 | let embedding = Vector([1, 1, 1]) 37 | let rows = try await client.query("SELECT id, embedding FROM nio_items ORDER BY embedding <-> \(embedding) LIMIT 5") 38 | for try await (id, embedding) in rows.decode((Int, Vector).self) { 39 | print(id, embedding) 40 | } 41 | 42 | try await client.query("CREATE INDEX ON nio_items USING hnsw (embedding vector_l2_ops)") 43 | 44 | let halfEmbedding = HalfVector([1, 2, 3]) 45 | let sparseEmbedding = SparseVector([1, 0, 2, 0, 3, 0]) 46 | let typeRows = try await client.query("SELECT \(embedding), \(halfEmbedding), \(sparseEmbedding)") 47 | for try await (v, h, s) in typeRows.decode((Vector, HalfVector, SparseVector).self) { 48 | print(v, h, s) 49 | } 50 | 51 | taskGroup.cancelAll() 52 | } 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /Tests/PgvectorTests/SparseVectorTests.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import Testing 3 | 4 | final class SparseVectorTests { 5 | @Test func equatable() { 6 | let a = SparseVector([1, 0, 2, 0, 3, 0]) 7 | let b = SparseVector([1, 0, 2, 0, 3, 0]) 8 | let c = SparseVector([1, 0, 2, 0, 4, 0]) 9 | #expect(a == a) 10 | #expect(a == b) 11 | #expect(a != c) 12 | } 13 | 14 | @Test func fromDictionary() { 15 | let a = SparseVector([2: 2, 4: 3, 1: 0, 0: 1], dim: 6)! 16 | #expect(a.dim == 6) 17 | #expect(a.indices == [0, 2, 4]) 18 | #expect(a.values == [1, 2, 3]) 19 | } 20 | } 21 | -------------------------------------------------------------------------------- /Tests/PgvectorTests/VectorTests.swift: -------------------------------------------------------------------------------- 1 | import Pgvector 2 | import Testing 3 | 4 | final class VectorTests { 5 | @Test func equatable() { 6 | let a = Vector([1, 2, 3]) 7 | let b = Vector([1, 2, 3]) 8 | let c = Vector([1, 2, 4]) 9 | #expect(a == a) 10 | #expect(a == b) 11 | #expect(a != c) 12 | } 13 | } 14 | --------------------------------------------------------------------------------