Skip to content

Commit 8c60685

Browse files
committed
Optimize PrimaryKey memory: serialize CqlValue into compact byte buffer
Replace PrimaryKey(Vec<CqlValue>) with a serialized byte buffer representation backed by a new InvariantKey type. CqlValue is 72 bytes per element due to enum sizing, wasting significant memory for typical primary key types like Int (4 bytes of data, 68 bytes of padding). InvariantKey stores values as Arc<[u8]> with a 1-byte type tag followed by the minimal binary encoding for each value: Int: 5 bytes (was 72) — 14× smaller BigInt: 9 bytes (was 72) — 8× smaller Uuid: 17 bytes (was 72) — 4× smaller Text(s): 5+len (was 72) — variable For a single Int primary key column, total per-row memory drops from ~96 bytes (24 Vec overhead + 72 CqlValue) to ~22 bytes (16 Arc<[u8]> + 6 heap), a 4× improvement. With millions of indexed rows stored in the BiMap, this substantially reduces RSS. Design: - invariant_key.rs: reusable encoding engine with InvariantKey, InvariantKeyBuilder, and InvariantKeyIter types. - primary_key.rs: thin newtype providing primary-key-specific semantics. - O(1) clone via Arc; values decoded on demand via get(index) or iter(), acceptable for primary keys with 1–3 columns. - Hash/Eq operate on raw bytes instead of format!("{:?}"), which is both faster and avoids Debug format instability. - hash_prefix() uses a separate hash domain from Hash::hash (documented). - Variable-length encoding guards against u32 overflow via try_into(). - InvariantKeyIter implements ExactSizeIterator and FusedIterator. Changes: - Add invariant_key.rs with encode/decode, Hash, Eq, Debug, Iterator - Add primary_key.rs as a newtype over InvariantKey - Update usearch.rs: closures return owned CqlValue instead of refs - Update httproutes.rs: use PrimaryKey::len()/get() API - Remove old PrimaryKey struct and format-based Hash impl from lib.rs Fixes: VECTOR-526
1 parent eb3c999 commit 8c60685

File tree

5 files changed

+1098
-42
lines changed

5 files changed

+1098
-42
lines changed

