Swift PGVector

Vector.swift

import PostgresNIO
import Foundation

/// A vector of floats. This is a custom type mirroring the `vector` type from
/// [pgvector](https://github.com/pgvector/pgvector/).
struct Vector: Codable {
    /// The vector's values.
    let value: [Float]

    /// Creates a new `Vector`.
    ///
    /// - Parameter value: The vector's values.
    /// - Returns: A new `Vector`.
    init(value: [Float]) {
        self.value = value
    }

    /// Loads the `PostgresDataType` for `Vector` from the database.
    /// This is a required step, because the `PostgresDataType`/ `oid` is not known at compile
    /// time (since it is a custom type).
    static func loadPsqlType(from database: PostgresDatabase) async throws {
        let queryString = "SELECT 'vector'::regtype::integer"
        guard let oidColumn = try await database.query(queryString).get().first?.first else {
            throw PostgresDecodingError.Code.missingData
        }
        Self.postgresDataType = try oidColumn.decode(PostgresDataType.self, context: .default)
    }
}

// From: https://forums.swift.org/t/how-to-create-a-postgresdata-for-one-of-postgres-point-types/56939/2
// and: https://github.com/vapor/postgres-nio/issues/280
extension Vector: PostgresDataConvertible {
    // initially set to .null, but will be updated in `loadPsqlType` below
    static var postgresDataType: PostgresNIO.PostgresDataType = .null

    var postgresData: PostgresNIO.PostgresData? {
        var byteBuffer = ByteBufferAllocator().buffer(capacity: 0)
        byteBuffer.writeMultipleIntegers(UInt16(value.count), (0 as UInt16))
        for element in value {
            byteBuffer.postgresWriteFloat(element)
        }
        return PostgresData(
            type: Vector.postgresDataType,
            typeModifier: Int32(value.count),
            formatCode: .binary,
            value: byteBuffer
        )
    }

    /// Creates a `Vector` from a `PostgresData`.
    ///
    /// - Parameter postgresData: The `PostgresData` to decode.
    /// - Returns: The decoded `Vector` or `nil` if the `PostgresData` could not be decoded.
    init?(postgresData: PostgresNIO.PostgresData) {
        guard var byteBuffer = postgresData.value else {
            return nil
        }

        guard let (dimension, _) = byteBuffer.readMultipleIntegers(endianness: .big, as: (UInt16, UInt16).self) else {
            return nil
        }

        var floatArray: [Float] = []
        for _ in 0..<dimension {
            guard let element: Float = byteBuffer.postgresReadFloat() else {
                return nil
            }
            floatArray.append(element)
        }
        self.value = floatArray
    }
}