Skip to main content

mito2/memtable/
bulk.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Memtable implementation for bulk load
16
17pub(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
28/// Reads an environment variable as usize, returning default if not set or invalid.
29fn 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
71/// Default merge threshold for triggering compaction.
72const DEFAULT_MERGE_THRESHOLD: usize = 16;
73
74/// Threshold for triggering merge of parts. Configurable via `GREPTIME_BULK_MERGE_THRESHOLD`.
75static MERGE_THRESHOLD: LazyLock<usize> =
76    LazyLock::new(|| env_usize("GREPTIME_BULK_MERGE_THRESHOLD", DEFAULT_MERGE_THRESHOLD));
77
78/// Default maximum number of groups for parallel merging.
79const DEFAULT_MAX_MERGE_GROUPS: usize = 32;
80
81/// Maximum merge groups. Configurable via `GREPTIME_BULK_MAX_MERGE_GROUPS`.
82static MAX_MERGE_GROUPS: LazyLock<usize> =
83    LazyLock::new(|| env_usize("GREPTIME_BULK_MAX_MERGE_GROUPS", DEFAULT_MAX_MERGE_GROUPS));
84
85/// Row threshold for encoding parts. Configurable via `GREPTIME_BULK_ENCODE_ROW_THRESHOLD`.
86/// When estimated rows exceed this threshold, parts are encoded as EncodedBulkPart.
87pub(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
94/// Default bytes threshold for encoding.
95const DEFAULT_ENCODE_BYTES_THRESHOLD: usize = 64 * 1024 * 1024;
96
97/// Maximum bytes threshold for encoding when adapting to the write buffer size.
98const MAX_ENCODE_BYTES_THRESHOLD: usize = 512 * 1024 * 1024;
99
100/// Divisor to derive the encode bytes threshold from the global write buffer size.
101const ENCODE_BYTES_THRESHOLD_DIVISOR: usize = 32;
102
103/// Optional bytes threshold override from `GREPTIME_BULK_ENCODE_BYTES_THRESHOLD`.
104static 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
110/// Computes the encode bytes threshold adapted to the global write buffer size:
111/// `max(64 MiB, min(global_write_buffer_size / 32, 512 MiB))`.
112fn 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
117/// Returns the default bytes threshold for encoding parts.
118fn default_encode_bytes_threshold() -> usize {
119    ENCODE_BYTES_THRESHOLD_OVERRIDE.unwrap_or(DEFAULT_ENCODE_BYTES_THRESHOLD)
120}
121
122/// Configuration for bulk memtable.
123#[serde_as]
124#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
125#[serde(default)]
126pub struct BulkMemtableConfig {
127    /// Threshold for triggering merge of parts.
128    #[serde_as(as = "DisplayFromStr")]
129    pub merge_threshold: usize,
130    /// Row threshold for encoding parts.
131    #[serde_as(as = "DisplayFromStr")]
132    pub encode_row_threshold: usize,
133    /// Bytes threshold for encoding parts.
134    #[serde_as(as = "DisplayFromStr")]
135    pub encode_bytes_threshold: usize,
136    /// Maximum number of groups for parallel merging.
137    #[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    /// Returns the default config adapted to `global_write_buffer_bytes`.
155    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
171/// Result of merging parts - either a MultiBulkPart or an EncodedBulkPart
172enum MergedPart {
173    /// Merged part stored as MultiBulkPart (when rows < DEFAULT_ROW_GROUP_SIZE)
174    Multi(MultiBulkPart),
175    /// Merged part stored as EncodedBulkPart (when rows >= DEFAULT_ROW_GROUP_SIZE)
176    Encoded(EncodedBulkPart),
177}
178
179/// Result of collecting parts to merge
180struct CollectedParts {
181    /// Consecutive groups of parts ready for merging.
182    groups: Vec<Vec<PartToMerge>>,
183}
184
185/// All parts in a bulk memtable.
186#[derive(Default)]
187struct BulkParts {
188    /// Unordered small parts.
189    unordered_part: UnorderedPart,
190    /// All parts (raw and encoded).
191    parts: Vec<BulkPartWrapper>,
192}
193
194impl BulkParts {
195    /// Total number of parts (including unordered).
196    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    /// Returns true if there is no part.
202    fn is_empty(&self) -> bool {
203        self.unordered_part.is_empty() && self.parts.is_empty()
204    }
205
206    /// Returns a conservative scheduling hint without sorting or allocating.
207    /// The picker checks sequence barriers and sizes when compaction runs.
208    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    /// Returns true if the unordered_part should be compacted into a BulkPart.
229    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    /// Converts buffered small writes into a sorted part once a threshold is reached.
235    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    /// Selects consecutive sequence components without crossing unordered, busy,
250    /// or oversized parts. Overlapping components are indivisible, even above the size limit.
251    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        // The unordered interval participates in overlap grouping as a barrier,
261        // so neither overlapping parts nor parts on opposite sides can be merged.
262        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    /// Installs merged parts and removes the original parts by file ids.
336    /// Returns the total number of rows in the merged parts.
337    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    /// Resets merging flag for parts with the given file ids.
379    /// Used when merging fails or is cancelled.
380    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
389/// RAII guard for managing merging flags.
390/// Automatically resets merging flags when dropped if the merge operation wasn't successful.
391struct MergingFlagsGuard<'a> {
392    bulk_parts: &'a RwLock<BulkParts>,
393    file_ids: &'a HashSet<FileId>,
394    success: bool,
395}
396
397impl<'a> MergingFlagsGuard<'a> {
398    /// Creates a new guard for the given file ids.
399    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    /// Marks the merge operation as successful.
408    /// When this is called, the guard will not reset the flags on drop.
409    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
424/// Memtable that ingests and scans parts directly.
425pub struct BulkMemtable {
426    id: MemtableId,
427    /// Configuration for the bulk memtable.
428    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 for merging bulk parts
438    compactor: Arc<Mutex<MemtableCompactor>>,
439    /// Dispatcher for scheduling compaction tasks
440    compact_dispatcher: Option<Arc<CompactDispatcher>>,
441    /// Whether a write-triggered compaction is queued or running.
442    compacting: Arc<AtomicBool>,
443    /// Whether the append mode is enabled
444    append_mode: bool,
445    /// Mode to handle duplicate rows while merging
446    merge_mode: MergeMode,
447    /// Max number of rows in a parquet row group for encoded parts.
448    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            // Routes small parts to unordered_part based on row and byte thresholds.
499            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                // Compacts unordered_part if the row or byte threshold is exceeded.
505                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            // Since this operation should be fast, we do it in parts lock scope.
519            // This ensure the statistics in `ranges()` are correct. What's more,
520            // it guarantees no rows are out of the time range so we don't need to
521            // prune rows by time range again in the iterator of the MemtableRange.
522            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        // TODO(yingwen): Filter ranges by sequence.
541        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        // Adds ranges for regular parts and encoded parts
551        {
552            let bulk_parts = self.parts.read().unwrap();
553
554            // Adds range for unordered part if not empty
555            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            // Adds ranges for all parts
577            for part_wrapper in bulk_parts.parts.iter() {
578                // Skips empty parts
579                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    /// Creates a new BulkMemtable with the default row group size.
727    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    /// Creates a new BulkMemtable with the given `row_group_size`.
749    #[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    /// Creates a new BulkMemtable with the given row group size and float encoding.
774    #[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    /// Sets the unordered part threshold (for testing).
816    #[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    /// Sets the unordered part compact threshold (for testing).
826    #[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    /// Updates memtable stats.
836    ///
837    /// Please update this inside the write lock scope.
838    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    /// Returns the estimated time series count.
854    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    /// Schedules at most one write-triggered compaction at a time.
864    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            // Uses synchronous compaction if no dispatcher is available.
892            if let Err(e) = self.compact(false) {
893                common_telemetry::error!(e; "Failed to compact table");
894            }
895        }
896    }
897}
898
899/// Iterator builder for bulk range
900pub struct BulkRangeIterBuilder {
901    pub part: BulkPart,
902    pub context: Arc<BulkIterContext>,
903    pub sequence: Option<SequenceRange>,
904}
905
906/// Iterator builder for multi bulk range
907struct 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            // All batches were pruned by the predicate. Return an empty iterator.
979            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
992/// Iterator builder for encoded bulk range
993struct 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            // Return an empty iterator if no data to read
1024            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    /// The part to store. It already contains the file id.
1042    part: PartToMerge,
1043    /// Whether this part is currently being merged.
1044    merging: bool,
1045}
1046
1047impl BulkPartWrapper {
1048    /// Returns the file id of this part.
1049    fn file_id(&self) -> FileId {
1050        self.part.file_id()
1051    }
1052}
1053
1054/// Enum to wrap different types of parts for unified merging.
1055#[derive(Clone)]
1056enum PartToMerge {
1057    /// Raw bulk part.
1058    Bulk { part: BulkPart, file_id: FileId },
1059    /// Multiple bulk parts.
1060    Multi {
1061        part: MultiBulkPart,
1062        file_id: FileId,
1063    },
1064    /// Encoded bulk part.
1065    Encoded {
1066        part: EncodedBulkPart,
1067        file_id: FileId,
1068    },
1069}
1070
1071impl PartToMerge {
1072    /// Gets the file ID of this part.
1073    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    /// Gets the minimum timestamp of this part.
1082    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    /// Gets the maximum timestamp of this part.
1091    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    /// Gets the number of rows in this part.
1100    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    /// Returns the lower bound of input sequences.
1109    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    /// Gets the maximum sequence number of this part.
1118    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    /// Gets the estimated series count in this part.
1127    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    /// Returns true if this is an encoded part.
1136    fn is_encoded(&self) -> bool {
1137        matches!(self, PartToMerge::Encoded { .. })
1138    }
1139
1140    /// Gets the estimated size in bytes of this part.
1141    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    /// Returns `(num_rows, estimated_decoded_bytes)` for batch sizing.
1150    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    /// Converts this part to `MemtableStats`.
1176    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    /// Returns the Arrow schema of the record batches contained in this [`PartToMerge`].
1185    fn arrow_schema(&self) -> SchemaRef {
1186        match self {
1187            PartToMerge::Bulk { part, .. } => part.schema(),
1188            // A MultiBulkPart is built from batches that have already been aligned, so
1189            // all contained batches are expected to share the same arrow schema.
1190            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    /// Creates a record batch iterator for this part.
1199    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, // No sequence filter for merging
1211                    series_count,
1212                    None, // No metrics for merging
1213                );
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    /// Configuration for the bulk memtable.
1234    config: BulkMemtableConfig,
1235    /// Max number of rows in a parquet row group for encoded parts.
1236    row_group_size: usize,
1237    float_field_encoding: FloatFieldEncoding,
1238}
1239
1240impl MemtableCompactor {
1241    /// Creates a new MemtableCompactor.
1242    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    /// Merges parts (bulk and encoded) and then encodes the result.
1259    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        // Collect all file IDs for tracking
1282        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        // Merge all groups in parallel
1299        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        // Install all merged parts
1317        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    /// Merges a group of parts into a single part (either MultiBulkPart or EncodedBulkPart).
1338    #[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        // Calculates timestamp bounds and statistics for merged data
1354        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        // Collects statistics from parts before creating iterators
1376        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, // No column projection for merging
1392            None, // No predicate for merging
1393            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        // Parts in one merge group may differ only in their JSON2 physical layouts. So every source
1411        // schema is therefore a valid template for producing the final target schema that has
1412        // the union JSON2 types (rewritten).
1413        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        // Encode as EncodedBulkPart if rows exceed row threshold OR bytes exceed bytes threshold
1450        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            // Otherwise, collect into MultiBulkPart
1473            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
1507/// Keeps the compaction flag set while queued or running and clears it on drop,
1508/// including failures and cancellation while waiting for a dispatcher permit.
1509struct CompactTaskGuard(Arc<AtomicBool>);
1510
1511impl CompactTaskGuard {
1512    fn try_new(compacting: Arc<AtomicBool>) -> Option<Self> {
1513        // The compactor mutex protects merging; this flag only avoids duplicate tasks.
1514        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
1527/// A memtable compact task to run in background.
1528struct MemCompactTask {
1529    metadata: RegionMetadataRef,
1530    parts: Arc<RwLock<BulkParts>>,
1531    _guard: CompactTaskGuard,
1532    /// Compactor for merging bulk parts
1533    compactor: Arc<Mutex<MemtableCompactor>>,
1534    /// Whether the append mode is enabled
1535    append_mode: bool,
1536    /// Mode to handle duplicate rows while merging
1537    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/// Scheduler to run compact tasks in background.
1552#[derive(Debug)]
1553pub struct CompactDispatcher {
1554    semaphore: Arc<Semaphore>,
1555}
1556
1557impl CompactDispatcher {
1558    /// Creates a new dispatcher with the given number of max concurrent tasks.
1559    pub fn new(permits: usize) -> Self {
1560        Self {
1561            semaphore: Arc::new(Semaphore::new(permits)),
1562        }
1563    }
1564
1565    /// Dispatches a compact task to run in background.
1566    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                // The async task above returns as soon as the blocking task is
1575                // submitted, so the permit has to travel with it to bound the
1576                // compaction that actually runs.
1577                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/// Builder to build a [BulkMemtable].
1587#[derive(Debug)]
1588pub struct BulkMemtableBuilder {
1589    /// Configuration for the bulk memtable.
1590    config: BulkMemtableConfig,
1591    write_buffer_manager: Option<WriteBufferManagerRef>,
1592    compact_dispatcher: Option<Arc<CompactDispatcher>>,
1593    append_mode: bool,
1594    merge_mode: MergeMode,
1595    /// Max number of rows in a parquet row group for encoded parts.
1596    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    /// Creates a new builder with specific `write_buffer_manager`.
1616    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    /// Sets the bulk memtable config.
1630    pub fn with_config(mut self, config: BulkMemtableConfig) -> Self {
1631        self.config = config;
1632        self
1633    }
1634
1635    /// Sets the max number of rows in a parquet row group for encoded parts.
1636    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    /// Sets the compact dispatcher.
1647    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        // Below the lower bound: clamped to the default 64 MiB.
1743        assert_eq!(
1744            DEFAULT_ENCODE_BYTES_THRESHOLD,
1745            adaptive_encode_bytes_threshold(1024 * 1024 * 1024) // 1 GiB / 32 = 32 MiB
1746        );
1747        // In range: global_write_buffer_size / 32.
1748        assert_eq!(
1749            256 * 1024 * 1024,
1750            adaptive_encode_bytes_threshold(8 * 1024 * 1024 * 1024) // 8 GiB / 32 = 256 MiB
1751        );
1752        // Above the upper bound: clamped to 512 MiB.
1753        assert_eq!(
1754            MAX_ENCODE_BYTES_THRESHOLD,
1755            adaptive_encode_bytes_threshold(64 * 1024 * 1024 * 1024) // 64 GiB / 32 = 2 GiB
1756        );
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            // Regular row writes produce heterogeneous parts, sorted by timestamp
1807            // rather than sequence. Neither the scalar nor the first row is a minimum.
1808            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        // The third write synchronously merges all parts without a dispatcher.
1822        // Dedup can remove the oldest rows, but their lower bound remains safe.
1823        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        // Disable unordered_part for this test
1853        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        // Neither path is explicit in every source; keep both opaque without losing values.
2049        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        // Disable unordered_part for this test
2364        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 }); // Filters out rows with sequence > 400
2435        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        // Disable unordered_part for this test
2468        memtable.set_unordered_part_threshold(0);
2469
2470        // Adds enough bulk parts to trigger encoding
2471        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        // Should have ranges for both bulk parts and encoded parts
2494        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        // Set smaller thresholds for testing with smaller inputs
2527        // Accept parts with < 5 rows into unordered_part
2528        memtable.set_unordered_part_threshold(5);
2529        // Compact when total rows >= 10
2530        memtable.set_unordered_part_compact_threshold(10);
2531
2532        // Write 3 small parts (each has 2 rows), should be collected in unordered_part
2533        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        // Total rows = 6, not yet reaching compact threshold
2547        let stats = memtable.stats();
2548        assert_eq!(6, stats.num_rows);
2549
2550        // Write 2 more small parts (each has 2 rows)
2551        // This should trigger compaction when total >= 10
2552        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        // Total rows = 10, should have compacted unordered_part into a regular part
2565        let stats = memtable.stats();
2566        assert_eq!(10, stats.num_rows);
2567
2568        // Verify we can read all data correctly
2569        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        // Should have at least 1 range (the compacted part)
2578        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        // Read all data and verify
2583        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        // Set threshold to 4 rows - parts with < 4 rows go to unordered_part
2610        memtable.set_unordered_part_threshold(4);
2611        memtable.set_unordered_part_compact_threshold(8);
2612
2613        // Write small parts (3 rows each) - should go to unordered_part
2614        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        // Write a large part (5 rows) - should go directly to regular parts
2628        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        // A later small write fills the unordered part to its compact threshold.
2646        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); // 3 + 3 + 5 + 2 = 13
2658
2659        // Verify all data can be read
2660        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        // A part exactly at or below the threshold is accepted.
2705        memtable.write_bulk(first).unwrap();
2706        assert_eq!(1, memtable.parts.read().unwrap().unordered_part.num_parts());
2707
2708        // Compact after the accumulated size becomes larger than the threshold.
2709        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        // Set small thresholds
2750        memtable.set_unordered_part_threshold(3);
2751        memtable.set_unordered_part_compact_threshold(100); // High threshold to prevent auto-compaction
2752
2753        // Write several small parts that stay in unordered_part
2754        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        // Test that ranges() can correctly read from unordered_part
2771        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        // Should have 1 range for the unordered_part
2780        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        // Verify data is sorted correctly in the range
2785        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            // Verify data is properly sorted by primary key
2793            assert!(batch.num_rows() > 0);
2794        }
2795        assert_eq!(3, total_rows);
2796    }
2797
2798    /// Helper to create a BulkPartWrapper from a BulkPart.
2799    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        // Add one fewer part than the merge threshold.
2814        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        // There are not enough parts to merge.
2831        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        // Add 8 bulk parts (at merge_threshold)
2840        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        // Should trigger merge since we have 8 parts
2853        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        // Add 10 bulk parts
2889        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        // Should trigger merge since we have 10 parts
2902        assert!(bulk_parts.may_merge_parts(merge_threshold));
2903
2904        // Mark first 3 parts as merging
2905        for wrapper in bulk_parts.parts.iter_mut().take(3) {
2906            wrapper.merging = true;
2907        }
2908
2909        // Now only 7 parts are available for merging, should not trigger
2910        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        // Deliberately use installation order rather than sequence order.
2962        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        // The oversized encoded part connects three intervals transitively;
2985        // neither the merge threshold nor its size may split that component.
2986        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        // The hint may schedule a merge even when unordered writes block all parts.
2992        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        // Touching the unordered interval blocks the entire transitive component.
2998        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        // Even without overlap, a group cannot straddle pending small writes.
3003        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        // Disable unordered_part for this test
3162        memtable.set_unordered_part_threshold(0);
3163
3164        // Write enough bulk parts to trigger merge (merge_threshold = 8)
3165        // Each part has small number of rows so total < DEFAULT_ROW_GROUP_SIZE
3166        // This will result in MultiBulkPart after compaction
3167        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        // Compact to trigger MultiBulkPart creation (since total rows < DEFAULT_ROW_GROUP_SIZE)
3180        memtable.compact(false).unwrap();
3181
3182        // Verify we can read from the memtable
3183        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; // Each part has 2 rows
3193        let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum();
3194        assert_eq!(expected_rows, total_rows);
3195
3196        // Read all data
3197        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        // Write enough bulk parts to trigger merge into MultiBulkPart.
3230        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        // Use a predicate that matches no rows so all batches are pruned.
3244        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        // Should return ranges but each range should produce an empty iterator
3254        // instead of an error.
3255        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}