mito2/sst/index/bloom_filter/
creator.rs1use std::collections::HashMap;
16use std::sync::atomic::AtomicUsize;
17use std::sync::Arc;
18
19use common_telemetry::{debug, warn};
20use datatypes::schema::SkippingIndexType;
21use index::bloom_filter::creator::BloomFilterCreator;
22use puffin::puffin_manager::{PuffinWriter, PutOptions};
23use snafu::{ensure, ResultExt};
24use store_api::metadata::RegionMetadataRef;
25use store_api::storage::ColumnId;
26use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
27
28use crate::error::{
29 BiErrorsSnafu, BloomFilterFinishSnafu, IndexOptionsSnafu, OperateAbortedIndexSnafu,
30 PuffinAddBlobSnafu, PushBloomFilterValueSnafu, Result,
31};
32use crate::read::Batch;
33use crate::row_converter::SortField;
34use crate::sst::file::FileId;
35use crate::sst::index::bloom_filter::INDEX_BLOB_TYPE;
36use crate::sst::index::codec::{IndexValueCodec, IndexValuesCodec};
37use crate::sst::index::intermediate::{
38 IntermediateLocation, IntermediateManager, TempFileProvider,
39};
40use crate::sst::index::puffin_manager::SstPuffinWriter;
41use crate::sst::index::statistics::{ByteCount, RowCount, Statistics};
42use crate::sst::index::TYPE_BLOOM_FILTER_INDEX;
43
44const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192;
46
47pub struct BloomFilterIndexer {
49 creators: HashMap<ColumnId, BloomFilterCreator>,
51
52 temp_file_provider: Arc<TempFileProvider>,
54
55 codec: IndexValuesCodec,
57
58 aborted: bool,
60
61 stats: Statistics,
63
64 global_memory_usage: Arc<AtomicUsize>,
66}
67
68impl BloomFilterIndexer {
69 pub fn new(
71 sst_file_id: FileId,
72 metadata: &RegionMetadataRef,
73 intermediate_manager: IntermediateManager,
74 memory_usage_threshold: Option<usize>,
75 ) -> Result<Option<Self>> {
76 let mut creators = HashMap::new();
77
78 let temp_file_provider = Arc::new(TempFileProvider::new(
79 IntermediateLocation::new(&metadata.region_id, &sst_file_id),
80 intermediate_manager,
81 ));
82 let global_memory_usage = Arc::new(AtomicUsize::new(0));
83
84 for column in &metadata.column_metadatas {
85 let options =
86 column
87 .column_schema
88 .skipping_index_options()
89 .context(IndexOptionsSnafu {
90 column_name: &column.column_schema.name,
91 })?;
92
93 let options = match options {
94 Some(options) if options.index_type == SkippingIndexType::BloomFilter => options,
95 _ => continue,
96 };
97
98 let creator = BloomFilterCreator::new(
99 options.granularity as _,
100 temp_file_provider.clone(),
101 global_memory_usage.clone(),
102 memory_usage_threshold,
103 );
104 creators.insert(column.column_id, creator);
105 }
106
107 if creators.is_empty() {
108 return Ok(None);
109 }
110
111 let codec = IndexValuesCodec::from_tag_columns(
112 metadata.primary_key_encoding,
113 metadata.primary_key_columns(),
114 );
115 let indexer = Self {
116 creators,
117 temp_file_provider,
118 codec,
119 aborted: false,
120 stats: Statistics::new(TYPE_BLOOM_FILTER_INDEX),
121 global_memory_usage,
122 };
123 Ok(Some(indexer))
124 }
125
126 pub async fn update(&mut self, batch: &mut Batch) -> Result<()> {
131 ensure!(!self.aborted, OperateAbortedIndexSnafu);
132
133 if self.creators.is_empty() {
134 return Ok(());
135 }
136
137 if let Err(update_err) = self.do_update(batch).await {
138 if let Err(err) = self.do_cleanup().await {
140 if cfg!(any(test, feature = "test")) {
141 panic!("Failed to clean up index creator, err: {err:?}",);
142 } else {
143 warn!(err; "Failed to clean up index creator");
144 }
145 }
146 return Err(update_err);
147 }
148
149 Ok(())
150 }
151
152 pub async fn finish(
157 &mut self,
158 puffin_writer: &mut SstPuffinWriter,
159 ) -> Result<(RowCount, ByteCount)> {
160 ensure!(!self.aborted, OperateAbortedIndexSnafu);
161
162 if self.stats.row_count() == 0 {
163 return Ok((0, 0));
165 }
166
167 let finish_res = self.do_finish(puffin_writer).await;
168 if let Err(err) = self.do_cleanup().await {
170 if cfg!(any(test, feature = "test")) {
171 panic!("Failed to clean up index creator, err: {err:?}",);
172 } else {
173 warn!(err; "Failed to clean up index creator");
174 }
175 }
176
177 finish_res.map(|_| (self.stats.row_count(), self.stats.byte_count()))
178 }
179
180 pub async fn abort(&mut self) -> Result<()> {
184 if self.aborted {
185 return Ok(());
186 }
187 self.aborted = true;
188
189 self.do_cleanup().await
190 }
191
192 async fn do_update(&mut self, batch: &mut Batch) -> Result<()> {
193 let mut guard = self.stats.record_update();
194
195 let n = batch.num_rows();
196 guard.inc_row_count(n);
197
198 for (col_id, creator) in &mut self.creators {
199 match self.codec.pk_col_info(*col_id) {
200 Some(col_info) => {
202 let pk_idx = col_info.idx;
203 let field = &col_info.field;
204 let elems = batch
205 .pk_col_value(self.codec.decoder(), pk_idx, *col_id)?
206 .filter(|v| !v.is_null())
207 .map(|v| {
208 let mut buf = vec![];
209 IndexValueCodec::encode_nonnull_value(
210 v.as_value_ref(),
211 field,
212 &mut buf,
213 )?;
214 Ok(buf)
215 })
216 .transpose()?;
217 creator
218 .push_n_row_elems(n, elems)
219 .await
220 .context(PushBloomFilterValueSnafu)?;
221 }
222 None => {
224 let Some(values) = batch.field_col_value(*col_id) else {
225 debug!(
226 "Column {} not found in the batch during building bloom filter index",
227 col_id
228 );
229 continue;
230 };
231 let sort_field = SortField::new(values.data.data_type());
232 for i in 0..n {
233 let value = values.data.get_ref(i);
234 let elems = (!value.is_null())
235 .then(|| {
236 let mut buf = vec![];
237 IndexValueCodec::encode_nonnull_value(
238 value,
239 &sort_field,
240 &mut buf,
241 )?;
242 Ok(buf)
243 })
244 .transpose()?;
245
246 creator
247 .push_row_elems(elems)
248 .await
249 .context(PushBloomFilterValueSnafu)?;
250 }
251 }
252 }
253 }
254
255 Ok(())
256 }
257
258 async fn do_finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<()> {
260 let mut guard = self.stats.record_finish();
261
262 for (id, creator) in &mut self.creators {
263 let written_bytes = Self::do_finish_single_creator(id, creator, puffin_writer).await?;
264 guard.inc_byte_count(written_bytes);
265 }
266
267 Ok(())
268 }
269
270 async fn do_cleanup(&mut self) -> Result<()> {
271 let mut _guard = self.stats.record_cleanup();
272
273 self.creators.clear();
274 self.temp_file_provider.cleanup().await
275 }
276
277 async fn do_finish_single_creator(
298 col_id: &ColumnId,
299 creator: &mut BloomFilterCreator,
300 puffin_writer: &mut SstPuffinWriter,
301 ) -> Result<ByteCount> {
302 let (tx, rx) = tokio::io::duplex(PIPE_BUFFER_SIZE_FOR_SENDING_BLOB);
303
304 let blob_name = format!("{}-{}", INDEX_BLOB_TYPE, col_id);
305 let (index_finish, puffin_add_blob) = futures::join!(
306 creator.finish(tx.compat_write()),
307 puffin_writer.put_blob(
308 &blob_name,
309 rx.compat(),
310 PutOptions::default(),
311 Default::default(),
312 )
313 );
314
315 match (
316 puffin_add_blob.context(PuffinAddBlobSnafu),
317 index_finish.context(BloomFilterFinishSnafu),
318 ) {
319 (Err(e1), Err(e2)) => BiErrorsSnafu {
320 first: Box::new(e1),
321 second: Box::new(e2),
322 }
323 .fail()?,
324
325 (Ok(_), e @ Err(_)) => e?,
326 (e @ Err(_), Ok(_)) => e.map(|_| ())?,
327 (Ok(written_bytes), Ok(_)) => {
328 return Ok(written_bytes);
329 }
330 }
331
332 Ok(0)
333 }
334
335 pub fn memory_usage(&self) -> usize {
337 self.global_memory_usage
338 .load(std::sync::atomic::Ordering::Relaxed)
339 }
340
341 pub fn column_ids(&self) -> impl Iterator<Item = ColumnId> + use<'_> {
343 self.creators.keys().copied()
344 }
345}
346
347#[cfg(test)]
348pub(crate) mod tests {
349
350 use api::v1::SemanticType;
351 use datatypes::data_type::ConcreteDataType;
352 use datatypes::schema::{ColumnSchema, SkippingIndexOptions};
353 use datatypes::value::ValueRef;
354 use datatypes::vectors::{UInt64Vector, UInt8Vector};
355 use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl};
356 use object_store::services::Memory;
357 use object_store::ObjectStore;
358 use puffin::puffin_manager::{PuffinManager, PuffinReader};
359 use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
360 use store_api::storage::RegionId;
361
362 use super::*;
363 use crate::access_layer::FilePathProvider;
364 use crate::read::BatchColumn;
365 use crate::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt};
366 use crate::sst::index::puffin_manager::PuffinManagerFactory;
367
368 pub fn mock_object_store() -> ObjectStore {
369 ObjectStore::new(Memory::default()).unwrap().finish()
370 }
371
372 pub async fn new_intm_mgr(path: impl AsRef<str>) -> IntermediateManager {
373 IntermediateManager::init_fs(path).await.unwrap()
374 }
375
376 pub struct TestPathProvider;
377
378 impl FilePathProvider for TestPathProvider {
379 fn build_index_file_path(&self, file_id: FileId) -> String {
380 file_id.to_string()
381 }
382
383 fn build_sst_file_path(&self, file_id: FileId) -> String {
384 file_id.to_string()
385 }
386 }
387
388 pub fn mock_region_metadata() -> RegionMetadataRef {
405 let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
406 builder
407 .push_column_metadata(ColumnMetadata {
408 column_schema: ColumnSchema::new(
409 "tag_str",
410 ConcreteDataType::string_datatype(),
411 false,
412 )
413 .with_skipping_options(SkippingIndexOptions {
414 index_type: SkippingIndexType::BloomFilter,
415 granularity: 2,
416 })
417 .unwrap(),
418 semantic_type: SemanticType::Tag,
419 column_id: 1,
420 })
421 .push_column_metadata(ColumnMetadata {
422 column_schema: ColumnSchema::new(
423 "ts",
424 ConcreteDataType::timestamp_millisecond_datatype(),
425 false,
426 ),
427 semantic_type: SemanticType::Timestamp,
428 column_id: 2,
429 })
430 .push_column_metadata(ColumnMetadata {
431 column_schema: ColumnSchema::new(
432 "field_u64",
433 ConcreteDataType::uint64_datatype(),
434 false,
435 )
436 .with_skipping_options(SkippingIndexOptions {
437 index_type: SkippingIndexType::BloomFilter,
438 granularity: 4,
439 })
440 .unwrap(),
441 semantic_type: SemanticType::Field,
442 column_id: 3,
443 })
444 .primary_key(vec![1]);
445
446 Arc::new(builder.build().unwrap())
447 }
448
449 pub fn new_batch(str_tag: impl AsRef<str>, u64_field: impl IntoIterator<Item = u64>) -> Batch {
450 let fields = vec![(0, SortField::new(ConcreteDataType::string_datatype()))];
451 let codec = DensePrimaryKeyCodec::with_fields(fields);
452 let row: [ValueRef; 1] = [str_tag.as_ref().into()];
453 let primary_key = codec.encode(row.into_iter()).unwrap();
454
455 let u64_field = BatchColumn {
456 column_id: 3,
457 data: Arc::new(UInt64Vector::from_iter_values(u64_field)),
458 };
459 let num_rows = u64_field.data.len();
460
461 Batch::new(
462 primary_key,
463 Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
464 0, num_rows,
465 ))),
466 Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(
467 0, num_rows,
468 ))),
469 Arc::new(UInt8Vector::from_iter_values(std::iter::repeat_n(
470 1, num_rows,
471 ))),
472 vec![u64_field],
473 )
474 .unwrap()
475 }
476
477 #[tokio::test]
478 async fn test_bloom_filter_indexer() {
479 let prefix = "test_bloom_filter_indexer_";
480 let tempdir = common_test_util::temp_dir::create_temp_dir(prefix);
481 let object_store = mock_object_store();
482 let intm_mgr = new_intm_mgr(tempdir.path().to_string_lossy()).await;
483 let region_metadata = mock_region_metadata();
484 let memory_usage_threshold = Some(1024);
485
486 let mut indexer = BloomFilterIndexer::new(
487 FileId::random(),
488 ®ion_metadata,
489 intm_mgr,
490 memory_usage_threshold,
491 )
492 .unwrap()
493 .unwrap();
494
495 let mut batch = new_batch("tag1", 0..10);
497 indexer.update(&mut batch).await.unwrap();
498
499 let mut batch = new_batch("tag2", 10..20);
500 indexer.update(&mut batch).await.unwrap();
501
502 let (_d, factory) = PuffinManagerFactory::new_for_test_async(prefix).await;
503 let puffin_manager = factory.build(object_store, TestPathProvider);
504
505 let file_id = FileId::random();
506 let mut puffin_writer = puffin_manager.writer(&file_id).await.unwrap();
507 let (row_count, byte_count) = indexer.finish(&mut puffin_writer).await.unwrap();
508 assert_eq!(row_count, 20);
509 assert!(byte_count > 0);
510 puffin_writer.finish().await.unwrap();
511
512 let puffin_reader = puffin_manager.reader(&file_id).await.unwrap();
513
514 {
516 let blob_guard = puffin_reader
517 .blob("greptime-bloom-filter-v1-1")
518 .await
519 .unwrap();
520 let reader = blob_guard.reader().await.unwrap();
521 let bloom_filter = BloomFilterReaderImpl::new(reader);
522 let metadata = bloom_filter.metadata().await.unwrap();
523
524 assert_eq!(metadata.segment_count, 10);
525 for i in 0..5 {
526 let loc = &metadata.bloom_filter_locs[metadata.segment_loc_indices[i] as usize];
527 let bf = bloom_filter.bloom_filter(loc).await.unwrap();
528 assert!(bf.contains(b"tag1"));
529 }
530 for i in 5..10 {
531 let loc = &metadata.bloom_filter_locs[metadata.segment_loc_indices[i] as usize];
532 let bf = bloom_filter.bloom_filter(loc).await.unwrap();
533 assert!(bf.contains(b"tag2"));
534 }
535 }
536
537 {
539 let sort_field = SortField::new(ConcreteDataType::uint64_datatype());
540
541 let blob_guard = puffin_reader
542 .blob("greptime-bloom-filter-v1-3")
543 .await
544 .unwrap();
545 let reader = blob_guard.reader().await.unwrap();
546 let bloom_filter = BloomFilterReaderImpl::new(reader);
547 let metadata = bloom_filter.metadata().await.unwrap();
548
549 assert_eq!(metadata.segment_count, 5);
550 for i in 0u64..20 {
551 let idx = i as usize / 4;
552 let loc = &metadata.bloom_filter_locs[metadata.segment_loc_indices[idx] as usize];
553 let bf = bloom_filter.bloom_filter(loc).await.unwrap();
554 let mut buf = vec![];
555 IndexValueCodec::encode_nonnull_value(ValueRef::UInt64(i), &sort_field, &mut buf)
556 .unwrap();
557
558 assert!(bf.contains(&buf));
559 }
560 }
561 }
562}