From d25fc155b2f1c93957bb414d4b7c3508bf240f88 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Mon, 27 Feb 2023 15:34:47 +0900 Subject: [PATCH] Making some of the column/termdict operations async-friendly (#1902) --- columnar/src/column_index/mod.rs | 4 +- columnar/src/columnar/reader/mod.rs | 77 ++++++++++++------- columnar/src/dynamic_column.rs | 7 +- src/aggregation/bucket/term_agg.rs | 12 +-- .../phrase_prefix_query.rs | 4 +- .../phrase_prefix_weight.rs | 10 +-- sstable/src/dictionary.rs | 10 +++ sstable/src/streamer.rs | 47 ++++++++--- 8 files changed, 113 insertions(+), 58 deletions(-) diff --git a/columnar/src/column_index/mod.rs b/columnar/src/column_index/mod.rs index d1a5ae81e..fa7a0408b 100644 --- a/columnar/src/column_index/mod.rs +++ b/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) { match self { - ColumnIndex::Empty { .. } => {} + ColumnIndex::Empty { .. } => { + rank_ids.clear(); + } ColumnIndex::Full => { // No need to do anything: // value_idx and row_idx are the same. diff --git a/columnar/src/columnar/reader/mod.rs b/columnar/src/columnar/reader/mod.rs index 4d7677565..d11a4dfb7 100644 --- a/columnar/src/columnar/reader/mod.rs +++ b/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> { + 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(file_slice: F) -> io::Result @@ -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> { + fn stream_for_column_range(&self, column_name: &str) -> sstable::StreamerBuilder { // 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> { + 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> { + 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. diff --git a/columnar/src/dynamic_column.rs b/columnar/src/dynamic_column.rs index 0d3cfe0e5..675ad99e7 100644 --- a/columnar/src/dynamic_column.rs +++ b/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 { - 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 diff --git a/src/aggregation/bucket/term_agg.rs b/src/aggregation/bucket/term_agg.rs index 580c4a938..6f884c1ac 100644 --- a/src/aggregation/bucket/term_agg.rs +++ b/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 = 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 = 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; diff --git a/src/query/phrase_prefix_query/phrase_prefix_query.rs b/src/query/phrase_prefix_query/phrase_prefix_query.rs index 43cc33739..f00e8e50f 100644 --- a/src/query/phrase_prefix_query/phrase_prefix_query.rs +++ b/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); diff --git a/src/query/phrase_prefix_query/phrase_prefix_weight.rs b/src/query/phrase_prefix_query/phrase_prefix_weight.rs index f9292a327..acf964468 100644 --- a/src/query/phrase_prefix_query/phrase_prefix_weight.rs +++ b/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); } } diff --git a/sstable/src/dictionary.rs b/sstable/src/dictionary.rs index 272aecf40..b907d55ba 100644 --- a/sstable/src/dictionary.rs +++ b/sstable/src/dictionary.rs @@ -76,6 +76,16 @@ impl Dictionary { Ok(TSSTable::reader(data)) } + pub(crate) async fn sstable_delta_reader_for_key_range_async( + &self, + key_range: impl RangeBounds<[u8]>, + limit: Option, + ) -> io::Result> { + 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]>, diff --git a/sstable/src/streamer.rs b/sstable/src/streamer.rs index 889f0a4b4..32414910e 100644 --- a/sstable/src/streamer.rs +++ b/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> { - // TODO Optimize by skipping to the right first block. - let start_state = self.automaton.start(); - + fn delta_reader(&self) -> io::Result> { 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> { + 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, ::ValueReader>, + ) -> io::Result> { + 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> { + 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> { + 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.