mito2/sst/range_index/
writer.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::cmp::Ordering;
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18
19use datatypes::arrow::array::{
20    Array, ArrayRef, BinaryArray, DictionaryArray, Int64Array, UInt32Array, UInt64Array,
21};
22use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type};
23use datatypes::arrow::record_batch::RecordBatch;
24use datatypes::prelude::ConcreteDataType;
25use mito_codec::row_converter::SparsePrimaryKeyCodec;
26use object_store::ObjectStore;
27use snafu::{OptionExt, ResultExt, ensure};
28use store_api::codec::PrimaryKeyEncoding;
29use store_api::metadata::RegionMetadataRef;
30use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId};
31
32use crate::error::{
33    DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result,
34    UnexpectedSnafu,
35};
36use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE;
37use crate::sst::parquet::index_writer::ParquetIndexWriter;
38use crate::sst::range_index::{
39    END_COLUMN, ROW_GROUP_ID_COLUMN, START_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN,
40};
41
42const WRITE_BATCH_SIZE: usize = 1024;
43
44/// Options for writing a per-SST range index.
45#[derive(Debug, Clone)]
46pub struct SstRangeIndexWriterOptions {
47    /// Maximum number of range entries in an index Parquet row group.
48    pub index_row_group_size: usize,
49}
50
51impl Default for SstRangeIndexWriterOptions {
52    fn default() -> Self {
53        Self {
54            index_row_group_size: DEFAULT_ROW_GROUP_SIZE,
55        }
56    }
57}
58
59/// Metrics collected by an [`SstRangeIndexWriter`].
60#[derive(Debug, Clone, Default)]
61pub struct SstRangeIndexWriterMetrics {
62    /// Number of input record batches passed to the writer.
63    pub input_batches: usize,
64    /// Number of source SST rows passed to the writer.
65    pub input_rows: usize,
66    /// Number of non-empty source SST row groups passed to the writer.
67    pub num_source_row_groups: usize,
68    /// Number of series ranges written to the index.
69    pub num_ranges: usize,
70    /// Size of the committed index file. This remains zero for an aborted writer.
71    pub output_bytes: u64,
72    /// Time spent opening the object-store and Parquet writers.
73    pub open_elapsed: Duration,
74    /// Time spent validating input and aggregating ranges.
75    pub aggregate_elapsed: Duration,
76    /// Time spent encoding and writing Parquet batches.
77    pub write_elapsed: Duration,
78    /// Time spent closing a completed file.
79    pub finish_elapsed: Duration,
80    /// Time spent removing incomplete output.
81    pub cleanup_elapsed: Duration,
82    /// Whether this writer was explicitly aborted.
83    pub aborted: bool,
84}
85
86impl SstRangeIndexWriterMetrics {
87    /// Returns time spent by writer-owned work.
88    pub fn total_elapsed(&self) -> Duration {
89        self.open_elapsed
90            + self.aggregate_elapsed
91            + self.write_elapsed
92            + self.finish_elapsed
93            + self.cleanup_elapsed
94    }
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98struct RangeIndexRow {
99    row_group_id: u32,
100    table_id: u32,
101    tsid: u64,
102    start: i64,
103    end: i64,
104}
105
106/// Incrementally writes the series row ranges of one SST to a Parquet file.
107///
108/// Each non-empty input batch is tagged with its source SST row-group ID. A row
109/// group may be supplied in multiple batches, but IDs must start at zero and be
110/// contiguous. Input primary keys must preserve the SST's global sort order.
111pub struct SstRangeIndexWriter {
112    codec: SparsePrimaryKeyCodec,
113    schema: SchemaRef,
114    writer: ParquetIndexWriter,
115    current_row_group_id: Option<u32>,
116    current_row_group_offset: i64,
117    last_primary_key: Option<Vec<u8>>,
118    current_row: Option<RangeIndexRow>,
119    buffered_rows: Vec<RangeIndexRow>,
120    metrics: SstRangeIndexWriterMetrics,
121    failed: bool,
122}
123
124impl SstRangeIndexWriter {
125    /// Creates a writer for `path` in `object_store`.
126    ///
127    /// The file name in `path` must be unique in the object store so aborting
128    /// the writer only removes temporary files that belong to this writer.
129    pub async fn try_new(
130        metadata: RegionMetadataRef,
131        object_store: ObjectStore,
132        path: &str,
133        options: SstRangeIndexWriterOptions,
134    ) -> Result<Self> {
135        let open_start = Instant::now();
136        ensure!(
137            options.index_row_group_size > 0,
138            InvalidMetaSnafu {
139                reason: "range index row group size must be greater than zero",
140            }
141        );
142        validate_metadata(&metadata)?;
143        let schema = range_index_schema();
144        let writer = ParquetIndexWriter::try_new(
145            "range index",
146            object_store,
147            path,
148            &schema,
149            options.index_row_group_size,
150            None,
151        )
152        .await?;
153        let codec = SparsePrimaryKeyCodec::new(&metadata);
154
155        Ok(Self {
156            codec,
157            schema,
158            writer,
159            current_row_group_id: None,
160            current_row_group_offset: 0,
161            last_primary_key: None,
162            current_row: None,
163            buffered_rows: Vec::with_capacity(WRITE_BATCH_SIZE),
164            metrics: SstRangeIndexWriterMetrics {
165                open_elapsed: open_start.elapsed(),
166                ..Default::default()
167            },
168            failed: false,
169        })
170    }
171
172    /// Returns the metrics collected so far.
173    pub fn metrics(&self) -> &SstRangeIndexWriterMetrics {
174        &self.metrics
175    }
176
177    /// Adds one batch from `row_group_id` of the source SST.
178    ///
179    /// The batch may be a primary-key-only projection or a full flat batch. Its
180    /// schema must contain a named `__primary_key` column. After this method
181    /// returns an error, callers must call [`Self::abort`].
182    pub async fn write(&mut self, row_group_id: u32, batch: &RecordBatch) -> Result<()> {
183        ensure!(
184            !self.failed,
185            InvalidRecordBatchSnafu {
186                reason: "cannot write to a failed range index writer",
187            }
188        );
189
190        self.metrics.input_batches += 1;
191        self.metrics.input_rows += batch.num_rows();
192        let aggregate_start = Instant::now();
193        let write_before = self.metrics.write_elapsed;
194        let result = self.write_inner(row_group_id, batch).await;
195        let write_cost = self.metrics.write_elapsed.saturating_sub(write_before);
196        self.metrics.aggregate_elapsed += aggregate_start.elapsed().saturating_sub(write_cost);
197        if result.is_err() {
198            self.failed = true;
199        }
200        result
201    }
202
203    /// Finishes and commits the index file.
204    pub async fn finish(mut self) -> Result<SstRangeIndexWriterMetrics> {
205        if self.failed {
206            let error = InvalidRecordBatchSnafu {
207                reason: "cannot finish a failed range index writer",
208            }
209            .build();
210            self.cleanup().await;
211            return Err(error);
212        }
213
214        let result = self.finish_inner().await;
215        if result.is_err() {
216            self.cleanup().await;
217        }
218        result.map(|_| self.metrics)
219    }
220
221    /// Aborts the writer and removes incomplete output files.
222    pub async fn abort(mut self) -> Result<SstRangeIndexWriterMetrics> {
223        self.metrics.aborted = true;
224        self.cleanup().await;
225        Ok(self.metrics)
226    }
227
228    async fn write_inner(&mut self, row_group_id: u32, batch: &RecordBatch) -> Result<()> {
229        if batch.num_rows() == 0 {
230            return Ok(());
231        }
232
233        self.enter_row_group(row_group_id).await?;
234        let pk_idx = batch
235            .schema()
236            .column_with_name(PRIMARY_KEY_COLUMN_NAME)
237            .map(|(idx, _)| idx)
238            .context(InvalidRecordBatchSnafu {
239                reason: "range index input does not contain __primary_key",
240            })?;
241        let primary_keys = batch.column(pk_idx);
242        let batch_rows = i64::try_from(batch.num_rows()).map_err(|_| {
243            InvalidRecordBatchSnafu {
244                reason: format!(
245                    "range index input batch exceeds Int64: {}",
246                    batch.num_rows()
247                ),
248            }
249            .build()
250        })?;
251        self.current_row_group_offset
252            .checked_add(batch_rows)
253            .context(InvalidRecordBatchSnafu {
254                reason: "source SST row-group offset exceeds Int64",
255            })?;
256
257        if let Some(array) = primary_keys.as_any().downcast_ref::<BinaryArray>() {
258            ensure!(
259                array.null_count() == 0,
260                InvalidRecordBatchSnafu {
261                    reason: "range index input contains null primary keys",
262                }
263            );
264            self.write_binary_primary_keys(array).await
265        } else if let Some(array) = primary_keys
266            .as_any()
267            .downcast_ref::<DictionaryArray<UInt32Type>>()
268        {
269            ensure!(
270                array.null_count() == 0,
271                InvalidRecordBatchSnafu {
272                    reason: "range index input contains null primary keys",
273                }
274            );
275            self.write_dictionary_primary_keys(array).await
276        } else {
277            InvalidRecordBatchSnafu {
278                reason: format!(
279                    "range index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}",
280                    primary_keys.data_type()
281                ),
282            }
283            .fail()
284        }
285    }
286
287    async fn enter_row_group(&mut self, row_group_id: u32) -> Result<()> {
288        let Some(current) = self.current_row_group_id else {
289            ensure!(
290                row_group_id == 0,
291                InvalidRecordBatchSnafu {
292                    reason: format!(
293                        "range index row groups must start at zero, got {row_group_id}"
294                    ),
295                }
296            );
297            self.current_row_group_id = Some(row_group_id);
298            self.metrics.num_source_row_groups = 1;
299            return Ok(());
300        };
301
302        ensure!(
303            row_group_id == current || current.checked_add(1) == Some(row_group_id),
304            InvalidRecordBatchSnafu {
305                reason: format!(
306                    "range index row groups must be contiguous, current {current}, got {row_group_id}"
307                ),
308            }
309        );
310        if row_group_id != current {
311            self.finish_current_row().await?;
312            self.current_row_group_id = Some(row_group_id);
313            self.current_row_group_offset = 0;
314            self.metrics.num_source_row_groups += 1;
315        }
316        Ok(())
317    }
318
319    async fn write_binary_primary_keys(&mut self, primary_keys: &BinaryArray) -> Result<()> {
320        let mut start = 0;
321        while start < primary_keys.len() {
322            let primary_key = primary_keys.value(start);
323            let mut end = start + 1;
324            while end < primary_keys.len() && primary_keys.value(end) == primary_key {
325                end += 1;
326            }
327            self.update_primary_key(primary_key, end - start).await?;
328            start = end;
329        }
330        Ok(())
331    }
332
333    async fn write_dictionary_primary_keys(
334        &mut self,
335        primary_keys: &DictionaryArray<UInt32Type>,
336    ) -> Result<()> {
337        let values = primary_keys
338            .values()
339            .as_any()
340            .downcast_ref::<BinaryArray>()
341            .context(InvalidRecordBatchSnafu {
342                reason: "range index primary-key dictionary values are not binary",
343            })?;
344        ensure!(
345            values.null_count() == 0,
346            InvalidRecordBatchSnafu {
347                reason: "range index input contains null primary-key dictionary values",
348            }
349        );
350        let keys = primary_keys.keys().values();
351        let mut start = 0;
352        while start < keys.len() {
353            let key = keys[start];
354            let mut end = start + 1;
355            while end < keys.len() && keys[end] == key {
356                end += 1;
357            }
358            self.update_primary_key(values.value(key as usize), end - start)
359                .await?;
360            start = end;
361        }
362        Ok(())
363    }
364
365    async fn update_primary_key(&mut self, primary_key: &[u8], run_len: usize) -> Result<()> {
366        if let Some(last) = self.last_primary_key.as_deref() {
367            ensure!(
368                primary_key.cmp(last) != Ordering::Less,
369                InvalidRecordBatchSnafu {
370                    reason: "range index input is not sorted by primary key",
371                }
372            );
373        }
374
375        let (table_id, tsid) = self.codec.decode_ids(primary_key).context(DecodeSnafu)?;
376        let run_len = i64::try_from(run_len).map_err(|_| {
377            InvalidRecordBatchSnafu {
378                reason: format!("range index primary-key run is too large: {run_len}"),
379            }
380            .build()
381        })?;
382        let start = self.current_row_group_offset;
383        let end = start
384            .checked_add(run_len)
385            .context(InvalidRecordBatchSnafu {
386                reason: "source SST row-group offset exceeds Int64",
387            })?;
388        let row_group_id = self.current_row_group_id.context(UnexpectedSnafu {
389            reason: "range index writer has no active source row group",
390        })?;
391
392        match self.current_row.as_mut() {
393            Some(row)
394                if row.row_group_id == row_group_id
395                    && row.table_id == table_id
396                    && row.tsid == tsid =>
397            {
398                row.end = end;
399            }
400            _ => {
401                self.finish_current_row().await?;
402                self.current_row = Some(RangeIndexRow {
403                    row_group_id,
404                    table_id,
405                    tsid,
406                    start,
407                    end,
408                });
409            }
410        }
411        self.current_row_group_offset = end;
412        self.last_primary_key = Some(primary_key.to_vec());
413        Ok(())
414    }
415
416    async fn finish_current_row(&mut self) -> Result<()> {
417        if let Some(row) = self.current_row.take() {
418            self.buffered_rows.push(row);
419            self.metrics.num_ranges += 1;
420        }
421        if self.buffered_rows.len() >= WRITE_BATCH_SIZE {
422            self.flush_rows().await?;
423        }
424        Ok(())
425    }
426
427    async fn flush_rows(&mut self) -> Result<()> {
428        if self.buffered_rows.is_empty() {
429            return Ok(());
430        }
431        let batch = rows_to_batch(&self.schema, &self.buffered_rows)?;
432        let start = Instant::now();
433        let result = self.writer.write(&batch).await;
434        self.metrics.write_elapsed += start.elapsed();
435        result?;
436        self.buffered_rows.clear();
437        Ok(())
438    }
439
440    async fn finish_inner(&mut self) -> Result<()> {
441        let aggregate_start = Instant::now();
442        let write_before = self.metrics.write_elapsed;
443        self.finish_current_row().await?;
444        self.flush_rows().await?;
445        let write_cost = self.metrics.write_elapsed.saturating_sub(write_before);
446        self.metrics.aggregate_elapsed += aggregate_start.elapsed().saturating_sub(write_cost);
447
448        let finish_start = Instant::now();
449        self.metrics.output_bytes = self.writer.finish().await?;
450        self.metrics.finish_elapsed += finish_start.elapsed();
451        Ok(())
452    }
453
454    async fn cleanup(&mut self) {
455        let start = Instant::now();
456        self.writer.abort().await;
457        self.current_row = None;
458        self.buffered_rows.clear();
459
460        self.metrics.output_bytes = 0;
461        self.metrics.cleanup_elapsed += start.elapsed();
462    }
463}
464
465/// Returns the Arrow schema of a per-SST range index.
466pub fn range_index_schema() -> SchemaRef {
467    Arc::new(Schema::new(vec![
468        Field::new(ROW_GROUP_ID_COLUMN, DataType::UInt32, false),
469        Field::new(TABLE_ID_COLUMN, DataType::UInt32, false),
470        Field::new(TSID_COLUMN, DataType::UInt64, false),
471        Field::new(START_COLUMN, DataType::Int64, false),
472        Field::new(END_COLUMN, DataType::Int64, false),
473    ]))
474}
475
476fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> {
477    ensure!(
478        metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse,
479        InvalidMetaSnafu {
480            reason: "range index only supports sparse primary-key encoding",
481        }
482    );
483    ensure!(
484        metadata
485            .primary_key
486            .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]),
487        InvalidMetaSnafu {
488            reason: "range index requires (__table_id, __tsid) as the primary-key prefix",
489        }
490    );
491    let table_id = metadata
492        .column_by_id(ReservedColumnId::table_id())
493        .context(InvalidMetaSnafu {
494            reason: "range index metadata is missing __table_id",
495        })?;
496    let tsid = metadata
497        .column_by_id(ReservedColumnId::tsid())
498        .context(InvalidMetaSnafu {
499            reason: "range index metadata is missing __tsid",
500        })?;
501    ensure!(
502        table_id.column_schema.data_type == ConcreteDataType::uint32_datatype()
503            && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype(),
504        InvalidMetaSnafu {
505            reason: "range index requires UInt32 __table_id and UInt64 __tsid",
506        }
507    );
508    Ok(())
509}
510
511fn rows_to_batch(schema: &SchemaRef, rows: &[RangeIndexRow]) -> Result<RecordBatch> {
512    let arrays: Vec<ArrayRef> = vec![
513        Arc::new(UInt32Array::from_iter_values(
514            rows.iter().map(|row| row.row_group_id),
515        )),
516        Arc::new(UInt32Array::from_iter_values(
517            rows.iter().map(|row| row.table_id),
518        )),
519        Arc::new(UInt64Array::from_iter_values(
520            rows.iter().map(|row| row.tsid),
521        )),
522        Arc::new(Int64Array::from_iter_values(
523            rows.iter().map(|row| row.start),
524        )),
525        Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.end))),
526    ];
527    RecordBatch::try_new(schema.clone(), arrays).context(NewRecordBatchSnafu)
528}
529
530#[cfg(test)]
531mod tests {
532    use datatypes::arrow::array::BinaryDictionaryBuilder;
533    use object_store::ErrorKind;
534    use object_store::services::Memory;
535    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
536
537    use super::*;
538    use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding};
539
540    fn object_store() -> ObjectStore {
541        ObjectStore::new(Memory::default()).unwrap()
542    }
543
544    fn pk_schema(primary_key_type: DataType) -> SchemaRef {
545        Arc::new(Schema::new(vec![Field::new(
546            PRIMARY_KEY_COLUMN_NAME,
547            primary_key_type,
548            false,
549        )]))
550    }
551
552    fn binary_batch(primary_keys: &[&[u8]]) -> RecordBatch {
553        RecordBatch::try_new(
554            pk_schema(DataType::Binary),
555            vec![Arc::new(BinaryArray::from_iter_values(
556                primary_keys.iter().copied(),
557            ))],
558        )
559        .unwrap()
560    }
561
562    fn dictionary_batch(primary_keys: &[&[u8]]) -> RecordBatch {
563        let mut builder = BinaryDictionaryBuilder::<UInt32Type>::new();
564        for primary_key in primary_keys {
565            builder.append(*primary_key).unwrap();
566        }
567        RecordBatch::try_new(
568            pk_schema(DataType::Dictionary(
569                Box::new(DataType::UInt32),
570                Box::new(DataType::Binary),
571            )),
572            vec![Arc::new(builder.finish())],
573        )
574        .unwrap()
575    }
576
577    async fn read_index(store: &ObjectStore, path: &str) -> (u64, usize, Vec<RangeIndexRow>) {
578        let bytes = store.read(path).await.unwrap().to_bytes();
579        let output_bytes = bytes.len() as u64;
580        let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
581        let row_groups = builder.metadata().num_row_groups();
582        let batches = builder
583            .build()
584            .unwrap()
585            .collect::<std::result::Result<Vec<_>, _>>()
586            .unwrap();
587        let mut rows = Vec::new();
588        for batch in batches {
589            let row_group_ids = batch
590                .column(0)
591                .as_any()
592                .downcast_ref::<UInt32Array>()
593                .unwrap();
594            let table_ids = batch
595                .column(1)
596                .as_any()
597                .downcast_ref::<UInt32Array>()
598                .unwrap();
599            let tsids = batch
600                .column(2)
601                .as_any()
602                .downcast_ref::<UInt64Array>()
603                .unwrap();
604            let starts = batch
605                .column(3)
606                .as_any()
607                .downcast_ref::<Int64Array>()
608                .unwrap();
609            let ends = batch
610                .column(4)
611                .as_any()
612                .downcast_ref::<Int64Array>()
613                .unwrap();
614            for row in 0..batch.num_rows() {
615                rows.push(RangeIndexRow {
616                    row_group_id: row_group_ids.value(row),
617                    table_id: table_ids.value(row),
618                    tsid: tsids.value(row),
619                    start: starts.value(row),
620                    end: ends.value(row),
621                });
622            }
623        }
624        (output_bytes, row_groups, rows)
625    }
626
627    #[tokio::test]
628    async fn test_write_ranges_and_metrics() {
629        let metadata = Arc::new(sst_region_metadata_with_encoding(
630            PrimaryKeyEncoding::Sparse,
631        ));
632        let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10);
633        let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20);
634        let primary_key_3 = new_sparse_primary_key(&["c", "z"], &metadata, 2, 30);
635        let store = object_store();
636        let mut writer = SstRangeIndexWriter::try_new(
637            metadata,
638            store.clone(),
639            "ranges.parquet",
640            SstRangeIndexWriterOptions {
641                index_row_group_size: 2,
642            },
643        )
644        .await
645        .unwrap();
646
647        writer
648            .write(
649                0,
650                &dictionary_batch(&[primary_key_1.as_slice(), primary_key_1.as_slice()]),
651            )
652            .await
653            .unwrap();
654        writer
655            .write(
656                0,
657                &binary_batch(&[
658                    primary_key_1.as_slice(),
659                    primary_key_2.as_slice(),
660                    primary_key_2.as_slice(),
661                ]),
662            )
663            .await
664            .unwrap();
665        writer
666            .write(
667                1,
668                &binary_batch(&[
669                    primary_key_2.as_slice(),
670                    primary_key_2.as_slice(),
671                    primary_key_3.as_slice(),
672                ]),
673            )
674            .await
675            .unwrap();
676
677        assert_eq!(writer.metrics().input_batches, 3);
678        assert_eq!(writer.metrics().input_rows, 8);
679        assert_eq!(writer.metrics().num_source_row_groups, 2);
680        let metrics = writer.finish().await.unwrap();
681        assert_eq!(metrics.num_ranges, 4);
682        assert!(!metrics.aborted);
683
684        let (output_bytes, row_groups, rows) = read_index(&store, "ranges.parquet").await;
685        assert_eq!(metrics.output_bytes, output_bytes);
686        assert_eq!(row_groups, 2);
687        assert_eq!(
688            rows,
689            vec![
690                RangeIndexRow {
691                    row_group_id: 0,
692                    table_id: 1,
693                    tsid: 10,
694                    start: 0,
695                    end: 3,
696                },
697                RangeIndexRow {
698                    row_group_id: 0,
699                    table_id: 1,
700                    tsid: 20,
701                    start: 3,
702                    end: 5,
703                },
704                RangeIndexRow {
705                    row_group_id: 1,
706                    table_id: 1,
707                    tsid: 20,
708                    start: 0,
709                    end: 2,
710                },
711                RangeIndexRow {
712                    row_group_id: 1,
713                    table_id: 2,
714                    tsid: 30,
715                    start: 2,
716                    end: 3,
717                },
718            ]
719        );
720    }
721
722    #[tokio::test]
723    async fn test_reject_invalid_row_groups_and_sorted_order() {
724        let metadata = Arc::new(sst_region_metadata_with_encoding(
725            PrimaryKeyEncoding::Sparse,
726        ));
727        let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10);
728        let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20);
729        let store = object_store();
730
731        let mut writer = SstRangeIndexWriter::try_new(
732            metadata.clone(),
733            store.clone(),
734            "groups-gap.parquet",
735            SstRangeIndexWriterOptions::default(),
736        )
737        .await
738        .unwrap();
739        writer
740            .write(0, &binary_batch(&[primary_key_1.as_slice()]))
741            .await
742            .unwrap();
743        let error = writer
744            .write(2, &binary_batch(&[primary_key_2.as_slice()]))
745            .await
746            .unwrap_err();
747        assert!(error.to_string().contains("must be contiguous"), "{error}");
748        writer.abort().await.unwrap();
749
750        let mut writer = SstRangeIndexWriter::try_new(
751            metadata,
752            store.clone(),
753            "sort-error.parquet",
754            SstRangeIndexWriterOptions::default(),
755        )
756        .await
757        .unwrap();
758        let error = writer
759            .write(
760                0,
761                &dictionary_batch(&[primary_key_2.as_slice(), primary_key_1.as_slice()]),
762            )
763            .await
764            .unwrap_err();
765        assert!(error.to_string().contains("not sorted"), "{error}");
766        let metrics = writer.abort().await.unwrap();
767        assert!(metrics.aborted);
768        assert_eq!(metrics.output_bytes, 0);
769        assert_eq!(
770            store.stat("sort-error.parquet").await.unwrap_err().kind(),
771            ErrorKind::NotFound
772        );
773    }
774}
mito2/sst/range_index/writer.rs

mito2/sst/range_index/
writer.rs
