mito2/
read.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//! Common structs and utilities for reading data.
16
17pub mod batch_adapter;
18pub mod compat;
19pub mod dedup;
20pub mod flat_dedup;
21pub mod flat_merge;
22pub mod flat_projection;
23pub mod last_row;
24pub mod projection;
25pub(crate) mod prune;
26pub(crate) mod pruner;
27pub mod range;
28#[cfg(feature = "test")]
29pub mod range_cache;
30#[cfg(not(feature = "test"))]
31pub(crate) mod range_cache;
32pub mod read_columns;
33pub mod scan_region;
34pub mod scan_util;
35pub(crate) mod seq_scan;
36pub(crate) mod series_candidate;
37pub(crate) mod series_reader;
38pub mod series_scan;
39pub mod stream;
40pub(crate) mod unordered_scan;
41
42use std::collections::HashMap;
43use std::sync::Arc;
44use std::time::Duration;
45
46use api::v1::OpType;
47use arrow_schema::SchemaRef;
48use async_trait::async_trait;
49use common_time::Timestamp;
50use datafusion_common::arrow::array::UInt8Array;
51use datatypes::arrow;
52use datatypes::arrow::array::{Array, ArrayRef};
53use datatypes::arrow::compute::SortOptions;
54use datatypes::arrow::record_batch::RecordBatch;
55use datatypes::arrow::row::{RowConverter, SortField};
56use datatypes::prelude::{ConcreteDataType, DataType, ScalarVector};
57use datatypes::scalars::ScalarVectorBuilder;
58use datatypes::types::TimestampType;
59use datatypes::value::{Value, ValueRef};
60use datatypes::vectors::{
61    BooleanVector, Helper, TimestampMicrosecondVector, TimestampMillisecondVector,
62    TimestampMillisecondVectorBuilder, TimestampNanosecondVector, TimestampSecondVector,
63    UInt8Vector, UInt8VectorBuilder, UInt32Vector, UInt64Vector, UInt64VectorBuilder, Vector,
64    VectorRef,
65};
66use futures::TryStreamExt;
67use futures::stream::BoxStream;
68use mito_codec::row_converter::{CompositeValues, PrimaryKeyCodec};
69use snafu::{OptionExt, ResultExt, ensure};
70use store_api::storage::{ColumnId, SequenceNumber, SequenceRange};
71
72use crate::error::{
73    ComputeArrowSnafu, ComputeVectorSnafu, ConvertVectorSnafu, DecodeSnafu, InvalidBatchSnafu,
74    Result,
75};
76use crate::memtable::{BoxedBatchIterator, BoxedRecordBatchIterator};
77
78pub(crate) fn timestamp_array_to_i64_slice(arr: &ArrayRef) -> &[i64] {
79    use datatypes::arrow::array::{
80        TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
81        TimestampSecondArray,
82    };
83    use datatypes::arrow::datatypes::{DataType, TimeUnit};
84
85    match arr.data_type() {
86        DataType::Timestamp(t, _) => match t {
87            TimeUnit::Second => arr
88                .as_any()
89                .downcast_ref::<TimestampSecondArray>()
90                .unwrap()
91                .values(),
92            TimeUnit::Millisecond => arr
93                .as_any()
94                .downcast_ref::<TimestampMillisecondArray>()
95                .unwrap()
96                .values(),
97            TimeUnit::Microsecond => arr
98                .as_any()
99                .downcast_ref::<TimestampMicrosecondArray>()
100                .unwrap()
101                .values(),
102            TimeUnit::Nanosecond => arr
103                .as_any()
104                .downcast_ref::<TimestampNanosecondArray>()
105                .unwrap()
106                .values(),
107        },
108        _ => unreachable!(),
109    }
110}
111
112/// Storage internal representation of a batch of rows for a primary key (time series).
113///
114/// Rows are sorted by primary key, timestamp, sequence desc, op_type desc. Fields
115/// always keep the same relative order as fields in [RegionMetadata](store_api::metadata::RegionMetadata).
116#[derive(Debug, PartialEq, Clone)]
117pub struct Batch {
118    /// Primary key encoded in a comparable form.
119    primary_key: Vec<u8>,
120    /// Possibly decoded `primary_key` values. Some places would decode it in advance.
121    pk_values: Option<CompositeValues>,
122    /// Timestamps of rows, should be sorted and not null.
123    timestamps: VectorRef,
124    /// Sequences of rows
125    ///
126    /// UInt64 type, not null.
127    sequences: Arc<UInt64Vector>,
128    /// Op types of rows
129    ///
130    /// UInt8 type, not null.
131    op_types: Arc<UInt8Vector>,
132    /// Fields organized in columnar format.
133    fields: Vec<BatchColumn>,
134    /// Cache for field index lookup.
135    fields_idx: Option<HashMap<ColumnId, usize>>,
136}
137
138impl Batch {
139    /// Creates a new batch.
140    pub fn new(
141        primary_key: Vec<u8>,
142        timestamps: VectorRef,
143        sequences: Arc<UInt64Vector>,
144        op_types: Arc<UInt8Vector>,
145        fields: Vec<BatchColumn>,
146    ) -> Result<Batch> {
147        BatchBuilder::with_required_columns(primary_key, timestamps, sequences, op_types)
148            .with_fields(fields)
149            .build()
150    }
151
152    /// Tries to set fields for the batch.
153    pub fn with_fields(self, fields: Vec<BatchColumn>) -> Result<Batch> {
154        Batch::new(
155            self.primary_key,
156            self.timestamps,
157            self.sequences,
158            self.op_types,
159            fields,
160        )
161    }
162
163    /// Returns primary key of the batch.
164    pub fn primary_key(&self) -> &[u8] {
165        &self.primary_key
166    }
167
168    /// Returns possibly decoded primary-key values.
169    pub fn pk_values(&self) -> Option<&CompositeValues> {
170        self.pk_values.as_ref()
171    }
172
173    /// Sets possibly decoded primary-key values.
174    pub fn set_pk_values(&mut self, pk_values: CompositeValues) {
175        self.pk_values = Some(pk_values);
176    }
177
178    /// Removes possibly decoded primary-key values. For testing only.
179    #[cfg(any(test, feature = "test"))]
180    pub fn remove_pk_values(&mut self) {
181        self.pk_values = None;
182    }
183
184    /// Returns fields in the batch.
185    pub fn fields(&self) -> &[BatchColumn] {
186        &self.fields
187    }
188
189    /// Returns timestamps of the batch.
190    pub fn timestamps(&self) -> &VectorRef {
191        &self.timestamps
192    }
193
194    /// Returns sequences of the batch.
195    pub fn sequences(&self) -> &Arc<UInt64Vector> {
196        &self.sequences
197    }
198
199    /// Returns op types of the batch.
200    pub fn op_types(&self) -> &Arc<UInt8Vector> {
201        &self.op_types
202    }
203
204    /// Returns the number of rows in the batch.
205    pub fn num_rows(&self) -> usize {
206        // All vectors have the same length. We use the length of sequences vector
207        // since it has static type.
208        self.sequences.len()
209    }
210
211    /// Create an empty [`Batch`].
212    #[allow(dead_code)]
213    pub(crate) fn empty() -> Self {
214        Self {
215            primary_key: vec![],
216            pk_values: None,
217            timestamps: Arc::new(TimestampMillisecondVectorBuilder::with_capacity(0).finish()),
218            sequences: Arc::new(UInt64VectorBuilder::with_capacity(0).finish()),
219            op_types: Arc::new(UInt8VectorBuilder::with_capacity(0).finish()),
220            fields: vec![],
221            fields_idx: None,
222        }
223    }
224
225    /// Returns true if the number of rows in the batch is 0.
226    pub fn is_empty(&self) -> bool {
227        self.num_rows() == 0
228    }
229
230    /// Returns the first timestamp in the batch or `None` if the batch is empty.
231    pub fn first_timestamp(&self) -> Option<Timestamp> {
232        if self.timestamps.is_empty() {
233            return None;
234        }
235
236        Some(self.get_timestamp(0))
237    }
238
239    /// Returns the last timestamp in the batch or `None` if the batch is empty.
240    pub fn last_timestamp(&self) -> Option<Timestamp> {
241        if self.timestamps.is_empty() {
242            return None;
243        }
244
245        Some(self.get_timestamp(self.timestamps.len() - 1))
246    }
247
248    /// Returns the first sequence in the batch or `None` if the batch is empty.
249    pub fn first_sequence(&self) -> Option<SequenceNumber> {
250        if self.sequences.is_empty() {
251            return None;
252        }
253
254        Some(self.get_sequence(0))
255    }
256
257    /// Returns the last sequence in the batch or `None` if the batch is empty.
258    pub fn last_sequence(&self) -> Option<SequenceNumber> {
259        if self.sequences.is_empty() {
260            return None;
261        }
262
263        Some(self.get_sequence(self.sequences.len() - 1))
264    }
265
266    /// Replaces the primary key of the batch.
267    ///
268    /// Notice that this [Batch] also contains a maybe-exist `pk_values`.
269    /// Be sure to update that field as well.
270    pub fn set_primary_key(&mut self, primary_key: Vec<u8>) {
271        self.primary_key = primary_key;
272    }
273
274    /// Slice the batch, returning a new batch.
275    ///
276    /// # Panics
277    /// Panics if `offset + length > self.num_rows()`.
278    pub fn slice(&self, offset: usize, length: usize) -> Batch {
279        let fields = self
280            .fields
281            .iter()
282            .map(|column| BatchColumn {
283                column_id: column.column_id,
284                data: column.data.slice(offset, length),
285            })
286            .collect();
287        // We skip using the builder to avoid validating the batch again.
288        Batch {
289            // Now we need to clone the primary key. We could try `Bytes` if
290            // this becomes a bottleneck.
291            primary_key: self.primary_key.clone(),
292            pk_values: self.pk_values.clone(),
293            timestamps: self.timestamps.slice(offset, length),
294            sequences: Arc::new(self.sequences.get_slice(offset, length)),
295            op_types: Arc::new(self.op_types.get_slice(offset, length)),
296            fields,
297            fields_idx: self.fields_idx.clone(),
298        }
299    }
300
301    /// Takes `batches` and concat them into one batch.
302    ///
303    /// All `batches` must have the same primary key.
304    pub fn concat(mut batches: Vec<Batch>) -> Result<Batch> {
305        ensure!(
306            !batches.is_empty(),
307            InvalidBatchSnafu {
308                reason: "empty batches",
309            }
310        );
311        if batches.len() == 1 {
312            // Now we own the `batches` so we could pop it directly.
313            return Ok(batches.pop().unwrap());
314        }
315
316        let primary_key = std::mem::take(&mut batches[0].primary_key);
317        let first = &batches[0];
318        // We took the primary key from the first batch so we don't use `first.primary_key()`.
319        ensure!(
320            batches
321                .iter()
322                .skip(1)
323                .all(|b| b.primary_key() == primary_key),
324            InvalidBatchSnafu {
325                reason: "batches have different primary key",
326            }
327        );
328        for b in batches.iter().skip(1) {
329            ensure!(
330                b.fields.len() == first.fields.len(),
331                InvalidBatchSnafu {
332                    reason: "batches have different field num",
333                }
334            );
335            for (l, r) in b.fields.iter().zip(&first.fields) {
336                ensure!(
337                    l.column_id == r.column_id,
338                    InvalidBatchSnafu {
339                        reason: "batches have different fields",
340                    }
341                );
342            }
343        }
344
345        // We take the primary key from the first batch.
346        let mut builder = BatchBuilder::new(primary_key);
347        // Concat timestamps, sequences, op_types, fields.
348        let array = concat_arrays(batches.iter().map(|b| b.timestamps().to_arrow_array()))?;
349        builder.timestamps_array(array)?;
350        let array = concat_arrays(batches.iter().map(|b| b.sequences().to_arrow_array()))?;
351        builder.sequences_array(array)?;
352        let array = concat_arrays(batches.iter().map(|b| b.op_types().to_arrow_array()))?;
353        builder.op_types_array(array)?;
354        for (i, batch_column) in first.fields.iter().enumerate() {
355            let array = concat_arrays(batches.iter().map(|b| b.fields()[i].data.to_arrow_array()))?;
356            builder.push_field_array(batch_column.column_id, array)?;
357        }
358
359        builder.build()
360    }
361
362    /// Removes rows whose op type is delete.
363    pub fn filter_deleted(&mut self) -> Result<()> {
364        // Safety: op type column is not null.
365        let array = self.op_types.as_arrow();
366        // Find rows with non-delete op type.
367        let rhs = UInt8Array::new_scalar(OpType::Delete as u8);
368        let predicate =
369            arrow::compute::kernels::cmp::neq(array, &rhs).context(ComputeArrowSnafu)?;
370        self.filter(&BooleanVector::from(predicate))
371    }
372
373    // Applies the `predicate` to the batch.
374    // Safety: We know the array type so we unwrap on casting.
375    pub fn filter(&mut self, predicate: &BooleanVector) -> Result<()> {
376        self.timestamps = self
377            .timestamps
378            .filter(predicate)
379            .context(ComputeVectorSnafu)?;
380        self.sequences = Arc::new(
381            UInt64Vector::try_from_arrow_array(
382                arrow::compute::filter(self.sequences.as_arrow(), predicate.as_boolean_array())
383                    .context(ComputeArrowSnafu)?,
384            )
385            .unwrap(),
386        );
387        self.op_types = Arc::new(
388            UInt8Vector::try_from_arrow_array(
389                arrow::compute::filter(self.op_types.as_arrow(), predicate.as_boolean_array())
390                    .context(ComputeArrowSnafu)?,
391            )
392            .unwrap(),
393        );
394        for batch_column in &mut self.fields {
395            batch_column.data = batch_column
396                .data
397                .filter(predicate)
398                .context(ComputeVectorSnafu)?;
399        }
400
401        Ok(())
402    }
403
404    /// Filters rows by the given `sequence`. Only preserves rows with sequence less than or equal to `sequence`.
405    pub fn filter_by_sequence(&mut self, sequence: Option<SequenceRange>) -> Result<()> {
406        let seq_range = match sequence {
407            None => return Ok(()),
408            Some(seq_range) => {
409                let (Some(first), Some(last)) = (self.first_sequence(), self.last_sequence())
410                else {
411                    return Ok(());
412                };
413                let is_subset = match seq_range {
414                    SequenceRange::Gt { min } => min < first,
415                    SequenceRange::LtEq { max } => max >= last,
416                    SequenceRange::GtLtEq { min, max } => min < first && max >= last,
417                };
418                if is_subset {
419                    return Ok(());
420                }
421                seq_range
422            }
423        };
424
425        let seqs = self.sequences.as_arrow();
426        let predicate = seq_range.filter(seqs).context(ComputeArrowSnafu)?;
427
428        let predicate = BooleanVector::from(predicate);
429        self.filter(&predicate)?;
430
431        Ok(())
432    }
433
434    /// Sorts rows in the batch. If `dedup` is true, it also removes
435    /// duplicated rows according to primary keys.
436    ///
437    /// It orders rows by timestamp, sequence desc and only keep the latest
438    /// row for the same timestamp. It doesn't consider op type as sequence
439    /// should already provide uniqueness for a row.
440    pub fn sort(&mut self, dedup: bool) -> Result<()> {
441        // If building a converter each time is costly, we may allow passing a
442        // converter.
443        let converter = RowConverter::new(vec![
444            SortField::new(self.timestamps.data_type().as_arrow_type()),
445            SortField::new_with_options(
446                self.sequences.data_type().as_arrow_type(),
447                SortOptions {
448                    descending: true,
449                    ..Default::default()
450                },
451            ),
452        ])
453        .context(ComputeArrowSnafu)?;
454        // Columns to sort.
455        let columns = [
456            self.timestamps.to_arrow_array(),
457            self.sequences.to_arrow_array(),
458        ];
459        let rows = converter.convert_columns(&columns).unwrap();
460        let mut to_sort: Vec<_> = rows.iter().enumerate().collect();
461
462        let was_sorted = to_sort.is_sorted_by_key(|x| x.1);
463        if !was_sorted {
464            to_sort.sort_unstable_by_key(|x| x.1);
465        }
466
467        let num_rows = to_sort.len();
468        if dedup {
469            // Dedup by timestamps.
470            to_sort.dedup_by(|left, right| {
471                debug_assert_eq!(18, left.1.as_ref().len());
472                debug_assert_eq!(18, right.1.as_ref().len());
473                let (left_key, right_key) = (left.1.as_ref(), right.1.as_ref());
474                // We only compare the timestamp part and ignore sequence.
475                left_key[..TIMESTAMP_KEY_LEN] == right_key[..TIMESTAMP_KEY_LEN]
476            });
477        }
478        let no_dedup = to_sort.len() == num_rows;
479
480        if was_sorted && no_dedup {
481            return Ok(());
482        }
483        let indices = UInt32Vector::from_iter_values(to_sort.iter().map(|v| v.0 as u32));
484        self.take_in_place(&indices)
485    }
486
487    /// Merges duplicated timestamps in the batch by keeping the latest non-null field values.
488    ///
489    /// Rows must already be sorted by timestamp (ascending) and sequence (descending).
490    ///
491    /// This method deduplicates rows with the same timestamp (keeping the first row in each
492    /// timestamp range as the base row) and fills null fields from subsequent rows until all
493    /// fields are filled or a delete operation is encountered.
494    pub(crate) fn merge_last_non_null(&mut self) -> Result<()> {
495        let num_rows = self.num_rows();
496        if num_rows < 2 {
497            return Ok(());
498        }
499
500        let Some(timestamps) = self.timestamps_native() else {
501            return Ok(());
502        };
503
504        // Fast path: check if there are any duplicate timestamps.
505        let mut has_dup = false;
506        let mut group_count = 1;
507        for i in 1..num_rows {
508            has_dup |= timestamps[i] == timestamps[i - 1];
509            group_count += (timestamps[i] != timestamps[i - 1]) as usize;
510        }
511        if !has_dup {
512            return Ok(());
513        }
514
515        let num_fields = self.fields.len();
516        let op_types = self.op_types.as_arrow().values();
517
518        let mut base_indices: Vec<u32> = Vec::with_capacity(group_count);
519        let mut field_indices: Vec<Vec<u32>> = (0..num_fields)
520            .map(|_| Vec::with_capacity(group_count))
521            .collect();
522
523        let mut start = 0;
524        while start < num_rows {
525            let ts = timestamps[start];
526            let mut end = start + 1;
527            while end < num_rows && timestamps[end] == ts {
528                end += 1;
529            }
530
531            let group_pos = base_indices.len();
532            base_indices.push(start as u32);
533
534            if num_fields > 0 {
535                // Default: take the base row for all fields.
536                for idx in &mut field_indices {
537                    idx.push(start as u32);
538                }
539
540                let base_deleted = op_types[start] == OpType::Delete as u8;
541                if !base_deleted {
542                    // Track fields that are null in the base row and try to fill them from older
543                    // rows in the same timestamp range.
544                    let mut missing_fields = Vec::new();
545                    for (field_idx, col) in self.fields.iter().enumerate() {
546                        if col.data.is_null(start) {
547                            missing_fields.push(field_idx);
548                        }
549                    }
550
551                    if !missing_fields.is_empty() {
552                        for row_idx in (start + 1)..end {
553                            if op_types[row_idx] == OpType::Delete as u8 {
554                                break;
555                            }
556
557                            missing_fields.retain(|&field_idx| {
558                                if self.fields[field_idx].data.is_null(row_idx) {
559                                    true
560                                } else {
561                                    field_indices[field_idx][group_pos] = row_idx as u32;
562                                    false
563                                }
564                            });
565
566                            if missing_fields.is_empty() {
567                                break;
568                            }
569                        }
570                    }
571                }
572            }
573
574            start = end;
575        }
576
577        let base_indices = UInt32Vector::from_vec(base_indices);
578        self.timestamps = self
579            .timestamps
580            .take(&base_indices)
581            .context(ComputeVectorSnafu)?;
582        let array = arrow::compute::take(self.sequences.as_arrow(), base_indices.as_arrow(), None)
583            .context(ComputeArrowSnafu)?;
584        // Safety: We know the array and vector type.
585        self.sequences = Arc::new(UInt64Vector::try_from_arrow_array(array).unwrap());
586        let array = arrow::compute::take(self.op_types.as_arrow(), base_indices.as_arrow(), None)
587            .context(ComputeArrowSnafu)?;
588        // Safety: We know the array and vector type.
589        self.op_types = Arc::new(UInt8Vector::try_from_arrow_array(array).unwrap());
590
591        for (field_idx, batch_column) in self.fields.iter_mut().enumerate() {
592            let idx = UInt32Vector::from_vec(std::mem::take(&mut field_indices[field_idx]));
593            batch_column.data = batch_column.data.take(&idx).context(ComputeVectorSnafu)?;
594        }
595
596        Ok(())
597    }
598
599    /// Returns the estimated memory size of the batch.
600    pub fn memory_size(&self) -> usize {
601        let mut size = std::mem::size_of::<Self>();
602        size += self.primary_key.len();
603        size += self.timestamps.memory_size();
604        size += self.sequences.memory_size();
605        size += self.op_types.memory_size();
606        for batch_column in &self.fields {
607            size += batch_column.data.memory_size();
608        }
609        size
610    }
611
612    /// Returns timestamps in a native slice or `None` if the batch is empty.
613    pub(crate) fn timestamps_native(&self) -> Option<&[i64]> {
614        if self.timestamps.is_empty() {
615            return None;
616        }
617
618        let values = match self.timestamps.data_type() {
619            ConcreteDataType::Timestamp(TimestampType::Second(_)) => self
620                .timestamps
621                .as_any()
622                .downcast_ref::<TimestampSecondVector>()
623                .unwrap()
624                .as_arrow()
625                .values(),
626            ConcreteDataType::Timestamp(TimestampType::Millisecond(_)) => self
627                .timestamps
628                .as_any()
629                .downcast_ref::<TimestampMillisecondVector>()
630                .unwrap()
631                .as_arrow()
632                .values(),
633            ConcreteDataType::Timestamp(TimestampType::Microsecond(_)) => self
634                .timestamps
635                .as_any()
636                .downcast_ref::<TimestampMicrosecondVector>()
637                .unwrap()
638                .as_arrow()
639                .values(),
640            ConcreteDataType::Timestamp(TimestampType::Nanosecond(_)) => self
641                .timestamps
642                .as_any()
643                .downcast_ref::<TimestampNanosecondVector>()
644                .unwrap()
645                .as_arrow()
646                .values(),
647            other => panic!("timestamps in a Batch has other type {:?}", other),
648        };
649
650        Some(values)
651    }
652
653    /// Takes the batch in place.
654    fn take_in_place(&mut self, indices: &UInt32Vector) -> Result<()> {
655        self.timestamps = self.timestamps.take(indices).context(ComputeVectorSnafu)?;
656        let array = arrow::compute::take(self.sequences.as_arrow(), indices.as_arrow(), None)
657            .context(ComputeArrowSnafu)?;
658        // Safety: we know the array and vector type.
659        self.sequences = Arc::new(UInt64Vector::try_from_arrow_array(array).unwrap());
660        let array = arrow::compute::take(self.op_types.as_arrow(), indices.as_arrow(), None)
661            .context(ComputeArrowSnafu)?;
662        self.op_types = Arc::new(UInt8Vector::try_from_arrow_array(array).unwrap());
663        for batch_column in &mut self.fields {
664            batch_column.data = batch_column
665                .data
666                .take(indices)
667                .context(ComputeVectorSnafu)?;
668        }
669
670        Ok(())
671    }
672
673    /// Gets a timestamp at given `index`.
674    ///
675    /// # Panics
676    /// Panics if `index` is out-of-bound or the timestamp vector returns null.
677    fn get_timestamp(&self, index: usize) -> Timestamp {
678        match self.timestamps.get_ref(index) {
679            ValueRef::Timestamp(timestamp) => timestamp,
680
681            // We have check the data type is timestamp compatible in the [BatchBuilder] so it's safe to panic.
682            value => panic!("{:?} is not a timestamp", value),
683        }
684    }
685
686    /// Gets a sequence at given `index`.
687    ///
688    /// # Panics
689    /// Panics if `index` is out-of-bound or the sequence vector returns null.
690    pub(crate) fn get_sequence(&self, index: usize) -> SequenceNumber {
691        // Safety: sequences is not null so it actually returns Some.
692        self.sequences.get_data(index).unwrap()
693    }
694
695    /// Checks the batch is monotonic by timestamps.
696    #[cfg(debug_assertions)]
697    #[allow(dead_code)]
698    pub(crate) fn check_monotonic(&self) -> Result<(), String> {
699        use std::cmp::Ordering;
700        if self.timestamps_native().is_none() {
701            return Ok(());
702        }
703
704        let timestamps = self.timestamps_native().unwrap();
705        let sequences = self.sequences.as_arrow().values();
706        for (i, window) in timestamps.windows(2).enumerate() {
707            let current = window[0];
708            let next = window[1];
709            let current_sequence = sequences[i];
710            let next_sequence = sequences[i + 1];
711            match current.cmp(&next) {
712                Ordering::Less => {
713                    // The current timestamp is less than the next timestamp.
714                    continue;
715                }
716                Ordering::Equal => {
717                    // The current timestamp is equal to the next timestamp.
718                    if current_sequence < next_sequence {
719                        return Err(format!(
720                            "sequence are not monotonic: ts {} == {} but current sequence {} < {}, index: {}",
721                            current, next, current_sequence, next_sequence, i
722                        ));
723                    }
724                }
725                Ordering::Greater => {
726                    // The current timestamp is greater than the next timestamp.
727                    return Err(format!(
728                        "timestamps are not monotonic: {} > {}, index: {}",
729                        current, next, i
730                    ));
731                }
732            }
733        }
734
735        Ok(())
736    }
737
738    /// Returns Ok if the given batch is behind the current batch.
739    #[cfg(debug_assertions)]
740    #[allow(dead_code)]
741    pub(crate) fn check_next_batch(&self, other: &Batch) -> Result<(), String> {
742        // Checks the primary key
743        if self.primary_key() < other.primary_key() {
744            return Ok(());
745        }
746        if self.primary_key() > other.primary_key() {
747            return Err(format!(
748                "primary key is not monotonic: {:?} > {:?}",
749                self.primary_key(),
750                other.primary_key()
751            ));
752        }
753        // Checks the timestamp.
754        if self.last_timestamp() < other.first_timestamp() {
755            return Ok(());
756        }
757        if self.last_timestamp() > other.first_timestamp() {
758            return Err(format!(
759                "timestamps are not monotonic: {:?} > {:?}",
760                self.last_timestamp(),
761                other.first_timestamp()
762            ));
763        }
764        // Checks the sequence.
765        if self.last_sequence() >= other.first_sequence() {
766            return Ok(());
767        }
768        Err(format!(
769            "sequences are not monotonic: {:?} < {:?}",
770            self.last_sequence(),
771            other.first_sequence()
772        ))
773    }
774
775    /// Returns the value of the column in the primary key.
776    ///
777    /// Lazily decodes the primary key and caches the result.
778    pub fn pk_col_value(
779        &mut self,
780        codec: &dyn PrimaryKeyCodec,
781        col_idx_in_pk: usize,
782        column_id: ColumnId,
783    ) -> Result<Option<&Value>> {
784        if self.pk_values.is_none() {
785            self.pk_values = Some(codec.decode(&self.primary_key).context(DecodeSnafu)?);
786        }
787
788        let pk_values = self.pk_values.as_ref().unwrap();
789        Ok(match pk_values {
790            CompositeValues::Dense(values) => values.get(col_idx_in_pk).map(|(_, v)| v),
791            CompositeValues::Sparse(values) => values.get(&column_id),
792        })
793    }
794
795    /// Returns values of the field in the batch.
796    ///
797    /// Lazily caches the field index.
798    pub fn field_col_value(&mut self, column_id: ColumnId) -> Option<&BatchColumn> {
799        if self.fields_idx.is_none() {
800            self.fields_idx = Some(
801                self.fields
802                    .iter()
803                    .enumerate()
804                    .map(|(i, c)| (c.column_id, i))
805                    .collect(),
806            );
807        }
808
809        self.fields_idx
810            .as_ref()
811            .unwrap()
812            .get(&column_id)
813            .map(|&idx| &self.fields[idx])
814    }
815}
816
817/// A struct to check the batch is monotonic.
818#[cfg(debug_assertions)]
819#[derive(Default)]
820#[allow(dead_code)]
821pub(crate) struct BatchChecker {
822    last_batch: Option<Batch>,
823    start: Option<Timestamp>,
824    end: Option<Timestamp>,
825}
826
827#[cfg(debug_assertions)]
828#[allow(dead_code)]
829impl BatchChecker {
830    /// Attaches the given start timestamp to the checker.
831    pub(crate) fn with_start(mut self, start: Option<Timestamp>) -> Self {
832        self.start = start;
833        self
834    }
835
836    /// Attaches the given end timestamp to the checker.
837    pub(crate) fn with_end(mut self, end: Option<Timestamp>) -> Self {
838        self.end = end;
839        self
840    }
841
842    /// Returns true if the given batch is monotonic and behind
843    /// the last batch.
844    pub(crate) fn check_monotonic(&mut self, batch: &Batch) -> Result<(), String> {
845        batch.check_monotonic()?;
846
847        if let (Some(start), Some(first)) = (self.start, batch.first_timestamp())
848            && start > first
849        {
850            return Err(format!(
851                "batch's first timestamp is before the start timestamp: {:?} > {:?}",
852                start, first
853            ));
854        }
855        if let (Some(end), Some(last)) = (self.end, batch.last_timestamp())
856            && end <= last
857        {
858            return Err(format!(
859                "batch's last timestamp is after the end timestamp: {:?} <= {:?}",
860                end, last
861            ));
862        }
863
864        // Checks the batch is behind the last batch.
865        // Then Updates the last batch.
866        let res = self
867            .last_batch
868            .as_ref()
869            .map(|last| last.check_next_batch(batch))
870            .unwrap_or(Ok(()));
871        self.last_batch = Some(batch.clone());
872        res
873    }
874
875    /// Formats current batch and last batch for debug.
876    pub(crate) fn format_batch(&self, batch: &Batch) -> String {
877        use std::fmt::Write;
878
879        let mut message = String::new();
880        if let Some(last) = &self.last_batch {
881            write!(
882                message,
883                "last_pk: {:?}, last_ts: {:?}, last_seq: {:?}, ",
884                last.primary_key(),
885                last.last_timestamp(),
886                last.last_sequence()
887            )
888            .unwrap();
889        }
890        write!(
891            message,
892            "batch_pk: {:?}, batch_ts: {:?}, batch_seq: {:?}",
893            batch.primary_key(),
894            batch.timestamps(),
895            batch.sequences()
896        )
897        .unwrap();
898
899        message
900    }
901
902    /// Checks batches from the part range are monotonic. Otherwise, panics.
903    pub(crate) fn ensure_part_range_batch(
904        &mut self,
905        scanner: &str,
906        region_id: store_api::storage::RegionId,
907        partition: usize,
908        part_range: store_api::region_engine::PartitionRange,
909        batch: &Batch,
910    ) {
911        if let Err(e) = self.check_monotonic(batch) {
912            let err_msg = format!(
913                "{}: batch is not sorted, {}, region_id: {}, partition: {}, part_range: {:?}",
914                scanner, e, region_id, partition, part_range,
915            );
916            common_telemetry::error!("{err_msg}, {}", self.format_batch(batch));
917            // Only print the number of row in the panic message.
918            panic!("{err_msg}, batch rows: {}", batch.num_rows());
919        }
920    }
921}
922
923/// Len of timestamp in arrow row format.
924const TIMESTAMP_KEY_LEN: usize = 9;
925
926/// Helper function to concat arrays from `iter`.
927fn concat_arrays(iter: impl Iterator<Item = ArrayRef>) -> Result<ArrayRef> {
928    let arrays: Vec<_> = iter.collect();
929    let dyn_arrays: Vec<_> = arrays.iter().map(|array| array.as_ref()).collect();
930    arrow::compute::concat(&dyn_arrays).context(ComputeArrowSnafu)
931}
932
933/// A column in a [Batch].
934#[derive(Debug, PartialEq, Eq, Clone)]
935pub struct BatchColumn {
936    /// Id of the column.
937    pub column_id: ColumnId,
938    /// Data of the column.
939    pub data: VectorRef,
940}
941
942/// Builder to build [Batch].
943pub struct BatchBuilder {
944    primary_key: Vec<u8>,
945    timestamps: Option<VectorRef>,
946    sequences: Option<Arc<UInt64Vector>>,
947    op_types: Option<Arc<UInt8Vector>>,
948    fields: Vec<BatchColumn>,
949}
950
951impl BatchBuilder {
952    /// Creates a new [BatchBuilder] with primary key.
953    pub fn new(primary_key: Vec<u8>) -> BatchBuilder {
954        BatchBuilder {
955            primary_key,
956            timestamps: None,
957            sequences: None,
958            op_types: None,
959            fields: Vec::new(),
960        }
961    }
962
963    /// Creates a new [BatchBuilder] with all required columns.
964    pub fn with_required_columns(
965        primary_key: Vec<u8>,
966        timestamps: VectorRef,
967        sequences: Arc<UInt64Vector>,
968        op_types: Arc<UInt8Vector>,
969    ) -> BatchBuilder {
970        BatchBuilder {
971            primary_key,
972            timestamps: Some(timestamps),
973            sequences: Some(sequences),
974            op_types: Some(op_types),
975            fields: Vec::new(),
976        }
977    }
978
979    /// Set all field columns.
980    pub fn with_fields(mut self, fields: Vec<BatchColumn>) -> Self {
981        self.fields = fields;
982        self
983    }
984
985    /// Push a field column.
986    pub fn push_field(&mut self, column: BatchColumn) -> &mut Self {
987        self.fields.push(column);
988        self
989    }
990
991    /// Push an array as a field.
992    pub fn push_field_array(&mut self, column_id: ColumnId, array: ArrayRef) -> Result<&mut Self> {
993        let vector = Helper::try_into_vector(array).context(ConvertVectorSnafu)?;
994        self.fields.push(BatchColumn {
995            column_id,
996            data: vector,
997        });
998
999        Ok(self)
1000    }
1001
1002    /// Try to set an array as timestamps.
1003    pub fn timestamps_array(&mut self, array: ArrayRef) -> Result<&mut Self> {
1004        let vector = Helper::try_into_vector(array).context(ConvertVectorSnafu)?;
1005        ensure!(
1006            vector.data_type().is_timestamp(),
1007            InvalidBatchSnafu {
1008                reason: format!("{:?} is not a timestamp type", vector.data_type()),
1009            }
1010        );
1011
1012        self.timestamps = Some(vector);
1013        Ok(self)
1014    }
1015
1016    /// Try to set an array as sequences.
1017    pub fn sequences_array(&mut self, array: ArrayRef) -> Result<&mut Self> {
1018        ensure!(
1019            *array.data_type() == arrow::datatypes::DataType::UInt64,
1020            InvalidBatchSnafu {
1021                reason: "sequence array is not UInt64 type",
1022            }
1023        );
1024        // Safety: The cast must success as we have ensured it is uint64 type.
1025        let vector = Arc::new(UInt64Vector::try_from_arrow_array(array).unwrap());
1026        self.sequences = Some(vector);
1027
1028        Ok(self)
1029    }
1030
1031    /// Try to set an array as op types.
1032    pub fn op_types_array(&mut self, array: ArrayRef) -> Result<&mut Self> {
1033        ensure!(
1034            *array.data_type() == arrow::datatypes::DataType::UInt8,
1035            InvalidBatchSnafu {
1036                reason: "sequence array is not UInt8 type",
1037            }
1038        );
1039        // Safety: The cast must success as we have ensured it is uint64 type.
1040        let vector = Arc::new(UInt8Vector::try_from_arrow_array(array).unwrap());
1041        self.op_types = Some(vector);
1042
1043        Ok(self)
1044    }
1045
1046    /// Builds the [Batch].
1047    pub fn build(self) -> Result<Batch> {
1048        let timestamps = self.timestamps.context(InvalidBatchSnafu {
1049            reason: "missing timestamps",
1050        })?;
1051        let sequences = self.sequences.context(InvalidBatchSnafu {
1052            reason: "missing sequences",
1053        })?;
1054        let op_types = self.op_types.context(InvalidBatchSnafu {
1055            reason: "missing op_types",
1056        })?;
1057        // Our storage format ensure these columns are not nullable so
1058        // we use assert here.
1059        assert_eq!(0, timestamps.null_count());
1060        assert_eq!(0, sequences.null_count());
1061        assert_eq!(0, op_types.null_count());
1062
1063        let ts_len = timestamps.len();
1064        ensure!(
1065            sequences.len() == ts_len,
1066            InvalidBatchSnafu {
1067                reason: format!(
1068                    "sequence have different len {} != {}",
1069                    sequences.len(),
1070                    ts_len
1071                ),
1072            }
1073        );
1074        ensure!(
1075            op_types.len() == ts_len,
1076            InvalidBatchSnafu {
1077                reason: format!(
1078                    "op type have different len {} != {}",
1079                    op_types.len(),
1080                    ts_len
1081                ),
1082            }
1083        );
1084        for column in &self.fields {
1085            ensure!(
1086                column.data.len() == ts_len,
1087                InvalidBatchSnafu {
1088                    reason: format!(
1089                        "column {} has different len {} != {}",
1090                        column.column_id,
1091                        column.data.len(),
1092                        ts_len
1093                    ),
1094                }
1095            );
1096        }
1097
1098        Ok(Batch {
1099            primary_key: self.primary_key,
1100            pk_values: None,
1101            timestamps,
1102            sequences,
1103            op_types,
1104            fields: self.fields,
1105            fields_idx: None,
1106        })
1107    }
1108}
1109
1110impl From<Batch> for BatchBuilder {
1111    fn from(batch: Batch) -> Self {
1112        Self {
1113            primary_key: batch.primary_key,
1114            timestamps: Some(batch.timestamps),
1115            sequences: Some(batch.sequences),
1116            op_types: Some(batch.op_types),
1117            fields: batch.fields,
1118        }
1119    }
1120}
1121
1122/// Async [Batch] reader and iterator wrapper.
1123///
1124/// This is the data source for SST writers or internal readers.
1125pub enum Source {
1126    /// Source from a [BoxedBatchReader].
1127    Reader(BoxedBatchReader),
1128    /// Source from a [BoxedBatchIterator].
1129    Iter(BoxedBatchIterator),
1130    /// Source from a [BoxedBatchStream].
1131    Stream(BoxedBatchStream),
1132}
1133
1134impl Source {
1135    /// Returns next [Batch] from this data source.
1136    pub async fn next_batch(&mut self) -> Result<Option<Batch>> {
1137        match self {
1138            Source::Reader(reader) => reader.next_batch().await,
1139            Source::Iter(iter) => iter.next().transpose(),
1140            Source::Stream(stream) => stream.try_next().await,
1141        }
1142    }
1143}
1144
1145/// Async [RecordBatch] reader and iterator wrapper for flat format.
1146pub struct FlatSource {
1147    schema: SchemaRef,
1148    inner: FlatSourceInner,
1149}
1150
1151impl FlatSource {
1152    /// Create a [FlatSource] from a [BoxedRecordBatchIterator] and its schema.
1153    pub fn new_iter(schema: SchemaRef, iter: BoxedRecordBatchIterator) -> Self {
1154        Self {
1155            schema,
1156            inner: FlatSourceInner::Iter(iter),
1157        }
1158    }
1159
1160    /// Create a [FlatSource] from a [BoxedRecordBatchStream] and its schema.
1161    pub fn new_stream(schema: SchemaRef, stream: BoxedRecordBatchStream) -> Self {
1162        Self {
1163            schema,
1164            inner: FlatSourceInner::Stream(stream),
1165        }
1166    }
1167
1168    pub(crate) fn schema(&self) -> &SchemaRef {
1169        &self.schema
1170    }
1171
1172    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
1173        self.inner.next_batch().await
1174    }
1175
1176    #[cfg(test)]
1177    pub(crate) fn take_iter(self) -> BoxedRecordBatchIterator {
1178        match self.inner {
1179            FlatSourceInner::Iter(iter) => iter,
1180            FlatSourceInner::Stream(_) => unreachable!(),
1181        }
1182    }
1183}
1184
1185enum FlatSourceInner {
1186    /// Source from a [BoxedRecordBatchIterator].
1187    Iter(BoxedRecordBatchIterator),
1188    /// Source from a [BoxedRecordBatchStream].
1189    Stream(BoxedRecordBatchStream),
1190}
1191
1192impl FlatSourceInner {
1193    /// Returns next [RecordBatch] from this data source.
1194    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
1195        match self {
1196            Self::Iter(iter) => iter.next().transpose(),
1197            Self::Stream(stream) => stream.try_next().await,
1198        }
1199    }
1200}
1201
1202/// Async batch reader.
1203///
1204/// The reader must guarantee [Batch]es returned by it have the same schema.
1205#[async_trait]
1206pub trait BatchReader: Send {
1207    /// Fetch next [Batch].
1208    ///
1209    /// Returns `Ok(None)` when the reader has reached its end and calling `next_batch()`
1210    /// again won't return batch again.
1211    ///
1212    /// If `Err` is returned, caller should not call this method again, the implementor
1213    /// may or may not panic in such case.
1214    async fn next_batch(&mut self) -> Result<Option<Batch>>;
1215}
1216
1217/// Pointer to [BatchReader].
1218pub type BoxedBatchReader = Box<dyn BatchReader>;
1219
1220/// Pointer to a stream that yields [Batch].
1221pub type BoxedBatchStream = BoxStream<'static, Result<Batch>>;
1222
1223/// Pointer to a stream that yields [RecordBatch].
1224pub type BoxedRecordBatchStream = BoxStream<'static, Result<RecordBatch>>;
1225
1226#[async_trait::async_trait]
1227impl<T: BatchReader + ?Sized> BatchReader for Box<T> {
1228    async fn next_batch(&mut self) -> Result<Option<Batch>> {
1229        (**self).next_batch().await
1230    }
1231}
1232
1233/// Local metrics for scanners.
1234#[derive(Debug, Default)]
1235pub(crate) struct ScannerMetrics {
1236    /// Duration to scan data.
1237    scan_cost: Duration,
1238    /// Duration while waiting for `yield`.
1239    yield_cost: Duration,
1240    /// Number of batches returned.
1241    num_batches: usize,
1242    /// Number of rows returned.
1243    num_rows: usize,
1244}
1245
1246#[cfg(test)]
1247mod tests {
1248    use datatypes::arrow::array::{TimestampMillisecondArray, UInt8Array, UInt64Array};
1249    use mito_codec::row_converter::{self, build_primary_key_codec_with_fields};
1250    use store_api::codec::PrimaryKeyEncoding;
1251    use store_api::storage::consts::ReservedColumnId;
1252
1253    use super::*;
1254    use crate::error::Error;
1255    use crate::test_util::new_batch_builder;
1256
1257    fn new_batch(
1258        timestamps: &[i64],
1259        sequences: &[u64],
1260        op_types: &[OpType],
1261        field: &[u64],
1262    ) -> Batch {
1263        new_batch_builder(b"test", timestamps, sequences, op_types, 1, field)
1264            .build()
1265            .unwrap()
1266    }
1267
1268    fn new_batch_with_u64_fields(
1269        timestamps: &[i64],
1270        sequences: &[u64],
1271        op_types: &[OpType],
1272        fields: &[(ColumnId, &[Option<u64>])],
1273    ) -> Batch {
1274        assert_eq!(timestamps.len(), sequences.len());
1275        assert_eq!(timestamps.len(), op_types.len());
1276        for (_, values) in fields {
1277            assert_eq!(timestamps.len(), values.len());
1278        }
1279
1280        let mut builder = BatchBuilder::new(b"test".to_vec());
1281        builder
1282            .timestamps_array(Arc::new(TimestampMillisecondArray::from_iter_values(
1283                timestamps.iter().copied(),
1284            )))
1285            .unwrap()
1286            .sequences_array(Arc::new(UInt64Array::from_iter_values(
1287                sequences.iter().copied(),
1288            )))
1289            .unwrap()
1290            .op_types_array(Arc::new(UInt8Array::from_iter_values(
1291                op_types.iter().map(|v| *v as u8),
1292            )))
1293            .unwrap();
1294
1295        for (col_id, values) in fields {
1296            builder
1297                .push_field_array(*col_id, Arc::new(UInt64Array::from(values.to_vec())))
1298                .unwrap();
1299        }
1300
1301        builder.build().unwrap()
1302    }
1303
1304    fn new_batch_without_fields(
1305        timestamps: &[i64],
1306        sequences: &[u64],
1307        op_types: &[OpType],
1308    ) -> Batch {
1309        assert_eq!(timestamps.len(), sequences.len());
1310        assert_eq!(timestamps.len(), op_types.len());
1311
1312        let mut builder = BatchBuilder::new(b"test".to_vec());
1313        builder
1314            .timestamps_array(Arc::new(TimestampMillisecondArray::from_iter_values(
1315                timestamps.iter().copied(),
1316            )))
1317            .unwrap()
1318            .sequences_array(Arc::new(UInt64Array::from_iter_values(
1319                sequences.iter().copied(),
1320            )))
1321            .unwrap()
1322            .op_types_array(Arc::new(UInt8Array::from_iter_values(
1323                op_types.iter().map(|v| *v as u8),
1324            )))
1325            .unwrap();
1326
1327        builder.build().unwrap()
1328    }
1329
1330    #[test]
1331    fn test_empty_batch() {
1332        let batch = Batch::empty();
1333        assert!(batch.is_empty());
1334        assert_eq!(None, batch.first_timestamp());
1335        assert_eq!(None, batch.last_timestamp());
1336        assert_eq!(None, batch.first_sequence());
1337        assert_eq!(None, batch.last_sequence());
1338        assert!(batch.timestamps_native().is_none());
1339    }
1340
1341    #[test]
1342    fn test_first_last_one() {
1343        let batch = new_batch(&[1], &[2], &[OpType::Put], &[4]);
1344        assert_eq!(
1345            Timestamp::new_millisecond(1),
1346            batch.first_timestamp().unwrap()
1347        );
1348        assert_eq!(
1349            Timestamp::new_millisecond(1),
1350            batch.last_timestamp().unwrap()
1351        );
1352        assert_eq!(2, batch.first_sequence().unwrap());
1353        assert_eq!(2, batch.last_sequence().unwrap());
1354    }
1355
1356    #[test]
1357    fn test_first_last_multiple() {
1358        let batch = new_batch(
1359            &[1, 2, 3],
1360            &[11, 12, 13],
1361            &[OpType::Put, OpType::Put, OpType::Put],
1362            &[21, 22, 23],
1363        );
1364        assert_eq!(
1365            Timestamp::new_millisecond(1),
1366            batch.first_timestamp().unwrap()
1367        );
1368        assert_eq!(
1369            Timestamp::new_millisecond(3),
1370            batch.last_timestamp().unwrap()
1371        );
1372        assert_eq!(11, batch.first_sequence().unwrap());
1373        assert_eq!(13, batch.last_sequence().unwrap());
1374    }
1375
1376    #[test]
1377    fn test_slice() {
1378        let batch = new_batch(
1379            &[1, 2, 3, 4],
1380            &[11, 12, 13, 14],
1381            &[OpType::Put, OpType::Delete, OpType::Put, OpType::Put],
1382            &[21, 22, 23, 24],
1383        );
1384        let batch = batch.slice(1, 2);
1385        let expect = new_batch(
1386            &[2, 3],
1387            &[12, 13],
1388            &[OpType::Delete, OpType::Put],
1389            &[22, 23],
1390        );
1391        assert_eq!(expect, batch);
1392    }
1393
1394    #[test]
1395    fn test_timestamps_native() {
1396        let batch = new_batch(
1397            &[1, 2, 3, 4],
1398            &[11, 12, 13, 14],
1399            &[OpType::Put, OpType::Delete, OpType::Put, OpType::Put],
1400            &[21, 22, 23, 24],
1401        );
1402        assert_eq!(&[1, 2, 3, 4], batch.timestamps_native().unwrap());
1403    }
1404
1405    #[test]
1406    fn test_concat_empty() {
1407        let err = Batch::concat(vec![]).unwrap_err();
1408        assert!(
1409            matches!(err, Error::InvalidBatch { .. }),
1410            "unexpected err: {err}"
1411        );
1412    }
1413
1414    #[test]
1415    fn test_concat_one() {
1416        let batch = new_batch(&[], &[], &[], &[]);
1417        let actual = Batch::concat(vec![batch.clone()]).unwrap();
1418        assert_eq!(batch, actual);
1419
1420        let batch = new_batch(&[1, 2], &[11, 12], &[OpType::Put, OpType::Put], &[21, 22]);
1421        let actual = Batch::concat(vec![batch.clone()]).unwrap();
1422        assert_eq!(batch, actual);
1423    }
1424
1425    #[test]
1426    fn test_concat_multiple() {
1427        let batches = vec![
1428            new_batch(&[1, 2], &[11, 12], &[OpType::Put, OpType::Put], &[21, 22]),
1429            new_batch(
1430                &[3, 4, 5],
1431                &[13, 14, 15],
1432                &[OpType::Put, OpType::Delete, OpType::Put],
1433                &[23, 24, 25],
1434            ),
1435            new_batch(&[], &[], &[], &[]),
1436            new_batch(&[6], &[16], &[OpType::Put], &[26]),
1437        ];
1438        let batch = Batch::concat(batches).unwrap();
1439        let expect = new_batch(
1440            &[1, 2, 3, 4, 5, 6],
1441            &[11, 12, 13, 14, 15, 16],
1442            &[
1443                OpType::Put,
1444                OpType::Put,
1445                OpType::Put,
1446                OpType::Delete,
1447                OpType::Put,
1448                OpType::Put,
1449            ],
1450            &[21, 22, 23, 24, 25, 26],
1451        );
1452        assert_eq!(expect, batch);
1453    }
1454
1455    #[test]
1456    fn test_concat_different() {
1457        let batch1 = new_batch(&[1], &[1], &[OpType::Put], &[1]);
1458        let mut batch2 = new_batch(&[2], &[2], &[OpType::Put], &[2]);
1459        batch2.primary_key = b"hello".to_vec();
1460        let err = Batch::concat(vec![batch1, batch2]).unwrap_err();
1461        assert!(
1462            matches!(err, Error::InvalidBatch { .. }),
1463            "unexpected err: {err}"
1464        );
1465    }
1466
1467    #[test]
1468    fn test_concat_different_fields() {
1469        let batch1 = new_batch(&[1], &[1], &[OpType::Put], &[1]);
1470        let fields = vec![
1471            batch1.fields()[0].clone(),
1472            BatchColumn {
1473                column_id: 2,
1474                data: Arc::new(UInt64Vector::from_slice([2])),
1475            },
1476        ];
1477        // Batch 2 has more fields.
1478        let batch2 = batch1.clone().with_fields(fields).unwrap();
1479        let err = Batch::concat(vec![batch1.clone(), batch2]).unwrap_err();
1480        assert!(
1481            matches!(err, Error::InvalidBatch { .. }),
1482            "unexpected err: {err}"
1483        );
1484
1485        // Batch 2 has different field.
1486        let fields = vec![BatchColumn {
1487            column_id: 2,
1488            data: Arc::new(UInt64Vector::from_slice([2])),
1489        }];
1490        let batch2 = batch1.clone().with_fields(fields).unwrap();
1491        let err = Batch::concat(vec![batch1, batch2]).unwrap_err();
1492        assert!(
1493            matches!(err, Error::InvalidBatch { .. }),
1494            "unexpected err: {err}"
1495        );
1496    }
1497
1498    #[test]
1499    fn test_filter_deleted_empty() {
1500        let mut batch = new_batch(&[], &[], &[], &[]);
1501        batch.filter_deleted().unwrap();
1502        assert!(batch.is_empty());
1503    }
1504
1505    #[test]
1506    fn test_filter_deleted() {
1507        let mut batch = new_batch(
1508            &[1, 2, 3, 4],
1509            &[11, 12, 13, 14],
1510            &[OpType::Delete, OpType::Put, OpType::Delete, OpType::Put],
1511            &[21, 22, 23, 24],
1512        );
1513        batch.filter_deleted().unwrap();
1514        let expect = new_batch(&[2, 4], &[12, 14], &[OpType::Put, OpType::Put], &[22, 24]);
1515        assert_eq!(expect, batch);
1516
1517        let mut batch = new_batch(
1518            &[1, 2, 3, 4],
1519            &[11, 12, 13, 14],
1520            &[OpType::Put, OpType::Put, OpType::Put, OpType::Put],
1521            &[21, 22, 23, 24],
1522        );
1523        let expect = batch.clone();
1524        batch.filter_deleted().unwrap();
1525        assert_eq!(expect, batch);
1526    }
1527
1528    #[test]
1529    fn test_filter_by_sequence() {
1530        // Filters put only.
1531        let mut batch = new_batch(
1532            &[1, 2, 3, 4],
1533            &[11, 12, 13, 14],
1534            &[OpType::Put, OpType::Put, OpType::Put, OpType::Put],
1535            &[21, 22, 23, 24],
1536        );
1537        batch
1538            .filter_by_sequence(Some(SequenceRange::LtEq { max: 13 }))
1539            .unwrap();
1540        let expect = new_batch(
1541            &[1, 2, 3],
1542            &[11, 12, 13],
1543            &[OpType::Put, OpType::Put, OpType::Put],
1544            &[21, 22, 23],
1545        );
1546        assert_eq!(expect, batch);
1547
1548        // Filters to empty.
1549        let mut batch = new_batch(
1550            &[1, 2, 3, 4],
1551            &[11, 12, 13, 14],
1552            &[OpType::Put, OpType::Delete, OpType::Put, OpType::Put],
1553            &[21, 22, 23, 24],
1554        );
1555
1556        batch
1557            .filter_by_sequence(Some(SequenceRange::LtEq { max: 10 }))
1558            .unwrap();
1559        assert!(batch.is_empty());
1560
1561        // None filter.
1562        let mut batch = new_batch(
1563            &[1, 2, 3, 4],
1564            &[11, 12, 13, 14],
1565            &[OpType::Put, OpType::Delete, OpType::Put, OpType::Put],
1566            &[21, 22, 23, 24],
1567        );
1568        let expect = batch.clone();
1569        batch.filter_by_sequence(None).unwrap();
1570        assert_eq!(expect, batch);
1571
1572        // Filter a empty batch
1573        let mut batch = new_batch(&[], &[], &[], &[]);
1574        batch
1575            .filter_by_sequence(Some(SequenceRange::LtEq { max: 10 }))
1576            .unwrap();
1577        assert!(batch.is_empty());
1578
1579        // Filter a empty batch with None
1580        let mut batch = new_batch(&[], &[], &[], &[]);
1581        batch.filter_by_sequence(None).unwrap();
1582        assert!(batch.is_empty());
1583
1584        // Test From variant - exclusive lower bound
1585        let mut batch = new_batch(
1586            &[1, 2, 3, 4],
1587            &[11, 12, 13, 14],
1588            &[OpType::Put, OpType::Put, OpType::Put, OpType::Put],
1589            &[21, 22, 23, 24],
1590        );
1591        batch
1592            .filter_by_sequence(Some(SequenceRange::Gt { min: 12 }))
1593            .unwrap();
1594        let expect = new_batch(&[3, 4], &[13, 14], &[OpType::Put, OpType::Put], &[23, 24]);
1595        assert_eq!(expect, batch);
1596
1597        // Test From variant with no matches
1598        let mut batch = new_batch(
1599            &[1, 2, 3, 4],
1600            &[11, 12, 13, 14],
1601            &[OpType::Put, OpType::Delete, OpType::Put, OpType::Put],
1602            &[21, 22, 23, 24],
1603        );
1604        batch
1605            .filter_by_sequence(Some(SequenceRange::Gt { min: 20 }))
1606            .unwrap();
1607        assert!(batch.is_empty());
1608
1609        // Test Range variant - exclusive lower bound, inclusive upper bound
1610        let mut batch = new_batch(
1611            &[1, 2, 3, 4, 5],
1612            &[11, 12, 13, 14, 15],
1613            &[
1614                OpType::Put,
1615                OpType::Put,
1616                OpType::Put,
1617                OpType::Put,
1618                OpType::Put,
1619            ],
1620            &[21, 22, 23, 24, 25],
1621        );
1622        batch
1623            .filter_by_sequence(Some(SequenceRange::GtLtEq { min: 12, max: 14 }))
1624            .unwrap();
1625        let expect = new_batch(&[3, 4], &[13, 14], &[OpType::Put, OpType::Put], &[23, 24]);
1626        assert_eq!(expect, batch);
1627
1628        // Test Range variant with mixed operations
1629        let mut batch = new_batch(
1630            &[1, 2, 3, 4, 5],
1631            &[11, 12, 13, 14, 15],
1632            &[
1633                OpType::Put,
1634                OpType::Delete,
1635                OpType::Put,
1636                OpType::Delete,
1637                OpType::Put,
1638            ],
1639            &[21, 22, 23, 24, 25],
1640        );
1641        batch
1642            .filter_by_sequence(Some(SequenceRange::GtLtEq { min: 11, max: 13 }))
1643            .unwrap();
1644        let expect = new_batch(
1645            &[2, 3],
1646            &[12, 13],
1647            &[OpType::Delete, OpType::Put],
1648            &[22, 23],
1649        );
1650        assert_eq!(expect, batch);
1651
1652        // Test Range variant with no matches
1653        let mut batch = new_batch(
1654            &[1, 2, 3, 4],
1655            &[11, 12, 13, 14],
1656            &[OpType::Put, OpType::Put, OpType::Put, OpType::Put],
1657            &[21, 22, 23, 24],
1658        );
1659        batch
1660            .filter_by_sequence(Some(SequenceRange::GtLtEq { min: 20, max: 25 }))
1661            .unwrap();
1662        assert!(batch.is_empty());
1663    }
1664
1665    #[test]
1666    fn test_merge_last_non_null_no_dup() {
1667        let mut batch = new_batch_with_u64_fields(
1668            &[1, 2],
1669            &[2, 1],
1670            &[OpType::Put, OpType::Put],
1671            &[(1, &[Some(10), None]), (2, &[Some(100), Some(200)])],
1672        );
1673        let expect = batch.clone();
1674        batch.merge_last_non_null().unwrap();
1675        assert_eq!(expect, batch);
1676    }
1677
1678    #[test]
1679    fn test_merge_last_non_null_fill_null_fields() {
1680        // Rows are already sorted by timestamp asc and sequence desc.
1681        let mut batch = new_batch_with_u64_fields(
1682            &[1, 1, 1],
1683            &[3, 2, 1],
1684            &[OpType::Put, OpType::Put, OpType::Put],
1685            &[
1686                (1, &[None, Some(10), Some(11)]),
1687                (2, &[Some(100), Some(200), Some(300)]),
1688            ],
1689        );
1690        batch.merge_last_non_null().unwrap();
1691
1692        // Field 1 is filled from the first older row (seq=2). Field 2 keeps the base value.
1693        // Filled fields must not be overwritten by even older duplicates.
1694        let expect = new_batch_with_u64_fields(
1695            &[1],
1696            &[3],
1697            &[OpType::Put],
1698            &[(1, &[Some(10)]), (2, &[Some(100)])],
1699        );
1700        assert_eq!(expect, batch);
1701    }
1702
1703    #[test]
1704    fn test_merge_last_non_null_stop_at_delete_row() {
1705        // A delete row in older duplicates should stop filling to avoid resurrecting values before
1706        // deletion.
1707        let mut batch = new_batch_with_u64_fields(
1708            &[1, 1, 1],
1709            &[3, 2, 1],
1710            &[OpType::Put, OpType::Delete, OpType::Put],
1711            &[
1712                (1, &[None, Some(10), Some(11)]),
1713                (2, &[Some(100), Some(200), Some(300)]),
1714            ],
1715        );
1716        batch.merge_last_non_null().unwrap();
1717
1718        let expect = new_batch_with_u64_fields(
1719            &[1],
1720            &[3],
1721            &[OpType::Put],
1722            &[(1, &[None]), (2, &[Some(100)])],
1723        );
1724        assert_eq!(expect, batch);
1725    }
1726
1727    #[test]
1728    fn test_merge_last_non_null_base_delete_no_merge() {
1729        let mut batch = new_batch_with_u64_fields(
1730            &[1, 1],
1731            &[3, 2],
1732            &[OpType::Delete, OpType::Put],
1733            &[(1, &[None, Some(10)]), (2, &[None, Some(200)])],
1734        );
1735        batch.merge_last_non_null().unwrap();
1736
1737        // Base row is delete, keep it as is and don't merge fields from older rows.
1738        let expect =
1739            new_batch_with_u64_fields(&[1], &[3], &[OpType::Delete], &[(1, &[None]), (2, &[None])]);
1740        assert_eq!(expect, batch);
1741    }
1742
1743    #[test]
1744    fn test_merge_last_non_null_multiple_timestamp_groups() {
1745        let mut batch = new_batch_with_u64_fields(
1746            &[1, 1, 2, 3, 3],
1747            &[5, 4, 3, 2, 1],
1748            &[
1749                OpType::Put,
1750                OpType::Put,
1751                OpType::Put,
1752                OpType::Put,
1753                OpType::Put,
1754            ],
1755            &[
1756                (1, &[None, Some(10), Some(20), None, Some(30)]),
1757                (2, &[Some(100), Some(110), Some(120), None, Some(130)]),
1758            ],
1759        );
1760        batch.merge_last_non_null().unwrap();
1761
1762        let expect = new_batch_with_u64_fields(
1763            &[1, 2, 3],
1764            &[5, 3, 2],
1765            &[OpType::Put, OpType::Put, OpType::Put],
1766            &[
1767                (1, &[Some(10), Some(20), Some(30)]),
1768                (2, &[Some(100), Some(120), Some(130)]),
1769            ],
1770        );
1771        assert_eq!(expect, batch);
1772    }
1773
1774    #[test]
1775    fn test_merge_last_non_null_no_fields() {
1776        let mut batch = new_batch_without_fields(
1777            &[1, 1, 2],
1778            &[3, 2, 1],
1779            &[OpType::Put, OpType::Put, OpType::Put],
1780        );
1781        batch.merge_last_non_null().unwrap();
1782
1783        let expect = new_batch_without_fields(&[1, 2], &[3, 1], &[OpType::Put, OpType::Put]);
1784        assert_eq!(expect, batch);
1785    }
1786
1787    #[test]
1788    fn test_filter() {
1789        // Filters put only.
1790        let mut batch = new_batch(
1791            &[1, 2, 3, 4],
1792            &[11, 12, 13, 14],
1793            &[OpType::Put, OpType::Put, OpType::Put, OpType::Put],
1794            &[21, 22, 23, 24],
1795        );
1796        let predicate = BooleanVector::from_vec(vec![false, false, true, true]);
1797        batch.filter(&predicate).unwrap();
1798        let expect = new_batch(&[3, 4], &[13, 14], &[OpType::Put, OpType::Put], &[23, 24]);
1799        assert_eq!(expect, batch);
1800
1801        // Filters deletion.
1802        let mut batch = new_batch(
1803            &[1, 2, 3, 4],
1804            &[11, 12, 13, 14],
1805            &[OpType::Put, OpType::Delete, OpType::Put, OpType::Put],
1806            &[21, 22, 23, 24],
1807        );
1808        let predicate = BooleanVector::from_vec(vec![false, false, true, true]);
1809        batch.filter(&predicate).unwrap();
1810        let expect = new_batch(&[3, 4], &[13, 14], &[OpType::Put, OpType::Put], &[23, 24]);
1811        assert_eq!(expect, batch);
1812
1813        // Filters to empty.
1814        let predicate = BooleanVector::from_vec(vec![false, false]);
1815        batch.filter(&predicate).unwrap();
1816        assert!(batch.is_empty());
1817    }
1818
1819    #[test]
1820    fn test_sort_and_dedup() {
1821        let original = new_batch(
1822            &[2, 3, 1, 4, 5, 2],
1823            &[1, 2, 3, 4, 5, 6],
1824            &[
1825                OpType::Put,
1826                OpType::Put,
1827                OpType::Put,
1828                OpType::Put,
1829                OpType::Put,
1830                OpType::Put,
1831            ],
1832            &[21, 22, 23, 24, 25, 26],
1833        );
1834
1835        let mut batch = original.clone();
1836        batch.sort(true).unwrap();
1837        // It should only keep one timestamp 2.
1838        assert_eq!(
1839            new_batch(
1840                &[1, 2, 3, 4, 5],
1841                &[3, 6, 2, 4, 5],
1842                &[
1843                    OpType::Put,
1844                    OpType::Put,
1845                    OpType::Put,
1846                    OpType::Put,
1847                    OpType::Put,
1848                ],
1849                &[23, 26, 22, 24, 25],
1850            ),
1851            batch
1852        );
1853
1854        let mut batch = original.clone();
1855        batch.sort(false).unwrap();
1856
1857        // It should only keep one timestamp 2.
1858        assert_eq!(
1859            new_batch(
1860                &[1, 2, 2, 3, 4, 5],
1861                &[3, 6, 1, 2, 4, 5],
1862                &[
1863                    OpType::Put,
1864                    OpType::Put,
1865                    OpType::Put,
1866                    OpType::Put,
1867                    OpType::Put,
1868                    OpType::Put,
1869                ],
1870                &[23, 26, 21, 22, 24, 25],
1871            ),
1872            batch
1873        );
1874
1875        let original = new_batch(
1876            &[2, 2, 1],
1877            &[1, 6, 1],
1878            &[OpType::Delete, OpType::Put, OpType::Put],
1879            &[21, 22, 23],
1880        );
1881
1882        let mut batch = original.clone();
1883        batch.sort(true).unwrap();
1884        let expect = new_batch(&[1, 2], &[1, 6], &[OpType::Put, OpType::Put], &[23, 22]);
1885        assert_eq!(expect, batch);
1886
1887        let mut batch = original.clone();
1888        batch.sort(false).unwrap();
1889        let expect = new_batch(
1890            &[1, 2, 2],
1891            &[1, 6, 1],
1892            &[OpType::Put, OpType::Put, OpType::Delete],
1893            &[23, 22, 21],
1894        );
1895        assert_eq!(expect, batch);
1896    }
1897
1898    #[test]
1899    fn test_get_value() {
1900        let encodings = [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse];
1901
1902        for encoding in encodings {
1903            let codec = build_primary_key_codec_with_fields(
1904                encoding,
1905                [
1906                    (
1907                        ReservedColumnId::table_id(),
1908                        row_converter::SortField::new(ConcreteDataType::uint32_datatype()),
1909                    ),
1910                    (
1911                        ReservedColumnId::tsid(),
1912                        row_converter::SortField::new(ConcreteDataType::uint64_datatype()),
1913                    ),
1914                    (
1915                        100,
1916                        row_converter::SortField::new(ConcreteDataType::string_datatype()),
1917                    ),
1918                    (
1919                        200,
1920                        row_converter::SortField::new(ConcreteDataType::string_datatype()),
1921                    ),
1922                ]
1923                .into_iter(),
1924            );
1925
1926            let values = [
1927                Value::UInt32(1000),
1928                Value::UInt64(2000),
1929                Value::String("abcdefgh".into()),
1930                Value::String("zyxwvu".into()),
1931            ];
1932            let mut buf = vec![];
1933            codec
1934                .encode_values(
1935                    &[
1936                        (ReservedColumnId::table_id(), values[0].clone()),
1937                        (ReservedColumnId::tsid(), values[1].clone()),
1938                        (100, values[2].clone()),
1939                        (200, values[3].clone()),
1940                    ],
1941                    &mut buf,
1942                )
1943                .unwrap();
1944
1945            let field_col_id = 2;
1946            let mut batch = new_batch_builder(
1947                &buf,
1948                &[1, 2, 3],
1949                &[1, 1, 1],
1950                &[OpType::Put, OpType::Put, OpType::Put],
1951                field_col_id,
1952                &[42, 43, 44],
1953            )
1954            .build()
1955            .unwrap();
1956
1957            let v = batch
1958                .pk_col_value(&*codec, 0, ReservedColumnId::table_id())
1959                .unwrap()
1960                .unwrap();
1961            assert_eq!(values[0], *v);
1962
1963            let v = batch
1964                .pk_col_value(&*codec, 1, ReservedColumnId::tsid())
1965                .unwrap()
1966                .unwrap();
1967            assert_eq!(values[1], *v);
1968
1969            let v = batch.pk_col_value(&*codec, 2, 100).unwrap().unwrap();
1970            assert_eq!(values[2], *v);
1971
1972            let v = batch.pk_col_value(&*codec, 3, 200).unwrap().unwrap();
1973            assert_eq!(values[3], *v);
1974
1975            let v = batch.field_col_value(field_col_id).unwrap();
1976            assert_eq!(v.data.get(0), Value::UInt64(42));
1977            assert_eq!(v.data.get(1), Value::UInt64(43));
1978            assert_eq!(v.data.get(2), Value::UInt64(44));
1979        }
1980    }
1981}
mito2/read.rs

mito2/
read.rs
