mito2/read/
flat_merge.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
15use std::cmp::Ordering;
16use std::collections::BinaryHeap;
17use std::fmt;
18use std::sync::Arc;
19use std::time::{Duration, Instant};
20
21use async_stream::try_stream;
22use common_telemetry::debug;
23use datatypes::arrow::array::{Array, AsArray, Int64Array, UInt64Array};
24use datatypes::arrow::compute::interleave;
25use datatypes::arrow::datatypes::{ArrowNativeType, BinaryType, DataType, SchemaRef, Utf8Type};
26use datatypes::arrow::error::ArrowError;
27use datatypes::arrow::record_batch::RecordBatch;
28use datatypes::arrow_array::BinaryArray;
29use datatypes::timestamp::timestamp_array_to_primitive;
30use futures::{Stream, TryStreamExt};
31use snafu::ResultExt;
32use store_api::storage::SequenceNumber;
33
34use crate::error::{ComputeArrowSnafu, Result};
35use crate::memtable::BoxedRecordBatchIterator;
36use crate::metrics::READ_STAGE_ELAPSED;
37use crate::read::BoxedRecordBatchStream;
38use crate::sst::parquet::flat_format::{
39    primary_key_column_index, sequence_column_index, time_index_column_index,
40};
41use crate::sst::parquet::format::PrimaryKeyArray;
42
43/// Checks whether interleaving the selected rows from byte columns would overflow
44/// i32 offsets. Similar to arrow-rs `interleave_bytes()`, accumulates offsets and
45/// returns an error if the capacity exceeds `i32::MAX`.
46///
47/// TODO(yingwen): Remove this after upgrading to arrow >= 58.1.0, which handles
48/// offset overflow in `interleave_bytes()` natively.
49///
50/// See: <https://github.com/apache/arrow-rs/blob/65ad652f2410fc51ad77da1805e85c0a76d9a7ea/arrow-select/src/interleave.rs#L208-L225>
51fn check_interleave_bytes_overflow<T: datatypes::arrow::datatypes::ByteArrayType>(
52    batches: &[(usize, RecordBatch)],
53    col_idx: usize,
54    indices: &[(usize, usize)],
55) -> std::result::Result<(), ArrowError> {
56    // Quick check: if concatenating all value data won't overflow, interleaving
57    // a subset of rows definitely won't either.
58    let total: usize = batches
59        .iter()
60        .map(|(_, batch)| batch.column(col_idx).as_bytes::<T>().value_data().len())
61        .sum();
62    if T::Offset::from_usize(total).is_some() {
63        return Ok(());
64    }
65    // Total exceeds the offset limit, do the precise per-row check.
66    let mut capacity: usize = 0;
67    for &(a, b) in indices {
68        let array = batches[a].1.column(col_idx).as_bytes::<T>();
69        let o = array.value_offsets();
70        let element_len = o[b + 1].as_usize() - o[b].as_usize();
71        capacity += element_len;
72        T::Offset::from_usize(capacity).ok_or(ArrowError::OffsetOverflowError(capacity))?;
73    }
74    Ok(())
75}
76
77/// Checks whether `interleave()` would overflow i32 offsets for `Utf8` or `Binary` columns.
78fn check_interleave_overflow(
79    batches: &[(usize, RecordBatch)],
80    schema: &SchemaRef,
81    indices: &[(usize, usize)],
82) -> Result<()> {
83    for (col_idx, field) in schema.fields.iter().enumerate() {
84        match field.data_type() {
85            DataType::Utf8 => {
86                check_interleave_bytes_overflow::<Utf8Type>(batches, col_idx, indices)
87                    .context(ComputeArrowSnafu)?;
88            }
89            DataType::Binary => {
90                check_interleave_bytes_overflow::<BinaryType>(batches, col_idx, indices)
91                    .context(ComputeArrowSnafu)?;
92            }
93            _ => continue,
94        }
95    }
96    Ok(())
97}
98
99/// Keeps track of the current position in a batch
100#[derive(Debug, Copy, Clone, Default)]
101struct BatchCursor {
102    /// The index into BatchBuilder::batches
103    batch_idx: usize,
104    /// The row index within the given batch
105    row_idx: usize,
106}
107
108/// Trait for reporting merge metrics.
109pub trait MergeMetricsReport: Send + Sync {
110    /// Reports and resets the metrics.
111    fn report(&self, metrics: &mut MergeMetrics);
112}
113
114/// Metrics for the merge reader.
115#[derive(Default)]
116pub struct MergeMetrics {
117    /// Cost to initialize the reader.
118    pub(crate) init_cost: Duration,
119    /// Total scan cost of the reader.
120    pub(crate) scan_cost: Duration,
121    /// Number of times to fetch batches.
122    pub(crate) num_fetch_by_batches: usize,
123    /// Number of times to fetch rows.
124    pub(crate) num_fetch_by_rows: usize,
125    /// Cost to fetch batches from sources.
126    pub(crate) fetch_cost: Duration,
127}
128
129impl fmt::Debug for MergeMetrics {
130    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
131        if self.scan_cost.is_zero() {
132            return write!(f, "{{}}");
133        }
134
135        write!(f, r#"{{"scan_cost":"{:?}""#, self.scan_cost)?;
136
137        if !self.init_cost.is_zero() {
138            write!(f, r#", "init_cost":"{:?}""#, self.init_cost)?;
139        }
140        if self.num_fetch_by_batches > 0 {
141            write!(
142                f,
143                r#", "num_fetch_by_batches":{}"#,
144                self.num_fetch_by_batches
145            )?;
146        }
147        if self.num_fetch_by_rows > 0 {
148            write!(f, r#", "num_fetch_by_rows":{}"#, self.num_fetch_by_rows)?;
149        }
150        if !self.fetch_cost.is_zero() {
151            write!(f, r#", "fetch_cost":"{:?}""#, self.fetch_cost)?;
152        }
153
154        write!(f, "}}")
155    }
156}
157
158impl MergeMetrics {
159    /// Merges metrics from another MergeMetrics instance.
160    pub(crate) fn merge(&mut self, other: &MergeMetrics) {
161        let MergeMetrics {
162            init_cost,
163            scan_cost,
164            num_fetch_by_batches,
165            num_fetch_by_rows,
166            fetch_cost,
167        } = other;
168
169        self.init_cost += *init_cost;
170        self.scan_cost += *scan_cost;
171        self.num_fetch_by_batches += *num_fetch_by_batches;
172        self.num_fetch_by_rows += *num_fetch_by_rows;
173        self.fetch_cost += *fetch_cost;
174    }
175
176    /// Reports the metrics if scan_cost exceeds 10ms and resets them.
177    pub(crate) fn maybe_report(&mut self, reporter: &Option<Arc<dyn MergeMetricsReport>>) {
178        if self.scan_cost.as_millis() > 10
179            && let Some(r) = reporter
180        {
181            r.report(self);
182        }
183    }
184}
185
186/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
187// Ports from https://github.com/apache/datafusion/blob/49.0.0/datafusion/physical-plan/src/sorts/builder.rs
188// Adds the `take_remaining_rows()` method.
189#[derive(Debug)]
190pub struct BatchBuilder {
191    /// The schema of the RecordBatches yielded by this stream
192    schema: SchemaRef,
193
194    /// Maintain a list of [`RecordBatch`] and their corresponding stream
195    batches: Vec<(usize, RecordBatch)>,
196
197    /// The current [`BatchCursor`] for each stream
198    cursors: Vec<BatchCursor>,
199
200    /// The accumulated stream indexes from which to pull rows
201    /// Consists of a tuple of `(batch_idx, row_idx)`
202    indices: Vec<(usize, usize)>,
203}
204
205impl BatchBuilder {
206    /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size`
207    pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self {
208        Self {
209            schema,
210            batches: Vec::with_capacity(stream_count * 2),
211            cursors: vec![BatchCursor::default(); stream_count],
212            indices: Vec::with_capacity(batch_size),
213        }
214    }
215
216    /// Append a new batch in `stream_idx`
217    pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) {
218        let batch_idx = self.batches.len();
219        self.batches.push((stream_idx, batch));
220        self.cursors[stream_idx] = BatchCursor {
221            batch_idx,
222            row_idx: 0,
223        };
224    }
225
226    /// Append the next row from `stream_idx`
227    pub fn push_row(&mut self, stream_idx: usize) {
228        let cursor = &mut self.cursors[stream_idx];
229        let row_idx = cursor.row_idx;
230        cursor.row_idx += 1;
231        self.indices.push((cursor.batch_idx, row_idx));
232    }
233
234    /// Returns the number of in-progress rows in this [`BatchBuilder`]
235    pub fn len(&self) -> usize {
236        self.indices.len()
237    }
238
239    /// Returns `true` if this [`BatchBuilder`] contains no in-progress rows
240    pub fn is_empty(&self) -> bool {
241        self.indices.is_empty()
242    }
243
244    /// Returns the schema of this [`BatchBuilder`]
245    pub fn schema(&self) -> &SchemaRef {
246        &self.schema
247    }
248
249    /// Drains the in_progress row indexes, and builds a new RecordBatch from them
250    ///
251    /// Will then drop any batches for which all rows have been yielded to the output
252    ///
253    /// Returns `None` if no pending rows
254    pub fn build_record_batch(&mut self) -> Result<Option<RecordBatch>> {
255        if self.is_empty() {
256            return Ok(None);
257        }
258
259        check_interleave_overflow(&self.batches, &self.schema, &self.indices)?;
260
261        let columns = (0..self.schema.fields.len())
262            .map(|column_idx| {
263                let arrays: Vec<_> = self
264                    .batches
265                    .iter()
266                    .map(|(_, batch)| batch.column(column_idx).as_ref())
267                    .collect();
268                interleave(&arrays, &self.indices).context(ComputeArrowSnafu)
269            })
270            .collect::<Result<Vec<_>>>()?;
271
272        self.indices.clear();
273
274        // New cursors are only created once the previous cursor for the stream
275        // is finished. This means all remaining rows from all but the last batch
276        // for each stream have been yielded to the newly created record batch
277        //
278        // We can therefore drop all but the last batch for each stream
279        self.retain_batches();
280
281        RecordBatch::try_new(Arc::clone(&self.schema), columns)
282            .context(ComputeArrowSnafu)
283            .map(Some)
284    }
285
286    /// Slice and take remaining rows from the last batch of `stream_idx` and push
287    /// the next batch if available.
288    pub fn take_remaining_rows(
289        &mut self,
290        stream_idx: usize,
291        next: Option<RecordBatch>,
292    ) -> RecordBatch {
293        let cursor = &mut self.cursors[stream_idx];
294        let batch = &self.batches[cursor.batch_idx];
295        let output = batch
296            .1
297            .slice(cursor.row_idx, batch.1.num_rows() - cursor.row_idx);
298        cursor.row_idx = batch.1.num_rows();
299
300        if let Some(b) = next {
301            self.push_batch(stream_idx, b);
302            self.retain_batches();
303        }
304
305        output
306    }
307
308    fn retain_batches(&mut self) {
309        let mut batch_idx = 0;
310        let mut retained = 0;
311        self.batches.retain(|(stream_idx, _)| {
312            let stream_cursor = &mut self.cursors[*stream_idx];
313            let retain = stream_cursor.batch_idx == batch_idx;
314            batch_idx += 1;
315
316            if retain {
317                stream_cursor.batch_idx = retained;
318                retained += 1;
319            }
320            retain
321        });
322    }
323}
324
325/// A comparable node of the heap.
326trait NodeCmp: Eq + Ord {
327    /// Returns whether the node still has batch to read.
328    fn is_eof(&self) -> bool;
329
330    /// Returns true if the key range of current batch in `self` is behind (exclusive) current
331    /// batch in `other`.
332    ///
333    /// # Panics
334    /// Panics if either `self` or `other` is EOF.
335    fn is_behind(&self, other: &Self) -> bool;
336}
337
338/// Common algorithm of merging sorted batches from multiple nodes.
339struct MergeAlgo<T> {
340    /// Holds nodes whose key range of current batch **is** overlapped with the merge window.
341    /// Each node yields batches from a `source`.
342    ///
343    /// Node in this heap **MUST** not be empty. A `merge window` is the (primary key, timestamp)
344    /// range of the **root node** in the `hot` heap.
345    hot: BinaryHeap<T>,
346    /// Holds nodes whose key range of current batch **isn't** overlapped with the merge window.
347    ///
348    /// Nodes in this heap **MUST** not be empty.
349    cold: BinaryHeap<T>,
350}
351
352impl<T: NodeCmp> MergeAlgo<T> {
353    /// Creates a new merge algorithm from `nodes`.
354    ///
355    /// All nodes must be initialized.
356    fn new(mut nodes: Vec<T>) -> Self {
357        // Skips EOF nodes.
358        nodes.retain(|node| !node.is_eof());
359        let hot = BinaryHeap::with_capacity(nodes.len());
360        let cold = BinaryHeap::from(nodes);
361
362        let mut algo = MergeAlgo { hot, cold };
363        // Initializes the algorithm.
364        algo.refill_hot();
365
366        algo
367    }
368
369    /// Moves nodes in `cold` heap, whose key range is overlapped with current merge
370    /// window to `hot` heap.
371    fn refill_hot(&mut self) {
372        while !self.cold.is_empty() {
373            if let Some(merge_window) = self.hot.peek() {
374                let warmest = self.cold.peek().unwrap();
375                if warmest.is_behind(merge_window) {
376                    // if the warmest node in the `cold` heap is totally after the
377                    // `merge_window`, then no need to add more nodes into the `hot`
378                    // heap for merge sorting.
379                    break;
380                }
381            }
382
383            let warmest = self.cold.pop().unwrap();
384            self.hot.push(warmest);
385        }
386    }
387
388    /// Push the node popped from `hot` back to a proper heap.
389    fn reheap(&mut self, node: T) {
390        if node.is_eof() {
391            // If the node is EOF, don't put it into the heap again.
392            // The merge window would be updated, need to refill the hot heap.
393            self.refill_hot();
394        } else {
395            // Find a proper heap for this node.
396            let node_is_cold = if let Some(hottest) = self.hot.peek() {
397                // If key range of this node is behind the hottest node's then we can
398                // push it to the cold heap. Otherwise we should push it to the hot heap.
399                node.is_behind(hottest)
400            } else {
401                // The hot heap is empty, but we don't known whether the current
402                // batch of this node is still the hottest.
403                true
404            };
405
406            if node_is_cold {
407                self.cold.push(node);
408            } else {
409                self.hot.push(node);
410            }
411            // Anyway, the merge window has been changed, we need to refill the hot heap.
412            self.refill_hot();
413        }
414    }
415
416    /// Pops the hottest node.
417    fn pop_hot(&mut self) -> Option<T> {
418        self.hot.pop()
419    }
420
421    /// Returns true if there are rows in the hot heap.
422    fn has_rows(&self) -> bool {
423        !self.hot.is_empty()
424    }
425
426    /// Returns true if we can fetch a batch directly instead of a row.
427    fn can_fetch_batch(&self) -> bool {
428        self.hot.len() == 1
429    }
430}
431
432// TODO(yingwen): Further downcast and store arrays in this struct.
433/// Columns to compare for a [RecordBatch].
434struct SortColumns {
435    primary_key: PrimaryKeyArray,
436    timestamp: Int64Array,
437    sequence: UInt64Array,
438}
439
440impl SortColumns {
441    /// Creates a new [SortColumns] from a [RecordBatch] and the position of the time index column.
442    ///
443    /// # Panics
444    /// Panics if the input batch doesn't have correct internal columns.
445    fn new(batch: &RecordBatch) -> Self {
446        let num_columns = batch.num_columns();
447        let primary_key = batch
448            .column(primary_key_column_index(num_columns))
449            .as_any()
450            .downcast_ref::<PrimaryKeyArray>()
451            .unwrap()
452            .clone();
453        let timestamp = batch.column(time_index_column_index(num_columns));
454        let (timestamp, _unit) = timestamp_array_to_primitive(timestamp).unwrap();
455        let sequence = batch
456            .column(sequence_column_index(num_columns))
457            .as_any()
458            .downcast_ref::<UInt64Array>()
459            .unwrap()
460            .clone();
461
462        Self {
463            primary_key,
464            timestamp,
465            sequence,
466        }
467    }
468
469    fn primary_key_at(&self, index: usize) -> &[u8] {
470        let key = self.primary_key.keys().value(index);
471        let binary_values = self
472            .primary_key
473            .values()
474            .as_any()
475            .downcast_ref::<BinaryArray>()
476            .unwrap();
477        binary_values.value(key as usize)
478    }
479
480    fn timestamp_at(&self, index: usize) -> i64 {
481        self.timestamp.value(index)
482    }
483
484    fn sequence_at(&self, index: usize) -> SequenceNumber {
485        self.sequence.value(index)
486    }
487
488    fn num_rows(&self) -> usize {
489        self.timestamp.len()
490    }
491}
492
493/// Cursor to a row in the [RecordBatch].
494///
495/// It compares batches by rows. During comparison, it ignores op type as sequence is enough to
496/// distinguish different rows.
497struct RowCursor {
498    /// Current row offset.
499    offset: usize,
500    /// Keys of the batch.
501    columns: SortColumns,
502}
503
504impl RowCursor {
505    fn new(columns: SortColumns) -> Self {
506        debug_assert!(columns.num_rows() > 0);
507
508        Self { offset: 0, columns }
509    }
510
511    fn is_finished(&self) -> bool {
512        self.offset >= self.columns.num_rows()
513    }
514
515    /// Returns whether advancing this cursor will finish the current batch.
516    fn is_last_row(&self) -> bool {
517        self.offset.checked_add(1) == Some(self.columns.num_rows())
518    }
519
520    fn advance(&mut self) {
521        self.offset += 1;
522    }
523
524    fn first_primary_key(&self) -> &[u8] {
525        self.columns.primary_key_at(self.offset)
526    }
527
528    fn first_timestamp(&self) -> i64 {
529        self.columns.timestamp_at(self.offset)
530    }
531
532    fn first_sequence(&self) -> SequenceNumber {
533        self.columns.sequence_at(self.offset)
534    }
535
536    fn last_primary_key(&self) -> &[u8] {
537        self.columns.primary_key_at(self.columns.num_rows() - 1)
538    }
539
540    fn last_timestamp(&self) -> i64 {
541        self.columns.timestamp_at(self.columns.num_rows() - 1)
542    }
543}
544
545impl PartialEq for RowCursor {
546    fn eq(&self, other: &Self) -> bool {
547        self.first_primary_key() == other.first_primary_key()
548            && self.first_timestamp() == other.first_timestamp()
549            && self.first_sequence() == other.first_sequence()
550    }
551}
552
553impl Eq for RowCursor {}
554
555impl PartialOrd for RowCursor {
556    fn partial_cmp(&self, other: &RowCursor) -> Option<Ordering> {
557        Some(self.cmp(other))
558    }
559}
560
561impl Ord for RowCursor {
562    /// Compares by primary key, time index, sequence desc.
563    fn cmp(&self, other: &RowCursor) -> Ordering {
564        self.first_primary_key()
565            .cmp(other.first_primary_key())
566            .then_with(|| self.first_timestamp().cmp(&other.first_timestamp()))
567            .then_with(|| other.first_sequence().cmp(&self.first_sequence()))
568    }
569}
570
571/// Iterator to merge multiple sorted iterators into a single sorted iterator.
572///
573/// All iterators must be sorted by primary key, time index, sequence desc.
574pub struct FlatMergeIterator {
575    /// The merge algorithm to maintain heaps.
576    algo: MergeAlgo<IterNode>,
577    /// Current buffered rows to output.
578    in_progress: BatchBuilder,
579    /// Non-empty batch to output.
580    output_batch: Option<RecordBatch>,
581    /// Batch size to merge rows.
582    /// This is not a hard limit, the iterator may return smaller batches to avoid concatenating
583    /// rows.
584    batch_size: usize,
585}
586
587impl FlatMergeIterator {
588    /// Creates a new iterator to merge sorted `iters`.
589    pub fn new(
590        schema: SchemaRef,
591        iters: Vec<BoxedRecordBatchIterator>,
592        batch_size: usize,
593    ) -> Result<Self> {
594        let mut in_progress = BatchBuilder::new(schema, iters.len(), batch_size);
595        let mut nodes = Vec::with_capacity(iters.len());
596        // Initialize nodes and the buffer.
597        for (node_index, iter) in iters.into_iter().enumerate() {
598            let mut node = IterNode {
599                node_index,
600                iter,
601                cursor: None,
602            };
603            if let Some(batch) = node.advance_batch()? {
604                in_progress.push_batch(node_index, batch);
605                nodes.push(node);
606            }
607        }
608
609        let algo = MergeAlgo::new(nodes);
610
611        let iter = Self {
612            algo,
613            in_progress,
614            output_batch: None,
615            batch_size,
616        };
617
618        Ok(iter)
619    }
620
621    /// Fetches next sorted batch.
622    pub fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
623        while self.algo.has_rows() && self.output_batch.is_none() {
624            if self.algo.can_fetch_batch() && !self.in_progress.is_empty() {
625                // Only one batch in the hot heap, but we have pending rows, output the pending rows first.
626                self.output_batch = self.in_progress.build_record_batch()?;
627                debug_assert!(self.output_batch.is_some());
628            } else if self.algo.can_fetch_batch() {
629                self.fetch_batch_from_hottest()?;
630            } else {
631                self.fetch_row_from_hottest()?;
632            }
633        }
634
635        Ok(self.output_batch.take())
636    }
637
638    /// Fetches a batch from the hottest node.
639    fn fetch_batch_from_hottest(&mut self) -> Result<()> {
640        debug_assert!(self.in_progress.is_empty());
641
642        // Safety: next_batch() ensures the heap is not empty.
643        let mut hottest = self.algo.pop_hot().unwrap();
644        debug_assert!(!hottest.current_cursor().is_finished());
645        let next = hottest.advance_batch()?;
646        // The node is the heap is not empty, so it must have existing rows in the builder.
647        let batch = self
648            .in_progress
649            .take_remaining_rows(hottest.node_index, next);
650        Self::maybe_output_batch(batch, &mut self.output_batch);
651        self.algo.reheap(hottest);
652
653        Ok(())
654    }
655
656    /// Fetches a row from the hottest node.
657    fn fetch_row_from_hottest(&mut self) -> Result<()> {
658        // Safety: next_batch() ensures the heap has more than 1 element.
659        let mut hottest = self.algo.pop_hot().unwrap();
660        debug_assert!(!hottest.current_cursor().is_finished());
661        self.in_progress.push_row(hottest.node_index);
662        if self.in_progress.len() >= self.batch_size {
663            // We buffered enough rows.
664            if let Some(output) = self.in_progress.build_record_batch()? {
665                Self::maybe_output_batch(output, &mut self.output_batch);
666            }
667        }
668
669        if let Some(next) = hottest.advance_row()? {
670            self.in_progress.push_batch(hottest.node_index, next);
671        }
672
673        self.algo.reheap(hottest);
674        Ok(())
675    }
676
677    /// Adds the batch to the output batch if it is not empty.
678    fn maybe_output_batch(batch: RecordBatch, output_batch: &mut Option<RecordBatch>) {
679        debug_assert!(output_batch.is_none());
680        if batch.num_rows() > 0 {
681            *output_batch = Some(batch);
682        }
683    }
684}
685
686impl Iterator for FlatMergeIterator {
687    type Item = Result<RecordBatch>;
688
689    fn next(&mut self) -> Option<Self::Item> {
690        self.next_batch().transpose()
691    }
692}
693
694/// Iterator to merge multiple sorted iterators into a single sorted iterator.
695///
696/// All iterators must be sorted by primary key, time index, sequence desc.
697pub struct FlatMergeReader {
698    /// The merge algorithm to maintain heaps.
699    algo: MergeAlgo<StreamNode>,
700    /// Current buffered rows to output.
701    in_progress: BatchBuilder,
702    /// Non-empty batch to output.
703    output_batch: Option<RecordBatch>,
704    /// Batch size to merge rows.
705    /// This is not a hard limit, the iterator may return smaller batches to avoid concatenating
706    /// rows.
707    batch_size: usize,
708    /// Local metrics.
709    metrics: MergeMetrics,
710    /// Optional metrics reporter.
711    metrics_reporter: Option<Arc<dyn MergeMetricsReport>>,
712}
713
714impl FlatMergeReader {
715    /// Creates a new iterator to merge sorted `iters`.
716    pub async fn new(
717        schema: SchemaRef,
718        iters: Vec<BoxedRecordBatchStream>,
719        batch_size: usize,
720        metrics_reporter: Option<Arc<dyn MergeMetricsReport>>,
721    ) -> Result<Self> {
722        let start = Instant::now();
723        let metrics = MergeMetrics::default();
724        let mut in_progress = BatchBuilder::new(schema, iters.len(), batch_size);
725        let mut nodes = Vec::with_capacity(iters.len());
726        // Initialize nodes and the buffer.
727        for (node_index, iter) in iters.into_iter().enumerate() {
728            let mut node = StreamNode {
729                node_index,
730                iter,
731                cursor: None,
732            };
733            if let Some(batch) = node.advance_batch().await? {
734                in_progress.push_batch(node_index, batch);
735                nodes.push(node);
736            }
737        }
738
739        let algo = MergeAlgo::new(nodes);
740
741        let mut reader = Self {
742            algo,
743            in_progress,
744            output_batch: None,
745            batch_size,
746            metrics,
747            metrics_reporter,
748        };
749        let elapsed = start.elapsed();
750        reader.metrics.init_cost += elapsed;
751        reader.metrics.scan_cost += elapsed;
752
753        Ok(reader)
754    }
755
756    /// Fetches next sorted batch.
757    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
758        let start = Instant::now();
759        while self.algo.has_rows() && self.output_batch.is_none() {
760            if self.algo.can_fetch_batch() && !self.in_progress.is_empty() {
761                // Only one batch in the hot heap, but we have pending rows, output the pending rows first.
762                self.output_batch = self.in_progress.build_record_batch()?;
763                debug_assert!(self.output_batch.is_some());
764            } else if self.algo.can_fetch_batch() {
765                self.fetch_batch_from_hottest().await?;
766                self.metrics.num_fetch_by_batches += 1;
767            } else {
768                self.fetch_row_from_hottest().await?;
769                self.metrics.num_fetch_by_rows += 1;
770            }
771        }
772
773        if let Some(batch) = self.output_batch.take() {
774            self.metrics.scan_cost += start.elapsed();
775            self.metrics.maybe_report(&self.metrics_reporter);
776            Ok(Some(batch))
777        } else {
778            // No more batches.
779            self.metrics.scan_cost += start.elapsed();
780            self.metrics.maybe_report(&self.metrics_reporter);
781            Ok(None)
782        }
783    }
784
785    /// Converts the reader into a stream.
786    pub fn into_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
787        try_stream! {
788            while let Some(batch) = self.next_batch().await? {
789                yield batch;
790            }
791        }
792    }
793
794    /// Fetches a batch from the hottest node.
795    async fn fetch_batch_from_hottest(&mut self) -> Result<()> {
796        debug_assert!(self.in_progress.is_empty());
797
798        // Safety: next_batch() ensures the heap is not empty.
799        let mut hottest = self.algo.pop_hot().unwrap();
800        debug_assert!(!hottest.current_cursor().is_finished());
801        let start = Instant::now();
802        let next = hottest.advance_batch().await?;
803        self.metrics.fetch_cost += start.elapsed();
804        // The node is the heap is not empty, so it must have existing rows in the builder.
805        let batch = self
806            .in_progress
807            .take_remaining_rows(hottest.node_index, next);
808        Self::maybe_output_batch(batch, &mut self.output_batch);
809        self.algo.reheap(hottest);
810
811        Ok(())
812    }
813
814    /// Fetches a row from the hottest node.
815    async fn fetch_row_from_hottest(&mut self) -> Result<()> {
816        // Safety: next_batch() ensures the heap has more than 1 element.
817        let mut hottest = self.algo.pop_hot().unwrap();
818        debug_assert!(!hottest.current_cursor().is_finished());
819        self.in_progress.push_row(hottest.node_index);
820        if self.in_progress.len() >= self.batch_size {
821            // We buffered enough rows.
822            if let Some(output) = self.in_progress.build_record_batch()? {
823                Self::maybe_output_batch(output, &mut self.output_batch);
824            }
825        }
826
827        // Only read the clock when advancing will attempt to fetch the next batch.
828        // Check first because `None` means either no fetch or EOF after a fetch attempt.
829        let start = hottest.current_cursor().is_last_row().then(Instant::now);
830        let next = hottest.advance_row().await?;
831        if let Some(start) = start {
832            self.metrics.fetch_cost += start.elapsed();
833        }
834        if let Some(next) = next {
835            self.in_progress.push_batch(hottest.node_index, next);
836        }
837
838        self.algo.reheap(hottest);
839        Ok(())
840    }
841
842    /// Adds the batch to the output batch if it is not empty.
843    fn maybe_output_batch(batch: RecordBatch, output_batch: &mut Option<RecordBatch>) {
844        debug_assert!(output_batch.is_none());
845        if batch.num_rows() > 0 {
846            *output_batch = Some(batch);
847        }
848    }
849}
850
851impl Drop for FlatMergeReader {
852    fn drop(&mut self) {
853        debug!("Flat merge reader finished, metrics: {:?}", self.metrics);
854
855        READ_STAGE_ELAPSED
856            .with_label_values(&["flat_merge"])
857            .observe(self.metrics.scan_cost.as_secs_f64());
858        READ_STAGE_ELAPSED
859            .with_label_values(&["flat_merge_fetch"])
860            .observe(self.metrics.fetch_cost.as_secs_f64());
861
862        // Report any remaining metrics.
863        if let Some(reporter) = &self.metrics_reporter {
864            reporter.report(&mut self.metrics);
865        }
866    }
867}
868
869/// A sync node in the merge iterator.
870struct GenericNode<T> {
871    /// Index of the node.
872    node_index: usize,
873    /// Iterator of this `Node`.
874    iter: T,
875    /// Current batch to be read. The node should ensure the batch is not empty (The
876    /// cursor is not finished).
877    ///
878    /// `None` means the `iter` has reached EOF.
879    cursor: Option<RowCursor>,
880}
881
882impl<T> NodeCmp for GenericNode<T> {
883    fn is_eof(&self) -> bool {
884        self.cursor.is_none()
885    }
886
887    fn is_behind(&self, other: &Self) -> bool {
888        debug_assert!(!self.current_cursor().is_finished());
889        debug_assert!(!other.current_cursor().is_finished());
890
891        // We only compare pk and timestamp so nodes in the cold
892        // heap don't have overlapping timestamps with the hottest node
893        // in the hot heap.
894        self.current_cursor()
895            .first_primary_key()
896            .cmp(other.current_cursor().last_primary_key())
897            .then_with(|| {
898                self.current_cursor()
899                    .first_timestamp()
900                    .cmp(&other.current_cursor().last_timestamp())
901            })
902            == Ordering::Greater
903    }
904}
905
906impl<T> PartialEq for GenericNode<T> {
907    fn eq(&self, other: &GenericNode<T>) -> bool {
908        self.cursor == other.cursor
909    }
910}
911
912impl<T> Eq for GenericNode<T> {}
913
914impl<T> PartialOrd for GenericNode<T> {
915    fn partial_cmp(&self, other: &GenericNode<T>) -> Option<Ordering> {
916        Some(self.cmp(other))
917    }
918}
919
920impl<T> Ord for GenericNode<T> {
921    fn cmp(&self, other: &GenericNode<T>) -> Ordering {
922        // The std binary heap is a max heap, but we want the nodes are ordered in
923        // ascend order, so we compare the nodes in reverse order.
924        other.cursor.cmp(&self.cursor)
925    }
926}
927
928impl<T> GenericNode<T> {
929    /// Returns current cursor.
930    ///
931    /// # Panics
932    /// Panics if the node has reached EOF.
933    fn current_cursor(&self) -> &RowCursor {
934        self.cursor.as_ref().unwrap()
935    }
936}
937
938impl GenericNode<BoxedRecordBatchIterator> {
939    /// Fetches a new batch from the iter and updates the cursor.
940    /// It advances the current batch.
941    /// Returns the fetched new batch.
942    fn advance_batch(&mut self) -> Result<Option<RecordBatch>> {
943        let batch = self.advance_inner_iter()?;
944        let columns = batch.as_ref().map(SortColumns::new);
945        self.cursor = columns.map(RowCursor::new);
946
947        Ok(batch)
948    }
949
950    /// Skips one row.
951    /// Returns the next batch if the current batch is finished.
952    fn advance_row(&mut self) -> Result<Option<RecordBatch>> {
953        let cursor = self.cursor.as_mut().unwrap();
954        cursor.advance();
955        if !cursor.is_finished() {
956            return Ok(None);
957        }
958
959        // Finished current batch, need to fetch a new batch.
960        self.advance_batch()
961    }
962
963    /// Fetches a non-empty batch from the iter.
964    fn advance_inner_iter(&mut self) -> Result<Option<RecordBatch>> {
965        while let Some(batch) = self.iter.next().transpose()? {
966            if batch.num_rows() > 0 {
967                return Ok(Some(batch));
968            }
969        }
970        Ok(None)
971    }
972}
973
974type StreamNode = GenericNode<BoxedRecordBatchStream>;
975type IterNode = GenericNode<BoxedRecordBatchIterator>;
976
977impl GenericNode<BoxedRecordBatchStream> {
978    /// Fetches a new batch from the iter and updates the cursor.
979    /// It advances the current batch.
980    /// Returns the fetched new batch.
981    async fn advance_batch(&mut self) -> Result<Option<RecordBatch>> {
982        let batch = self.advance_inner_iter().await?;
983        let columns = batch.as_ref().map(SortColumns::new);
984        self.cursor = columns.map(RowCursor::new);
985
986        Ok(batch)
987    }
988
989    /// Skips one row.
990    /// Returns the next batch if the current batch is finished.
991    async fn advance_row(&mut self) -> Result<Option<RecordBatch>> {
992        let cursor = self.cursor.as_mut().unwrap();
993        cursor.advance();
994        if !cursor.is_finished() {
995            return Ok(None);
996        }
997
998        // Finished current batch, need to fetch a new batch.
999        self.advance_batch().await
1000    }
1001
1002    /// Fetches a non-empty batch from the iter.
1003    async fn advance_inner_iter(&mut self) -> Result<Option<RecordBatch>> {
1004        while let Some(batch) = self.iter.try_next().await? {
1005            if batch.num_rows() > 0 {
1006                return Ok(Some(batch));
1007            }
1008        }
1009        Ok(None)
1010    }
1011}
1012
1013#[cfg(test)]
1014mod tests {
1015    use std::sync::Arc;
1016
1017    use api::v1::OpType;
1018    use datatypes::arrow::array::builder::BinaryDictionaryBuilder;
1019    use datatypes::arrow::array::{Int64Array, TimestampMillisecondArray, UInt8Array, UInt64Array};
1020    use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type};
1021    use datatypes::arrow::record_batch::RecordBatch;
1022
1023    use super::*;
1024
1025    /// Creates a test RecordBatch with the specified data.
1026    fn create_test_record_batch(
1027        primary_keys: &[&[u8]],
1028        timestamps: &[i64],
1029        sequences: &[u64],
1030        op_types: &[OpType],
1031        field_values: &[i64],
1032    ) -> RecordBatch {
1033        let schema = Arc::new(Schema::new(vec![
1034            Field::new("field1", DataType::Int64, false),
1035            Field::new(
1036                "timestamp",
1037                DataType::Timestamp(TimeUnit::Millisecond, None),
1038                false,
1039            ),
1040            Field::new(
1041                "__primary_key",
1042                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
1043                false,
1044            ),
1045            Field::new("__sequence", DataType::UInt64, false),
1046            Field::new("__op_type", DataType::UInt8, false),
1047        ]));
1048
1049        let field1 = Arc::new(Int64Array::from_iter_values(field_values.iter().copied()));
1050        let timestamp = Arc::new(TimestampMillisecondArray::from_iter_values(
1051            timestamps.iter().copied(),
1052        ));
1053
1054        // Create primary key dictionary array using BinaryDictionaryBuilder
1055        let mut builder = BinaryDictionaryBuilder::<UInt32Type>::new();
1056        for &key in primary_keys {
1057            builder.append(key).unwrap();
1058        }
1059        let primary_key = Arc::new(builder.finish());
1060
1061        let sequence = Arc::new(UInt64Array::from_iter_values(sequences.iter().copied()));
1062        let op_type = Arc::new(UInt8Array::from_iter_values(
1063            op_types.iter().map(|&v| v as u8),
1064        ));
1065
1066        RecordBatch::try_new(
1067            schema,
1068            vec![field1, timestamp, primary_key, sequence, op_type],
1069        )
1070        .unwrap()
1071    }
1072
1073    fn new_test_iter(batches: Vec<RecordBatch>) -> BoxedRecordBatchIterator {
1074        Box::new(batches.into_iter().map(Ok))
1075    }
1076
1077    #[test]
1078    fn test_row_cursor_last_row() {
1079        let batch = create_test_record_batch(
1080            &[b"k1", b"k1"],
1081            &[1000, 2000],
1082            &[21, 22],
1083            &[OpType::Put, OpType::Put],
1084            &[11, 12],
1085        );
1086        let mut cursor = RowCursor::new(SortColumns::new(&batch));
1087
1088        assert!(!cursor.is_last_row());
1089        cursor.advance();
1090        assert!(cursor.is_last_row());
1091        cursor.advance();
1092        assert!(!cursor.is_last_row());
1093    }
1094
1095    /// Helper function to check if two record batches are equivalent.
1096    fn assert_record_batches_eq(expected: &[RecordBatch], actual: &[RecordBatch]) {
1097        for (exp, act) in expected.iter().zip(actual.iter()) {
1098            assert_eq!(exp, act,);
1099        }
1100    }
1101
1102    /// Helper function to collect all batches from a FlatMergeIterator.
1103    fn collect_merge_iterator_batches(iter: FlatMergeIterator) -> Vec<RecordBatch> {
1104        iter.map(|result| result.unwrap()).collect()
1105    }
1106
1107    #[test]
1108    fn test_merge_iterator_empty() {
1109        let schema = Arc::new(Schema::new(vec![
1110            Field::new("field1", DataType::Int64, false),
1111            Field::new(
1112                "timestamp",
1113                DataType::Timestamp(TimeUnit::Millisecond, None),
1114                false,
1115            ),
1116            Field::new(
1117                "__primary_key",
1118                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
1119                false,
1120            ),
1121            Field::new("__sequence", DataType::UInt64, false),
1122            Field::new("__op_type", DataType::UInt8, false),
1123        ]));
1124
1125        let mut merge_iter = FlatMergeIterator::new(schema, vec![], 1024).unwrap();
1126        assert!(merge_iter.next_batch().unwrap().is_none());
1127    }
1128
1129    #[test]
1130    fn test_merge_iterator_single_batch() {
1131        let batch = create_test_record_batch(
1132            &[b"k1", b"k1"],
1133            &[1000, 2000],
1134            &[21, 22],
1135            &[OpType::Put, OpType::Put],
1136            &[11, 12],
1137        );
1138
1139        let schema = batch.schema();
1140        let iter = Box::new(new_test_iter(vec![batch.clone()]));
1141
1142        let merge_iter = FlatMergeIterator::new(schema, vec![iter], 1024).unwrap();
1143        let result = collect_merge_iterator_batches(merge_iter);
1144
1145        assert_eq!(result.len(), 1);
1146        assert_record_batches_eq(&[batch], &result);
1147    }
1148
1149    #[test]
1150    fn test_merge_iterator_non_overlapping() {
1151        let batch1 = create_test_record_batch(
1152            &[b"k1", b"k1"],
1153            &[1000, 2000],
1154            &[21, 22],
1155            &[OpType::Put, OpType::Put],
1156            &[11, 12],
1157        );
1158        let batch2 = create_test_record_batch(
1159            &[b"k1", b"k1"],
1160            &[4000, 5000],
1161            &[24, 25],
1162            &[OpType::Put, OpType::Put],
1163            &[14, 15],
1164        );
1165        let batch3 = create_test_record_batch(
1166            &[b"k2", b"k2"],
1167            &[2000, 3000],
1168            &[22, 23],
1169            &[OpType::Delete, OpType::Put],
1170            &[12, 13],
1171        );
1172
1173        let schema = batch1.schema();
1174        let iter1 = Box::new(new_test_iter(vec![batch1.clone(), batch3.clone()]));
1175        let iter2 = Box::new(new_test_iter(vec![batch2.clone()]));
1176
1177        let merge_iter = FlatMergeIterator::new(schema, vec![iter1, iter2], 1024).unwrap();
1178        let result = collect_merge_iterator_batches(merge_iter);
1179
1180        // Results should be sorted by primary key, timestamp, sequence desc
1181        let expected = vec![batch1, batch2, batch3];
1182        assert_record_batches_eq(&expected, &result);
1183    }
1184
1185    #[test]
1186    fn test_merge_iterator_overlapping_timestamps() {
1187        // Create batches with overlapping timestamps but different sequences
1188        let batch1 = create_test_record_batch(
1189            &[b"k1", b"k1"],
1190            &[1000, 2000],
1191            &[21, 22],
1192            &[OpType::Put, OpType::Put],
1193            &[11, 12],
1194        );
1195        let batch2 = create_test_record_batch(
1196            &[b"k1", b"k1"],
1197            &[1500, 2500],
1198            &[31, 32],
1199            &[OpType::Put, OpType::Put],
1200            &[15, 25],
1201        );
1202
1203        let schema = batch1.schema();
1204        let iter1 = Box::new(new_test_iter(vec![batch1]));
1205        let iter2 = Box::new(new_test_iter(vec![batch2]));
1206
1207        let merge_iter = FlatMergeIterator::new(schema, vec![iter1, iter2], 1024).unwrap();
1208        let result = collect_merge_iterator_batches(merge_iter);
1209
1210        let expected = vec![
1211            create_test_record_batch(
1212                &[b"k1", b"k1"],
1213                &[1000, 1500],
1214                &[21, 31],
1215                &[OpType::Put, OpType::Put],
1216                &[11, 15],
1217            ),
1218            create_test_record_batch(&[b"k1"], &[2000], &[22], &[OpType::Put], &[12]),
1219            create_test_record_batch(&[b"k1"], &[2500], &[32], &[OpType::Put], &[25]),
1220        ];
1221        assert_record_batches_eq(&expected, &result);
1222    }
1223
1224    #[test]
1225    fn test_merge_iterator_duplicate_keys_sequences() {
1226        // Test with same primary key and timestamp but different sequences
1227        let batch1 = create_test_record_batch(
1228            &[b"k1", b"k1"],
1229            &[1000, 1000],
1230            &[20, 10],
1231            &[OpType::Put, OpType::Put],
1232            &[1, 2],
1233        );
1234        let batch2 = create_test_record_batch(
1235            &[b"k1"],
1236            &[1000],
1237            &[15], // Middle sequence
1238            &[OpType::Put],
1239            &[3],
1240        );
1241
1242        let schema = batch1.schema();
1243        let iter1 = Box::new(new_test_iter(vec![batch1]));
1244        let iter2 = Box::new(new_test_iter(vec![batch2]));
1245
1246        let merge_iter = FlatMergeIterator::new(schema, vec![iter1, iter2], 1024).unwrap();
1247        let result = collect_merge_iterator_batches(merge_iter);
1248
1249        // Should be sorted by sequence descending for same key/timestamp
1250        let expected = vec![
1251            create_test_record_batch(
1252                &[b"k1", b"k1"],
1253                &[1000, 1000],
1254                &[20, 15],
1255                &[OpType::Put, OpType::Put],
1256                &[1, 3],
1257            ),
1258            create_test_record_batch(&[b"k1"], &[1000], &[10], &[OpType::Put], &[2]),
1259        ];
1260        assert_record_batches_eq(&expected, &result);
1261    }
1262
1263    #[test]
1264    fn test_batch_builder_basic() {
1265        let schema = Arc::new(Schema::new(vec![
1266            Field::new("field1", DataType::Int64, false),
1267            Field::new(
1268                "timestamp",
1269                DataType::Timestamp(TimeUnit::Millisecond, None),
1270                false,
1271            ),
1272        ]));
1273
1274        let mut builder = BatchBuilder::new(schema.clone(), 2, 1024);
1275        assert!(builder.is_empty());
1276
1277        let batch = RecordBatch::try_new(
1278            schema,
1279            vec![
1280                Arc::new(Int64Array::from(vec![1, 2])),
1281                Arc::new(TimestampMillisecondArray::from(vec![1000, 2000])),
1282            ],
1283        )
1284        .unwrap();
1285
1286        builder.push_batch(0, batch);
1287        builder.push_row(0);
1288        builder.push_row(0);
1289
1290        assert!(!builder.is_empty());
1291        assert_eq!(builder.len(), 2);
1292
1293        let result_batch = builder.build_record_batch().unwrap().unwrap();
1294        assert_eq!(result_batch.num_rows(), 2);
1295    }
1296
1297    #[test]
1298    fn test_row_cursor_comparison() {
1299        // Create test batches for cursor comparison
1300        let batch1 = create_test_record_batch(
1301            &[b"k1", b"k1"],
1302            &[1000, 2000],
1303            &[22, 21],
1304            &[OpType::Put, OpType::Put],
1305            &[11, 12],
1306        );
1307        let batch2 = create_test_record_batch(
1308            &[b"k1", b"k1"],
1309            &[1000, 2000],
1310            &[23, 20], // Different sequences
1311            &[OpType::Put, OpType::Put],
1312            &[11, 12],
1313        );
1314
1315        let columns1 = SortColumns::new(&batch1);
1316        let columns2 = SortColumns::new(&batch2);
1317
1318        let cursor1 = RowCursor::new(columns1);
1319        let cursor2 = RowCursor::new(columns2);
1320
1321        // cursors with same pk and timestamp should be ordered by sequence desc
1322        // cursor1 has sequence 22, cursor2 has sequence 23, so cursor2 < cursor1 (higher sequence comes first)
1323        assert!(cursor2 < cursor1);
1324    }
1325}
mito2/read/flat_merge.rs

mito2/read/
flat_merge.rs
