Skip to content

Commit bb6a0a2

Browse files
Jordan MaplesCopilot
andcommitted
Use checked tag->slot conversion in direct streaming path
The direct (non-Managed) streaming path cast runbook tag IDs (usize) to u32 provider slot IDs with `as u32`, silently truncating on overflow and corrupting the ID mapping. Convert with a checked `u32::try_from` that errors on overflow, matching the inmem streaming path's behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74602942-339e-45b2-a3b8-18636427763a
1 parent ebf125f commit bb6a0a2

1 file changed

Lines changed: 15 additions & 4 deletions

File tree

  • diskann-benchmark/src/index/streaming

diskann-benchmark/src/index/streaming/runner.rs

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -223,10 +223,21 @@ where
223223
// Direct Stream (no ID management) //
224224
/////////////////////////////////////
225225

226+
/// Convert runbook tag IDs to `u32` provider slot IDs, erroring on overflow rather
227+
/// than silently truncating.
228+
fn tags_to_slots(tags: Range<usize>) -> anyhow::Result<Vec<u32>> {
229+
tags.map(|t| {
230+
u32::try_from(t).map_err(|_| {
231+
anyhow::anyhow!("runbook tag id {t} exceeds the u32 provider slot ID space")
232+
})
233+
})
234+
.collect()
235+
}
236+
226237
/// Direct [`streaming::Stream`] implementation for providers that manage their own IDs.
227238
///
228239
/// In this mode, the external tag IDs from the runbook are used directly as slot IDs
229-
/// (cast to `u32`). No [`super::Managed`] layer is needed.
240+
/// (checked conversion to `u32`). No [`super::Managed`] layer is needed.
230241
impl<DP, T, S, M> streaming::Stream<bigann::DataArgs<T, u32>> for StreamRunner<DP, T, S, M>
231242
where
232243
DP: DataProvider<Context: Default, ExternalId = u32, InternalId = u32>
@@ -274,20 +285,20 @@ where
274285
&mut self,
275286
(data, tags): (MatrixView<'_, T>, Range<usize>),
276287
) -> anyhow::Result<Self::Output> {
277-
let slots: Vec<u32> = tags.map(|t| t as u32).collect();
288+
let slots = tags_to_slots(tags)?;
278289
Ok(StreamStats::Insert(self.build(data, &slots)?))
279290
}
280291

281292
fn replace(
282293
&mut self,
283294
(data, tags): (MatrixView<'_, T>, Range<usize>),
284295
) -> anyhow::Result<Self::Output> {
285-
let slots: Vec<u32> = tags.map(|t| t as u32).collect();
296+
let slots = tags_to_slots(tags)?;
286297
Ok(StreamStats::Replace(self.build(data, &slots)?))
287298
}
288299

289300
fn delete(&mut self, tags: Range<usize>) -> anyhow::Result<Self::Output> {
290-
let slots: Vec<u32> = tags.map(|t| t as u32).collect();
301+
let slots = tags_to_slots(tags)?;
291302
let runner = InplaceDelete::new(
292303
self.index.clone(),
293304
self.strategy.clone(),

0 commit comments

Comments
 (0)