add save_path handling to bftree streaming path

Author: Jordan Maples

diskann-benchmark/src/index/bftree/full_precision_streaming.rs

@@ -16,7 +16,10 @@ use diskann_benchmark_runner::{
     Benchmark, Checkpoint,
 };
 use diskann_bftree::{BfTreeProvider, NoStore};
-use diskann_providers::model::graph::provider::async_::common::FullPrecision;
+use diskann_providers::{
+    model::graph::provider::async_::common::FullPrecision,
+    storage::{FileStorageProvider, SaveWith},
+};
 use diskann_utils::{
     sampling::WithApproximateNorm,
     views::{Matrix, MatrixView},
@@ -199,18 +202,36 @@ where
     ) -> anyhow::Result<Self::Output> {
         writeln!(output, "{}", input)?;
 
-        super::streaming_utils::run_streaming::<T, _>(
+        let mut index_for_save: Option<BfTreeFPIndex<T>> = None;
+
+        let results = super::streaming_utils::run_streaming::<T, _>(
             input.runbook_params(),
-            |max_points| bftree_streaming::<T>(input, max_points),
+            |max_points| {
+                let (streamer, index) = bftree_streaming::<T>(input, max_points)?;
+                index_for_save = Some(index);
+                Ok(streamer)
+            },
             output,
-        )
+        )?;
+
+        // save the index if requested
+        if let Some(save_path) = input.build().save_path() {
+            let index = index_for_save.expect("index should have been set by make_streamer");
+            crate::utils::tokio::block_on(
+                index
+                    .provider()
+                    .save_with(&FileStorageProvider, &save_path.to_string()),
+            )?;
+        }
+
+        Ok(results)
     }
 }
 
 fn bftree_streaming<T>(
     input: &BfTreeDynamicRun,
     max_points: usize,
-) -> anyhow::Result<bigann::WithData<T, u32, Managed<T, StreamStats>>>
+) -> anyhow::Result<(bigann::WithData<T, u32, Managed<T, StreamStats>>, BfTreeFPIndex<T>)>
 where
     T: bytemuck::Pod + VectorRepr + WithApproximateNorm + SampleableForStart,
 {
@@ -232,6 +253,7 @@ where
         .compute(data.as_view())?;
     let provider = BfTreeProvider::new(params, start_points.as_view(), NoStore)?;
     let index = Arc::new(DiskANNIndex::new(config, provider, None));
+    let index_handle = index.clone();
 
     let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap();
     let managed_stream = BfTreeStream {
@@ -258,5 +280,5 @@ where
         )?))
     });
 
-    Ok(layered)
+    Ok((layered, index_handle))
 }

diskann-benchmark/src/index/bftree/spherical_streaming.rs

@@ -19,7 +19,10 @@ use diskann_benchmark_runner::{
     Benchmark, Checkpoint,
 };
 use diskann_bftree::{quant::QuantVectorProvider, BfTreeProvider};
-use diskann_providers::model::graph::provider::async_::common::Quantized;
+use diskann_providers::{
+    model::graph::provider::async_::common::Quantized,
+    storage::{FileStorageProvider, SaveWith},
+};
 use diskann_quantization::alloc::{AllocatorError, GlobalAllocator, Poly};
 use diskann_quantization::spherical::{
     iface::{self as spherical_iface, Quantizer},
@@ -204,18 +207,36 @@ impl Benchmark for StreamingSpherical {
     ) -> anyhow::Result<Self::Output> {
         writeln!(output, "{}", input)?;
 
-        super::streaming_utils::run_streaming::<f32, _>(
+        let mut index_for_save: Option<BfTreeSQIndex> = None;
+
+        let results = super::streaming_utils::run_streaming::<f32, _>(
             input.runbook_params(),
-            |max_points| bftree_sq_streaming_impl(input, max_points),
+            |max_points| {
+                let (streamer, index) = bftree_sq_streaming_impl(input, max_points)?;
+                index_for_save = Some(index);
+                Ok(streamer)
+            },
             output,
-        )
+        )?;
+
+        // save the index if requested
+        if let Some(save_path) = input.build().save_path() {
+            let index = index_for_save.expect("index should have been set by make_streamer");
+            crate::utils::tokio::block_on(
+                index
+                    .provider()
+                    .save_with(&FileStorageProvider, &save_path.to_string()),
+            )?;
+        }
+
+        Ok(results)
     }
 }
 
 fn bftree_sq_streaming_impl(
     input: &BfTreeSphericalDynamicRun,
     max_points: usize,
-) -> anyhow::Result<bigann::WithData<f32, u32, Managed<f32, StreamStats>>> {
+) -> anyhow::Result<(bigann::WithData<f32, u32, Managed<f32, StreamStats>>, BfTreeSQIndex)> {
     let topk = match input.search_phase() {
         SearchPhase::Topk(topk) => topk,
         _ => anyhow::bail!("Only TopK is currently supported by the streaming index"),
@@ -257,6 +278,7 @@ fn bftree_sq_streaming_impl(
         .compute(data.as_view())?;
     let provider = BfTreeProvider::new(params, start_points.as_view(), quantizer_poly)?;
     let index = Arc::new(DiskANNIndex::new(config, provider, None));
+    let index_handle = index.clone();
 
     let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap();
     let managed_stream = BfTreeSQStream {
@@ -283,5 +305,5 @@ fn bftree_sq_streaming_impl(
         )?))
     });
 
-    Ok(layered)
+    Ok((layered, index_handle))
 }