1pub(crate) mod chunk_reader;
18pub mod context;
19pub mod part;
20pub mod part_reader;
21mod row_group_reader;
22
23use std::collections::{BTreeMap, HashSet};
24use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering};
25use std::sync::{Arc, LazyLock, Mutex, RwLock};
26use std::time::Instant;
27
28fn env_usize(name: &str, default: usize) -> usize {
30 std::env::var(name)
31 .ok()
32 .and_then(|v| v.parse().ok())
33 .unwrap_or(default)
34}
35
36use common_time::Timestamp;
37use datatypes::arrow::datatypes::SchemaRef;
38use mito_codec::key_values::KeyValue;
39use rayon::prelude::*;
40use serde::{Deserialize, Serialize};
41use serde_with::{DisplayFromStr, serde_as};
42use store_api::metadata::RegionMetadataRef;
43use store_api::mito_engine_options::FloatFieldEncoding;
44use store_api::storage::{ColumnId, FileId, RegionId, SequenceNumber, SequenceRange};
45use tokio::sync::Semaphore;
46
47use crate::compaction::{
48 Json2RewritePlans, collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema,
49};
50use crate::error::{Result, UnsupportedOperationSnafu};
51use crate::flush::WriteBufferManagerRef;
52use crate::memtable::bulk::context::BulkIterContext;
53use crate::memtable::bulk::part::{
54 BulkPart, BulkPartEncodeMetrics, BulkPartEncoder, MultiBulkPart, UnorderedPart,
55 should_prune_bulk_part,
56};
57use crate::memtable::bulk::part_reader::BulkPartBatchIter;
58use crate::memtable::stats::WriteMetrics;
59use crate::memtable::{
60 AllocTracker, BoxedBatchIterator, BoxedRecordBatchIterator, EncodedBulkPart, EncodedRange,
61 IterBuilder, KeyValues, MemScanMetrics, Memtable, MemtableBuilder, MemtableId, MemtableRange,
62 MemtableRangeContext, MemtableRanges, MemtableRef, MemtableStats, RangesOptions,
63 group_by_sequence,
64};
65use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow};
66use crate::read::flat_merge::FlatMergeIterator;
67use crate::region::options::MergeMode;
68use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE;
69use crate::sst::parquet::flat_format::field_column_start;
70
71const DEFAULT_MERGE_THRESHOLD: usize = 16;
73
74static MERGE_THRESHOLD: LazyLock<usize> =
76 LazyLock::new(|| env_usize("GREPTIME_BULK_MERGE_THRESHOLD", DEFAULT_MERGE_THRESHOLD));
77
78const DEFAULT_MAX_MERGE_GROUPS: usize = 32;
80
81static MAX_MERGE_GROUPS: LazyLock<usize> =
83 LazyLock::new(|| env_usize("GREPTIME_BULK_MAX_MERGE_GROUPS", DEFAULT_MAX_MERGE_GROUPS));
84
85pub(crate) static ENCODE_ROW_THRESHOLD: LazyLock<usize> = LazyLock::new(|| {
88 env_usize(
89 "GREPTIME_BULK_ENCODE_ROW_THRESHOLD",
90 10 * DEFAULT_ROW_GROUP_SIZE,
91 )
92});
93
94const DEFAULT_ENCODE_BYTES_THRESHOLD: usize = 64 * 1024 * 1024;
96
97const MAX_ENCODE_BYTES_THRESHOLD: usize = 512 * 1024 * 1024;
99
100const ENCODE_BYTES_THRESHOLD_DIVISOR: usize = 32;
102
103static ENCODE_BYTES_THRESHOLD_OVERRIDE: LazyLock<Option<usize>> = LazyLock::new(|| {
105 std::env::var("GREPTIME_BULK_ENCODE_BYTES_THRESHOLD")
106 .ok()
107 .and_then(|v| v.parse().ok())
108});
109
110fn adaptive_encode_bytes_threshold(global_write_buffer_bytes: usize) -> usize {
113 (global_write_buffer_bytes / ENCODE_BYTES_THRESHOLD_DIVISOR)
114 .clamp(DEFAULT_ENCODE_BYTES_THRESHOLD, MAX_ENCODE_BYTES_THRESHOLD)
115}
116
117fn default_encode_bytes_threshold() -> usize {
119 ENCODE_BYTES_THRESHOLD_OVERRIDE.unwrap_or(DEFAULT_ENCODE_BYTES_THRESHOLD)
120}
121
122#[serde_as]
124#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
125#[serde(default)]
126pub struct BulkMemtableConfig {
127 #[serde_as(as = "DisplayFromStr")]
129 pub merge_threshold: usize,
130 #[serde_as(as = "DisplayFromStr")]
132 pub encode_row_threshold: usize,
133 #[serde_as(as = "DisplayFromStr")]
135 pub encode_bytes_threshold: usize,
136 #[serde_as(as = "DisplayFromStr")]
138 pub max_merge_groups: usize,
139}
140
141impl Default for BulkMemtableConfig {
142 fn default() -> Self {
143 Self {
144 merge_threshold: *MERGE_THRESHOLD,
145 encode_row_threshold: *ENCODE_ROW_THRESHOLD,
146 encode_bytes_threshold: default_encode_bytes_threshold(),
147 max_merge_groups: *MAX_MERGE_GROUPS,
148 }
149 .sanitize()
150 }
151}
152
153impl BulkMemtableConfig {
154 pub(crate) fn default_for_write_buffer_size(global_write_buffer_bytes: usize) -> Self {
156 Self {
157 encode_bytes_threshold: ENCODE_BYTES_THRESHOLD_OVERRIDE
158 .unwrap_or_else(|| adaptive_encode_bytes_threshold(global_write_buffer_bytes)),
159 ..Default::default()
160 }
161 }
162
163 fn sanitize(mut self) -> Self {
164 if self.merge_threshold == 0 {
165 self.merge_threshold = DEFAULT_MERGE_THRESHOLD;
166 }
167 self
168 }
169}
170
171enum MergedPart {
173 Multi(MultiBulkPart),
175 Encoded(EncodedBulkPart),
177}
178
179struct CollectedParts {
181 groups: Vec<Vec<PartToMerge>>,
183}
184
185#[derive(Default)]
187struct BulkParts {
188 unordered_part: UnorderedPart,
190 parts: Vec<BulkPartWrapper>,
192}
193
194impl BulkParts {
195 fn num_parts(&self) -> usize {
197 let unordered_count = if self.unordered_part.is_empty() { 0 } else { 1 };
198 self.parts.len() + unordered_count
199 }
200
201 fn is_empty(&self) -> bool {
203 self.unordered_part.is_empty() && self.parts.is_empty()
204 }
205
206 fn may_merge_parts(&self, merge_threshold: usize) -> bool {
209 let merge_threshold = merge_threshold.max(2);
210 self.parts.len() >= merge_threshold
211 && self
212 .parts
213 .iter()
214 .filter(|wrapper| !wrapper.merging)
215 .take(merge_threshold)
216 .count()
217 == merge_threshold
218 }
219
220 fn is_merge_candidate_by_size(
221 is_encoded: bool,
222 estimated_size: usize,
223 encode_bytes_threshold: usize,
224 ) -> bool {
225 !is_encoded || estimated_size <= encode_bytes_threshold
226 }
227
228 fn should_compact_unordered_part(&self, encode_bytes_threshold: usize) -> bool {
230 self.unordered_part.should_compact()
231 || self.unordered_part.estimated_bytes() > encode_bytes_threshold
232 }
233
234 fn seal_unordered(&mut self, metadata: &RegionMetadataRef) -> Result<()> {
236 if let Some(part) = self.unordered_part.to_bulk_part(metadata)? {
237 self.parts.push(BulkPartWrapper {
238 part: PartToMerge::Bulk {
239 part,
240 file_id: FileId::random(),
241 },
242 merging: false,
243 });
244 self.unordered_part.clear();
245 }
246 Ok(())
247 }
248
249 fn pick_merge_groups(
252 &self,
253 merge_threshold: usize,
254 max_merge_groups: usize,
255 encode_bytes_threshold: usize,
256 ) -> Vec<Vec<usize>> {
257 if self.parts.len() < merge_threshold.max(2) || max_merge_groups == 0 {
258 return Vec::new();
259 }
260 let intervals = self
263 .parts
264 .iter()
265 .enumerate()
266 .map(|(index, wrapper)| {
267 let part = &wrapper.part;
268 (Some(index), (part.min_sequence(), part.max_sequence()))
269 })
270 .chain(
271 self.unordered_part
272 .sequence_bounds()
273 .map(|bounds| (None, bounds)),
274 )
275 .collect();
276 let components = group_by_sequence(intervals, |(_, bounds)| *bounds);
277 let mut groups = Vec::new();
278 let mut pending = Vec::new();
279 for component in components {
280 if groups.len() >= max_merge_groups {
281 break;
282 }
283 let Some(component) = component
284 .into_iter()
285 .map(|(index, _)| index)
286 .collect::<Option<Vec<_>>>()
287 else {
288 pending.clear();
289 continue;
290 };
291 let busy = component.iter().any(|&i| self.parts[i].merging);
292 let oversized = component.len() == 1 && {
293 let part = &self.parts[component[0]].part;
294 !Self::is_merge_candidate_by_size(
295 part.is_encoded(),
296 part.estimated_size(),
297 encode_bytes_threshold,
298 )
299 };
300 if busy || oversized {
301 pending.clear();
302 continue;
303 }
304 pending.extend(component);
305 if pending.len() >= merge_threshold.max(2) {
306 groups.push(std::mem::take(&mut pending));
307 }
308 }
309 groups
310 }
311
312 fn collect_parts_to_merge(
313 &mut self,
314 merge_threshold: usize,
315 max_merge_groups: usize,
316 encode_bytes_threshold: usize,
317 ) -> CollectedParts {
318 let groups = self
319 .pick_merge_groups(merge_threshold, max_merge_groups, encode_bytes_threshold)
320 .into_iter()
321 .map(|group| {
322 group
323 .into_iter()
324 .map(|i| {
325 let wrapper = &mut self.parts[i];
326 wrapper.merging = true;
327 wrapper.part.clone()
328 })
329 .collect()
330 })
331 .collect();
332 CollectedParts { groups }
333 }
334
335 fn install_merged_parts<I>(
338 &mut self,
339 merged_parts: I,
340 merged_file_ids: &HashSet<FileId>,
341 ) -> usize
342 where
343 I: IntoIterator<Item = MergedPart>,
344 {
345 let mut total_output_rows = 0;
346
347 for merged_part in merged_parts {
348 match merged_part {
349 MergedPart::Encoded(encoded_part) => {
350 total_output_rows += encoded_part.metadata().num_rows;
351 self.parts.push(BulkPartWrapper {
352 part: PartToMerge::Encoded {
353 part: encoded_part,
354 file_id: FileId::random(),
355 },
356 merging: false,
357 });
358 }
359 MergedPart::Multi(multi_part) => {
360 total_output_rows += multi_part.num_rows();
361 self.parts.push(BulkPartWrapper {
362 part: PartToMerge::Multi {
363 part: multi_part,
364 file_id: FileId::random(),
365 },
366 merging: false,
367 });
368 }
369 }
370 }
371
372 self.parts
373 .retain(|wrapper| !merged_file_ids.contains(&wrapper.file_id()));
374
375 total_output_rows
376 }
377
378 fn reset_merging_flags(&mut self, file_ids: &HashSet<FileId>) {
381 for wrapper in &mut self.parts {
382 if file_ids.contains(&wrapper.file_id()) {
383 wrapper.merging = false;
384 }
385 }
386 }
387}
388
389struct MergingFlagsGuard<'a> {
392 bulk_parts: &'a RwLock<BulkParts>,
393 file_ids: &'a HashSet<FileId>,
394 success: bool,
395}
396
397impl<'a> MergingFlagsGuard<'a> {
398 fn new(bulk_parts: &'a RwLock<BulkParts>, file_ids: &'a HashSet<FileId>) -> Self {
400 Self {
401 bulk_parts,
402 file_ids,
403 success: false,
404 }
405 }
406
407 fn mark_success(&mut self) {
410 self.success = true;
411 }
412}
413
414impl<'a> Drop for MergingFlagsGuard<'a> {
415 fn drop(&mut self) {
416 if !self.success
417 && let Ok(mut parts) = self.bulk_parts.write()
418 {
419 parts.reset_merging_flags(self.file_ids);
420 }
421 }
422}
423
424pub struct BulkMemtable {
426 id: MemtableId,
427 config: BulkMemtableConfig,
429 parts: Arc<RwLock<BulkParts>>,
430 metadata: RegionMetadataRef,
431 alloc_tracker: AllocTracker,
432 max_timestamp: AtomicI64,
433 min_timestamp: AtomicI64,
434 max_sequence: AtomicU64,
435 min_sequence: AtomicU64,
436 num_rows: AtomicUsize,
437 compactor: Arc<Mutex<MemtableCompactor>>,
439 compact_dispatcher: Option<Arc<CompactDispatcher>>,
441 compacting: Arc<AtomicBool>,
443 append_mode: bool,
445 merge_mode: MergeMode,
447 row_group_size: usize,
449 float_field_encoding: FloatFieldEncoding,
450}
451
452impl std::fmt::Debug for BulkMemtable {
453 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
454 f.debug_struct("BulkMemtable")
455 .field("id", &self.id)
456 .field("num_rows", &self.num_rows.load(Ordering::Relaxed))
457 .field("min_timestamp", &self.min_timestamp.load(Ordering::Relaxed))
458 .field("max_timestamp", &self.max_timestamp.load(Ordering::Relaxed))
459 .field("max_sequence", &self.max_sequence.load(Ordering::Relaxed))
460 .finish()
461 }
462}
463
464impl Memtable for BulkMemtable {
465 fn id(&self) -> MemtableId {
466 self.id
467 }
468
469 fn write(&self, _kvs: &KeyValues) -> Result<()> {
470 UnsupportedOperationSnafu {
471 err_msg: "write() is not supported for bulk memtable",
472 }
473 .fail()
474 }
475
476 fn write_one(&self, _key_value: KeyValue) -> Result<()> {
477 UnsupportedOperationSnafu {
478 err_msg: "write_one() is not supported for bulk memtable",
479 }
480 .fail()
481 }
482
483 fn write_bulk(&self, fragment: BulkPart) -> Result<()> {
484 let local_metrics = WriteMetrics {
485 key_bytes: 0,
486 value_bytes: fragment.estimated_size(),
487 min_ts: fragment.min_timestamp,
488 max_ts: fragment.max_timestamp,
489 num_rows: fragment.num_rows(),
490 max_sequence: fragment.sequence,
491 min_sequence: fragment.min_sequence,
492 };
493
494 {
495 let mut bulk_parts = self.parts.write().unwrap();
496
497 let fragment_size = fragment.estimated_size();
498 if bulk_parts.unordered_part.should_accept(fragment.num_rows())
500 && fragment_size <= self.config.encode_bytes_threshold
501 {
502 bulk_parts.unordered_part.push(fragment);
503
504 if bulk_parts.should_compact_unordered_part(self.config.encode_bytes_threshold) {
506 bulk_parts.seal_unordered(&self.metadata)?;
507 }
508 } else {
509 bulk_parts.parts.push(BulkPartWrapper {
510 part: PartToMerge::Bulk {
511 part: fragment,
512 file_id: FileId::random(),
513 },
514 merging: false,
515 });
516 }
517
518 self.update_stats(local_metrics);
523 }
524
525 self.schedule_compact();
526
527 Ok(())
528 }
529
530 fn ranges(
531 &self,
532 projection: Option<&[ColumnId]>,
533 options: RangesOptions,
534 ) -> Result<MemtableRanges> {
535 let predicate = options.predicate;
536 let sequence = options.sequence;
537 let mut ranges = BTreeMap::new();
538 let mut range_id = 0;
539
540 let context = Arc::new(BulkIterContext::new_with_pre_filter_mode(
542 self.metadata.clone(),
543 projection,
544 predicate.predicate().cloned(),
545 options.for_flush,
546 options.pre_filter_mode,
547 options.batch_size,
548 )?);
549
550 {
552 let bulk_parts = self.parts.read().unwrap();
553
554 if !bulk_parts.unordered_part.is_empty()
556 && let Some(unordered_bulk_part) =
557 bulk_parts.unordered_part.to_bulk_part(&self.metadata)?
558 {
559 let part_stats = unordered_bulk_part.to_memtable_stats(&self.metadata);
560 let range = MemtableRange::new(
561 Arc::new(MemtableRangeContext::new(
562 self.id,
563 Box::new(BulkRangeIterBuilder {
564 part: unordered_bulk_part,
565 context: context.clone(),
566 sequence,
567 }),
568 predicate.clone(),
569 )),
570 part_stats,
571 );
572 ranges.insert(range_id, range);
573 range_id += 1;
574 }
575
576 for part_wrapper in bulk_parts.parts.iter() {
578 if part_wrapper.part.num_rows() == 0 {
580 continue;
581 }
582
583 let part_stats = part_wrapper.part.to_memtable_stats(&self.metadata);
584 let iter_builder: Box<dyn IterBuilder> = match &part_wrapper.part {
585 PartToMerge::Bulk { part, .. } => Box::new(BulkRangeIterBuilder {
586 part: part.clone(),
587 context: context.clone(),
588 sequence,
589 }),
590 PartToMerge::Multi { part, .. } => Box::new(MultiBulkRangeIterBuilder {
591 part: part.clone(),
592 context: context.clone(),
593 sequence,
594 }),
595 PartToMerge::Encoded { part, file_id } => {
596 Box::new(EncodedBulkRangeIterBuilder {
597 file_id: *file_id,
598 part: part.clone(),
599 context: context.clone(),
600 sequence,
601 })
602 }
603 };
604
605 let range = MemtableRange::new(
606 Arc::new(MemtableRangeContext::new(
607 self.id,
608 iter_builder,
609 predicate.clone(),
610 )),
611 part_stats,
612 );
613 ranges.insert(range_id, range);
614 range_id += 1;
615 }
616 }
617
618 Ok(MemtableRanges { ranges })
619 }
620
621 fn is_empty(&self) -> bool {
622 let bulk_parts = self.parts.read().unwrap();
623 bulk_parts.is_empty()
624 }
625
626 fn freeze(&self) -> Result<()> {
627 self.alloc_tracker.done_allocating();
628 Ok(())
629 }
630
631 fn stats(&self) -> MemtableStats {
632 let estimated_bytes = self.alloc_tracker.bytes_allocated();
633
634 if estimated_bytes == 0 || self.num_rows.load(Ordering::Relaxed) == 0 {
635 return MemtableStats {
636 estimated_bytes,
637 time_range: None,
638 num_rows: 0,
639 num_ranges: 0,
640 max_sequence: 0,
641 min_sequence: 0,
642 series_count: 0,
643 };
644 }
645
646 let ts_type = self
647 .metadata
648 .time_index_column()
649 .column_schema
650 .data_type
651 .clone()
652 .as_timestamp()
653 .expect("Timestamp column must have timestamp type");
654 let max_timestamp = ts_type.create_timestamp(self.max_timestamp.load(Ordering::Relaxed));
655 let min_timestamp = ts_type.create_timestamp(self.min_timestamp.load(Ordering::Relaxed));
656
657 let num_ranges = self.parts.read().unwrap().num_parts();
658
659 MemtableStats {
660 estimated_bytes,
661 time_range: Some((min_timestamp, max_timestamp)),
662 num_rows: self.num_rows.load(Ordering::Relaxed),
663 num_ranges,
664 max_sequence: self.max_sequence.load(Ordering::Relaxed),
665 min_sequence: self.min_sequence.load(Ordering::Relaxed),
666 series_count: self.estimated_series_count(),
667 }
668 }
669
670 fn min_sequence(&self) -> SequenceNumber {
671 if self.alloc_tracker.bytes_allocated() == 0 || self.num_rows.load(Ordering::Relaxed) == 0 {
672 return 0;
673 }
674 self.min_sequence.load(Ordering::Relaxed)
675 }
676
677 fn fork(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef {
678 self.fork_inner(id, metadata)
679 }
680
681 fn compact(&self, for_flush: bool) -> Result<()> {
682 let mut compactor = self.compactor.lock().unwrap();
683
684 if for_flush {
685 return Ok(());
686 }
687
688 compactor.merge_parts(
689 &self.parts,
690 &self.metadata,
691 !self.append_mode,
692 self.merge_mode,
693 )
694 }
695}
696
697impl BulkMemtable {
698 fn fork_inner(&self, id: MemtableId, metadata: &RegionMetadataRef) -> Arc<BulkMemtable> {
699 Arc::new(Self {
700 id,
701 config: self.config.clone(),
702 parts: Arc::new(RwLock::new(BulkParts::default())),
703 metadata: metadata.clone(),
704 alloc_tracker: AllocTracker::new(self.alloc_tracker.write_buffer_manager()),
705 max_timestamp: AtomicI64::new(i64::MIN),
706 min_timestamp: AtomicI64::new(i64::MAX),
707 max_sequence: AtomicU64::new(0),
708 min_sequence: AtomicU64::new(u64::MAX),
709 num_rows: AtomicUsize::new(0),
710 compactor: Arc::new(Mutex::new(MemtableCompactor::new(
711 metadata.region_id,
712 id,
713 self.config.clone(),
714 self.row_group_size,
715 self.float_field_encoding,
716 ))),
717 compact_dispatcher: self.compact_dispatcher.clone(),
718 compacting: Arc::new(AtomicBool::new(false)),
719 append_mode: self.append_mode,
720 merge_mode: self.merge_mode,
721 row_group_size: self.row_group_size,
722 float_field_encoding: self.float_field_encoding,
723 })
724 }
725
726 pub fn new(
728 id: MemtableId,
729 config: BulkMemtableConfig,
730 metadata: RegionMetadataRef,
731 write_buffer_manager: Option<WriteBufferManagerRef>,
732 compact_dispatcher: Option<Arc<CompactDispatcher>>,
733 append_mode: bool,
734 merge_mode: MergeMode,
735 ) -> Self {
736 Self::new_with_row_group_size(
737 id,
738 config,
739 metadata,
740 write_buffer_manager,
741 compact_dispatcher,
742 append_mode,
743 merge_mode,
744 DEFAULT_ROW_GROUP_SIZE,
745 )
746 }
747
748 #[allow(clippy::too_many_arguments)]
750 pub fn new_with_row_group_size(
751 id: MemtableId,
752 config: BulkMemtableConfig,
753 metadata: RegionMetadataRef,
754 write_buffer_manager: Option<WriteBufferManagerRef>,
755 compact_dispatcher: Option<Arc<CompactDispatcher>>,
756 append_mode: bool,
757 merge_mode: MergeMode,
758 row_group_size: usize,
759 ) -> Self {
760 Self::new_with_row_group_size_and_encoding(
761 id,
762 config,
763 metadata,
764 write_buffer_manager,
765 compact_dispatcher,
766 append_mode,
767 merge_mode,
768 row_group_size,
769 FloatFieldEncoding::default(),
770 )
771 }
772
773 #[allow(clippy::too_many_arguments)]
775 fn new_with_row_group_size_and_encoding(
776 id: MemtableId,
777 config: BulkMemtableConfig,
778 metadata: RegionMetadataRef,
779 write_buffer_manager: Option<WriteBufferManagerRef>,
780 compact_dispatcher: Option<Arc<CompactDispatcher>>,
781 append_mode: bool,
782 merge_mode: MergeMode,
783 row_group_size: usize,
784 float_field_encoding: FloatFieldEncoding,
785 ) -> Self {
786 let config = config.sanitize();
787 let region_id = metadata.region_id;
788 Self {
789 id,
790 config: config.clone(),
791 parts: Arc::new(RwLock::new(BulkParts::default())),
792 metadata,
793 alloc_tracker: AllocTracker::new(write_buffer_manager),
794 max_timestamp: AtomicI64::new(i64::MIN),
795 min_timestamp: AtomicI64::new(i64::MAX),
796 max_sequence: AtomicU64::new(0),
797 min_sequence: AtomicU64::new(u64::MAX),
798 num_rows: AtomicUsize::new(0),
799 compactor: Arc::new(Mutex::new(MemtableCompactor::new(
800 region_id,
801 id,
802 config,
803 row_group_size,
804 float_field_encoding,
805 ))),
806 compact_dispatcher,
807 compacting: Arc::new(AtomicBool::new(false)),
808 append_mode,
809 merge_mode,
810 row_group_size,
811 float_field_encoding,
812 }
813 }
814
815 #[cfg(test)]
817 pub fn set_unordered_part_threshold(&self, threshold: usize) {
818 self.parts
819 .write()
820 .unwrap()
821 .unordered_part
822 .set_threshold(threshold);
823 }
824
825 #[cfg(test)]
827 pub fn set_unordered_part_compact_threshold(&self, compact_threshold: usize) {
828 self.parts
829 .write()
830 .unwrap()
831 .unordered_part
832 .set_compact_threshold(compact_threshold);
833 }
834
835 fn update_stats(&self, stats: WriteMetrics) {
839 self.alloc_tracker
840 .on_allocation(stats.key_bytes + stats.value_bytes);
841
842 self.max_timestamp
843 .fetch_max(stats.max_ts, Ordering::Relaxed);
844 self.min_timestamp
845 .fetch_min(stats.min_ts, Ordering::Relaxed);
846 self.max_sequence
847 .fetch_max(stats.max_sequence, Ordering::Relaxed);
848 self.min_sequence
849 .fetch_min(stats.min_sequence, Ordering::Relaxed);
850 self.num_rows.fetch_add(stats.num_rows, Ordering::Relaxed);
851 }
852
853 fn estimated_series_count(&self) -> usize {
855 let bulk_parts = self.parts.read().unwrap();
856 bulk_parts
857 .parts
858 .iter()
859 .map(|part_wrapper| part_wrapper.part.series_count())
860 .sum()
861 }
862
863 fn schedule_compact(&self) {
865 if self.config.max_merge_groups == 0 || self.compacting.load(Ordering::Relaxed) {
866 return;
867 }
868 if !self
869 .parts
870 .read()
871 .unwrap()
872 .may_merge_parts(self.config.merge_threshold)
873 {
874 return;
875 }
876 let Some(guard) = CompactTaskGuard::try_new(self.compacting.clone()) else {
877 return;
878 };
879 if let Some(dispatcher) = &self.compact_dispatcher {
880 let task = MemCompactTask {
881 metadata: self.metadata.clone(),
882 parts: self.parts.clone(),
883 _guard: guard,
884 compactor: self.compactor.clone(),
885 append_mode: self.append_mode,
886 merge_mode: self.merge_mode,
887 };
888
889 dispatcher.dispatch_compact(task);
890 } else {
891 if let Err(e) = self.compact(false) {
893 common_telemetry::error!(e; "Failed to compact table");
894 }
895 }
896 }
897}
898
899pub struct BulkRangeIterBuilder {
901 pub part: BulkPart,
902 pub context: Arc<BulkIterContext>,
903 pub sequence: Option<SequenceRange>,
904}
905
906struct MultiBulkRangeIterBuilder {
908 part: MultiBulkPart,
909 context: Arc<BulkIterContext>,
910 sequence: Option<SequenceRange>,
911}
912
913impl IterBuilder for BulkRangeIterBuilder {
914 fn build(&self, _metrics: Option<MemScanMetrics>) -> Result<BoxedBatchIterator> {
915 UnsupportedOperationSnafu {
916 err_msg: "BatchIterator is not supported for bulk memtable",
917 }
918 .fail()
919 }
920
921 fn is_record_batch(&self) -> bool {
922 true
923 }
924
925 fn build_record_batch(
926 &self,
927 _time_range: Option<(Timestamp, Timestamp)>,
928 metrics: Option<MemScanMetrics>,
929 ) -> Result<BoxedRecordBatchIterator> {
930 let metadata = self.context.read_format().metadata();
931 if should_prune_bulk_part(&self.part.batch, &self.context, metadata) {
932 return Ok(Box::new(std::iter::empty()));
933 }
934
935 let series_count = self.part.estimated_series_count();
936 let iter = BulkPartBatchIter::from_single(
937 self.part.batch.clone(),
938 self.context.clone(),
939 self.sequence,
940 series_count,
941 metrics,
942 );
943
944 Ok(Box::new(iter))
945 }
946
947 fn record_batch_schema_hint(&self) -> Option<SchemaRef> {
948 Some(self.part.schema())
949 }
950
951 fn encoded_range(&self) -> Option<EncodedRange> {
952 None
953 }
954}
955
956impl IterBuilder for MultiBulkRangeIterBuilder {
957 fn build(&self, _metrics: Option<MemScanMetrics>) -> Result<BoxedBatchIterator> {
958 UnsupportedOperationSnafu {
959 err_msg: "BatchIterator is not supported for multi bulk memtable",
960 }
961 .fail()
962 }
963
964 fn is_record_batch(&self) -> bool {
965 true
966 }
967
968 fn build_record_batch(
969 &self,
970 _time_range: Option<(Timestamp, Timestamp)>,
971 metrics: Option<MemScanMetrics>,
972 ) -> Result<BoxedRecordBatchIterator> {
973 match self
974 .part
975 .read(self.context.clone(), self.sequence, metrics)?
976 {
977 Some(iter) => Ok(iter),
978 None => Ok(Box::new(std::iter::empty())),
980 }
981 }
982
983 fn record_batch_schema_hint(&self) -> Option<SchemaRef> {
984 self.part.schemas().next()
985 }
986
987 fn encoded_range(&self) -> Option<EncodedRange> {
988 None
989 }
990}
991
992struct EncodedBulkRangeIterBuilder {
994 file_id: FileId,
995 part: EncodedBulkPart,
996 context: Arc<BulkIterContext>,
997 sequence: Option<SequenceRange>,
998}
999
1000impl IterBuilder for EncodedBulkRangeIterBuilder {
1001 fn build(&self, _metrics: Option<MemScanMetrics>) -> Result<BoxedBatchIterator> {
1002 UnsupportedOperationSnafu {
1003 err_msg: "BatchIterator is not supported for encoded bulk memtable",
1004 }
1005 .fail()
1006 }
1007
1008 fn is_record_batch(&self) -> bool {
1009 true
1010 }
1011
1012 fn build_record_batch(
1013 &self,
1014 _time_range: Option<(Timestamp, Timestamp)>,
1015 metrics: Option<MemScanMetrics>,
1016 ) -> Result<BoxedRecordBatchIterator> {
1017 if let Some(iter) = self
1018 .part
1019 .read(self.context.clone(), self.sequence, metrics)?
1020 {
1021 Ok(iter)
1022 } else {
1023 Ok(Box::new(std::iter::empty()))
1025 }
1026 }
1027
1028 fn record_batch_schema_hint(&self) -> Option<SchemaRef> {
1029 Some(self.part.schema())
1030 }
1031
1032 fn encoded_range(&self) -> Option<EncodedRange> {
1033 Some(EncodedRange {
1034 data: self.part.data().clone(),
1035 sst_info: self.part.to_sst_info(self.file_id),
1036 })
1037 }
1038}
1039
1040struct BulkPartWrapper {
1041 part: PartToMerge,
1043 merging: bool,
1045}
1046
1047impl BulkPartWrapper {
1048 fn file_id(&self) -> FileId {
1050 self.part.file_id()
1051 }
1052}
1053
1054#[derive(Clone)]
1056enum PartToMerge {
1057 Bulk { part: BulkPart, file_id: FileId },
1059 Multi {
1061 part: MultiBulkPart,
1062 file_id: FileId,
1063 },
1064 Encoded {
1066 part: EncodedBulkPart,
1067 file_id: FileId,
1068 },
1069}
1070
1071impl PartToMerge {
1072 fn file_id(&self) -> FileId {
1074 match self {
1075 PartToMerge::Bulk { file_id, .. } => *file_id,
1076 PartToMerge::Multi { file_id, .. } => *file_id,
1077 PartToMerge::Encoded { file_id, .. } => *file_id,
1078 }
1079 }
1080
1081 fn min_timestamp(&self) -> i64 {
1083 match self {
1084 PartToMerge::Bulk { part, .. } => part.min_timestamp,
1085 PartToMerge::Multi { part, .. } => part.min_timestamp(),
1086 PartToMerge::Encoded { part, .. } => part.metadata().min_timestamp,
1087 }
1088 }
1089
1090 fn max_timestamp(&self) -> i64 {
1092 match self {
1093 PartToMerge::Bulk { part, .. } => part.max_timestamp,
1094 PartToMerge::Multi { part, .. } => part.max_timestamp(),
1095 PartToMerge::Encoded { part, .. } => part.metadata().max_timestamp,
1096 }
1097 }
1098
1099 fn num_rows(&self) -> usize {
1101 match self {
1102 PartToMerge::Bulk { part, .. } => part.num_rows(),
1103 PartToMerge::Multi { part, .. } => part.num_rows(),
1104 PartToMerge::Encoded { part, .. } => part.metadata().num_rows,
1105 }
1106 }
1107
1108 fn min_sequence(&self) -> SequenceNumber {
1110 match self {
1111 PartToMerge::Bulk { part, .. } => part.min_sequence,
1112 PartToMerge::Multi { part, .. } => part.min_sequence(),
1113 PartToMerge::Encoded { part, .. } => part.metadata().min_sequence,
1114 }
1115 }
1116
1117 fn max_sequence(&self) -> u64 {
1119 match self {
1120 PartToMerge::Bulk { part, .. } => part.sequence,
1121 PartToMerge::Multi { part, .. } => part.max_sequence(),
1122 PartToMerge::Encoded { part, .. } => part.metadata().max_sequence,
1123 }
1124 }
1125
1126 fn series_count(&self) -> usize {
1128 match self {
1129 PartToMerge::Bulk { part, .. } => part.estimated_series_count(),
1130 PartToMerge::Multi { part, .. } => part.series_count(),
1131 PartToMerge::Encoded { part, .. } => part.metadata().num_series as usize,
1132 }
1133 }
1134
1135 fn is_encoded(&self) -> bool {
1137 matches!(self, PartToMerge::Encoded { .. })
1138 }
1139
1140 fn estimated_size(&self) -> usize {
1142 match self {
1143 PartToMerge::Bulk { part, .. } => part.estimated_size(),
1144 PartToMerge::Multi { part, .. } => part.estimated_size(),
1145 PartToMerge::Encoded { part, .. } => part.size_bytes(),
1146 }
1147 }
1148
1149 fn batch_size_statistic(&self) -> Option<(u64, u64)> {
1151 match self {
1152 PartToMerge::Bulk { part, .. } => {
1153 Some((part.num_rows() as u64, part.estimated_size() as u64))
1154 }
1155 PartToMerge::Multi { part, .. } => {
1156 Some((part.num_rows() as u64, part.estimated_size() as u64))
1157 }
1158 PartToMerge::Encoded { part, .. } => part
1159 .metadata()
1160 .parquet_metadata
1161 .row_groups()
1162 .iter()
1163 .map(|row_group| {
1164 let uncompressed_bytes = row_group
1165 .columns()
1166 .iter()
1167 .map(|column| column.uncompressed_size() as u64)
1168 .sum();
1169 (row_group.num_rows() as u64, uncompressed_bytes)
1170 })
1171 .max_by_key(|(_, uncompressed_bytes)| *uncompressed_bytes),
1172 }
1173 }
1174
1175 fn to_memtable_stats(&self, region_metadata: &RegionMetadataRef) -> MemtableStats {
1177 match self {
1178 PartToMerge::Bulk { part, .. } => part.to_memtable_stats(region_metadata),
1179 PartToMerge::Multi { part, .. } => part.to_memtable_stats(region_metadata),
1180 PartToMerge::Encoded { part, .. } => part.to_memtable_stats(),
1181 }
1182 }
1183
1184 fn arrow_schema(&self) -> SchemaRef {
1186 match self {
1187 PartToMerge::Bulk { part, .. } => part.schema(),
1188 PartToMerge::Multi { part, .. } => part
1191 .schemas()
1192 .next()
1193 .expect("MultiBulkPart must contain at least one record batch"),
1194 PartToMerge::Encoded { part, .. } => part.schema(),
1195 }
1196 }
1197
1198 fn create_iterator(
1200 self,
1201 context: Arc<BulkIterContext>,
1202 plans: Arc<Json2RewritePlans>,
1203 ) -> Result<Option<BoxedRecordBatchIterator>> {
1204 let iter = match self {
1205 PartToMerge::Bulk { part, .. } => {
1206 let series_count = part.estimated_series_count();
1207 let iter = BulkPartBatchIter::from_single(
1208 part.batch,
1209 context,
1210 None, series_count,
1212 None, );
1214 Some(Box::new(iter) as BoxedRecordBatchIterator)
1215 }
1216 PartToMerge::Multi { part, .. } => part.read(context, None, None)?,
1217 PartToMerge::Encoded { part, .. } => part.read(context, None, None)?,
1218 };
1219 if plans.is_empty() {
1220 Ok(iter)
1221 } else {
1222 Ok(iter.map(|x| {
1223 Box::new(x.map(move |batch| rewrite_json2_batch(batch?, &plans)))
1224 as BoxedRecordBatchIterator
1225 }))
1226 }
1227 }
1228}
1229
1230struct MemtableCompactor {
1231 region_id: RegionId,
1232 memtable_id: MemtableId,
1233 config: BulkMemtableConfig,
1235 row_group_size: usize,
1237 float_field_encoding: FloatFieldEncoding,
1238}
1239
1240impl MemtableCompactor {
1241 fn new(
1243 region_id: RegionId,
1244 memtable_id: MemtableId,
1245 config: BulkMemtableConfig,
1246 row_group_size: usize,
1247 float_field_encoding: FloatFieldEncoding,
1248 ) -> Self {
1249 Self {
1250 region_id,
1251 memtable_id,
1252 config,
1253 row_group_size,
1254 float_field_encoding,
1255 }
1256 }
1257
1258 fn merge_parts(
1260 &mut self,
1261 bulk_parts: &RwLock<BulkParts>,
1262 metadata: &RegionMetadataRef,
1263 dedup: bool,
1264 merge_mode: MergeMode,
1265 ) -> Result<()> {
1266 let start = Instant::now();
1267
1268 let collected = {
1269 let mut parts = bulk_parts.write().unwrap();
1270 parts.collect_parts_to_merge(
1271 self.config.merge_threshold,
1272 self.config.max_merge_groups,
1273 self.config.encode_bytes_threshold,
1274 )
1275 };
1276
1277 if collected.groups.is_empty() {
1278 return Ok(());
1279 }
1280
1281 let merged_file_ids: HashSet<FileId> = collected
1283 .groups
1284 .iter()
1285 .flatten()
1286 .map(|part| part.file_id())
1287 .collect();
1288 let mut guard = MergingFlagsGuard::new(bulk_parts, &merged_file_ids);
1289
1290 let num_groups = collected.groups.len();
1291 let num_parts: usize = collected.groups.iter().map(|g| g.len()).sum();
1292
1293 let encode_row_threshold = self.config.encode_row_threshold;
1294 let encode_bytes_threshold = self.config.encode_bytes_threshold;
1295 let row_group_size = self.row_group_size;
1296 let float_field_encoding = self.float_field_encoding;
1297
1298 let merged_parts = collected
1300 .groups
1301 .into_par_iter()
1302 .map(|group| {
1303 Self::merge_parts_group(
1304 group,
1305 metadata,
1306 dedup,
1307 merge_mode,
1308 encode_row_threshold,
1309 encode_bytes_threshold,
1310 row_group_size,
1311 float_field_encoding,
1312 )
1313 })
1314 .collect::<Result<Vec<Option<MergedPart>>>>()?;
1315
1316 let total_output_rows = {
1318 let mut parts = bulk_parts.write().unwrap();
1319 parts.install_merged_parts(merged_parts.into_iter().flatten(), &merged_file_ids)
1320 };
1321
1322 guard.mark_success();
1323
1324 common_telemetry::debug!(
1325 "BulkMemtable {} {} concurrent compact {} groups, {} parts, {} rows, cost: {:?}",
1326 self.region_id,
1327 self.memtable_id,
1328 num_groups,
1329 num_parts,
1330 total_output_rows,
1331 start.elapsed()
1332 );
1333
1334 Ok(())
1335 }
1336
1337 #[allow(clippy::too_many_arguments)]
1339 fn merge_parts_group(
1340 parts_to_merge: Vec<PartToMerge>,
1341 metadata: &RegionMetadataRef,
1342 dedup: bool,
1343 merge_mode: MergeMode,
1344 encode_row_threshold: usize,
1345 encode_bytes_threshold: usize,
1346 row_group_size: usize,
1347 float_field_encoding: FloatFieldEncoding,
1348 ) -> Result<Option<MergedPart>> {
1349 if parts_to_merge.is_empty() {
1350 return Ok(None);
1351 }
1352
1353 let min_timestamp = parts_to_merge
1355 .iter()
1356 .map(|p| p.min_timestamp())
1357 .min()
1358 .unwrap_or(i64::MAX);
1359 let max_timestamp = parts_to_merge
1360 .iter()
1361 .map(|p| p.max_timestamp())
1362 .max()
1363 .unwrap_or(i64::MIN);
1364 let min_sequence = parts_to_merge
1365 .iter()
1366 .map(PartToMerge::min_sequence)
1367 .min()
1368 .unwrap_or(0);
1369 let max_sequence = parts_to_merge
1370 .iter()
1371 .map(|p| p.max_sequence())
1372 .max()
1373 .unwrap_or(0);
1374
1375 let estimated_total_rows: usize = parts_to_merge.iter().map(|p| p.num_rows()).sum();
1377 let estimated_total_bytes: usize = parts_to_merge.iter().map(|p| p.estimated_size()).sum();
1378 let estimated_series_count = parts_to_merge
1379 .iter()
1380 .map(|p| p.series_count())
1381 .max()
1382 .unwrap_or(0);
1383
1384 let batch_size = crate::batch_size::estimate_batch_size(
1385 parts_to_merge
1386 .iter()
1387 .filter_map(PartToMerge::batch_size_statistic),
1388 );
1389 let context = Arc::new(BulkIterContext::new(
1390 metadata.clone(),
1391 None, None, true,
1394 batch_size,
1395 )?);
1396
1397 let schemas = parts_to_merge
1398 .iter()
1399 .map(|part| (part.arrow_schema(), part.num_rows() as u64))
1400 .collect::<Vec<_>>();
1401 let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?);
1402
1403 debug_assert!(parts_to_merge.windows(2).all(|w| rewrite_json2_schema(
1404 &w[0].arrow_schema(),
1405 &plans
1406 ) == rewrite_json2_schema(
1407 &w[1].arrow_schema(),
1408 &plans
1409 )));
1410 let schema = rewrite_json2_schema(&parts_to_merge[0].arrow_schema(), &plans);
1414
1415 let iterators: Vec<BoxedRecordBatchIterator> = parts_to_merge
1416 .into_iter()
1417 .map(|part| part.create_iterator(context.clone(), plans.clone()))
1418 .collect::<Result<Vec<_>>>()?
1419 .into_iter()
1420 .flatten()
1421 .collect();
1422
1423 if iterators.is_empty() {
1424 return Ok(None);
1425 }
1426
1427 let merged_iter = FlatMergeIterator::new(schema.clone(), iterators, batch_size)?;
1428
1429 let boxed_iter: BoxedRecordBatchIterator = if dedup {
1430 match merge_mode {
1431 MergeMode::LastRow => {
1432 let dedup_iter = FlatDedupIterator::new(merged_iter, FlatLastRow::new(false));
1433 Box::new(dedup_iter)
1434 }
1435 MergeMode::LastNonNull => {
1436 let field_column_start = field_column_start(metadata, schema.fields().len());
1437
1438 let dedup_iter = FlatDedupIterator::new(
1439 merged_iter,
1440 FlatLastNonNull::new(field_column_start, false),
1441 );
1442 Box::new(dedup_iter)
1443 }
1444 }
1445 } else {
1446 Box::new(merged_iter)
1447 };
1448
1449 if estimated_total_rows > encode_row_threshold
1451 || estimated_total_bytes > encode_bytes_threshold
1452 {
1453 let encoder = BulkPartEncoder::new_with_float_field_encoding(
1454 metadata.clone(),
1455 row_group_size,
1456 float_field_encoding,
1457 )?;
1458 let mut metrics = BulkPartEncodeMetrics::default();
1459 let encoded_part = encoder.encode_record_batch_iter(
1460 boxed_iter,
1461 schema,
1462 min_timestamp,
1463 max_timestamp,
1464 min_sequence..=max_sequence,
1465 &mut metrics,
1466 )?;
1467
1468 common_telemetry::trace!("merge_parts_group metrics: {:?}", metrics);
1469
1470 Ok(encoded_part.map(MergedPart::Encoded))
1471 } else {
1472 let mut batches = Vec::new();
1474 let mut actual_total_rows = 0;
1475
1476 for batch_result in boxed_iter {
1477 let batch = batch_result?;
1478 actual_total_rows += batch.num_rows();
1479 batches.push(batch);
1480 }
1481
1482 if actual_total_rows == 0 {
1483 return Ok(None);
1484 }
1485
1486 let multi_part = MultiBulkPart::new(
1487 batches,
1488 min_timestamp,
1489 max_timestamp,
1490 min_sequence,
1491 max_sequence,
1492 estimated_series_count,
1493 metadata,
1494 );
1495
1496 common_telemetry::trace!(
1497 "merge_parts_group created MultiBulkPart: rows={}, batches={}",
1498 actual_total_rows,
1499 multi_part.num_batches()
1500 );
1501
1502 Ok(Some(MergedPart::Multi(multi_part)))
1503 }
1504 }
1505}
1506
1507struct CompactTaskGuard(Arc<AtomicBool>);
1510
1511impl CompactTaskGuard {
1512 fn try_new(compacting: Arc<AtomicBool>) -> Option<Self> {
1513 compacting
1515 .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
1516 .ok()
1517 .map(|_| Self(compacting))
1518 }
1519}
1520
1521impl Drop for CompactTaskGuard {
1522 fn drop(&mut self) {
1523 self.0.store(false, Ordering::Relaxed);
1524 }
1525}
1526
1527struct MemCompactTask {
1529 metadata: RegionMetadataRef,
1530 parts: Arc<RwLock<BulkParts>>,
1531 _guard: CompactTaskGuard,
1532 compactor: Arc<Mutex<MemtableCompactor>>,
1534 append_mode: bool,
1536 merge_mode: MergeMode,
1538}
1539
1540impl MemCompactTask {
1541 fn compact(&self) -> Result<()> {
1542 self.compactor.lock().unwrap().merge_parts(
1543 &self.parts,
1544 &self.metadata,
1545 !self.append_mode,
1546 self.merge_mode,
1547 )
1548 }
1549}
1550
1551#[derive(Debug)]
1553pub struct CompactDispatcher {
1554 semaphore: Arc<Semaphore>,
1555}
1556
1557impl CompactDispatcher {
1558 pub fn new(permits: usize) -> Self {
1560 Self {
1561 semaphore: Arc::new(Semaphore::new(permits)),
1562 }
1563 }
1564
1565 fn dispatch_compact(&self, task: MemCompactTask) {
1567 let semaphore = self.semaphore.clone();
1568 common_runtime::spawn_global(async move {
1569 let Ok(permit) = semaphore.acquire_owned().await else {
1570 return;
1571 };
1572
1573 common_runtime::spawn_blocking_global(move || {
1574 let _permit = permit;
1578 if let Err(e) = task.compact() {
1579 common_telemetry::error!(e; "Failed to compact memtable, region: {}", task.metadata.region_id);
1580 }
1581 });
1582 });
1583 }
1584}
1585
1586#[derive(Debug)]
1588pub struct BulkMemtableBuilder {
1589 config: BulkMemtableConfig,
1591 write_buffer_manager: Option<WriteBufferManagerRef>,
1592 compact_dispatcher: Option<Arc<CompactDispatcher>>,
1593 append_mode: bool,
1594 merge_mode: MergeMode,
1595 row_group_size: usize,
1597 float_field_encoding: FloatFieldEncoding,
1598}
1599
1600impl Default for BulkMemtableBuilder {
1601 fn default() -> Self {
1602 Self {
1603 config: BulkMemtableConfig::default(),
1604 write_buffer_manager: None,
1605 compact_dispatcher: None,
1606 append_mode: false,
1607 merge_mode: MergeMode::default(),
1608 row_group_size: DEFAULT_ROW_GROUP_SIZE,
1609 float_field_encoding: FloatFieldEncoding::default(),
1610 }
1611 }
1612}
1613
1614impl BulkMemtableBuilder {
1615 pub fn new(
1617 write_buffer_manager: Option<WriteBufferManagerRef>,
1618 append_mode: bool,
1619 merge_mode: MergeMode,
1620 ) -> Self {
1621 Self {
1622 write_buffer_manager,
1623 append_mode,
1624 merge_mode,
1625 ..Default::default()
1626 }
1627 }
1628
1629 pub fn with_config(mut self, config: BulkMemtableConfig) -> Self {
1631 self.config = config;
1632 self
1633 }
1634
1635 pub fn with_row_group_size(mut self, row_group_size: usize) -> Self {
1637 self.row_group_size = row_group_size;
1638 self
1639 }
1640
1641 pub(super) fn with_float_field_encoding(mut self, encoding: FloatFieldEncoding) -> Self {
1642 self.float_field_encoding = encoding;
1643 self
1644 }
1645
1646 pub fn with_compact_dispatcher(mut self, compact_dispatcher: Arc<CompactDispatcher>) -> Self {
1648 self.compact_dispatcher = Some(compact_dispatcher);
1649 self
1650 }
1651
1652 #[cfg(test)]
1653 pub(crate) fn config(&self) -> &BulkMemtableConfig {
1654 &self.config
1655 }
1656}
1657
1658impl MemtableBuilder for BulkMemtableBuilder {
1659 fn build(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef {
1660 Arc::new(BulkMemtable::new_with_row_group_size_and_encoding(
1661 id,
1662 self.config.clone(),
1663 metadata.clone(),
1664 self.write_buffer_manager.clone(),
1665 self.compact_dispatcher.clone(),
1666 self.append_mode,
1667 self.merge_mode,
1668 self.row_group_size,
1669 self.float_field_encoding,
1670 ))
1671 }
1672
1673 fn use_bulk_insert(&self, _metadata: &RegionMetadataRef) -> bool {
1674 true
1675 }
1676}
1677
1678#[cfg(test)]
1679mod tests {
1680 use api::helper::encode_json_value;
1681 use api::v1::value::ValueData;
1682 use api::v1::{Mutation, Row, Rows, SemanticType};
1683 use common_error::ext::WhateverResult;
1684 use datatypes::arrow::array::AsArray;
1685 use datatypes::arrow::datatypes::{DataType as ArrowDataType, UInt64Type};
1686 use datatypes::data_type::ConcreteDataType;
1687 use datatypes::extension::json::{
1688 JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, Json2PhysicalLayout, JsonMetadata,
1689 };
1690 use datatypes::json::value::JsonValue;
1691 use datatypes::json::{JsonSettings, JsonTypeHint};
1692 use datatypes::schema::ColumnSchema;
1693 use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
1694 use mito_codec::row_converter::build_primary_key_codec;
1695 use parquet::basic::{Compression, Encoding, Type as PhysicalType};
1696 use serde_json::json;
1697 use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef};
1698 use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME};
1699
1700 use super::*;
1701 use crate::memtable::bulk::part::BulkPartConverter;
1702 use crate::read::scan_region::PredicateGroup;
1703 use crate::sst::parquet::flat_format::sequence_column_index;
1704 use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema};
1705 use crate::test_util::memtable_util::{
1706 build_key_values_with_ts_seq_values, metadata_for_test, region_metadata_to_row_schema,
1707 };
1708
1709 fn create_bulk_part_with_converter(
1710 k0: &str,
1711 k1: u32,
1712 timestamps: Vec<i64>,
1713 values: Vec<Option<f64>>,
1714 sequence: u64,
1715 ) -> Result<BulkPart> {
1716 let metadata = metadata_for_test();
1717 let capacity = 100;
1718 let primary_key_codec = build_primary_key_codec(&metadata);
1719 let schema = to_flat_sst_arrow_schema(
1720 &metadata,
1721 &FlatSchemaOptions::from_encoding(metadata.primary_key_encoding),
1722 );
1723
1724 let mut converter =
1725 BulkPartConverter::new(&metadata, schema, capacity, primary_key_codec, true);
1726
1727 let key_values = build_key_values_with_ts_seq_values(
1728 &metadata,
1729 k0.to_string(),
1730 k1,
1731 timestamps.into_iter(),
1732 values.into_iter(),
1733 sequence,
1734 );
1735
1736 converter.append_key_values(&key_values)?;
1737 converter.convert()
1738 }
1739
1740 #[test]
1741 fn test_adaptive_encode_bytes_threshold() {
1742 assert_eq!(
1744 DEFAULT_ENCODE_BYTES_THRESHOLD,
1745 adaptive_encode_bytes_threshold(1024 * 1024 * 1024) );
1747 assert_eq!(
1749 256 * 1024 * 1024,
1750 adaptive_encode_bytes_threshold(8 * 1024 * 1024 * 1024) );
1752 assert_eq!(
1754 MAX_ENCODE_BYTES_THRESHOLD,
1755 adaptive_encode_bytes_threshold(64 * 1024 * 1024 * 1024) );
1757 }
1758
1759 #[test]
1760 fn test_bulk_memtable_sanitizes_zero_merge_threshold() {
1761 let metadata = metadata_for_test();
1762 let config = BulkMemtableConfig {
1763 merge_threshold: 0,
1764 ..Default::default()
1765 };
1766
1767 let memtable =
1768 BulkMemtable::new(999, config, metadata, None, None, false, MergeMode::LastRow);
1769
1770 assert_eq!(DEFAULT_MERGE_THRESHOLD, memtable.config.merge_threshold);
1771 assert_eq!(
1772 DEFAULT_MERGE_THRESHOLD,
1773 memtable.compactor.lock().unwrap().config.merge_threshold
1774 );
1775 }
1776
1777 #[test]
1778 fn test_min_sequence_survives_mixed_parts_compaction_and_fork() {
1779 let metadata = metadata_for_test();
1780 let config = BulkMemtableConfig {
1781 merge_threshold: 3,
1782 encode_row_threshold: 1,
1783 encode_bytes_threshold: 1,
1784 ..Default::default()
1785 };
1786 let memtable = BulkMemtable::new(
1787 1,
1788 config,
1789 metadata.clone(),
1790 None,
1791 None,
1792 false,
1793 MergeMode::LastNonNull,
1794 );
1795 memtable.set_unordered_part_threshold(0);
1796 for (sequence, expected_min, expected_ranges) in [(100, 100, 1), (10, 10, 2), (200, 10, 1)]
1797 {
1798 let part = create_bulk_part_with_converter(
1799 "a",
1800 1,
1801 vec![2000, 1000],
1802 vec![Some(1.0), Some(2.0)],
1803 sequence,
1804 )
1805 .unwrap();
1806 assert_eq!(sequence + 1, part.sequence);
1809 assert_eq!(sequence, part.min_sequence);
1810 let sequences = part
1811 .batch
1812 .column(sequence_column_index(part.batch.num_columns()))
1813 .as_primitive::<UInt64Type>();
1814 assert_eq!(&[sequence + 1, sequence], sequences.values().as_ref());
1815 memtable.write_bulk(part).unwrap();
1816 assert_eq!(expected_min, memtable.stats().min_sequence);
1817 assert_eq!(expected_min, memtable.min_sequence());
1818 assert_eq!(expected_ranges, memtable.stats().num_ranges);
1819 assert!(!memtable.compacting.load(Ordering::Relaxed));
1820 }
1821 assert_eq!(10, memtable.stats().min_sequence);
1824 assert_eq!(201, memtable.stats().max_sequence);
1825 let ranges = memtable.ranges(None, RangesOptions::default()).unwrap();
1826 let merged = ranges.ranges.values().next().unwrap();
1827 assert!(merged.encoded().is_some());
1828 assert_eq!(10, merged.stats().min_sequence);
1829 let fork = memtable.fork(2, &metadata);
1830 assert!(fork.is_empty());
1831 assert_eq!(0, fork.min_sequence());
1832 fork.write_bulk(
1833 create_bulk_part_with_converter("a", 1, vec![1000], vec![Some(1.0)], 50).unwrap(),
1834 )
1835 .unwrap();
1836 assert_eq!(50, fork.stats().min_sequence);
1837 assert_eq!(50, fork.min_sequence());
1838 }
1839
1840 #[test]
1841 fn test_bulk_memtable_write_read() {
1842 let metadata = metadata_for_test();
1843 let memtable = BulkMemtable::new(
1844 999,
1845 BulkMemtableConfig::default(),
1846 metadata.clone(),
1847 None,
1848 None,
1849 false,
1850 MergeMode::LastRow,
1851 );
1852 memtable.set_unordered_part_threshold(0);
1854
1855 let test_data = [
1856 (
1857 "key_a",
1858 1u32,
1859 vec![1000i64, 2000i64],
1860 vec![Some(10.5), Some(20.5)],
1861 100u64,
1862 ),
1863 (
1864 "key_b",
1865 2u32,
1866 vec![1500i64, 2500i64],
1867 vec![Some(15.5), Some(25.5)],
1868 200u64,
1869 ),
1870 ("key_c", 3u32, vec![3000i64], vec![Some(30.5)], 300u64),
1871 ];
1872
1873 for (k0, k1, timestamps, values, seq) in test_data.iter() {
1874 let part =
1875 create_bulk_part_with_converter(k0, *k1, timestamps.clone(), values.clone(), *seq)
1876 .unwrap();
1877 memtable.write_bulk(part).unwrap();
1878 }
1879
1880 let stats = memtable.stats();
1881 assert_eq!(5, stats.num_rows);
1882 assert_eq!(3, stats.num_ranges);
1883 assert_eq!(300, stats.max_sequence);
1884 assert_eq!(100, stats.min_sequence);
1885
1886 let (min_ts, max_ts) = stats.time_range.unwrap();
1887 assert_eq!(1000, min_ts.value());
1888 assert_eq!(3000, max_ts.value());
1889
1890 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
1891 let ranges = memtable
1892 .ranges(
1893 None,
1894 RangesOptions::default().with_predicate(predicate_group),
1895 )
1896 .unwrap();
1897
1898 assert_eq!(3, ranges.ranges.len());
1899 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
1900 assert_eq!(5, total_rows);
1901
1902 for range in ranges.ranges.values() {
1903 assert!(range.num_rows() > 0);
1904 assert!(range.is_record_batch());
1905
1906 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
1907
1908 let mut total_rows = 0;
1909 for batch_result in record_batch_iter {
1910 let batch = batch_result.unwrap();
1911 total_rows += batch.num_rows();
1912 assert!(batch.num_rows() > 0);
1913 assert_eq!(8, batch.num_columns());
1914 }
1915 assert_eq!(total_rows, range.num_rows());
1916 }
1917 }
1918
1919 #[test]
1920 fn test_bulk_memtable_compact_parts_with_json2() {
1921 let metadata = mock_metadata_with_json2(JsonSettings::default());
1922
1923 let config = BulkMemtableConfig {
1924 merge_threshold: 2,
1925 encode_row_threshold: 1,
1926 encode_bytes_threshold: 1,
1927 ..Default::default()
1928 };
1929 let memtable = BulkMemtable::new(
1930 999,
1931 config,
1932 metadata.clone(),
1933 None,
1934 None,
1935 true,
1936 MergeMode::LastRow,
1937 );
1938 memtable.set_unordered_part_threshold(0);
1939
1940 let part1 = mock_bulk_part_with_json2(&metadata, vec![1000, 2000], 100).unwrap();
1941 let part2 = mock_bulk_part_with_json2(&metadata, vec![3000, 4000], 200).unwrap();
1942
1943 memtable.write_bulk(part1).unwrap();
1944 memtable.write_bulk(part2).unwrap();
1945 memtable.compact(false).unwrap();
1946
1947 let stats = memtable.stats();
1948 assert_eq!(4, stats.num_rows);
1949 assert_eq!(201, stats.max_sequence);
1950
1951 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
1952 let opts = RangesOptions::default().with_predicate(predicate_group);
1953 let ranges = memtable.ranges(None, opts).unwrap();
1954
1955 assert_eq!(1, ranges.ranges.len());
1956 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
1957 assert_eq!(4, total_rows);
1958 }
1959
1960 #[test]
1961 fn test_bulk_memtable_merge_bounds_json2_paths() -> WhateverResult<()> {
1962 let metadata = mock_metadata_with_json2(JsonSettings::try_new(vec![], Some(1))?);
1963 let first = mock_bulk_part_with_json2_values(
1964 &metadata,
1965 vec![1000, 2000],
1966 vec![json!({"a": 1}), json!({"a": 2})],
1967 100,
1968 )?;
1969 let second = mock_bulk_part_with_json2_values(
1970 &metadata,
1971 vec![3000, 4000],
1972 vec![json!({"a": 3, "b": 3}), json!({"a": 4, "b": 4})],
1973 200,
1974 )?;
1975 let parts = vec![
1976 PartToMerge::Bulk {
1977 part: first,
1978 file_id: FileId::random(),
1979 },
1980 PartToMerge::Bulk {
1981 part: second,
1982 file_id: FileId::random(),
1983 },
1984 ];
1985
1986 let merged = MemtableCompactor::merge_parts_group(
1987 parts,
1988 &metadata,
1989 false,
1990 MergeMode::LastRow,
1991 usize::MAX,
1992 usize::MAX,
1993 DEFAULT_ROW_GROUP_SIZE,
1994 FloatFieldEncoding::default(),
1995 )?
1996 .unwrap();
1997 let MergedPart::Multi(part) = merged else {
1998 unreachable!()
1999 };
2000 let schema = part.schemas().next().unwrap();
2001 let ArrowDataType::Struct(fields) = schema.field(0).data_type() else {
2002 unreachable!()
2003 };
2004 assert_eq!(
2005 vec![JSON2_REMAINDER_FIELD_NAME, "a"],
2006 fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
2007 );
2008 Ok(())
2009 }
2010
2011 #[test]
2012 fn test_unordered_parts_align_json2_layouts() -> WhateverResult<()> {
2013 let metadata = mock_metadata_with_json2(JsonSettings::try_new(vec![], Some(2))?);
2014 let memtable = BulkMemtable::new(
2015 42,
2016 BulkMemtableConfig::default(),
2017 metadata.clone(),
2018 None,
2019 None,
2020 false,
2021 MergeMode::LastRow,
2022 );
2023 memtable.write_bulk(mock_bulk_part_with_json2_values(
2024 &metadata,
2025 vec![1000, 2000],
2026 vec![json!({"a": 1}), json!({"a": 2})],
2027 100,
2028 )?)?;
2029 memtable.write_bulk(mock_bulk_part_with_json2_values(
2030 &metadata,
2031 vec![3000, 4000],
2032 vec![json!({"b": 3}), json!({"b": 4})],
2033 200,
2034 )?)?;
2035
2036 let predicate = PredicateGroup::new(&metadata, &[])?;
2037 let ranges = memtable.ranges(None, RangesOptions::default().with_predicate(predicate))?;
2038 let range = ranges.ranges.values().next().unwrap();
2039 let batch = range.build_record_batch_iter(None, None)?.next().unwrap()?;
2040 let schema = batch.schema();
2041 let ArrowDataType::Struct(fields) = schema.field(0).data_type() else {
2042 unreachable!()
2043 };
2044 assert_eq!(
2045 vec![JSON2_REMAINDER_FIELD_NAME],
2046 fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
2047 );
2048 assert_eq!(4, batch.num_rows());
2050 let json = datatypes::vectors::json::array::JsonArray::from(batch.column(0));
2051 let restored = json.project_to_v2(schema.field(0), &ArrowDataType::Binary)?;
2052 let restored = datatypes::vectors::json::array::JsonArray::from(&restored);
2053 assert_eq!(json!({"a": 1}), restored.try_get_value(0)?);
2054 assert_eq!(json!({"a": 2}), restored.try_get_value(1)?);
2055 assert_eq!(json!({"b": 3}), restored.try_get_value(2)?);
2056 assert_eq!(json!({"b": 4}), restored.try_get_value(3)?);
2057 Ok(())
2058 }
2059
2060 #[test]
2061 fn test_bulk_part_converter_uses_json2_v2_layout() -> WhateverResult<()> {
2062 let settings = JsonSettings::try_new(
2063 vec![JsonTypeHint {
2064 path: vec!["id".to_string()],
2065 data_type: ConcreteDataType::int64_datatype(),
2066 inverted_index: false,
2067 }],
2068 Some(0),
2069 )?;
2070 let metadata = mock_metadata_with_json2(settings);
2071 let part = mock_bulk_part_with_json2(&metadata, vec![1000, 2000], 100)?;
2072 let schema = part.batch.schema();
2073 let field = schema.field(0);
2074 let layout = Json2PhysicalLayout::try_from_root(field)?;
2075
2076 assert!(layout.is_version_2());
2077 let ArrowDataType::Struct(fields) = field.data_type() else {
2078 unreachable!()
2079 };
2080 assert_eq!(
2081 vec![JSON2_REMAINDER_FIELD_NAME, "id"],
2082 fields.iter().map(|x| x.name().as_str()).collect::<Vec<_>>()
2083 );
2084 Ok(())
2085 }
2086
2087 fn mock_metadata_with_json2(settings: JsonSettings) -> RegionMetadataRef {
2088 let col_meta_1 = ColumnMetadata {
2089 column_schema: ColumnSchema::new(
2090 "ts",
2091 ConcreteDataType::timestamp_millisecond_datatype(),
2092 false,
2093 ),
2094 semantic_type: SemanticType::Timestamp,
2095 column_id: 0,
2096 };
2097
2098 let data_type = ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new()));
2099 let mut col_schema = ColumnSchema::new("data", data_type, true);
2100 col_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new(
2101 settings,
2102 ))));
2103
2104 let col_meta_2 = ColumnMetadata {
2105 column_schema: col_schema,
2106 semantic_type: SemanticType::Field,
2107 column_id: 1,
2108 };
2109 let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 789));
2110 builder
2111 .push_column_metadata(col_meta_1)
2112 .push_column_metadata(col_meta_2)
2113 .primary_key(vec![]);
2114 Arc::new(builder.build().unwrap())
2115 }
2116
2117 fn mock_bulk_part_with_json2(
2118 metadata: &RegionMetadataRef,
2119 timestamps: Vec<i64>,
2120 sequence: u64,
2121 ) -> Result<BulkPart> {
2122 let values = timestamps
2123 .iter()
2124 .map(|ts| {
2125 json!({
2126 "id": ts,
2127 "payload": {
2128 "message": format!("row-{ts}"),
2129 },
2130 })
2131 })
2132 .collect();
2133 mock_bulk_part_with_json2_values(metadata, timestamps, values, sequence)
2134 }
2135
2136 fn mock_bulk_part_with_json2_values(
2137 metadata: &RegionMetadataRef,
2138 timestamps: Vec<i64>,
2139 values: Vec<serde_json::Value>,
2140 sequence: u64,
2141 ) -> Result<BulkPart> {
2142 let capacity = timestamps.len();
2143 debug_assert_eq!(capacity, values.len());
2144 let primary_key_codec = build_primary_key_codec(metadata);
2145 let json_type = JsonNativeType::Object(JsonObjectType::from([
2146 ("id".to_string(), JsonNativeType::i64()),
2147 (
2148 "payload".to_string(),
2149 JsonNativeType::Object(JsonObjectType::from([(
2150 "message".to_string(),
2151 JsonNativeType::String,
2152 )])),
2153 ),
2154 ]));
2155 let mut options = FlatSchemaOptions::from_encoding(metadata.primary_key_encoding);
2156 options
2157 .concretized_json_types
2158 .insert("data".to_string(), json_type.as_arrow_type());
2159 let schema = to_flat_sst_arrow_schema(metadata, &options);
2160
2161 let mut converter =
2162 BulkPartConverter::new(metadata, schema, capacity, primary_key_codec, true);
2163
2164 let rows = timestamps
2165 .into_iter()
2166 .zip(values)
2167 .map(|(ts, value)| {
2168 let val1 = api::v1::Value {
2169 value_data: Some(ValueData::TimestampMillisecondValue(ts)),
2170 };
2171 let value_data = ValueData::JsonValue(encode_json_value(JsonValue::from(value)));
2172 let val2 = api::v1::Value {
2173 value_data: Some(value_data),
2174 };
2175 Row {
2176 values: vec![val1, val2],
2177 }
2178 })
2179 .collect();
2180
2181 let mutation = Mutation {
2182 op_type: 1,
2183 sequence,
2184 rows: Some(Rows {
2185 schema: region_metadata_to_row_schema(metadata),
2186 rows,
2187 }),
2188 write_hint: None,
2189 };
2190 let key_values = KeyValues::new(metadata.as_ref(), mutation).unwrap();
2191
2192 converter.append_key_values(&key_values)?;
2193 converter.convert()
2194 }
2195
2196 #[test]
2197 fn test_bulk_memtable_ranges_with_projection() {
2198 let metadata = metadata_for_test();
2199 let memtable = BulkMemtable::new(
2200 111,
2201 BulkMemtableConfig::default(),
2202 metadata.clone(),
2203 None,
2204 None,
2205 false,
2206 MergeMode::LastRow,
2207 );
2208
2209 let bulk_part = create_bulk_part_with_converter(
2210 "projection_test",
2211 5,
2212 vec![5000, 6000, 7000],
2213 vec![Some(50.0), Some(60.0), Some(70.0)],
2214 500,
2215 )
2216 .unwrap();
2217
2218 memtable.write_bulk(bulk_part).unwrap();
2219
2220 let projection = vec![4u32];
2221 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2222 let ranges = memtable
2223 .ranges(
2224 Some(&projection),
2225 RangesOptions::default().with_predicate(predicate_group),
2226 )
2227 .unwrap();
2228
2229 assert_eq!(1, ranges.ranges.len());
2230 let range = ranges.ranges.get(&0).unwrap();
2231
2232 assert!(range.is_record_batch());
2233 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
2234
2235 let mut total_rows = 0;
2236 for batch_result in record_batch_iter {
2237 let batch = batch_result.unwrap();
2238 assert!(batch.num_rows() > 0);
2239 assert_eq!(5, batch.num_columns());
2240 total_rows += batch.num_rows();
2241 }
2242 assert_eq!(3, total_rows);
2243 }
2244
2245 #[test]
2246 fn test_bulk_memtable_unsupported_operations() {
2247 let metadata = metadata_for_test();
2248 let memtable = BulkMemtable::new(
2249 111,
2250 BulkMemtableConfig::default(),
2251 metadata.clone(),
2252 None,
2253 None,
2254 false,
2255 MergeMode::LastRow,
2256 );
2257
2258 let key_values = build_key_values_with_ts_seq_values(
2259 &metadata,
2260 "test".to_string(),
2261 1,
2262 vec![1000].into_iter(),
2263 vec![Some(1.0)].into_iter(),
2264 1,
2265 );
2266
2267 let err = memtable.write(&key_values).unwrap_err();
2268 assert!(err.to_string().contains("not supported"));
2269
2270 let kv = key_values.iter().next().unwrap();
2271 let err = memtable.write_one(kv).unwrap_err();
2272 assert!(err.to_string().contains("not supported"));
2273 }
2274
2275 #[test]
2276 fn test_bulk_memtable_freeze() {
2277 let metadata = metadata_for_test();
2278 let memtable = BulkMemtable::new(
2279 222,
2280 BulkMemtableConfig::default(),
2281 metadata.clone(),
2282 None,
2283 None,
2284 false,
2285 MergeMode::LastRow,
2286 );
2287
2288 let bulk_part = create_bulk_part_with_converter(
2289 "freeze_test",
2290 10,
2291 vec![10000],
2292 vec![Some(100.0)],
2293 1000,
2294 )
2295 .unwrap();
2296
2297 memtable.write_bulk(bulk_part).unwrap();
2298 memtable.freeze().unwrap();
2299
2300 let stats_after_freeze = memtable.stats();
2301 assert_eq!(1, stats_after_freeze.num_rows);
2302 }
2303
2304 #[test]
2305 fn test_bulk_memtable_fork() {
2306 let metadata = metadata_for_test();
2307 let original_memtable = BulkMemtable::new_with_row_group_size_and_encoding(
2308 333,
2309 BulkMemtableConfig::default(),
2310 metadata.clone(),
2311 None,
2312 None,
2313 false,
2314 MergeMode::LastRow,
2315 DEFAULT_ROW_GROUP_SIZE,
2316 FloatFieldEncoding::ByteStreamSplit,
2317 );
2318
2319 let bulk_part =
2320 create_bulk_part_with_converter("fork_test", 15, vec![15000], vec![Some(150.0)], 1500)
2321 .unwrap();
2322
2323 original_memtable.write_bulk(bulk_part).unwrap();
2324
2325 let forked_memtable = original_memtable.fork_inner(444, &metadata);
2326
2327 assert_eq!(forked_memtable.id(), 444);
2328 assert_eq!(
2329 FloatFieldEncoding::ByteStreamSplit,
2330 original_memtable.float_field_encoding
2331 );
2332 assert_eq!(
2333 original_memtable.float_field_encoding,
2334 forked_memtable.float_field_encoding
2335 );
2336 assert_eq!(
2337 original_memtable.float_field_encoding,
2338 forked_memtable
2339 .compactor
2340 .lock()
2341 .unwrap()
2342 .float_field_encoding
2343 );
2344 assert!(forked_memtable.is_empty());
2345 assert_eq!(0, forked_memtable.stats().num_rows);
2346
2347 assert!(!original_memtable.is_empty());
2348 assert_eq!(1, original_memtable.stats().num_rows);
2349 }
2350
2351 #[test]
2352 fn test_bulk_memtable_ranges_multiple_parts() {
2353 let metadata = metadata_for_test();
2354 let memtable = BulkMemtable::new(
2355 777,
2356 BulkMemtableConfig::default(),
2357 metadata.clone(),
2358 None,
2359 None,
2360 false,
2361 MergeMode::LastRow,
2362 );
2363 memtable.set_unordered_part_threshold(0);
2365
2366 let parts_data = vec![
2367 (
2368 "part1",
2369 1u32,
2370 vec![1000i64, 1100i64],
2371 vec![Some(10.0), Some(11.0)],
2372 100u64,
2373 ),
2374 (
2375 "part2",
2376 2u32,
2377 vec![2000i64, 2100i64],
2378 vec![Some(20.0), Some(21.0)],
2379 200u64,
2380 ),
2381 ("part3", 3u32, vec![3000i64], vec![Some(30.0)], 300u64),
2382 ];
2383
2384 for (k0, k1, timestamps, values, seq) in parts_data {
2385 let part = create_bulk_part_with_converter(k0, k1, timestamps, values, seq).unwrap();
2386 memtable.write_bulk(part).unwrap();
2387 }
2388
2389 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2390 let ranges = memtable
2391 .ranges(
2392 None,
2393 RangesOptions::default().with_predicate(predicate_group),
2394 )
2395 .unwrap();
2396
2397 assert_eq!(3, ranges.ranges.len());
2398 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
2399 assert_eq!(5, total_rows);
2400 assert_eq!(3, ranges.ranges.len());
2401
2402 for (range_id, range) in ranges.ranges.iter() {
2403 assert!(*range_id < 3);
2404 assert!(range.num_rows() > 0);
2405 assert!(range.is_record_batch());
2406 }
2407 }
2408
2409 #[test]
2410 fn test_bulk_memtable_ranges_with_sequence_filter() {
2411 let metadata = metadata_for_test();
2412 let memtable = BulkMemtable::new(
2413 888,
2414 BulkMemtableConfig::default(),
2415 metadata.clone(),
2416 None,
2417 None,
2418 false,
2419 MergeMode::LastRow,
2420 );
2421
2422 let part = create_bulk_part_with_converter(
2423 "seq_test",
2424 1,
2425 vec![1000, 2000, 3000],
2426 vec![Some(10.0), Some(20.0), Some(30.0)],
2427 500,
2428 )
2429 .unwrap();
2430
2431 memtable.write_bulk(part).unwrap();
2432
2433 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2434 let sequence_filter = Some(SequenceRange::LtEq { max: 400 }); let ranges = memtable
2436 .ranges(
2437 None,
2438 RangesOptions::default()
2439 .with_predicate(predicate_group)
2440 .with_sequence(sequence_filter),
2441 )
2442 .unwrap();
2443
2444 assert_eq!(1, ranges.ranges.len());
2445 let range = ranges.ranges.get(&0).unwrap();
2446
2447 let mut record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
2448 assert!(record_batch_iter.next().is_none());
2449 }
2450
2451 #[test]
2452 fn test_bulk_memtable_ranges_with_encoded_parts() {
2453 let metadata = metadata_for_test();
2454 let config = BulkMemtableConfig {
2455 merge_threshold: 8,
2456 ..Default::default()
2457 };
2458 let memtable = BulkMemtable::new(
2459 999,
2460 config,
2461 metadata.clone(),
2462 None,
2463 None,
2464 false,
2465 MergeMode::LastRow,
2466 );
2467 memtable.set_unordered_part_threshold(0);
2469
2470 for i in 0..10 {
2472 let part = create_bulk_part_with_converter(
2473 &format!("key_{}", i),
2474 i,
2475 vec![1000 + i as i64 * 100],
2476 vec![Some(i as f64 * 10.0)],
2477 100 + i as u64,
2478 )
2479 .unwrap();
2480 memtable.write_bulk(part).unwrap();
2481 }
2482
2483 memtable.compact(false).unwrap();
2484
2485 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2486 let ranges = memtable
2487 .ranges(
2488 None,
2489 RangesOptions::default().with_predicate(predicate_group),
2490 )
2491 .unwrap();
2492
2493 assert_eq!(3, ranges.ranges.len());
2495 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
2496 assert_eq!(10, total_rows);
2497
2498 for range in ranges.ranges.values() {
2499 assert!(range.num_rows() > 0);
2500 assert!(range.is_record_batch());
2501
2502 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
2503 let mut total_rows = 0;
2504 for batch_result in record_batch_iter {
2505 let batch = batch_result.unwrap();
2506 total_rows += batch.num_rows();
2507 assert!(batch.num_rows() > 0);
2508 }
2509 assert_eq!(total_rows, range.num_rows());
2510 }
2511 }
2512
2513 #[test]
2514 fn test_bulk_memtable_unordered_part() {
2515 let metadata = metadata_for_test();
2516 let memtable = BulkMemtable::new(
2517 1001,
2518 BulkMemtableConfig::default(),
2519 metadata.clone(),
2520 None,
2521 None,
2522 false,
2523 MergeMode::LastRow,
2524 );
2525
2526 memtable.set_unordered_part_threshold(5);
2529 memtable.set_unordered_part_compact_threshold(10);
2531
2532 for i in 0..3 {
2534 let part = create_bulk_part_with_converter(
2535 &format!("key_{}", i),
2536 i,
2537 vec![1000 + i as i64 * 100, 1100 + i as i64 * 100],
2538 vec![Some(i as f64 * 10.0), Some(i as f64 * 10.0 + 1.0)],
2539 100 + i as u64,
2540 )
2541 .unwrap();
2542 assert_eq!(2, part.num_rows());
2543 memtable.write_bulk(part).unwrap();
2544 }
2545
2546 let stats = memtable.stats();
2548 assert_eq!(6, stats.num_rows);
2549
2550 for i in 3..5 {
2553 let part = create_bulk_part_with_converter(
2554 &format!("key_{}", i),
2555 i,
2556 vec![1000 + i as i64 * 100, 1100 + i as i64 * 100],
2557 vec![Some(i as f64 * 10.0), Some(i as f64 * 10.0 + 1.0)],
2558 100 + i as u64,
2559 )
2560 .unwrap();
2561 memtable.write_bulk(part).unwrap();
2562 }
2563
2564 let stats = memtable.stats();
2566 assert_eq!(10, stats.num_rows);
2567
2568 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2570 let ranges = memtable
2571 .ranges(
2572 None,
2573 RangesOptions::default().with_predicate(predicate_group),
2574 )
2575 .unwrap();
2576
2577 assert!(!ranges.ranges.is_empty());
2579 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
2580 assert_eq!(10, total_rows);
2581
2582 let mut total_rows_read = 0;
2584 for range in ranges.ranges.values() {
2585 assert!(range.is_record_batch());
2586 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
2587
2588 for batch_result in record_batch_iter {
2589 let batch = batch_result.unwrap();
2590 total_rows_read += batch.num_rows();
2591 }
2592 }
2593 assert_eq!(10, total_rows_read);
2594 }
2595
2596 #[test]
2597 fn test_bulk_memtable_unordered_part_mixed_sizes() {
2598 let metadata = metadata_for_test();
2599 let memtable = BulkMemtable::new(
2600 1002,
2601 BulkMemtableConfig::default(),
2602 metadata.clone(),
2603 None,
2604 None,
2605 false,
2606 MergeMode::LastRow,
2607 );
2608
2609 memtable.set_unordered_part_threshold(4);
2611 memtable.set_unordered_part_compact_threshold(8);
2612
2613 for i in 0..2 {
2615 let part = create_bulk_part_with_converter(
2616 &format!("small_{}", i),
2617 i,
2618 vec![1000 + i as i64, 2000 + i as i64, 3000 + i as i64],
2619 vec![Some(i as f64), Some(i as f64 + 1.0), Some(i as f64 + 2.0)],
2620 10 + i as u64,
2621 )
2622 .unwrap();
2623 assert_eq!(3, part.num_rows());
2624 memtable.write_bulk(part).unwrap();
2625 }
2626
2627 let large_part = create_bulk_part_with_converter(
2629 "large_key",
2630 100,
2631 vec![5000, 6000, 7000, 8000, 9000],
2632 vec![
2633 Some(100.0),
2634 Some(101.0),
2635 Some(102.0),
2636 Some(103.0),
2637 Some(104.0),
2638 ],
2639 50,
2640 )
2641 .unwrap();
2642 assert_eq!(5, large_part.num_rows());
2643 memtable.write_bulk(large_part).unwrap();
2644
2645 let part = create_bulk_part_with_converter(
2647 "small_2",
2648 2,
2649 vec![4000, 4100],
2650 vec![Some(20.0), Some(21.0)],
2651 30,
2652 )
2653 .unwrap();
2654 memtable.write_bulk(part).unwrap();
2655
2656 let stats = memtable.stats();
2657 assert_eq!(13, stats.num_rows); let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2661 let ranges = memtable
2662 .ranges(
2663 None,
2664 RangesOptions::default().with_predicate(predicate_group),
2665 )
2666 .unwrap();
2667
2668 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
2669 assert_eq!(13, total_rows);
2670
2671 let mut total_rows_read = 0;
2672 for range in ranges.ranges.values() {
2673 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
2674 for batch_result in record_batch_iter {
2675 let batch = batch_result.unwrap();
2676 total_rows_read += batch.num_rows();
2677 }
2678 }
2679 assert_eq!(13, total_rows_read);
2680 }
2681
2682 #[test]
2683 fn test_bulk_memtable_unordered_part_byte_threshold() {
2684 let metadata = metadata_for_test();
2685 let first =
2686 create_bulk_part_with_converter("first", 1, vec![1000], vec![Some(1.0)], 1).unwrap();
2687 let second =
2688 create_bulk_part_with_converter("second", 2, vec![2000], vec![Some(2.0)], 2).unwrap();
2689 let threshold = first.estimated_size().max(second.estimated_size());
2690 let memtable = BulkMemtable::new(
2691 1004,
2692 BulkMemtableConfig {
2693 encode_bytes_threshold: threshold,
2694 ..Default::default()
2695 },
2696 metadata.clone(),
2697 None,
2698 None,
2699 false,
2700 MergeMode::LastRow,
2701 );
2702 memtable.set_unordered_part_compact_threshold(usize::MAX);
2703
2704 memtable.write_bulk(first).unwrap();
2706 assert_eq!(1, memtable.parts.read().unwrap().unordered_part.num_parts());
2707
2708 memtable.write_bulk(second).unwrap();
2710 let parts = memtable.parts.read().unwrap();
2711 assert!(parts.unordered_part.is_empty());
2712 assert_eq!(1, parts.parts.len());
2713 drop(parts);
2714
2715 let oversized =
2716 create_bulk_part_with_converter("oversized", 3, vec![3000], vec![Some(3.0)], 3)
2717 .unwrap();
2718 let oversized_memtable = BulkMemtable::new(
2719 1005,
2720 BulkMemtableConfig {
2721 encode_bytes_threshold: oversized.estimated_size() - 1,
2722 ..Default::default()
2723 },
2724 metadata,
2725 None,
2726 None,
2727 false,
2728 MergeMode::LastRow,
2729 );
2730 oversized_memtable.write_bulk(oversized).unwrap();
2731 let parts = oversized_memtable.parts.read().unwrap();
2732 assert!(parts.unordered_part.is_empty());
2733 assert_eq!(1, parts.parts.len());
2734 }
2735
2736 #[test]
2737 fn test_bulk_memtable_unordered_part_with_ranges() {
2738 let metadata = metadata_for_test();
2739 let memtable = BulkMemtable::new(
2740 1003,
2741 BulkMemtableConfig::default(),
2742 metadata.clone(),
2743 None,
2744 None,
2745 false,
2746 MergeMode::LastRow,
2747 );
2748
2749 memtable.set_unordered_part_threshold(3);
2751 memtable.set_unordered_part_compact_threshold(100); for i in 0..3 {
2755 let part = create_bulk_part_with_converter(
2756 &format!("key_{}", i),
2757 i,
2758 vec![1000 + i as i64 * 100],
2759 vec![Some(i as f64 * 10.0)],
2760 100 + i as u64,
2761 )
2762 .unwrap();
2763 assert_eq!(1, part.num_rows());
2764 memtable.write_bulk(part).unwrap();
2765 }
2766
2767 let stats = memtable.stats();
2768 assert_eq!(3, stats.num_rows);
2769
2770 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
2772 let ranges = memtable
2773 .ranges(
2774 None,
2775 RangesOptions::default().with_predicate(predicate_group),
2776 )
2777 .unwrap();
2778
2779 assert_eq!(1, ranges.ranges.len());
2781 let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
2782 assert_eq!(3, total_rows);
2783
2784 let range = ranges.ranges.get(&0).unwrap();
2786 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
2787
2788 let mut total_rows = 0;
2789 for batch_result in record_batch_iter {
2790 let batch = batch_result.unwrap();
2791 total_rows += batch.num_rows();
2792 assert!(batch.num_rows() > 0);
2794 }
2795 assert_eq!(3, total_rows);
2796 }
2797
2798 fn create_bulk_part_wrapper(part: BulkPart) -> BulkPartWrapper {
2800 BulkPartWrapper {
2801 part: PartToMerge::Bulk {
2802 part,
2803 file_id: FileId::random(),
2804 },
2805 merging: false,
2806 }
2807 }
2808
2809 #[test]
2810 fn test_may_merge_parts_below_threshold() {
2811 let mut bulk_parts = BulkParts::default();
2812
2813 for i in 0..DEFAULT_MERGE_THRESHOLD - 1 {
2815 let part = create_bulk_part_with_converter(
2816 &format!("key_{}", i),
2817 i as u32,
2818 vec![1000 + i as i64 * 100],
2819 vec![Some(i as f64 * 10.0)],
2820 100 + i as u64,
2821 )
2822 .unwrap();
2823 bulk_parts.parts.push(create_bulk_part_wrapper(part));
2824 if i == 0 {
2825 assert!(!bulk_parts.may_merge_parts(0));
2826 assert!(!bulk_parts.may_merge_parts(1));
2827 }
2828 }
2829
2830 assert!(!bulk_parts.may_merge_parts(DEFAULT_MERGE_THRESHOLD));
2832 }
2833
2834 #[test]
2835 fn test_may_merge_parts_at_threshold() {
2836 let mut bulk_parts = BulkParts::default();
2837 let merge_threshold = 8;
2838
2839 for i in 0..merge_threshold {
2841 let part = create_bulk_part_with_converter(
2842 &format!("key_{}", i),
2843 i as u32,
2844 vec![1000 + i as i64 * 100],
2845 vec![Some(i as f64 * 10.0)],
2846 100 + i as u64,
2847 )
2848 .unwrap();
2849 bulk_parts.parts.push(create_bulk_part_wrapper(part));
2850 }
2851
2852 assert!(bulk_parts.may_merge_parts(merge_threshold));
2854 }
2855
2856 #[test]
2857 fn test_raw_parts_do_not_use_encoded_merge_limit() {
2858 let mut bulk_parts = BulkParts::default();
2859 let merge_threshold = 4;
2860 for i in 0..merge_threshold {
2861 let part = create_bulk_part_with_converter(
2862 &format!("key_{}", i),
2863 i as u32,
2864 vec![1000 + i as i64],
2865 vec![Some(i as f64)],
2866 100 + i as u64,
2867 )
2868 .unwrap();
2869 bulk_parts.parts.push(create_bulk_part_wrapper(part));
2870 }
2871
2872 let max_size = bulk_parts.parts[0].part.estimated_size();
2873 assert!(bulk_parts.may_merge_parts(merge_threshold));
2874 assert_eq!(
2875 1,
2876 bulk_parts
2877 .collect_parts_to_merge(merge_threshold, 1, max_size)
2878 .groups
2879 .len()
2880 );
2881 }
2882
2883 #[test]
2884 fn test_may_merge_parts_with_merging_flag() {
2885 let mut bulk_parts = BulkParts::default();
2886 let merge_threshold = 8;
2887
2888 for i in 0..10 {
2890 let part = create_bulk_part_with_converter(
2891 &format!("key_{}", i),
2892 i as u32,
2893 vec![1000 + i as i64 * 100],
2894 vec![Some(i as f64 * 10.0)],
2895 100 + i as u64,
2896 )
2897 .unwrap();
2898 bulk_parts.parts.push(create_bulk_part_wrapper(part));
2899 }
2900
2901 assert!(bulk_parts.may_merge_parts(merge_threshold));
2903
2904 for wrapper in bulk_parts.parts.iter_mut().take(3) {
2906 wrapper.merging = true;
2907 }
2908
2909 assert!(!bulk_parts.may_merge_parts(merge_threshold));
2911 }
2912
2913 #[test]
2914 fn test_compact_task_guard() {
2915 let compacting = Arc::new(AtomicBool::new(false));
2916 let guard = CompactTaskGuard::try_new(compacting.clone()).unwrap();
2917 assert!(compacting.load(Ordering::Relaxed));
2918 assert!(CompactTaskGuard::try_new(compacting.clone()).is_none());
2919 drop(guard);
2920 assert!(!compacting.load(Ordering::Relaxed));
2921 assert!(CompactTaskGuard::try_new(compacting.clone()).is_some());
2922 assert!(!compacting.load(Ordering::Relaxed));
2923 }
2924
2925 #[test]
2926 fn test_compact_scheduling_disabled() {
2927 let memtable = BulkMemtable::new(
2928 0,
2929 BulkMemtableConfig {
2930 merge_threshold: 2,
2931 max_merge_groups: 0,
2932 ..Default::default()
2933 },
2934 metadata_for_test(),
2935 None,
2936 Some(Arc::new(CompactDispatcher::new(0))),
2937 false,
2938 MergeMode::LastRow,
2939 );
2940 let part =
2941 create_bulk_part_with_converter("key", 0, vec![1000], vec![Some(1.0)], 0).unwrap();
2942 memtable.parts.write().unwrap().parts = vec![
2943 create_bulk_part_wrapper(part.clone()),
2944 create_bulk_part_wrapper(part),
2945 ];
2946 memtable.schedule_compact();
2947 assert!(!memtable.compacting.load(Ordering::Relaxed));
2948 }
2949
2950 #[test]
2951 fn test_collect_parts_to_merge_sequence_components() {
2952 let make_part = |min, max| {
2953 let mut part =
2954 create_bulk_part_with_converter("key", 0, vec![1000], vec![Some(1.0)], max)
2955 .unwrap();
2956 part.min_sequence = min;
2957 part
2958 };
2959 let metadata = metadata_for_test();
2960 let encoder = BulkPartEncoder::new(metadata.clone(), DEFAULT_ROW_GROUP_SIZE).unwrap();
2961 let mut parts = BulkParts {
2963 parts: vec![
2964 create_bulk_part_wrapper(make_part(70, 70)),
2965 create_bulk_part_wrapper(make_part(10, 30)),
2966 BulkPartWrapper {
2967 part: PartToMerge::Multi {
2968 part: MultiBulkPart::from_bulk_part(make_part(35, 50), &metadata),
2969 file_id: FileId::random(),
2970 },
2971 merging: false,
2972 },
2973 create_bulk_part_wrapper(make_part(60, 60)),
2974 BulkPartWrapper {
2975 part: PartToMerge::Encoded {
2976 part: encoder.encode_part(&make_part(20, 40)).unwrap().unwrap(),
2977 file_id: FileId::random(),
2978 },
2979 merging: false,
2980 },
2981 ],
2982 ..Default::default()
2983 };
2984 assert!(parts.may_merge_parts(parts.parts.len()));
2987 assert_eq!(
2988 vec![vec![1, 4, 2], vec![3, 0]],
2989 parts.pick_merge_groups(2, 2, 1)
2990 );
2991 parts.unordered_part.push(make_part(0, 100));
2993 assert!(parts.may_merge_parts(2));
2994 assert!(parts.pick_merge_groups(2, 2, 1).is_empty());
2995 parts.unordered_part.clear();
2996
2997 parts.unordered_part.push(make_part(50, 55));
2999 assert_eq!(vec![vec![3, 0]], parts.pick_merge_groups(2, 2, 1));
3000 parts.unordered_part.clear();
3001
3002 parts.unordered_part.push(make_part(65, 65));
3004 assert_eq!(vec![vec![1, 4, 2]], parts.pick_merge_groups(2, 2, 1));
3005 parts.unordered_part.clear();
3006
3007 parts.parts[4].merging = true;
3008 let selected = parts.collect_parts_to_merge(2, 2, 1);
3009 assert_eq!(1, selected.groups.len());
3010 assert_eq!(
3011 vec![60, 70],
3012 selected.groups[0]
3013 .iter()
3014 .map(PartToMerge::max_sequence)
3015 .collect::<Vec<_>>()
3016 );
3017 assert!(!parts.parts[1].merging);
3018 assert!(!parts.parts[2].merging);
3019 }
3020
3021 #[test]
3022 fn test_encoded_merge_candidate_size_limit() {
3023 assert!(BulkParts::is_merge_candidate_by_size(false, usize::MAX, 8));
3024 assert!(BulkParts::is_merge_candidate_by_size(true, 8, 8));
3025 assert!(!BulkParts::is_merge_candidate_by_size(true, 9, 8));
3026 }
3027
3028 #[test]
3029 fn test_bulk_part_encoder_preserves_byte_stream_split() -> WhateverResult<()> {
3030 let metadata = metadata_for_test();
3031 let bulk_part = create_bulk_part_with_converter(
3032 "byte_stream_split",
3033 0,
3034 vec![1000, 2000],
3035 vec![Some(10.0), Some(20.0)],
3036 100,
3037 )?;
3038 let encoder = BulkPartEncoder::new_with_float_field_encoding(
3039 metadata,
3040 DEFAULT_ROW_GROUP_SIZE,
3041 FloatFieldEncoding::ByteStreamSplit,
3042 )?;
3043 let encoded_part = encoder.encode_part(&bulk_part)?.unwrap();
3044 let field_column = encoded_part
3045 .metadata()
3046 .parquet_metadata
3047 .row_groups()
3048 .iter()
3049 .flat_map(|row_group| row_group.columns())
3050 .find(|column| {
3051 column.column_path().string() == "v1"
3052 && column.column_type() == PhysicalType::DOUBLE
3053 })
3054 .expect("Float64 field should be present in encoded bulk part");
3055
3056 assert!(
3057 field_column
3058 .encodings()
3059 .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT)
3060 );
3061 Ok(())
3062 }
3063
3064 #[test]
3065 fn test_bulk_part_encoder_matches_sst_writer_internal_encodings() -> WhateverResult<()> {
3066 let metadata = metadata_for_test();
3067 let bulk_part = create_bulk_part_with_converter(
3068 "key",
3069 0,
3070 vec![1000, 2000, 3000],
3071 vec![Some(1.0), Some(2.0), Some(3.0)],
3072 100,
3073 )?;
3074 let encoder = BulkPartEncoder::new(metadata, DEFAULT_ROW_GROUP_SIZE)?;
3075 let encoded_part = encoder.encode_part(&bulk_part)?.unwrap();
3076 let columns = encoded_part
3077 .metadata()
3078 .parquet_metadata
3079 .row_group(0)
3080 .columns();
3081
3082 for name in ["ts", SEQUENCE_COLUMN_NAME] {
3083 let column = columns
3084 .iter()
3085 .find(|column| column.column_path().string() == name)
3086 .unwrap();
3087 assert!(
3088 column
3089 .encodings()
3090 .any(|encoding| encoding == Encoding::DELTA_BINARY_PACKED),
3091 "{name}"
3092 );
3093 assert!(column.dictionary_page_offset().is_none(), "{name}");
3094 }
3095 let op_type = columns
3096 .iter()
3097 .find(|column| column.column_path().string() == OP_TYPE_COLUMN_NAME)
3098 .unwrap();
3099 assert_eq!(Compression::UNCOMPRESSED, op_type.compression());
3100 Ok(())
3101 }
3102
3103 #[test]
3104 fn test_encoded_part_batch_size_uses_largest_uncompressed_row_group() {
3105 const NUM_ROWS: usize = 14;
3106 const ROW_GROUP_SIZE: usize = 4;
3107
3108 let metadata = metadata_for_test();
3109 let timestamps = (0..NUM_ROWS as i64).collect::<Vec<_>>();
3110 let field_values = (0..NUM_ROWS)
3111 .map(|value| Some(value as f64))
3112 .collect::<Vec<_>>();
3113 let bulk_part =
3114 create_bulk_part_with_converter("key", 0, timestamps, field_values, 0).unwrap();
3115 let encoder = BulkPartEncoder::new(metadata, ROW_GROUP_SIZE).unwrap();
3116 let encoded_part = encoder.encode_part(&bulk_part).unwrap().unwrap();
3117 let max_row_group = encoded_part
3118 .metadata()
3119 .parquet_metadata
3120 .row_groups()
3121 .iter()
3122 .max_by_key(|row_group| {
3123 row_group
3124 .columns()
3125 .iter()
3126 .map(|column| column.uncompressed_size() as u64)
3127 .sum::<u64>()
3128 })
3129 .unwrap();
3130 let max_uncompressed_size = max_row_group
3131 .columns()
3132 .iter()
3133 .map(|column| column.uncompressed_size() as u64)
3134 .sum();
3135 let expected = (max_row_group.num_rows() as u64, max_uncompressed_size);
3136 let part = PartToMerge::Encoded {
3137 part: encoded_part,
3138 file_id: FileId::random(),
3139 };
3140
3141 assert_eq!(Some(expected), part.batch_size_statistic());
3142 }
3143
3144 #[test]
3145 fn test_bulk_memtable_ranges_with_multi_bulk_part() {
3146 let metadata = metadata_for_test();
3147 let merge_threshold = 8;
3148 let config = BulkMemtableConfig {
3149 merge_threshold,
3150 ..Default::default()
3151 };
3152 let memtable = BulkMemtable::new(
3153 2005,
3154 config,
3155 metadata.clone(),
3156 None,
3157 None,
3158 false,
3159 MergeMode::LastRow,
3160 );
3161 memtable.set_unordered_part_threshold(0);
3163
3164 for i in 0..merge_threshold {
3168 let part = create_bulk_part_with_converter(
3169 &format!("key_{}", i),
3170 i as u32,
3171 vec![1000 + i as i64 * 100, 2000 + i as i64 * 100],
3172 vec![Some(i as f64 * 10.0), Some(i as f64 * 10.0 + 1.0)],
3173 100 + i as u64,
3174 )
3175 .unwrap();
3176 memtable.write_bulk(part).unwrap();
3177 }
3178
3179 memtable.compact(false).unwrap();
3181
3182 let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap();
3184 let ranges = memtable
3185 .ranges(
3186 None,
3187 RangesOptions::default().with_predicate(predicate_group),
3188 )
3189 .unwrap();
3190
3191 assert_eq!(1, ranges.ranges.len());
3192 let expected_rows = merge_threshold * 2; let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
3194 assert_eq!(expected_rows, total_rows);
3195
3196 let mut total_rows_read = 0;
3198 for range in ranges.ranges.values() {
3199 assert!(range.is_record_batch());
3200 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
3201
3202 for batch_result in record_batch_iter {
3203 let batch = batch_result.unwrap();
3204 total_rows_read += batch.num_rows();
3205 }
3206 }
3207 assert_eq!(expected_rows, total_rows_read);
3208 }
3209
3210 #[test]
3211 fn test_multi_bulk_range_iter_builder_all_pruned() {
3212 let metadata = metadata_for_test();
3213 let merge_threshold = 8;
3214 let config = BulkMemtableConfig {
3215 merge_threshold,
3216 ..Default::default()
3217 };
3218 let memtable = BulkMemtable::new(
3219 2006,
3220 config,
3221 metadata.clone(),
3222 None,
3223 None,
3224 false,
3225 MergeMode::LastRow,
3226 );
3227 memtable.set_unordered_part_threshold(0);
3228
3229 for i in 0..merge_threshold {
3231 let part = create_bulk_part_with_converter(
3232 &format!("key_{}", i),
3233 i as u32,
3234 vec![1000 + i as i64 * 100, 2000 + i as i64 * 100],
3235 vec![Some(i as f64 * 10.0), Some(i as f64 * 10.0 + 1.0)],
3236 100 + i as u64,
3237 )
3238 .unwrap();
3239 memtable.write_bulk(part).unwrap();
3240 }
3241 memtable.compact(false).unwrap();
3242
3243 let filter = datafusion_expr::col("k0").eq(datafusion_expr::lit("nonexistent"));
3245 let predicate_group = PredicateGroup::new(&metadata, &[filter]).unwrap();
3246 let ranges = memtable
3247 .ranges(
3248 None,
3249 RangesOptions::default().with_predicate(predicate_group),
3250 )
3251 .unwrap();
3252
3253 for range in ranges.ranges.values() {
3256 assert!(range.is_record_batch());
3257 let record_batch_iter = range.build_record_batch_iter(None, None).unwrap();
3258 let total_rows: usize = record_batch_iter.map(|r| r.unwrap().num_rows()).sum();
3259 assert_eq!(0, total_rows);
3260 }
3261 }
3262}