perf(index): retain Bytes backing for FST readers

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/index/src/inverted_index.rs

@@ -17,4 +17,7 @@ pub mod error;
 pub mod format;
 pub mod search;
 
-pub type FstMap = fst::Map<Vec<u8>>;
+/// A finite state transducer map that shares ownership of its backing bytes
+/// with the reader that produced it, so the FST payload is not copied into
+/// the map.
+pub type FstMap = fst::Map<bytes::Bytes>;

src/index/src/inverted_index/format/reader.rs

@@ -129,7 +129,8 @@ pub trait InvertedIndexReader: Send + Sync {
         metrics: Option<&'a mut InvertedIndexReadMetrics>,
     ) -> Result<FstMap> {
         let fst_data = self.range_read(offset, size, metrics).await?;
-        FstMap::new(fst_data).context(DecodeFstSnafu)
+        // `Bytes::from` takes ownership of the buffer without copying it.
+        FstMap::new(Bytes::from(fst_data)).context(DecodeFstSnafu)
     }
 
     /// Retrieves the multiple finite state transducer (FST) maps from the given ranges.
@@ -138,10 +139,12 @@ pub trait InvertedIndexReader: Send + Sync {
         ranges: &[Range<u64>],
         metrics: Option<&'a mut InvertedIndexReadMetrics>,
     ) -> Result<Vec<FstMap>> {
+        // The read buffers are reference-counted, so building the maps directly
+        // from them avoids copying each FST payload again.
         self.read_vec(ranges, metrics)
             .await?
             .into_iter()
-            .map(|bytes| FstMap::new(bytes.to_vec()).context(DecodeFstSnafu))
+            .map(|bytes| FstMap::new(bytes).context(DecodeFstSnafu))
             .collect::<Result<Vec<_>>>()
     }
 
@@ -177,3 +180,90 @@ pub trait InvertedIndexReader: Send + Sync {
             .collect::<Result<VecDeque<_>>>()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use async_trait::async_trait;
+    use fst::MapBuilder;
+
+    use super::*;
+    use crate::inverted_index::error::Error;
+
+    /// A reader whose `read_vec` returns the byte buffers it was created with.
+    struct BytesReader(Vec<Bytes>);
+
+    #[async_trait]
+    impl InvertedIndexReader for BytesReader {
+        async fn range_read<'a>(
+            &self,
+            _offset: u64,
+            _size: u32,
+            _metrics: Option<&'a mut InvertedIndexReadMetrics>,
+        ) -> Result<Vec<u8>> {
+            unimplemented!("only `read_vec` is used by the tests")
+        }
+
+        async fn read_vec<'a>(
+            &self,
+            _ranges: &[Range<u64>],
+            _metrics: Option<&'a mut InvertedIndexReadMetrics>,
+        ) -> Result<Vec<Bytes>> {
+            Ok(self.0.clone())
+        }
+
+        async fn metadata<'a>(
+            &self,
+            _metrics: Option<&'a mut InvertedIndexReadMetrics>,
+        ) -> Result<Arc<InvertedIndexMetas>> {
+            unimplemented!("only `read_vec` is used by the tests")
+        }
+    }
+
+    fn mock_fst_bytes() -> Vec<u8> {
+        let mut build = MapBuilder::memory();
+        build.insert("key1".as_bytes(), 1).unwrap();
+        build.insert("key2".as_bytes(), 2).unwrap();
+        build.into_inner().unwrap()
+    }
+
+    #[tokio::test]
+    async fn test_fst_vec_shares_the_reader_bytes() {
+        let fst_bytes = Bytes::from(mock_fst_bytes());
+        let data_ptr = fst_bytes.as_ptr();
+
+        let mut reader = BytesReader(vec![fst_bytes.clone()]);
+        let ranges = [0..fst_bytes.len() as u64];
+        let mut metrics = InvertedIndexReadMetrics::default();
+        let fsts = reader.fst_vec(&ranges, Some(&mut metrics)).await.unwrap();
+
+        assert_eq!(fsts.len(), 1);
+        let fst = &fsts[0];
+        // The map is backed by the exact buffer the reader returned instead of
+        // a copy of the FST payload: the pointer is compared while the caller
+        // side handle is still alive, so a copy could not alias it.
+        assert_eq!(fst.as_fst().as_inner().as_ptr(), data_ptr);
+        assert_eq!(fst.as_fst().as_bytes().as_ptr(), data_ptr);
+        assert_eq!(fst.as_fst().as_bytes().len(), fst_bytes.len());
+        assert_eq!(fst.get("key1"), Some(1));
+        assert_eq!(fst.get("key2"), Some(2));
+
+        // Dropping the caller-side references keeps the map usable: the map
+        // holds a shared handle to the same allocation.
+        drop(reader);
+        drop(fst_bytes);
+        assert_eq!(fst.get("key1"), Some(1));
+        assert_eq!(fst.get("key2"), Some(2));
+        assert_eq!(fst.get("key3"), None);
+    }
+
+    #[tokio::test]
+    async fn test_fst_vec_rejects_empty_and_malformed_bytes() {
+        for payload in [vec![], b"not an fst".to_vec()] {
+            let size = payload.len() as u64;
+            let mut reader = BytesReader(vec![Bytes::from(payload)]);
+            let mut metrics = InvertedIndexReadMetrics::default();
+            let result = reader.fst_vec(&[0..size], Some(&mut metrics)).await;
+            assert!(matches!(result, Err(Error::DecodeFst { .. })));
+        }
+    }
+}

src/index/src/inverted_index/search/fst_apply/intersection_apply.rs

