mito2/memtable/partition_tree/
tree.rs1use std::collections::{BTreeMap, HashSet, VecDeque};
18use std::sync::{Arc, RwLock};
19use std::time::{Duration, Instant};
20
21use api::v1::OpType;
22use common_recordbatch::filter::SimpleFilterEvaluator;
23use common_time::Timestamp;
24use datafusion_common::ScalarValue;
25use datatypes::prelude::ValueRef;
26use mito_codec::key_values::KeyValue;
27use mito_codec::primary_key_filter::is_partition_column;
28use mito_codec::row_converter::sparse::{FieldWithId, SparseEncoder};
29use mito_codec::row_converter::{PrimaryKeyCodec, SortField};
30use snafu::{ensure, ResultExt};
31use store_api::codec::PrimaryKeyEncoding;
32use store_api::metadata::RegionMetadataRef;
33use store_api::storage::{ColumnId, SequenceNumber};
34use table::predicate::Predicate;
35
36use crate::error::{
37 EncodeSnafu, EncodeSparsePrimaryKeySnafu, PrimaryKeyLengthMismatchSnafu, Result,
38};
39use crate::flush::WriteBufferManagerRef;
40use crate::memtable::partition_tree::partition::{
41 Partition, PartitionKey, PartitionReader, PartitionRef, ReadPartitionContext,
42};
43use crate::memtable::partition_tree::PartitionTreeConfig;
44use crate::memtable::stats::WriteMetrics;
45use crate::memtable::{BoxedBatchIterator, KeyValues};
46use crate::metrics::{PARTITION_TREE_READ_STAGE_ELAPSED, READ_ROWS_TOTAL, READ_STAGE_ELAPSED};
47use crate::read::dedup::LastNonNullIter;
48use crate::read::Batch;
49use crate::region::options::MergeMode;
50
51pub struct PartitionTree {
53 config: PartitionTreeConfig,
55 pub(crate) metadata: RegionMetadataRef,
57 row_codec: Arc<dyn PrimaryKeyCodec>,
59 partitions: RwLock<BTreeMap<PartitionKey, PartitionRef>>,
61 is_partitioned: bool,
63 write_buffer_manager: Option<WriteBufferManagerRef>,
65 sparse_encoder: Arc<SparseEncoder>,
66}
67
68impl PartitionTree {
69 pub fn new(
71 row_codec: Arc<dyn PrimaryKeyCodec>,
72 metadata: RegionMetadataRef,
73 config: &PartitionTreeConfig,
74 write_buffer_manager: Option<WriteBufferManagerRef>,
75 ) -> Self {
76 let sparse_encoder = SparseEncoder::new(
77 metadata
78 .primary_key_columns()
79 .map(|c| FieldWithId {
80 field: SortField::new(c.column_schema.data_type.clone()),
81 column_id: c.column_id,
82 })
83 .collect(),
84 );
85 let is_partitioned = Partition::has_multi_partitions(&metadata);
86 let mut config = config.clone();
87 if config.merge_mode == MergeMode::LastNonNull {
88 config.dedup = false;
89 }
90
91 PartitionTree {
92 config,
93 metadata,
94 row_codec,
95 partitions: Default::default(),
96 is_partitioned,
97 write_buffer_manager,
98 sparse_encoder: Arc::new(sparse_encoder),
99 }
100 }
101
102 fn verify_primary_key_length(&self, kv: &KeyValue) -> Result<()> {
103 if let Some(expected_num_fields) = self.row_codec.num_fields() {
105 ensure!(
106 expected_num_fields == kv.num_primary_keys(),
107 PrimaryKeyLengthMismatchSnafu {
108 expect: expected_num_fields,
109 actual: kv.num_primary_keys(),
110 }
111 );
112 }
113 Ok(())
115 }
116
117 fn encode_sparse_primary_key(&self, kv: &KeyValue, buffer: &mut Vec<u8>) -> Result<()> {
119 if kv.primary_key_encoding() == PrimaryKeyEncoding::Sparse {
120 let ValueRef::Binary(primary_key) = kv.primary_keys().next().unwrap() else {
123 return EncodeSparsePrimaryKeySnafu {
124 reason: "sparse primary key is not binary".to_string(),
125 }
126 .fail();
127 };
128 buffer.extend_from_slice(primary_key);
129 } else {
130 self.sparse_encoder
132 .encode_to_vec(kv.primary_keys(), buffer)
133 .context(EncodeSnafu)?;
134 }
135 Ok(())
136 }
137
138 pub fn write(
144 &self,
145 kvs: &KeyValues,
146 pk_buffer: &mut Vec<u8>,
147 metrics: &mut WriteMetrics,
148 ) -> Result<()> {
149 let has_pk = !self.metadata.primary_key.is_empty();
150
151 for kv in kvs.iter() {
152 self.verify_primary_key_length(&kv)?;
153 let ts = kv.timestamp().as_timestamp().unwrap().unwrap().value();
155 metrics.min_ts = metrics.min_ts.min(ts);
156 metrics.max_ts = metrics.max_ts.max(ts);
157 metrics.value_bytes += kv.fields().map(|v| v.data_size()).sum::<usize>();
158
159 if !has_pk {
160 self.write_no_key(kv)?;
162 continue;
163 }
164
165 pk_buffer.clear();
167 if self.is_partitioned {
168 self.encode_sparse_primary_key(&kv, pk_buffer)?;
169 } else {
170 self.row_codec
171 .encode_key_value(&kv, pk_buffer)
172 .context(EncodeSnafu)?;
173 }
174
175 self.write_with_key(pk_buffer, kv, metrics)?;
177 }
178
179 metrics.value_bytes +=
180 kvs.num_rows() * (std::mem::size_of::<Timestamp>() + std::mem::size_of::<OpType>());
181
182 Ok(())
183 }
184
185 pub fn write_one(
190 &self,
191 kv: KeyValue,
192 pk_buffer: &mut Vec<u8>,
193 metrics: &mut WriteMetrics,
194 ) -> Result<()> {
195 let has_pk = !self.metadata.primary_key.is_empty();
196
197 self.verify_primary_key_length(&kv)?;
198 let ts = kv.timestamp().as_timestamp().unwrap().unwrap().value();
200 metrics.min_ts = metrics.min_ts.min(ts);
201 metrics.max_ts = metrics.max_ts.max(ts);
202 metrics.value_bytes += kv.fields().map(|v| v.data_size()).sum::<usize>();
203
204 if !has_pk {
205 return self.write_no_key(kv);
207 }
208
209 pk_buffer.clear();
211 if self.is_partitioned {
212 self.encode_sparse_primary_key(&kv, pk_buffer)?;
213 } else {
214 self.row_codec
215 .encode_key_value(&kv, pk_buffer)
216 .context(EncodeSnafu)?;
217 }
218
219 self.write_with_key(pk_buffer, kv, metrics)?;
221
222 metrics.value_bytes += std::mem::size_of::<Timestamp>() + std::mem::size_of::<OpType>();
223
224 Ok(())
225 }
226
227 pub fn read(
229 &self,
230 projection: Option<&[ColumnId]>,
231 predicate: Option<Predicate>,
232 sequence: Option<SequenceNumber>,
233 mem_scan_metrics: Option<crate::memtable::MemScanMetrics>,
234 ) -> Result<BoxedBatchIterator> {
235 let start = Instant::now();
236 let projection: HashSet<_> = if let Some(projection) = projection {
238 projection.iter().copied().collect()
239 } else {
240 self.metadata.field_columns().map(|c| c.column_id).collect()
241 };
242
243 let filters = predicate
244 .map(|predicate| {
245 predicate
246 .exprs()
247 .iter()
248 .filter_map(SimpleFilterEvaluator::try_new)
249 .collect::<Vec<_>>()
250 })
251 .unwrap_or_default();
252
253 let mut tree_iter_metric = TreeIterMetrics::default();
254 let partitions = self.prune_partitions(&filters, &mut tree_iter_metric);
255
256 let mut iter = TreeIter {
257 sequence,
258 partitions,
259 current_reader: None,
260 metrics: tree_iter_metric,
261 mem_scan_metrics,
262 };
263 let context = ReadPartitionContext::new(
264 self.metadata.clone(),
265 self.row_codec.clone(),
266 projection,
267 Arc::new(filters),
268 );
269 iter.fetch_next_partition(context)?;
270
271 iter.metrics.iter_elapsed += start.elapsed();
272
273 if self.config.merge_mode == MergeMode::LastNonNull {
274 let iter = LastNonNullIter::new(iter);
275 Ok(Box::new(iter))
276 } else {
277 Ok(Box::new(iter))
278 }
279 }
280
281 pub fn is_empty(&self) -> bool {
285 let partitions = self.partitions.read().unwrap();
286 partitions.values().all(|part| !part.has_data())
287 }
288
289 pub fn freeze(&self) -> Result<()> {
293 let partitions = self.partitions.read().unwrap();
294 for partition in partitions.values() {
295 partition.freeze()?;
296 }
297 Ok(())
298 }
299
300 pub fn fork(&self, metadata: RegionMetadataRef) -> PartitionTree {
303 if self.metadata.schema_version != metadata.schema_version
304 || self.metadata.column_metadatas != metadata.column_metadatas
305 {
306 return PartitionTree::new(
308 self.row_codec.clone(),
309 metadata,
310 &self.config,
311 self.write_buffer_manager.clone(),
312 );
313 }
314
315 let mut total_shared_size = 0;
316 let mut part_infos = {
317 let partitions = self.partitions.read().unwrap();
318 partitions
319 .iter()
320 .filter_map(|(part_key, part)| {
321 let stats = part.stats();
322 if stats.num_rows > 0 {
323 total_shared_size += stats.shared_memory_size;
325 Some((*part_key, part.clone(), stats))
326 } else {
327 None
328 }
329 })
330 .collect::<Vec<_>>()
331 };
332
333 let fork_size = self.config.fork_dictionary_bytes.as_bytes() as usize;
335 if total_shared_size > fork_size {
336 part_infos.sort_unstable_by_key(|info| info.2.shared_memory_size);
338 while total_shared_size > fork_size {
339 let Some(info) = part_infos.pop() else {
340 break;
341 };
342
343 common_telemetry::debug!(
344 "Evict partition {} with memory size {}, {} shards",
345 info.0,
346 info.2.shared_memory_size,
347 info.2.shard_num,
348 );
349
350 total_shared_size -= info.2.shared_memory_size;
351 }
352 }
353
354 let mut forked = BTreeMap::new();
355 for (part_key, part, _) in part_infos {
356 let forked_part = part.fork(&metadata, &self.config);
357 forked.insert(part_key, Arc::new(forked_part));
358 }
359
360 PartitionTree {
361 config: self.config.clone(),
362 metadata,
363 row_codec: self.row_codec.clone(),
364 partitions: RwLock::new(forked),
365 is_partitioned: self.is_partitioned,
366 write_buffer_manager: self.write_buffer_manager.clone(),
367 sparse_encoder: self.sparse_encoder.clone(),
368 }
369 }
370
371 pub(crate) fn write_buffer_manager(&self) -> Option<WriteBufferManagerRef> {
373 self.write_buffer_manager.clone()
374 }
375
376 fn write_with_key(
377 &self,
378 primary_key: &mut Vec<u8>,
379 key_value: KeyValue,
380 metrics: &mut WriteMetrics,
381 ) -> Result<()> {
382 let partition_key = Partition::get_partition_key(&key_value, self.is_partitioned);
383 let partition = self.get_or_create_partition(partition_key);
384
385 partition.write_with_key(
386 primary_key,
387 self.row_codec.as_ref(),
388 key_value,
389 self.is_partitioned, metrics,
391 )
392 }
393
394 fn write_no_key(&self, key_value: KeyValue) -> Result<()> {
395 let partition_key = Partition::get_partition_key(&key_value, self.is_partitioned);
396 let partition = self.get_or_create_partition(partition_key);
397
398 partition.write_no_key(key_value)
399 }
400
401 fn get_or_create_partition(&self, partition_key: PartitionKey) -> PartitionRef {
402 let mut partitions = self.partitions.write().unwrap();
403 partitions
404 .entry(partition_key)
405 .or_insert_with(|| Arc::new(Partition::new(self.metadata.clone(), &self.config)))
406 .clone()
407 }
408
409 fn prune_partitions(
410 &self,
411 filters: &[SimpleFilterEvaluator],
412 metrics: &mut TreeIterMetrics,
413 ) -> VecDeque<PartitionRef> {
414 let partitions = self.partitions.read().unwrap();
415 metrics.partitions_total = partitions.len();
416 if !self.is_partitioned {
417 return partitions.values().cloned().collect();
418 }
419
420 let mut pruned = VecDeque::new();
421 for (key, partition) in partitions.iter() {
423 let mut is_needed = true;
424 for filter in filters {
425 if !is_partition_column(filter.column_name()) {
426 continue;
427 }
428
429 if !filter
430 .evaluate_scalar(&ScalarValue::UInt32(Some(*key)))
431 .unwrap_or(true)
432 {
433 is_needed = false;
434 }
435 }
436
437 if is_needed {
438 pruned.push_back(partition.clone());
439 }
440 }
441 metrics.partitions_after_pruning = pruned.len();
442 pruned
443 }
444
445 pub(crate) fn series_count(&self) -> usize {
447 self.partitions
448 .read()
449 .unwrap()
450 .values()
451 .map(|p| p.series_count())
452 .sum()
453 }
454}
455
456#[derive(Default)]
457struct TreeIterMetrics {
458 iter_elapsed: Duration,
459 fetch_partition_elapsed: Duration,
460 rows_fetched: usize,
461 batches_fetched: usize,
462 partitions_total: usize,
463 partitions_after_pruning: usize,
464}
465
466struct TreeIter {
467 sequence: Option<SequenceNumber>,
469 partitions: VecDeque<PartitionRef>,
470 current_reader: Option<PartitionReader>,
471 metrics: TreeIterMetrics,
472 mem_scan_metrics: Option<crate::memtable::MemScanMetrics>,
473}
474
475impl TreeIter {
476 fn report_mem_scan_metrics(&mut self) {
477 if let Some(mem_scan_metrics) = self.mem_scan_metrics.take() {
478 let inner = crate::memtable::MemScanMetricsData {
479 total_series: 0, num_rows: self.metrics.rows_fetched,
481 num_batches: self.metrics.batches_fetched,
482 scan_cost: self.metrics.iter_elapsed,
483 };
484 mem_scan_metrics.merge_inner(&inner);
485 }
486 }
487}
488
489impl Drop for TreeIter {
490 fn drop(&mut self) {
491 self.report_mem_scan_metrics();
493
494 READ_ROWS_TOTAL
495 .with_label_values(&["partition_tree_memtable"])
496 .inc_by(self.metrics.rows_fetched as u64);
497 PARTITION_TREE_READ_STAGE_ELAPSED
498 .with_label_values(&["fetch_next_partition"])
499 .observe(self.metrics.fetch_partition_elapsed.as_secs_f64());
500 let scan_elapsed = self.metrics.iter_elapsed.as_secs_f64();
501 READ_STAGE_ELAPSED
502 .with_label_values(&["scan_memtable"])
503 .observe(scan_elapsed);
504 common_telemetry::debug!(
505 "TreeIter partitions total: {}, partitions after prune: {}, rows fetched: {}, batches fetched: {}, scan elapsed: {}",
506 self.metrics.partitions_total,
507 self.metrics.partitions_after_pruning,
508 self.metrics.rows_fetched,
509 self.metrics.batches_fetched,
510 scan_elapsed
511 );
512 }
513}
514
515impl Iterator for TreeIter {
516 type Item = Result<Batch>;
517
518 fn next(&mut self) -> Option<Self::Item> {
519 let start = Instant::now();
520 let res = self.next_batch().transpose();
521 self.metrics.iter_elapsed += start.elapsed();
522 res
523 }
524}
525
526impl TreeIter {
527 fn fetch_next_partition(&mut self, mut context: ReadPartitionContext) -> Result<()> {
529 let start = Instant::now();
530 while let Some(partition) = self.partitions.pop_front() {
531 let part_reader = partition.read(context)?;
532 if !part_reader.is_valid() {
533 context = part_reader.into_context();
534 continue;
535 }
536 self.current_reader = Some(part_reader);
537 break;
538 }
539 self.metrics.fetch_partition_elapsed += start.elapsed();
540 Ok(())
541 }
542
543 fn next_batch(&mut self) -> Result<Option<Batch>> {
545 let Some(part_reader) = &mut self.current_reader else {
546 self.report_mem_scan_metrics();
548 return Ok(None);
549 };
550
551 debug_assert!(part_reader.is_valid());
552 let batch = part_reader.convert_current_batch()?;
553 part_reader.next()?;
554 if part_reader.is_valid() {
555 self.metrics.rows_fetched += batch.num_rows();
556 self.metrics.batches_fetched += 1;
557 let mut batch = batch;
558 batch.filter_by_sequence(self.sequence)?;
559 return Ok(Some(batch));
560 }
561
562 let part_reader = self.current_reader.take().unwrap();
564 let context = part_reader.into_context();
565 self.fetch_next_partition(context)?;
566
567 self.metrics.rows_fetched += batch.num_rows();
568 self.metrics.batches_fetched += 1;
569 let mut batch = batch;
570 batch.filter_by_sequence(self.sequence)?;
571 Ok(Some(batch))
572 }
573}