diff --git a/src/index/src/inverted_index.rs b/src/index/src/inverted_index.rs index 9dc5c87014d..06672897621 100644 --- a/src/index/src/inverted_index.rs +++ b/src/index/src/inverted_index.rs @@ -17,4 +17,7 @@ pub mod error; pub mod format; pub mod search; -pub type FstMap = fst::Map>; +/// 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; diff --git a/src/index/src/inverted_index/format/reader.rs b/src/index/src/inverted_index/format/reader.rs index ff67284e51a..ae5443089f4 100644 --- a/src/index/src/inverted_index/format/reader.rs +++ b/src/index/src/inverted_index/format/reader.rs @@ -129,7 +129,8 @@ pub trait InvertedIndexReader: Send + Sync { metrics: Option<&'a mut InvertedIndexReadMetrics>, ) -> Result { 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], metrics: Option<&'a mut InvertedIndexReadMetrics>, ) -> Result> { + // 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::>>() } @@ -177,3 +180,90 @@ pub trait InvertedIndexReader: Send + Sync { .collect::>>() } } + +#[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); + + #[async_trait] + impl InvertedIndexReader for BytesReader { + async fn range_read<'a>( + &self, + _offset: u64, + _size: u32, + _metrics: Option<&'a mut InvertedIndexReadMetrics>, + ) -> Result> { + unimplemented!("only `read_vec` is used by the tests") + } + + async fn read_vec<'a>( + &self, + _ranges: &[Range], + _metrics: Option<&'a mut InvertedIndexReadMetrics>, + ) -> Result> { + Ok(self.0.clone()) + } + + async fn metadata<'a>( + &self, + _metrics: Option<&'a mut InvertedIndexReadMetrics>, + ) -> Result> { + unimplemented!("only `read_vec` is used by the tests") + } + } + + fn mock_fst_bytes() -> Vec { + 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 { .. }))); + } + } +} diff --git a/src/index/src/inverted_index/search/fst_apply/intersection_apply.rs b/src/index/src/inverted_index/search/fst_apply/intersection_apply.rs index 7b551808e3c..368dc43e88d 100644 --- a/src/index/src/inverted_index/search/fst_apply/intersection_apply.rs +++ b/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>(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 { diff --git a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs index b68f58cdfd3..6c958e2023b 100644 --- a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs +++ b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs @@ -173,21 +173,22 @@ impl TryFrom> 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> { + /// 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::>(); 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 { diff --git a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs index 5fcc962ca7d..ec2f2af93ba 100644 --- a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs +++ b/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>(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!(), } }