Move some callers to PersistenceReader::load_documents_from_table (#36354)

GitOrigin-RevId: 35cf2f8cf142d6ab231a84d6a8ac331c58420ce0

Author: goffrie

crates/common/src/persistence.rs

@@ -529,12 +529,6 @@ impl RepeatablePersistence {
         self.upper_bound
     }
 
-    /// Same as [`Persistence::load_all_documents`] but only including documents
-    /// in the snapshot range.
-    pub fn load_all_documents(&self, order: Order) -> DocumentStream<'_> {
-        self.load_documents(TimestampRange::snapshot(*self.upper_bound), order)
-    }
-
     /// Same as [`Persistence::load_documents`] but only including documents in
     /// the snapshot range.
     pub fn load_documents(&self, range: TimestampRange, order: Order) -> DocumentStream<'_> {
@@ -547,6 +541,24 @@ impl RepeatablePersistence {
         Box::pin(stream.try_filter(|entry| future::ready(entry.ts <= *self.upper_bound)))
     }
 
+    /// Same as [`Persistence::load_documents_from_table`] but only including
+    /// documents in the snapshot range.
+    pub fn load_documents_from_table(
+        &self,
+        tablet_id: TabletId,
+        range: TimestampRange,
+        order: Order,
+    ) -> DocumentStream<'_> {
+        let stream = self.reader.load_documents_from_table(
+            tablet_id,
+            range,
+            order,
+            *DEFAULT_DOCUMENTS_PAGE_SIZE,
+            self.retention_validator.clone(),
+        );
+        Box::pin(stream.try_filter(|entry| future::ready(entry.ts <= *self.upper_bound)))
+    }
+
     /// Same as `load_documents` but doesn't use the `RetentionValidator` from
     /// this `RepeatablePersistence`. Instead, the caller can choose its
     /// own validator.

crates/database/src/index_worker.rs

@@ -76,7 +76,6 @@ use common::{
     },
 };
 use futures::{
-    future,
     pin_mut,
     stream::FusedStream,
     Future,
@@ -92,7 +91,6 @@ use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
 use value::{
     DeveloperDocumentId,
-    InternalDocumentId,
     TableNamespace,
 };
 
@@ -195,11 +193,11 @@ impl IndexSelector {
         tables.into_iter()
     }
 
-    fn filter_id(&self, id: InternalDocumentId) -> bool {
+    fn tablet_id(&self) -> Option<TabletId> {
         match self {
-            Self::All(_) => true,
-            Self::Index { name, .. } => *name.table() == id.table(),
-            Self::ManyIndexes { tablet_id, .. } => *tablet_id == id.table(),
+            Self::All(_) => None,
+            Self::Index { name, .. } => Some(*name.table()),
+            Self::ManyIndexes { tablet_id, .. } => Some(*tablet_id),
         }
     }
 }
@@ -940,9 +938,11 @@ impl<RT: Runtime> IndexWriter<RT> {
         order: Order,
         index_selector: &'a IndexSelector,
     ) -> impl Stream<Item = anyhow::Result<RevisionPair>> + 'a {
-        let document_stream = reader
-            .load_documents(range, order)
-            .try_filter(|entry| future::ready(index_selector.filter_id(entry.id)));
+        let document_stream = if let Some(tablet_id) = index_selector.tablet_id() {
+            reader.load_documents_from_table(tablet_id, range, order)
+        } else {
+            reader.load_documents(range, order)
+        };
         stream_revision_pairs(document_stream, reader)
     }
 

crates/database/src/retention.rs

@@ -1432,9 +1432,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         let Some(doc) = maybe_doc else {
             return Ok(());
         };
-        if doc.id().tablet_id != index_tablet_id {
-            return Ok(());
-        }
+        anyhow::ensure!(doc.id().tablet_id == index_tablet_id);
         let index_id = doc.id().internal_id();
         let index: ParsedDocument<IndexMetadata<TabletId>> = doc.parse()?;
         let index = index.into_value();
@@ -1471,7 +1469,8 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         retention_validator: Arc<dyn RetentionValidator>,
     ) -> anyhow::Result<()> {
         let reader = persistence.reader();
-        let mut document_stream = reader.load_documents(
+        let mut document_stream = reader.load_documents_from_table(
+            index_table_id,
             TimestampRange::new(**cursor..*latest_ts)?,
             Order::Asc,
             *DEFAULT_DOCUMENTS_PAGE_SIZE,