OR27N4EKS57UDYZR6SWQCY74P772H7WOR3FAJKXF6YADOQTACW7QC
/*
amroutine.ambuildempty = todo!
amroutine.aminsert = todo!;
amroutine.ambulkdelete = todo!;
amroutine.amvacuumcleanup = todo!;
amroutine.amcostestimate = todo!;
amroutine.amoptions = todo!;
amroutine.amvalidate = todo!;
amroutine.ambeginscan = todo!;
amroutine.amrescan = todo!;
amroutine.amgettuple = todo!;
amroutine.amendscan = todo!;
*/
amroutine.ambuildempty = Some(build::vec_mem_cache_buildempty);
amroutine.aminsert = None; // todo!();
amroutine.ambulkdelete = None; // todo!();
amroutine.amvacuumcleanup = None; // todo!();
amroutine.amcostestimate = None; //todo!();
amroutine.amoptions = Some(options::vec_mem_cache_options);
amroutine.amvalidate = Some(amvalidate);
amroutine.ambeginscan = None; //todo!();
amroutine.amrescan = None; //todo!();
amroutine.amgettuple = None; //todo!();
amroutine.amendscan = None; //todo!();
}
use pgx::*;
#[pg_guard]
pub extern "C" fn amvalidate(_opclassoid: pg_sys::Oid) -> bool {
true
}
#[cfg(test)]
mod test {
// We want to try out something like
// ```sql
// CREATE TABLE public.vecs_2
//(
// idx bigserial,
// vec real[] CONSTRAINT dim512 CHECK (cardinality(vec) = 512),
// PRIMARY KEY (idx)
//);
//
// ```
fn test_constraints() {
todo!() // YOLO
}
}
struct BuildState<'a> {
tupdesc: &'a PgTupleDesc<'a>,
dimension: u16,
// attributes: Vec<CategorizedAttribute<'a>>,
memcxt: PgMemoryContexts,
pub rows_added: u32,
}
impl<'a> BuildState<'a> {
fn new(
tupdesc: &'a PgTupleDesc,
dimension: u16,
// attributes: Vec<CategorizedAttribute<'a>>,
) -> Self {
BuildState {
tupdesc,
dimension,
memcxt: PgMemoryContexts::new("pgvec_rs context"),
rows_added: 0,
}
}
}
#[cfg(feature = "pg13")]
#[pg_guard]
unsafe extern "C" fn build_callback(
_index: pg_sys::Relation,
ctid: pg_sys::ItemPointer,
values: *mut pg_sys::Datum,
_isnull: *mut bool,
_tuple_is_alive: bool,
state: *mut std::os::raw::c_void,
) {
build_callback_internal(*ctid, values, state);
}
#[inline(always)]
unsafe extern "C" fn build_callback_internal(
ctid: pg_sys::ItemPointerData,
values: *mut pg_sys::Datum,
state: *mut std::os::raw::c_void,
) {
check_for_interrupts!();
let state = (state as *mut BuildState).as_mut().unwrap();
let old_context = state.memcxt.set_as_current();
// TODO: Add check for structure
let values = std::slice::from_raw_parts(values, 1);
let VecRow { id: _, data: vec } = row_to_vec_iter(&state.tupdesc, values[0]).unwrap();
let mut vec_allocator = VECTOR_ALLOCATOR.exclusive();
vec_allocator.add_vec(&vec);
state.rows_added += 1;
}
fn do_heap_scan<'a>(
index_info: *mut pg_sys::IndexInfo,
heap_relation: &'a PgRelation,
index_relation: &'a PgRelation,
tupdesc: &'a PgTupleDesc,
) -> u32 {
// Should be able a to get a dimension from the description, but we'll hard code it now YOLO
let dimension = 512;
let mut state = BuildState::new(&tupdesc, dimension);
unsafe {
pg_sys::IndexBuildHeapScan(
heap_relation.as_ptr(),
index_relation.as_ptr(),
index_info,
Some(build_callback),
&mut state,
);
}
state.rows_added
}
#[derive(Default, Clone)]
struct VecRowPartial {
id: Option<i64>,
data: Option<Vec<f32>>,
}
impl VecRowPartial {
fn check(self) -> Result<VecRow, &'static str> {
match self {
VecRowPartial {
id: Some(id),
data: Some(data),
} => Ok(VecRow { id, data }),
VecRowPartial {
id: None,
data: Some(_),
} => Err("Missing Id for vec row"),
VecRowPartial {
id: Some(_),
data: None,
} => Err("Missing vector data for vec row"),
VecRowPartial {
id: None,
data: None,
} => Err("Missing all data for vector"),
}
}
}
struct VecRow {
id: i64,
data: Vec<f32>,
#[inline]
unsafe fn row_to_vec_iter<'a>(
tupdesc: &'a PgTupleDesc,
row: pg_sys::Datum,
) -> Result<VecRow, &'static str> {
let td = pg_sys::pg_detoast_datum(row as *mut pg_sys::varlena) as pg_sys::HeapTupleHeader;
let mut tmptup = pg_sys::HeapTupleData {
t_len: varsize(td as *mut pg_sys::varlena) as u32,
t_self: Default::default(),
t_tableOid: 0,
t_data: td,
};
let mut datums = vec![0 as pg_sys::Datum; tupdesc.natts as usize];
let mut nulls = vec![false; tupdesc.natts as usize];
pg_sys::heap_deform_tuple(
&mut tmptup,
tupdesc.as_ptr(),
datums.as_mut_ptr(),
nulls.as_mut_ptr(),
);
let mut drop_cnt = 0;
let mut vec_row = VecRowPartial::default();
tupdesc
.iter()
.map(|attribute| {
let is_dropped = attribute.is_dropped();
let array_type = unsafe { pg_sys::get_element_type(attribute.type_oid().value()) };
let (base_oid, is_array) = if array_type != pg_sys::InvalidOid {
(PgOid::from(array_type), true)
} else {
(attribute.type_oid(), false)
};
let typoid = base_oid;
let fill_process: Result<(), &'static str> = match &typoid {
PgOid::BuiltIn(builtin) => match (builtin, is_array) {
(PgBuiltInOids::FLOAT4OID, true) => {
if let None = vec_row.data {
let data =
unsafe { Vec::<f32>::from_datum(row, false, builtin.value()) }
.unwrap();
vec_row.data = Some(data);
Ok(())
} else {
Err("Received more than one vector data entries in row")
}
}
(PgBuiltInOids::INT8OID, false) => {
if let None = vec_row.id {
let id =
unsafe { i64::from_datum(row, false, builtin.value()) }.unwrap();
vec_row.id = Some(id);
Ok(())
} else {
Err("Received more than one id in the row")
}
}
(_, _) => {
Err("This row is not the right shape. It should be (bigint, real[])")
}
},
_ => Err("This element is not an accepted type"), // todo: communicate the right and wrong types
};
fill_process
})
.collect::<Result<Vec<()>, &'static str>>()?;
vec_row.check()
}
#[pg_guard]
pub extern "C" fn vec_mem_cache_buildempty(_index_relation: pg_sys::Relation) {}
You'll need [pgx](https://github.com/zombodb/pgx). Follow ths instructions there.
```bash
cargo install-pgx
```
But importantly if you don't have and postgres installation, you will need to run
```bash
cargo pgx init
```
This installs a `.pgx/` directory in your home directorym and downloads some linking requirements... or something
This currently only targets Postgres 13, as such the only way to run it is
```bash
cargo pgx run pg13
```
This also installs the extension, so if you run your `.pgx` version of Postgres 13 it will have the extension binaries.
To actually install it run
```sql
CREATE EXTENSION pgvector_rs
```
If you update the provided functions or types you will need to drop the extension and recreate it.