Generalize logging text/vector segment/index metadata after a flush (#26326)

We're going to care about similar things for text / vector search and I'd like to generalize the remainder of the flushers. So I've created a combined set of metrics that can be used with both.

I ended up expanding the metrics slightly to include both the total number of documents as well as the number of non-deleted documents in each new segment and each index. Total number of documents gives us some idea of how big an index is (actual byte size would be nice too). Non-deleted documents lets us track how much overhead there is for each segment due to deletions.

I also removed logging the size of each segment after a flush. Now we only log the size of the new segment and the total size of the index across all segments. I'll have to fix a few graphs, but this seems more coherent.

GitOrigin-RevId: 1be03cb43651c98477ee7166c829527645da17ce

Author: sjudd

crates/common/src/bootstrap_model/index/text_index/index_snapshot.rs

@@ -143,7 +143,7 @@ pub struct SerializedFragmentedTextSegment {
     pub id_tracker_key: String,
     pub deleted_terms_table_key: String,
     pub alive_bitset_key: String,
-    pub num_indexed_documents: u32,
+    pub num_indexed_documents: u64,
     pub id: String,
 }
 
@@ -184,7 +184,12 @@ pub struct FragmentedTextSegment {
     pub id_tracker_key: ObjectKey,
     pub deleted_terms_table_key: ObjectKey,
     pub alive_bitset_key: ObjectKey,
-    pub num_indexed_documents: u32,
+    // 2^63 ~= 9.2 * 10^18. We only support i64 in Convex.
+    #[cfg_attr(
+        any(test, feature = "testing"),
+        proptest(strategy = "1u64..9223372000000000000")
+    )]
+    pub num_indexed_documents: u64,
     // A random UUID that can be used to identify a segment to determine if the
     // segment has changed during non-transactional index changes (compaction).
     pub id: String,

crates/common/src/knobs.rs

@@ -784,7 +784,7 @@ pub static MULTI_SEGMENT_FULL_SCAN_THRESHOLD_KB: LazyLock<usize> =
 pub static SEGMENT_MAX_SIZE_BYTES: LazyLock<u64> =
     LazyLock::new(|| env_config("SEGMENT_MAX_SIZE_BYTES", (1100000 * 2048 * 4) / 3_u64));
 /// The minimum number of segments we will compact in one pass.
-pub static MIN_COMPACTION_SEGMENTS: LazyLock<usize> =
+pub static MIN_COMPACTION_SEGMENTS: LazyLock<u64> =
     LazyLock::new(|| env_config("MIN_COMPACTION_SEGMENTS", 3));
 /// The maximum percentage of a Segment that can be deleted before we will
 /// recompact that segment to remove deleted vectors

crates/database/src/index_workers/index_meta.rs

