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

mito2/
read.rs
