mito2/memtable/bulk/
part_reader.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::VecDeque;
16use std::time::Instant;
17
18use datatypes::arrow::array::BooleanArray;
19use datatypes::arrow::record_batch::RecordBatch;
20use mito_codec::row_converter::PrimaryKeyFilter;
21use parquet::arrow::ProjectionMask;
22use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
23use snafu::ResultExt;
24use store_api::storage::SequenceRange;
25
26use crate::error::{self, ComputeArrowSnafu, DecodeArrowRowGroupSnafu};
27use crate::memtable::bulk::context::{BulkIterContext, BulkIterContextRef};
28use crate::memtable::bulk::part::EncodedBulkPart;
29use crate::memtable::bulk::row_group_reader::MemtableRowGroupReaderBuilder;
30use crate::memtable::{MemScanMetrics, MemScanMetricsData};
31use crate::metrics::{READ_ROWS_TOTAL, READ_STAGE_ELAPSED};
32use crate::sst::parquet::file_range::TagDecodeState;
33use crate::sst::parquet::flat_format::{primary_key_column_index, sequence_column_index};
34use crate::sst::parquet::prefilter::{CachedPrimaryKeyFilter, prefilter_flat_batch_by_primary_key};
35
36/// Iterator for reading data inside a bulk part.
37pub struct EncodedBulkPartIter {
38    context: BulkIterContextRef,
39    row_groups_to_read: VecDeque<usize>,
40    current_reader: Option<ParquetRecordBatchReader>,
41    builder: MemtableRowGroupReaderBuilder,
42    /// Sequence number filter.
43    sequence: Option<SequenceRange>,
44    /// Cached skip_fields for current row group.
45    current_skip_fields: bool,
46    /// Primary key filter for prefiltering before convert_batch.
47    pk_filter: Option<CachedPrimaryKeyFilter>,
48    /// Metrics for this iterator.
49    metrics: MemScanMetricsData,
50    /// Optional memory scan metrics to report to.
51    mem_scan_metrics: Option<MemScanMetrics>,
52}
53
54impl EncodedBulkPartIter {
55    /// Creates a new [BulkPartIter].
56    pub fn try_new(
57        encoded_part: &EncodedBulkPart,
58        context: BulkIterContextRef,
59        mut row_groups_to_read: VecDeque<usize>,
60        sequence: Option<SequenceRange>,
61        mem_scan_metrics: Option<MemScanMetrics>,
62    ) -> error::Result<Self> {
63        let parquet_meta = encoded_part.metadata().parquet_metadata.clone();
64        let data = encoded_part.data().clone();
65        let series_count = encoded_part.metadata().num_series as usize;
66
67        // TODO(fys): Nested projection pushdown to the memtable layer is not supported yet.
68        let root_indices = context
69            .read_format()
70            .parquet_read_columns()
71            .root_indices_iter();
72        let projection_mask =
73            ProjectionMask::roots(parquet_meta.file_metadata().schema_descr(), root_indices);
74        let builder =
75            MemtableRowGroupReaderBuilder::try_new(&context, projection_mask, parquet_meta, data)?;
76
77        // Build PK filter if applicable (flat format with dictionary-encoded PKs).
78        let pk_filter = context.build_pk_filter();
79
80        let (init_reader, current_skip_fields) = match row_groups_to_read.pop_front() {
81            Some(first_row_group) => {
82                let skip_fields = context.pre_filter_mode().skip_fields();
83                let reader = builder.build_row_group_reader(first_row_group, None)?;
84                (Some(reader), skip_fields)
85            }
86            None => (None, false),
87        };
88
89        Ok(Self {
90            context,
91            row_groups_to_read,
92            current_reader: init_reader,
93            builder,
94            sequence,
95            current_skip_fields,
96            pk_filter,
97            metrics: MemScanMetricsData {
98                total_series: series_count,
99                ..Default::default()
100            },
101            mem_scan_metrics,
102        })
103    }
104
105    fn report_mem_scan_metrics(&mut self) {
106        if let Some(mem_scan_metrics) = self.mem_scan_metrics.take() {
107            mem_scan_metrics.merge_inner(&self.metrics);
108        }
109    }
110
111    /// Fetches next non-empty record batch.
112    pub(crate) fn next_record_batch(&mut self) -> error::Result<Option<RecordBatch>> {
113        let start = Instant::now();
114
115        let Some(current) = &mut self.current_reader else {
116            // All row group exhausted.
117            self.metrics.scan_cost += start.elapsed();
118            return Ok(None);
119        };
120
121        for batch in current {
122            let batch = batch.context(DecodeArrowRowGroupSnafu)?;
123            if let Some(batch) = apply_combined_filters(
124                &self.context,
125                &self.sequence,
126                batch,
127                self.current_skip_fields,
128                self.pk_filter
129                    .as_mut()
130                    .map(|f| f as &mut dyn PrimaryKeyFilter),
131                &mut self.metrics,
132            )? {
133                // Update metrics
134                self.metrics.num_batches += 1;
135                self.metrics.num_rows += batch.num_rows();
136                self.metrics.scan_cost += start.elapsed();
137                return Ok(Some(batch));
138            }
139        }
140
141        // Previous row group exhausted, read next row group
142        while let Some(next_row_group) = self.row_groups_to_read.pop_front() {
143            // Compute skip_fields for this row group
144            self.current_skip_fields = self.context.pre_filter_mode().skip_fields();
145
146            let next_reader = self.builder.build_row_group_reader(next_row_group, None)?;
147            let current = self.current_reader.insert(next_reader);
148
149            for batch in current {
150                let batch = batch.context(DecodeArrowRowGroupSnafu)?;
151                if let Some(batch) = apply_combined_filters(
152                    &self.context,
153                    &self.sequence,
154                    batch,
155                    self.current_skip_fields,
156                    self.pk_filter
157                        .as_mut()
158                        .map(|f| f as &mut dyn PrimaryKeyFilter),
159                    &mut self.metrics,
160                )? {
161                    // Update metrics
162                    self.metrics.num_batches += 1;
163                    self.metrics.num_rows += batch.num_rows();
164                    self.metrics.scan_cost += start.elapsed();
165                    return Ok(Some(batch));
166                }
167            }
168        }
169
170        self.metrics.scan_cost += start.elapsed();
171        Ok(None)
172    }
173}
174
175impl Iterator for EncodedBulkPartIter {
176    type Item = error::Result<RecordBatch>;
177
178    fn next(&mut self) -> Option<Self::Item> {
179        let result = self.next_record_batch().transpose();
180
181        // Report metrics when iteration is complete
182        if result.is_none() {
183            self.report_mem_scan_metrics();
184        }
185
186        result
187    }
188}
189
190impl Drop for EncodedBulkPartIter {
191    fn drop(&mut self) {
192        common_telemetry::debug!(
193            "EncodedBulkPartIter region: {}, metrics: total_series={}, num_rows={}, num_batches={}, scan_cost={:?}, prefilter_cost={:?}, prefilter_rows_filtered={}",
194            self.context.region_id(),
195            self.metrics.total_series,
196            self.metrics.num_rows,
197            self.metrics.num_batches,
198            self.metrics.scan_cost,
199            self.metrics.prefilter_cost,
200            self.metrics.prefilter_rows_filtered
201        );
202
203        // Report MemScanMetrics if not already reported
204        self.report_mem_scan_metrics();
205
206        READ_ROWS_TOTAL
207            .with_label_values(&["bulk_memtable"])
208            .inc_by(self.metrics.num_rows as u64);
209        READ_STAGE_ELAPSED
210            .with_label_values(&["scan_memtable"])
211            .observe(self.metrics.scan_cost.as_secs_f64());
212    }
213}
214
215/// Iterator for reading record batches from a bulk part.
216///
217/// Iterates through one or more RecordBatches, applying filters and projections.
218pub struct BulkPartBatchIter {
219    /// Queue of RecordBatches to process.
220    batches: VecDeque<RecordBatch>,
221    /// Iterator context for filtering and projection.
222    context: BulkIterContextRef,
223    /// Sequence number filter.
224    sequence: Option<SequenceRange>,
225    /// Primary key filter for prefiltering before convert_batch.
226    pk_filter: Option<CachedPrimaryKeyFilter>,
227    /// Metrics for this iterator.
228    metrics: MemScanMetricsData,
229    /// Optional memory scan metrics to report to.
230    mem_scan_metrics: Option<MemScanMetrics>,
231}
232
233impl BulkPartBatchIter {
234    /// Creates a new [BulkPartBatchIter] from multiple RecordBatches.
235    pub fn new(
236        batches: Vec<RecordBatch>,
237        context: BulkIterContextRef,
238        sequence: Option<SequenceRange>,
239        series_count: usize,
240        mem_scan_metrics: Option<MemScanMetrics>,
241    ) -> Self {
242        let pk_filter = context.build_pk_filter();
243
244        Self {
245            batches: VecDeque::from(batches),
246            context,
247            sequence,
248            pk_filter,
249            metrics: MemScanMetricsData {
250                total_series: series_count,
251                ..Default::default()
252            },
253            mem_scan_metrics,
254        }
255    }
256
257    /// Creates a new [BulkPartBatchIter] from a single RecordBatch.
258    pub fn from_single(
259        record_batch: RecordBatch,
260        context: BulkIterContextRef,
261        sequence: Option<SequenceRange>,
262        series_count: usize,
263        mem_scan_metrics: Option<MemScanMetrics>,
264    ) -> Self {
265        Self::new(
266            vec![record_batch],
267            context,
268            sequence,
269            series_count,
270            mem_scan_metrics,
271        )
272    }
273
274    fn report_mem_scan_metrics(&mut self) {
275        if let Some(mem_scan_metrics) = self.mem_scan_metrics.take() {
276            mem_scan_metrics.merge_inner(&self.metrics);
277        }
278    }
279
280    /// Applies projection to the RecordBatch if needed.
281    fn apply_projection(&self, record_batch: RecordBatch) -> error::Result<RecordBatch> {
282        let projection_indices = self
283            .context
284            .read_format()
285            .parquet_read_columns()
286            .root_indices();
287        if projection_indices.len() == record_batch.num_columns() {
288            return Ok(record_batch);
289        }
290
291        record_batch
292            .project(projection_indices)
293            .context(ComputeArrowSnafu)
294    }
295
296    fn process_batch(&mut self, record_batch: RecordBatch) -> error::Result<Option<RecordBatch>> {
297        let start = Instant::now();
298
299        // Apply projection first.
300        let projected_batch = self.apply_projection(record_batch)?;
301
302        // Apply combined filtering (both predicate and sequence filters)
303        let skip_fields = self.context.pre_filter_mode().skip_fields();
304
305        let Some(filtered_batch) = apply_combined_filters(
306            &self.context,
307            &self.sequence,
308            projected_batch,
309            skip_fields,
310            self.pk_filter
311                .as_mut()
312                .map(|f| f as &mut dyn PrimaryKeyFilter),
313            &mut self.metrics,
314        )?
315        else {
316            self.metrics.scan_cost += start.elapsed();
317            return Ok(None);
318        };
319
320        // Update metrics
321        self.metrics.num_batches += 1;
322        self.metrics.num_rows += filtered_batch.num_rows();
323        self.metrics.scan_cost += start.elapsed();
324
325        Ok(Some(filtered_batch))
326    }
327}
328
329impl Iterator for BulkPartBatchIter {
330    type Item = error::Result<RecordBatch>;
331
332    fn next(&mut self) -> Option<Self::Item> {
333        // Process batches until we find a non-empty one or run out
334        while let Some(batch) = self.batches.pop_front() {
335            match self.process_batch(batch) {
336                Ok(Some(result)) => return Some(Ok(result)),
337                Ok(None) => continue, // This batch was filtered out, try next
338                Err(e) => {
339                    self.report_mem_scan_metrics();
340                    return Some(Err(e));
341                }
342            }
343        }
344
345        // No more batches
346        self.report_mem_scan_metrics();
347        None
348    }
349}
350
351impl Drop for BulkPartBatchIter {
352    fn drop(&mut self) {
353        common_telemetry::debug!(
354            "BulkPartBatchIter region: {}, metrics: total_series={}, num_rows={}, num_batches={}, scan_cost={:?}, prefilter_cost={:?}, prefilter_rows_filtered={}",
355            self.context.region_id(),
356            self.metrics.total_series,
357            self.metrics.num_rows,
358            self.metrics.num_batches,
359            self.metrics.scan_cost,
360            self.metrics.prefilter_cost,
361            self.metrics.prefilter_rows_filtered
362        );
363
364        // Report MemScanMetrics if not already reported
365        self.report_mem_scan_metrics();
366
367        READ_ROWS_TOTAL
368            .with_label_values(&["bulk_memtable"])
369            .inc_by(self.metrics.num_rows as u64);
370        READ_STAGE_ELAPSED
371            .with_label_values(&["scan_memtable"])
372            .observe(self.metrics.scan_cost.as_secs_f64());
373    }
374}
375
376/// Applies both predicate filtering and sequence filtering in a single pass.
377/// Returns None if the filtered batch is empty.
378///
379/// # Panics
380/// Panics if the format is not flat.
381fn apply_combined_filters(
382    context: &BulkIterContext,
383    sequence: &Option<SequenceRange>,
384    record_batch: RecordBatch,
385    skip_fields: bool,
386    pk_filter: Option<&mut dyn PrimaryKeyFilter>,
387    metrics: &mut MemScanMetricsData,
388) -> error::Result<Option<RecordBatch>> {
389    // Apply PK prefilter on raw batch before convert_batch to reduce conversion overhead.
390    let record_batch = if let Some(pk_filter) = pk_filter {
391        let rows_before = record_batch.num_rows();
392        let prefilter_start = Instant::now();
393        let pk_col_idx = primary_key_column_index(record_batch.num_columns());
394        match prefilter_flat_batch_by_primary_key(record_batch, pk_col_idx, pk_filter)? {
395            Some(batch) => {
396                metrics.prefilter_cost += prefilter_start.elapsed();
397                metrics.prefilter_rows_filtered += rows_before - batch.num_rows();
398                batch
399            }
400            None => {
401                metrics.prefilter_cost += prefilter_start.elapsed();
402                metrics.prefilter_rows_filtered += rows_before;
403                return Ok(None);
404            }
405        }
406    } else {
407        record_batch
408    };
409
410    // Converts the format to the flat format.
411    let record_batch = context.read_format().convert_batch(record_batch, None)?;
412
413    let num_rows = record_batch.num_rows();
414    let mut combined_filter = None;
415    let mut tag_decode_state = TagDecodeState::new();
416
417    // Apply predicate filters using the shared method.
418    if !context.base.filters.is_empty() {
419        let predicate_mask = context.base.compute_filter_mask_flat(
420            &record_batch,
421            skip_fields,
422            false,
423            &mut tag_decode_state,
424        )?;
425        // If predicate filters out the entire batch, return None early
426        let Some(mask) = predicate_mask else {
427            return Ok(None);
428        };
429        combined_filter = Some(BooleanArray::from(mask));
430    }
431
432    // Filters rows by the given `sequence`. Only preserves rows with sequence less than or equal to `sequence`.
433    if let Some(sequence) = sequence {
434        let sequence_column =
435            record_batch.column(sequence_column_index(record_batch.num_columns()));
436        let sequence_filter = sequence
437            .filter(&sequence_column)
438            .context(ComputeArrowSnafu)?;
439        // Combine with existing filter using AND operation
440        combined_filter = match combined_filter {
441            None => Some(sequence_filter),
442            Some(existing_filter) => {
443                let and_result = datatypes::arrow::compute::and(&existing_filter, &sequence_filter)
444                    .context(ComputeArrowSnafu)?;
445                Some(and_result)
446            }
447        };
448    }
449
450    // Apply the combined filter if any filters were applied
451    let Some(filter_array) = combined_filter else {
452        // No filters applied, return original batch
453        return Ok(Some(record_batch));
454    };
455    let select_count = filter_array.true_count();
456    if select_count == 0 {
457        return Ok(None);
458    }
459    if select_count == num_rows {
460        return Ok(Some(record_batch));
461    }
462    let filtered_batch =
463        datatypes::arrow::compute::filter_record_batch(&record_batch, &filter_array)
464            .context(ComputeArrowSnafu)?;
465
466    Ok(Some(filtered_batch))
467}
468
469#[cfg(test)]
470mod tests {
471    use std::sync::Arc;
472
473    use api::v1::SemanticType;
474    use datafusion_expr::{col, lit};
475    use datatypes::arrow::array::{
476        ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, UInt8Array, UInt32Array,
477        UInt64Array,
478    };
479    use datatypes::arrow::datatypes::{DataType, Field, Schema};
480    use datatypes::data_type::ConcreteDataType;
481    use datatypes::schema::ColumnSchema;
482    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
483    use store_api::storage::RegionId;
484    use table::predicate::Predicate;
485
486    use super::*;
487    use crate::memtable::bulk::context::BulkIterContext;
488    use crate::test_util::sst_util::new_primary_key;
489
490    #[test]
491    fn test_bulk_part_batch_iter() {
492        // Create a simple schema
493        let schema = Arc::new(Schema::new(vec![
494            Field::new("key1", DataType::Utf8, false),
495            Field::new("field1", DataType::Int64, false),
496            Field::new(
497                "timestamp",
498                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None),
499                false,
500            ),
501            Field::new(
502                "__primary_key",
503                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
504                false,
505            ),
506            Field::new("__sequence", DataType::UInt64, false),
507            Field::new("__op_type", DataType::UInt8, false),
508        ]));
509
510        // Create test data
511        let key1 = Arc::new(StringArray::from_iter_values(["key1", "key2", "key3"]));
512        let field1 = Arc::new(Int64Array::from(vec![11, 12, 13]));
513        let timestamp = Arc::new(datatypes::arrow::array::TimestampMillisecondArray::from(
514            vec![1000, 2000, 3000],
515        ));
516
517        // Create primary key dictionary array with properly encoded PKs
518        use datatypes::arrow::array::{BinaryArray, DictionaryArray, UInt32Array};
519        let pk1 = new_primary_key(&["key1"]);
520        let pk2 = new_primary_key(&["key2"]);
521        let pk3 = new_primary_key(&["key3"]);
522        let values = Arc::new(BinaryArray::from_iter_values([
523            pk1.as_slice(),
524            pk2.as_slice(),
525            pk3.as_slice(),
526        ]));
527        let keys = UInt32Array::from(vec![0, 1, 2]);
528        let primary_key = Arc::new(DictionaryArray::new(keys, values));
529
530        let sequence = Arc::new(UInt64Array::from(vec![1, 2, 3]));
531        let op_type = Arc::new(UInt8Array::from(vec![1, 1, 1])); // PUT operations
532
533        let record_batch = RecordBatch::try_new(
534            schema,
535            vec![
536                key1,
537                field1,
538                timestamp,
539                primary_key.clone(),
540                sequence,
541                op_type,
542            ],
543        )
544        .unwrap();
545
546        // Create a minimal region metadata for testing
547        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
548        builder
549            .push_column_metadata(ColumnMetadata {
550                column_schema: ColumnSchema::new(
551                    "key1",
552                    ConcreteDataType::string_datatype(),
553                    false,
554                ),
555                semantic_type: SemanticType::Tag,
556                column_id: 0,
557            })
558            .push_column_metadata(ColumnMetadata {
559                column_schema: ColumnSchema::new(
560                    "field1",
561                    ConcreteDataType::int64_datatype(),
562                    false,
563                ),
564                semantic_type: SemanticType::Field,
565                column_id: 1,
566            })
567            .push_column_metadata(ColumnMetadata {
568                column_schema: ColumnSchema::new(
569                    "timestamp",
570                    ConcreteDataType::timestamp_millisecond_datatype(),
571                    false,
572                ),
573                semantic_type: SemanticType::Timestamp,
574                column_id: 2,
575            })
576            .primary_key(vec![0]);
577
578        let region_metadata = builder.build().unwrap();
579
580        // Create context
581        let context = Arc::new(
582            BulkIterContext::new(
583                Arc::new(region_metadata.clone()),
584                None, // No projection
585                None, // No predicate
586                false,
587                crate::sst::parquet::DEFAULT_READ_BATCH_SIZE,
588            )
589            .unwrap(),
590        );
591        // Iterates all rows.
592        let iter =
593            BulkPartBatchIter::from_single(record_batch.clone(), context.clone(), None, 0, None);
594        let result: Vec<_> = iter.map(|rb| rb.unwrap()).collect();
595        assert_eq!(1, result.len());
596        assert_eq!(3, result[0].num_rows());
597        assert_eq!(6, result[0].num_columns(),);
598
599        // Creates iter with sequence filter (only include sequences <= 2)
600        let iter = BulkPartBatchIter::from_single(
601            record_batch.clone(),
602            context,
603            Some(SequenceRange::LtEq { max: 2 }),
604            0,
605            None,
606        );
607        let result: Vec<_> = iter.map(|rb| rb.unwrap()).collect();
608        assert_eq!(1, result.len());
609        let expect_sequence = Arc::new(UInt64Array::from(vec![1, 2])) as ArrayRef;
610        assert_eq!(
611            &expect_sequence,
612            result[0].column(result[0].num_columns() - 2)
613        );
614        assert_eq!(6, result[0].num_columns());
615
616        let context = Arc::new(
617            BulkIterContext::new(
618                Arc::new(region_metadata),
619                Some(&[0, 2]),
620                Some(Predicate::new(vec![col("key1").eq(lit("key2"))])),
621                false,
622                crate::sst::parquet::DEFAULT_READ_BATCH_SIZE,
623            )
624            .unwrap(),
625        );
626        // Creates iter with projection and predicate.
627        let iter =
628            BulkPartBatchIter::from_single(record_batch.clone(), context.clone(), None, 0, None);
629        let result: Vec<_> = iter.map(|rb| rb.unwrap()).collect();
630        assert_eq!(1, result.len());
631        assert_eq!(1, result[0].num_rows());
632        assert_eq!(5, result[0].num_columns());
633        let expect_sequence = Arc::new(UInt64Array::from(vec![2])) as ArrayRef;
634        assert_eq!(
635            &expect_sequence,
636            result[0].column(result[0].num_columns() - 2)
637        );
638    }
639
640    #[test]
641    fn test_bulk_part_batch_iter_multiple_batches() {
642        // Create a simple schema
643        let schema = Arc::new(Schema::new(vec![
644            Field::new("key1", DataType::Utf8, false),
645            Field::new("field1", DataType::Int64, false),
646            Field::new(
647                "timestamp",
648                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None),
649                false,
650            ),
651            Field::new(
652                "__primary_key",
653                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
654                false,
655            ),
656            Field::new("__sequence", DataType::UInt64, false),
657            Field::new("__op_type", DataType::UInt8, false),
658        ]));
659
660        // Create first batch with 2 rows
661        let pk1 = new_primary_key(&["key1"]);
662        let pk2 = new_primary_key(&["key2"]);
663        let key1_1 = Arc::new(StringArray::from_iter_values(["key1", "key2"]));
664        let field1_1 = Arc::new(Int64Array::from(vec![11, 12]));
665        let timestamp_1 = Arc::new(datatypes::arrow::array::TimestampMillisecondArray::from(
666            vec![1000, 2000],
667        ));
668        let values_1 = Arc::new(BinaryArray::from_iter_values([
669            pk1.as_slice(),
670            pk2.as_slice(),
671        ]));
672        let keys_1 = UInt32Array::from(vec![0, 1]);
673        let primary_key_1 = Arc::new(DictionaryArray::new(keys_1, values_1));
674        let sequence_1 = Arc::new(UInt64Array::from(vec![1, 2]));
675        let op_type_1 = Arc::new(UInt8Array::from(vec![1, 1]));
676
677        let batch1 = RecordBatch::try_new(
678            schema.clone(),
679            vec![
680                key1_1,
681                field1_1,
682                timestamp_1,
683                primary_key_1,
684                sequence_1,
685                op_type_1,
686            ],
687        )
688        .unwrap();
689
690        // Create second batch with 3 rows
691        let pk3 = new_primary_key(&["key3"]);
692        let pk4 = new_primary_key(&["key4"]);
693        let pk5 = new_primary_key(&["key5"]);
694        let key1_2 = Arc::new(StringArray::from_iter_values(["key3", "key4", "key5"]));
695        let field1_2 = Arc::new(Int64Array::from(vec![13, 14, 15]));
696        let timestamp_2 = Arc::new(datatypes::arrow::array::TimestampMillisecondArray::from(
697            vec![3000, 4000, 5000],
698        ));
699        let values_2 = Arc::new(BinaryArray::from_iter_values([
700            pk3.as_slice(),
701            pk4.as_slice(),
702            pk5.as_slice(),
703        ]));
704        let keys_2 = UInt32Array::from(vec![0, 1, 2]);
705        let primary_key_2 = Arc::new(DictionaryArray::new(keys_2, values_2));
706        let sequence_2 = Arc::new(UInt64Array::from(vec![3, 4, 5]));
707        let op_type_2 = Arc::new(UInt8Array::from(vec![1, 1, 1]));
708
709        let batch2 = RecordBatch::try_new(
710            schema.clone(),
711            vec![
712                key1_2,
713                field1_2,
714                timestamp_2,
715                primary_key_2,
716                sequence_2,
717                op_type_2,
718            ],
719        )
720        .unwrap();
721
722        // Create region metadata
723        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
724        builder
725            .push_column_metadata(ColumnMetadata {
726                column_schema: ColumnSchema::new(
727                    "key1",
728                    ConcreteDataType::string_datatype(),
729                    false,
730                ),
731                semantic_type: SemanticType::Tag,
732                column_id: 0,
733            })
734            .push_column_metadata(ColumnMetadata {
735                column_schema: ColumnSchema::new(
736                    "field1",
737                    ConcreteDataType::int64_datatype(),
738                    false,
739                ),
740                semantic_type: SemanticType::Field,
741                column_id: 1,
742            })
743            .push_column_metadata(ColumnMetadata {
744                column_schema: ColumnSchema::new(
745                    "timestamp",
746                    ConcreteDataType::timestamp_millisecond_datatype(),
747                    false,
748                ),
749                semantic_type: SemanticType::Timestamp,
750                column_id: 2,
751            })
752            .primary_key(vec![0]);
753
754        let region_metadata = builder.build().unwrap();
755
756        // Create context
757        let context = Arc::new(
758            BulkIterContext::new(
759                Arc::new(region_metadata),
760                None, // No projection
761                None, // No predicate
762                false,
763                1,
764            )
765            .unwrap(),
766        );
767
768        // Create iterator with multiple batches
769        let expect_batches = vec![batch1, batch2];
770        let iter = BulkPartBatchIter::new(expect_batches.clone(), context.clone(), None, 0, None);
771
772        // Existing valid Arrow batches retain their boundaries even when the merge hint is lower.
773        let result: Vec<_> = iter.map(|rb| rb.unwrap()).collect();
774        assert_eq!(expect_batches, result);
775    }
776}
mito2/memtable/bulk/part_reader.rs

mito2/memtable/bulk/
part_reader.rs
