WIP: remove ids buffering and add support for json

Author: imotov

src/fastfield/readers.rs

@@ -234,6 +234,22 @@ impl FastFieldReaders {
         Ok(dynamic_column_handle_opt)
     }
 
+    /// Returning all `dynamic_column_handle`.
+    pub fn dynamic_column_handles(
+        &self,
+        field_name: &str,
+    ) -> crate::Result<Vec<DynamicColumnHandle>> {
+        let Some(resolved_field_name) = self.resolve_field(field_name)? else {
+            return Ok(Vec::new());
+        };
+        let dynamic_column_handles = self
+            .columnar
+            .read_columns(&resolved_field_name)?
+            .into_iter()
+            .collect();
+        Ok(dynamic_column_handles)
+    }
+
     #[doc(hidden)]
     pub async fn list_dynamic_column_handles(
         &self,

src/query/exist_query.rs

@@ -1,22 +1,20 @@
 use core::fmt::Debug;
 
-use columnar::{ColumnType, DynamicColumn};
+use columnar::{ColumnIndex, DynamicColumn};
 
-use super::range_query::VecCursor;
-use super::ConstScorer;
+use super::{ConstScorer, EmptyScorer};
 use crate::core::SegmentReader;
 use crate::docset::{DocSet, TERMINATED};
 use crate::query::explanation::does_not_match;
 use crate::query::{EnableScoring, Explanation, Query, Scorer, Weight};
-use crate::schema::FieldType;
 use crate::{DocId, Score, TantivyError};
 
 /// Query that matches all of the documents.
 ///
 /// All of the document get the score 1.0.
 #[derive(Clone, Debug)]
 pub struct ExistsQuery {
-    field: String,
+    field_name: String,
 }
 
 impl ExistsQuery {
@@ -25,62 +23,55 @@ impl ExistsQuery {
     /// If the value type is not correct, something may go terribly wrong when
     /// the `Weight` object is created.
     pub fn new_exists_query(field: String) -> ExistsQuery {
-        ExistsQuery { field }
+        ExistsQuery { field_name: field }
     }
 }
 
 impl Query for ExistsQuery {
     fn weight(&self, enable_scoring: EnableScoring) -> crate::Result<Box<dyn Weight>> {
         let schema = enable_scoring.schema();
-        let field_type = schema
-            .get_field_entry(schema.get_field(&self.field)?)
-            .field_type();
+        let Some((field, _path)) = schema.find_field(&self.field_name) else {
+            return Err(TantivyError::FieldNotFound(self.field_name.clone()));
+        };
+        let field_type = schema.get_field_entry(field).field_type();
         if !field_type.is_fast() {
             return Err(TantivyError::SchemaError(format!(
-                "Field {:?} is not a fast field.",
-                self.field
+                "Field {} is not a fast field.",
+                self.field_name
             )));
         }
         Ok(Box::new(ExistsWeight {
-            field: self.field.clone(),
-            field_type: field_type.clone(),
+            field_name: self.field_name.clone(),
         }))
     }
 }
 
 /// Weight associated with the `ExistsQuery` query.
 pub struct ExistsWeight {
-    field: String,
-    field_type: FieldType,
+    field_name: String,
 }
 
 impl ExistsWeight {}
 
 impl Weight for ExistsWeight {
     fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
         let fast_field_reader = reader.fast_fields();
-        let column_type = match self.field_type {
-            FieldType::Str(_) => Some(ColumnType::Str),
-            FieldType::U64(_) => Some(ColumnType::U64),
-            FieldType::I64(_) => Some(ColumnType::I64),
-            FieldType::F64(_) => Some(ColumnType::F64),
-            FieldType::Bool(_) => Some(ColumnType::Bool),
-            FieldType::Date(_) => Some(ColumnType::DateTime),
-            FieldType::Bytes(_) => Some(ColumnType::Bytes),
-            FieldType::IpAddr(_) => Some(ColumnType::IpAddr),
-            _ => None,
+        let dynamic_columns: crate::Result<Vec<DynamicColumn>> = fast_field_reader
+            .dynamic_column_handles(&self.field_name)?
+            .into_iter()
+            .map(|handle| handle.open().map_err(|io_error| io_error.into()))
+            .collect();
+        let mut non_empty_columns = Vec::new();
+        for column in dynamic_columns? {
+            if !matches!(column.column_index(), ColumnIndex::Empty { .. }) {
+                non_empty_columns.push(column)
+            }
         }
-        .expect("Should be here");
-        if let Some(dynamic_column) =
-            fast_field_reader.dynamic_column_handle(&self.field, column_type)?
-        {
-            let docset = ExistsDocSet::new(dynamic_column.open()?);
+        if !non_empty_columns.is_empty() {
+            let docset = ExistsDocSet::new(non_empty_columns, reader.max_doc());
             return Ok(Box::new(ConstScorer::new(docset, boost)));
         } else {
-            return Err(TantivyError::SchemaError(format!(
-                "Field {:?} with type {:?} is not supported by exists query.",
-                self.field, self.field_type,
-            )));
+            return Ok(Box::new(EmptyScorer));
         }
     }
 
@@ -94,147 +85,64 @@ impl Weight for ExistsWeight {
 }
 
 pub(crate) struct ExistsDocSet {
-    column: DynamicColumn,
-    /// The next docid start range to fetch (inclusive).
-    next_fetch_start: u32,
-    /// Number of docs range checked in a batch.
-    ///
-    /// There are two patterns.
-    /// - We do a full scan. => We can load large chunks. We don't know in advance if seek call
-    /// will come, so we start with small chunks
-    /// - We load docs, interspersed with seek calls. When there are big jumps in the seek, we
-    /// should load small chunks. When the seeks are small, we can employ the same strategy as on a
-    /// full scan.
-    fetch_horizon: u32,
-    /// Current batch of loaded docs.
-    loaded_docs: VecCursor,
-    last_seek_pos_opt: Option<u32>,
+    columns: Vec<DynamicColumn>,
+    doc: DocId,
+    max_doc: DocId,
 }
 
-const DEFAULT_FETCH_HORIZON: u32 = 128;
 impl ExistsDocSet {
-    pub(crate) fn new(column: DynamicColumn) -> Self {
-        let mut exists_doc_set = Self {
-            column,
-            loaded_docs: VecCursor::new(),
-            next_fetch_start: 0,
-            fetch_horizon: DEFAULT_FETCH_HORIZON,
-            last_seek_pos_opt: None,
+    pub(crate) fn new(columns: Vec<DynamicColumn>, max_doc: DocId) -> Self {
+        let mut set = Self {
+            columns,
+            doc: 0u32,
+            max_doc,
         };
-        exists_doc_set.reset_fetch_range();
-        exists_doc_set.fetch_block();
-        exists_doc_set
+        set.find_next();
+        set
     }
 
-    /// Returns true if more data could be fetched
-    fn fetch_block(&mut self) {
-        const MAX_HORIZON: u32 = 100_000;
-        while self.loaded_docs.is_empty() {
-            let finished_to_end = self.fetch_horizon(self.fetch_horizon);
-            if finished_to_end {
-                break;
+    fn find_next(&mut self) -> DocId {
+        // TODO: can this be optimized?
+        while self.doc < self.max_doc {
+            if self
+                .columns
+                .iter()
+                .find(|col| col.column_index().has_value(self.doc))
+                .is_some()
+            {
+                return self.doc;
             }
-            // Fetch more data, increase horizon. Horizon only gets reset when doing a seek.
-            self.fetch_horizon = (self.fetch_horizon * 2).min(MAX_HORIZON);
+            self.doc += 1;
         }
-    }
-
-    /// Fetches a block for docid range [next_fetch_start .. next_fetch_start + HORIZON]
-    fn fetch_horizon(&mut self, horizon: u32) -> bool {
-        let mut finished_to_end = false;
-
-        let limit = self.column.num_docs();
-        let mut end = self.next_fetch_start + horizon;
-        if end >= limit {
-            end = limit;
-            finished_to_end = true;
-        }
-
-        let last_value = self.loaded_docs.last_value();
-        let doc_buffer: &mut Vec<DocId> = self.loaded_docs.get_cleared_data();
-        self.column
-            .column_index()
-            .get_docids_with_existing_value(self.next_fetch_start..end, doc_buffer);
-        if let Some(last_value) = last_value {
-            while self.loaded_docs.current() == Some(last_value) {
-                self.loaded_docs.next();
-            }
-        }
-        self.next_fetch_start = end;
-
-        finished_to_end
-    }
-
-    /// check if the distance between the seek calls is large
-    fn is_last_seek_distance_large(&self, new_seek: DocId) -> bool {
-        if let Some(last_seek_pos) = self.last_seek_pos_opt {
-            (new_seek - last_seek_pos) >= 128
-        } else {
-            true
-        }
-    }
-
-    fn reset_fetch_range(&mut self) {
-        self.fetch_horizon = DEFAULT_FETCH_HORIZON;
+        self.doc = TERMINATED;
+        return TERMINATED;
     }
 }
 
 impl DocSet for ExistsDocSet {
     fn advance(&mut self) -> DocId {
-        if let Some(docid) = self.loaded_docs.next() {
-            return docid;
-        }
-        if self.next_fetch_start >= self.column.num_values() {
-            return TERMINATED;
-        }
-        self.fetch_block();
-        self.loaded_docs.current().unwrap_or(TERMINATED)
-    }
-
-    #[inline]
-    fn doc(&self) -> DocId {
-        self.loaded_docs.current().unwrap_or(TERMINATED)
-    }
-
-    /// Advances the `DocSet` forward until reaching the target, or going to the
-    /// lowest [`DocId`] greater than the target.
-    ///
-    /// If the end of the `DocSet` is reached, [`TERMINATED`] is returned.
-    ///
-    /// Calling `.seek(target)` on a terminated `DocSet` is legal. Implementation
-    /// of `DocSet` should support it.
-    ///
-    /// Calling `seek(TERMINATED)` is also legal and is the normal way to consume a `DocSet`.
-    fn seek(&mut self, target: DocId) -> DocId {
-        if self.is_last_seek_distance_large(target) {
-            self.reset_fetch_range();
-        }
-        if target > self.next_fetch_start {
-            self.next_fetch_start = target;
-        }
-        let mut doc = self.doc();
-        debug_assert!(doc <= target);
-        while doc < target {
-            doc = self.advance();
-        }
-        self.last_seek_pos_opt = Some(target);
-        doc
+        self.doc += 1;
+        self.find_next()
     }
 
     fn size_hint(&self) -> u32 {
         0 // heuristic possible by checking number of hits when fetching a block
     }
+
+    fn doc(&self) -> DocId {
+        self.doc
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use crate::collector::Count;
     use crate::query::exist_query::ExistsQuery;
-    use crate::schema::{Schema, FAST, INDEXED, STRING};
+    use crate::schema::{Schema, FAST, INDEXED, STRING, TEXT};
     use crate::{doc, Index};
 
     #[test]
-    fn test_range_query_simple() -> crate::Result<()> {
+    fn test_exists_query_simple() -> crate::Result<()> {
         let mut schema_builder = Schema::builder();
         let all_field = schema_builder.add_u64_field("all", INDEXED | FAST);
         let even_field = schema_builder.add_u64_field("even", INDEXED | FAST);
@@ -284,4 +192,41 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_exists_query_json() -> crate::Result<()> {
+        let mut schema_builder = Schema::builder();
+        let json = schema_builder.add_json_field("json", TEXT | FAST);
+        let schema = schema_builder.build();
+
+        let index = Index::create_in_ram(schema);
+        {
+            let mut index_writer = index.writer_for_tests()?;
+            for i in 0u64..100u64 {
+                if i % 2 == 0 {
+                    index_writer.add_document(doc!(json => json!({"all": i, "even": true})))?;
+                } else {
+                    index_writer
+                        .add_document(doc!(json => json!({"all": i.to_string(), "odd": true})))?;
+                }
+            }
+            index_writer.commit()?;
+        }
+        let reader = index.reader()?;
+        let searcher = reader.searcher();
+
+        let alldocs = ExistsQuery::new_exists_query("json.all".to_string());
+        let count = searcher.search(&alldocs, &Count)?;
+        assert_eq!(count, 100);
+
+        let even_docs = ExistsQuery::new_exists_query("json.even".to_string());
+        let count = searcher.search(&even_docs, &Count)?;
+        assert_eq!(count, 50);
+
+        let odd_docs = ExistsQuery::new_exists_query("json.odd".to_string());
+        let count = searcher.search(&odd_docs, &Count)?;
+        assert_eq!(count, 50);
+
+        Ok(())
+    }
 }

src/query/range_query/fast_field_range_query.rs

@@ -7,34 +7,34 @@ use crate::{DocId, DocSet, TERMINATED};
 
 /// Helper to have a cursor over a vec of docids
 #[derive(Debug)]
-pub(crate) struct VecCursor {
+struct VecCursor {
     docs: Vec<u32>,
     current_pos: usize,
 }
 impl VecCursor {
-    pub fn new() -> Self {
+    fn new() -> Self {
         Self {
             docs: Vec::with_capacity(32),
             current_pos: 0,
         }
     }
-    pub fn next(&mut self) -> Option<u32> {
+    fn next(&mut self) -> Option<u32> {
         self.current_pos += 1;
         self.current()
     }
     #[inline]
-    pub fn current(&self) -> Option<u32> {
+    fn current(&self) -> Option<u32> {
         self.docs.get(self.current_pos).copied()
     }
-    pub fn get_cleared_data(&mut self) -> &mut Vec<u32> {
+    fn get_cleared_data(&mut self) -> &mut Vec<u32> {
         self.docs.clear();
         self.current_pos = 0;
         &mut self.docs
     }
-    pub fn last_value(&self) -> Option<u32> {
+    fn last_value(&self) -> Option<u32> {
         self.docs.iter().last().cloned()
     }
-    pub fn is_empty(&self) -> bool {
+    fn is_empty(&self) -> bool {
         self.current().is_none()
     }
 }