Skip to content

Commit c970ffa

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 ba15a46 commit c970ffa

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
@@ -647,16 +647,20 @@ async fn post_index_ann(
647647
let primary_keys: anyhow::Result<_> = primary_keys
648648
.iter()
649649
.map(|primary_key| {
650-
if primary_key.0.len() != primary_key_columns.len() {
650+
if primary_key.len() != primary_key_columns.len() {
651651
bail!(
652652
"wrong size of a primary key: {}, {}",
653653
primary_key_columns.len(),
654-
primary_key.0.len()
654+
primary_key.len()
655655
);
656656
}
657657
Ok(primary_key)
658658
})
659-
.map_ok(|primary_key| primary_key.0[idx_column].clone())
659+
.map_ok(|primary_key| {
660+
primary_key
661+
.get(idx_column)
662+
.expect("primary key index out of bounds after length check")
663+
})
660664
.map_ok(try_to_json)
661665
.map(|primary_key| primary_key.flatten())
662666
.collect();

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

Lines changed: 50 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -944,15 +944,15 @@ fn cql_cmp(lhs: &CqlValue, rhs: &CqlValue) -> Option<std::cmp::Ordering> {
944944

945945
/// Lexicographically compare tuple values.
946946
/// Returns the ordering of the first non-equal pair, or Equal if all pairs are equal.
947-
fn cql_cmp_tuple<'a>(
948-
primary_key: &'a PrimaryKey,
949-
primary_key_value: impl Fn(&'a PrimaryKey, &ColumnName) -> Option<&'a CqlValue>,
947+
fn cql_cmp_tuple(
948+
primary_key: &PrimaryKey,
949+
primary_key_value: impl Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue>,
950950
lhs: &[ColumnName],
951951
rhs: &[CqlValue],
952952
) -> Option<std::cmp::Ordering> {
953953
for (col, rhs_val) in lhs.iter().zip(rhs.iter()) {
954954
let lhs_val = primary_key_value(primary_key, col)?;
955-
match cql_cmp(lhs_val, rhs_val)? {
955+
match cql_cmp(&lhs_val, rhs_val)? {
956956
std::cmp::Ordering::Equal => continue,
957957
other => return Some(other),
958958
}
@@ -971,17 +971,17 @@ fn filtered_ann(
971971
) {
972972
fn annotate<F>(f: F) -> F
973973
where
974-
F: for<'a, 'b> Fn(&'a PrimaryKey, &'b ColumnName) -> Option<&'a CqlValue>,
974+
F: Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue>,
975975
{
976976
f
977977
}
978978

979979
let primary_key_value = annotate(
980-
|primary_key: &PrimaryKey, name: &ColumnName| -> Option<&CqlValue> {
980+
|primary_key: &PrimaryKey, name: &ColumnName| -> Option<CqlValue> {
981981
primary_key_columns
982982
.iter()
983983
.position(|key_column| key_column == name)
984-
.and_then(move |idx| primary_key.0.get(idx))
984+
.and_then(move |idx| primary_key.get(idx))
985985
},
986986
);
987987

@@ -993,27 +993,29 @@ fn filtered_ann(
993993
.restrictions
994994
.iter()
995995
.all(|restriction| match restriction {
996-
Restriction::Eq { lhs, rhs } => primary_key_value(&primary_key, lhs) == Some(rhs),
996+
Restriction::Eq { lhs, rhs } => {
997+
primary_key_value(&primary_key, lhs).as_ref() == Some(rhs)
998+
}
997999
Restriction::In { lhs, rhs } => {
9981000
let value = primary_key_value(&primary_key, lhs);
999-
rhs.iter().any(|rhs| value == Some(rhs))
1001+
rhs.iter().any(|rhs| value.as_ref() == Some(rhs))
10001002
}
10011003
Restriction::Lt { lhs, rhs } => primary_key_value(&primary_key, lhs)
1002-
.and_then(|value| cql_cmp(value, rhs))
1004+
.and_then(|value| cql_cmp(&value, rhs))
10031005
.is_some_and(|ord| ord.is_lt()),
10041006
Restriction::Lte { lhs, rhs } => primary_key_value(&primary_key, lhs)
1005-
.and_then(|value| cql_cmp(value, rhs))
1007+
.and_then(|value| cql_cmp(&value, rhs))
10061008
.is_some_and(|ord| ord.is_le()),
10071009
Restriction::Gt { lhs, rhs } => primary_key_value(&primary_key, lhs)
1008-
.and_then(|value| cql_cmp(value, rhs))
1010+
.and_then(|value| cql_cmp(&value, rhs))
10091011
.is_some_and(|ord| ord.is_gt()),
10101012
Restriction::Gte { lhs, rhs } => primary_key_value(&primary_key, lhs)
1011-
.and_then(|value| cql_cmp(value, rhs))
1013+
.and_then(|value| cql_cmp(&value, rhs))
10121014
.is_some_and(|ord| ord.is_ge()),
10131015
Restriction::EqTuple { lhs, rhs } => lhs
10141016
.iter()
10151017
.zip(rhs.iter())
1016-
.all(|(lhs, rhs)| primary_key_value(&primary_key, lhs) == Some(rhs)),
1018+
.all(|(lhs, rhs)| primary_key_value(&primary_key, lhs).as_ref() == Some(rhs)),
10171019
Restriction::InTuple { lhs, rhs } => {
10181020
let values: Vec<_> = lhs
10191021
.iter()
@@ -1023,7 +1025,7 @@ fn filtered_ann(
10231025
values
10241026
.iter()
10251027
.zip(rhs.iter())
1026-
.all(|(value, rhs)| value == &Some(rhs))
1028+
.all(|(value, rhs)| value.as_ref() == Some(rhs))
10271029
})
10281030
}
10291031
Restriction::LtTuple { lhs, rhs } => {
@@ -1132,6 +1134,7 @@ mod tests {
11321134
use crate::ExpansionSearch;
11331135
use crate::IndexId;
11341136
use crate::index::IndexExt;
1137+
use crate::invariant_key::InvariantKey;
11351138
use crate::memory;
11361139
use scylla::value::CqlValue;
11371140
use std::num::NonZeroUsize;
@@ -1156,7 +1159,7 @@ mod tests {
11561159
let id = worker * adds_per_worker + offset;
11571160
actor
11581161
.add(
1159-
vec![CqlValue::Int(id as i32)].into(),
1162+
InvariantKey::new(vec![CqlValue::Int(id as i32)]).into(),
11601163
vec![0.0f32; dimensions.get()].into(),
11611164
None,
11621165
)
@@ -1215,21 +1218,22 @@ mod tests {
12151218

12161219
actor
12171220
.add(
1218-
vec![CqlValue::Int(1), CqlValue::Text("one".to_string())].into(),
1221+
InvariantKey::new(vec![CqlValue::Int(1), CqlValue::Text("one".to_string())]).into(),
12191222
vec![1., 1., 1.].into(),
12201223
None,
12211224
)
12221225
.await;
12231226
actor
12241227
.add(
1225-
vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
1228+
InvariantKey::new(vec![CqlValue::Int(2), CqlValue::Text("two".to_string())]).into(),
12261229
vec![2., -2., 2.].into(),
12271230
None,
12281231
)
12291232
.await;
12301233
actor
12311234
.add(
1232-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1235+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1236+
.into(),
12331237
vec![3., 3., 3.].into(),
12341238
None,
12351239
)
@@ -1254,18 +1258,20 @@ mod tests {
12541258
assert_eq!(distances.len(), 1);
12551259
assert_eq!(
12561260
primary_keys.first().unwrap(),
1257-
&vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
1261+
&InvariantKey::new(vec![CqlValue::Int(2), CqlValue::Text("two".to_string())]).into(),
12581262
);
12591263

12601264
actor
12611265
.remove(
1262-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1266+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1267+
.into(),
12631268
None,
12641269
)
12651270
.await;
12661271
actor
12671272
.add(
1268-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1273+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1274+
.into(),
12691275
vec![2.1, -2.1, 2.1].into(),
12701276
None,
12711277
)
@@ -1282,7 +1288,8 @@ mod tests {
12821288
.0
12831289
.first()
12841290
.unwrap()
1285-
!= &vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into()
1291+
!= &InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1292+
.into()
12861293
{
12871294
task::yield_now().await;
12881295
}
@@ -1292,7 +1299,8 @@ mod tests {
12921299

12931300
actor
12941301
.remove(
1295-
vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
1302+
InvariantKey::new(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())])
1303+
.into(),
12961304
None,
12971305
)
12981306
.await;
@@ -1316,7 +1324,7 @@ mod tests {
13161324
assert_eq!(distances.len(), 1);
13171325
assert_eq!(
13181326
primary_keys.first().unwrap(),
1319-
&vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
1327+
&InvariantKey::new(vec![CqlValue::Int(2), CqlValue::Text("two".to_string())]).into(),
13201328
);
13211329
}
13221330

@@ -1351,7 +1359,11 @@ mod tests {
13511359
memory_rx
13521360
});
13531361
actor
1354-
.add(vec![CqlValue::Int(1)].into(), vec![1., 1., 1.].into(), None)
1362+
.add(
1363+
InvariantKey::new(vec![CqlValue::Int(1)]).into(),
1364+
vec![1., 1., 1.].into(),
1365+
None,
1366+
)
13551367
.await;
13561368
let mut memory_rx = memory_respond.await.unwrap();
13571369
assert_eq!(actor.count().await.unwrap(), 0);
@@ -1361,7 +1373,11 @@ mod tests {
13611373
_ = tx.send(Allocate::Can);
13621374
});
13631375
actor
1364-
.add(vec![CqlValue::Int(1)].into(), vec![1., 1., 1.].into(), None)
1376+
.add(
1377+
InvariantKey::new(vec![CqlValue::Int(1)]).into(),
1378+
vec![1., 1., 1.].into(),
1379+
None,
1380+
)
13651381
.await;
13661382
memory_respond.await.unwrap();
13671383

@@ -1506,21 +1522,22 @@ mod tests {
15061522

15071523
mod cql_cmp_tuple_tests {
15081524
use super::super::{ColumnName, PrimaryKey, cql_cmp_tuple};
1525+
use crate::invariant_key::InvariantKey;
15091526
use scylla::value::CqlValue;
15101527
use std::cmp::Ordering;
15111528

15121529
fn make_primary_key(values: Vec<CqlValue>) -> PrimaryKey {
1513-
values.into()
1530+
InvariantKey::new(values).into()
15141531
}
15151532

1516-
fn primary_key_value_fn<'a>(
1517-
columns: &'a [ColumnName],
1518-
) -> impl Fn(&'a PrimaryKey, &ColumnName) -> Option<&'a CqlValue> {
1519-
move |pk: &'a PrimaryKey, name: &ColumnName| {
1533+
fn primary_key_value_fn(
1534+
columns: &[ColumnName],
1535+
) -> impl Fn(&PrimaryKey, &ColumnName) -> Option<CqlValue> + use<'_> {
1536+
move |pk: &PrimaryKey, name: &ColumnName| {
15201537
columns
15211538
.iter()
15221539
.position(|col| col == name)
1523-
.and_then(|idx| pk.0.get(idx))
1540+
.and_then(|idx| pk.get(idx))
15241541
}
15251542
}
15261543

0 commit comments

Comments
 (0)