mito2/read/
flat_dedup.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//! Dedup implementation for flat format.
16
17use std::ops::Range;
18use std::sync::Arc;
19use std::time::Instant;
20
21use api::v1::OpType;
22use async_stream::try_stream;
23use common_telemetry::debug;
24use datatypes::arrow::array::{
25    Array, ArrayRef, BinaryArray, BooleanArray, BooleanBufferBuilder, UInt8Array, UInt64Array,
26    make_comparator,
27};
28use datatypes::arrow::buffer::BooleanBuffer;
29use datatypes::arrow::compute::kernels::cmp::distinct;
30use datatypes::arrow::compute::kernels::partition::{Partitions, partition};
31use datatypes::arrow::compute::kernels::take::take;
32use datatypes::arrow::compute::{
33    SortOptions, TakeOptions, concat_batches, filter_record_batch, take_record_batch,
34};
35use datatypes::arrow::error::ArrowError;
36use datatypes::arrow::record_batch::RecordBatch;
37use futures::{Stream, TryStreamExt};
38use snafu::ResultExt;
39
40use crate::error::{ComputeArrowSnafu, NewRecordBatchSnafu, Result};
41use crate::metrics::MERGE_FILTER_ROWS_TOTAL;
42use crate::read::dedup::{DedupMetrics, DedupMetricsReport};
43use crate::read::timestamp_array_to_i64_slice;
44use crate::sst::parquet::flat_format::{
45    op_type_column_index, primary_key_column_index, time_index_column_index,
46};
47use crate::sst::parquet::format::{FIXED_POS_COLUMN_NUM, PrimaryKeyArray};
48
49/// An iterator to dedup sorted batches from an iterator based on the dedup strategy.
50pub struct FlatDedupIterator<I, S> {
51    iter: I,
52    strategy: S,
53    metrics: DedupMetrics,
54}
55
56impl<I, S> FlatDedupIterator<I, S> {
57    /// Creates a new dedup iterator.
58    pub fn new(iter: I, strategy: S) -> Self {
59        Self {
60            iter,
61            strategy,
62            metrics: DedupMetrics::default(),
63        }
64    }
65}
66
67impl<I: Iterator<Item = Result<RecordBatch>>, S: RecordBatchDedupStrategy> FlatDedupIterator<I, S> {
68    /// Returns the next deduplicated batch.
69    fn fetch_next_batch(&mut self) -> Result<Option<RecordBatch>> {
70        while let Some(batch) = self.iter.next().transpose()? {
71            if let Some(batch) = self.strategy.push_batch(batch, &mut self.metrics)? {
72                return Ok(Some(batch));
73            }
74        }
75
76        self.strategy.finish(&mut self.metrics)
77    }
78}
79
80impl<I: Iterator<Item = Result<RecordBatch>>, S: RecordBatchDedupStrategy> Iterator
81    for FlatDedupIterator<I, S>
82{
83    type Item = Result<RecordBatch>;
84
85    fn next(&mut self) -> Option<Self::Item> {
86        self.fetch_next_batch().transpose()
87    }
88}
89
90/// An async reader to dedup sorted record batches from a stream based on the dedup strategy.
91pub struct FlatDedupReader<I, S> {
92    stream: I,
93    strategy: S,
94    metrics: DedupMetrics,
95    /// Optional metrics reporter.
96    metrics_reporter: Option<Arc<dyn DedupMetricsReport>>,
97}
98
99impl<I, S> FlatDedupReader<I, S> {
100    /// Creates a new dedup reader.
101    pub fn new(
102        stream: I,
103        strategy: S,
104        metrics_reporter: Option<Arc<dyn DedupMetricsReport>>,
105    ) -> Self {
106        Self {
107            stream,
108            strategy,
109            metrics: DedupMetrics::default(),
110            metrics_reporter,
111        }
112    }
113}
114
115impl<I: Stream<Item = Result<RecordBatch>> + Unpin, S: RecordBatchDedupStrategy>
116    FlatDedupReader<I, S>
117{
118    /// Returns the next deduplicated batch.
119    async fn fetch_next_batch(&mut self) -> Result<Option<RecordBatch>> {
120        while let Some(batch) = self.stream.try_next().await? {
121            if let Some(batch) = self.strategy.push_batch(batch, &mut self.metrics)? {
122                self.metrics.maybe_report(&self.metrics_reporter);
123                return Ok(Some(batch));
124            }
125        }
126
127        let result = self.strategy.finish(&mut self.metrics)?;
128        self.metrics.maybe_report(&self.metrics_reporter);
129        Ok(result)
130    }
131
132    /// Converts the reader into a stream.
133    pub fn into_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
134        try_stream! {
135            while let Some(batch) = self.fetch_next_batch().await? {
136                yield batch;
137            }
138        }
139    }
140}
141
142impl<I, S> Drop for FlatDedupReader<I, S> {
143    fn drop(&mut self) {
144        debug!("Flat dedup reader finished, metrics: {:?}", self.metrics);
145
146        MERGE_FILTER_ROWS_TOTAL
147            .with_label_values(&["dedup"])
148            .inc_by(self.metrics.num_unselected_rows as u64);
149        MERGE_FILTER_ROWS_TOTAL
150            .with_label_values(&["delete"])
151            .inc_by(self.metrics.num_deleted_rows as u64);
152
153        // Report any remaining metrics.
154        if let Some(reporter) = &self.metrics_reporter {
155            reporter.report(&mut self.metrics);
156        }
157    }
158}
159
160/// Strategy to remove duplicate rows from sorted record batches.
161pub trait RecordBatchDedupStrategy: Send {
162    /// Pushes a batch to the dedup strategy.
163    /// Returns rows that no longer need to be updated by subsequent input.
164    /// A strategy for partial merges may retain duplicates separated by Deletes.
165    fn push_batch(
166        &mut self,
167        batch: RecordBatch,
168        metrics: &mut DedupMetrics,
169    ) -> Result<Option<RecordBatch>>;
170
171    /// Finishes the deduplication process and returns any remaining batch.
172    ///
173    /// Users must ensure that `push_batch` is called for all batches before
174    /// calling this method.
175    fn finish(&mut self, metrics: &mut DedupMetrics) -> Result<Option<RecordBatch>>;
176}
177
178/// Dedup strategy that keeps the row with latest sequence of each key.
179pub struct FlatLastRow {
180    /// Meta of the last row in the previous batch that has the same key
181    /// as the batch to push.
182    prev_batch: Option<BatchLastRow>,
183    /// Filter deleted rows.
184    filter_deleted: bool,
185}
186
187impl FlatLastRow {
188    /// Creates a new strategy with the given `filter_deleted` flag.
189    pub fn new(filter_deleted: bool) -> Self {
190        Self {
191            prev_batch: None,
192            filter_deleted,
193        }
194    }
195
196    /// Remove duplications from the batch without considering previous rows.
197    fn dedup_one_batch(batch: RecordBatch) -> Result<RecordBatch> {
198        let num_rows = batch.num_rows();
199        if num_rows < 2 {
200            return Ok(batch);
201        }
202
203        let num_columns = batch.num_columns();
204        let timestamps = batch.column(time_index_column_index(num_columns));
205        // Checks duplications based on the timestamp.
206        let mask = find_boundaries(timestamps).context(ComputeArrowSnafu)?;
207        if mask.count_set_bits() == num_rows - 1 {
208            // Fast path: No duplication.
209            return Ok(batch);
210        }
211
212        // The batch has duplicated timestamps, but it doesn't mean it must
213        // has duplicated rows.
214        // Partitions the batch by the primary key and time index.
215        let columns: Vec<_> = [
216            primary_key_column_index(num_columns),
217            time_index_column_index(num_columns),
218        ]
219        .iter()
220        .map(|index| batch.column(*index).clone())
221        .collect();
222        let partitions = partition(&columns).context(ComputeArrowSnafu)?;
223
224        Self::dedup_by_partitions(batch, &partitions)
225    }
226
227    /// Remove duplications for each partition.
228    fn dedup_by_partitions(batch: RecordBatch, partitions: &Partitions) -> Result<RecordBatch> {
229        let ranges = partitions.ranges();
230        // Each range at least has 1 row.
231        let num_duplications: usize = ranges.iter().map(|r| r.end - r.start - 1).sum();
232        if num_duplications == 0 {
233            // Fast path, no duplications.
234            return Ok(batch);
235        }
236
237        // Always takes the first row in each range.
238        let take_indices: UInt64Array = ranges.iter().map(|r| Some(r.start as u64)).collect();
239        take_record_batch(&batch, &take_indices).context(ComputeArrowSnafu)
240    }
241}
242
243impl RecordBatchDedupStrategy for FlatLastRow {
244    fn push_batch(
245        &mut self,
246        batch: RecordBatch,
247        metrics: &mut DedupMetrics,
248    ) -> Result<Option<RecordBatch>> {
249        let start = Instant::now();
250
251        if batch.num_rows() == 0 {
252            return Ok(None);
253        }
254
255        // Dedup current batch to ensure no duplication before we checking the previous row.
256        let row_before_dedup = batch.num_rows();
257        let mut batch = Self::dedup_one_batch(batch)?;
258
259        if let Some(prev_batch) = &self.prev_batch {
260            // If we have previous batch.
261            if prev_batch.is_last_row_duplicated(&batch) {
262                // Duplicated with the last batch, skip the first row.
263                batch = batch.slice(1, batch.num_rows() - 1);
264            }
265        }
266        metrics.num_unselected_rows += row_before_dedup - batch.num_rows();
267
268        let Some(batch_last_row) = BatchLastRow::try_new(batch.clone()) else {
269            // The batch after dedup is empty.
270            // We don't need to update `prev_batch` because they have the same
271            // key and timestamp.
272            metrics.dedup_cost += start.elapsed();
273            return Ok(None);
274        };
275
276        // Store current batch to `prev_batch` so we could compare the next batch
277        // with this batch. We store batch before filtering it as rows with `OpType::Delete`
278        // would be removed from the batch after filter, then we may store an incorrect `last row`
279        // of previous batch.
280        // Safety: We checked the batch is not empty before.
281        self.prev_batch = Some(batch_last_row);
282
283        // Filters deleted rows at last.
284        let result = maybe_filter_deleted(batch, self.filter_deleted, metrics);
285
286        metrics.dedup_cost += start.elapsed();
287
288        result
289    }
290
291    fn finish(&mut self, _metrics: &mut DedupMetrics) -> Result<Option<RecordBatch>> {
292        Ok(None)
293    }
294}
295
296/// Dedup strategy that keeps the last non-null field for the same key.
297pub struct FlatLastNonNull {
298    /// The start index of field columns:
299    field_column_start: usize,
300    /// Filter deleted rows; otherwise preserve every delete barrier.
301    filter_deleted: bool,
302    /// Buffered batch to check whether the next batch have duplicated rows with this batch.
303    /// Fields in the last row of this batch may be updated by the next batch.
304    /// In preserve mode, duplicate keys may be separated by delete markers.
305    buffer: Option<BatchLastRow>,
306    /// Whether the last row range contains a delete operation.
307    /// If so, we don't need to update null fields.
308    contains_delete: bool,
309}
310
311impl RecordBatchDedupStrategy for FlatLastNonNull {
312    fn push_batch(
313        &mut self,
314        batch: RecordBatch,
315        metrics: &mut DedupMetrics,
316    ) -> Result<Option<RecordBatch>> {
317        let start = Instant::now();
318
319        if batch.num_rows() == 0 {
320            return Ok(None);
321        }
322
323        let row_before_dedup = batch.num_rows();
324
325        let Some(buffer) = self.buffer.take() else {
326            // If the buffer is None, dedup the batch, put the batch into the buffer and return.
327            // There is no previous batch with the same key, we can pass contains_delete as false.
328            let (record_batch, contains_delete) =
329                Self::dedup_one_batch(batch, self.field_column_start, false, self.filter_deleted)?;
330            metrics.num_unselected_rows += row_before_dedup - record_batch.num_rows();
331            self.buffer = BatchLastRow::try_new(record_batch);
332            self.contains_delete = contains_delete;
333
334            metrics.dedup_cost += start.elapsed();
335            return Ok(None);
336        };
337
338        if !buffer.is_last_row_duplicated(&batch) {
339            // The first row of batch has different key from the buffer.
340            // We can replace the buffer with the new batch.
341            // Dedup the batch.
342            // There is no previous batch with the same key, we can pass contains_delete as false.
343            let (record_batch, contains_delete) =
344                Self::dedup_one_batch(batch, self.field_column_start, false, self.filter_deleted)?;
345            metrics.num_unselected_rows += row_before_dedup - record_batch.num_rows();
346            debug_assert!(record_batch.num_rows() > 0);
347            self.buffer = BatchLastRow::try_new(record_batch);
348            self.contains_delete = contains_delete;
349
350            let result = maybe_filter_deleted(buffer.last_batch, self.filter_deleted, metrics);
351            metrics.dedup_cost += start.elapsed();
352            return result;
353        }
354
355        // The next batch has duplicated rows.
356        // We can return rows except the last row in the buffer.
357        let output = if buffer.last_batch.num_rows() > 1 {
358            let dedup_batch = buffer.last_batch.slice(0, buffer.last_batch.num_rows() - 1);
359            debug_assert_eq!(buffer.last_batch.num_rows() - 1, dedup_batch.num_rows());
360
361            maybe_filter_deleted(dedup_batch, self.filter_deleted, metrics)?
362        } else {
363            None
364        };
365        let last_row = buffer.last_batch.slice(buffer.last_batch.num_rows() - 1, 1);
366
367        // We concat the last row with the next batch.
368        let schema = batch.schema();
369        let merged = concat_batches(&schema, &[last_row, batch]).context(ComputeArrowSnafu)?;
370        let merged_row_count = merged.num_rows();
371        // Dedup the merged batch and update the buffer.
372        let (record_batch, contains_delete) = Self::dedup_one_batch(
373            merged,
374            self.field_column_start,
375            self.contains_delete,
376            self.filter_deleted,
377        )?;
378        metrics.num_unselected_rows += merged_row_count - record_batch.num_rows();
379        debug_assert!(record_batch.num_rows() > 0);
380        self.buffer = BatchLastRow::try_new(record_batch);
381        self.contains_delete = contains_delete;
382
383        metrics.dedup_cost += start.elapsed();
384
385        Ok(output)
386    }
387
388    fn finish(&mut self, metrics: &mut DedupMetrics) -> Result<Option<RecordBatch>> {
389        let Some(buffer) = self.buffer.take() else {
390            return Ok(None);
391        };
392
393        let start = Instant::now();
394
395        let result = maybe_filter_deleted(buffer.last_batch, self.filter_deleted, metrics);
396
397        metrics.dedup_cost += start.elapsed();
398
399        result
400    }
401}
402
403impl FlatLastNonNull {
404    /// Creates a strategy with the given `filter_deleted` flag. When false,
405    /// every delete barrier is preserved and Puts are folded only between
406    /// barriers so later reads cannot fill fields from a deleted row.
407    /// Inputs for partial merges must form a sequence-contiguous group to avoid
408    /// folding fields across an intermediate version left outside the merge.
409    pub fn new(field_column_start: usize, filter_deleted: bool) -> Self {
410        Self {
411            field_column_start,
412            filter_deleted,
413            buffer: None,
414            contains_delete: false,
415        }
416    }
417
418    /// Remove duplications from the batch without considering the previous and next rows.
419    /// Returns a tuple containing the deduplicated batch and a boolean indicating whether the last range contains deleted rows.
420    fn dedup_one_batch(
421        batch: RecordBatch,
422        field_column_start: usize,
423        prev_batch_contains_delete: bool,
424        filter_deleted: bool,
425    ) -> Result<(RecordBatch, bool)> {
426        // Get op type array for checking delete operations
427        let op_type_column = batch
428            .column(op_type_column_index(batch.num_columns()))
429            .clone();
430        let op_types = op_type_column
431            .as_any()
432            .downcast_ref::<UInt8Array>()
433            .unwrap();
434        let num_rows = batch.num_rows();
435        if num_rows < 2 {
436            let contains_delete = if num_rows > 0 {
437                op_types.value(0) == OpType::Delete as u8
438            } else {
439                false
440            };
441            return Ok((batch, contains_delete));
442        }
443
444        let num_columns = batch.num_columns();
445        let timestamps = batch.column(time_index_column_index(num_columns));
446        // Checks duplications based on the timestamp.
447        let mask = find_boundaries(timestamps).context(ComputeArrowSnafu)?;
448        if mask.count_set_bits() == num_rows - 1 {
449            let contains_delete = op_types.value(num_rows - 1) == OpType::Delete as u8;
450            // Fast path: No duplication.
451            return Ok((batch, contains_delete));
452        }
453
454        // The batch has duplicated timestamps, but it doesn't mean it must
455        // has duplicated rows.
456        // Partitions the batch by the primary key and time index.
457        let columns: Vec<_> = [
458            primary_key_column_index(num_columns),
459            time_index_column_index(num_columns),
460        ]
461        .iter()
462        .map(|index| batch.column(*index).clone())
463        .collect();
464        let partitions = partition(&columns).context(ComputeArrowSnafu)?;
465
466        Self::dedup_by_partitions(
467            batch,
468            &partitions,
469            field_column_start,
470            op_types,
471            prev_batch_contains_delete,
472            filter_deleted,
473        )
474    }
475
476    /// Remove depulications for each partition.
477    /// Returns a tuple containing the deduplicated batch and a boolean indicating whether the last range contains deleted rows.
478    fn dedup_by_partitions(
479        batch: RecordBatch,
480        partitions: &Partitions,
481        field_column_start: usize,
482        op_types: &UInt8Array,
483        first_range_contains_delete: bool,
484        filter_deleted: bool,
485    ) -> Result<(RecordBatch, bool)> {
486        let mut ranges = partitions.ranges();
487        let contains_delete = filter_deleted && Self::last_range_has_delete(&ranges, op_types);
488        if !filter_deleted && op_types.values().contains(&(OpType::Delete as u8)) {
489            let mut segments = Vec::with_capacity(ranges.len());
490            for range in ranges {
491                let mut start = range.start;
492                for i in range.clone() {
493                    if op_types.value(i) == OpType::Delete as u8 {
494                        if start < i {
495                            segments.push(start..i);
496                        }
497                        segments.push(i..i + 1);
498                        start = i + 1;
499                    }
500                }
501                if start < range.end {
502                    segments.push(start..range.end);
503                }
504            }
505            ranges = segments;
506        }
507
508        // Each range at least has 1 row.
509        let num_duplications: usize = ranges.iter().map(|r| r.end - r.start - 1).sum();
510        if num_duplications == 0 {
511            // Fast path, no duplication.
512            return Ok((batch, contains_delete));
513        }
514
515        let field_column_end = batch.num_columns() - FIXED_POS_COLUMN_NUM;
516        let take_options = Some(TakeOptions {
517            check_bounds: false,
518        });
519        // Always takes the first value for non-field columns in each range.
520        let non_field_indices: UInt64Array = ranges.iter().map(|r| Some(r.start as u64)).collect();
521        let new_columns = batch
522            .columns()
523            .iter()
524            .enumerate()
525            .map(|(col_idx, column)| {
526                if col_idx >= field_column_start && col_idx < field_column_end {
527                    let field_indices = Self::compute_field_indices(
528                        &ranges,
529                        column,
530                        op_types,
531                        first_range_contains_delete && filter_deleted,
532                    );
533                    take(column, &field_indices, take_options.clone()).context(ComputeArrowSnafu)
534                } else {
535                    take(column, &non_field_indices, take_options.clone())
536                        .context(ComputeArrowSnafu)
537                }
538            })
539            .collect::<Result<Vec<ArrayRef>>>()?;
540
541        let record_batch =
542            RecordBatch::try_new(batch.schema(), new_columns).context(NewRecordBatchSnafu)?;
543        Ok((record_batch, contains_delete))
544    }
545
546    /// Returns an array of indices of the latest non null value for
547    /// each input range.
548    /// If all values in a range are null, the returned index is unspecific.
549    /// Stops when encountering a delete operation and ignores all subsequent rows.
550    fn compute_field_indices(
551        ranges: &[Range<usize>],
552        field_array: &ArrayRef,
553        op_types: &UInt8Array,
554        first_range_contains_delete: bool,
555    ) -> UInt64Array {
556        ranges
557            .iter()
558            .enumerate()
559            .map(|(range_idx, r)| {
560                let mut value_index = r.start as u64;
561                if range_idx == 0 && first_range_contains_delete {
562                    return Some(value_index);
563                }
564
565                // Iterate through the range to find the first valid non-null value
566                // but stop if we encounter a delete operation.
567                for i in r.clone() {
568                    if op_types.value(i) == OpType::Delete as u8 {
569                        break;
570                    }
571                    if field_array.is_valid(i) {
572                        value_index = i as u64;
573                        break;
574                    }
575                }
576
577                Some(value_index)
578            })
579            .collect()
580    }
581
582    /// Checks whether the last range contains a delete operation.
583    fn last_range_has_delete(ranges: &[Range<usize>], op_types: &UInt8Array) -> bool {
584        if let Some(last_range) = ranges.last() {
585            last_range
586                .clone()
587                .any(|i| op_types.value(i) == OpType::Delete as u8)
588        } else {
589            false
590        }
591    }
592}
593
594/// State of the batch with the last row for dedup.
595struct BatchLastRow {
596    /// The record batch that contains the last row.
597    /// It must has at least one row.
598    last_batch: RecordBatch,
599    /// Primary keys of the last batch.
600    primary_key: PrimaryKeyArray,
601    /// Last timestamp value.
602    timestamp: i64,
603}
604
605impl BatchLastRow {
606    /// Returns a new [BatchLastRow] if the record batch is not empty.
607    fn try_new(record_batch: RecordBatch) -> Option<Self> {
608        if record_batch.num_rows() > 0 {
609            let num_columns = record_batch.num_columns();
610            let primary_key = record_batch
611                .column(primary_key_column_index(num_columns))
612                .as_any()
613                .downcast_ref::<PrimaryKeyArray>()
614                .unwrap()
615                .clone();
616            let timestamp_array = record_batch.column(time_index_column_index(num_columns));
617            let timestamp = timestamp_value(timestamp_array, timestamp_array.len() - 1);
618
619            Some(Self {
620                last_batch: record_batch,
621                primary_key,
622                timestamp,
623            })
624        } else {
625            None
626        }
627    }
628
629    /// Returns true if the first row of the input `batch` is duplicated with the last row.
630    fn is_last_row_duplicated(&self, batch: &RecordBatch) -> bool {
631        if batch.num_rows() == 0 {
632            return false;
633        }
634
635        // The first timestamp in the batch.
636        let batch_timestamp = timestamp_value(
637            batch.column(time_index_column_index(batch.num_columns())),
638            0,
639        );
640        if batch_timestamp != self.timestamp {
641            return false;
642        }
643
644        let last_key = primary_key_at(&self.primary_key, self.last_batch.num_rows() - 1);
645        let primary_key = batch
646            .column(primary_key_column_index(batch.num_columns()))
647            .as_any()
648            .downcast_ref::<PrimaryKeyArray>()
649            .unwrap();
650        // Primary key of the first row in the batch.
651        let batch_key = primary_key_at(primary_key, 0);
652
653        last_key == batch_key
654    }
655}
656
657// TODO(yingwen): We only compares timestamp arrays, we can modify this function
658// to simplify the comparator.
659// Port from https://github.com/apache/arrow-rs/blob/55.0.0/arrow-ord/src/partition.rs#L155-L168
660/// Returns a mask with bits set whenever the value or nullability changes
661fn find_boundaries(v: &dyn Array) -> Result<BooleanBuffer, ArrowError> {
662    let slice_len = v.len() - 1;
663    let v1 = v.slice(0, slice_len);
664    let v2 = v.slice(1, slice_len);
665
666    if !v.data_type().is_nested() {
667        return Ok(distinct(&v1, &v2)?.values().clone());
668    }
669    // Given that we're only comparing values, null ordering in the input or
670    // sort options do not matter.
671    let cmp = make_comparator(&v1, &v2, SortOptions::default())?;
672    Ok((0..slice_len).map(|i| !cmp(i, i).is_eq()).collect())
673}
674
675/// Filters deleted rows from the record batch if `filter_deleted` is true.
676fn maybe_filter_deleted(
677    record_batch: RecordBatch,
678    filter_deleted: bool,
679    metrics: &mut DedupMetrics,
680) -> Result<Option<RecordBatch>> {
681    if !filter_deleted {
682        return Ok(Some(record_batch));
683    }
684    let batch = filter_deleted_from_batch(record_batch, metrics)?;
685    // Skips empty batches.
686    if batch.num_rows() == 0 {
687        return Ok(None);
688    }
689    Ok(Some(batch))
690}
691
692/// Removes deleted rows from the batch and updates metrics.
693fn filter_deleted_from_batch(
694    batch: RecordBatch,
695    metrics: &mut DedupMetrics,
696) -> Result<RecordBatch> {
697    let num_rows = batch.num_rows();
698    let op_type_column = batch.column(op_type_column_index(batch.num_columns()));
699    // Safety: The column should be op type.
700    let op_types = op_type_column
701        .as_any()
702        .downcast_ref::<UInt8Array>()
703        .unwrap();
704    let has_delete = op_types
705        .values()
706        .iter()
707        .any(|op_type| *op_type != OpType::Put as u8);
708    if !has_delete {
709        return Ok(batch);
710    }
711
712    let mut builder = BooleanBufferBuilder::new(op_types.len());
713    for op_type in op_types.values() {
714        if *op_type == OpType::Delete as u8 {
715            builder.append(false);
716        } else {
717            builder.append(true);
718        }
719    }
720    let predicate = BooleanArray::new(builder.into(), None);
721    let new_batch = filter_record_batch(&batch, &predicate).context(ComputeArrowSnafu)?;
722    let num_deleted = num_rows - new_batch.num_rows();
723    metrics.num_deleted_rows += num_deleted;
724    metrics.num_unselected_rows += num_deleted;
725
726    Ok(new_batch)
727}
728
729/// Gets the primary key at `index`.
730fn primary_key_at(primary_key: &PrimaryKeyArray, index: usize) -> &[u8] {
731    let key = primary_key.keys().value(index);
732    let binary_values = primary_key
733        .values()
734        .as_any()
735        .downcast_ref::<BinaryArray>()
736        .unwrap();
737    binary_values.value(key as usize)
738}
739
740/// Gets the timestamp value from the timestamp array.
741///
742/// # Panics
743/// Panics if the array is not a timestamp array or
744/// the index is out of bound.
745pub(crate) fn timestamp_value(array: &ArrayRef, idx: usize) -> i64 {
746    timestamp_array_to_i64_slice(array)[idx]
747}
748
749#[cfg(test)]
750mod tests {
751    use std::sync::Arc;
752
753    use api::v1::OpType;
754    use datatypes::arrow::array::{
755        ArrayRef, BinaryDictionaryBuilder, Int64Array, StringDictionaryBuilder,
756        TimestampMillisecondArray, UInt8Array, UInt64Array,
757    };
758    use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit, UInt32Type};
759    use datatypes::arrow::record_batch::RecordBatch;
760
761    use super::*;
762
763    /// Creates a test RecordBatch in flat format with given parameters.
764    fn new_record_batch(
765        primary_keys: &[&[u8]],
766        timestamps: &[i64],
767        sequences: &[u64],
768        op_types: &[OpType],
769        fields: &[u64],
770    ) -> RecordBatch {
771        let num_rows = timestamps.len();
772        debug_assert_eq!(sequences.len(), num_rows);
773        debug_assert_eq!(op_types.len(), num_rows);
774        debug_assert_eq!(fields.len(), num_rows);
775        debug_assert_eq!(primary_keys.len(), num_rows);
776
777        let columns: Vec<ArrayRef> = vec![
778            // k0 column (primary key as string dictionary)
779            build_test_pk_string_dict_array(primary_keys),
780            // field0 column
781            Arc::new(Int64Array::from_iter(
782                fields.iter().map(|v| Some(*v as i64)),
783            )),
784            // ts column (time index)
785            Arc::new(TimestampMillisecondArray::from_iter_values(
786                timestamps.iter().copied(),
787            )),
788            // __primary_key column
789            build_test_pk_array(primary_keys),
790            // __sequence column
791            Arc::new(UInt64Array::from_iter_values(sequences.iter().copied())),
792            // __op_type column
793            Arc::new(UInt8Array::from_iter_values(
794                op_types.iter().map(|v| *v as u8),
795            )),
796        ];
797
798        RecordBatch::try_new(build_test_flat_schema(), columns).unwrap()
799    }
800
801    /// Creates a test RecordBatch in flat format with multiple fields for testing FlatLastNonNull.
802    fn new_record_batch_multi_fields(
803        primary_keys: &[&[u8]],
804        timestamps: &[i64],
805        sequences: &[u64],
806        op_types: &[OpType],
807        fields: &[(Option<u64>, Option<u64>)],
808    ) -> RecordBatch {
809        let num_rows = timestamps.len();
810        debug_assert_eq!(sequences.len(), num_rows);
811        debug_assert_eq!(op_types.len(), num_rows);
812        debug_assert_eq!(fields.len(), num_rows);
813        debug_assert_eq!(primary_keys.len(), num_rows);
814
815        let columns: Vec<ArrayRef> = vec![
816            // k0 column (primary key as string dictionary)
817            build_test_pk_string_dict_array(primary_keys),
818            // field0 column
819            Arc::new(Int64Array::from_iter(
820                fields.iter().map(|field| field.0.map(|v| v as i64)),
821            )),
822            // field1 column
823            Arc::new(Int64Array::from_iter(
824                fields.iter().map(|field| field.1.map(|v| v as i64)),
825            )),
826            // ts column (time index)
827            Arc::new(TimestampMillisecondArray::from_iter_values(
828                timestamps.iter().copied(),
829            )),
830            // __primary_key column
831            build_test_pk_array(primary_keys),
832            // __sequence column
833            Arc::new(UInt64Array::from_iter_values(sequences.iter().copied())),
834            // __op_type column
835            Arc::new(UInt8Array::from_iter_values(
836                op_types.iter().map(|v| *v as u8),
837            )),
838        ];
839
840        RecordBatch::try_new(build_test_multi_field_schema(), columns).unwrap()
841    }
842
843    /// Creates a test string dictionary primary key array for given primary keys.
844    fn build_test_pk_string_dict_array(primary_keys: &[&[u8]]) -> ArrayRef {
845        let mut builder = StringDictionaryBuilder::<UInt32Type>::new();
846        for &pk in primary_keys {
847            let pk_str = std::str::from_utf8(pk).unwrap();
848            builder.append(pk_str).unwrap();
849        }
850        Arc::new(builder.finish())
851    }
852
853    /// Creates a test primary key array for given primary keys.
854    fn build_test_pk_array(primary_keys: &[&[u8]]) -> ArrayRef {
855        let mut builder = BinaryDictionaryBuilder::<UInt32Type>::new();
856        for &pk in primary_keys {
857            builder.append(pk).unwrap();
858        }
859        Arc::new(builder.finish())
860    }
861
862    /// Builds the arrow schema for test flat format.
863    fn build_test_flat_schema() -> SchemaRef {
864        let fields = vec![
865            Field::new(
866                "k0",
867                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
868                false,
869            ),
870            Field::new("field0", DataType::Int64, true),
871            Field::new(
872                "ts",
873                DataType::Timestamp(TimeUnit::Millisecond, None),
874                false,
875            ),
876            Field::new(
877                "__primary_key",
878                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
879                false,
880            ),
881            Field::new("__sequence", DataType::UInt64, false),
882            Field::new("__op_type", DataType::UInt8, false),
883        ];
884        Arc::new(Schema::new(fields))
885    }
886
887    /// Builds the arrow schema for test flat format with multiple fields.
888    fn build_test_multi_field_schema() -> SchemaRef {
889        let fields = vec![
890            Field::new(
891                "k0",
892                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
893                false,
894            ),
895            Field::new("field0", DataType::Int64, true),
896            Field::new("field1", DataType::Int64, true),
897            Field::new(
898                "ts",
899                DataType::Timestamp(TimeUnit::Millisecond, None),
900                false,
901            ),
902            Field::new(
903                "__primary_key",
904                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
905                false,
906            ),
907            Field::new("__sequence", DataType::UInt64, false),
908            Field::new("__op_type", DataType::UInt8, false),
909        ];
910        Arc::new(Schema::new(fields))
911    }
912
913    /// Asserts that two RecordBatch vectors are equal.
914    fn check_record_batches_equal(expected: &[RecordBatch], actual: &[RecordBatch]) {
915        for (i, (exp, act)) in expected.iter().zip(actual.iter()).enumerate() {
916            assert_eq!(exp, act, "RecordBatch {} differs", i);
917        }
918        assert_eq!(
919            expected.len(),
920            actual.len(),
921            "Number of batches don't match"
922        );
923    }
924
925    /// Helper function to collect iterator results.
926    fn collect_iterator_results<I>(iter: I) -> Vec<RecordBatch>
927    where
928        I: Iterator<Item = Result<RecordBatch>>,
929    {
930        iter.map(|result| result.unwrap()).collect()
931    }
932
933    #[test]
934    fn test_flat_last_row_no_duplications() {
935        let input = vec![
936            new_record_batch(
937                &[b"k1", b"k1"],
938                &[1, 2],
939                &[11, 12],
940                &[OpType::Put, OpType::Put],
941                &[21, 22],
942            ),
943            new_record_batch(&[b"k1"], &[3], &[13], &[OpType::Put], &[23]),
944            new_record_batch(
945                &[b"k2", b"k2"],
946                &[1, 2],
947                &[111, 112],
948                &[OpType::Put, OpType::Put],
949                &[31, 32],
950            ),
951        ];
952
953        // Test with filter_deleted = true
954        let iter = input.clone().into_iter().map(Ok);
955        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastRow::new(true));
956        let result = collect_iterator_results(&mut dedup_iter);
957        check_record_batches_equal(&input, &result);
958        assert_eq!(0, dedup_iter.metrics.num_unselected_rows);
959        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
960
961        // Test with filter_deleted = false
962        let iter = input.clone().into_iter().map(Ok);
963        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastRow::new(false));
964        let result = collect_iterator_results(&mut dedup_iter);
965        check_record_batches_equal(&input, &result);
966        assert_eq!(0, dedup_iter.metrics.num_unselected_rows);
967        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
968    }
969
970    #[test]
971    fn test_flat_last_row_duplications() {
972        let input = vec![
973            new_record_batch(
974                &[b"k1", b"k1"],
975                &[1, 2],
976                &[13, 11],
977                &[OpType::Put, OpType::Put],
978                &[11, 12],
979            ),
980            // empty batch.
981            new_record_batch(&[], &[], &[], &[], &[]),
982            // Duplicate with the previous batch.
983            new_record_batch(
984                &[b"k1", b"k1", b"k1"],
985                &[2, 3, 4],
986                &[10, 13, 13],
987                &[OpType::Put, OpType::Put, OpType::Delete],
988                &[2, 13, 14],
989            ),
990            new_record_batch(
991                &[b"k2", b"k2"],
992                &[1, 2],
993                &[20, 20],
994                &[OpType::Put, OpType::Delete],
995                &[101, 0],
996            ),
997            new_record_batch(&[b"k2"], &[2], &[19], &[OpType::Put], &[102]),
998            new_record_batch(&[b"k3"], &[2], &[20], &[OpType::Put], &[202]),
999            // This batch won't increase the deleted rows count as it
1000            // is filtered out by the previous batch.
1001            new_record_batch(&[b"k3"], &[2], &[19], &[OpType::Delete], &[0]),
1002        ];
1003
1004        // Test with filter_deleted = true
1005        let expected_filter_deleted = vec![
1006            new_record_batch(
1007                &[b"k1", b"k1"],
1008                &[1, 2],
1009                &[13, 11],
1010                &[OpType::Put, OpType::Put],
1011                &[11, 12],
1012            ),
1013            new_record_batch(&[b"k1"], &[3], &[13], &[OpType::Put], &[13]),
1014            new_record_batch(&[b"k2"], &[1], &[20], &[OpType::Put], &[101]),
1015            new_record_batch(&[b"k3"], &[2], &[20], &[OpType::Put], &[202]),
1016        ];
1017
1018        let iter = input.clone().into_iter().map(Ok);
1019        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastRow::new(true));
1020        let result = collect_iterator_results(&mut dedup_iter);
1021        check_record_batches_equal(&expected_filter_deleted, &result);
1022        assert_eq!(5, dedup_iter.metrics.num_unselected_rows);
1023        assert_eq!(2, dedup_iter.metrics.num_deleted_rows);
1024
1025        // Test with filter_deleted = false
1026        let expected_no_filter = vec![
1027            new_record_batch(
1028                &[b"k1", b"k1"],
1029                &[1, 2],
1030                &[13, 11],
1031                &[OpType::Put, OpType::Put],
1032                &[11, 12],
1033            ),
1034            new_record_batch(
1035                &[b"k1", b"k1"],
1036                &[3, 4],
1037                &[13, 13],
1038                &[OpType::Put, OpType::Delete],
1039                &[13, 14],
1040            ),
1041            new_record_batch(
1042                &[b"k2", b"k2"],
1043                &[1, 2],
1044                &[20, 20],
1045                &[OpType::Put, OpType::Delete],
1046                &[101, 0],
1047            ),
1048            new_record_batch(&[b"k3"], &[2], &[20], &[OpType::Put], &[202]),
1049        ];
1050
1051        let iter = input.clone().into_iter().map(Ok);
1052        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastRow::new(false));
1053        let result = collect_iterator_results(&mut dedup_iter);
1054        check_record_batches_equal(&expected_no_filter, &result);
1055        assert_eq!(3, dedup_iter.metrics.num_unselected_rows);
1056        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
1057    }
1058
1059    #[test]
1060    fn exact_sequence_filter_component_composition_covers_tombstones_and_last_non_null() {
1061        use store_api::storage::SequenceRange;
1062
1063        use crate::read::scan_util::filter_flat_batch_by_sequence;
1064
1065        let filter = |batch| {
1066            filter_flat_batch_by_sequence(
1067                batch,
1068                Some(SequenceRange::GtLtEq { min: 1, max: 3 }),
1069                true,
1070            )
1071            .unwrap()
1072            .unwrap()
1073        };
1074
1075        // The shared merge presents duplicate rows newest-first: the in-range
1076        // tombstone (seq 3) reaches shared dedup and suppresses the eligible
1077        // put (seq 2) across batches.
1078        let output = FlatDedupIterator::new(
1079            vec![
1080                Ok(filter(new_record_batch_multi_fields(
1081                    &[b"series"],
1082                    &[1000],
1083                    &[3],
1084                    &[OpType::Delete],
1085                    &[(None, None)],
1086                ))),
1087                Ok(filter(new_record_batch_multi_fields(
1088                    &[b"series"],
1089                    &[1000],
1090                    &[2],
1091                    &[OpType::Put],
1092                    &[(Some(10), None)],
1093                ))),
1094            ]
1095            .into_iter(),
1096            FlatLastNonNull::new(1, true),
1097        )
1098        .collect::<Result<Vec<_>>>()
1099        .unwrap();
1100        assert!(output.is_empty());
1101
1102        // A newest outside-range tombstone is removed before shared dedup and
1103        // cannot suppress the eligible put.
1104        let output = FlatDedupIterator::new(
1105            vec![Ok(filter(new_record_batch_multi_fields(
1106                &[b"series", b"series"],
1107                &[1000, 1000],
1108                &[4, 2],
1109                &[OpType::Delete, OpType::Put],
1110                &[(None, None), (Some(10), None)],
1111            )))]
1112            .into_iter(),
1113            FlatLastNonNull::new(1, true),
1114        )
1115        .collect::<Result<Vec<_>>>()
1116        .unwrap();
1117        assert_eq!(1, output[0].num_rows());
1118
1119        // A newest outside-range non-null value is removed before
1120        // LastNonNull, so it cannot fill the eligible row's null field.
1121        let output = FlatDedupIterator::new(
1122            vec![Ok(filter(new_record_batch_multi_fields(
1123                &[b"series", b"series"],
1124                &[1000, 1000],
1125                &[4, 2],
1126                &[OpType::Put, OpType::Put],
1127                &[(None, Some(40)), (Some(10), None)],
1128            )))]
1129            .into_iter(),
1130            FlatLastNonNull::new(1, true),
1131        )
1132        .collect::<Result<Vec<_>>>()
1133        .unwrap();
1134        assert_eq!(1, output[0].num_rows());
1135        assert!(output[0].column(2).is_null(0));
1136    }
1137
1138    #[test]
1139    fn test_flat_last_non_null_no_duplications() {
1140        let input = vec![
1141            new_record_batch(
1142                &[b"k1", b"k1"],
1143                &[1, 2],
1144                &[11, 12],
1145                &[OpType::Put, OpType::Put],
1146                &[21, 22],
1147            ),
1148            new_record_batch(&[b"k1"], &[3], &[13], &[OpType::Put], &[23]),
1149            new_record_batch(
1150                &[b"k2", b"k2"],
1151                &[1, 2],
1152                &[111, 112],
1153                &[OpType::Put, OpType::Put],
1154                &[31, 32],
1155            ),
1156        ];
1157
1158        // Test with filter_deleted = true
1159        let iter = input.clone().into_iter().map(Ok);
1160        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, true));
1161        let result = collect_iterator_results(&mut dedup_iter);
1162        check_record_batches_equal(&input, &result);
1163        assert_eq!(0, dedup_iter.metrics.num_unselected_rows);
1164        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
1165
1166        // Test with filter_deleted = false
1167        let iter = input.clone().into_iter().map(Ok);
1168        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, false));
1169        let result = collect_iterator_results(&mut dedup_iter);
1170        check_record_batches_equal(&input, &result);
1171        assert_eq!(0, dedup_iter.metrics.num_unselected_rows);
1172        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
1173    }
1174
1175    #[test]
1176    fn test_flat_last_non_null_field_merging() {
1177        let input = vec![
1178            new_record_batch_multi_fields(
1179                &[b"k1", b"k1"],
1180                &[1, 2],
1181                &[13, 11],
1182                &[OpType::Put, OpType::Put],
1183                &[(Some(11), Some(11)), (None, None)],
1184            ),
1185            // empty batch
1186            new_record_batch_multi_fields(&[], &[], &[], &[], &[]),
1187            // Duplicate with the previous batch - should merge fields
1188            new_record_batch_multi_fields(
1189                &[b"k1"],
1190                &[2],
1191                &[10],
1192                &[OpType::Put],
1193                &[(Some(12), None)],
1194            ),
1195            new_record_batch_multi_fields(
1196                &[b"k1", b"k1", b"k1"],
1197                &[2, 3, 4],
1198                &[10, 13, 13],
1199                &[OpType::Put, OpType::Put, OpType::Delete],
1200                &[(Some(2), Some(22)), (Some(13), None), (None, Some(14))],
1201            ),
1202            new_record_batch_multi_fields(
1203                &[b"k2", b"k2"],
1204                &[1, 2],
1205                &[20, 20],
1206                &[OpType::Put, OpType::Delete],
1207                &[(Some(101), Some(101)), (None, None)],
1208            ),
1209            new_record_batch_multi_fields(
1210                &[b"k2"],
1211                &[2],
1212                &[19],
1213                &[OpType::Put],
1214                &[(Some(102), Some(102))],
1215            ),
1216            new_record_batch_multi_fields(
1217                &[b"k3"],
1218                &[2],
1219                &[20],
1220                &[OpType::Put],
1221                &[(Some(202), Some(202))],
1222            ),
1223            // This batch won't increase the deleted rows count as it
1224            // is filtered out by the previous batch. (All fields are null).
1225            new_record_batch_multi_fields(
1226                &[b"k3"],
1227                &[2],
1228                &[19],
1229                &[OpType::Delete],
1230                &[(None, None)],
1231            ),
1232        ];
1233
1234        // Test with filter_deleted = true
1235        let expected_filter_deleted = vec![
1236            new_record_batch_multi_fields(
1237                &[b"k1"],
1238                &[1],
1239                &[13],
1240                &[OpType::Put],
1241                &[(Some(11), Some(11))],
1242            ),
1243            new_record_batch_multi_fields(
1244                &[b"k1", b"k1"],
1245                &[2, 3],
1246                &[11, 13],
1247                &[OpType::Put, OpType::Put],
1248                &[(Some(12), Some(22)), (Some(13), None)],
1249            ),
1250            new_record_batch_multi_fields(
1251                &[b"k2"],
1252                &[1],
1253                &[20],
1254                &[OpType::Put],
1255                &[(Some(101), Some(101))],
1256            ),
1257            new_record_batch_multi_fields(
1258                &[b"k3"],
1259                &[2],
1260                &[20],
1261                &[OpType::Put],
1262                &[(Some(202), Some(202))],
1263            ),
1264        ];
1265
1266        let iter = input.clone().into_iter().map(Ok);
1267        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, true));
1268        let result = collect_iterator_results(&mut dedup_iter);
1269        check_record_batches_equal(&expected_filter_deleted, &result);
1270        assert_eq!(6, dedup_iter.metrics.num_unselected_rows);
1271        assert_eq!(2, dedup_iter.metrics.num_deleted_rows);
1272    }
1273
1274    #[test]
1275    fn test_flat_last_non_null_skip_merge_no_null() {
1276        let input = vec![
1277            new_record_batch_multi_fields(
1278                &[b"k1", b"k1"],
1279                &[1, 2],
1280                &[13, 11],
1281                &[OpType::Put, OpType::Put],
1282                &[(Some(11), Some(11)), (Some(12), Some(12))],
1283            ),
1284            new_record_batch_multi_fields(
1285                &[b"k1"],
1286                &[2],
1287                &[10],
1288                &[OpType::Put],
1289                &[(None, Some(22))],
1290            ),
1291            new_record_batch_multi_fields(
1292                &[b"k1", b"k1"],
1293                &[2, 3],
1294                &[9, 13],
1295                &[OpType::Put, OpType::Put],
1296                &[(Some(32), None), (Some(13), Some(13))],
1297            ),
1298        ];
1299
1300        let expected = vec![
1301            new_record_batch_multi_fields(
1302                &[b"k1"],
1303                &[1],
1304                &[13],
1305                &[OpType::Put],
1306                &[(Some(11), Some(11))],
1307            ),
1308            new_record_batch_multi_fields(
1309                &[b"k1", b"k1"],
1310                &[2, 3],
1311                &[11, 13],
1312                &[OpType::Put, OpType::Put],
1313                &[(Some(12), Some(12)), (Some(13), Some(13))],
1314            ),
1315        ];
1316
1317        let iter = input.into_iter().map(Ok);
1318        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, true));
1319        let result = collect_iterator_results(&mut dedup_iter);
1320        check_record_batches_equal(&expected, &result);
1321        assert_eq!(2, dedup_iter.metrics.num_unselected_rows);
1322        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
1323    }
1324
1325    #[test]
1326    fn test_flat_last_non_null_merge_null() {
1327        let input = vec![
1328            new_record_batch_multi_fields(
1329                &[b"k1", b"k1"],
1330                &[1, 2],
1331                &[13, 11],
1332                &[OpType::Put, OpType::Put],
1333                &[(Some(11), Some(11)), (None, None)],
1334            ),
1335            new_record_batch_multi_fields(
1336                &[b"k1"],
1337                &[2],
1338                &[10],
1339                &[OpType::Put],
1340                &[(None, Some(22))],
1341            ),
1342            new_record_batch_multi_fields(
1343                &[b"k1"],
1344                &[3],
1345                &[13],
1346                &[OpType::Put],
1347                &[(Some(33), None)],
1348            ),
1349        ];
1350
1351        let expected = vec![
1352            new_record_batch_multi_fields(
1353                &[b"k1"],
1354                &[1],
1355                &[13],
1356                &[OpType::Put],
1357                &[(Some(11), Some(11))],
1358            ),
1359            new_record_batch_multi_fields(
1360                &[b"k1"],
1361                &[2],
1362                &[11],
1363                &[OpType::Put],
1364                &[(None, Some(22))],
1365            ),
1366            new_record_batch_multi_fields(
1367                &[b"k1"],
1368                &[3],
1369                &[13],
1370                &[OpType::Put],
1371                &[(Some(33), None)],
1372            ),
1373        ];
1374
1375        let iter = input.into_iter().map(Ok);
1376        let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, true));
1377        let result = collect_iterator_results(&mut dedup_iter);
1378        check_record_batches_equal(&expected, &result);
1379        assert_eq!(1, dedup_iter.metrics.num_unselected_rows);
1380        assert_eq!(0, dedup_iter.metrics.num_deleted_rows);
1381    }
1382
1383    /// Helper function to check dedup strategy behavior directly.
1384    fn check_flat_dedup_strategy(
1385        input: &[RecordBatch],
1386        strategy: &mut dyn RecordBatchDedupStrategy,
1387        expect: &[RecordBatch],
1388    ) {
1389        let mut actual = Vec::new();
1390        let mut metrics = DedupMetrics::default();
1391        for batch in input {
1392            if let Some(out) = strategy.push_batch(batch.clone(), &mut metrics).unwrap() {
1393                actual.push(out);
1394            }
1395        }
1396        if let Some(out) = strategy.finish(&mut metrics).unwrap() {
1397            actual.push(out);
1398        }
1399
1400        check_record_batches_equal(expect, &actual);
1401    }
1402
1403    #[test]
1404    fn test_flat_last_non_null_strategy_delete_last() {
1405        let input = vec![
1406            new_record_batch_multi_fields(
1407                &[b"k1"],
1408                &[1],
1409                &[6],
1410                &[OpType::Put],
1411                &[(Some(11), None)],
1412            ),
1413            new_record_batch_multi_fields(
1414                &[b"k1", b"k1"],
1415                &[1, 2],
1416                &[1, 7],
1417                &[OpType::Put, OpType::Put],
1418                &[(Some(1), None), (Some(22), Some(222))],
1419            ),
1420            new_record_batch_multi_fields(
1421                &[b"k1"],
1422                &[2],
1423                &[4],
1424                &[OpType::Put],
1425                &[(Some(12), None)],
1426            ),
1427            new_record_batch_multi_fields(
1428                &[b"k2", b"k2"],
1429                &[2, 3],
1430                &[2, 5],
1431                &[OpType::Put, OpType::Delete],
1432                &[(None, None), (Some(13), None)],
1433            ),
1434            new_record_batch_multi_fields(&[b"k2"], &[3], &[3], &[OpType::Put], &[(None, Some(3))]),
1435        ];
1436
1437        let mut strategy = FlatLastNonNull::new(1, true);
1438        check_flat_dedup_strategy(
1439            &input,
1440            &mut strategy,
1441            &[
1442                new_record_batch_multi_fields(
1443                    &[b"k1"],
1444                    &[1],
1445                    &[6],
1446                    &[OpType::Put],
1447                    &[(Some(11), None)],
1448                ),
1449                new_record_batch_multi_fields(
1450                    &[b"k1"],
1451                    &[2],
1452                    &[7],
1453                    &[OpType::Put],
1454                    &[(Some(22), Some(222))],
1455                ),
1456                new_record_batch_multi_fields(
1457                    &[b"k2"],
1458                    &[2],
1459                    &[2],
1460                    &[OpType::Put],
1461                    &[(None, None)],
1462                ),
1463            ],
1464        );
1465    }
1466
1467    #[test]
1468    fn test_flat_last_non_null_strategy_delete_one() {
1469        let input = vec![
1470            new_record_batch_multi_fields(&[b"k1"], &[1], &[1], &[OpType::Delete], &[(None, None)]),
1471            new_record_batch_multi_fields(
1472                &[b"k2"],
1473                &[1],
1474                &[6],
1475                &[OpType::Put],
1476                &[(Some(11), None)],
1477            ),
1478        ];
1479
1480        let mut strategy = FlatLastNonNull::new(1, true);
1481        check_flat_dedup_strategy(
1482            &input,
1483            &mut strategy,
1484            &[new_record_batch_multi_fields(
1485                &[b"k2"],
1486                &[1],
1487                &[6],
1488                &[OpType::Put],
1489                &[(Some(11), None)],
1490            )],
1491        );
1492    }
1493
1494    #[test]
1495    fn test_flat_last_non_null_strategy_delete_all() {
1496        let input = vec![
1497            new_record_batch_multi_fields(&[b"k1"], &[1], &[1], &[OpType::Delete], &[(None, None)]),
1498            new_record_batch_multi_fields(
1499                &[b"k2"],
1500                &[1],
1501                &[6],
1502                &[OpType::Delete],
1503                &[(Some(11), None)],
1504            ),
1505        ];
1506
1507        let mut strategy = FlatLastNonNull::new(1, true);
1508        check_flat_dedup_strategy(&input, &mut strategy, &[]);
1509    }
1510
1511    #[test]
1512    fn test_flat_last_non_null_strategy_same_batch() {
1513        let input = vec![
1514            new_record_batch_multi_fields(
1515                &[b"k1"],
1516                &[1],
1517                &[6],
1518                &[OpType::Put],
1519                &[(Some(11), None)],
1520            ),
1521            new_record_batch_multi_fields(
1522                &[b"k1", b"k1"],
1523                &[1, 2],
1524                &[1, 7],
1525                &[OpType::Put, OpType::Put],
1526                &[(Some(1), None), (Some(22), Some(222))],
1527            ),
1528            new_record_batch_multi_fields(
1529                &[b"k1"],
1530                &[2],
1531                &[4],
1532                &[OpType::Put],
1533                &[(Some(12), None)],
1534            ),
1535            new_record_batch_multi_fields(
1536                &[b"k1", b"k1"],
1537                &[2, 3],
1538                &[2, 5],
1539                &[OpType::Put, OpType::Put],
1540                &[(None, None), (Some(13), None)],
1541            ),
1542            new_record_batch_multi_fields(&[b"k1"], &[3], &[3], &[OpType::Put], &[(None, Some(3))]),
1543        ];
1544
1545        let mut strategy = FlatLastNonNull::new(1, true);
1546        check_flat_dedup_strategy(
1547            &input,
1548            &mut strategy,
1549            &[
1550                new_record_batch_multi_fields(
1551                    &[b"k1"],
1552                    &[1],
1553                    &[6],
1554                    &[OpType::Put],
1555                    &[(Some(11), None)],
1556                ),
1557                new_record_batch_multi_fields(
1558                    &[b"k1"],
1559                    &[2],
1560                    &[7],
1561                    &[OpType::Put],
1562                    &[(Some(22), Some(222))],
1563                ),
1564                new_record_batch_multi_fields(
1565                    &[b"k1"],
1566                    &[3],
1567                    &[5],
1568                    &[OpType::Put],
1569                    &[(Some(13), Some(3))],
1570                ),
1571            ],
1572        );
1573    }
1574
1575    #[test]
1576    fn test_last_non_null_preserves_delete_barriers() {
1577        let input = new_record_batch_multi_fields(
1578            &[b"k1".as_slice(); 9],
1579            &[1; 9],
1580            &[9, 8, 7, 6, 5, 4, 3, 2, 1],
1581            &[
1582                OpType::Put,
1583                OpType::Put,
1584                OpType::Delete,
1585                OpType::Delete,
1586                OpType::Put,
1587                OpType::Put,
1588                OpType::Delete,
1589                OpType::Put,
1590                OpType::Put,
1591            ],
1592            &[
1593                (None, Some(9)),
1594                (Some(8), None),
1595                (None, None),
1596                (None, None),
1597                (None, Some(5)),
1598                (Some(4), None),
1599                (None, None),
1600                (None, None),
1601                (Some(1), Some(1)),
1602            ],
1603        );
1604        let expected = new_record_batch_multi_fields(
1605            &[b"k1".as_slice(); 6],
1606            &[1; 6],
1607            &[9, 7, 6, 5, 3, 2],
1608            &[
1609                OpType::Put,
1610                OpType::Delete,
1611                OpType::Delete,
1612                OpType::Put,
1613                OpType::Delete,
1614                OpType::Put,
1615            ],
1616            &[
1617                (Some(8), Some(9)),
1618                (None, None),
1619                (None, None),
1620                (Some(4), Some(5)),
1621                (None, None),
1622                (Some(1), Some(1)),
1623            ],
1624        );
1625        // Check both intra-batch partitioning and barriers split across batches.
1626        for batch_size in [1, input.num_rows()] {
1627            let batches = (0..input.num_rows())
1628                .step_by(batch_size)
1629                .map(|start| Ok(input.slice(start, batch_size.min(input.num_rows() - start))));
1630            let output = FlatDedupIterator::new(batches, FlatLastNonNull::new(1, false))
1631                .collect::<Result<Vec<_>>>()
1632                .unwrap();
1633            let actual = concat_batches(&input.schema(), &output).unwrap();
1634            assert_eq!(expected, actual);
1635        }
1636    }
1637
1638    #[test]
1639    fn test_flat_last_non_null_strategy_delete_middle() {
1640        let input = vec![
1641            new_record_batch_multi_fields(
1642                &[b"k1"],
1643                &[1],
1644                &[7],
1645                &[OpType::Put],
1646                &[(Some(11), None)],
1647            ),
1648            new_record_batch_multi_fields(&[b"k1"], &[1], &[4], &[OpType::Delete], &[(None, None)]),
1649            new_record_batch_multi_fields(
1650                &[b"k1"],
1651                &[1],
1652                &[1],
1653                &[OpType::Put],
1654                &[(Some(12), Some(1))],
1655            ),
1656            new_record_batch_multi_fields(
1657                &[b"k1"],
1658                &[2],
1659                &[8],
1660                &[OpType::Put],
1661                &[(Some(21), None)],
1662            ),
1663            new_record_batch_multi_fields(&[b"k1"], &[2], &[5], &[OpType::Delete], &[(None, None)]),
1664            new_record_batch_multi_fields(
1665                &[b"k1"],
1666                &[2],
1667                &[2],
1668                &[OpType::Put],
1669                &[(Some(22), Some(2))],
1670            ),
1671            new_record_batch_multi_fields(
1672                &[b"k1"],
1673                &[3],
1674                &[9],
1675                &[OpType::Put],
1676                &[(Some(31), None)],
1677            ),
1678            new_record_batch_multi_fields(&[b"k1"], &[3], &[6], &[OpType::Delete], &[(None, None)]),
1679            new_record_batch_multi_fields(
1680                &[b"k1"],
1681                &[3],
1682                &[3],
1683                &[OpType::Put],
1684                &[(Some(32), Some(3))],
1685            ),
1686        ];
1687
1688        let mut strategy = FlatLastNonNull::new(1, true);
1689        check_flat_dedup_strategy(
1690            &input,
1691            &mut strategy,
1692            &[
1693                new_record_batch_multi_fields(
1694                    &[b"k1"],
1695                    &[1],
1696                    &[7],
1697                    &[OpType::Put],
1698                    &[(Some(11), None)],
1699                ),
1700                new_record_batch_multi_fields(
1701                    &[b"k1"],
1702                    &[2],
1703                    &[8],
1704                    &[OpType::Put],
1705                    &[(Some(21), None)],
1706                ),
1707                new_record_batch_multi_fields(
1708                    &[b"k1"],
1709                    &[3],
1710                    &[9],
1711                    &[OpType::Put],
1712                    &[(Some(31), None)],
1713                ),
1714            ],
1715        );
1716    }
1717}
mito2/read/flat_dedup.rs

mito2/read/
flat_dedup.rs
