Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: move algorithm to a crate #172

Merged
merged 1 commit into from
Feb 5, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 13 additions & 6 deletions .taplo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,16 @@
indent_string = " "

[[rule]]
keys = ["dependencies", "*-denpendencies", "lints", "patch.*", "profile.*"]

[rule.formatting]
reorder_keys = true
reorder_arrays = true
align_comments = true
include = ["**/Cargo.toml"]
keys = [
"dependencies",
"dev-dependencies",
"build-dependencies",
"target.*.dependencies",
"lints",
"patch.*",
"profile.*",
"workspace.dependencies",
"lints.dependencies",
]
formatting = { reorder_keys = true, reorder_arrays = true, align_comments = true }
52 changes: 40 additions & 12 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

17 changes: 9 additions & 8 deletions Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[package]
name = "vchord"
version.workspace = true
edition.workspace = true
edition = "2021"

[lib]
name = "vchord"
Expand All @@ -20,25 +20,21 @@ pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"]
pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"]

[dependencies]
always_equal = { path = "./crates/always_equal" }
algorithm = { path = "./crates/algorithm" }
distance = { path = "./crates/distance" }
rabitq = { path = "./crates/rabitq" }
k_means = { path = "./crates/k_means" }
random_orthogonal_matrix = { path = "./crates/random_orthogonal_matrix" }
simd = { path = "./crates/simd" }
vector = { path = "./crates/vector" }

half.workspace = true
log = "0.4.25"
paste = "1"
pgrx = { version = "=0.12.9", default-features = false, features = ["cshim"] }
pgrx-catalog = "0.1.0"
rand.workspace = true
rayon = "1.10.0"
serde.workspace = true
toml = "0.8.19"
validator = { version = "0.19.0", features = ["derive"] }
zerocopy = "0.8.14"
zerocopy-derive = "0.8.14"
validator.workspace = true

[patch.crates-io]
half = { git = "https://github.com/tensorchord/half-rs.git", rev = "3f9a8843d6722bd1833de2289347640ad8770146" }
Expand All @@ -56,14 +52,19 @@ edition = "2021"

[workspace.dependencies]
half = { version = "2.4.1", features = ["serde", "zerocopy"] }
log = "0.4.25"
rand = "0.8.5"
serde = "1"
validator = { version = "0.20.0", features = ["derive"] }
zerocopy = "0.8.14"
zerocopy-derive = "0.8.14"

[workspace.lints]
clippy.identity_op = "allow"
clippy.int_plus_one = "allow"
clippy.needless_range_loop = "allow"
clippy.nonminimal_bool = "allow"
rust.unsafe_code = "deny"
rust.unsafe_op_in_unsafe_fn = "deny"
rust.unused_lifetimes = "warn"
rust.unused_qualifications = "warn"
Expand Down
25 changes: 25 additions & 0 deletions crates/algorithm/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
[package]
name = "algorithm"
version.workspace = true
edition.workspace = true

[dependencies]
always_equal = { path = "../always_equal" }
distance = { path = "../distance" }
k_means = { path = "../k_means" }
rabitq = { path = "../rabitq" }
random_orthogonal_matrix = { path = "../random_orthogonal_matrix" }
simd = { path = "../simd" }
vector = { path = "../vector" }

half.workspace = true
paste = "1"
rand.workspace = true
serde.workspace = true
toml = "0.8.19"
validator.workspace = true
zerocopy.workspace = true
zerocopy-derive.workspace = true

[lints]
workspace = true
101 changes: 101 additions & 0 deletions crates/algorithm/src/build.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
use crate::RelationWrite;
use crate::operator::{Accessor2, Operator, Vector};
use crate::tape::*;
use crate::tuples::*;
use crate::types::*;
use vector::VectorOwned;

pub fn build<O: Operator>(
vector_options: VectorOptions,
vchordrq_options: VchordrqIndexingOptions,
index: impl RelationWrite,
structures: Vec<Structure<O::Vector>>,
) {
let dims = vector_options.dims;
let is_residual = vchordrq_options.residual_quantization && O::SUPPORTS_RESIDUAL;
let mut meta = TapeWriter::<_, _, MetaTuple>::create(|| index.extend(false));
assert_eq!(meta.first(), 0);
let freepage = TapeWriter::<_, _, FreepageTuple>::create(|| index.extend(false));
let mut vectors = TapeWriter::<_, _, VectorTuple<O::Vector>>::create(|| index.extend(true));
let mut pointer_of_means = Vec::<Vec<IndexPointer>>::new();
for i in 0..structures.len() {
let mut level = Vec::new();
for j in 0..structures[i].len() {
let vector = structures[i].means[j].as_borrowed();
let (metadata, slices) = O::Vector::vector_split(vector);
let mut chain = Ok(metadata);
for i in (0..slices.len()).rev() {
chain = Err(vectors.push(match chain {
Ok(metadata) => VectorTuple::_0 {
payload: None,
elements: slices[i].to_vec(),
metadata,
},
Err(pointer) => VectorTuple::_1 {
payload: None,
elements: slices[i].to_vec(),
pointer,
},
}));
}
level.push(chain.err().unwrap());
}
pointer_of_means.push(level);
}
let mut pointer_of_firsts = Vec::<Vec<u32>>::new();
for i in 0..structures.len() {
let mut level = Vec::new();
for j in 0..structures[i].len() {
if i == 0 {
let tape = TapeWriter::<_, _, H0Tuple>::create(|| index.extend(false));
let mut jump = TapeWriter::<_, _, JumpTuple>::create(|| index.extend(false));
jump.push(JumpTuple {
first: tape.first(),
});
level.push(jump.first());
} else {
let mut tape = H1TapeWriter::<_, _>::create(|| index.extend(false));
let h2_mean = structures[i].means[j].as_borrowed();
let h2_children = structures[i].children[j].as_slice();
for child in h2_children.iter().copied() {
let h1_mean = structures[i - 1].means[child as usize].as_borrowed();
let code = if is_residual {
let mut residual_accessor = O::ResidualAccessor::default();
residual_accessor.push(
O::Vector::elements_and_metadata(h1_mean).0,
O::Vector::elements_and_metadata(h2_mean).0,
);
let residual = residual_accessor.finish(
O::Vector::elements_and_metadata(h1_mean).1,
O::Vector::elements_and_metadata(h2_mean).1,
);
O::Vector::code(residual.as_borrowed())
} else {
O::Vector::code(h1_mean)
};
tape.push(H1Branch {
mean: pointer_of_means[i - 1][child as usize],
dis_u_2: code.dis_u_2,
factor_ppc: code.factor_ppc,
factor_ip: code.factor_ip,
factor_err: code.factor_err,
signs: code.signs,
first: pointer_of_firsts[i - 1][child as usize],
});
}
let tape = tape.into_inner();
level.push(tape.first());
}
}
pointer_of_firsts.push(level);
}
meta.push(MetaTuple {
dims,
height_of_root: structures.len() as u32,
is_residual,
vectors_first: vectors.first(),
root_mean: pointer_of_means.last().unwrap()[0],
root_first: pointer_of_firsts.last().unwrap()[0],
freepage_first: freepage.first(),
});
}
Loading
Loading