fix: Use smaller buffers during merging (#71)

`MergeOptimizedInvertedIndexReader` was added in #32 in order to avoid making small reads to our underlying `FileHandle`. It did so by reading the entire content of the posting lists and positions at open time.

As that PR says:
> A likely downside to this approach is that now pg_search will be, indirectly, holding onto a lot of heap-allocated memory that was read from its block storage. Perhaps in the (near) future we can further optimize the new `MergeOptimizedInvertedIndexReader` such that it pages in blocks of a few megabytes at a time, on demand, rather than the whole file.

This PR makes that change. But it additionally removes code that was later added in #47 to borrow individual entries rather than creating `OwnedBytes` for them. I believe that this code was added due to a misunderstanding:

`OwnedBytes` is a total misnomer: the bytes are not "owned": they are immutably borrowed and reference counted. An `OwnedBytes` object can be created for any type which derefs to a slice of bytes, and can be cheaply cloned and sliced. So there is no need to actually borrow _or_ copy the buffer under the `OwnedBytes`. Removing the code that was doing so allows us to safely recreate our buffer without worrying about the lifetimes of buffers that we've handed out.

Author: stuhood

common/src/buffered_file_slice.rs

@@ -0,0 +1,106 @@
+use std::cell::RefCell;
+use std::cmp::min;
+use std::io;
+use std::ops::Range;
+
+use super::file_slice::FileSlice;
+use super::{HasLen, OwnedBytes};
+
+const DEFAULT_BUFFER_MAX_SIZE: usize = 4 * 1024 * 1024; // 4 MB
+
+/// A buffered reader for a FileSlice.
+///
+/// Reads the underlying `FileSlice` in large, sequential chunks to amortize
+/// the cost of `read_bytes` calls, while keeping peak memory usage under control.
+///
+/// TODO: Rather than wrapping a `FileSlice` in buffering, it will usually be better to adjust a
+/// `FileHandle` to directly handle buffering itself (as that allows separate `FileSlice`s read
+/// from the same `FileHandle` to share buffers.)
+pub struct BufferedFileSlice {
+    file_slice: FileSlice,
+    buffer: RefCell<OwnedBytes>,
+    buffer_range: RefCell<Range<u64>>,
+    buffer_max_size: usize,
+}
+
+impl BufferedFileSlice {
+    /// Creates a new `BufferedFileSlice`.
+    ///
+    /// The `buffer_max_size` is the amount of data that will be read from the
+    /// `FileSlice` on a buffer miss.
+    pub fn new(file_slice: FileSlice, buffer_max_size: usize) -> Self {
+        Self {
+            file_slice,
+            buffer: RefCell::new(OwnedBytes::empty()),
+            buffer_range: RefCell::new(0..0),
+            buffer_max_size,
+        }
+    }
+
+    /// Creates a new `BufferedFileSlice` with a default buffer max size.
+    pub fn new_with_default_buffer_size(file_slice: FileSlice) -> Self {
+        Self::new(file_slice, DEFAULT_BUFFER_MAX_SIZE)
+    }
+
+    /// Creates an empty `BufferedFileSlice`.
+    pub fn empty() -> Self {
+        Self::new(FileSlice::empty(), 0)
+    }
+
+    /// Returns an `OwnedBytes` corresponding to the given `required_range`.
+    ///
+    /// If the requested range is not in the buffer, this will trigger a read
+    /// from the underlying `FileSlice`.
+    ///
+    /// If the requested range is larger than the buffer_max_size, it will be read directly from the
+    /// source without buffering.
+    ///
+    /// # Errors
+    ///
+    /// Returns an `io::Error` if the underlying read fails or the range is
+    /// out of bounds.
+    pub fn get_bytes(&self, required_range: Range<u64>) -> io::Result<OwnedBytes> {
+        let buffer_range = self.buffer_range.borrow();
+
+        // Cache miss condition: the required range is not fully contained in the current buffer.
+        if required_range.start < buffer_range.start || required_range.end > buffer_range.end {
+            drop(buffer_range); // release borrow before mutating
+
+            if required_range.end > self.file_slice.len() as u64 {
+                return Err(io::Error::new(
+                    io::ErrorKind::UnexpectedEof,
+                    "Requested range extends beyond the end of the file slice.",
+                ));
+            }
+
+            if (required_range.end - required_range.start) as usize > self.buffer_max_size {
+                // This read is larger than our buffer max size.
+                // Read it directly and bypass the buffer to avoid churning.
+                return self
+                    .file_slice
+                    .read_bytes_slice(required_range.start as usize..required_range.end as usize);
+            }
+
+            let new_buffer_start = required_range.start;
+            let new_buffer_end = min(
+                new_buffer_start + self.buffer_max_size as u64,
+                self.file_slice.len() as u64,
+            );
+            let read_range = new_buffer_start..new_buffer_end;
+
+            let new_buffer = self
+                .file_slice
+                .read_bytes_slice(read_range.start as usize..read_range.end as usize)?;
+
+            self.buffer.replace(new_buffer);
+            self.buffer_range.replace(read_range);
+        }
+
+        // Now the data is guaranteed to be in the buffer.
+        let buffer = self.buffer.borrow();
+        let buffer_range = self.buffer_range.borrow();
+        let local_start = (required_range.start - buffer_range.start) as usize;
+        let local_end = (required_range.end - buffer_range.start) as usize;
+        Ok(buffer.slice(local_start..local_end))
+    }
+}

common/src/lib.rs

@@ -6,6 +6,7 @@ pub use byteorder::LittleEndian as Endianness;
 
 mod bitset;
 pub mod bounds;
+pub mod buffered_file_slice;
 mod byte_count;
 mod datetime;
 pub mod file_slice;

src/directory/mod.rs

@@ -19,6 +19,7 @@ mod composite_file;
 use std::io::BufWriter;
 use std::path::PathBuf;
 
+pub use common::buffered_file_slice::BufferedFileSlice;
 pub use common::file_slice::{FileHandle, FileSlice};
 pub use common::{AntiCallToken, OwnedBytes, TerminatingWrite};
 

src/index/merge_optimized_inverted_index_reader.rs

@@ -1,28 +1,25 @@
 use std::io;
 
-use common::OwnedBytes;
-
-use crate::directory::FileSlice;
-use crate::positions::borrowed_position_reader::BorrowedPositionReader;
-use crate::postings::borrowed_block_segment_postings::BorrowedBlockSegmentPostings;
-use crate::postings::borrowed_segment_postings::BorrowedSegmentPostings;
-use crate::postings::TermInfo;
+use crate::directory::{BufferedFileSlice, FileSlice};
+use crate::positions::PositionReader;
+use crate::postings::{BlockSegmentPostings, SegmentPostings, TermInfo};
 use crate::schema::IndexRecordOption;
 use crate::termdict::TermDictionary;
 
 /// The inverted index reader is in charge of accessing
 /// the inverted index associated with a specific field.
 ///
-/// This is optimized for merging in that it full reads
-/// the postings and positions files into memory when opened.
-/// This eliminates all disk I/O to these files during merging.
+/// This is optimized for merging in that it uses a buffered reader
+/// for the postings and positions files.
+/// This eliminates most disk I/O to these files during merging, without
+/// reading the entire file into memory at once.
 ///
 /// NB:  This is a copy/paste from [`InvertedIndexReader`] and trimmed
 /// down to only include the methods required by the merge process.
 pub(crate) struct MergeOptimizedInvertedIndexReader {
     termdict: TermDictionary,
-    postings_bytes: OwnedBytes,
-    positions_bytes: OwnedBytes,
+    postings_reader: BufferedFileSlice,
+    positions_reader: BufferedFileSlice,
     record_option: IndexRecordOption,
 }
 
@@ -36,8 +33,8 @@ impl MergeOptimizedInvertedIndexReader {
         let (_, postings_body) = postings_file_slice.split(8);
         Ok(MergeOptimizedInvertedIndexReader {
             termdict,
-            postings_bytes: postings_body.read_bytes()?,
-            positions_bytes: positions_file_slice.read_bytes()?,
+            postings_reader: BufferedFileSlice::new_with_default_buffer_size(postings_body),
+            positions_reader: BufferedFileSlice::new_with_default_buffer_size(positions_file_slice),
             record_option,
         })
     }
@@ -47,8 +44,8 @@ impl MergeOptimizedInvertedIndexReader {
     pub fn empty(record_option: IndexRecordOption) -> MergeOptimizedInvertedIndexReader {
         MergeOptimizedInvertedIndexReader {
             termdict: TermDictionary::empty(),
-            postings_bytes: FileSlice::empty().read_bytes().unwrap(),
-            positions_bytes: FileSlice::empty().read_bytes().unwrap(),
+            postings_reader: BufferedFileSlice::empty(),
+            positions_reader: BufferedFileSlice::empty(),
             record_option,
         }
     }
@@ -66,9 +63,11 @@ impl MergeOptimizedInvertedIndexReader {
         &self,
         term_info: &TermInfo,
         requested_option: IndexRecordOption,
-    ) -> io::Result<BorrowedBlockSegmentPostings> {
-        let postings_data = &self.postings_bytes[term_info.postings_range.clone()];
-        BorrowedBlockSegmentPostings::open(
+    ) -> io::Result<BlockSegmentPostings> {
+        let postings_data = self.postings_reader.get_bytes(
+            term_info.postings_range.start as u64..term_info.postings_range.end as u64,
+        )?;
+        BlockSegmentPostings::open(
             term_info.doc_freq,
             postings_data,
             self.record_option,
@@ -84,20 +83,22 @@ impl MergeOptimizedInvertedIndexReader {
         &self,
         term_info: &TermInfo,
         option: IndexRecordOption,
-    ) -> io::Result<BorrowedSegmentPostings> {
+    ) -> io::Result<SegmentPostings> {
         let option = option.downgrade(self.record_option);
 
         let block_postings = self.read_block_postings_from_terminfo(term_info, option)?;
         let position_reader = {
             if option.has_positions() {
-                let positions_data = &self.positions_bytes[term_info.positions_range.clone()];
-                let position_reader = BorrowedPositionReader::open(positions_data)?;
+                let positions_data = self.positions_reader.get_bytes(
+                    term_info.positions_range.start as u64..term_info.positions_range.end as u64,
+                )?;
+                let position_reader = PositionReader::open(positions_data)?;
                 Some(position_reader)
             } else {
                 None
             }
         };
-        Ok(BorrowedSegmentPostings::from_block_postings(
+        Ok(SegmentPostings::from_block_postings(
             block_postings,
             position_reader,
         ))

src/indexer/merger.rs

@@ -17,8 +17,7 @@ use crate::index::{Segment, SegmentComponent, SegmentReader};
 use crate::indexer::doc_id_mapping::{MappingType, SegmentDocIdMapping};
 use crate::indexer::segment_updater::CancelSentinel;
 use crate::indexer::SegmentSerializer;
-use crate::postings::borrowed_segment_postings::BorrowedSegmentPostings;
-use crate::postings::InvertedIndexSerializer;
+use crate::postings::{InvertedIndexSerializer, Postings, SegmentPostings};
 use crate::schema::{value_type_to_column_type, Field, FieldType, Schema};
 use crate::store::StoreWriter;
 use crate::termdict::{TermMerger, TermOrdinal};
@@ -373,8 +372,7 @@ impl IndexMerger {
                          indexed. Have you modified the schema?",
         );
 
-        let mut segment_postings_containing_the_term: Vec<(usize, BorrowedSegmentPostings)> =
-            vec![];
+        let mut segment_postings_containing_the_term: Vec<(usize, SegmentPostings)> = vec![];
 
         let mut cnt = 0;
         while merged_terms.advance() {