Making some of the column/termdict operations async-friendly (#1902)

columnar/src/column_index/mod.rs

@@ -97,7 +97,9 @@ impl ColumnIndex {
 
     pub fn select_batch_in_place(&self, doc_id_start: DocId, rank_ids: &mut Vec<RowId>) {
         match self {
-            ColumnIndex::Empty { .. } => {}
+            ColumnIndex::Empty { .. } => {
+                rank_ids.clear();
+            }
             ColumnIndex::Full => {
                 // No need to do anything:
                 // value_idx and row_idx are the same.

columnar/src/columnar/reader/mod.rs

@@ -21,6 +21,32 @@ pub struct ColumnarReader {
     num_rows: RowId,
 }
 
+/// Functions by both the async/sync code listing columns.
+/// It takes a stream from the column sstable and return the list of
+/// `DynamicColumn` available in it.
+fn read_all_columns_in_stream(
+    mut stream: sstable::Streamer<'_, RangeSSTable>,
+    column_data: &FileSlice,
+) -> io::Result<Vec<DynamicColumnHandle>> {
+    let mut results = Vec::new();
+    while stream.advance() {
+        let key_bytes: &[u8] = stream.key();
+        let Some(column_code) = key_bytes.last().copied() else {
+            return Err(io_invalid_data("Empty column name.".to_string()));
+        };
+        let column_type = ColumnType::try_from_code(column_code)
+            .map_err(|_| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
+        let range = stream.value();
+        let file_slice = column_data.slice(range.start as usize..range.end as usize);
+        let dynamic_column_handle = DynamicColumnHandle {
+            file_slice,
+            column_type,
+        };
+        results.push(dynamic_column_handle);
+    }
+    Ok(results)
+}
+
 impl ColumnarReader {
     /// Opens a new Columnar file.
     pub fn open<F>(file_slice: F) -> io::Result<ColumnarReader>
@@ -76,11 +102,7 @@ impl ColumnarReader {
         Ok(results)
     }
 
-    /// Get all columns for the given column name.
-    ///
-    /// There can be more than one column associated to a given column name, provided they have
-    /// different types.
-    pub fn read_columns(&self, column_name: &str) -> io::Result<Vec<DynamicColumnHandle>> {
+    fn stream_for_column_range(&self, column_name: &str) -> sstable::StreamerBuilder<RangeSSTable> {
         // Each column is a associated to a given `column_key`,
         // that starts by `column_name\0column_header`.
         //
@@ -89,36 +111,35 @@ impl ColumnarReader {
         //
         // This is in turn equivalent to searching for the range
         // `[column_name,\0`..column_name\1)`.
-
-        // TODO can we get some more generic `prefix(..)` logic in the dictioanry.
+        // TODO can we get some more generic `prefix(..)` logic in the dictionary.
         let mut start_key = column_name.to_string();
         start_key.push('\0');
         let mut end_key = column_name.to_string();
         end_key.push(1u8 as char);
-        let mut stream = self
-            .column_dictionary
+        self.column_dictionary
             .range()
             .ge(start_key.as_bytes())
             .lt(end_key.as_bytes())
-            .into_stream()?;
-        let mut results = Vec::new();
-        while stream.advance() {
-            let key_bytes: &[u8] = stream.key();
-            assert!(key_bytes.starts_with(start_key.as_bytes()));
-            let column_code: u8 = key_bytes.last().cloned().unwrap();
-            let column_type = ColumnType::try_from_code(column_code)
-                .map_err(|_| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
-            let range = stream.value().clone();
-            let file_slice = self
-                .column_data
-                .slice(range.start as usize..range.end as usize);
-            let dynamic_column_handle = DynamicColumnHandle {
-                file_slice,
-                column_type,
-            };
-            results.push(dynamic_column_handle);
-        }
-        Ok(results)
+    }
+
+    pub async fn read_columns_async(
+        &self,
+        column_name: &str,
+    ) -> io::Result<Vec<DynamicColumnHandle>> {
+        let stream = self
+            .stream_for_column_range(column_name)
+            .into_stream_async()
+            .await?;
+        read_all_columns_in_stream(stream, &self.column_data)
+    }
+
+    /// Get all columns for the given column name.
+    ///
+    /// There can be more than one column associated to a given column name, provided they have
+    /// different types.
+    pub fn read_columns(&self, column_name: &str) -> io::Result<Vec<DynamicColumnHandle>> {
+        let stream = self.stream_for_column_range(column_name).into_stream()?;
+        read_all_columns_in_stream(stream, &self.column_data)
     }
 
     /// Return the number of columns in the columnar.

columnar/src/dynamic_column.rs

@@ -206,10 +206,9 @@ impl DynamicColumnHandle {
         self.open_internal(column_bytes)
     }
 
-    // TODO rename load_async
-    pub async fn open_async(&self) -> io::Result<DynamicColumn> {
-        let column_bytes: OwnedBytes = self.file_slice.read_bytes_async().await?;
-        self.open_internal(column_bytes)
+    #[doc(hidden)]
+    pub fn file_slice(&self) -> &FileSlice {
+        &self.file_slice
     }
 
     /// Returns the `u64` fast field reader reader associated with `fields` of types

src/aggregation/bucket/term_agg.rs

@@ -423,16 +423,16 @@ impl SegmentTermCollector {
             cut_off_buckets(&mut entries, self.req.segment_size as usize)
         };
 
-        let inverted_index = agg_with_accessor
+        let mut dict: FxHashMap<String, IntermediateTermBucketEntry> = Default::default();
+
+        let str_column = agg_with_accessor
             .str_dict_column
             .as_ref()
-            .expect("internal error: inverted index not loaded for term aggregation");
-        let term_dict = inverted_index;
+            .expect("Missing str column"); //< TODO Fixme
 
-        let mut dict: FxHashMap<String, IntermediateTermBucketEntry> = Default::default();
         let mut buffer = String::new();
         for (term_id, entry) in entries {
-            if !term_dict.ord_to_str(term_id as u64, &mut buffer)? {
+            if !str_column.ord_to_str(term_id as u64, &mut buffer)? {
                 return Err(TantivyError::InternalError(format!(
                     "Couldn't find term_id {} in dict",
                     term_id
@@ -445,7 +445,7 @@ impl SegmentTermCollector {
         }
         if self.req.min_doc_count == 0 {
             // TODO: Handle rev streaming for descending sorting by keys
-            let mut stream = term_dict.dictionary().stream()?;
+            let mut stream = str_column.dictionary().stream()?;
             while let Some((key, _ord)) = stream.next() {
                 if dict.len() >= self.req.segment_size as usize {
                     break;

src/query/phrase_prefix_query/phrase_prefix_query.rs

@@ -88,7 +88,7 @@ impl PhrasePrefixQuery {
     /// a specialized type [`PhraseQueryWeight`] instead of a Boxed trait.
     /// If the query was only one term long, this returns `None` wherease [`Query::weight`]
     /// returns a boxed [`RangeWeight`]
-    ///    
+    ///
     /// Returns `None`, if phrase_terms is empty, which happens if the phrase prefix query was
     /// built with a single term.
     pub(crate) fn phrase_prefix_query_weight(
@@ -138,7 +138,7 @@ impl Query for PhrasePrefixQuery {
             Ok(Box::new(phrase_weight))
         } else {
             // There are no prefix. Let's just match the suffix.
-            let end_term = if let Some(end_value) = prefix_end(&self.prefix.1.value_bytes()) {
+            let end_term = if let Some(end_value) = prefix_end(self.prefix.1.value_bytes()) {
                 let mut end_term = Term::with_capacity(end_value.len());
                 end_term.set_field_and_type(self.field, self.prefix.1.typ());
                 end_term.append_bytes(&end_value);

src/query/phrase_prefix_query/phrase_prefix_weight.rs

@@ -106,12 +106,10 @@ impl PhrasePrefixWeight {
                 {
                     suffixes.push(postings);
                 }
-            } else {
-                if let Some(postings) = inv_index
-                    .read_postings_no_deletes(&new_term, IndexRecordOption::WithFreqsAndPositions)?
-                {
-                    suffixes.push(postings);
-                }
+            } else if let Some(postings) = inv_index
+                .read_postings_no_deletes(&new_term, IndexRecordOption::WithFreqsAndPositions)?
+            {
+                suffixes.push(postings);
             }
         }
 

sstable/src/dictionary.rs

@@ -76,6 +76,16 @@ impl<TSSTable: SSTable> Dictionary<TSSTable> {
         Ok(TSSTable::reader(data))
     }
 
+    pub(crate) async fn sstable_delta_reader_for_key_range_async(
+        &self,
+        key_range: impl RangeBounds<[u8]>,
+        limit: Option<u64>,
+    ) -> io::Result<DeltaReader<'static, TSSTable::ValueReader>> {
+        let slice = self.file_slice_for_range(key_range, limit);
+        let data = slice.read_bytes_async().await?;
+        Ok(TSSTable::delta_reader(data))
+    }
+
     pub(crate) fn sstable_delta_reader_for_key_range(
         &self,
         key_range: impl RangeBounds<[u8]>,

sstable/src/streamer.rs

@@ -5,7 +5,7 @@ use tantivy_fst::automaton::AlwaysMatch;
 use tantivy_fst::Automaton;
 
 use crate::dictionary::Dictionary;
-use crate::{SSTable, TermOrdinal};
+use crate::{DeltaReader, SSTable, TermOrdinal};
 
 /// `StreamerBuilder` is a helper object used to define
 /// a range of terms that should be streamed.
@@ -80,18 +80,33 @@ where
         self
     }
 
-    /// Creates the stream corresponding to the range
-    /// of terms defined using the `StreamerBuilder`.
-    pub fn into_stream(self) -> io::Result<Streamer<'a, TSSTable, A>> {
-        // TODO Optimize by skipping to the right first block.
-        let start_state = self.automaton.start();
-
+    fn delta_reader(&self) -> io::Result<DeltaReader<'a, TSSTable::ValueReader>> {
         let key_range = (
             bound_as_byte_slice(&self.lower),
             bound_as_byte_slice(&self.upper),
         );
+        self.term_dict
+            .sstable_delta_reader_for_key_range(key_range, self.limit)
+    }
 
-        let first_term = match &key_range.0 {
+    async fn delta_reader_async(&self) -> io::Result<DeltaReader<'a, TSSTable::ValueReader>> {
+        let key_range = (
+            bound_as_byte_slice(&self.lower),
+            bound_as_byte_slice(&self.upper),
+        );
+        self.term_dict
+            .sstable_delta_reader_for_key_range_async(key_range, self.limit)
+            .await
+    }
+
+    fn into_stream_given_delta_reader(
+        self,
+        delta_reader: DeltaReader<'a, <TSSTable as SSTable>::ValueReader>,
+    ) -> io::Result<Streamer<'a, TSSTable, A>> {
+        let start_state = self.automaton.start();
+        let start_key = bound_as_byte_slice(&self.lower);
+
+        let first_term = match start_key {
             Bound::Included(key) | Bound::Excluded(key) => self
                 .term_dict
                 .sstable_index
@@ -101,9 +116,6 @@ where
             Bound::Unbounded => 0,
         };
 
-        let delta_reader = self
-            .term_dict
-            .sstable_delta_reader_for_key_range(key_range, self.limit)?;
         Ok(Streamer {
             automaton: self.automaton,
             states: vec![start_state],
@@ -114,6 +126,19 @@ where
             upper_bound: self.upper,
         })
     }
+
+    /// See `into_stream(..)`
+    pub async fn into_stream_async(self) -> io::Result<Streamer<'a, TSSTable, A>> {
+        let delta_reader = self.delta_reader_async().await?;
+        self.into_stream_given_delta_reader(delta_reader)
+    }
+
+    /// Creates the stream corresponding to the range
+    /// of terms defined using the `StreamerBuilder`.
+    pub fn into_stream(self) -> io::Result<Streamer<'a, TSSTable, A>> {
+        let delta_reader = self.delta_reader()?;
+        self.into_stream_given_delta_reader(delta_reader)
+    }
 }
 
 /// `Streamer` acts as a cursor over a range of terms of a segment.