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

mito2/
read.rs
