Skip to content

Commit 7f76c37

Browse files
swasikswasik
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 7f76c37

File tree

11 files changed

+1018
-92
lines changed

11 files changed

+1018
-92
lines changed

‎crates/vector-store/src/db_index.rs‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,12 @@ use crate::DbEmbedding;
1010
use crate::IndexMetadata;
1111
use crate::KeyspaceName;
1212
use crate::Percentage;
13+
use crate::PrimaryKey;
1314
use crate::Progress;
1415
use crate::TableName;
1516
use crate::internals::Internals;
1617
use crate::internals::InternalsExt;
18+
use crate::invariant_key::InvariantKey;
1719
use crate::node_state::Event;
1820
use crate::node_state::NodeState;
1921
use crate::node_state::NodeStateExt;
@@ -451,6 +453,15 @@ impl Statements {
451453
.collect_vec(),
452454
);
453455

456+
anyhow::ensure!(
457+
primary_key_columns.len() <= InvariantKey::MAX_COLUMNS,
458+
"table {}.{} has {} primary key columns, but at most {} are supported",
459+
metadata.keyspace_name,
460+
metadata.table_name,
461+
primary_key_columns.len(),
462+
InvariantKey::MAX_COLUMNS,
463+
);
464+
454465
let table_columns = Arc::new(
455466
table
456467
.columns
@@ -752,7 +763,11 @@ impl Statements {
752763
else {
753764
return None;
754765
};
755-
let primary_key = primary_key.into();
766+
let primary_key = PrimaryKey::from(
767+
InvariantKey::try_new(primary_key)
768+
.inspect_err(|err| debug!("range_scan_stream: {err}"))
769+
.ok()?,
770+
);
756771

757772
Some(DbEmbedding {
758773
primary_key,
@@ -825,8 +840,8 @@ impl Consumer for CdcConsumer {
825840
"CDC error: primary key column {column} value should exist"
826841
))
827842
})
828-
.collect::<anyhow::Result<Vec<_>>>()?
829-
.into();
843+
.collect::<anyhow::Result<Vec<_>>>()?;
844+
let primary_key = PrimaryKey::from(InvariantKey::try_new(primary_key)?);
830845

831846
const HUNDREDS_NANOS_TO_MICROS: u64 = 10;
832847
let timestamp = (self.0.gregorian_epoch

‎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: 50 additions & 33 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 } => {
@@ -1083,6 +1085,7 @@ mod tests {
10831085
use crate::ExpansionSearch;
10841086
use crate::IndexId;
10851087
use crate::index::IndexExt;
1088+
use crate::invariant_key::InvariantKey;
10861089
use crate::memory;
10871090
use scylla::value::CqlValue;
10881091
use std::num::NonZeroUsize;
@@ -1107,7 +1110,7 @@ mod tests {
11071110
let id = worker * adds_per_worker + offset;
11081111
actor
11091112
.add(
1110-
vec![CqlValue::Int(id as i32)].into(),
1113+
InvariantKey::new(vec![CqlValue::Int(id as i32)]).into(),
11111114
vec![0.0f32; dimensions.get()].into(),
11121115
None,
11131116
)
@@ -1166,21 +1169,22 @@ mod tests {
11661169

11671170
actor
11681171
.add(
1169-
vec![CqlValue::Int(1), CqlValue::Text("one".to_string())].into(),
1172+
InvariantKey::new(vec![CqlValue::Int(1), CqlValue::Text("one".to_string())]).into(),
11701173
vec![1., 1., 1.].into(),
11711174
None,
11721175
)
11731176
.await;
11741177
actor
11751178
.add(
1176-
vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
1179+
InvariantKey::new(vec![CqlValue::Int(2), CqlValue::Text("two".to_string())]).into(),
11771180
vec![2., -2., 2.].into(),
11781181
None,
11791182
)
11801183
.await;
11811184
actor
11821185
.add(
1183-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1186+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1187+
.into(),
11841188
vec![3., 3., 3.].into(),
11851189
None,
11861190
)
@@ -1205,18 +1209,20 @@ mod tests {
12051209
assert_eq!(distances.len(), 1);
12061210
assert_eq!(
12071211
primary_keys.first().unwrap(),
1208-
&vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
1212+
&InvariantKey::new(vec![CqlValue::Int(2), CqlValue::Text("two".to_string())]).into(),
12091213
);
12101214

12111215
actor
12121216
.remove(
1213-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1217+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1218+
.into(),
12141219
None,
12151220
)
12161221
.await;
12171222
actor
12181223
.add(
1219-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1224+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1225+
.into(),
12201226
vec![2.1, -2.1, 2.1].into(),
12211227
None,
12221228
)
@@ -1233,7 +1239,8 @@ mod tests {
12331239
.0
12341240
.first()
12351241
.unwrap()
1236-
!= &vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into()
1242+
!= &InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1243+
.into()
12371244
{
12381245
task::yield_now().await;
12391246
}
@@ -1243,7 +1250,8 @@ mod tests {
12431250

12441251
actor
12451252
.remove(
1246-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1253+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1254+
.into(),
12471255
None,
12481256
)
12491257
.await;
@@ -1267,7 +1275,7 @@ mod tests {
12671275
assert_eq!(distances.len(), 1);
12681276
assert_eq!(
12691277
primary_keys.first().unwrap(),
1270-
&vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
1278+
&InvariantKey::new(vec![CqlValue::Int(2), CqlValue::Text("two".to_string())]).into(),
12711279
);
12721280
}
12731281

@@ -1302,7 +1310,11 @@ mod tests {
13021310
memory_rx
13031311
});
13041312
actor
1305-
.add(vec![CqlValue::Int(1)].into(), vec![1., 1., 1.].into(), None)
1313+
.add(
1314+
InvariantKey::new(vec![CqlValue::Int(1)]).into(),
1315+
vec![1., 1., 1.].into(),
1316+
None,
1317+
)
13061318
.await;
13071319
let mut memory_rx = memory_respond.await.unwrap();
13081320
assert_eq!(actor.count().await.unwrap(), 0);
@@ -1312,7 +1324,11 @@ mod tests {
13121324
_ = tx.send(Allocate::Can);
13131325
});
13141326
actor
1315-
.add(vec![CqlValue::Int(1)].into(), vec![1., 1., 1.].into(), None)
1327+
.add(
1328+
InvariantKey::new(vec![CqlValue::Int(1)]).into(),
1329+
vec![1., 1., 1.].into(),
1330+
None,
1331+
)
13161332
.await;
13171333
memory_respond.await.unwrap();
13181334

@@ -1457,21 +1473,22 @@ mod tests {
14571473

14581474
mod cql_cmp_tuple_tests {
14591475
use super::super::{ColumnName, PrimaryKey, cql_cmp_tuple};
1476+
use crate::invariant_key::InvariantKey;
14601477
use scylla::value::CqlValue;
14611478
use std::cmp::Ordering;
14621479

14631480
fn make_primary_key(values: Vec<CqlValue>) -> PrimaryKey {
1464-
values.into()
1481+
InvariantKey::new(values).into()
14651482
}
14661483

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| {
1484+
fn primary_key_value_fn(
1485+
columns: &[ColumnName],
1486+
) -> impl Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue> + use<'_> {
1487+
move |pk: &PrimaryKey, name: &ColumnName| {
14711488
columns
14721489
.iter()
14731490
.position(|col| col == name)
1474-
.and_then(|idx| pk.0.get(idx))
1491+
.and_then(|idx| pk.get(idx))
14751492
}
14761493
}
14771494

0 commit comments

Comments
 (0)