Add multi segment metadata to text search's backfilling state (#25752)

GitOrigin-RevId: 9afd42bb5610db0f02a294c09f70bfb3ba36c861

Author: sjudd

crates/common/src/bootstrap_model/index/index_config.rs

@@ -75,7 +75,7 @@ impl IndexConfig {
                 matches!(on_disk_state, DatabaseIndexState::Backfilling(_))
             },
             IndexConfig::Search { on_disk_state, .. } => {
-                matches!(on_disk_state, SearchIndexState::Backfilling)
+                matches!(on_disk_state, SearchIndexState::Backfilling(_))
             },
             IndexConfig::Vector { on_disk_state, .. } => {
                 matches!(on_disk_state, VectorIndexState::Backfilling(_))

crates/common/src/bootstrap_model/index/index_metadata.rs

@@ -36,6 +36,7 @@ use super::{
     IndexConfig,
 };
 use crate::{
+    bootstrap_model::index::search_index::TextIndexBackfillState,
     document::{
         ParsedDocument,
         ResolvedDocument,
@@ -90,7 +91,7 @@ impl<T: TableIdentifier> IndexMetadata<T> {
                 search_field,
                 filter_fields,
             },
-            SearchIndexState::Backfilling,
+            SearchIndexState::Backfilling(TextIndexBackfillState::new()),
         )
     }
 

crates/common/src/bootstrap_model/index/search_index/backfill_state.rs

@@ -0,0 +1,113 @@
+use std::str::FromStr;
+
+use serde::{
+    Deserialize,
+    Serialize,
+};
+use sync_types::Timestamp;
+use value::{
+    codegen_convex_serialization,
+    InternalId,
+};
+
+use crate::bootstrap_model::index::search_index::{
+    index_snapshot::SerializedFragmentedSearchSegment,
+    FragmentedSearchSegment,
+};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
+pub struct TextIndexBackfillState {
+    pub segments: Vec<FragmentedSearchSegment>,
+    // None at the start of backfill, then set after the first backfill iteration.
+    pub cursor: Option<TextBackfillCursor>,
+}
+
+impl TextIndexBackfillState {
+    pub fn new() -> Self {
+        Self {
+            segments: vec![],
+            cursor: None,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
+pub struct TextBackfillCursor {
+    pub cursor: InternalId,
+    pub backfill_snapshot_ts: Timestamp,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct SerializedTextBackfillCursor {
+    pub document_cursor: String,
+    pub backfill_snapshot_ts: i64,
+}
+
+impl From<TextBackfillCursor> for SerializedTextBackfillCursor {
+    fn from(value: TextBackfillCursor) -> Self {
+        Self {
+            document_cursor: value.cursor.to_string(),
+            backfill_snapshot_ts: value.backfill_snapshot_ts.into(),
+        }
+    }
+}
+
+impl TryFrom<SerializedTextBackfillCursor> for TextBackfillCursor {
+    type Error = anyhow::Error;
+
+    fn try_from(value: SerializedTextBackfillCursor) -> Result<Self, Self::Error> {
+        Ok(Self {
+            cursor: InternalId::from_str(&value.document_cursor)?,
+            backfill_snapshot_ts: Timestamp::try_from(value.backfill_snapshot_ts)?,
+        })
+    }
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct SerializedTextIndexBackfillState {
+    segments: Option<Vec<SerializedFragmentedSearchSegment>>,
+    cursor: Option<SerializedTextBackfillCursor>,
+}
+
+impl TryFrom<TextIndexBackfillState> for SerializedTextIndexBackfillState {
+    type Error = anyhow::Error;
+
+    fn try_from(backfill_state: TextIndexBackfillState) -> Result<Self, Self::Error> {
+        Ok(SerializedTextIndexBackfillState {
+            segments: Some(
+                backfill_state
+                    .segments
+                    .into_iter()
+                    .map(|s| s.try_into())
+                    .collect::<anyhow::Result<Vec<_>>>()?,
+            ),
+            cursor: backfill_state
+                .cursor
+                .map(|cursor| cursor.try_into())
+                .transpose()?,
+        })
+    }
+}
+
+impl TryFrom<SerializedTextIndexBackfillState> for TextIndexBackfillState {
+    type Error = anyhow::Error;
+
+    fn try_from(serialized: SerializedTextIndexBackfillState) -> Result<Self, Self::Error> {
+        Ok(TextIndexBackfillState {
+            segments: serialized
+                .segments
+                .unwrap_or_default()
+                .into_iter()
+                .map(|s| s.try_into())
+                .collect::<anyhow::Result<Vec<_>>>()?,
+            cursor: serialized
+                .cursor
+                .map(TextBackfillCursor::try_from)
+                .transpose()?,
+        })
+    }
+}
+
+codegen_convex_serialization!(TextIndexBackfillState, SerializedTextIndexBackfillState);

crates/common/src/bootstrap_model/index/search_index/index_state.rs

@@ -8,6 +8,10 @@ use super::{
     index_snapshot::SerializedSearchIndexSnapshot,
     SearchIndexSnapshot,
 };
+use crate::bootstrap_model::index::search_index::backfill_state::{
+    SerializedTextIndexBackfillState,
+    TextIndexBackfillState,
+};
 
 /// The state of a search index.
 /// Search indexes begin in `Backfilling`.
@@ -16,7 +20,7 @@ use super::{
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
 pub enum SearchIndexState {
-    Backfilling,
+    Backfilling(TextIndexBackfillState),
     Backfilled(SearchIndexSnapshot),
     SnapshottedAt(SearchIndexSnapshot),
 }
@@ -25,6 +29,10 @@ pub enum SearchIndexState {
 #[serde(tag = "state", rename_all = "camelCase")]
 pub enum SerializedSearchIndexState {
     Backfilling,
+    Backfilling2 {
+        #[serde(flatten)]
+        backfill_state: SerializedTextIndexBackfillState,
+    },
     Backfilled {
         #[serde(flatten)]
         snapshot: SerializedSearchIndexSnapshot,
@@ -40,7 +48,9 @@ impl TryFrom<SearchIndexState> for SerializedSearchIndexState {
 
     fn try_from(state: SearchIndexState) -> Result<Self, Self::Error> {
         Ok(match state {
-            SearchIndexState::Backfilling => SerializedSearchIndexState::Backfilling,
+            SearchIndexState::Backfilling(state) => SerializedSearchIndexState::Backfilling2 {
+                backfill_state: state.try_into()?,
+            },
             SearchIndexState::Backfilled(snapshot) => SerializedSearchIndexState::Backfilled {
                 snapshot: snapshot.try_into()?,
             },
@@ -56,7 +66,12 @@ impl TryFrom<SerializedSearchIndexState> for SearchIndexState {
 
     fn try_from(serialized: SerializedSearchIndexState) -> Result<Self, Self::Error> {
         Ok(match serialized {
-            SerializedSearchIndexState::Backfilling => SearchIndexState::Backfilling,
+            SerializedSearchIndexState::Backfilling => {
+                SearchIndexState::Backfilling(TextIndexBackfillState::new())
+            },
+            SerializedSearchIndexState::Backfilling2 { backfill_state } => {
+                SearchIndexState::Backfilling(backfill_state.try_into()?)
+            },
             SerializedSearchIndexState::Backfilled { snapshot } => {
                 SearchIndexState::Backfilled(snapshot.try_into()?)
             },

crates/common/src/bootstrap_model/index/search_index/mod.rs

@@ -1,8 +1,10 @@
+mod backfill_state;
 mod index_config;
 mod index_snapshot;
 mod index_state;
 
 pub use self::{
+    backfill_state::TextIndexBackfillState,
     index_config::{
         DeveloperSearchIndexConfig,
         SerializedDeveloperSearchIndexConfig,

crates/database/src/bootstrap_model/index.rs

@@ -210,7 +210,7 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
                 SearchIndexState::Backfilled(snapshot) => {
                     *on_disk_state = SearchIndexState::SnapshottedAt(snapshot.clone());
                 },
-                SearchIndexState::Backfilling | SearchIndexState::SnapshottedAt(_) => {
+                SearchIndexState::Backfilling(_) | SearchIndexState::SnapshottedAt(_) => {
                     anyhow::bail!(
                         "Expected backfilled index, but found: {on_disk_state:?} for {:?}",
                         backfilled_index.name.descriptor()

crates/database/src/index_workers/index_meta.rs

@@ -80,7 +80,7 @@ impl SearchIndexConfigParser for TextIndexConfigParser {
         Some(SearchIndexConfig {
             developer_config,
             on_disk_state: match on_disk_state {
-                SearchIndexState::Backfilling => {
+                SearchIndexState::Backfilling(_) => {
                     // TODO(sam): Add support for a backfilling partial state to text search
                     SearchOnDiskState::Backfilling(BackfillState {
                         segments: vec![],

crates/database/src/index_workers/mod.rs

@@ -16,6 +16,7 @@ use common::{
     runtime::Runtime,
 };
 use rand::Rng;
+use value::ResolvedDocumentId;
 
 pub const MAX_BACKOFF: Duration = Duration::from_secs(30);
 
@@ -48,3 +49,9 @@ pub async fn timeout_with_jitter<RT: Runtime>(rt: &RT, duration: Duration) {
     let sleep = rt.with_rng(|rng| half_timer + duration.mul_f32(rng.gen::<f32>()));
     rt.wait(sleep).await;
 }
+
+#[derive(Debug)]
+pub struct MultiSegmentBackfillResult {
+    pub new_cursor: Option<ResolvedDocumentId>,
+    pub is_backfill_complete: bool,
+}

crates/database/src/search_and_vector_bootstrap.rs

@@ -165,7 +165,7 @@ impl IndexesToBootstrap {
                     on_disk_state,
                 } => {
                     let search_index = match on_disk_state {
-                        SearchIndexState::Backfilling => {
+                        SearchIndexState::Backfilling(_) => {
                             // We'll start a new memory search index starting at the next commit
                             // after our persistence upper bound. After
                             // bootstrapping, all commits after

crates/database/src/search_index_worker/fast_forward.rs

@@ -43,7 +43,7 @@ impl<RT: Runtime> IndexFastForward<RT, SearchSnapshotVersion> for SearchFastForw
             SearchIndexState::SnapshottedAt(snapshot) | SearchIndexState::Backfilled(snapshot) => {
                 snapshot
             },
-            SearchIndexState::Backfilling => return None,
+            SearchIndexState::Backfilling(_) => return None,
         };
         Some((*ts, *version))
     }

crates/database/src/search_index_worker/flusher.rs

@@ -146,7 +146,7 @@ impl<RT: Runtime> SearchIndexFlusher<RT> {
             // If the index is in the `Backfilling` state, or is already `SnapshottedAt` but
             // has grown too large or has the wrong format, it needs to be backfilled.
             let needs_backfill = match &on_disk_state {
-                SearchIndexState::Backfilling => Some(BuildReason::Backfilling),
+                SearchIndexState::Backfilling(_) => Some(BuildReason::Backfilling),
                 SearchIndexState::SnapshottedAt(SearchIndexSnapshot { version, .. })
                 | SearchIndexState::Backfilled(SearchIndexSnapshot { version, .. })
                     if *version != expected_version =>
@@ -226,7 +226,7 @@ impl<RT: Runtime> SearchIndexFlusher<RT> {
             version: SearchSnapshotVersion::new(tx.persistence_version()),
         };
         let new_on_disk_state = match job.on_disk_state {
-            SearchIndexState::Backfilling | SearchIndexState::Backfilled(_) => {
+            SearchIndexState::Backfilling(_) | SearchIndexState::Backfilled(_) => {
                 SearchIndexState::Backfilled(snapshot_data)
             },
             SearchIndexState::SnapshottedAt(_) => SearchIndexState::SnapshottedAt(snapshot_data),

crates/database/src/test_helpers/index_utils.rs

@@ -69,7 +69,7 @@ pub fn assert_backfilling(
             assert_matches!(on_disk_state, DatabaseIndexState::Backfilling(_))
         },
         IndexConfig::Search { on_disk_state, .. } => {
-            assert_eq!(on_disk_state, SearchIndexState::Backfilling)
+            assert_matches!(on_disk_state, SearchIndexState::Backfilling(_))
         },
         IndexConfig::Vector { on_disk_state, .. } => {
             assert_matches!(on_disk_state, VectorIndexState::Backfilling(_))

crates/database/src/vector_index_worker/flusher.rs

@@ -73,6 +73,7 @@ use crate::{
             VectorSearchIndex,
         },
         BuildReason,
+        MultiSegmentBackfillResult,
     },
     metrics::{
         self,
@@ -518,7 +519,7 @@ impl<RT: Runtime> VectorIndexFlusher<RT> {
 
         let index_backfill_result =
             if let MultipartBuildType::IncrementalComplete { .. } = build_type {
-                Some(VectorIndexMultiSegmentBackfillResult {
+                Some(MultiSegmentBackfillResult {
                     new_cursor,
                     is_backfill_complete,
                 })
@@ -622,7 +623,7 @@ struct MultiSegmentBuildResult {
     new_segment: Option<DiskSegmentValues>,
     updated_previous_segments: Vec<FragmentedVectorSegment>,
     // This is set only if the build iteration created a segment for a backfilling index
-    backfill_result: Option<VectorIndexMultiSegmentBackfillResult>,
+    backfill_result: Option<MultiSegmentBackfillResult>,
 }
 
 #[cfg(any(test, feature = "testing"))]
@@ -646,13 +647,7 @@ struct IndexBuildResult {
     vectors_in_new_segment: Option<u32>,
     new_segment_id: Option<String>,
     // If this is set, this iteration made progress on backfilling an index
-    backfill_result: Option<VectorIndexMultiSegmentBackfillResult>,
-}
-
-#[derive(Debug)]
-pub struct VectorIndexMultiSegmentBackfillResult {
-    pub new_cursor: Option<ResolvedDocumentId>,
-    pub is_backfill_complete: bool,
+    backfill_result: Option<MultiSegmentBackfillResult>,
 }
 
 #[cfg(test)]

crates/database/src/vector_index_worker/writer.rs

@@ -49,7 +49,10 @@ use vector::QdrantExternalId;
 
 use super::IndexBuild;
 use crate::{
-    index_workers::index_meta::VectorSearchIndex,
+    index_workers::{
+        index_meta::VectorSearchIndex,
+        MultiSegmentBackfillResult,
+    },
     metrics::vector::{
         finish_vector_index_merge_timer,
         vector_compaction_merge_commit_timer,
@@ -58,7 +61,6 @@ use crate::{
         VectorIndexMergeType,
         VectorWriterLockWaiter,
     },
-    vector_index_worker::flusher::VectorIndexMultiSegmentBackfillResult,
     Database,
     IndexModel,
     SystemMetadataModel,
@@ -142,7 +144,7 @@ impl<RT: Runtime> VectorMetadataWriter<RT> {
         new_ts: Timestamp,
         new_and_modified_segments: Vec<FragmentedVectorSegment>,
         new_segment_id: Option<String>,
-        index_backfill_result: Option<VectorIndexMultiSegmentBackfillResult>,
+        index_backfill_result: Option<MultiSegmentBackfillResult>,
     ) -> anyhow::Result<()> {
         let inner = self.inner(VectorWriterLockWaiter::Flusher).await;
 
@@ -392,7 +394,7 @@ impl<RT: Runtime> Inner<RT> {
         backfill_complete_ts: Timestamp,
         mut new_and_modified_segments: Vec<FragmentedVectorSegment>,
         new_segment_id: Option<String>,
-        backfill_result: VectorIndexMultiSegmentBackfillResult,
+        backfill_result: MultiSegmentBackfillResult,
     ) -> anyhow::Result<()> {
         let timer = vector_flush_merge_commit_timer();
         let mut tx: Transaction<RT> = self.database.begin(Identity::system()).await?;

crates/local_backend/src/schema.rs

@@ -152,7 +152,7 @@ impl TryFrom<IndexMetadata<TableName>> for IndexMetadataResponse {
                     },
             } => {
                 let backfill_state = match on_disk_state {
-                    SearchIndexState::Backfilling => "in_progress".to_string(),
+                    SearchIndexState::Backfilling(_) => "in_progress".to_string(),
                     // TODO(CX-3851): The result of this is used to poll for state in the CLI and
                     // also for display in the dashboard. We might consider a new value that would
                     // let us differentiate between Backfilled and SnapshottedAt in the dashboard.

crates/model/src/config/index_tests.rs

@@ -650,7 +650,8 @@ async fn apply_config_with_backfilling_search_index_throws(rt: TestRuntime) -> a
     // commit the schema.
     assert_root_cause_contains(
         result,
-        "Expected backfilled index, but found: Backfilling for \"index\"",
+        "Expected backfilled index, but found: Backfilling(TextIndexBackfillState { segments: [], \
+         cursor: None }) for \"index\"",
     );
 
     Ok(())
@@ -1099,7 +1100,7 @@ fn assert_index_data(actual: Vec<IndexConfig>, expected: Vec<TestIndexConfig>) {
                 on_disk_state,
             } => {
                 let search_state = match on_disk_state {
-                    SearchIndexState::Backfilling => TestIndexState::Backfilling,
+                    SearchIndexState::Backfilling(_) => TestIndexState::Backfilling,
                     SearchIndexState::Backfilled(_) => TestIndexState::Backfilled,
                     SearchIndexState::SnapshottedAt(_) => TestIndexState::Enabled,
                 };