Can I model a vector type using pgvector?

[Daniel-A-Gutierrez]

SeaORM doesn't yet have this functionality - it seems to be just about complete but they're waiting on this to be merged.

It'd be doable in sqlx but I'm just wondering how I'd go about it without screwing things up with loco.

[fivehanz]

I made it custom for now.

pg_vector.rs custom type file

use sea_orm_migration::sea_query::types::Iden;
use std::fmt;

// default vector size
pub const DEFAULT_VECTOR_SIZE: usize = 1536;

// define a vector type
pub type VectorEmbedding = Vec<f32>;

// SeaORM vector type definition
pub struct PGVectorType<const SIZE: usize>;

// implement Iden for SeaORM
impl<const SIZE: usize> Iden for PGVectorType<SIZE> {
    fn unquoted(&self, s: &mut dyn fmt::Write) {
        write!(s, "VECTOR({})", SIZE).unwrap(); // Dynamically set the dimension, default is 1536
    }
}

in the migrations file

                    .col(
                        ColumnDef::new(Chunks::Embedding)
                            .custom(PGVectorType::<DEFAULT_VECTOR_SIZE>) // entity type VectorEmbedding
                            .not_null(),
                    )

this creates expected tables properly.

as for entities, the type needs to be added to generated entities. for now manually.

    #[sea_orm(column_type = "custom(\"vector\")")]
    pub embedding: VectorEmbedding,

I am yet to test and use the generated entities.

[Daniel-A-Gutierrez]

Nice, thanks

[cpunion]

@fivehanz Hi, thanks for your code. I used it and do migration success, but test failed. When call:

let created = chunk.insert(&boot.app_context.db).await.unwrap();

the record wrote, but not returned, seems decode failed:

called `Result::unwrap()` on an `Err` value: Query(SqlxError(ColumnDecode { index: "\"embedding\"", source: "mismatched types; Rust type `core::option::Option<alloc::vec::Vec<f32>>` (as SQL type `FLOAT4[]`) is not compatible with SQL type `vector`" }))

Do you meet this issue? how can I fix it?

The full log:

test models::chunks::test_update_chunk ... FAILED
thread 'models::chunks::test_create_chunk' panicked at tests/models/chunks.rs:124:59:
called `Result::unwrap()` on an `Err` value: Query(SqlxError(ColumnDecode { index: "\"embedding\"", source: "mismatched types; Rust type `core::option::Option<alloc::vec::Vec<f32>>` (as SQL type `FLOAT4[]`) is not compatible with SQL type `vector`" }))
stack backtrace:
   0: _rust_begin_unwind
   1: core::panicking::panic_fmt
   2: core::result::unwrap_failed
   3: core::result::Result<T,E>::unwrap
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/core/src/result.rs:1104:23
   4: mod::models::chunks::test_create_chunk::{{closure}}::{{closure}}
             at ./tests/models/chunks.rs:124:18
   5: <core::pin::Pin<P> as core::future::future::Future>::poll
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/core/src/future/future.rs:123:9
   6: <core::pin::Pin<P> as core::future::future::Future>::poll
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/core/src/future/future.rs:123:9
   7: tokio::runtime::scheduler::current_thread::CoreGuard::block_on::{{closure}}::{{closure}}::{{closure}}
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:729:57
   8: tokio::runtime::coop::with_budget
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/coop.rs:107:5
   9: tokio::runtime::coop::budget
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/coop.rs:73:5
  10: tokio::runtime::scheduler::current_thread::CoreGuard::block_on::{{closure}}::{{closure}}
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:729:25
  11: tokio::runtime::scheduler::current_thread::Context::enter
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:428:19
  12: tokio::runtime::scheduler::current_thread::CoreGuard::block_on::{{closure}}
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:728:36
  13: tokio::runtime::scheduler::current_thread::CoreGuard::enter::{{closure}}
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:807:68
  14: tokio::runtime::context::scoped::Scoped<T>::set
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/context/scoped.rs:40:9
  15: tokio::runtime::context::set_scheduler::{{closure}}
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/context.rs:180:26
  16: std::thread::local::LocalKey<T>::try_with
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/std/src/thread/local.rs:283:12
  17: std::thread::local::LocalKey<T>::with
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/std/src/thread/local.rs:260:9
  18: tokio::runtime::context::set_scheduler
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/context.rs:180:9
  19: tokio::runtime::scheduler::current_thread::CoreGuard::enter
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:807:27
  20: tokio::runtime::scheduler::current_thread::CoreGuard::block_on
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:716:19
  21: tokio::runtime::scheduler::current_thread::CurrentThread::block_on::{{closure}}
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:196:28
  22: tokio::runtime::context::runtime::enter_runtime
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/context/runtime.rs:65:16
  23: tokio::runtime::scheduler::current_thread::CurrentThread::block_on
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/scheduler/current_thread/mod.rs:184:9
  24: tokio::runtime::runtime::Runtime::block_on_inner
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/runtime.rs:368:47
  25: tokio::runtime::runtime::Runtime::block_on
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tokio-1.42.0/src/runtime/runtime.rs:342:13
  26: mod::models::chunks::test_create_chunk::{{closure}}
             at ./tests/models/chunks.rs:125:5
  27: core::ops::function::FnOnce::call_once
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/core/src/ops/function.rs:250:5
  28: serial_test::serial_code_lock::local_serial_core
             at /Users/lijie/.cargo/registry/src/index.crates.io-6f17d22bba15001f/serial_test-3.2.0/src/serial_code_lock.rs:36:5
  29: mod::models::chunks::test_create_chunk
             at ./tests/models/chunks.rs:96:1
  30: mod::models::chunks::test_create_chunk::{{closure}}
             at ./tests/models/chunks.rs:96:10
  31: core::ops::function::FnOnce::call_once
             at /private/tmp/rust-20241128-8064-v5d6qj/rustc-1.83.0-src/library/core/src/ops/function.rs:250:5
note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.
test models::chunks::test_create_chunk ... FAILED

[cpunion]

Just fixed by add select_as:

#[sea_orm(column_type = "custom(\"vector\")", select_as = "float4[]")]
pub embedding: VectorEmbedding,

But most cases we doesn't need to query this field, and I try to find how to exclude it. I'm a newbie in Rust and loco.