crates/vector-store/src/httproutes.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -606,16 +606,20 @@ async fn post_index_ann(
606606
let primary_keys: anyhow::Result<_> = primary_keys
607607
.iter()
608608
.map(|primary_key| {
609-
if primary_key.0.len() != primary_key_columns.len() {
609+
if primary_key.len() != primary_key_columns.len() {
610610
bail!(
611611
"wrong size of a primary key: {}, {}",
612612
primary_key_columns.len(),
613-
primary_key.0.len()
613+
primary_key.len()
614614
);
615615
}
616616
Ok(primary_key)
617617
})
618-
.map_ok(|primary_key| primary_key.0[idx_column].clone())
618+
.map_ok(|primary_key| {
619+
primary_key
620+
.get(idx_column)
621+
.expect("primary key index out of bounds after length check")
622+
})
619623
.map_ok(try_to_json)
620624
.map(|primary_key| primary_key.flatten())
621625
.collect();

crates/vector-store/src/index/usearch.rs

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -897,15 +897,15 @@ fn cql_cmp(lhs: &CqlValue, rhs: &CqlValue) -> Option<std::cmp::Ordering> {
897897

898898
/// Lexicographically compare tuple values.
899899
/// Returns the ordering of the first non-equal pair, or Equal if all pairs are equal.
900-
fn cql_cmp_tuple<'a>(
901-
primary_key: &'a PrimaryKey,
902-
primary_key_value: impl Fn(&'a PrimaryKey, &ColumnName) -> Option<&'a CqlValue>,
900+
fn cql_cmp_tuple(
901+
primary_key: &PrimaryKey,
902+
primary_key_value: impl Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue>,
903903
lhs: &[ColumnName],
904904
rhs: &[CqlValue],
905905
) -> Option<std::cmp::Ordering> {
906906
for (col, rhs_val) in lhs.iter().zip(rhs.iter()) {
907907
let lhs_val = primary_key_value(primary_key, col)?;
908-
match cql_cmp(lhs_val, rhs_val)? {
908+
match cql_cmp(&lhs_val, rhs_val)? {
909909
std::cmp::Ordering::Equal => continue,
910910
other => return Some(other),
911911
}
@@ -924,17 +924,17 @@ fn filtered_ann(
924924
) {
925925
fn annotate<F>(f: F) -> F
926926
where
927-
F: for<'a, 'b> Fn(&'a PrimaryKey, &'b ColumnName) -> Option<&'a CqlValue>,
927+
F: Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue>,
928928
{
929929
f
930930
}
931931

932932
let primary_key_value = annotate(
933-
|primary_key: &PrimaryKey, name: &ColumnName| -> Option<&CqlValue> {
933+
|primary_key: &PrimaryKey, name: &ColumnName| -> Option<CqlValue> {
934934
primary_key_columns
935935
.iter()
936936
.position(|key_column| key_column == name)
937-
.and_then(move |idx| primary_key.0.get(idx))
937+
.and_then(move |idx| primary_key.get(idx))
938938
},
939939
);
940940

@@ -946,27 +946,29 @@ fn filtered_ann(
946946
.restrictions
947947
.iter()
948948
.all(|restriction| match restriction {
949-
Restriction::Eq { lhs, rhs } => primary_key_value(&primary_key, lhs) == Some(rhs),
949+
Restriction::Eq { lhs, rhs } => {
950+
primary_key_value(&primary_key, lhs).as_ref() == Some(rhs)
951+
}
950952
Restriction::In { lhs, rhs } => {
951953
let value = primary_key_value(&primary_key, lhs);
952-
rhs.iter().any(|rhs| value == Some(rhs))
954+
rhs.iter().any(|rhs| value.as_ref() == Some(rhs))
953955
}
954956
Restriction::Lt { lhs, rhs } => primary_key_value(&primary_key, lhs)
955-
.and_then(|value| cql_cmp(value, rhs))
957+
.and_then(|value| cql_cmp(&value, rhs))
956958
.is_some_and(|ord| ord.is_lt()),
957959
Restriction::Lte { lhs, rhs } => primary_key_value(&primary_key, lhs)
958-
.and_then(|value| cql_cmp(value, rhs))
960+
.and_then(|value| cql_cmp(&value, rhs))
959961
.is_some_and(|ord| ord.is_le()),
960962
Restriction::Gt { lhs, rhs } => primary_key_value(&primary_key, lhs)
961-
.and_then(|value| cql_cmp(value, rhs))
963+
.and_then(|value| cql_cmp(&value, rhs))
962964
.is_some_and(|ord| ord.is_gt()),
963965
Restriction::Gte { lhs, rhs } => primary_key_value(&primary_key, lhs)
964-
.and_then(|value| cql_cmp(value, rhs))
966+
.and_then(|value| cql_cmp(&value, rhs))
965967
.is_some_and(|ord| ord.is_ge()),
966968
Restriction::EqTuple { lhs, rhs } => lhs
967969
.iter()
968970
.zip(rhs.iter())
969-
.all(|(lhs, rhs)| primary_key_value(&primary_key, lhs) == Some(rhs)),
971+
.all(|(lhs, rhs)| primary_key_value(&primary_key, lhs).as_ref() == Some(rhs)),
970972
Restriction::InTuple { lhs, rhs } => {
971973
let values: Vec<_> = lhs
972974
.iter()
@@ -976,7 +978,7 @@ fn filtered_ann(
976978
values
977979
.iter()
978980
.zip(rhs.iter())
979-
.all(|(value, rhs)| value == &Some(rhs))
981+
.all(|(value, rhs)| value.as_ref() == Some(rhs))
980982
})
981983
}
982984
Restriction::LtTuple { lhs, rhs } => {
@@ -1464,14 +1466,14 @@ mod tests {
14641466
values.into()
14651467
}
14661468

1467-
fn primary_key_value_fn<'a>(
1468-
columns: &'a [ColumnName],
1469-
) -> impl Fn(&'a PrimaryKey, &ColumnName) -> Option<&'a CqlValue> {
1470-
move |pk: &'a PrimaryKey, name: &ColumnName| {
1469+
fn primary_key_value_fn(
1470+
columns: &[ColumnName],
1471+
) -> impl Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue> + use<'_> {
1472+
move |pk: &PrimaryKey, name: &ColumnName| {
14711473
columns
14721474
.iter()
14731475
.position(|col| col == name)
1474-
.and_then(|idx| pk.0.get(idx))
1476+
.and_then(|idx| pk.get(idx))
14751477
}
14761478
}
14771479

0 commit comments

Comments
 (0)