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 structurelet 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 YOLOlet 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.```bashcargo install-pgx```But importantly if you don't have and postgres installation, you will need to run```bashcargo pgx init```This installs a `.pgx/` directory in your home directorym and downloads some linking requirements... or somethingThis currently only targets Postgres 13, as such the only way to run it is```bashcargo 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```sqlCREATE EXTENSION pgvector_rs```If you update the provided functions or types you will need to drop the extension and recreate it.