@@ -128,7 +128,10 @@ pub trait SearchIndex: Clone {
 
 pub trait SegmentStatistics: Default {
     fn add(lhs: anyhow::Result<Self>, rhs: anyhow::Result<Self>) -> anyhow::Result<Self>;
-    fn log(&self);
+
+    fn num_documents(&self) -> u64;
+
+    fn num_non_deleted_documents(&self) -> u64;
 }
 pub struct SearchIndexConfig<T: SearchIndex> {
     pub developer_config: T::DeveloperConfig,

crates/database/src/index_workers/search_flusher.rs

@@ -299,11 +299,7 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
 
         let total_stats = new_and_updated_parts
             .iter()
-            .map(|segment| {
-                let segment_stats = segment.statistics()?;
-                segment_stats.log();
-                Ok(segment_stats)
-            })
+            .map(|segment| segment.statistics())
             .reduce(SegmentStatistics::add)
             .transpose()?
             .unwrap_or_default();

crates/database/src/index_workers/writer.rs

@@ -76,6 +76,16 @@ pub(crate) struct SearchIndexMetadataWriter<RT: Runtime, T: SearchIndex> {
     search_type: SearchType,
 }
 
+pub struct SearchIndexWriteResult<T: SearchIndex> {
+    /// Stats summing metrics across all segments in the index.
+    pub index_stats: T::Statistics,
+    /// Stats for just the new segment if a new segment was built.
+    pub new_segment_stats: Option<T::Statistics>,
+    /// Stats for all segments in the index, including the new segment if it was
+    /// built.
+    pub per_segment_stats: Vec<T::Statistics>,
+}
+
 impl<RT: Runtime, T: SearchIndex> SearchIndexMetadataWriter<RT, T> {
     pub(crate) fn new(
         runtime: RT,
@@ -159,7 +169,7 @@ impl<RT: Runtime, T: SearchIndex> SearchIndexMetadataWriter<RT, T> {
         &self,
         job: &IndexBuild<T>,
         result: IndexBuildResult<T>,
-    ) -> anyhow::Result<(T::Statistics, Option<T::Statistics>)> {
+    ) -> anyhow::Result<SearchIndexWriteResult<T>> {
         let IndexBuildResult {
             snapshot_ts,
             data,
@@ -171,6 +181,10 @@ impl<RT: Runtime, T: SearchIndex> SearchIndexMetadataWriter<RT, T> {
 
         let inner = self.inner(SearchWriterLockWaiter::Flusher).await;
         let segments = data.require_multi_segment()?;
+        let per_segment_stats = segments
+            .iter()
+            .map(|segment| segment.statistics())
+            .collect::<anyhow::Result<Vec<_>>>()?;
 
         if let Some(index_backfill_result) = backfill_result {
             inner
@@ -188,7 +202,11 @@ impl<RT: Runtime, T: SearchIndex> SearchIndexMetadataWriter<RT, T> {
                 .await?
         }
 
-        Ok((total_stats, new_segment_stats))
+        Ok(SearchIndexWriteResult {
+            index_stats: total_stats,
+            new_segment_stats,
+            per_segment_stats,
+        })
     }
 
     async fn inner(&self, waiter: SearchWriterLockWaiter) -> MutexGuard<Inner<RT, T>> {

crates/database/src/metrics.rs

@@ -808,6 +808,75 @@ pub fn finish_search_index_merge_timer(mut timer: StatusTimer, merge_type: Searc
     timer.finish();
 }
 
+register_convex_histogram!(
+    DOCUMENTS_PER_NEW_SEARCH_SEGMENT_TOTAL,
+    "Total number of documents in a newly built search index segment.",
+    &[SEARCH_TYPE_LABEL],
+);
+pub fn log_documents_per_new_search_segment(count: u64, search_type: SearchType) {
+    log_distribution_with_labels(
+        &DOCUMENTS_PER_NEW_SEARCH_SEGMENT_TOTAL,
+        count as f64,
+        vec![search_type.tag()],
+    );
+}
+
+register_convex_histogram!(
+    DOCUMENTS_PER_SEARCH_SEGMENT_TOTAL,
+    "Total number of documents in a specific search segment, including documents that were added \
+     to the segment but are deleted and excluded from any search results",
+    &[SEARCH_TYPE_LABEL],
+);
+pub fn log_documents_per_search_segment(count: u64, search_type: SearchType) {
+    log_distribution_with_labels(
+        &DOCUMENTS_PER_SEARCH_SEGMENT_TOTAL,
+        count as f64,
+        vec![search_type.tag()],
+    );
+}
+
+register_convex_histogram!(
+    NON_DELETED_DOCUMENTS_PER_SEARCH_SEGMENT_TOTAL,
+    "Total number of non-deleted documents in a specific search segment, excluding documents that \
+     were added to the segment but are deleted and excluded from any search results",
+    &[SEARCH_TYPE_LABEL],
+);
+pub fn log_non_deleted_documents_per_search_segment(count: u64, search_type: SearchType) {
+    log_distribution_with_labels(
+        &NON_DELETED_DOCUMENTS_PER_SEARCH_SEGMENT_TOTAL,
+        count as f64,
+        vec![search_type.tag()],
+    );
+}
+
+register_convex_histogram!(
+    DOCUMENTS_PER_SEARCH_INDEX_TOTAL,
+    "Total number of documents across all segments in a search index, including documents that \
+     were added to the index but are deleted and excluded from any search results",
+    &[SEARCH_TYPE_LABEL],
+);
+pub fn log_documents_per_search_index(count: u64, search_type: SearchType) {
+    log_distribution_with_labels(
+        &DOCUMENTS_PER_SEARCH_INDEX_TOTAL,
+        count as f64,
+        vec![search_type.tag()],
+    );
+}
+register_convex_histogram!(
+    NON_DELETED_DOCUMENTS_PER_SEARCH_INDEX_TOTAL,
+    "Total number of non-deleted documents across all segments in a search index segment, \
+     excluding documents that were added to the index but are deleted and excluded from any \
+     search results",
+    &[SEARCH_TYPE_LABEL],
+);
+pub fn log_non_deleted_documents_per_search_index(count: u64, search_type: SearchType) {
+    log_distribution_with_labels(
+        &NON_DELETED_DOCUMENTS_PER_SEARCH_INDEX_TOTAL,
+        count as f64,
+        vec![search_type.tag()],
+    );
+}
+
 pub mod search {
 
     use metrics::{
@@ -860,14 +929,6 @@ pub mod search {
         log_distribution(&DATABASE_SEARCH_DOCUMENTS_PER_INDEX_TOTAL, count as f64);
     }
 
-    register_convex_histogram!(
-        DATABASE_TEXT_DOCUMENTS_PER_NEW_SEGMENT_TOTAL,
-        "Number of documents in a newly built text search index segment",
-    );
-    pub fn log_documents_per_new_segment(count: u32) {
-        log_distribution(&DATABASE_TEXT_DOCUMENTS_PER_NEW_SEGMENT_TOTAL, count as f64);
-    }
-
     register_convex_histogram!(
         DATABASE_SEARCH_ITERATOR_NEXT_SECONDS,
         "Time to fetch the next document in a search query iterator",
@@ -914,33 +975,6 @@ pub mod vector {
         StatusTimer::new(&DATABASE_VECTOR_BUILD_ONE_SECONDS)
     }
 
-    register_convex_histogram!(
-        DATABASE_VECTOR_DOCUMENTS_PER_INDEX_TOTAL,
-        "Number of documents per vector index",
-    );
-    pub fn log_documents_per_index(count: u64) {
-        log_distribution(&DATABASE_VECTOR_DOCUMENTS_PER_INDEX_TOTAL, count as f64);
-    }
-
-    register_convex_histogram!(
-        DATABASE_VECTOR_DOCUMENTS_PER_SEGMENT_TOTAL,
-        "Number of documents per vector index segment",
-    );
-    pub fn log_documents_per_segment(count: u64) {
-        log_distribution(&DATABASE_VECTOR_DOCUMENTS_PER_SEGMENT_TOTAL, count as f64);
-    }
-
-    register_convex_histogram!(
-        DATABASE_VECTOR_DOCUMENTS_PER_NEW_SEGMENT_TOTAL,
-        "Number of documents in a newly built vector index segment",
-    );
-    pub fn log_documents_per_new_segment(count: u32) {
-        log_distribution(
-            &DATABASE_VECTOR_DOCUMENTS_PER_NEW_SEGMENT_TOTAL,
-            count as f64,
-        );
-    }
-
     register_convex_histogram!(
         DATABASE_VECTOR_SEARCH_QUERY_SECONDS,
         "Time to run a single vector search, not including retries due to bootstrapping",

crates/database/src/text_index_worker/flusher.rs

@@ -95,7 +95,7 @@ impl<RT: Runtime> TextIndexFlusher<RT> {
     ///
     /// Returns a map of IndexName to number of documents indexed for each
     /// index that was built.
-    pub(crate) async fn step(&mut self) -> anyhow::Result<(BTreeMap<TabletIndexName, u32>, Token)> {
+    pub(crate) async fn step(&mut self) -> anyhow::Result<(BTreeMap<TabletIndexName, u64>, Token)> {
         let mut metrics = BTreeMap::new();
 
         let (to_build, token) = self.needs_backfill().await?;
@@ -107,7 +107,7 @@ impl<RT: Runtime> TextIndexFlusher<RT> {
         for job in to_build {
             let index_name = job.index_name.clone();
             let num_documents_indexed = self.build_one(job).await?;
-            metrics.insert(index_name, num_documents_indexed as u32);
+            metrics.insert(index_name, num_documents_indexed as u64);
         }
 
         Ok((metrics, token))

crates/database/src/text_index_worker/flusher2.rs

@@ -19,16 +19,22 @@ use storage::Storage;
 
 use crate::{
     index_workers::{
+        index_meta::SegmentStatistics,
         search_flusher::{
             IndexBuild,
             SearchFlusher,
             SearchIndexLimits,
         },
-        writer::SearchIndexMetadataWriter,
+        writer::{
+            SearchIndexMetadataWriter,
+            SearchIndexWriteResult,
+        },
     },
-    metrics::search::{
-        log_documents_per_index,
-        log_documents_per_new_segment,
+    metrics::{
+        log_documents_per_new_search_segment,
+        log_documents_per_search_segment,
+        log_non_deleted_documents_per_search_index,
+        log_non_deleted_documents_per_search_segment,
     },
     text_index_worker::text_meta::{
         BuildTextIndexArgs,
@@ -158,7 +164,7 @@ impl<RT: Runtime> TextIndexFlusher2<RT> {
     ///
     /// Returns a map of IndexName to number of documents indexed for each
     /// index that was built.
-    pub(crate) async fn step(&mut self) -> anyhow::Result<(BTreeMap<TabletIndexName, u32>, Token)> {
+    pub(crate) async fn step(&mut self) -> anyhow::Result<(BTreeMap<TabletIndexName, u64>, Token)> {
         let mut metrics = BTreeMap::new();
 
         let (to_build, token) = self.flusher.needs_backfill().await?;
@@ -185,7 +191,7 @@ impl<RT: Runtime> TextIndexFlusher2<RT> {
         Ok((metrics, token))
     }
 
-    async fn build_one(&self, job: IndexBuild<TextSearchIndex>) -> anyhow::Result<u32> {
+    async fn build_one(&self, job: IndexBuild<TextSearchIndex>) -> anyhow::Result<u64> {
         let timer = crate::metrics::search::build_one_timer();
 
         let build_index_args = BuildTextIndexArgs {
@@ -198,13 +204,30 @@ impl<RT: Runtime> TextIndexFlusher2<RT> {
             .await?;
         tracing::debug!("Built a text segment for: {result:#?}");
 
-        let (total_stats, new_segment_stats) = self.writer.commit_flush(&job, result).await?;
-
-        let num_indexed_documents = new_segment_stats.unwrap_or_default().num_indexed_documents;
-        log_documents_per_new_segment(num_indexed_documents);
-        log_documents_per_index(total_stats.num_indexed_documents as usize);
+        let SearchIndexWriteResult {
+            index_stats,
+            new_segment_stats,
+            per_segment_stats,
+        } = self.writer.commit_flush(&job, result).await?;
+
+        let new_segment_stats = new_segment_stats.unwrap_or_default();
+        log_documents_per_new_search_segment(new_segment_stats.num_documents(), SearchType::Text);
+
+        per_segment_stats.into_iter().for_each(|stats| {
+            log_documents_per_search_segment(stats.num_documents(), SearchType::Text);
+            log_non_deleted_documents_per_search_segment(
+                stats.num_non_deleted_documents(),
+                SearchType::Text,
+            );
+        });
+
+        log_documents_per_new_search_segment(index_stats.num_documents(), SearchType::Text);
+        log_non_deleted_documents_per_search_index(
+            index_stats.num_non_deleted_documents(),
+            SearchType::Text,
+        );
         timer.finish();
-        Ok(num_indexed_documents)
+        Ok(new_segment_stats.num_documents())
     }
 }
 

crates/database/src/text_index_worker/text_meta.rs

@@ -260,7 +260,7 @@ impl SearchIndex for TextSearchIndex {
 
 #[derive(Debug, Default)]
 pub struct TextStatistics {
-    pub num_indexed_documents: u32,
+    pub num_indexed_documents: u64,
 }
 
 impl From<SearchOnDiskState<TextSearchIndex>> for TextIndexState {
@@ -290,7 +290,14 @@ impl SegmentStatistics for TextStatistics {
         })
     }
 
-    fn log(&self) {}
+    fn num_documents(&self) -> u64 {
+        self.num_indexed_documents
+    }
+
+    fn num_non_deleted_documents(&self) -> u64 {
+        // TODO(sam): Add a non-deleted number of documents.
+        self.num_indexed_documents
+    }
 }
 
 impl From<TextIndexBackfillState> for BackfillState<TextSearchIndex> {