mito2/sst/index/inverted_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::num::NonZeroUsize;
16use std::sync::Arc;
17use std::sync::atomic::AtomicUsize;
18
19use api::v1::SemanticType;
20use common_telemetry::{debug, warn};
21use datatypes::arrow::record_batch::RecordBatch;
22use datatypes::vectors::Helper;
23use index::inverted_index::create::InvertedIndexCreator;
24use index::inverted_index::create::sort::external_sort::ExternalSorter;
25use index::inverted_index::create::sort_create::SortIndexCreator;
26use index::inverted_index::format::writer::InvertedIndexBlobWriter;
27use index::target::IndexTarget;
28use mito_codec::index::{IndexValueCodec, IndexValuesCodec};
29use mito_codec::row_converter::sparse::SparsePrimaryKeyView;
30use mito_codec::row_converter::{SortField, SparseOffsetsCache};
31use puffin::puffin_manager::{PuffinWriter, PutOptions};
32use smallvec::SmallVec;
33use snafu::{OptionExt, ResultExt, ensure};
34use store_api::codec::PrimaryKeyEncoding;
35use store_api::metadata::RegionMetadataRef;
36use store_api::storage::{ColumnId, FileId};
37use tokio::io::duplex;
38use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
39
40use crate::error::{
41    BiErrorsSnafu, ConvertVectorSnafu, DecodeSnafu, EncodeSnafu, IndexFinishSnafu,
42    InvalidRecordBatchSnafu, OperateAbortedIndexSnafu, PuffinAddBlobSnafu, PushIndexValueSnafu,
43    Result,
44};
45use crate::read::Batch;
46use crate::sst::index::TYPE_INVERTED_INDEX;
47use crate::sst::index::column::{column_index_rows, extract_json_type_hint_values};
48use crate::sst::index::intermediate::{
49    IntermediateLocation, IntermediateManager, TempFileProvider,
50};
51use crate::sst::index::inverted_index::INDEX_BLOB_TYPE;
52use crate::sst::index::primary_key::PrimaryKeyRuns;
53use crate::sst::index::puffin_manager::SstPuffinWriter;
54use crate::sst::index::statistics::{ByteCount, RowCount, Statistics};
55
56/// The minimum memory usage threshold for one column.
57const MIN_MEMORY_USAGE_THRESHOLD_PER_COLUMN: usize = 1024 * 1024; // 1MB
58
59/// The buffer size for the pipe used to send index data to the puffin blob.
60const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192;
61
62/// `InvertedIndexer` creates inverted index for SST files.
63pub struct InvertedIndexer {
64    /// The index creator.
65    index_creator: Box<dyn InvertedIndexCreator>,
66    /// The provider of intermediate files.
67    temp_file_provider: Arc<TempFileProvider>,
68
69    /// Codec for decoding primary keys.
70    codec: IndexValuesCodec,
71    /// Reusable buffer for encoding index values.
72    value_buf: Vec<u8>,
73    /// Scratch offsets shared by indexed tags of one sparse primary key.
74    pk_offsets: SparseOffsetsCache,
75
76    /// Statistics of index creation.
77    stats: Statistics,
78    /// Whether the index creation is aborted.
79    aborted: bool,
80
81    /// The memory usage of the index creator.
82    memory_usage: Arc<AtomicUsize>,
83
84    /// Index targets and their encoded keys.
85    index_targets: Vec<(IndexTarget, String)>,
86
87    /// Region metadata for column lookups.
88    metadata: RegionMetadataRef,
89}
90
91impl InvertedIndexer {
92    /// Creates a new `InvertedIndexer`.
93    pub fn new(
94        sst_file_id: FileId,
95        metadata: &RegionMetadataRef,
96        intermediate_manager: IntermediateManager,
97        memory_usage_threshold: Option<usize>,
98        segment_row_count: NonZeroUsize,
99        index_targets: Vec<IndexTarget>,
100    ) -> Self {
101        let temp_file_provider = Arc::new(TempFileProvider::new(
102            IntermediateLocation::new(&metadata.region_id, &sst_file_id),
103            intermediate_manager,
104        ));
105
106        let memory_usage = Arc::new(AtomicUsize::new(0));
107
108        let sorter = ExternalSorter::factory(
109            temp_file_provider.clone() as _,
110            Some(MIN_MEMORY_USAGE_THRESHOLD_PER_COLUMN),
111            memory_usage.clone(),
112            memory_usage_threshold,
113        );
114        let index_creator = Box::new(SortIndexCreator::new(sorter, segment_row_count));
115
116        let codec = IndexValuesCodec::from_tag_columns(
117            metadata.primary_key_encoding,
118            metadata.primary_key_columns(),
119        );
120        let index_targets = index_targets
121            .into_iter()
122            .map(|target| {
123                let key = target.encode();
124                (target, key)
125            })
126            .collect();
127        Self {
128            codec,
129            index_creator,
130            temp_file_provider,
131            value_buf: vec![],
132            pk_offsets: SparseOffsetsCache::new(),
133            stats: Statistics::new(TYPE_INVERTED_INDEX),
134            aborted: false,
135            memory_usage,
136            index_targets,
137            metadata: metadata.clone(),
138        }
139    }
140
141    /// Updates index with a batch of rows.
142    /// Garbage will be cleaned up if failed to update.
143    pub async fn update(&mut self, batch: &mut Batch) -> Result<()> {
144        ensure!(!self.aborted, OperateAbortedIndexSnafu);
145
146        if batch.is_empty() {
147            return Ok(());
148        }
149
150        if let Err(update_err) = self.do_update(batch).await {
151            // clean up garbage if failed to update
152            if let Err(err) = self.do_cleanup().await {
153                if cfg!(any(test, feature = "test")) {
154                    panic!("Failed to clean up index creator, err: {err}",);
155                } else {
156                    warn!(err; "Failed to clean up index creator");
157                }
158            }
159            return Err(update_err);
160        }
161
162        Ok(())
163    }
164
165    /// Updates the inverted index with the given flat format RecordBatch.
166    pub async fn update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
167        ensure!(!self.aborted, OperateAbortedIndexSnafu);
168
169        if batch.num_rows() == 0 {
170            return Ok(());
171        }
172
173        self.do_update_flat(batch).await
174    }
175
176    async fn do_update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
177        let mut guard = self.stats.record_update();
178
179        guard.inc_row_count(batch.num_rows());
180
181        let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse;
182        let mut sparse_columns: SmallVec<[(ColumnId, &str); 8]> = SmallVec::new();
183
184        for (target, target_key) in &self.index_targets {
185            let (vector, sort_field, semantic_type) = match target {
186                IndexTarget::ColumnId(col_id) => {
187                    let Some(col_meta) = self.metadata.column_by_id(*col_id) else {
188                        debug!(
189                            "Column {} not found in the metadata during building inverted index",
190                            col_id
191                        );
192                        continue;
193                    };
194                    let col_name = &col_meta.column_schema.name;
195                    let Some(col_array) = batch.column_by_name(col_name) else {
196                        if is_sparse && col_meta.semantic_type == SemanticType::Tag {
197                            if self.codec.pk_col_info(*col_id).is_some() {
198                                sparse_columns.push((*col_id, target_key));
199                            }
200                        } else {
201                            debug!(
202                                "Column {} not found in the batch during building inverted index",
203                                col_id
204                            );
205                        }
206                        continue;
207                    };
208                    let vector =
209                        Helper::try_into_vector(col_array.clone()).context(ConvertVectorSnafu)?;
210                    let sort_field = SortField::new(vector.data_type());
211                    (vector, sort_field, col_meta.semantic_type)
212                }
213                IndexTarget::JsonPath(target) => {
214                    let col_id = target.column_id();
215                    let col = self.metadata.column_by_id(col_id).with_context(|| {
216                        InvalidRecordBatchSnafu {
217                            reason: format!("Missing JSON index column {col_id}"),
218                        }
219                    })?;
220                    let root =
221                        batch
222                            .column_by_name(&col.column_schema.name)
223                            .with_context(|| InvalidRecordBatchSnafu {
224                                reason: format!("Missing JSON index column {col_id} in batch"),
225                            })?;
226                    let leaf = extract_json_type_hint_values(root, target.path())?;
227                    let field = SortField::new(leaf.data_type());
228                    (leaf, field, SemanticType::Field)
229                }
230            };
231            for (row, count) in column_index_rows(batch, semantic_type) {
232                let value = IndexValueCodec::encode_value(
233                    vector.get_ref(row),
234                    &sort_field,
235                    &mut self.value_buf,
236                )
237                .context(EncodeSnafu)?;
238                if self
239                    .index_creator
240                    .push_with_name_n(target_key, value, count)
241                {
242                    self.index_creator
243                        .spill()
244                        .await
245                        .context(PushIndexValueSnafu)?;
246                }
247            }
248        }
249
250        if !sparse_columns.is_empty() {
251            for (pk, count) in PrimaryKeyRuns::try_new(batch)? {
252                let mut view =
253                    SparsePrimaryKeyView::new(pk, &mut self.pk_offsets).context(DecodeSnafu)?;
254                // Visit all needed tags before moving to the next PK so offset discovery is shared.
255                for &(col_id, target_key) in &sparse_columns {
256                    let value = IndexValueCodec::encode_sparse_value(
257                        &mut view,
258                        col_id,
259                        &mut self.value_buf,
260                    )
261                    .context(DecodeSnafu)?;
262                    if self
263                        .index_creator
264                        .push_with_name_n(target_key, value, count)
265                    {
266                        self.index_creator
267                            .spill()
268                            .await
269                            .context(PushIndexValueSnafu)?;
270                    }
271                }
272            }
273        }
274
275        Ok(())
276    }
277
278    /// Finishes index creation and cleans up garbage.
279    /// Returns the number of rows and bytes written.
280    pub(crate) async fn finish(
281        &mut self,
282        puffin_writer: &mut SstPuffinWriter,
283    ) -> Result<(RowCount, ByteCount)> {
284        ensure!(!self.aborted, OperateAbortedIndexSnafu);
285
286        if self.stats.row_count() == 0 {
287            // no IO is performed, no garbage to clean up, just return
288            return Ok((0, 0));
289        }
290
291        let finish_res = self.do_finish(puffin_writer).await;
292        // clean up garbage no matter finish successfully or not
293        if let Err(err) = self.do_cleanup().await {
294            if cfg!(any(test, feature = "test")) {
295                panic!("Failed to clean up index creator, err: {err}",);
296            } else {
297                warn!(err; "Failed to clean up index creator");
298            }
299        }
300
301        finish_res.map(|_| (self.stats.row_count(), self.stats.byte_count()))
302    }
303
304    /// Aborts index creation and clean up garbage.
305    pub async fn abort(&mut self) -> Result<()> {
306        if self.aborted {
307            return Ok(());
308        }
309        self.aborted = true;
310
311        self.do_cleanup().await
312    }
313
314    async fn do_update(&mut self, batch: &mut Batch) -> Result<()> {
315        let mut guard = self.stats.record_update();
316
317        let n = batch.num_rows();
318        guard.inc_row_count(n);
319
320        for (target, target_key) in &self.index_targets {
321            let (vector, sort_field) = match target {
322                IndexTarget::ColumnId(col_id) => {
323                    if let Some(col_info) = self.codec.pk_col_info(*col_id) {
324                        let pk_idx = col_info.idx;
325                        let field = &col_info.field;
326                        let value = batch
327                            .pk_col_value(self.codec.decoder(), pk_idx, *col_id)?
328                            .filter(|v| !v.is_null())
329                            .map(|v| {
330                                self.value_buf.clear();
331                                IndexValueCodec::encode_nonnull_value(
332                                    v.as_value_ref(),
333                                    field,
334                                    &mut self.value_buf,
335                                )
336                                .context(EncodeSnafu)?;
337                                Ok(self.value_buf.as_slice())
338                            })
339                            .transpose()?;
340
341                        if self.index_creator.push_with_name_n(target_key, value, n) {
342                            self.index_creator
343                                .spill()
344                                .await
345                                .context(PushIndexValueSnafu)?;
346                        }
347                        continue;
348                    }
349                    let Some(values) = batch.field_col_value(*col_id) else {
350                        debug!(
351                            "Column {} not found in the batch during building inverted index",
352                            col_id
353                        );
354                        continue;
355                    };
356                    let sort_field = SortField::new(values.data.data_type());
357                    (values.data.clone(), sort_field)
358                }
359                IndexTarget::JsonPath(target) => {
360                    let col_id = target.column_id();
361                    let root = batch
362                        .field_col_value(col_id)
363                        .with_context(|| InvalidRecordBatchSnafu {
364                            reason: format!("Missing JSON index column {col_id} in batch"),
365                        })?
366                        .data
367                        .to_arrow_array();
368                    let leaf = extract_json_type_hint_values(&root, target.path())?;
369                    let field = SortField::new(leaf.data_type());
370                    (leaf, field)
371                }
372            };
373            for row in 0..n {
374                let value = IndexValueCodec::encode_value(
375                    vector.get_ref(row),
376                    &sort_field,
377                    &mut self.value_buf,
378                )
379                .context(EncodeSnafu)?;
380                if self.index_creator.push_with_name(target_key, value) {
381                    self.index_creator
382                        .spill()
383                        .await
384                        .context(PushIndexValueSnafu)?;
385                }
386            }
387        }
388
389        Ok(())
390    }
391
392    /// Data flow of finishing index:
393    ///
394    /// ```text
395    ///                               (In Memory Buffer)
396    ///                                    ┌──────┐
397    ///  ┌─────────────┐                   │ PIPE │
398    ///  │             │ write index data  │      │
399    ///  │ IndexWriter ├──────────────────►│ tx   │
400    ///  │             │                   │      │
401    ///  └─────────────┘                   │      │
402    ///                  ┌─────────────────┤ rx   │
403    ///  ┌─────────────┐ │ read as blob    └──────┘
404    ///  │             │ │
405    ///  │ PuffinWriter├─┤
406    ///  │             │ │ copy to file    ┌──────┐
407    ///  └─────────────┘ └────────────────►│ File │
408    ///                                    └──────┘
409    /// ```
410    async fn do_finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<()> {
411        let mut guard = self.stats.record_finish();
412
413        let (tx, rx) = duplex(PIPE_BUFFER_SIZE_FOR_SENDING_BLOB);
414        let mut index_writer = InvertedIndexBlobWriter::new(tx.compat_write());
415
416        let (index_finish, puffin_add_blob) = futures::join!(
417            // TODO(zhongzc): config bitmap type
418            self.index_creator
419                .finish(&mut index_writer, index::bitmap::BitmapType::Roaring),
420            puffin_writer.put_blob(
421                INDEX_BLOB_TYPE,
422                rx.compat(),
423                PutOptions::default(),
424                Default::default(),
425            )
426        );
427
428        match (
429            puffin_add_blob.context(PuffinAddBlobSnafu),
430            index_finish.context(IndexFinishSnafu),
431        ) {
432            (Err(e1), Err(e2)) => BiErrorsSnafu {
433                first: Box::new(e1),
434                second: Box::new(e2),
435            }
436            .fail()?,
437
438            (Ok(_), e @ Err(_)) => e?,
439            (e @ Err(_), Ok(_)) => e.map(|_| ())?,
440            (Ok(written_bytes), Ok(_)) => {
441                guard.inc_byte_count(written_bytes);
442            }
443        }
444
445        Ok(())
446    }
447
448    async fn do_cleanup(&mut self) -> Result<()> {
449        let _guard = self.stats.record_cleanup();
450
451        self.temp_file_provider.cleanup().await
452    }
453
454    /// Returns the targets indexed by this creator.
455    pub fn indexed_targets(&self) -> impl Iterator<Item = &IndexTarget> {
456        self.index_targets.iter().map(|(target, _)| target)
457    }
458
459    pub fn memory_usage(&self) -> usize {
460        self.memory_usage.load(std::sync::atomic::Ordering::Relaxed)
461    }
462}
463
464#[cfg(test)]
465mod tests {
466    use std::collections::{BTreeMap, BTreeSet, HashSet};
467
468    use api::v1::SemanticType;
469    use datafusion_expr::{Expr as DfExpr, Operator, binary_expr, col, lit};
470    use datatypes::arrow::array::{Array, ArrayRef, Int64Array, StructArray};
471    use datatypes::arrow::buffer::NullBuffer;
472    use datatypes::arrow::datatypes::{Field, Schema};
473    use datatypes::data_type::ConcreteDataType;
474    use datatypes::schema::ColumnSchema;
475    use datatypes::value::ValueRef;
476    use datatypes::vectors::{UInt8Vector, UInt64Vector};
477    use futures::future::BoxFuture;
478    use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader};
479    use index::inverted_index::search::predicate::{InListPredicate, Predicate};
480    use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt};
481    use object_store::ObjectStore;
482    use object_store::services::Memory;
483    use puffin::puffin_manager::cache::PuffinMetadataCache;
484    use puffin::puffin_manager::{PuffinManager, PuffinReader};
485    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
486    use store_api::region_request::PathType;
487    use store_api::storage::RegionId;
488
489    use super::*;
490    use crate::access_layer::RegionFilePathFactory;
491    use crate::cache::index::inverted_index::InvertedIndexCache;
492    use crate::metrics::CACHE_BYTES;
493    use crate::read::BatchColumn;
494    use crate::sst::file::{RegionFileId, RegionIndexId};
495    use crate::sst::index::inverted_index::applier::InvertedIndexApplier;
496    use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
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() -> RegionMetadataRef {
508        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
509        builder
510            .push_column_metadata(ColumnMetadata {
511                column_schema: ColumnSchema::new(
512                    "tag_str",
513                    ConcreteDataType::string_datatype(),
514                    false,
515                ),
516                semantic_type: SemanticType::Tag,
517                column_id: 1,
518            })
519            .push_column_metadata(ColumnMetadata {
520                column_schema: ColumnSchema::new(
521                    "tag_i32",
522                    ConcreteDataType::int32_datatype(),
523                    false,
524                ),
525                semantic_type: SemanticType::Tag,
526                column_id: 2,
527            })
528            .push_column_metadata(ColumnMetadata {
529                column_schema: ColumnSchema::new(
530                    "ts",
531                    ConcreteDataType::timestamp_millisecond_datatype(),
532                    false,
533                ),
534                semantic_type: SemanticType::Timestamp,
535                column_id: 3,
536            })
537            .push_column_metadata(ColumnMetadata {
538                column_schema: ColumnSchema::new(
539                    "field_u64",
540                    ConcreteDataType::uint64_datatype(),
541                    false,
542                ),
543                semantic_type: SemanticType::Field,
544                column_id: 4,
545            })
546            .primary_key(vec![1, 2]);
547
548        Arc::new(builder.build().unwrap())
549    }
550
551    fn new_batch(
552        str_tag: impl AsRef<str>,
553        i32_tag: impl Into<i32>,
554        u64_field: impl IntoIterator<Item = u64>,
555    ) -> Batch {
556        let fields = vec![
557            (0, SortField::new(ConcreteDataType::string_datatype())),
558            (1, SortField::new(ConcreteDataType::int32_datatype())),
559        ];
560        let codec = DensePrimaryKeyCodec::with_fields(fields);
561        let row: [ValueRef; 2] = [str_tag.as_ref().into(), i32_tag.into().into()];
562        let primary_key = codec.encode(row.into_iter()).unwrap();
563
564        let u64_field = BatchColumn {
565            column_id: 4,
566            data: Arc::new(UInt64Vector::from_iter_values(u64_field)),
567        };
568        let num_rows = u64_field.data.len();
569
570        Batch::new(
571            primary_key,
572            Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
573                0, num_rows,
574            ))),
575            Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
576                0, num_rows,
577            ))),
578            Arc::new(UInt8Vector::from_iter_values(std::iter::repeat_n(
579                1, num_rows,
580            ))),
581            vec![u64_field],
582        )
583        .unwrap()
584    }
585
586    async fn build_applier_factory(
587        prefix: &str,
588        rows: BTreeSet<(&'static str, i32, [u64; 2])>,
589        sst_indexed_column_ids: HashSet<ColumnId>,
590    ) -> impl Fn(DfExpr) -> BoxFuture<'static, Vec<usize>> {
591        let (d, factory) = PuffinManagerFactory::new_for_test_async(prefix).await;
592        let table_dir = "table0".to_string();
593        let sst_file_id = FileId::random();
594        let object_store = mock_object_store();
595        let region_metadata = mock_region_metadata();
596        let intm_mgr = new_intm_mgr(d.path().to_string_lossy()).await;
597        let memory_threshold = None;
598        let segment_row_count = 2;
599        let indexed_column_ids = HashSet::from_iter([1, 2, 4]);
600
601        let mut creator = InvertedIndexer::new(
602            sst_file_id,
603            &region_metadata,
604            intm_mgr,
605            memory_threshold,
606            NonZeroUsize::new(segment_row_count).unwrap(),
607            sst_indexed_column_ids
608                .iter()
609                .copied()
610                .map(IndexTarget::ColumnId)
611                .collect(),
612        );
613
614        for (str_tag, i32_tag, u64_field) in &rows {
615            let mut batch = new_batch(str_tag, *i32_tag, u64_field.iter().copied());
616            creator.update(&mut batch).await.unwrap();
617        }
618
619        let puffin_manager = factory.build(
620            object_store.clone(),
621            RegionFilePathFactory::new(table_dir.clone(), PathType::Bare),
622        );
623
624        let sst_file_id = RegionFileId::new(region_metadata.region_id, sst_file_id);
625        let index_id = RegionIndexId::new(sst_file_id, 0);
626        let mut writer = puffin_manager.writer(&index_id).await.unwrap();
627        let (row_count, _) = creator.finish(&mut writer).await.unwrap();
628        assert_eq!(row_count, rows.len() * segment_row_count);
629        writer.finish().await.unwrap();
630
631        move |expr| {
632            let _d = &d;
633            let cache = Arc::new(InvertedIndexCache::new(10, 10, 100));
634            let puffin_metadata_cache = Arc::new(PuffinMetadataCache::new(10, &CACHE_BYTES));
635            let applier = InvertedIndexApplierBuilder::new(
636                table_dir.clone(),
637                PathType::Bare,
638                object_store.clone(),
639                &region_metadata,
640                indexed_column_ids.clone(),
641                factory.clone(),
642            )
643            .with_inverted_index_cache(Some(cache))
644            .with_puffin_metadata_cache(Some(puffin_metadata_cache))
645            .build(&[expr])
646            .unwrap()
647            .unwrap();
648            let sst_metadata = Arc::new(region_metadata.clone());
649            let plan = applier.plan_for_sst(&sst_metadata).unwrap().unwrap();
650            Box::pin(async move {
651                applier
652                    .apply(index_id, None, &plan.index_applier, None)
653                    .await
654                    .unwrap()
655                    .matched_segment_ids
656                    .iter_ones()
657                    .collect()
658            })
659        }
660    }
661
662    #[tokio::test]
663    async fn test_create_and_query_missing_index() {
664        let rows = BTreeSet::from_iter([("aaa", 1, [1, 2]), ("bbb", 2, [3, 4])]);
665        // The SST predates enabling indexes on tag_i32 and field_u64.
666        let applier_factory = build_applier_factory(
667            "test_create_and_query_missing_index_",
668            rows,
669            HashSet::from([1]),
670        )
671        .await;
672
673        let cases = [
674            (col("tag_i32").eq(lit(1)), vec![0, 1]),
675            (
676                col("tag_i32")
677                    .eq(lit(1))
678                    .and(col("field_u64").eq(lit(1u64))),
679                vec![0, 1],
680            ),
681            (
682                col("tag_str").eq(lit("aaa")).and(col("tag_i32").eq(lit(1))),
683                vec![0],
684            ),
685            (
686                col("tag_str")
687                    .eq(lit("missing"))
688                    .and(col("tag_i32").eq(lit(1))),
689                vec![],
690            ),
691        ];
692        for (expr, expected) in cases {
693            let actual = applier_factory(expr.clone()).await;
694            assert_eq!(actual, expected, "{expr}");
695        }
696    }
697
698    #[tokio::test]
699    async fn test_create_json_hint_index_both_formats() {
700        for flat in [false, true] {
701            let (dir, factory) = PuffinManagerFactory::new_for_test_async("json_hint_index").await;
702            let mut metadata = (*mock_region_metadata()).clone();
703            let column = &mut metadata.column_metadatas[3].column_schema;
704            column.data_type =
705                ConcreteDataType::json2(datatypes::types::json_type::JsonNativeType::Object(
706                    datatypes::types::json_type::JsonObjectType::new(),
707                ));
708            column.with_extension_type(&datatypes::extension::json::Json2ExtensionType::new(
709                Arc::new(datatypes::extension::json::JsonMetadata::new(
710                    datatypes::json::JsonSettings::try_new(
711                        vec![datatypes::json::JsonTypeHint {
712                            path: vec!["value".into()],
713                            data_type: ConcreteDataType::int64_datatype(),
714                            inverted_index: true,
715                            skipping_index: None,
716                            fulltext_index: None,
717                        }],
718                        Some(0),
719                    )
720                    .unwrap(),
721                )),
722            ));
723            let metadata = Arc::new(metadata);
724            let file_id = FileId::random();
725            let target = IndexTarget::new_json_path(4, vec!["value".into()]).unwrap();
726            let mut creator = InvertedIndexer::new(
727                file_id,
728                &metadata,
729                new_intm_mgr(dir.path().to_string_lossy()).await,
730                None,
731                NonZeroUsize::new(2).unwrap(),
732                vec![target.clone()],
733            );
734            let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(1), None]));
735            let root: ArrayRef = Arc::new(StructArray::new(
736                vec![Field::new("value", values.data_type().clone(), true)].into(),
737                vec![values],
738                Some(NullBuffer::from(vec![true, false, true, true])),
739            ));
740            for start in [0, 2] {
741                let array = root.slice(start, 2);
742                if flat {
743                    let batch = RecordBatch::try_new(
744                        Arc::new(Schema::new(vec![Field::new(
745                            "field_u64",
746                            array.data_type().clone(),
747                            true,
748                        )])),
749                        vec![array],
750                    )
751                    .unwrap();
752                    creator.update_flat(&batch).await.unwrap();
753                } else {
754                    let mut batch = Batch::new(
755                        vec![],
756                        Arc::new(UInt64Vector::from_vec(vec![0, 1])),
757                        Arc::new(UInt64Vector::from_vec(vec![0, 0])),
758                        Arc::new(UInt8Vector::from_vec(vec![1, 1])),
759                        vec![BatchColumn {
760                            column_id: 4,
761                            data: Helper::try_into_vector(array).unwrap(),
762                        }],
763                    )
764                    .unwrap();
765                    creator.update(&mut batch).await.unwrap();
766                }
767            }
768            assert_eq!(creator.indexed_targets().collect::<Vec<_>>(), vec![&target]);
769            let store = mock_object_store();
770            let manager = factory.build(
771                store.clone(),
772                RegionFilePathFactory::new("table0".into(), PathType::Bare),
773            );
774            let index_id = RegionIndexId::new(RegionFileId::new(metadata.region_id, file_id), 0);
775            let mut writer = manager.writer(&index_id).await.unwrap();
776            let (rows, bytes) = creator.finish(&mut writer).await.unwrap();
777            assert_eq!(rows, 4);
778            assert!(bytes > 0);
779            writer.finish().await.unwrap();
780            let reader = manager.reader(&index_id).await.unwrap();
781            let blob = reader
782                .blob(INDEX_BLOB_TYPE)
783                .await
784                .unwrap()
785                .reader()
786                .await
787                .unwrap();
788            let reader = InvertedIndexBlobReader::new(blob);
789            let metas = reader.metadata(None).await.unwrap();
790            assert_eq!(metas.total_row_count, 4);
791            assert_eq!(metas.metas.len(), 1);
792            let meta = &metas.metas[&target.encode()];
793            let fst = reader
794                .fst(
795                    meta.base_offset + u64::from(meta.relative_fst_offset),
796                    meta.fst_size,
797                    None,
798                )
799                .await
800                .unwrap();
801            let mut encoded = Vec::new();
802            IndexValueCodec::encode_nonnull_value(
803                ValueRef::Int64(1),
804                &SortField::new(ConcreteDataType::int64_datatype()),
805                &mut encoded,
806            )
807            .unwrap();
808            assert!(fst.get(&encoded).is_some());
809            // Null parent and leaf rows must not add values to the FST.
810            assert_eq!(fst.len(), 1);
811            assert!(meta.null_bitmap_size > 0);
812            // Apply the generated index to both matching and non-matching values.
813            for (value, expected) in [(2, vec![]), (1, vec![0, 1])] {
814                let mut bytes = Vec::new();
815                IndexValueCodec::encode_nonnull_value(
816                    ValueRef::Int64(value),
817                    &SortField::new(ConcreteDataType::int64_datatype()),
818                    &mut bytes,
819                )
820                .unwrap();
821                let predicates = BTreeMap::from([(
822                    target.clone(),
823                    vec![Predicate::InList(InListPredicate {
824                        list: BTreeSet::from([bytes]),
825                    })],
826                )]);
827                let applier = InvertedIndexApplier::new(
828                    "table0".into(),
829                    PathType::Bare,
830                    store.clone(),
831                    factory.clone(),
832                    predicates,
833                    BTreeMap::from([(target.clone(), ConcreteDataType::int64_datatype())]),
834                )
835                .unwrap();
836                let plan = applier.plan_for_sst(&metadata).unwrap().unwrap();
837                let output = applier
838                    .apply(index_id, None, &plan.index_applier, None)
839                    .await
840                    .unwrap();
841                assert_eq!(
842                    output.matched_segment_ids.iter_ones().collect::<Vec<_>>(),
843                    expected
844                );
845            }
846        }
847    }
848
849    #[tokio::test]
850    async fn test_create_and_query_get_key() {
851        let rows = BTreeSet::from_iter([
852            ("aaa", 1, [1, 2]),
853            ("aaa", 2, [2, 3]),
854            ("aaa", 3, [3, 4]),
855            ("aab", 1, [4, 5]),
856            ("aab", 2, [5, 6]),
857            ("aab", 3, [6, 7]),
858            ("abc", 1, [7, 8]),
859            ("abc", 2, [8, 9]),
860            ("abc", 3, [9, 10]),
861        ]);
862
863        let applier_factory = build_applier_factory(
864            "test_create_and_query_get_key_",
865            rows,
866            HashSet::from([1, 2, 4]),
867        )
868        .await;
869
870        let expr = col("tag_str").eq(lit("aaa"));
871        let res = applier_factory(expr).await;
872        assert_eq!(res, vec![0, 1, 2]);
873
874        let expr = col("tag_i32").eq(lit(2));
875        let res = applier_factory(expr).await;
876        assert_eq!(res, vec![1, 4, 7]);
877
878        let expr = col("tag_str").eq(lit("aaa")).and(col("tag_i32").eq(lit(2)));
879        let res = applier_factory(expr).await;
880        assert_eq!(res, vec![1]);
881
882        let expr = col("tag_str")
883            .eq(lit("aaa"))
884            .or(col("tag_str").eq(lit("abc")));
885        let res = applier_factory(expr).await;
886        assert_eq!(res, vec![0, 1, 2, 6, 7, 8]);
887
888        let expr = col("tag_str").in_list(vec![lit("aaa"), lit("abc")], false);
889        let res = applier_factory(expr).await;
890        assert_eq!(res, vec![0, 1, 2, 6, 7, 8]);
891
892        let expr = col("field_u64").eq(lit(2u64));
893        let res = applier_factory(expr).await;
894        assert_eq!(res, vec![0, 1]);
895    }
896
897    #[tokio::test]
898    async fn test_create_and_query_range() {
899        let rows = BTreeSet::from_iter([
900            ("aaa", 1, [1, 2]),
901            ("aaa", 2, [2, 3]),
902            ("aaa", 3, [3, 4]),
903            ("aab", 1, [4, 5]),
904            ("aab", 2, [5, 6]),
905            ("aab", 3, [6, 7]),
906            ("abc", 1, [7, 8]),
907            ("abc", 2, [8, 9]),
908            ("abc", 3, [9, 10]),
909        ]);
910
911        let applier_factory = build_applier_factory(
912            "test_create_and_query_range_",
913            rows,
914            HashSet::from([1, 2, 4]),
915        )
916        .await;
917
918        let expr = col("tag_str").between(lit("aaa"), lit("aab"));
919        let res = applier_factory(expr).await;
920        assert_eq!(res, vec![0, 1, 2, 3, 4, 5]);
921
922        let expr = col("tag_i32").between(lit(2), lit(3));
923        let res = applier_factory(expr).await;
924        assert_eq!(res, vec![1, 2, 4, 5, 7, 8]);
925
926        let expr = col("tag_str").between(lit("aaa"), lit("aaa"));
927        let res = applier_factory(expr).await;
928        assert_eq!(res, vec![0, 1, 2]);
929
930        let expr = col("tag_i32").between(lit(2), lit(2));
931        let res = applier_factory(expr).await;
932        assert_eq!(res, vec![1, 4, 7]);
933
934        let expr = col("field_u64").between(lit(2u64), lit(5u64));
935        let res = applier_factory(expr).await;
936        assert_eq!(res, vec![0, 1, 2, 3, 4]);
937    }
938
939    #[tokio::test]
940    async fn test_create_and_query_comparison() {
941        let rows = BTreeSet::from_iter([
942            ("aaa", 1, [1, 2]),
943            ("aaa", 2, [2, 3]),
944            ("aaa", 3, [3, 4]),
945            ("aab", 1, [4, 5]),
946            ("aab", 2, [5, 6]),
947            ("aab", 3, [6, 7]),
948            ("abc", 1, [7, 8]),
949            ("abc", 2, [8, 9]),
950            ("abc", 3, [9, 10]),
951        ]);
952
953        let applier_factory = build_applier_factory(
954            "test_create_and_query_comparison_",
955            rows,
956            HashSet::from([1, 2, 4]),
957        )
958        .await;
959
960        let expr = col("tag_str").lt(lit("aab"));
961        let res = applier_factory(expr).await;
962        assert_eq!(res, vec![0, 1, 2]);
963
964        let expr = col("tag_i32").lt(lit(2));
965        let res = applier_factory(expr).await;
966        assert_eq!(res, vec![0, 3, 6]);
967
968        let expr = col("field_u64").lt(lit(2u64));
969        let res = applier_factory(expr).await;
970        assert_eq!(res, vec![0]);
971
972        let expr = col("tag_str").gt(lit("aab"));
973        let res = applier_factory(expr).await;
974        assert_eq!(res, vec![6, 7, 8]);
975
976        let expr = col("tag_i32").gt(lit(2));
977        let res = applier_factory(expr).await;
978        assert_eq!(res, vec![2, 5, 8]);
979
980        let expr = col("field_u64").gt(lit(8u64));
981        let res = applier_factory(expr).await;
982        assert_eq!(res, vec![7, 8]);
983
984        let expr = col("tag_str").lt_eq(lit("aab"));
985        let res = applier_factory(expr).await;
986        assert_eq!(res, vec![0, 1, 2, 3, 4, 5]);
987
988        let expr = col("tag_i32").lt_eq(lit(2));
989        let res = applier_factory(expr).await;
990        assert_eq!(res, vec![0, 1, 3, 4, 6, 7]);
991
992        let expr = col("field_u64").lt_eq(lit(2u64));
993        let res = applier_factory(expr).await;
994        assert_eq!(res, vec![0, 1]);
995
996        let expr = col("tag_str").gt_eq(lit("aab"));
997        let res = applier_factory(expr).await;
998        assert_eq!(res, vec![3, 4, 5, 6, 7, 8]);
999
1000        let expr = col("tag_i32").gt_eq(lit(2));
1001        let res = applier_factory(expr).await;
1002        assert_eq!(res, vec![1, 2, 4, 5, 7, 8]);
1003
1004        let expr = col("field_u64").gt_eq(lit(8u64));
1005        let res = applier_factory(expr).await;
1006        assert_eq!(res, vec![6, 7, 8]);
1007
1008        let expr = col("tag_str")
1009            .gt(lit("aaa"))
1010            .and(col("tag_str").lt(lit("abc")));
1011        let res = applier_factory(expr).await;
1012        assert_eq!(res, vec![3, 4, 5]);
1013
1014        let expr = col("tag_i32").gt(lit(1)).and(col("tag_i32").lt(lit(3)));
1015        let res = applier_factory(expr).await;
1016        assert_eq!(res, vec![1, 4, 7]);
1017
1018        let expr = col("field_u64")
1019            .gt(lit(2u64))
1020            .and(col("field_u64").lt(lit(9u64)));
1021        let res = applier_factory(expr).await;
1022        assert_eq!(res, vec![1, 2, 3, 4, 5, 6, 7]);
1023    }
1024
1025    #[tokio::test]
1026    async fn test_create_and_query_regex() {
1027        let rows = BTreeSet::from_iter([
1028            ("aaa", 1, [1, 2]),
1029            ("aaa", 2, [2, 3]),
1030            ("aaa", 3, [3, 4]),
1031            ("aab", 1, [4, 5]),
1032            ("aab", 2, [5, 6]),
1033            ("aab", 3, [6, 7]),
1034            ("abc", 1, [7, 8]),
1035            ("abc", 2, [8, 9]),
1036            ("abc", 3, [9, 10]),
1037        ]);
1038
1039        let applier_factory = build_applier_factory(
1040            "test_create_and_query_regex_",
1041            rows,
1042            HashSet::from([1, 2, 4]),
1043        )
1044        .await;
1045
1046        let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit(".*"));
1047        let res = applier_factory(expr).await;
1048        assert_eq!(res, vec![0, 1, 2, 3, 4, 5, 6, 7, 8]);
1049
1050        let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit("a.*c"));
1051        let res = applier_factory(expr).await;
1052        assert_eq!(res, vec![6, 7, 8]);
1053
1054        let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit("a.*b$"));
1055        let res = applier_factory(expr).await;
1056        assert_eq!(res, vec![3, 4, 5]);
1057
1058        let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit("\\w"));
1059        let res = applier_factory(expr).await;
1060        assert_eq!(res, vec![0, 1, 2, 3, 4, 5, 6, 7, 8]);
1061
1062        let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit("\\d"));
1063        let res = applier_factory(expr).await;
1064        assert!(res.is_empty());
1065
1066        let expr = binary_expr(col("tag_str"), Operator::RegexMatch, lit("^aaa$"));
1067        let res = applier_factory(expr).await;
1068        assert_eq!(res, vec![0, 1, 2]);
1069    }
1070}
mito2/sst/index/inverted_index/creator.rs

mito2/sst/index/inverted_index/
creator.rs
