perf: lazily open positions file (#54)

Not all queries use positions and it's okay if we (from the perspective of `pg_search`, anyways) defer opening/loading them until they're first needed.

This is probably completely wrong for a mmap-based Directory, but we (again, `pg_search`) decided long ago that we don't care about that use case.

This saves a lot of disk I/O when an index has lots of segments and the query doesn't need positions.

As a drive by, make sure a random Vec has enough space before pushing items to it.  This showed up in the profiler, believe it or not.

Author: eeeebbbbrrrr

common/src/file_slice.rs

@@ -1,7 +1,7 @@
 use std::fs::File;
 use std::ops::{Deref, Range, RangeBounds};
 use std::path::Path;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 use std::{fmt, io};
 
 use async_trait::async_trait;
@@ -339,6 +339,27 @@ impl FileHandle for OwnedBytes {
     }
 }
 
+pub struct DeferredFileSlice {
+    opener: Arc<dyn Fn() -> io::Result<FileSlice> + Send + Sync + 'static>,
+    file_slice: OnceLock<std::io::Result<FileSlice>>,
+}
+
+impl DeferredFileSlice {
+    pub fn new(opener: impl Fn() -> io::Result<FileSlice> + Send + Sync + 'static) -> Self {
+        DeferredFileSlice {
+            opener: Arc::new(opener),
+            file_slice: OnceLock::default(),
+        }
+    }
+
+    pub fn open(&self) -> io::Result<&FileSlice> {
+        match self.file_slice.get_or_init(|| (self.opener)()) {
+            Ok(file_slice) => Ok(file_slice),
+            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::io;

src/index/inverted_index_reader.rs

@@ -1,5 +1,6 @@
 use std::io;
 
+use common::file_slice::DeferredFileSlice;
 use common::json_path_writer::JSON_END_OF_PATH;
 use common::{BinarySerializable, ByteCount};
 #[cfg(feature = "quickwit")]
@@ -30,7 +31,7 @@ use crate::termdict::TermDictionary;
 pub struct InvertedIndexReader {
     termdict: TermDictionary,
     postings_file_slice: FileSlice,
-    positions_file_slice: FileSlice,
+    positions_file_slice: DeferredFileSlice,
     record_option: IndexRecordOption,
     total_num_tokens: u64,
 }
@@ -66,7 +67,7 @@ impl InvertedIndexReader {
     pub(crate) fn new(
         termdict: TermDictionary,
         postings_file_slice: FileSlice,
-        positions_file_slice: FileSlice,
+        positions_file_slice: DeferredFileSlice,
         record_option: IndexRecordOption,
     ) -> io::Result<InvertedIndexReader> {
         let (total_num_tokens_slice, postings_body) = postings_file_slice.split(8);
@@ -86,7 +87,7 @@ impl InvertedIndexReader {
         InvertedIndexReader {
             termdict: TermDictionary::empty(),
             postings_file_slice: FileSlice::empty(),
-            positions_file_slice: FileSlice::empty(),
+            positions_file_slice: DeferredFileSlice::new(|| Ok(FileSlice::empty())),
             record_option,
             total_num_tokens: 0u64,
         }
@@ -233,6 +234,7 @@ impl InvertedIndexReader {
             if option.has_positions() {
                 let positions_data = self
                     .positions_file_slice
+                    .open()?
                     .read_bytes_slice(term_info.positions_range.clone())?;
                 let position_reader = PositionReader::open(positions_data)?;
                 Some(position_reader)
@@ -342,6 +344,7 @@ impl InvertedIndexReader {
             if with_positions {
                 let positions = self
                     .positions_file_slice
+                    .open()?
                     .read_bytes_slice_async(term_info.positions_range.clone());
                 futures_util::future::try_join(postings, positions).await?;
             } else {
@@ -384,6 +387,7 @@ impl InvertedIndexReader {
         if with_positions {
             let positions = self
                 .positions_file_slice
+                .open()?
                 .read_bytes_slice_async(positions_range);
             futures_util::future::try_join(postings, positions).await?;
         } else {
@@ -478,7 +482,7 @@ impl InvertedIndexReader {
     pub async fn warm_postings_full(&self, with_positions: bool) -> io::Result<()> {
         self.postings_file_slice.read_bytes_async().await?;
         if with_positions {
-            self.positions_file_slice.read_bytes_async().await?;
+            self.positions_file_slice.open()?.read_bytes_async().await?;
         }
         Ok(())
     }

src/index/segment_reader.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use std::sync::{Arc, OnceLock, RwLock};
 use std::{fmt, io};
 
+use common::file_slice::DeferredFileSlice;
 use common::{ByteCount, HasLen};
 use fnv::FnvHashMap;
 use itertools::Itertools;
@@ -237,19 +238,38 @@ impl SegmentReader {
                 ))
             })?;
 
-        let positions_file = self.positions_composite().open_read(field).ok_or_else(|| {
-            let error_msg = format!(
-                "Failed to open field {:?}'s positions in the composite file. Has the schema been \
-                 modified?",
-                field_entry.name()
-            );
-            DataCorruption::comment_only(error_msg)
-        })?;
+        // not all queries require positions.
+        // we can defer opening the file until needed
+        let positions_file_opener = {
+            let path = self.relative_path(SegmentComponent::Positions);
+            let directory = self.index.directory().clone();
+            let field_entry = field_entry.clone();
+            move || {
+                let composite_file = if let Ok(positions_file) = &directory.open_read(&path) {
+                    CompositeFile::open(&positions_file)
+                        .expect("should be able to open positions composite component")
+                } else {
+                    CompositeFile::empty()
+                };
+
+                composite_file.open_read(field).ok_or_else(|| {
+                    let error_msg = format!(
+                        "Failed to open field {:?}'s positions in the composite file. Has the \
+                         schema been modified?",
+                        field_entry.name()
+                    );
+                    io::Error::new(
+                        io::ErrorKind::InvalidData,
+                        format!("{:?}", DataCorruption::comment_only(error_msg)),
+                    )
+                })
+            }
+        };
 
         let inv_idx_reader = Arc::new(InvertedIndexReader::new(
             TermDictionary::open(termdict_file)?,
             postings_file,
-            positions_file,
+            DeferredFileSlice::new(positions_file_opener),
             record_option,
         )?);
 

src/termdict/sstable_termdict/mod.rs

@@ -158,10 +158,12 @@ impl ValueReader for TermInfoValueReader {
 
     fn load(&mut self, mut data: &[u8]) -> io::Result<usize> {
         let len_before = data.len();
-        self.term_infos.clear();
         let num_els = VInt::deserialize_u64(&mut data)?;
         let mut postings_start = VInt::deserialize_u64(&mut data)? as usize;
         let mut positions_start = VInt::deserialize_u64(&mut data)? as usize;
+
+        self.term_infos.clear();
+        self.term_infos.reserve_exact(num_els as usize);
         for _ in 0..num_els {
             let doc_freq = VInt::deserialize_u64(&mut data)? as u32;
             let postings_num_bytes = VInt::deserialize_u64(&mut data)?;