1use std::collections::HashMap;
16use std::sync::Arc;
17use std::sync::atomic::AtomicUsize;
18
19use api::v1::SemanticType;
20use arrow_schema::extension::ExtensionType;
21use common_telemetry::{debug, warn};
22use datatypes::arrow::record_batch::RecordBatch;
23use datatypes::extension::json::Json2ExtensionType;
24use datatypes::vectors::Helper;
25use index::bloom_filter::creator::BloomFilterCreator;
26use index::target::IndexTarget;
27use mito_codec::index::{IndexValueCodec, IndexValuesCodec};
28use mito_codec::row_converter::sparse::SparsePrimaryKeyView;
29use mito_codec::row_converter::{SortField, SparseOffsetsCache};
30use puffin::puffin_manager::{PuffinWriter, PutOptions};
31use smallvec::SmallVec;
32use snafu::{OptionExt, ResultExt, ensure};
33use store_api::codec::PrimaryKeyEncoding;
34use store_api::metadata::RegionMetadataRef;
35use store_api::storage::{ColumnId, FileId};
36use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
37
38use crate::error::{
39 BiErrorsSnafu, BloomFilterFinishSnafu, DecodeSnafu, EncodeSnafu, IndexOptionsSnafu,
40 OperateAbortedIndexSnafu, PuffinAddBlobSnafu, PushBloomFilterValueSnafu, Result,
41};
42use crate::read::Batch;
43use crate::sst::index::bloom_filter::INDEX_BLOB_TYPE;
44use crate::sst::index::column::{column_index_rows, extract_json_type_hint_values};
45use crate::sst::index::intermediate::{
46 IntermediateLocation, IntermediateManager, TempFileProvider,
47};
48use crate::sst::index::primary_key::PrimaryKeyRuns;
49use crate::sst::index::puffin_manager::SstPuffinWriter;
50use crate::sst::index::statistics::{ByteCount, RowCount, Statistics};
51use crate::sst::index::{IndexBuildType, TYPE_BLOOM_FILTER_INDEX};
52
53const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192;
55
56pub struct BloomFilterIndexer {
58 creators: HashMap<IndexTarget, BloomFilterCreator>,
60
61 temp_file_provider: Arc<TempFileProvider>,
63
64 codec: IndexValuesCodec,
66 pk_offsets: SparseOffsetsCache,
68 value_buf: Vec<u8>,
70
71 aborted: bool,
73
74 stats: Statistics,
76
77 global_memory_usage: Arc<AtomicUsize>,
79
80 metadata: RegionMetadataRef,
82}
83
84impl BloomFilterIndexer {
85 pub fn new(
88 sst_file_id: FileId,
89 metadata: &RegionMetadataRef,
90 intermediate_manager: IntermediateManager,
91 memory_usage_threshold: Option<usize>,
92 build_type: IndexBuildType,
93 ) -> Result<Option<Self>> {
94 let mut creators = HashMap::new();
95
96 let temp_file_provider = Arc::new(TempFileProvider::new(
97 IntermediateLocation::new(&metadata.region_id, &sst_file_id),
98 intermediate_manager,
99 ));
100 let global_memory_usage = Arc::new(AtomicUsize::new(0));
101
102 for column in &metadata.column_metadatas {
103 let options =
104 column
105 .column_schema
106 .skipping_index_options()
107 .context(IndexOptionsSnafu {
108 column_name: &column.column_schema.name,
109 })?;
110
111 let mut targets = Vec::new();
112 if let Some(options) = options {
113 targets.push((IndexTarget::ColumnId(column.column_id), options));
114 }
115 if matches!(build_type, IndexBuildType::Flush | IndexBuildType::Compact)
119 && column.column_schema.data_type.is_json2()
120 && let Ok(Some(extension)) =
121 column.column_schema.extension_type::<Json2ExtensionType>()
122 {
123 for hint in extension.metadata().json_settings().type_hints() {
124 let Some(options) = &hint.skipping_index else {
125 continue;
126 };
127 if let Ok(target) =
128 IndexTarget::new_json_path(column.column_id, hint.path.clone())
129 {
130 targets.push((target, options.clone()));
131 }
132 }
133 }
134 for (target, options) in targets {
135 let creator = BloomFilterCreator::new(
136 options.granularity as _,
137 options.false_positive_rate(),
138 temp_file_provider.clone(),
139 global_memory_usage.clone(),
140 memory_usage_threshold,
141 );
142 creators.insert(target, creator);
143 }
144 }
145
146 if creators.is_empty() {
147 return Ok(None);
148 }
149
150 let codec = IndexValuesCodec::from_tag_columns(
151 metadata.primary_key_encoding,
152 metadata.primary_key_columns(),
153 );
154 let indexer = Self {
155 creators,
156 temp_file_provider,
157 codec,
158 pk_offsets: SparseOffsetsCache::new(),
159 value_buf: Vec::new(),
160 aborted: false,
161 stats: Statistics::new(TYPE_BLOOM_FILTER_INDEX),
162 global_memory_usage,
163 metadata: metadata.clone(),
164 };
165 Ok(Some(indexer))
166 }
167
168 pub async fn update(&mut self, batch: &mut Batch) -> Result<()> {
173 ensure!(!self.aborted, OperateAbortedIndexSnafu);
174
175 if self.creators.is_empty() {
176 return Ok(());
177 }
178
179 if let Err(update_err) = self.do_update(batch).await {
180 if let Err(err) = self.do_cleanup().await {
182 if cfg!(any(test, feature = "test")) {
183 panic!("Failed to clean up index creator, err: {err:?}",);
184 } else {
185 warn!(err; "Failed to clean up index creator");
186 }
187 }
188 return Err(update_err);
189 }
190
191 Ok(())
192 }
193
194 pub async fn update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
196 ensure!(!self.aborted, OperateAbortedIndexSnafu);
197
198 if self.creators.is_empty() || batch.num_rows() == 0 {
199 return Ok(());
200 }
201
202 if let Err(update_err) = self.do_update_flat(batch).await {
203 if let Err(err) = self.do_cleanup().await {
205 if cfg!(any(test, feature = "test")) {
206 panic!("Failed to clean up index creator, err: {err:?}",);
207 } else {
208 warn!(err; "Failed to clean up index creator");
209 }
210 }
211 return Err(update_err);
212 }
213
214 Ok(())
215 }
216
217 pub(crate) async fn finish(
222 &mut self,
223 puffin_writer: &mut SstPuffinWriter,
224 ) -> Result<(RowCount, ByteCount)> {
225 ensure!(!self.aborted, OperateAbortedIndexSnafu);
226
227 if self.stats.row_count() == 0 {
228 return Ok((0, 0));
230 }
231
232 let finish_res = self.do_finish(puffin_writer).await;
233 if let Err(err) = self.do_cleanup().await {
235 if cfg!(any(test, feature = "test")) {
236 panic!("Failed to clean up index creator, err: {err:?}",);
237 } else {
238 warn!(err; "Failed to clean up index creator");
239 }
240 }
241
242 finish_res.map(|_| (self.stats.row_count(), self.stats.byte_count()))
243 }
244
245 pub async fn abort(&mut self) -> Result<()> {
249 if self.aborted {
250 return Ok(());
251 }
252 self.aborted = true;
253
254 self.do_cleanup().await
255 }
256
257 async fn do_update(&mut self, batch: &mut Batch) -> Result<()> {
258 let mut guard = self.stats.record_update();
259
260 let n = batch.num_rows();
261 guard.inc_row_count(n);
262
263 for (target, creator) in &mut self.creators {
264 let (vector, field) = match target {
265 IndexTarget::ColumnId(col_id) => {
266 if let Some(col_info) = self.codec.pk_col_info(*col_id) {
267 let pk_idx = col_info.idx;
268 let field = &col_info.field;
269 let elems = batch
270 .pk_col_value(self.codec.decoder(), pk_idx, *col_id)?
271 .filter(|v| !v.is_null())
272 .map(|v| {
273 let mut buf = vec![];
274 IndexValueCodec::encode_nonnull_value(
275 v.as_value_ref(),
276 field,
277 &mut buf,
278 )
279 .context(EncodeSnafu)?;
280 Ok(buf)
281 })
282 .transpose()?;
283 creator
284 .push_n_row_elems(n, elems)
285 .await
286 .context(PushBloomFilterValueSnafu)?;
287 continue;
288 }
289
290 let Some(values) = batch.field_col_value(*col_id) else {
291 debug!(
292 "Column {} not found in the batch during building bloom filter index",
293 col_id
294 );
295 continue;
296 };
297 let field = SortField::new(values.data.data_type());
298 (values.data.clone(), field)
299 }
300 IndexTarget::JsonPath(path) => {
301 let root = batch
302 .field_col_value(path.column_id())
303 .with_context(|| crate::error::InvalidRecordBatchSnafu {
304 reason: format!(
305 "Missing JSON index column {} in batch",
306 path.column_id()
307 ),
308 })?
309 .data
310 .to_arrow_array();
311 let vector = extract_json_type_hint_values(&root, path.path())?;
312 let field = SortField::new(vector.data_type());
313 (vector, field)
314 }
315 };
316 for row in 0..n {
317 let value =
318 IndexValueCodec::encode_value(vector.get_ref(row), &field, &mut self.value_buf)
319 .context(EncodeSnafu)?;
320 creator
321 .push_n_row_elem(1, value)
322 .await
323 .context(PushBloomFilterValueSnafu)?;
324 }
325 }
326
327 Ok(())
328 }
329
330 async fn do_update_flat(&mut self, batch: &RecordBatch) -> Result<()> {
331 let mut guard = self.stats.record_update();
332
333 let n = batch.num_rows();
334 guard.inc_row_count(n);
335
336 let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse;
337 let mut sparse_columns: SmallVec<[(ColumnId, &mut BloomFilterCreator); 8]> =
338 SmallVec::new();
339
340 for (target, creator) in &mut self.creators {
341 let col_id = match target {
342 IndexTarget::ColumnId(id) => *id,
343 IndexTarget::JsonPath(path) => path.column_id(),
344 };
345 let column_meta = self.metadata.column_by_id(col_id).unwrap();
347 let column_name = &column_meta.column_schema.name;
348 if let Some(column_array) = batch.column_by_name(column_name) {
349 let vector = match target {
351 IndexTarget::ColumnId(_) => Helper::try_into_vector(column_array.clone())
352 .context(crate::error::ConvertVectorSnafu)?,
353 IndexTarget::JsonPath(path) => {
354 extract_json_type_hint_values(column_array, path.path())?
355 }
356 };
357 let sort_field = SortField::new(vector.data_type());
358
359 for (row, count) in column_index_rows(batch, column_meta.semantic_type) {
360 let elem = IndexValueCodec::encode_value(
361 vector.get_ref(row),
362 &sort_field,
363 &mut self.value_buf,
364 )
365 .context(EncodeSnafu)?;
366
367 creator
368 .push_n_row_elem(count, elem)
369 .await
370 .context(PushBloomFilterValueSnafu)?;
371 }
372 } else if matches!(target, IndexTarget::JsonPath(_)) {
373 return crate::error::InvalidRecordBatchSnafu {
374 reason: format!("Missing JSON index column {column_name} in batch"),
375 }
376 .fail();
377 } else if is_sparse && column_meta.semantic_type == SemanticType::Tag {
378 if self.codec.pk_col_info(col_id).is_some() {
379 sparse_columns.push((col_id, creator));
380 }
381 } else {
382 debug!(
383 "Column {} not found in the batch during building bloom filter index",
384 column_name
385 );
386 }
387 }
388
389 if !sparse_columns.is_empty() {
390 for (pk, count) in PrimaryKeyRuns::try_new(batch)? {
391 let mut view =
392 SparsePrimaryKeyView::new(pk, &mut self.pk_offsets).context(DecodeSnafu)?;
393 for (col_id, creator) in &mut sparse_columns {
394 let value = IndexValueCodec::encode_sparse_value(
395 &mut view,
396 *col_id,
397 &mut self.value_buf,
398 )
399 .context(DecodeSnafu)?;
400 creator
401 .push_n_row_elem(count, value)
402 .await
403 .context(PushBloomFilterValueSnafu)?;
404 }
405 }
406 }
407
408 Ok(())
409 }
410
411 async fn do_finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<()> {
413 let mut guard = self.stats.record_finish();
414
415 for (id, creator) in &mut self.creators {
416 let written_bytes = Self::do_finish_single_creator(id, creator, puffin_writer).await?;
417 guard.inc_byte_count(written_bytes);
418 }
419
420 Ok(())
421 }
422
423 async fn do_cleanup(&mut self) -> Result<()> {
424 let mut _guard = self.stats.record_cleanup();
425
426 self.creators.clear();
427 self.temp_file_provider.cleanup().await
428 }
429
430 async fn do_finish_single_creator(
451 target: &IndexTarget,
452 creator: &mut BloomFilterCreator,
453 puffin_writer: &mut SstPuffinWriter,
454 ) -> Result<ByteCount> {
455 let (tx, rx) = tokio::io::duplex(PIPE_BUFFER_SIZE_FOR_SENDING_BLOB);
456
457 let target_key = target.encode();
458 let blob_name = format!("{INDEX_BLOB_TYPE}-{target_key}");
459 let (index_finish, puffin_add_blob) = futures::join!(
460 creator.finish(tx.compat_write()),
461 puffin_writer.put_blob(
462 &blob_name,
463 rx.compat(),
464 PutOptions::default(),
465 Default::default(),
466 )
467 );
468
469 match (
470 puffin_add_blob.context(PuffinAddBlobSnafu),
471 index_finish.context(BloomFilterFinishSnafu),
472 ) {
473 (Err(e1), Err(e2)) => BiErrorsSnafu {
474 first: Box::new(e1),
475 second: Box::new(e2),
476 }
477 .fail()?,
478
479 (Ok(_), e @ Err(_)) => e?,
480 (e @ Err(_), Ok(_)) => e.map(|_| ())?,
481 (Ok(written_bytes), Ok(_)) => {
482 return Ok(written_bytes);
483 }
484 }
485
486 Ok(0)
487 }
488
489 pub fn memory_usage(&self) -> usize {
491 self.global_memory_usage
492 .load(std::sync::atomic::Ordering::Relaxed)
493 }
494
495 pub fn indexed_targets(&self) -> impl Iterator<Item = &IndexTarget> {
497 self.creators.keys()
498 }
499}
500
501#[cfg(test)]
502pub(crate) mod tests {
503
504 use api::v1::SemanticType;
505 use datatypes::data_type::ConcreteDataType;
506 use datatypes::schema::{ColumnSchema, SkippingIndexOptions, SkippingIndexType};
507 use datatypes::value::ValueRef;
508 use datatypes::vectors::{UInt8Vector, UInt64Vector};
509 use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl};
510 use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt};
511 use object_store::ObjectStore;
512 use object_store::services::Memory;
513 use puffin::puffin_manager::{PuffinManager, PuffinReader};
514 use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
515 use store_api::storage::RegionId;
516
517 use super::*;
518 use crate::access_layer::FilePathProvider;
519 use crate::read::BatchColumn;
520 use crate::sst::file::{RegionFileId, RegionIndexId};
521 use crate::sst::index::puffin_manager::PuffinManagerFactory;
522
523 pub fn mock_object_store() -> ObjectStore {
524 ObjectStore::new(Memory::default()).unwrap()
525 }
526
527 pub async fn new_intm_mgr(path: impl AsRef<str>) -> IntermediateManager {
528 IntermediateManager::init_fs(path).await.unwrap()
529 }
530
531 pub struct TestPathProvider;
532
533 impl FilePathProvider for TestPathProvider {
534 fn build_index_file_path(&self, file_id: RegionFileId) -> String {
535 file_id.file_id().to_string()
536 }
537
538 fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String {
539 index_id.file_id.file_id().to_string()
540 }
541
542 fn build_sst_file_path(&self, file_id: RegionFileId) -> String {
543 file_id.file_id().to_string()
544 }
545 }
546
547 pub fn mock_region_metadata() -> RegionMetadataRef {
564 let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
565 builder
566 .push_column_metadata(ColumnMetadata {
567 column_schema: ColumnSchema::new(
568 "tag_str",
569 ConcreteDataType::string_datatype(),
570 false,
571 )
572 .with_skipping_options(SkippingIndexOptions::new_unchecked(
573 2,
574 0.01,
575 SkippingIndexType::BloomFilter,
576 ))
577 .unwrap(),
578 semantic_type: SemanticType::Tag,
579 column_id: 1,
580 })
581 .push_column_metadata(ColumnMetadata {
582 column_schema: ColumnSchema::new(
583 "ts",
584 ConcreteDataType::timestamp_millisecond_datatype(),
585 false,
586 ),
587 semantic_type: SemanticType::Timestamp,
588 column_id: 2,
589 })
590 .push_column_metadata(ColumnMetadata {
591 column_schema: ColumnSchema::new(
592 "field_u64",
593 ConcreteDataType::uint64_datatype(),
594 false,
595 )
596 .with_skipping_options(SkippingIndexOptions::new_unchecked(
597 4,
598 0.01,
599 SkippingIndexType::BloomFilter,
600 ))
601 .unwrap(),
602 semantic_type: SemanticType::Field,
603 column_id: 3,
604 })
605 .primary_key(vec![1]);
606
607 Arc::new(builder.build().unwrap())
608 }
609
610 pub fn new_batch(str_tag: impl AsRef<str>, u64_field: impl IntoIterator<Item = u64>) -> Batch {
611 let fields = vec![(0, SortField::new(ConcreteDataType::string_datatype()))];
612 let codec = DensePrimaryKeyCodec::with_fields(fields);
613 let row: [ValueRef; 1] = [str_tag.as_ref().into()];
614 let primary_key = codec.encode(row.into_iter()).unwrap();
615
616 let u64_field = BatchColumn {
617 column_id: 3,
618 data: Arc::new(UInt64Vector::from_iter_values(u64_field)),
619 };
620 let num_rows = u64_field.data.len();
621
622 Batch::new(
623 primary_key,
624 Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
625 0, num_rows,
626 ))),
627 Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
628 0, num_rows,
629 ))),
630 Arc::new(UInt8Vector::from_iter_values(std::iter::repeat_n(
631 1, num_rows,
632 ))),
633 vec![u64_field],
634 )
635 .unwrap()
636 }
637
638 #[tokio::test]
639 async fn test_bloom_filter_indexer() {
640 let prefix = "test_bloom_filter_indexer_";
641 let tempdir = common_test_util::temp_dir::create_temp_dir(prefix);
642 let object_store = mock_object_store();
643 let intm_mgr = new_intm_mgr(tempdir.path().to_string_lossy()).await;
644 let region_metadata = mock_region_metadata();
645 let memory_usage_threshold = Some(1024);
646
647 let file_id = FileId::random();
648 let mut indexer = BloomFilterIndexer::new(
649 file_id,
650 ®ion_metadata,
651 intm_mgr,
652 memory_usage_threshold,
653 IndexBuildType::Flush,
654 )
655 .unwrap()
656 .unwrap();
657
658 let mut batch = new_batch("tag1", 0..10);
660 indexer.update(&mut batch).await.unwrap();
661
662 let mut batch = new_batch("tag2", 10..20);
663 indexer.update(&mut batch).await.unwrap();
664
665 let (_d, factory) = PuffinManagerFactory::new_for_test_async(prefix).await;
666 let puffin_manager = factory.build(object_store, TestPathProvider);
667
668 let file_id = RegionFileId::new(region_metadata.region_id, file_id);
669 let file_id = RegionIndexId::new(file_id, 0);
670 let mut puffin_writer = puffin_manager.writer(&file_id).await.unwrap();
671 let (row_count, byte_count) = indexer.finish(&mut puffin_writer).await.unwrap();
672 assert_eq!(row_count, 20);
673 assert!(byte_count > 0);
674 puffin_writer.finish().await.unwrap();
675
676 let puffin_reader = puffin_manager.reader(&file_id).await.unwrap();
677
678 {
680 let blob_guard = puffin_reader
681 .blob("greptime-bloom-filter-v1-1")
682 .await
683 .unwrap();
684 let reader = blob_guard.reader().await.unwrap();
685 let bloom_filter = BloomFilterReaderImpl::new(reader);
686 let metadata = bloom_filter.metadata(None).await.unwrap();
687
688 assert_eq!(metadata.segment_count, 10);
689 for i in 0..5 {
690 let loc = &metadata.bloom_filter_locs[metadata.segment_loc_indices[i] as usize];
691 let bf = bloom_filter.bloom_filter(loc, None).await.unwrap();
692 assert!(bf.contains(b"tag1"));
693 }
694 for i in 5..10 {
695 let loc = &metadata.bloom_filter_locs[metadata.segment_loc_indices[i] as usize];
696 let bf = bloom_filter.bloom_filter(loc, None).await.unwrap();
697 assert!(bf.contains(b"tag2"));
698 }
699 }
700
701 {
703 let sort_field = SortField::new(ConcreteDataType::uint64_datatype());
704
705 let blob_guard = puffin_reader
706 .blob("greptime-bloom-filter-v1-3")
707 .await
708 .unwrap();
709 let reader = blob_guard.reader().await.unwrap();
710 let bloom_filter = BloomFilterReaderImpl::new(reader);
711 let metadata = bloom_filter.metadata(None).await.unwrap();
712
713 assert_eq!(metadata.segment_count, 5);
714 for i in 0u64..20 {
715 let idx = i as usize / 4;
716 let loc = &metadata.bloom_filter_locs[metadata.segment_loc_indices[idx] as usize];
717 let bf = bloom_filter.bloom_filter(loc, None).await.unwrap();
718 let mut buf = vec![];
719 IndexValueCodec::encode_nonnull_value(ValueRef::UInt64(i), &sort_field, &mut buf)
720 .unwrap();
721
722 assert!(bf.contains(&buf));
723 }
724 }
725 }
726}