@@ -201,9 +201,18 @@ mod tests {
         })])
     }
 
+    /// Builds an `FstMap` sharing its backing bytes with an owned buffer.
+    fn create_fst_map<K: AsRef<[u8]>>(items: &[(K, u64)]) -> FstMap {
+        let mut builder = fst::MapBuilder::memory();
+        for (key, value) in items {
+            builder.insert(key, *value).unwrap();
+        }
+        FstMap::new(bytes::Bytes::from(builder.into_inner().unwrap())).unwrap()
+    }
+
     #[test]
     fn test_intersection_fst_applier_with_ranges() {
-        let test_fst = FstMap::from_iter([("aa", 1), ("bb", 2), ("cc", 3)]).unwrap();
+        let test_fst = create_fst_map(&[("aa", 1), ("bb", 2), ("cc", 3)]);
 
         let applier_inclusive_lower = create_applier_from_range(Range {
             lower: Some(Bound {
@@ -280,7 +289,7 @@ mod tests {
 
     #[test]
     fn test_intersection_fst_applier_with_valid_pattern() {
-        let test_fst = FstMap::from_iter([("123", 1), ("abc", 2)]).unwrap();
+        let test_fst = create_fst_map(&[("123", 1), ("abc", 2)]);
 
         let cases = vec![
             ("1", vec![1]),
@@ -345,7 +354,7 @@ mod tests {
 
     #[test]
     fn test_intersection_fst_applier_with_composite_predicates() {
-        let test_fst = FstMap::from_iter([("aa", 1), ("bb", 2), ("cc", 3)]).unwrap();
+        let test_fst = create_fst_map(&[("aa", 1), ("bb", 2), ("cc", 3)]);
 
         let applier = IntersectionFstApplier::try_from(vec![
             Predicate::Range(RangePredicate {

src/index/src/inverted_index/search/fst_apply/keys_apply.rs

@@ -173,21 +173,22 @@ impl TryFrom<Vec<Predicate>> for KeysFstApplier {
 
 #[cfg(test)]
 mod tests {
-    use fst::Map as FstMap;
-
     use super::*;
     use crate::inverted_index::error::Error;
     use crate::inverted_index::search::predicate::{
         Bound, InListPredicate, Predicate, Range, RangePredicate, RegexMatchPredicate,
     };
 
-    fn create_fst_map(items: &[(&[u8], u64)]) -> FstMap<Vec<u8>> {
+    /// Builds an `FstMap` sharing its backing bytes with an owned buffer.
+    fn create_fst_map(items: &[(&[u8], u64)]) -> FstMap {
         let mut items = items
             .iter()
             .map(|(k, v)| (k.to_vec(), *v))
             .collect::<Vec<_>>();
         items.sort();
-        FstMap::from_iter(items).unwrap()
+        let mut builder = fst::MapBuilder::memory();
+        builder.extend_iter(items).unwrap();
+        FstMap::new(bytes::Bytes::from(builder.into_inner().unwrap())).unwrap()
     }
 
     fn b(s: &str) -> Vec<u8> {

src/index/src/inverted_index/search/index_apply/predicates_apply.rs

@@ -213,6 +213,15 @@ mod tests {
         bytemuck::cast::<_, u64>([offset, size])
     }
 
+    /// Builds an `FstMap` sharing its backing bytes with an owned buffer.
+    fn create_fst_map<K: AsRef<[u8]>>(items: &[(K, u64)]) -> FstMap {
+        let mut builder = fst::MapBuilder::memory();
+        for (key, value) in items {
+            builder.insert(key, *value).unwrap();
+        }
+        FstMap::new(bytes::Bytes::from(builder.into_inner().unwrap())).unwrap()
+    }
+
     #[tokio::test]
     async fn test_index_applier_apply_get_key() {
         // An index applier that point-gets "tag-0_value-0" on tag "tag-0"
@@ -226,9 +235,7 @@ mod tests {
             .expect_metadata()
             .returning(|_| Ok(mock_metas([("tag-0", 0)])));
         mock_reader.expect_fst_vec().returning(|_ranges, _metrics| {
-            Ok(vec![
-                FstMap::from_iter([(b"tag-0_value-0", fst_value(2, 1))]).unwrap(),
-            ])
+            Ok(vec![create_fst_map(&[(b"tag-0_value-0", fst_value(2, 1))])])
         });
 
         mock_reader
@@ -259,9 +266,7 @@ mod tests {
             .expect_metadata()
             .returning(|_| Ok(mock_metas([("tag-0", 0)])));
         mock_reader.expect_fst_vec().returning(|_range, _metrics| {
-            Ok(vec![
-                FstMap::from_iter([(b"tag-0_value-1", fst_value(2, 1))]).unwrap(),
-            ])
+            Ok(vec![create_fst_map(&[(b"tag-0_value-1", fst_value(2, 1))])])
         });
         let output = applier
             .apply(SearchContext::default(), &mut mock_reader, None)
@@ -289,10 +294,8 @@ mod tests {
             let mut output = vec![];
             for range in ranges {
                 match range.start {
-                    0 => output
-                        .push(FstMap::from_iter([(b"tag-0_value-0", fst_value(1, 1))]).unwrap()),
-                    1 => output
-                        .push(FstMap::from_iter([(b"tag-1_value-a", fst_value(2, 1))]).unwrap()),
+                    0 => output.push(create_fst_map(&[(b"tag-0_value-0", fst_value(1, 1))])),
+                    1 => output.push(create_fst_map(&[(b"tag-1_value-a", fst_value(2, 1))])),
                     _ => unreachable!(),
                 }
             }