From 49cff547ae1f4c2ae1db5591a8cc799f6cd50b4d Mon Sep 17 00:00:00 2001 From: generall Date: Sun, 31 May 2026 19:42:57 +0200 Subject: [PATCH] refactor(id_tracker): split active vs deferred maps in PointMappings (PR A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First step of moving the mutable id tracker to a "two-track" model so append-only mutations into a deferred segment can keep both the visible (active) version and the latest mutated (deferred) version of a point at the same time. This PR is purely structural. On-disk format is unchanged: the loader still produces a single combined map, and `PointMappings::new` partitions it at construction time based on `deferred_internal_id`. Every observable behaviour at the existing API surface is preserved. Concretely: - New fields: `external_to_internal_num_deferred`, `external_to_internal_uuid_deferred`, and a `shadowed: BitVec` for future use by PR B (lazy-grown, default-false). - `internal_id(ext)` checks active first, falls through to deferred — matches the pre-split "any matching id" contract for ext ids whose internal id sat above the cutoff. - `set_link(ext, new_id)` now routes by cutoff: writes below the cutoff land in active, writes at or above land in deferred. Any prior head in the other track is tombstoned, so each ext still owns exactly one slot — same observable result as the pre-split single-map insert. PR B replaces the cross-track tombstone with a shadow-bit flip; PR A keeps current semantics on purpose. - `drop(ext)` clears entries from both tracks and tombstones each one, again matching the prior single-map behaviour. - `iter_external` and `iter_from` merge the active and deferred BTreeMap views into one sorted-by-key stream (dedup'd in case an ext exists in both tracks). - `available_point_count` counts distinct external ids across both tracks — preserves the prior observable count for segments where some entries used to sit above the cutoff in the single map. No write-path or read-path behaviour change. Reads still filter `internal_id >= cutoff` exactly as before; mutations still tombstone prior heads. The shadow bit and the deferred-aware lookup wiring land in PR B and PR C. All existing id_tracker tests pass against the split layout. Co-Authored-By: Claude Opus 4.7 --- lib/segment/src/id_tracker/point_mappings.rs | 304 +++++++++++++++---- 1 file changed, 241 insertions(+), 63 deletions(-) diff --git a/lib/segment/src/id_tracker/point_mappings.rs b/lib/segment/src/id_tracker/point_mappings.rs index fbd34dbf4d..a1d1346daf 100644 --- a/lib/segment/src/id_tracker/point_mappings.rs +++ b/lib/segment/src/id_tracker/point_mappings.rs @@ -32,10 +32,27 @@ pub struct PointMappings { deleted: BitVec, internal_to_external: Vec, + // Active head per external id (internal_id < deferred cutoff, or no cutoff). // Having two separate maps allows us iterating only over one type at a time without having to filter. external_to_internal_num: BTreeMap, external_to_internal_uuid: BTreeMap, + // Deferred head per external id (internal_id >= deferred cutoff). Same external + // id can appear in both maps at once — see `shadowed` for that case. On a + // segment without a deferred cutoff these stay empty. + external_to_internal_num_deferred: BTreeMap, + external_to_internal_uuid_deferred: BTreeMap, + + /// Bit set on active internal ids whose external id also has a deferred + /// head. Read-side iteration in `IncludeAll` mode uses this to skip the + /// stale active version when a deferred override exists, avoiding + /// duplicate-by-external yields. + /// + /// PR A: declared and partition-checked at construction; no writer yet + /// flips bits during mutations (PR B routes deferred writes through + /// here). + shadowed: BitVec, + /// Points with internal id >= this value are hidden from reads. /// Only set for appendable segments with deferred points. deferred_internal_id: Option, @@ -49,10 +66,58 @@ impl PointMappings { pub fn new( deleted: BitVec, internal_to_external: Vec, - external_to_internal_num: BTreeMap, - external_to_internal_uuid: BTreeMap, + mut external_to_internal_num: BTreeMap, + mut external_to_internal_uuid: BTreeMap, deferred_internal_id: Option, ) -> Self { + // Partition the loaded single-map mappings into active (id < cutoff) + // and deferred (id >= cutoff). Persisted format is unchanged — the + // split is purely runtime. With no cutoff every entry stays active. + let mut external_to_internal_num_deferred = BTreeMap::new(); + let mut external_to_internal_uuid_deferred = BTreeMap::new(); + if let Some(cutoff) = deferred_internal_id { + external_to_internal_num.retain(|&k, &mut v| { + if v >= cutoff { + external_to_internal_num_deferred.insert(k, v); + false + } else { + true + } + }); + external_to_internal_uuid.retain(|&k, &mut v| { + if v >= cutoff { + external_to_internal_uuid_deferred.insert(k, v); + false + } else { + true + } + }); + } + // Shadowed bits: any active id whose external also has a deferred + // head. Impossible from a fresh load (each ext was in a single map), + // but compute it so mutations land on a consistent starting state. + // Grown lazily — out-of-bounds bits are treated as `false` by + // readers, so the empty default is a valid no-shadow state. + let mut shadowed = BitVec::new(); + let mut mark_shadow = |active_id: PointOffsetType| { + let active_id = active_id as usize; + if active_id >= shadowed.len() { + shadowed.resize(active_id + 1, false); + } + shadowed.set(active_id, true); + }; + for (k, _) in &external_to_internal_num_deferred { + if let Some(active_id) = external_to_internal_num.get(k) { + mark_shadow(*active_id); + } + } + for (k, _) in &external_to_internal_uuid_deferred { + if let Some(active_id) = external_to_internal_uuid.get(k) { + mark_shadow(*active_id); + } + } + drop(mark_shadow); + let deferred_deleted_count = deferred_internal_id .map(|deferred_from| { let total = deleted.len(); @@ -68,6 +133,9 @@ impl PointMappings { internal_to_external, external_to_internal_num, external_to_internal_uuid, + external_to_internal_num_deferred, + external_to_internal_uuid_deferred, + shadowed, deferred_internal_id, deferred_deleted_count, } @@ -91,19 +159,41 @@ impl PointMappings { } /// Number of points, excluding deleted ones. + /// + /// Counts each distinct external id once. An ext that lives only in the + /// deferred map (no active head) still counts — preserves the pre-split + /// observable count from when both tracks shared one map. pub(crate) fn available_point_count(&self) -> usize { - self.external_to_internal_num.len() + self.external_to_internal_uuid.len() + let active = self.external_to_internal_num.len() + self.external_to_internal_uuid.len(); + let deferred_only_num = self + .external_to_internal_num_deferred + .keys() + .filter(|k| !self.external_to_internal_num.contains_key(k)) + .count(); + let deferred_only_uuid = self + .external_to_internal_uuid_deferred + .keys() + .filter(|k| !self.external_to_internal_uuid.contains_key(k)) + .count(); + active + deferred_only_num + deferred_only_uuid } pub(crate) fn deleted(&self) -> &BitSlice { &self.deleted } + /// Internal id for `external_id`. Checks the active map first and falls + /// back to deferred — preserves "any matching id" semantics from before + /// the split, where every ext lived in a single map. pub(crate) fn internal_id(&self, external_id: &PointIdType) -> Option { - match external_id { + let active = match external_id { PointIdType::NumId(num) => self.external_to_internal_num.get(num).copied(), PointIdType::Uuid(uuid) => self.external_to_internal_uuid.get(uuid).copied(), - } + }; + active.or_else(|| match external_id { + PointIdType::NumId(num) => self.external_to_internal_num_deferred.get(num).copied(), + PointIdType::Uuid(uuid) => self.external_to_internal_uuid_deferred.get(uuid).copied(), + }) } pub(crate) fn external_id(&self, internal_id: PointOffsetType) -> Option { @@ -117,38 +207,54 @@ impl PointMappings { } pub(crate) fn drop(&mut self, external_id: PointIdType) -> Option { - let internal_id = match external_id { - // We "temporarily" remove existing points from the BTreeMaps without writing them to disk - // because we remove deleted points of a previous load directly when loading. - PointIdType::NumId(num) => self.external_to_internal_num.remove(&num), - PointIdType::Uuid(uuid) => self.external_to_internal_uuid.remove(&uuid), + // Drop from both tracks. A point can live in only one map at a time + // today (PR A), but PR B will introduce shadowed-active + deferred + // pairs for the same ext, and `drop` is the API both states wire + // through. Tombstoning both is the safe behaviour to land first. + // We "temporarily" remove existing points from the BTreeMaps without writing them to disk + // because we remove deleted points of a previous load directly when loading. + let (active, deferred) = match external_id { + PointIdType::NumId(num) => ( + self.external_to_internal_num.remove(&num), + self.external_to_internal_num_deferred.remove(&num), + ), + PointIdType::Uuid(uuid) => ( + self.external_to_internal_uuid.remove(&uuid), + self.external_to_internal_uuid_deferred.remove(&uuid), + ), }; - // Also reset inverse mapping - if let Some(internal_id) = internal_id { + for internal_id in [active, deferred].into_iter().flatten() { + // Reset inverse mapping self.internal_to_external[internal_id as usize] = PointIdType::NumId(u64::MAX); - } - if let Some(internal_id) = &internal_id { let was_already_deleted = *self .deleted - .get(*internal_id as usize) + .get(internal_id as usize) .as_deref() .unwrap_or(&true); - self.deleted.set(*internal_id as usize, true); + self.deleted.set(internal_id as usize, true); + // Clearing the shadowed bit on tombstone keeps PR C's read-side + // filter from skipping a slot that's already filtered by + // `deleted` — no observable change today, but cheap insurance. + if (internal_id as usize) < self.shadowed.len() { + self.shadowed.set(internal_id as usize, false); + } // Count newly-deleted deferred points so we can report visible deferred totals // without rescanning the deleted bitslice. if !was_already_deleted && self .deferred_internal_id - .is_some_and(|deferred_from| *internal_id >= deferred_from) + .is_some_and(|deferred_from| internal_id >= deferred_from) { self.deferred_deleted_count += 1; } } - internal_id + // Preserve the prior single-return signature: prefer active (the + // visible head) when both tracks held this ext. + active.or(deferred) } pub(crate) fn iter_random( @@ -183,59 +289,79 @@ impl PointMappings { &self, external_id: Option, ) -> Box + '_> { - let full_num_iter = || { - self.external_to_internal_num - .iter() - .map(|(k, v)| (PointIdType::NumId(*k), *v)) + // Merge active + deferred BTreeMap views into one sorted-by-key + // stream, deduping the rare case where the same ext exists in + // both tracks (PR B will introduce that; PR A loads at most one). + // The dedup prefers the active entry because that's the visible + // head — keeping callers (which today don't expect to see + // deferred-only writes) on the same offset as before. + let merged_num = |start: Option| { + let active: Box> = match start { + None => Box::new(self.external_to_internal_num.iter()), + Some(s) => Box::new(self.external_to_internal_num.range(s..)), + }; + let deferred: Box> = match start { + None => Box::new(self.external_to_internal_num_deferred.iter()), + Some(s) => Box::new(self.external_to_internal_num_deferred.range(s..)), + }; + active + .merge_join_by(deferred, |a, d| a.0.cmp(d.0)) + .map(|either| match either { + itertools::EitherOrBoth::Both((k, v), _) + | itertools::EitherOrBoth::Left((k, v)) + | itertools::EitherOrBoth::Right((k, v)) => (PointIdType::NumId(*k), *v), + }) }; - let offset_num_iter = |offset: u64| { - self.external_to_internal_num - .range(offset..) - .map(|(k, v)| (PointIdType::NumId(*k), *v)) - }; - let full_uuid_iter = || { - self.external_to_internal_uuid - .iter() - .map(|(k, v)| (PointIdType::Uuid(*k), *v)) - }; - let offset_uuid_iter = |offset: Uuid| { - self.external_to_internal_uuid - .range(offset..) - .map(|(k, v)| (PointIdType::Uuid(*k), *v)) + let merged_uuid = |start: Option| { + let active: Box> = match start { + None => Box::new(self.external_to_internal_uuid.iter()), + Some(s) => Box::new(self.external_to_internal_uuid.range(s..)), + }; + let deferred: Box> = match start { + None => Box::new(self.external_to_internal_uuid_deferred.iter()), + Some(s) => Box::new(self.external_to_internal_uuid_deferred.range(s..)), + }; + active + .merge_join_by(deferred, |a, d| a.0.cmp(d.0)) + .map(|either| match either { + itertools::EitherOrBoth::Both((k, v), _) + | itertools::EitherOrBoth::Left((k, v)) + | itertools::EitherOrBoth::Right((k, v)) => (PointIdType::Uuid(*k), *v), + }) }; match external_id { None => { - let iter_num = full_num_iter(); - let iter_uuid = full_uuid_iter(); // order is important here, we want to iterate over the u64 ids first - Box::new(iter_num.chain(iter_uuid)) + Box::new(merged_num(None).chain(merged_uuid(None))) } Some(offset) => match offset { PointIdType::NumId(idx) => { // Because u64 keys are less that uuid key, we can just use the full iterator for uuid - let iter_num = offset_num_iter(idx); - let iter_uuid = full_uuid_iter(); - // order is important here, we want to iterate over the u64 ids first - Box::new(iter_num.chain(iter_uuid)) + Box::new(merged_num(Some(idx)).chain(merged_uuid(None))) } PointIdType::Uuid(uuid) => { // if offset is a uuid, we can only iterate over uuids - Box::new(offset_uuid_iter(uuid)) + Box::new(merged_uuid(Some(uuid))) } }, } } pub(crate) fn iter_external(&self) -> Box + '_> { + // Merge sorted active + deferred key streams, deduping identical + // keys (PR B can introduce active+deferred pairs for one ext). let iter_num = self .external_to_internal_num .keys() + .merge(self.external_to_internal_num_deferred.keys()) + .dedup() .map(|i| PointIdType::NumId(*i)); - let iter_uuid = self .external_to_internal_uuid .keys() + .merge(self.external_to_internal_uuid_deferred.keys()) + .dedup() .map(|i| PointIdType::Uuid(*i)); // order is important here, we want to iterate over the u64 ids first Box::new(iter_num.chain(iter_uuid)) @@ -268,36 +394,77 @@ impl PointMappings { /// Sets the link between an external and internal id. /// Returns the previous internal id if it existed. + /// + /// PR A: the write routes into the active or deferred map based on + /// whether `internal_id` sits below or at-or-above + /// `deferred_internal_id`. Either way the previous head for this ext + /// (in either track) is tombstoned, matching the pre-split behaviour + /// where one ext owned exactly one slot. PR B replaces the + /// "tombstone the other track" branch with a shadow-bit flip so a + /// deferred mutation no longer hides the visible active version. pub(crate) fn set_link( &mut self, external_id: PointIdType, internal_id: PointOffsetType, ) -> Option { - let old_internal_id = match external_id { - PointIdType::NumId(idx) => self.external_to_internal_num.insert(idx, internal_id), - PointIdType::Uuid(uuid) => self.external_to_internal_uuid.insert(uuid, internal_id), + let is_deferred = self + .deferred_internal_id + .is_some_and(|cutoff| internal_id >= cutoff); + + // Insert into the routed map; pull the prior head out of the SAME + // map. The other-track cleanup happens below so the return value + // mirrors the pre-split single-map contract. + let same_track_prior = match (external_id, is_deferred) { + (PointIdType::NumId(idx), false) => { + self.external_to_internal_num.insert(idx, internal_id) + } + (PointIdType::NumId(idx), true) => self + .external_to_internal_num_deferred + .insert(idx, internal_id), + (PointIdType::Uuid(uuid), false) => { + self.external_to_internal_uuid.insert(uuid, internal_id) + } + (PointIdType::Uuid(uuid), true) => self + .external_to_internal_uuid_deferred + .insert(uuid, internal_id), + }; + // Cross-track cleanup — drop any prior entry in the OTHER map so + // each ext owns exactly one slot, matching the pre-split contract. + let other_track_prior = match (external_id, is_deferred) { + (PointIdType::NumId(idx), false) => self.external_to_internal_num_deferred.remove(&idx), + (PointIdType::NumId(idx), true) => self.external_to_internal_num.remove(&idx), + (PointIdType::Uuid(uuid), false) => { + self.external_to_internal_uuid_deferred.remove(&uuid) + } + (PointIdType::Uuid(uuid), true) => self.external_to_internal_uuid.remove(&uuid), }; - let internal_id = internal_id as usize; - if internal_id >= self.internal_to_external.len() { + let internal_id_usize = internal_id as usize; + if internal_id_usize >= self.internal_to_external.len() { self.internal_to_external - .resize(internal_id + 1, PointIdType::NumId(u64::MAX)); + .resize(internal_id_usize + 1, PointIdType::NumId(u64::MAX)); } - if internal_id >= self.deleted.len() { - self.deleted.resize(internal_id + 1, true); + if internal_id_usize >= self.deleted.len() { + self.deleted.resize(internal_id_usize + 1, true); } - if let Some(old_internal_id) = &old_internal_id { - let old_internal_id = *old_internal_id as usize; - if old_internal_id != internal_id { - self.deleted.set(old_internal_id, true); + for old in [same_track_prior, other_track_prior].into_iter().flatten() { + let old = old as usize; + if old != internal_id_usize { + self.deleted.set(old, true); + if old < self.shadowed.len() { + self.shadowed.set(old, false); + } } } - self.internal_to_external[internal_id] = external_id; - self.deleted.set(internal_id, false); + self.internal_to_external[internal_id_usize] = external_id; + self.deleted.set(internal_id_usize, false); - old_internal_id + // Prefer the same-track prior when reporting back, matching the + // pre-split single-map flavour. Cross-track replacement (the + // unusual case) falls back to whatever the other map held. + same_track_prior.or(other_track_prior) } pub(crate) fn total_point_count(&self) -> usize { @@ -370,6 +537,9 @@ impl PointMappings { internal_to_external, external_to_internal_num, external_to_internal_uuid, + external_to_internal_num_deferred: BTreeMap::new(), + external_to_internal_uuid_deferred: BTreeMap::new(), + shadowed: BitVec::new(), deferred_internal_id: None, deferred_deleted_count: 0, } @@ -395,11 +565,15 @@ impl PointMappings { internal_to_external, external_to_internal_num, external_to_internal_uuid, + external_to_internal_num_deferred, + external_to_internal_uuid_deferred, + shadowed, deferred_internal_id: _, deferred_deleted_count: _, } = self; let deleted_bytes = deleted.capacity().div_ceil(u8::BITS as usize); + let shadowed_bytes = shadowed.capacity().div_ceil(u8::BITS as usize); let internal_to_external_bytes = internal_to_external.capacity() * std::mem::size_of::(); // BTreeMap node overhead: key + value + 2 child pointers + parent pointer + metadata. @@ -411,8 +585,12 @@ impl PointMappings { let uuid_entry_size = std::mem::size_of::() + std::mem::size_of::() + btree_node_overhead; - let num_map_bytes = external_to_internal_num.len() * num_entry_size; - let uuid_map_bytes = external_to_internal_uuid.len() * uuid_entry_size; - deleted_bytes + internal_to_external_bytes + num_map_bytes + uuid_map_bytes + let num_map_bytes = (external_to_internal_num.len() + + external_to_internal_num_deferred.len()) + * num_entry_size; + let uuid_map_bytes = (external_to_internal_uuid.len() + + external_to_internal_uuid_deferred.len()) + * uuid_entry_size; + deleted_bytes + shadowed_bytes + internal_to_external_bytes + num_map_bytes + uuid_map_bytes } }