mito2/sst/index/fulltext_index/
creator.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::sync::Arc;
17use std::sync::atomic::AtomicUsize;
18
19use api::v1::SemanticType;
20use common_telemetry::warn;
21use datatypes::arrow::array::{Array, LargeStringArray, StringArray};
22use datatypes::arrow::datatypes::DataType;
23use datatypes::arrow::record_batch::RecordBatch;
24use datatypes::schema::{FulltextAnalyzer, FulltextBackend};
25use index::fulltext_index::create::{
26    BloomFilterFulltextIndexCreator, FulltextIndexCreator, TantivyFulltextIndexCreator,
27};
28use index::fulltext_index::{Analyzer, Config};
29use index::target::IndexTarget;
30use puffin::blob_metadata::CompressionCodec;
31use puffin::puffin_manager::PutOptions;
32use snafu::{ResultExt, ensure};
33use store_api::metadata::RegionMetadataRef;
34use store_api::storage::{ColumnId, ConcreteDataType, FileId, RegionId};
35
36use crate::error::{
37    CastVectorSnafu, ComputeArrowSnafu, CreateFulltextCreatorSnafu, DataTypeMismatchSnafu,
38    FulltextFinishSnafu, FulltextPushTextSnafu, IndexOptionsSnafu, OperateAbortedIndexSnafu,
39    Result,
40};
41use crate::read::Batch;
42use crate::sst::index::TYPE_FULLTEXT_INDEX;
43use crate::sst::index::fulltext_index::{INDEX_BLOB_TYPE_BLOOM, INDEX_BLOB_TYPE_TANTIVY};
44use crate::sst::index::intermediate::{
45    IntermediateLocation, IntermediateManager, TempFileProvider,
46};
47use crate::sst::index::puffin_manager::SstPuffinWriter;
48use crate::sst::index::statistics::{ByteCount, RowCount, Statistics};
49
50/// `FulltextIndexer` is responsible for creating fulltext indexes for SST files.
51pub struct FulltextIndexer {
52    /// Creators for each column.
53    creators: HashMap<ColumnId, SingleCreator>,
54    /// Whether the index creation was aborted.
55    aborted: bool,
56    /// Statistics of index creation.
57    stats: Statistics,
58}
59
60impl FulltextIndexer {
61    /// Creates a new `FulltextIndexer`.
62    pub async fn new(
63        region_id: &RegionId,
64        sst_file_id: &FileId,
65        intermediate_manager: &IntermediateManager,
66        metadata: &RegionMetadataRef,
67        compress: bool,
68        mem_limit: usize,
69    ) -> Result<Option<Self>> {
70        let mut creators = HashMap::new();
71
72        for column in &metadata.column_metadatas {
73            // Tag columns don't support fulltext index now.
74            // If we need to support fulltext index for tag columns, we also need to parse
75            // the codec and handle sparse encoding for flat format specially.
76            if column.semantic_type == SemanticType::Tag {
77                common_telemetry::debug!(
78                    "Skip creating fulltext index for tag column {}",
79                    column.column_schema.name
80                );
81                continue;
82            }
83
84            let options = column
85                .column_schema
86                .fulltext_options()
87                .context(IndexOptionsSnafu {
88                    column_name: &column.column_schema.name,
89                })?;
90
91            // Relax the type constraint here as many types can be casted to string.
92
93            let options = match options {
94                Some(options) if options.enable => options,
95                _ => continue,
96            };
97
98            let column_id = column.column_id;
99            let intm_path = intermediate_manager.fulltext_path(region_id, sst_file_id, column_id);
100
101            let config = Config {
102                analyzer: match options.analyzer {
103                    FulltextAnalyzer::English => Analyzer::English,
104                    FulltextAnalyzer::Chinese => Analyzer::Chinese,
105                },
106                case_sensitive: options.case_sensitive,
107            };
108
109            let inner = match options.backend {
110                FulltextBackend::Tantivy => {
111                    let creator = TantivyFulltextIndexCreator::new(&intm_path, config, mem_limit)
112                        .await
113                        .context(CreateFulltextCreatorSnafu)?;
114                    AltFulltextCreator::Tantivy(creator)
115                }
116                FulltextBackend::Bloom => {
117                    let temp_file_provider = Arc::new(TempFileProvider::new(
118                        IntermediateLocation::new(&metadata.region_id, sst_file_id),
119                        intermediate_manager.clone(),
120                    ));
121                    let global_memory_usage = Arc::new(AtomicUsize::new(0));
122                    let creator = BloomFilterFulltextIndexCreator::new(
123                        config,
124                        options.granularity as _,
125                        options.false_positive_rate(),
126                        temp_file_provider,
127                        global_memory_usage,
128                        Some(mem_limit),
129                    );
130                    AltFulltextCreator::Bloom(creator)
131                }
132            };
133
134            creators.insert(
135                column_id,
136                SingleCreator {
137                    column_id,
138                    column_name: column.column_schema.name.clone(),
139                    inner,
140                    compress,
141                },
142            );
143        }
144
145        Ok((!creators.is_empty()).then(move || Self {
146            creators,
147            aborted: false,
148            stats: Statistics::new(TYPE_FULLTEXT_INDEX),
149        }))
150    }
151
152    /// Updates the index with the given batch.
153    pub async fn update(&mut self, batch: &mut Batch) -> Result<()> {
154        ensure!(!self.aborted, OperateAbortedIndexSnafu);
155
156        if let Err(update_err) = self.do_update(batch).await {
157            if let Err(err) = self.do_abort().await {
158                if cfg!(any(test, feature = "test")) {
159                    panic!("Failed to abort index creator, err: {err}");
160                } else {
161                    warn!(err; "Failed to abort index creator");
162                }
163            }
164            return Err(update_err);
165        }
166
167        Ok(())
168    }
169
170    /// Updates the fulltext index with the given flat format RecordBatch.
171    pub async fn update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
172        ensure!(!self.aborted, OperateAbortedIndexSnafu);
173
174        if batch.num_rows() == 0 {
175            return Ok(());
176        }
177
178        if let Err(update_err) = self.do_update_flat(batch).await {
179            if let Err(err) = self.do_abort().await {
180                if cfg!(any(test, feature = "test")) {
181                    panic!("Failed to abort index creator, err: {err}");
182                } else {
183                    warn!(err; "Failed to abort index creator");
184                }
185            }
186            return Err(update_err);
187        }
188
189        Ok(())
190    }
191
192    /// Finalizes the index creation.
193    pub async fn finish(
194        &mut self,
195        puffin_writer: &mut SstPuffinWriter,
196    ) -> Result<(RowCount, ByteCount)> {
197        ensure!(!self.aborted, OperateAbortedIndexSnafu);
198
199        match self.do_finish(puffin_writer).await {
200            Ok(()) => Ok((self.stats.row_count(), self.stats.byte_count())),
201            Err(finish_err) => {
202                if let Err(err) = self.do_abort().await {
203                    if cfg!(any(test, feature = "test")) {
204                        panic!("Failed to abort index creator, err: {err}");
205                    } else {
206                        warn!(err; "Failed to abort index creator");
207                    }
208                }
209                Err(finish_err)
210            }
211        }
212    }
213
214    /// Aborts the index creation.
215    pub async fn abort(&mut self) -> Result<()> {
216        if self.aborted {
217            return Ok(());
218        }
219
220        self.do_abort().await
221    }
222
223    /// Returns the memory usage of the index creator.
224    pub fn memory_usage(&self) -> usize {
225        self.creators.values().map(|c| c.inner.memory_usage()).sum()
226    }
227
228    /// Returns IDs of columns that the creator is responsible for.
229    pub fn column_ids(&self) -> impl Iterator<Item = ColumnId> + '_ {
230        self.creators.keys().copied()
231    }
232
233    /// Returns IDs of columns indexed with the bloom backend.
234    pub fn bloom_column_ids(&self) -> impl Iterator<Item = ColumnId> + '_ {
235        self.creators
236            .values()
237            .filter(|c| matches!(c.inner, AltFulltextCreator::Bloom(_)))
238            .map(|c| c.column_id)
239    }
240}
241
242impl FulltextIndexer {
243    async fn do_update(&mut self, batch: &mut Batch) -> Result<()> {
244        let mut guard = self.stats.record_update();
245        guard.inc_row_count(batch.num_rows());
246
247        for creator in self.creators.values_mut() {
248            creator.update(batch).await?;
249        }
250
251        Ok(())
252    }
253
254    async fn do_update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
255        let mut guard = self.stats.record_update();
256        guard.inc_row_count(batch.num_rows());
257
258        for creator in self.creators.values_mut() {
259            creator.update_flat(batch).await?;
260        }
261
262        Ok(())
263    }
264
265    async fn do_finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<()> {
266        let mut guard = self.stats.record_finish();
267
268        let mut written_bytes = 0;
269        for creator in self.creators.values_mut() {
270            written_bytes += creator.finish(puffin_writer).await?;
271        }
272
273        guard.inc_byte_count(written_bytes);
274        Ok(())
275    }
276
277    async fn do_abort(&mut self) -> Result<()> {
278        let _guard = self.stats.record_cleanup();
279
280        self.aborted = true;
281
282        for (_, mut creator) in self.creators.drain() {
283            creator.abort().await?;
284        }
285
286        Ok(())
287    }
288}
289
290/// `SingleCreator` is a creator for a single column.
291struct SingleCreator {
292    /// Column ID.
293    column_id: ColumnId,
294    /// Column name.
295    column_name: String,
296    /// Inner creator.
297    inner: AltFulltextCreator,
298    /// Whether the index should be compressed.
299    compress: bool,
300}
301
302impl SingleCreator {
303    async fn update(&mut self, batch: &mut Batch) -> Result<()> {
304        let text_column = batch
305            .fields()
306            .iter()
307            .find(|c| c.column_id == self.column_id);
308        match text_column {
309            Some(column) => {
310                let data = column
311                    .data
312                    .cast(&ConcreteDataType::string_datatype())
313                    .context(CastVectorSnafu {
314                        from: column.data.data_type(),
315                        to: ConcreteDataType::string_datatype(),
316                    })?;
317
318                for i in 0..batch.num_rows() {
319                    let data = data.get_ref(i);
320                    let text = data
321                        .try_into_string()
322                        .context(DataTypeMismatchSnafu)?
323                        .unwrap_or_default();
324                    self.inner.push_text(text).await?;
325                }
326            }
327            _ => {
328                // If the column is not found in the batch, push empty text.
329                // Ensure that the number of texts pushed is the same as the number of rows in the SST,
330                // so that the texts are aligned with the row ids.
331                for _ in 0..batch.num_rows() {
332                    self.inner.push_text("").await?;
333                }
334            }
335        }
336
337        Ok(())
338    }
339
340    async fn update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
341        // Find the column in the RecordBatch by name
342        if let Some(column_array) = batch.column_by_name(&self.column_name) {
343            // Convert Arrow array to string array.
344            // TODO(yingwen): Use Utf8View later if possible.
345            match column_array.data_type() {
346                DataType::Utf8 => {
347                    let string_array = column_array.as_any().downcast_ref::<StringArray>().unwrap();
348                    for text_opt in string_array.iter() {
349                        let text = text_opt.unwrap_or_default();
350                        self.inner.push_text(text).await?;
351                    }
352                }
353                DataType::LargeUtf8 => {
354                    let large_string_array = column_array
355                        .as_any()
356                        .downcast_ref::<LargeStringArray>()
357                        .unwrap();
358                    for text_opt in large_string_array.iter() {
359                        let text = text_opt.unwrap_or_default();
360                        self.inner.push_text(text).await?;
361                    }
362                }
363                _ => {
364                    // For other types, cast to Utf8 as before
365                    let array = datatypes::arrow::compute::cast(column_array, &DataType::Utf8)
366                        .context(ComputeArrowSnafu)?;
367                    let string_array = array.as_any().downcast_ref::<StringArray>().unwrap();
368                    for text_opt in string_array.iter() {
369                        let text = text_opt.unwrap_or_default();
370                        self.inner.push_text(text).await?;
371                    }
372                }
373            }
374        } else {
375            // If the column is not found in the batch, push empty text.
376            // Ensure that the number of texts pushed is the same as the number of rows in the SST,
377            // so that the texts are aligned with the row ids.
378            for _ in 0..batch.num_rows() {
379                self.inner.push_text("").await?;
380            }
381        }
382
383        Ok(())
384    }
385
386    async fn finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<ByteCount> {
387        let options = PutOptions {
388            compression: self.compress.then_some(CompressionCodec::Zstd),
389        };
390        self.inner
391            .finish(puffin_writer, &self.column_id, options)
392            .await
393    }
394
395    async fn abort(&mut self) -> Result<()> {
396        self.inner.abort(&self.column_id).await;
397        Ok(())
398    }
399}
400
401#[allow(dead_code, clippy::large_enum_variant)]
402/// `AltFulltextCreator` is an alternative fulltext index creator that can be either Tantivy or BloomFilter.
403enum AltFulltextCreator {
404    Tantivy(TantivyFulltextIndexCreator),
405    Bloom(BloomFilterFulltextIndexCreator),
406}
407
408impl AltFulltextCreator {
409    async fn push_text(&mut self, text: &str) -> Result<()> {
410        match self {
411            Self::Tantivy(creator) => creator.push_text(text).await.context(FulltextPushTextSnafu),
412            Self::Bloom(creator) => creator.push_text(text).await.context(FulltextPushTextSnafu),
413        }
414    }
415
416    fn memory_usage(&self) -> usize {
417        match self {
418            Self::Tantivy(creator) => creator.memory_usage(),
419            Self::Bloom(creator) => creator.memory_usage(),
420        }
421    }
422
423    async fn finish(
424        &mut self,
425        puffin_writer: &mut SstPuffinWriter,
426        column_id: &ColumnId,
427        put_options: PutOptions,
428    ) -> Result<ByteCount> {
429        match self {
430            Self::Tantivy(creator) => {
431                let blob_key = format!(
432                    "{INDEX_BLOB_TYPE_TANTIVY}-{}",
433                    IndexTarget::ColumnId(*column_id)
434                );
435                creator
436                    .finish(puffin_writer, &blob_key, put_options)
437                    .await
438                    .context(FulltextFinishSnafu)
439            }
440            Self::Bloom(creator) => {
441                let blob_key = format!(
442                    "{INDEX_BLOB_TYPE_BLOOM}-{}",
443                    IndexTarget::ColumnId(*column_id)
444                );
445                creator
446                    .finish(puffin_writer, &blob_key, put_options)
447                    .await
448                    .context(FulltextFinishSnafu)
449            }
450        }
451    }
452
453    async fn abort(&mut self, column_id: &ColumnId) {
454        match self {
455            Self::Tantivy(creator) => {
456                if let Err(err) = creator.abort().await {
457                    warn!(err; "Failed to abort the fulltext index creator in the Tantivy flavor, col_id: {:?}", column_id);
458                }
459            }
460            Self::Bloom(creator) => {
461                if let Err(err) = creator.abort().await {
462                    warn!(err; "Failed to abort the fulltext index creator in the Bloom Filter flavor, col_id: {:?}", column_id);
463                }
464            }
465        }
466    }
467}
468
469#[cfg(test)]
470mod tests {
471    use std::collections::{BTreeMap, BTreeSet};
472    use std::sync::Arc;
473
474    use api::v1::SemanticType;
475    use common_base::BitVec;
476    use datatypes::data_type::DataType;
477    use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextOptions};
478    use datatypes::vectors::{UInt8Vector, UInt64Vector};
479    use futures::FutureExt;
480    use futures::future::BoxFuture;
481    use index::fulltext_index::search::RowId;
482    use object_store::ObjectStore;
483    use object_store::services::Memory;
484    use puffin::puffin_manager::{PuffinManager, PuffinWriter};
485    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef};
486    use store_api::region_request::PathType;
487    use store_api::storage::{ConcreteDataType, FileId, RegionId};
488
489    use super::*;
490    use crate::access_layer::RegionFilePathFactory;
491    use crate::read::{Batch, BatchColumn};
492    use crate::sst::file::{RegionFileId, RegionIndexId};
493    use crate::sst::index::fulltext_index::applier::FulltextIndexApplier;
494    use crate::sst::index::fulltext_index::applier::builder::{
495        FulltextQuery, FulltextRequest, FulltextTerm,
496    };
497    use crate::sst::index::puffin_manager::PuffinManagerFactory;
498
499    fn mock_object_store() -> ObjectStore {
500        ObjectStore::new(Memory::default()).unwrap()
501    }
502
503    async fn new_intm_mgr(path: impl AsRef<str>) -> IntermediateManager {
504        IntermediateManager::init_fs(path).await.unwrap()
505    }
506
507    fn mock_region_metadata(backend: FulltextBackend) -> RegionMetadataRef {
508        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
509        builder
510            .push_column_metadata(ColumnMetadata {
511                column_schema: ColumnSchema::new(
512                    "text_english_case_sensitive",
513                    ConcreteDataType::string_datatype(),
514                    true,
515                )
516                .with_fulltext_options(FulltextOptions::new_unchecked(
517                    true,
518                    FulltextAnalyzer::English,
519                    true,
520                    backend.clone(),
521                    1,
522                    0.01,
523                ))
524                .unwrap(),
525                semantic_type: SemanticType::Field,
526                column_id: 1,
527            })
528            .push_column_metadata(ColumnMetadata {
529                column_schema: ColumnSchema::new(
530                    "text_english_case_insensitive",
531                    ConcreteDataType::string_datatype(),
532                    true,
533                )
534                .with_fulltext_options(FulltextOptions::new_unchecked(
535                    true,
536                    FulltextAnalyzer::English,
537                    false,
538                    backend.clone(),
539                    1,
540                    0.01,
541                ))
542                .unwrap(),
543                semantic_type: SemanticType::Field,
544                column_id: 2,
545            })
546            .push_column_metadata(ColumnMetadata {
547                column_schema: ColumnSchema::new(
548                    "text_chinese",
549                    ConcreteDataType::string_datatype(),
550                    true,
551                )
552                .with_fulltext_options(FulltextOptions::new_unchecked(
553                    true,
554                    FulltextAnalyzer::Chinese,
555                    false,
556                    backend.clone(),
557                    1,
558                    0.01,
559                ))
560                .unwrap(),
561                semantic_type: SemanticType::Field,
562                column_id: 3,
563            })
564            .push_column_metadata(ColumnMetadata {
565                column_schema: ColumnSchema::new(
566                    "ts",
567                    ConcreteDataType::timestamp_millisecond_datatype(),
568                    false,
569                ),
570                semantic_type: SemanticType::Timestamp,
571                column_id: 4,
572            });
573
574        Arc::new(builder.build().unwrap())
575    }
576
577    fn new_batch(
578        rows: &[(
579            Option<&str>, // text_english_case_sensitive
580            Option<&str>, // text_english_case_insensitive
581            Option<&str>, // text_chinese
582        )],
583    ) -> Batch {
584        let mut vec_english_sensitive =
585            ConcreteDataType::string_datatype().create_mutable_vector(0);
586        let mut vec_english_insensitive =
587            ConcreteDataType::string_datatype().create_mutable_vector(0);
588        let mut vec_chinese = ConcreteDataType::string_datatype().create_mutable_vector(0);
589
590        for (text_english_case_sensitive, text_english_case_insensitive, text_chinese) in rows {
591            match text_english_case_sensitive {
592                Some(s) => vec_english_sensitive.push_value_ref(&(*s).into()),
593                None => vec_english_sensitive.push_null(),
594            }
595            match text_english_case_insensitive {
596                Some(s) => vec_english_insensitive.push_value_ref(&(*s).into()),
597                None => vec_english_insensitive.push_null(),
598            }
599            match text_chinese {
600                Some(s) => vec_chinese.push_value_ref(&(*s).into()),
601                None => vec_chinese.push_null(),
602            }
603        }
604
605        let num_rows = vec_english_sensitive.len();
606        Batch::new(
607            vec![],
608            Arc::new(UInt64Vector::from_iter_values(
609                (0..num_rows).map(|n| n as u64),
610            )),
611            Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
612                0, num_rows,
613            ))),
614            Arc::new(UInt8Vector::from_iter_values(std::iter::repeat_n(
615                1, num_rows,
616            ))),
617            vec![
618                BatchColumn {
619                    column_id: 1,
620                    data: vec_english_sensitive.to_vector(),
621                },
622                BatchColumn {
623                    column_id: 2,
624                    data: vec_english_insensitive.to_vector(),
625                },
626                BatchColumn {
627                    column_id: 3,
628                    data: vec_chinese.to_vector(),
629                },
630            ],
631        )
632        .unwrap()
633    }
634
635    /// Applier factory that can handle both queries and terms.
636    ///
637    /// It builds a fulltext index with the given data rows, and returns a function
638    /// that can handle both queries and terms in a single request.
639    ///
640    /// The function takes two parameters:
641    /// - `queries`: A list of (ColumnId, query_string) pairs for fulltext queries
642    /// - `terms`: A list of (ColumnId, [(bool, String)]) for fulltext terms, where bool indicates if term is lowercased
643    async fn build_fulltext_applier_factory(
644        prefix: &str,
645        backend: FulltextBackend,
646        rows: &[(
647            Option<&str>, // text_english_case_sensitive
648            Option<&str>, // text_english_case_insensitive
649            Option<&str>, // text_chinese
650        )],
651    ) -> impl Fn(
652        Vec<(ColumnId, &str)>,
653        Vec<(ColumnId, Vec<(bool, &str)>)>,
654        Option<BitVec>,
655    ) -> BoxFuture<'static, Option<BTreeSet<RowId>>> {
656        let (d, factory) = PuffinManagerFactory::new_for_test_async(prefix).await;
657        let table_dir = "table0".to_string();
658        let sst_file_id = FileId::random();
659        let object_store = mock_object_store();
660        let region_metadata = mock_region_metadata(backend.clone());
661        let intm_mgr = new_intm_mgr(d.path().to_string_lossy()).await;
662
663        let mut indexer = FulltextIndexer::new(
664            &region_metadata.region_id,
665            &sst_file_id,
666            &intm_mgr,
667            &region_metadata,
668            true,
669            1024,
670        )
671        .await
672        .unwrap()
673        .unwrap();
674
675        let mut batch = new_batch(rows);
676        indexer.update(&mut batch).await.unwrap();
677
678        let puffin_manager = factory.build(
679            object_store.clone(),
680            RegionFilePathFactory::new(table_dir.clone(), PathType::Bare),
681        );
682        let region_file_id = RegionFileId::new(region_metadata.region_id, sst_file_id);
683        let index_id = RegionIndexId::new(region_file_id, 0);
684        let mut writer = puffin_manager.writer(&index_id).await.unwrap();
685        let _ = indexer.finish(&mut writer).await.unwrap();
686        writer.finish().await.unwrap();
687
688        move |queries: Vec<(ColumnId, &str)>,
689              terms_requests: Vec<(ColumnId, Vec<(bool, &str)>)>,
690              coarse_mask: Option<BitVec>| {
691            let _d = &d;
692            let table_dir = table_dir.clone();
693            let object_store = object_store.clone();
694            let factory = factory.clone();
695
696            let mut requests: BTreeMap<ColumnId, FulltextRequest> = BTreeMap::new();
697
698            // Add queries
699            for (column_id, query) in queries {
700                requests
701                    .entry(column_id)
702                    .or_default()
703                    .queries
704                    .push(FulltextQuery(query.to_string()));
705            }
706
707            // Add terms
708            for (column_id, terms) in terms_requests {
709                let fulltext_terms = terms
710                    .into_iter()
711                    .map(|(col_lowered, term)| FulltextTerm {
712                        col_lowered,
713                        term: term.to_string(),
714                    })
715                    .collect::<Vec<_>>();
716
717                requests
718                    .entry(column_id)
719                    .or_default()
720                    .terms
721                    .extend(fulltext_terms);
722            }
723
724            let applier = FulltextIndexApplier::new(
725                table_dir,
726                PathType::Bare,
727                object_store,
728                requests,
729                factory,
730            );
731
732            let backend = backend.clone();
733            async move {
734                match backend {
735                    FulltextBackend::Tantivy => {
736                        applier.apply_fine(index_id, None, None).await.unwrap()
737                    }
738                    FulltextBackend::Bloom => {
739                        let coarse_mask = coarse_mask.unwrap_or_default();
740                        let row_groups = (0..coarse_mask.len()).map(|i| (1, coarse_mask[i]));
741                        // row group id == row id
742                        let resp = applier
743                            .apply_coarse(index_id, None, row_groups, None)
744                            .await
745                            .unwrap();
746                        resp.map(|r| {
747                            r.into_iter()
748                                .filter(|(_, ranges)| !ranges.is_empty())
749                                .map(|(row_group_id, _)| row_group_id as RowId)
750                                .collect()
751                        })
752                    }
753                }
754            }
755            .boxed()
756        }
757    }
758
759    fn rows(row_ids: impl IntoIterator<Item = RowId>) -> BTreeSet<RowId> {
760        row_ids.into_iter().collect()
761    }
762
763    #[tokio::test]
764    async fn test_fulltext_index_basic_case_sensitive_tantivy() {
765        let applier_factory = build_fulltext_applier_factory(
766            "test_fulltext_index_basic_case_sensitive_tantivy_",
767            FulltextBackend::Tantivy,
768            &[
769                (Some("hello"), None, None),
770                (Some("world"), None, None),
771                (None, None, None),
772                (Some("Hello, World"), None, None),
773            ],
774        )
775        .await;
776
777        let row_ids = applier_factory(vec![(1, "hello")], vec![], None).await;
778        assert_eq!(row_ids, Some(rows([0])));
779
780        let row_ids = applier_factory(vec![(1, "world")], vec![], None).await;
781        assert_eq!(row_ids, Some(rows([1])));
782
783        let row_ids = applier_factory(vec![(1, "Hello")], vec![], None).await;
784        assert_eq!(row_ids, Some(rows([3])));
785
786        let row_ids = applier_factory(vec![(1, "World")], vec![], None).await;
787        assert_eq!(row_ids, Some(rows([3])));
788
789        // Terms are left to the bloom backend.
790        let row_ids = applier_factory(vec![], vec![(1, vec![(false, "hello")])], None).await;
791        assert_eq!(row_ids, None);
792    }
793
794    #[tokio::test]
795    async fn test_fulltext_index_basic_case_sensitive_bloom() {
796        let applier_factory = build_fulltext_applier_factory(
797            "test_fulltext_index_basic_case_sensitive_bloom_",
798            FulltextBackend::Bloom,
799            &[
800                (Some("hello"), None, None),
801                (Some("world"), None, None),
802                (None, None, None),
803                (Some("Hello, World"), None, None),
804            ],
805        )
806        .await;
807
808        let row_ids = applier_factory(
809            vec![],
810            vec![(1, vec![(false, "hello")])],
811            Some(BitVec::from_slice(&[0b1111])),
812        )
813        .await;
814        assert_eq!(row_ids, Some(rows([0])));
815
816        let row_ids = applier_factory(
817            vec![],
818            vec![(1, vec![(false, "hello")])],
819            Some(BitVec::from_slice(&[0b1110])), // row 0 is filtered out
820        )
821        .await;
822        assert_eq!(row_ids, Some(rows([])));
823
824        let row_ids = applier_factory(
825            vec![],
826            vec![(1, vec![(true, "hello")])],
827            Some(BitVec::from_slice(&[0b1111])),
828        )
829        .await;
830        assert_eq!(row_ids, None);
831
832        let row_ids = applier_factory(
833            vec![],
834            vec![(1, vec![(false, "world")])],
835            Some(BitVec::from_slice(&[0b1111])),
836        )
837        .await;
838        assert_eq!(row_ids, Some(rows([1])));
839
840        let row_ids = applier_factory(
841            vec![],
842            vec![(1, vec![(false, "world")])],
843            Some(BitVec::from_slice(&[0b1101])), // row 1 is filtered out
844        )
845        .await;
846        assert_eq!(row_ids, Some(rows([])));
847
848        let row_ids = applier_factory(
849            vec![],
850            vec![(1, vec![(true, "world")])],
851            Some(BitVec::from_slice(&[0b1111])),
852        )
853        .await;
854        assert_eq!(row_ids, None);
855
856        let row_ids = applier_factory(
857            vec![],
858            vec![(1, vec![(false, "Hello")])],
859            Some(BitVec::from_slice(&[0b1111])),
860        )
861        .await;
862        assert_eq!(row_ids, Some(rows([3])));
863
864        let row_ids = applier_factory(
865            vec![],
866            vec![(1, vec![(false, "Hello")])],
867            Some(BitVec::from_slice(&[0b0111])), // row 3 is filtered out
868        )
869        .await;
870        assert_eq!(row_ids, Some(rows([])));
871
872        let row_ids = applier_factory(
873            vec![],
874            vec![(1, vec![(true, "Hello")])],
875            Some(BitVec::from_slice(&[0b1111])),
876        )
877        .await;
878        assert_eq!(row_ids, None);
879
880        let row_ids = applier_factory(
881            vec![],
882            vec![(1, vec![(false, "Hello, World")])],
883            Some(BitVec::from_slice(&[0b1111])),
884        )
885        .await;
886        assert_eq!(row_ids, Some(rows([3])));
887
888        let row_ids = applier_factory(
889            vec![],
890            vec![(1, vec![(false, "Hello, World")])],
891            Some(BitVec::from_slice(&[0b0111])), // row 3 is filtered out
892        )
893        .await;
894        assert_eq!(row_ids, Some(rows([])));
895
896        let row_ids = applier_factory(
897            vec![],
898            vec![(1, vec![(true, "Hello, World")])],
899            Some(BitVec::from_slice(&[0b1111])),
900        )
901        .await;
902        assert_eq!(row_ids, None);
903    }
904
905    #[tokio::test]
906    async fn test_fulltext_index_basic_case_insensitive_tantivy() {
907        let applier_factory = build_fulltext_applier_factory(
908            "test_fulltext_index_basic_case_insensitive_tantivy_",
909            FulltextBackend::Tantivy,
910            &[
911                (None, Some("hello"), None),
912                (None, None, None),
913                (None, Some("world"), None),
914                (None, Some("Hello, World"), None),
915            ],
916        )
917        .await;
918
919        let row_ids = applier_factory(vec![(2, "hello")], vec![], None).await;
920        assert_eq!(row_ids, Some(rows([0, 3])));
921
922        let row_ids = applier_factory(vec![(2, "world")], vec![], None).await;
923        assert_eq!(row_ids, Some(rows([2, 3])));
924
925        let row_ids = applier_factory(vec![(2, "Hello")], vec![], None).await;
926        assert_eq!(row_ids, Some(rows([0, 3])));
927
928        let row_ids = applier_factory(vec![(2, "World")], vec![], None).await;
929        assert_eq!(row_ids, Some(rows([2, 3])));
930    }
931
932    #[tokio::test]
933    async fn test_fulltext_index_basic_case_insensitive_bloom() {
934        let applier_factory = build_fulltext_applier_factory(
935            "test_fulltext_index_basic_case_insensitive_bloom_",
936            FulltextBackend::Bloom,
937            &[
938                (None, Some("hello"), None),
939                (None, None, None),
940                (None, Some("world"), None),
941                (None, Some("Hello, World"), None),
942            ],
943        )
944        .await;
945
946        let row_ids = applier_factory(
947            vec![],
948            vec![(2, vec![(false, "hello")])],
949            Some(BitVec::from_slice(&[0b1111])),
950        )
951        .await;
952        assert_eq!(row_ids, Some(rows([0, 3])));
953
954        let row_ids = applier_factory(
955            vec![],
956            vec![(2, vec![(false, "hello")])],
957            Some(BitVec::from_slice(&[0b1110])), // row 0 is filtered out
958        )
959        .await;
960        assert_eq!(row_ids, Some(rows([3])));
961
962        let row_ids = applier_factory(
963            vec![],
964            vec![(2, vec![(false, "world")])],
965            Some(BitVec::from_slice(&[0b1111])),
966        )
967        .await;
968        assert_eq!(row_ids, Some(rows([2, 3])));
969
970        let row_ids = applier_factory(
971            vec![],
972            vec![(2, vec![(false, "world")])],
973            Some(BitVec::from_slice(&[0b1011])), // row 2 is filtered out
974        )
975        .await;
976        assert_eq!(row_ids, Some(rows([3])));
977
978        let row_ids = applier_factory(
979            vec![],
980            vec![(2, vec![(false, "Hello")])],
981            Some(BitVec::from_slice(&[0b1111])),
982        )
983        .await;
984        assert_eq!(row_ids, Some(rows([0, 3])));
985
986        let row_ids = applier_factory(
987            vec![],
988            vec![(2, vec![(false, "Hello")])],
989            Some(BitVec::from_slice(&[0b0111])), // row 3 is filtered out
990        )
991        .await;
992        assert_eq!(row_ids, Some(rows([0])));
993
994        let row_ids = applier_factory(
995            vec![],
996            vec![(2, vec![(false, "World")])],
997            Some(BitVec::from_slice(&[0b1111])),
998        )
999        .await;
1000        assert_eq!(row_ids, Some(rows([2, 3])));
1001
1002        let row_ids = applier_factory(
1003            vec![],
1004            vec![(2, vec![(false, "World")])],
1005            Some(BitVec::from_slice(&[0b0111])), // row 3 is filtered out
1006        )
1007        .await;
1008        assert_eq!(row_ids, Some(rows([2])));
1009
1010        let row_ids = applier_factory(
1011            vec![],
1012            vec![(2, vec![(true, "hello")])],
1013            Some(BitVec::from_slice(&[0b1111])),
1014        )
1015        .await;
1016        assert_eq!(row_ids, Some(rows([0, 3])));
1017    }
1018
1019    #[tokio::test]
1020    async fn test_fulltext_index_basic_chinese_tantivy() {
1021        let applier_factory = build_fulltext_applier_factory(
1022            "test_fulltext_index_basic_chinese_tantivy_",
1023            FulltextBackend::Tantivy,
1024            &[
1025                (None, None, Some("你好")),
1026                (None, None, None),
1027                (None, None, Some("世界")),
1028                (None, None, Some("你好，世界")),
1029            ],
1030        )
1031        .await;
1032
1033        let row_ids = applier_factory(vec![(3, "你好")], vec![], None).await;
1034        assert_eq!(row_ids, Some(rows([0, 3])));
1035
1036        let row_ids = applier_factory(vec![(3, "世界")], vec![], None).await;
1037        assert_eq!(row_ids, Some(rows([2, 3])));
1038    }
1039
1040    #[tokio::test]
1041    async fn test_fulltext_index_basic_chinese_bloom() {
1042        let applier_factory = build_fulltext_applier_factory(
1043            "test_fulltext_index_basic_chinese_bloom_",
1044            FulltextBackend::Bloom,
1045            &[
1046                (None, None, Some("你好")),
1047                (None, None, None),
1048                (None, None, Some("世界")),
1049                (None, None, Some("你好，世界")),
1050            ],
1051        )
1052        .await;
1053
1054        let row_ids = applier_factory(
1055            vec![],
1056            vec![(3, vec![(false, "你好")])],
1057            Some(BitVec::from_slice(&[0b1111])),
1058        )
1059        .await;
1060        assert_eq!(row_ids, Some(rows([0, 3])));
1061
1062        let row_ids = applier_factory(
1063            vec![],
1064            vec![(3, vec![(false, "你好")])],
1065            Some(BitVec::from_slice(&[0b1110])), // row 0 is filtered out
1066        )
1067        .await;
1068        assert_eq!(row_ids, Some(rows([3])));
1069
1070        let row_ids = applier_factory(
1071            vec![],
1072            vec![(3, vec![(false, "世界")])],
1073            Some(BitVec::from_slice(&[0b1111])),
1074        )
1075        .await;
1076        assert_eq!(row_ids, Some(rows([2, 3])));
1077
1078        let row_ids = applier_factory(
1079            vec![],
1080            vec![(3, vec![(false, "世界")])],
1081            Some(BitVec::from_slice(&[0b1011])), // row 2 is filtered out
1082        )
1083        .await;
1084        assert_eq!(row_ids, Some(rows([3])));
1085    }
1086
1087    #[tokio::test]
1088    async fn test_fulltext_index_multi_terms_case_sensitive_bloom() {
1089        let applier_factory = build_fulltext_applier_factory(
1090            "test_fulltext_index_multi_terms_case_sensitive_bloom_",
1091            FulltextBackend::Bloom,
1092            &[
1093                (Some("Hello"), None, None),
1094                (Some("World"), None, None),
1095                (None, None, None),
1096                (Some("Hello, World"), None, None),
1097            ],
1098        )
1099        .await;
1100
1101        let row_ids = applier_factory(
1102            vec![],
1103            vec![(1, vec![(false, "hello"), (false, "world")])],
1104            Some(BitVec::from_slice(&[0b1111])),
1105        )
1106        .await;
1107        assert_eq!(row_ids, Some(rows([])));
1108
1109        let row_ids = applier_factory(
1110            vec![],
1111            vec![(1, vec![(false, "Hello"), (false, "World")])],
1112            Some(BitVec::from_slice(&[0b1111])),
1113        )
1114        .await;
1115        assert_eq!(row_ids, Some(rows([3])));
1116
1117        let row_ids = applier_factory(
1118            vec![],
1119            vec![(1, vec![(true, "Hello"), (false, "World")])],
1120            Some(BitVec::from_slice(&[0b1111])),
1121        )
1122        .await;
1123        assert_eq!(row_ids, Some(rows([1, 3])));
1124
1125        let row_ids = applier_factory(
1126            vec![],
1127            vec![(1, vec![(false, "Hello"), (true, "World")])],
1128            Some(BitVec::from_slice(&[0b1111])),
1129        )
1130        .await;
1131        assert_eq!(row_ids, Some(rows([0, 3])));
1132
1133        let row_ids = applier_factory(
1134            vec![],
1135            vec![(1, vec![(true, "Hello"), (true, "World")])],
1136            Some(BitVec::from_slice(&[0b1111])),
1137        )
1138        .await;
1139        assert_eq!(row_ids, None);
1140    }
1141
1142    #[tokio::test]
1143    async fn test_fulltext_index_multi_terms_case_insensitive_bloom() {
1144        let applier_factory = build_fulltext_applier_factory(
1145            "test_fulltext_index_multi_terms_case_insensitive_bloom_",
1146            FulltextBackend::Bloom,
1147            &[
1148                (None, Some("hello"), None),
1149                (None, None, None),
1150                (None, Some("world"), None),
1151                (None, Some("Hello, World"), None),
1152            ],
1153        )
1154        .await;
1155
1156        let row_ids = applier_factory(
1157            vec![],
1158            vec![(2, vec![(false, "hello"), (false, "world")])],
1159            Some(BitVec::from_slice(&[0b1111])),
1160        )
1161        .await;
1162        assert_eq!(row_ids, Some(rows([3])));
1163
1164        let row_ids = applier_factory(
1165            vec![],
1166            vec![(2, vec![(true, "hello"), (false, "world")])],
1167            Some(BitVec::from_slice(&[0b1111])),
1168        )
1169        .await;
1170        assert_eq!(row_ids, Some(rows([3])));
1171
1172        let row_ids = applier_factory(
1173            vec![],
1174            vec![(2, vec![(false, "hello"), (true, "world")])],
1175            Some(BitVec::from_slice(&[0b1111])),
1176        )
1177        .await;
1178        assert_eq!(row_ids, Some(rows([3])));
1179
1180        let row_ids = applier_factory(
1181            vec![],
1182            vec![(2, vec![(true, "hello"), (true, "world")])],
1183            Some(BitVec::from_slice(&[0b1111])),
1184        )
1185        .await;
1186        assert_eq!(row_ids, Some(rows([3])));
1187    }
1188
1189    #[tokio::test]
1190    async fn test_fulltext_index_multi_columns_tantivy() {
1191        let applier_factory = build_fulltext_applier_factory(
1192            "test_fulltext_index_multi_columns_tantivy_",
1193            FulltextBackend::Tantivy,
1194            &[
1195                (Some("Hello"), None, Some("你好")),
1196                (Some("World"), Some("world"), None),
1197                (None, Some("World"), Some("世界")),
1198                (
1199                    Some("Hello, World"),
1200                    Some("Hello, World"),
1201                    Some("你好，世界"),
1202                ),
1203            ],
1204        )
1205        .await;
1206
1207        let row_ids = applier_factory(vec![(1, "Hello"), (3, "你好")], vec![], None).await;
1208        assert_eq!(row_ids, Some(rows([0, 3])));
1209
1210        let row_ids = applier_factory(vec![(1, "World"), (2, "World")], vec![], None).await;
1211        assert_eq!(row_ids, Some(rows([1, 3])));
1212    }
1213
1214    #[tokio::test]
1215    async fn test_fulltext_index_multi_columns_bloom() {
1216        let applier_factory = build_fulltext_applier_factory(
1217            "test_fulltext_index_multi_columns_bloom_",
1218            FulltextBackend::Bloom,
1219            &[
1220                (Some("Hello"), None, Some("你好")),
1221                (Some("World"), Some("world"), None),
1222                (None, Some("World"), Some("世界")),
1223                (
1224                    Some("Hello, World"),
1225                    Some("Hello, World"),
1226                    Some("你好，世界"),
1227                ),
1228            ],
1229        )
1230        .await;
1231
1232        let row_ids = applier_factory(
1233            vec![],
1234            vec![
1235                (1, vec![(false, "Hello")]),
1236                (2, vec![(false, "world")]),
1237                (3, vec![(false, "你好")]),
1238            ],
1239            Some(BitVec::from_slice(&[0b1111])),
1240        )
1241        .await;
1242        assert_eq!(row_ids, Some(rows([3])));
1243
1244        let row_ids = applier_factory(
1245            vec![],
1246            vec![(1, vec![(false, "World")]), (2, vec![(false, "World")])],
1247            Some(BitVec::from_slice(&[0b1111])),
1248        )
1249        .await;
1250        assert_eq!(row_ids, Some(rows([1, 3])));
1251    }
1252}
mito2/sst/index/fulltext_index/creator.rs

mito2/sst/index/fulltext_index/
creator.rs
