Generalize vector flush stats (#25811)

This removes the vector specific statistics from the generalized search flusher. Still a little awkward, but it's something.

GitOrigin-RevId: 030a62bff75b7dd3ea99f2cc39d0c78e202e7e4d

Author: sjudd

crates/database/src/index_workers/index_meta.rs

@@ -44,7 +44,10 @@ use vector::{
     QdrantSchema,
 };
 
-use crate::Snapshot;
+use crate::{
+    metrics::vector::log_documents_per_segment,
+    Snapshot,
+};
 
 pub trait SearchIndexConfigParser {
     type IndexType: SearchIndex;
@@ -127,10 +130,11 @@ pub trait SearchIndex {
 
     type PreviousSegments;
 
+    type Statistics: SegmentStatistics;
+
     type Schema: Send + Sync;
-    // TODO(sam): Convert these to something more like segment statistics
-    fn num_vectors(segment: &Self::Segment) -> u32;
-    fn non_deleted_vectors(segment: &Self::Segment) -> anyhow::Result<u64>;
+
+    fn statistics(segment: &Self::Segment) -> anyhow::Result<Self::Statistics>;
 
     async fn upload_new_segment<RT: Runtime>(
         rt: &RT,
@@ -170,6 +174,11 @@ pub trait SearchIndex {
     ) -> anyhow::Result<Vec<Self::Segment>>;
 }
 
+pub trait SegmentStatistics: Default {
+    fn add(lhs: anyhow::Result<Self>, rhs: anyhow::Result<Self>) -> anyhow::Result<Self>;
+    fn log(&self);
+}
+
 pub struct TextSearchIndex;
 #[async_trait]
 impl SearchIndex for TextSearchIndex {
@@ -178,6 +187,7 @@ impl SearchIndex for TextSearchIndex {
     type PreviousSegments = ();
     type Schema = ();
     type Segment = FragmentedSearchSegment;
+    type Statistics = TextStatistics;
 
     fn get_index_sizes(snapshot: Snapshot) -> anyhow::Result<BTreeMap<IndexId, usize>> {
         Ok(snapshot
@@ -230,19 +240,26 @@ impl SearchIndex for TextSearchIndex {
         anyhow::bail!("Not implemented")
     }
 
-    fn num_vectors(_segment: &Self::Segment) -> u32 {
-        0
-    }
-
     fn segment_id(_segment: &Self::Segment) -> String {
         "".to_string()
     }
 
-    fn non_deleted_vectors(_segment: &Self::Segment) -> anyhow::Result<u64> {
-        anyhow::bail!("Not implemented")
+    fn statistics(_segment: &Self::Segment) -> anyhow::Result<Self::Statistics> {
+        Ok(TextStatistics)
     }
 }
 
+#[derive(Default)]
+pub struct TextStatistics;
+
+impl SegmentStatistics for TextStatistics {
+    fn add(_: anyhow::Result<Self>, _: anyhow::Result<Self>) -> anyhow::Result<Self> {
+        Ok(Self)
+    }
+
+    fn log(&self) {}
+}
+
 #[derive(Debug)]
 pub struct VectorSearchIndex;
 
@@ -253,6 +270,7 @@ impl SearchIndex for VectorSearchIndex {
     type PreviousSegments = Vec<MutableFragmentedSegmentMetadata>;
     type Schema = QdrantSchema;
     type Segment = FragmentedVectorSegment;
+    type Statistics = VectorStatistics;
 
     fn get_index_sizes(snapshot: Snapshot) -> anyhow::Result<BTreeMap<IndexId, usize>> {
         Ok(snapshot
@@ -322,16 +340,16 @@ impl SearchIndex for VectorSearchIndex {
         upload_segment(rt, storage, new_segment).await
     }
 
-    fn num_vectors(segment: &Self::Segment) -> u32 {
-        segment.num_vectors
-    }
-
     fn segment_id(segment: &Self::Segment) -> String {
         segment.id.clone()
     }
 
-    fn non_deleted_vectors(segment: &Self::Segment) -> anyhow::Result<u64> {
-        segment.non_deleted_vectors()
+    fn statistics(segment: &Self::Segment) -> anyhow::Result<Self::Statistics> {
+        let non_deleted_vectors = segment.non_deleted_vectors()?;
+        Ok(VectorStatistics {
+            non_deleted_vectors,
+            num_vectors: segment.num_vectors,
+        })
     }
 }
 pub struct SearchIndexConfig<T: SearchIndex> {
@@ -350,6 +368,27 @@ pub struct BackfillState<T: SearchIndex> {
     pub backfill_snapshot_ts: Option<Timestamp>,
 }
 
+#[derive(Debug, Default)]
+pub struct VectorStatistics {
+    pub num_vectors: u32,
+    pub non_deleted_vectors: u64,
+}
+
+impl SegmentStatistics for VectorStatistics {
+    fn add(lhs: anyhow::Result<Self>, rhs: anyhow::Result<Self>) -> anyhow::Result<Self> {
+        let rhs = rhs?;
+        let lhs = lhs?;
+        Ok(Self {
+            num_vectors: lhs.num_vectors + rhs.num_vectors,
+            non_deleted_vectors: lhs.non_deleted_vectors + rhs.non_deleted_vectors,
+        })
+    }
+
+    fn log(&self) {
+        log_documents_per_segment(self.non_deleted_vectors);
+    }
+}
+
 impl From<VectorIndexBackfillState> for BackfillState<VectorSearchIndex> {
     fn from(value: VectorIndexBackfillState) -> Self {
         Self {

crates/database/src/index_workers/search_flusher.rs

@@ -50,12 +50,12 @@ use crate::{
             SearchIndexConfigParser,
             SearchOnDiskState,
             SearchSnapshot,
+            SegmentStatistics,
             SnapshotData,
         },
         BuildReason,
         MultiSegmentBackfillResult,
     },
-    metrics::vector::log_documents_per_segment,
     Database,
     IndexModel,
     Token,
@@ -238,7 +238,10 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
             None
         };
         let new_segment_id = new_segment.as_ref().map(T::IndexType::segment_id);
-        let vectors_in_new_segment = new_segment.as_ref().map(T::IndexType::num_vectors);
+        let new_segment_stats = new_segment
+            .as_ref()
+            .map(T::IndexType::statistics)
+            .transpose()?;
 
         let new_and_updated_parts = if let Some(new_segment) = new_segment {
             updated_previous_segments
@@ -249,21 +252,23 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
             updated_previous_segments
         };
 
-        let total_vectors = new_and_updated_parts
+        let total_stats = new_and_updated_parts
             .iter()
             .map(|segment| {
-                let total_vectors = T::IndexType::non_deleted_vectors(segment)?;
-                log_documents_per_segment(total_vectors);
-                Ok(total_vectors)
+                let segment_stats = T::IndexType::statistics(segment)?;
+                segment_stats.log();
+                Ok(segment_stats)
             })
-            .sum::<anyhow::Result<_>>()?;
+            .reduce(SegmentStatistics::add)
+            .transpose()?
+            .unwrap_or_default();
         let data = SnapshotData::MultiSegment(new_and_updated_parts);
 
         Ok(IndexBuildResult {
             snapshot_ts: *new_ts,
             data,
-            total_vectors,
-            vectors_in_new_segment,
+            total_stats,
+            new_segment_stats,
             new_segment_id,
             backfill_result,
         })
@@ -442,8 +447,8 @@ pub struct IndexBuild<T: SearchIndex> {
 pub struct IndexBuildResult<T: SearchIndex> {
     pub snapshot_ts: Timestamp,
     pub data: SnapshotData<T::Segment>,
-    pub total_vectors: u64,
-    pub vectors_in_new_segment: Option<u32>,
+    pub total_stats: T::Statistics,
+    pub new_segment_stats: Option<T::Statistics>,
     pub new_segment_id: Option<String>,
     // If this is set, this iteration made progress on backfilling an index
     pub backfill_result: Option<MultiSegmentBackfillResult>,

crates/database/src/vector_index_worker/flusher.rs

@@ -112,8 +112,8 @@ impl<RT: Runtime> VectorIndexFlusher<RT> {
         let IndexBuildResult {
             snapshot_ts,
             data,
-            total_vectors,
-            vectors_in_new_segment,
+            total_stats,
+            new_segment_stats,
             new_segment_id,
             backfill_result,
         } = result;
@@ -129,8 +129,8 @@ impl<RT: Runtime> VectorIndexFlusher<RT> {
             },
         }
 
-        let vectors_in_new_segment = vectors_in_new_segment.unwrap_or(0);
-        metrics::vector::log_documents_per_index(total_vectors);
+        let vectors_in_new_segment = new_segment_stats.unwrap_or_default().num_vectors;
+        metrics::vector::log_documents_per_index(total_stats.non_deleted_vectors);
         metrics::vector::log_documents_per_new_segment(vectors_in_new_segment);
         timer.finish();
         Ok(vectors_in_new_segment)