perf(mito2): lazily extract sparse primary key index values (#9176)

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

src/index/src/bloom_filter/creator.rs

@@ -149,6 +149,35 @@ impl BloomFilterCreator {
         Ok(())
     }
 
+    /// Adds `nrows` copies of a single borrowed value (or null), copying it only when
+    /// it is new to a segment. Row counts advance for nulls as well.
+    pub async fn push_n_row_elem(&mut self, mut nrows: usize, elem: Option<&[u8]>) -> Result<()> {
+        while nrows > 0 {
+            let rows_to_seg_end =
+                self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment);
+            let rows_to_push = nrows.min(rows_to_seg_end);
+            nrows -= rows_to_push;
+            self.accumulated_row_count += rows_to_push;
+
+            if let Some(elem) = elem
+                && !self.cur_seg_distinct_elems.contains(elem)
+            {
+                self.cur_seg_distinct_elems.insert(elem.to_vec());
+                self.cur_seg_distinct_elems_mem_usage += elem.len();
+                self.global_memory_usage
+                    .fetch_add(elem.len(), Ordering::Relaxed);
+            }
+            if self
+                .accumulated_row_count
+                .is_multiple_of(self.rows_per_segment)
+            {
+                self.finalize_segment().await?;
+                self.finalized_row_count = self.accumulated_row_count;
+            }
+        }
+        Ok(())
+    }
+
     /// Adds a row of elements to the bloom filter. If the number of accumulated rows
     /// reaches `rows_per_segment`, it finalizes the current segment.
     pub async fn push_row_elems(&mut self, elems: impl IntoIterator<Item = Bytes>) -> Result<()> {
@@ -420,6 +449,42 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn borrowed_single_value_matches_owned_rows_across_segments() {
+        let make_creator = || {
+            BloomFilterCreator::new(
+                3,
+                0.01,
+                Arc::new(MockExternalTempFileProvider::new()),
+                Arc::new(AtomicUsize::new(0)),
+                None,
+            )
+        };
+        let mut borrowed = make_creator();
+        let mut owned = make_creator();
+        // Zero rows, nulls, empty values, duplicates and runs crossing segment boundaries.
+        for (rows, elem) in [
+            (0, Some(b"ignored".as_slice())),
+            (1, None),
+            (5, Some(b"".as_slice())),
+            (1, Some(b"".as_slice())),
+            (8, Some(b"label".as_slice())),
+            (1, None),
+        ] {
+            borrowed.push_n_row_elem(rows, elem).await.unwrap();
+            owned
+                .push_n_row_elems(rows, elem.map(<[u8]>::to_vec))
+                .await
+                .unwrap();
+            assert_eq!(borrowed.memory_usage(), owned.memory_usage());
+        }
+        let mut borrowed_blob = Cursor::new(Vec::new());
+        let mut owned_blob = Cursor::new(Vec::new());
+        borrowed.finish(&mut borrowed_blob).await.unwrap();
+        owned.finish(&mut owned_blob).await.unwrap();
+        assert_eq!(borrowed_blob.into_inner(), owned_blob.into_inner());
+    }
+
     #[tokio::test]
     async fn test_final_seg_all_null() {
         let mut writer = Cursor::new(Vec::new());

src/mito-codec/src/index.rs

@@ -24,13 +24,56 @@ use store_api::codec::PrimaryKeyEncoding;
 use store_api::metadata::ColumnMetadata;
 use store_api::storage::ColumnId;
 
-use crate::error::{FieldTypeMismatchSnafu, IndexEncodeNullSnafu, Result};
+use crate::error::{
+    FieldTypeMismatchSnafu, IndexEncodeNullSnafu, InvalidSparsePrimaryKeySnafu, Result,
+};
+use crate::row_converter::sparse::{
+    RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView,
+};
 use crate::row_converter::{PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields};
 
 /// Encodes index values according to their data types for sorting and storage use.
 pub struct IndexValueCodec;
 
 impl IndexValueCodec {
+    /// Extracts one sparse PK column in index format, without constructing a Value.
+    /// Numeric reserved fields borrow the PK; strings are unchunked into the reusable buffer.
+    /// Missing and null labels return None, whereas an empty string returns Some(&[]).
+    pub fn encode_sparse_value<'a, 'pk: 'a>(
+        pk: &mut SparsePrimaryKeyView<'pk, '_>,
+        column_id: ColumnId,
+        buffer: &'a mut Vec<u8>,
+    ) -> Result<Option<&'a [u8]>> {
+        let Some(encoded) = pk.encoded_value(column_id)? else {
+            return Ok(None);
+        };
+        if encoded[0] == 0 {
+            return Ok(None);
+        }
+        if matches!(
+            column_id,
+            RESERVED_COLUMN_ID_TABLE_ID | RESERVED_COLUMN_ID_TSID
+        ) {
+            return Ok(Some(encoded));
+        }
+
+        buffer.clear();
+        // The view has checked the Option marker, bytes marker and every chunk length.
+        // Reserving once avoids growing the buffer for each 8-byte chunk.
+        buffer.reserve(encoded.len() - 2);
+        for chunk in encoded[2..].chunks_exact(9) {
+            let len = usize::from(chunk[8]).min(8);
+            buffer.extend_from_slice(&chunk[..len]);
+        }
+        std::str::from_utf8(buffer).map_err(|_| {
+            InvalidSparsePrimaryKeySnafu {
+                reason: "label is not valid UTF-8",
+            }
+            .build()
+        })?;
+        Ok(Some(buffer.as_slice()))
+    }
+
     /// Serializes a non-null `ValueRef` using the data type defined in `SortField` and writes
     /// the result into a buffer.
     ///

src/mito-codec/src/row_converter/sparse.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod checked;
+
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
@@ -34,6 +36,7 @@ use crate::error::{
 use crate::key_values::KeyValue;
 use crate::primary_key_filter::SparsePrimaryKeyFilter;
 use crate::row_converter::dense::SortField;
+pub use crate::row_converter::sparse::checked::SparsePrimaryKeyView;
 use crate::row_converter::{CompositeValues, PrimaryKeyCodec, PrimaryKeyFilter};
 
 /// A codec for sparse key of metrics.

src/mito-codec/src/row_converter/sparse/checked.rs

@@ -0,0 +1,244 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use snafu::{OptionExt, ensure};
+use store_api::storage::ColumnId;
+
+use crate::error::{InvalidSparsePrimaryKeySnafu, Result};
+use crate::row_converter::sparse::{
+    COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID,
+    SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET,
+};
+
+/// A borrowing view of one sparse primary key, with checked, lazy field lookup.
+///
+/// The scratch cache is reset on construction and retains its capacity between keys.
+/// Ordinary labels can be skipped without schema lookup: their encoding is always String.
+pub struct SparsePrimaryKeyView<'a, 'b> {
+    pk: &'a [u8],
+    cache: &'b mut SparseOffsetsCache,
+}
+
+impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> {
+    /// Validates the fixed prefix and binds fresh offset discovery to this key.
+    pub fn new(pk: &'a [u8], cache: &'b mut SparseOffsetsCache) -> Result<Self> {
+        ensure!(
+            pk.len() >= TAGS_START_OFFSET
+                && pk[..COLUMN_ID_ENCODE_SIZE] == RESERVED_COLUMN_ID_TABLE_ID.to_be_bytes()
+                && pk[TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE..TSID_VALUE_OFFSET]
+                    == RESERVED_COLUMN_ID_TSID.to_be_bytes()
+                && pk[TABLE_ID_VALUE_OFFSET] == 1
+                && pk[TSID_VALUE_OFFSET] == 1,
+            InvalidSparsePrimaryKeySnafu {
+                reason: "invalid table_id/tsid prefix"
+            }
+        );
+        cache.clear();
+        Ok(Self { pk, cache })
+    }
+
+    /// Returns the encoded value, including its null marker, or None for an absent label.
+    /// The result borrows the key, not the scratch cache.
+    pub fn encoded_value(&mut self, column_id: ColumnId) -> Result<Option<&'a [u8]>> {
+        match column_id {
+            RESERVED_COLUMN_ID_TABLE_ID => {
+                return Ok(Some(
+                    &self.pk[TABLE_ID_VALUE_OFFSET..TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE],
+                ));
+            }
+            RESERVED_COLUMN_ID_TSID => {
+                return Ok(Some(&self.pk[TSID_VALUE_OFFSET..TAGS_START_OFFSET]));
+            }
+            _ => {}
+        }
+        if let Some(offset) = self.cache.get(column_id) {
+            return encoded_label(&self.pk[offset..]).map(Some);
+        }
+        while !self.cache.finished && self.cache.cursor < self.pk.len() {
+            let bytes = &self.pk[self.cache.cursor..];
+            let id = bytes.first_chunk::<COLUMN_ID_ENCODE_SIZE>().context(
+                InvalidSparsePrimaryKeySnafu {
+                    reason: "truncated label column id",
+                },
+            )?;
+            let id = u32::from_be_bytes(*id);
+            let offset = self.cache.cursor + COLUMN_ID_ENCODE_SIZE;
+            let value = encoded_label(&self.pk[offset..])?;
+            self.cache.insert(id, offset);
+            self.cache.cursor = offset + value.len();
+            if id == column_id {
+                return Ok(Some(value));
+            }
+        }
+        self.cache.finished = true;
+        Ok(None)
+    }
+}
+
+/// Finds the end of an Option<String> without allocating or reading its payload.
+/// Unlike memcomparable's skip_bytes, all advances are checked for truncated input.
+fn encoded_label(bytes: &[u8]) -> Result<&[u8]> {
+    match bytes.first() {
+        Some(0) => return Ok(&bytes[..1]),
+        Some(1) => {}
+        _ => {
+            return InvalidSparsePrimaryKeySnafu {
+                reason: "invalid label null marker",
+            }
+            .fail();
+        }
+    }
+    match bytes.get(1) {
+        Some(0) => return Ok(&bytes[..2]),
+        Some(1) => {}
+        _ => {
+            return InvalidSparsePrimaryKeySnafu {
+                reason: "invalid label bytes marker",
+            }
+            .fail();
+        }
+    }
+    let mut end = 2;
+    loop {
+        let chunk = bytes
+            .get(end..end + 9)
+            .context(InvalidSparsePrimaryKeySnafu {
+                reason: "truncated label chunk",
+            })?;
+        end += 9;
+        match chunk[8] {
+            1..=8 => return Ok(&bytes[..end]),
+            9 => {}
+            _ => {
+                return InvalidSparsePrimaryKeySnafu {
+                    reason: "invalid label chunk length",
+                }
+                .fail();
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use datatypes::data_type::ConcreteDataType;
+    use datatypes::value::ValueRef;
+
+    use super::*;
+    use crate::index::IndexValueCodec;
+    use crate::row_converter::{PrimaryKeyCodec, SortField, SparsePrimaryKeyCodec};
+
+    #[test]
+    fn index_bytes_match_full_decode_across_chunks_and_lookup_orders() {
+        let codec = SparsePrimaryKeyCodec::schemaless();
+        let mut cache = SparseOffsetsCache::new();
+        let mut buffer = Vec::new();
+        // Reusing the cache across different keys must not reuse their field offsets.
+        for len in [0, 1, 7, 8, 9, 16, 24, 65] {
+            let label = "x".repeat(len);
+            let values = [
+                (RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)),
+                (RESERVED_COLUMN_ID_TSID, ValueRef::UInt64(u64::MAX)),
+                (90, ValueRef::String(&label)),
+                (7, ValueRef::String("")),
+                (1, ValueRef::Null),
+                (3, ValueRef::String("last中文")),
+            ];
+            let mut pk = Vec::new();
+            codec.encode_value_refs(&values, &mut pk).unwrap();
+            // Exercise the offset cache's overflow storage as well as its inline entries.
+            codec
+                .encode_raw_tag_value((100..140).map(|id| (id, label.as_bytes())), &mut pk)
+                .unwrap();
+            // Writers normally omit null labels; existing encoded nulls must remain null too.
+            pk.extend_from_slice(&2_u32.to_be_bytes());
+            pk.push(0);
+            let decoded = codec.decode(&pk).unwrap().into_sparse();
+            for order in [
+                [
+                    RESERVED_COLUMN_ID_TABLE_ID,
+                    RESERVED_COLUMN_ID_TSID,
+                    90,
+                    7,
+                    1,
+                    3,
+                ],
+                [
+                    3,
+                    1,
+                    7,
+                    90,
+                    RESERVED_COLUMN_ID_TSID,
+                    RESERVED_COLUMN_ID_TABLE_ID,
+                ],
+            ] {
+                let mut view = SparsePrimaryKeyView::new(&pk, &mut cache).unwrap();
+                for id in order.into_iter().chain([139, 100, 132, 2]) {
+                    let expected = decoded.get(&id).filter(|v| !v.is_null()).map(|value| {
+                        let field = SortField::new(match id {
+                            RESERVED_COLUMN_ID_TABLE_ID => ConcreteDataType::uint32_datatype(),
+                            RESERVED_COLUMN_ID_TSID => ConcreteDataType::uint64_datatype(),
+                            _ => ConcreteDataType::string_datatype(),
+                        });
+                        let mut bytes = Vec::new();
+                        IndexValueCodec::encode_nonnull_value(
+                            value.as_value_ref(),
+                            &field,
+                            &mut bytes,
+                        )
+                        .unwrap();
+                        bytes
+                    });
+                    let actual =
+                        IndexValueCodec::encode_sparse_value(&mut view, id, &mut buffer).unwrap();
+                    assert_eq!(actual, expected.as_deref(), "len={len}, column={id}");
+                }
+                assert!(view.encoded_value(999).unwrap().is_none());
+            }
+        }
+    }
+
+    #[test]
+    fn malformed_sparse_values_return_errors() {
+        let codec = SparsePrimaryKeyCodec::schemaless();
+        let mut pk = Vec::new();
+        codec.encode_internal(1, 2, &mut pk).unwrap();
+        codec
+            .encode_raw_tag_value([(1, b"0123456789".as_slice())].into_iter(), &mut pk)
+            .unwrap();
+        let mut cache = SparseOffsetsCache::new();
+        for end in 0..pk.len() {
+            // A complete reserved prefix is a valid key with no labels.
+            if end == TAGS_START_OFFSET {
+                continue;
+            }
+            let result = SparsePrimaryKeyView::new(&pk[..end], &mut cache)
+                .and_then(|mut view| view.encoded_value(1));
+            assert!(result.is_err(), "truncation at {end}");
+        }
+        for (offset, value) in [(0, 0), (4, 0), (13, 2), (26, 2), (27, 2), (36, 0), (45, 10)] {
+            let mut invalid = pk.clone();
+            invalid[offset] = value;
+            assert!(
+                SparsePrimaryKeyView::new(&invalid, &mut cache)
+                    .and_then(|mut view| view.encoded_value(1))
+                    .is_err()
+            );
+        }
+        let mut invalid_utf8 = pk;
+        invalid_utf8[28] = 0xff;
+        let mut view = SparsePrimaryKeyView::new(&invalid_utf8, &mut cache).unwrap();
+        assert!(IndexValueCodec::encode_sparse_value(&mut view, 1, &mut Vec::new()).is_err());
+    }
+}

src/mito2/Cargo.toml

@@ -145,5 +145,10 @@ required-features = ["test"]
 name = "bench_wal_encode"
 harness = false
 
+[[bench]]
+name = "bench_index_update"
+harness = false
+required-features = ["testing"]
+
 [package.metadata.cargo-udeps.ignore]
 normal = ["aquamarine"]

src/mito2/benches/bench_index_update.rs

@@ -0,0 +1,276 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Measures real index updates, excluding input generation, creator setup and cleanup.
+//!
+//! Run on the baseline revision with `--save-baseline before`, then on the candidate
+//! with `--baseline before` (arguments after `--` in the command below):
+//! `CARGO_PROFILE_BENCH_DEBUG=0 cargo bench -p mito2 --features testing --bench bench_index_update`.
+
+use std::collections::HashSet;
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use api::v1::SemanticType;
+use common_test_util::temp_dir::create_temp_dir;
+use criterion::{Criterion, Throughput, criterion_group, criterion_main};
+use datatypes::arrow::array::{
+    ArrayRef, BinaryDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array,
+};
+use datatypes::arrow::datatypes::UInt32Type;
+use datatypes::arrow::record_batch::RecordBatch;
+use datatypes::data_type::ConcreteDataType;
+use datatypes::schema::{ColumnSchema, SkippingIndexOptions};
+use mito_codec::row_converter::SparsePrimaryKeyCodec;
+use mito2::sst::index::intermediate::IntermediateManager;
+use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema};
+use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer};
+use store_api::codec::PrimaryKeyEncoding;
+use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef};
+use store_api::storage::consts::ReservedColumnId;
+use store_api::storage::{ColumnId, FileId, RegionId};
+
+const ROWS: usize = 4096;
+const SEGMENT_ROWS: usize = 1024;
+
+struct Shape {
+    name: &'static str,
+    tags: u32,
+    indexed_tags: u32,
+    inverted: bool,
+    bloom: bool,
+    rows_per_key: usize,
+}
+
+fn input(shape: &Shape) -> (RegionMetadataRef, HashSet<ColumnId>, RecordBatch) {
+    let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
+    let mut pk = vec![ReservedColumnId::table_id(), ReservedColumnId::tsid()];
+    let mut inverted_columns = HashSet::new();
+    let mut add_tag = |id, name: String, data_type, indexed| {
+        let mut schema = ColumnSchema::new(name, data_type, true).with_inverted_index(false);
+        if indexed && shape.inverted {
+            inverted_columns.insert(id);
+        }
+        if indexed && shape.bloom {
+            schema = schema
+                .with_skipping_options(SkippingIndexOptions {
+                    granularity: SEGMENT_ROWS as _,
+                    ..Default::default()
+                })
+                .unwrap();
+        }
+        builder.push_column_metadata(ColumnMetadata {
+            column_schema: schema,
+            semantic_type: SemanticType::Tag,
+            column_id: id,
+        });
+    };
+    add_tag(
+        pk[0],
+        "__table_id".into(),
+        ConcreteDataType::uint32_datatype(),
+        shape.indexed_tags == 0,
+    );
+    add_tag(
+        pk[1],
+        "__tsid".into(),
+        ConcreteDataType::uint64_datatype(),
+        false,
+    );
+    for id in 0..shape.tags {
+        // A single indexed label is deliberately last in the key.
+        add_tag(
+            id,
+            format!("tag_{id:03}"),
+            ConcreteDataType::string_datatype(),
+            id >= shape.tags - shape.indexed_tags,
+        );
+        pk.push(id);
+    }
+    builder
+        .push_column_metadata(ColumnMetadata {
+            column_schema: ColumnSchema::new(
+                "ts",
+                ConcreteDataType::timestamp_millisecond_datatype(),
+                false,
+            ),
+            semantic_type: SemanticType::Timestamp,
+            column_id: shape.tags,
+        })
+        .primary_key(pk)
+        .primary_key_encoding(PrimaryKeyEncoding::Sparse);
+    let metadata = Arc::new(builder.build().unwrap());
+    let codec = SparsePrimaryKeyCodec::schemaless();
+    let mut keys = BinaryDictionaryBuilder::<UInt32Type>::new();
+    for series in 0..ROWS / shape.rows_per_key {
+        let mut key = Vec::new();
+        codec
+            .encode_internal((series / 128) as u32, series as u64, &mut key)
+            .unwrap();
+        let labels: Vec<_> = (0..shape.tags)
+            .map(|id| (id, format!("tag-{id:03}-value-{series:010}")))
+            .collect();
+        codec
+            .encode_raw_tag_value(
+                labels.iter().map(|(id, value)| (*id, value.as_bytes())),
+                &mut key,
+            )
+            .unwrap();
+        for _ in 0..shape.rows_per_key {
+            keys.append(&key).unwrap();
+        }
+    }
+    let schema = to_flat_sst_arrow_schema(
+        &metadata,
+        &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse),
+    );
+    let columns: Vec<ArrayRef> = vec![
+        Arc::new(TimestampMillisecondArray::from_iter_values(0..ROWS as i64)),
+        Arc::new(keys.finish()),
+        Arc::new(UInt64Array::from(vec![1; ROWS])),
+        Arc::new(UInt8Array::from(vec![1; ROWS])),
+    ];
+    let batch = RecordBatch::try_new(schema, columns).unwrap();
+    (metadata, inverted_columns, batch)
+}
+
+fn bench_index_update(c: &mut Criterion) {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    let dir = create_temp_dir("bench_index_update");
+    let intermediate = runtime
+        .block_on(IntermediateManager::init_fs(dir.path().to_str().unwrap()))
+        .unwrap();
+    let mut group = c.benchmark_group("sparse_index_update");
+    group.sample_size(30);
+    group.warm_up_time(Duration::from_secs(1));
+    group.measurement_time(Duration::from_secs(3));
+    group.throughput(Throughput::Elements(ROWS as u64));
+    for shape in [
+        Shape {
+            name: "table_id_bloom_40tags",
+            tags: 40,
+            indexed_tags: 0,
+            inverted: false,
+            bloom: true,
+            rows_per_key: 1,
+        },
+        Shape {
+            name: "inverted_one_40tags",
+            tags: 40,
+            indexed_tags: 1,
+            inverted: true,
+            bloom: false,
+            rows_per_key: 1,
+        },
+        Shape {
+            name: "bloom_one_40tags",
+            tags: 40,
+            indexed_tags: 1,
+            inverted: false,
+            bloom: true,
+            rows_per_key: 1,
+        },
+        Shape {
+            name: "both_one_40tags",
+            tags: 40,
+            indexed_tags: 1,
+            inverted: true,
+            bloom: true,
+            rows_per_key: 1,
+        },
+        Shape {
+            name: "both_all_40tags",
+            tags: 40,
+            indexed_tags: 40,
+            inverted: true,
+            bloom: true,
+            rows_per_key: 1,
+        },
+        Shape {
+            name: "both_all_40tags_8rpk",
+            tags: 40,
+            indexed_tags: 40,
+            inverted: true,
+            bloom: true,
+            rows_per_key: 8,
+        },
+        Shape {
+            name: "both_one_10tags",
+            tags: 10,
+            indexed_tags: 1,
+            inverted: true,
+            bloom: true,
+            rows_per_key: 1,
+        },
+        Shape {
+            name: "both_all_10tags",
+            tags: 10,
+            indexed_tags: 10,
+            inverted: true,
+            bloom: true,
+            rows_per_key: 1,
+        },
+    ] {
+        let (metadata, inverted_columns, batch) = input(&shape);
+        group.bench_function(shape.name, |b| {
+            b.iter_custom(|iterations| {
+                runtime.block_on(async {
+                    let mut elapsed = Duration::ZERO;
+                    for _ in 0..iterations {
+                        let file_id = FileId::random();
+                        let mut inverted = shape.inverted.then(|| {
+                            InvertedIndexer::new(
+                                file_id,
+                                &metadata,
+                                intermediate.clone(),
+                                None,
+                                NonZeroUsize::new(SEGMENT_ROWS).unwrap(),
+                                inverted_columns.clone(),
+                            )
+                        });
+                        let mut bloom = if shape.bloom {
+                            BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None)
+                                .unwrap()
+                        } else {
+                            None
+                        };
+                        let start = Instant::now();
+                        if let Some(indexer) = &mut inverted {
+                            indexer.update_flat(&batch).await.unwrap();
+                        }
+                        if let Some(indexer) = &mut bloom {
+                            indexer.update_flat(&batch).await.unwrap();
+                        }
+                        elapsed += start.elapsed();
+                        if let Some(indexer) = &mut inverted {
+                            indexer.abort().await.unwrap();
+                        }
+                        if let Some(indexer) = &mut bloom {
+                            indexer.abort().await.unwrap();
+                        }
+                    }
+                    elapsed
+                })
+            });
+        });
+    }
+    group.finish();
+}
+
+criterion_group!(benches, bench_index_update);
+criterion_main!(benches);

src/mito2/src/sst/index.rs

@@ -17,7 +17,10 @@ pub(crate) mod fulltext_index;
 mod indexer;
 pub mod intermediate;
 pub(crate) mod inverted_index;
+mod primary_key;
 pub mod puffin_manager;
+#[cfg(test)]
+mod sparse_test;
 mod statistics;
 pub(crate) mod store;
 #[cfg(feature = "vector_index")]
@@ -30,14 +33,11 @@ use std::sync::Arc;
 
 use bloom_filter::creator::BloomFilterIndexer;
 use common_telemetry::{debug, error, info, warn};
-use datatypes::arrow::array::BinaryArray;
 use datatypes::arrow::record_batch::RecordBatch;
-use mito_codec::index::IndexValuesCodec;
-use mito_codec::row_converter::CompositeValues;
 use object_store::ObjectStore;
 use puffin_manager::SstPuffinManager;
 use smallvec::{SmallVec, smallvec};
-use snafu::{OptionExt, ResultExt};
+use snafu::ResultExt;
 use statistics::{ByteCount, RowCount};
 use store_api::metadata::RegionMetadataRef;
 use store_api::storage::{ColumnId, FileId, RegionId};
@@ -54,8 +54,8 @@ use crate::cache::{CacheManagerRef, CacheStrategy};
 use crate::config::VectorIndexConfig;
 use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig};
 use crate::error::{
-    BuildIndexAsyncSnafu, DecodeSnafu, Error, InvalidRecordBatchSnafu, RegionClosedSnafu,
-    RegionDroppedSnafu, RegionTruncatedSnafu, Result,
+    BuildIndexAsyncSnafu, Error, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu,
+    Result,
 };
 use crate::metrics::{
     INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL,
@@ -79,8 +79,6 @@ use crate::sst::index::fulltext_index::creator::FulltextIndexer;
 use crate::sst::index::intermediate::IntermediateManager;
 use crate::sst::index::inverted_index::creator::InvertedIndexer;
 use crate::sst::parquet::SstInfo;
-use crate::sst::parquet::flat_format::primary_key_column_index;
-use crate::sst::parquet::format::PrimaryKeyArray;
 use crate::worker::WorkerListener;
 
 pub(crate) const TYPE_INVERTED_INDEX: &str = "inverted_index";
@@ -1450,55 +1448,6 @@ impl IndexBuildScheduler {
     }
 }
 
-/// Decodes primary keys from a flat format RecordBatch.
-/// Returns a list of (decoded_pk_value, count) tuples where count is the number of occurrences.
-pub(crate) fn decode_primary_keys_with_counts(
-    batch: &RecordBatch,
-    codec: &IndexValuesCodec,
-) -> Result<Vec<(CompositeValues, usize)>> {
-    let primary_key_index = primary_key_column_index(batch.num_columns());
-    let pk_dict_array = batch
-        .column(primary_key_index)
-        .as_any()
-        .downcast_ref::<PrimaryKeyArray>()
-        .context(InvalidRecordBatchSnafu {
-            reason: "Primary key column is not a dictionary array",
-        })?;
-    let pk_values_array = pk_dict_array
-        .values()
-        .as_any()
-        .downcast_ref::<BinaryArray>()
-        .context(InvalidRecordBatchSnafu {
-            reason: "Primary key values are not binary array",
-        })?;
-    let keys = pk_dict_array.keys();
-
-    // Decodes primary keys and count consecutive occurrences
-    let mut result: Vec<(CompositeValues, usize)> = Vec::new();
-    let mut prev_key: Option<u32> = None;
-
-    let pk_indices = keys.values();
-    for &current_key in pk_indices.iter().take(keys.len()) {
-        // Checks if current key is the same as previous key
-        if let Some(prev) = prev_key
-            && prev == current_key
-        {
-            // Safety: We already have a key in the result vector.
-            result.last_mut().unwrap().1 += 1;
-            continue;
-        }
-
-        // New key, decodes it.
-        let pk_bytes = pk_values_array.value(current_key as usize);
-        let decoded_value = codec.decoder().decode(pk_bytes).context(DecodeSnafu)?;
-
-        result.push((decoded_value, 1));
-        prev_key = Some(current_key);
-    }
-
-    Ok(result)
-}
-
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

src/mito2/src/sst/index/bloom_filter/creator.rs

@@ -24,8 +24,10 @@ use datatypes::vectors::Helper;
 use index::bloom_filter::creator::BloomFilterCreator;
 use index::target::IndexTarget;
 use mito_codec::index::{IndexValueCodec, IndexValuesCodec};
-use mito_codec::row_converter::{CompositeValues, SortField};
+use mito_codec::row_converter::sparse::SparsePrimaryKeyView;
+use mito_codec::row_converter::{SortField, SparseOffsetsCache};
 use puffin::puffin_manager::{PuffinWriter, PutOptions};
+use smallvec::SmallVec;
 use snafu::{ResultExt, ensure};
 use store_api::codec::PrimaryKeyEncoding;
 use store_api::metadata::RegionMetadataRef;
@@ -33,17 +35,18 @@ use store_api::storage::{ColumnId, FileId};
 use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
 
 use crate::error::{
-    BiErrorsSnafu, BloomFilterFinishSnafu, EncodeSnafu, IndexOptionsSnafu,
+    BiErrorsSnafu, BloomFilterFinishSnafu, DecodeSnafu, EncodeSnafu, IndexOptionsSnafu,
     OperateAbortedIndexSnafu, PuffinAddBlobSnafu, PushBloomFilterValueSnafu, Result,
 };
 use crate::read::Batch;
+use crate::sst::index::TYPE_BLOOM_FILTER_INDEX;
 use crate::sst::index::bloom_filter::INDEX_BLOB_TYPE;
 use crate::sst::index::intermediate::{
     IntermediateLocation, IntermediateManager, TempFileProvider,
 };
+use crate::sst::index::primary_key::PrimaryKeyRuns;
 use crate::sst::index::puffin_manager::SstPuffinWriter;
 use crate::sst::index::statistics::{ByteCount, RowCount, Statistics};
-use crate::sst::index::{TYPE_BLOOM_FILTER_INDEX, decode_primary_keys_with_counts};
 
 /// The buffer size for the pipe used to send index data to the puffin blob.
 const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192;
@@ -58,6 +61,9 @@ pub struct BloomFilterIndexer {
 
     /// Codec for decoding primary keys.
     codec: IndexValuesCodec,
+    /// Scratch storage for extracting indexed tags from sparse primary keys.
+    pk_offsets: SparseOffsetsCache,
+    value_buf: Vec<u8>,
 
     /// Whether the indexing process has been aborted.
     aborted: bool,
@@ -124,6 +130,8 @@ impl BloomFilterIndexer {
             creators,
             temp_file_provider,
             codec,
+            pk_offsets: SparseOffsetsCache::new(),
+            value_buf: Vec::new(),
             aborted: false,
             stats: Statistics::new(TYPE_BLOOM_FILTER_INDEX),
             global_memory_usage,
@@ -185,7 +193,7 @@ impl BloomFilterIndexer {
     /// Returns the number of rows and bytes written.
     ///
     /// TODO(zhongzc): duplicate with `mito2::sst::index::inverted_index::creator::InvertedIndexCreator`
-    pub async fn finish(
+    pub(crate) async fn finish(
         &mut self,
         puffin_writer: &mut SstPuffinWriter,
     ) -> Result<(RowCount, ByteCount)> {
@@ -292,7 +300,8 @@ impl BloomFilterIndexer {
         guard.inc_row_count(n);
 
         let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse;
-        let mut decoded_pks: Option<Vec<(CompositeValues, usize)>> = None;
+        let mut sparse_columns: SmallVec<[(ColumnId, &mut BloomFilterCreator); 8]> =
+            SmallVec::new();
 
         for (col_id, creator) in &mut self.creators {
             // Safety: `creators` are created from the metadata so it won't be None.
@@ -321,45 +330,8 @@ impl BloomFilterIndexer {
                         .context(PushBloomFilterValueSnafu)?;
                 }
             } else if is_sparse && column_meta.semantic_type == SemanticType::Tag {
-                // Column not found in batch, tries to decode from primary keys for sparse encoding.
-                if decoded_pks.is_none() {
-                    decoded_pks = Some(decode_primary_keys_with_counts(batch, &self.codec)?);
-                }
-
-                let pk_values_with_counts = decoded_pks.as_ref().unwrap();
-                let Some(col_info) = self.codec.pk_col_info(*col_id) else {
-                    debug!(
-                        "Column {} not found in primary key during building bloom filter index",
-                        column_name
-                    );
-                    continue;
-                };
-                let pk_index = col_info.idx;
-                let field = &col_info.field;
-                for (decoded, count) in pk_values_with_counts {
-                    let value = match decoded {
-                        CompositeValues::Dense(dense) => dense.get(pk_index).map(|v| &v.1),
-                        CompositeValues::Sparse(sparse) => sparse.get(col_id),
-                    };
-
-                    let elems = value
-                        .filter(|v| !v.is_null())
-                        .map(|v| {
-                            let mut buf = vec![];
-                            IndexValueCodec::encode_nonnull_value(
-                                v.as_value_ref(),
-                                field,
-                                &mut buf,
-                            )
-                            .context(EncodeSnafu)?;
-                            Ok(buf)
-                        })
-                        .transpose()?;
-
-                    creator
-                        .push_n_row_elems(*count, elems)
-                        .await
-                        .context(PushBloomFilterValueSnafu)?;
+                if self.codec.pk_col_info(*col_id).is_some() {
+                    sparse_columns.push((*col_id, creator));
                 }
             } else {
                 debug!(
@@ -369,6 +341,25 @@ impl BloomFilterIndexer {
             }
         }
 
+        if !sparse_columns.is_empty() {
+            for (pk, count) in PrimaryKeyRuns::try_new(batch)? {
+                let mut view =
+                    SparsePrimaryKeyView::new(pk, &mut self.pk_offsets).context(DecodeSnafu)?;
+                for (col_id, creator) in &mut sparse_columns {
+                    let value = IndexValueCodec::encode_sparse_value(
+                        &mut view,
+                        *col_id,
+                        &mut self.value_buf,
+                    )
+                    .context(DecodeSnafu)?;
+                    creator
+                        .push_n_row_elem(count, value)
+                        .await
+                        .context(PushBloomFilterValueSnafu)?;
+                }
+            }
+        }
+
         Ok(())
     }
 

src/mito2/src/sst/index/inverted_index/creator.rs

@@ -27,8 +27,10 @@ use index::inverted_index::create::sort_create::SortIndexCreator;
 use index::inverted_index::format::writer::InvertedIndexBlobWriter;
 use index::target::IndexTarget;
 use mito_codec::index::{IndexValueCodec, IndexValuesCodec};
-use mito_codec::row_converter::{CompositeValues, SortField};
+use mito_codec::row_converter::sparse::SparsePrimaryKeyView;
+use mito_codec::row_converter::{SortField, SparseOffsetsCache};
 use puffin::puffin_manager::{PuffinWriter, PutOptions};
+use smallvec::SmallVec;
 use snafu::{ResultExt, ensure};
 use store_api::codec::PrimaryKeyEncoding;
 use store_api::metadata::RegionMetadataRef;
@@ -37,17 +39,18 @@ use tokio::io::duplex;
 use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
 
 use crate::error::{
-    BiErrorsSnafu, EncodeSnafu, IndexFinishSnafu, OperateAbortedIndexSnafu, PuffinAddBlobSnafu,
-    PushIndexValueSnafu, Result,
+    BiErrorsSnafu, DecodeSnafu, EncodeSnafu, IndexFinishSnafu, OperateAbortedIndexSnafu,
+    PuffinAddBlobSnafu, PushIndexValueSnafu, Result,
 };
 use crate::read::Batch;
+use crate::sst::index::TYPE_INVERTED_INDEX;
 use crate::sst::index::intermediate::{
     IntermediateLocation, IntermediateManager, TempFileProvider,
 };
 use crate::sst::index::inverted_index::INDEX_BLOB_TYPE;
+use crate::sst::index::primary_key::PrimaryKeyRuns;
 use crate::sst::index::puffin_manager::SstPuffinWriter;
 use crate::sst::index::statistics::{ByteCount, RowCount, Statistics};
-use crate::sst::index::{TYPE_INVERTED_INDEX, decode_primary_keys_with_counts};
 
 /// The minimum memory usage threshold for one column.
 const MIN_MEMORY_USAGE_THRESHOLD_PER_COLUMN: usize = 1024 * 1024; // 1MB
@@ -66,6 +69,8 @@ pub struct InvertedIndexer {
     codec: IndexValuesCodec,
     /// Reusable buffer for encoding index values.
     value_buf: Vec<u8>,
+    /// Scratch offsets shared by indexed tags of one sparse primary key.
+    pk_offsets: SparseOffsetsCache,
 
     /// Statistics of index creation.
     stats: Statistics,
@@ -124,6 +129,7 @@ impl InvertedIndexer {
             index_creator,
             temp_file_provider,
             value_buf: vec![],
+            pk_offsets: SparseOffsetsCache::new(),
             stats: Statistics::new(TYPE_INVERTED_INDEX),
             aborted: false,
             memory_usage,
@@ -173,7 +179,7 @@ impl InvertedIndexer {
         guard.inc_row_count(batch.num_rows());
 
         let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse;
-        let mut decoded_pks: Option<Vec<(CompositeValues, usize)>> = None;
+        let mut sparse_columns: SmallVec<[(ColumnId, &str); 8]> = SmallVec::new();
 
         for (col_id, target_key) in &self.indexed_column_ids {
             let Some(column_meta) = self.metadata.column_by_id(*col_id) else {
@@ -213,45 +219,8 @@ impl InvertedIndexer {
                     }
                 }
             } else if is_sparse && column_meta.semantic_type == SemanticType::Tag {
-                // Column not found in batch, tries to decode from primary keys for sparse encoding.
-                if decoded_pks.is_none() {
-                    decoded_pks = Some(decode_primary_keys_with_counts(batch, &self.codec)?);
-                }
-
-                let pk_values_with_counts = decoded_pks.as_ref().unwrap();
-                let Some(col_info) = self.codec.pk_col_info(*col_id) else {
-                    debug!(
-                        "Column {} not found in primary key during building bloom filter index",
-                        column_name
-                    );
-                    continue;
-                };
-                let pk_index = col_info.idx;
-                let field = &col_info.field;
-                for (decoded, count) in pk_values_with_counts {
-                    let value = match decoded {
-                        CompositeValues::Dense(dense) => dense.get(pk_index).map(|v| &v.1),
-                        CompositeValues::Sparse(sparse) => sparse.get(col_id),
-                    };
-
-                    let elem = value
-                        .filter(|v| !v.is_null())
-                        .map(|v| {
-                            self.value_buf.clear();
-                            IndexValueCodec::encode_nonnull_value(
-                                v.as_value_ref(),
-                                field,
-                                &mut self.value_buf,
-                            )
-                            .context(EncodeSnafu)?;
-                            Ok(self.value_buf.as_slice())
-                        })
-                        .transpose()?;
-
-                    self.index_creator
-                        .push_with_name_n(target_key, elem, *count)
-                        .await
-                        .context(PushIndexValueSnafu)?;
+                if self.codec.pk_col_info(*col_id).is_some() {
+                    sparse_columns.push((*col_id, target_key));
                 }
             } else {
                 debug!(
@@ -261,12 +230,32 @@ impl InvertedIndexer {
             }
         }
 
+        if !sparse_columns.is_empty() {
+            for (pk, count) in PrimaryKeyRuns::try_new(batch)? {
+                let mut view =
+                    SparsePrimaryKeyView::new(pk, &mut self.pk_offsets).context(DecodeSnafu)?;
+                // Visit all needed tags before moving to the next PK so offset discovery is shared.
+                for &(col_id, target_key) in &sparse_columns {
+                    let value = IndexValueCodec::encode_sparse_value(
+                        &mut view,
+                        col_id,
+                        &mut self.value_buf,
+                    )
+                    .context(DecodeSnafu)?;
+                    self.index_creator
+                        .push_with_name_n(target_key, value, count)
+                        .await
+                        .context(PushIndexValueSnafu)?;
+                }
+            }
+        }
+
         Ok(())
     }
 
     /// Finishes index creation and cleans up garbage.
     /// Returns the number of rows and bytes written.
-    pub async fn finish(
+    pub(crate) async fn finish(
         &mut self,
         puffin_writer: &mut SstPuffinWriter,
     ) -> Result<(RowCount, ByteCount)> {

src/mito2/src/sst/index/primary_key.rs

@@ -0,0 +1,108 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use datatypes::arrow::array::{Array, BinaryArray};
+use datatypes::arrow::record_batch::RecordBatch;
+use snafu::{OptionExt, ensure};
+
+use crate::error::{InvalidRecordBatchSnafu, Result};
+use crate::sst::parquet::flat_format::primary_key_column_index;
+use crate::sst::parquet::format::PrimaryKeyArray;
+
+/// Iterates consecutive dictionary-key runs, preserving their row order without decoding.
+pub(crate) struct PrimaryKeyRuns<'a> {
+    keys: &'a [u32],
+    values: &'a BinaryArray,
+}
+
+impl<'a> PrimaryKeyRuns<'a> {
+    pub(crate) fn try_new(batch: &'a RecordBatch) -> Result<Self> {
+        let pk = batch
+            .column(primary_key_column_index(batch.num_columns()))
+            .as_any()
+            .downcast_ref::<PrimaryKeyArray>()
+            .context(InvalidRecordBatchSnafu {
+                reason: "Primary key column is not a dictionary array",
+            })?;
+        let values = pk.values().as_any().downcast_ref::<BinaryArray>().context(
+            InvalidRecordBatchSnafu {
+                reason: "Primary key values are not binary array",
+            },
+        )?;
+        ensure!(
+            pk.null_count() == 0 && values.null_count() == 0,
+            InvalidRecordBatchSnafu {
+                reason: "Primary keys must not be null"
+            }
+        );
+        Ok(Self {
+            keys: pk.keys().values(),
+            values,
+        })
+    }
+}
+
+impl<'a> Iterator for PrimaryKeyRuns<'a> {
+    type Item = (&'a [u8], usize);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let &key = self.keys.first()?;
+        let count = self
+            .keys
+            .iter()
+            .take_while(|&&current| current == key)
+            .count();
+        self.keys = &self.keys[count..];
+        Some((self.values.value(key as usize), count))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use datatypes::arrow::array::{ArrayRef, BinaryDictionaryBuilder, UInt8Array};
+    use datatypes::arrow::datatypes::UInt32Type;
+
+    use super::*;
+
+    #[test]
+    fn sliced_runs_preserve_order_and_nonconsecutive_keys() {
+        let mut keys = BinaryDictionaryBuilder::<UInt32Type>::new();
+        for key in ["a", "a", "b", "b", "b", "a", "c"] {
+            keys.append(key).unwrap();
+        }
+        let batch = RecordBatch::try_from_iter([
+            ("pk", Arc::new(keys.finish()) as ArrayRef),
+            ("seq", Arc::new(UInt8Array::from(vec![0; 7])) as ArrayRef),
+            ("op", Arc::new(UInt8Array::from(vec![0; 7])) as ArrayRef),
+        ])
+        .unwrap();
+        let slice = batch.slice(1, 5);
+        assert_eq!(
+            PrimaryKeyRuns::try_new(&slice).unwrap().collect::<Vec<_>>(),
+            vec![
+                (b"a".as_slice(), 1),
+                (b"b".as_slice(), 3),
+                (b"a".as_slice(), 1)
+            ]
+        );
+        assert!(
+            PrimaryKeyRuns::try_new(&batch.slice(0, 0))
+                .unwrap()
+                .next()
+                .is_none()
+        );
+    }
+}

src/mito2/src/sst/index/sparse_test.rs

@@ -0,0 +1,198 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use api::v1::SemanticType;
+use datatypes::arrow::array::{
+    ArrayRef, BinaryDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array,
+};
+use datatypes::arrow::datatypes::UInt32Type;
+use datatypes::arrow::record_batch::RecordBatch;
+use datatypes::schema::SkippingIndexOptions;
+use datatypes::value::ValueRef;
+use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl};
+use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader};
+use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec};
+use object_store::ObjectStore;
+use object_store::services::Memory;
+use prost::Message;
+use puffin::puffin_manager::{PuffinManager, PuffinReader};
+use store_api::codec::PrimaryKeyEncoding;
+use store_api::metadata::{RegionMetadataBuilder, RegionMetadataRef};
+use store_api::storage::FileId;
+use store_api::storage::consts::ReservedColumnId;
+
+use super::{IndexBuildType, IndexerBuilder, IndexerBuilderImpl};
+use crate::region::options::IndexOptions;
+use crate::sst::file::{RegionFileId, RegionIndexId};
+use crate::sst::index::bloom_filter::creator::tests::TestPathProvider;
+use crate::sst::index::intermediate::IntermediateManager;
+use crate::sst::index::puffin_manager::PuffinManagerFactory;
+use crate::sst::parquet::flat_format::FlatReadFormat;
+use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema};
+use crate::test_util::sst_util::sst_region_metadata_with_encoding;
+
+fn sparse_input() -> (RegionMetadataRef, RecordBatch) {
+    let mut metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse);
+    for column in &mut metadata.column_metadatas {
+        if column.semantic_type == SemanticType::Tag {
+            column.column_schema.set_inverted_index(true);
+            column
+                .column_schema
+                .set_skipping_options(&SkippingIndexOptions {
+                    granularity: 3,
+                    ..Default::default()
+                })
+                .unwrap();
+        }
+    }
+    let metadata = Arc::new(
+        RegionMetadataBuilder::from_existing(metadata)
+            .build()
+            .unwrap(),
+    );
+    let codec = SparsePrimaryKeyCodec::new(&metadata);
+    let mut keys = BinaryDictionaryBuilder::<UInt32Type>::new();
+    for (series, (first, second, count)) in [
+        (Some("中文12345678"), Some(""), 5),
+        (None, Some("abcdefgh"), 2),
+        (Some(""), None, 4),
+    ]
+    .into_iter()
+    .enumerate()
+    {
+        let mut key = Vec::new();
+        codec
+            .encode_value_refs(
+                &[
+                    (ReservedColumnId::table_id(), ValueRef::UInt32(42)),
+                    (ReservedColumnId::tsid(), ValueRef::UInt64(series as u64)),
+                    (0, first.map_or(ValueRef::Null, ValueRef::String)),
+                    (1, second.map_or(ValueRef::Null, ValueRef::String)),
+                ],
+                &mut key,
+            )
+            .unwrap();
+        for _ in 0..count {
+            keys.append(&key).unwrap();
+        }
+    }
+    let schema = to_flat_sst_arrow_schema(
+        &metadata,
+        &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse),
+    );
+    let columns: Vec<ArrayRef> = vec![
+        Arc::new(UInt64Array::from(vec![1; 11])),
+        Arc::new(TimestampMillisecondArray::from_iter_values(0..11)),
+        Arc::new(keys.finish()),
+        Arc::new(UInt64Array::from(vec![1; 11])),
+        Arc::new(UInt8Array::from(vec![1; 11])),
+    ];
+    (metadata, RecordBatch::try_new(schema, columns).unwrap())
+}
+
+#[tokio::test]
+async fn sparse_and_materialized_tags_produce_identical_indexes() {
+    let (metadata, sparse) = sparse_input();
+    // Use the independent whole-PK read conversion as the oracle for both index types.
+    let materialized = FlatReadFormat::new_with_all_columns(metadata.clone())
+        .convert_batch(sparse.clone(), None)
+        .unwrap();
+    let projection = materialized
+        .schema()
+        .fields()
+        .iter()
+        .enumerate()
+        .filter_map(|(i, field)| {
+            (field.name() == "tag_0" || sparse.column_by_name(field.name()).is_some()).then_some(i)
+        })
+        .collect::<Vec<_>>();
+    let mixed = materialized.project(&projection).unwrap();
+
+    let (dir, factory) = PuffinManagerFactory::new_for_test_async("sparse_index_bytes").await;
+    let puffin = factory.build(
+        ObjectStore::new(Memory::default()).unwrap(),
+        TestPathProvider,
+    );
+    let mut options = IndexOptions::default();
+    options.inverted_index.segment_row_count = 3;
+    let builder = IndexerBuilderImpl {
+        build_type: IndexBuildType::Flush,
+        metadata: metadata.clone(),
+        puffin_manager: puffin.clone(),
+        write_cache_enabled: false,
+        intermediate_manager: IntermediateManager::init_fs(dir.path().to_str().unwrap())
+            .await
+            .unwrap(),
+        index_options: options,
+        inverted_index_config: Default::default(),
+        fulltext_index_config: Default::default(),
+        bloom_filter_index_config: Default::default(),
+        #[cfg(feature = "vector_index")]
+        vector_index_config: Default::default(),
+    };
+    let mut expected = None;
+    for batch in [materialized, sparse, mixed] {
+        let file = RegionFileId::new(metadata.region_id, FileId::random());
+        let mut indexer = builder.build(file, 0, None).await;
+        // Split a repeated PK across batches; runs also cross 3-row segment boundaries.
+        indexer.update_flat(&batch.slice(0, 2)).await;
+        indexer.update_flat(&batch.slice(2, 9)).await;
+        let output = indexer.finish().await;
+        assert_eq!(output.inverted_index.row_count, 11);
+        assert_eq!(output.bloom_filter.row_count, 11);
+        assert_eq!(output.inverted_index.columns.len(), 4);
+        assert_eq!(output.bloom_filter.columns.len(), 4);
+
+        let reader = puffin.reader(&RegionIndexId::new(file, 0)).await.unwrap();
+        let blob = reader.blob("greptime-inverted-index-v1").await.unwrap();
+        let inverted = InvertedIndexBlobReader::new(blob.reader().await.unwrap());
+        let metas = inverted.metadata(None).await.unwrap();
+        let mut index_bytes = BTreeMap::new();
+        // Columns can be emitted in a different order; compare each column's FST and bitmaps.
+        for (name, meta) in &metas.metas {
+            index_bytes.insert(
+                format!("inverted/{name}"),
+                inverted
+                    .range_read(meta.base_offset, meta.inverted_index_size as u32, None)
+                    .await
+                    .unwrap(),
+            );
+        }
+        for id in &metadata.primary_key {
+            let blob = reader
+                .blob(&format!("greptime-bloom-filter-v1-{id}"))
+                .await
+                .unwrap();
+            let bloom = BloomFilterReaderImpl::new(blob.reader().await.unwrap());
+            let meta = bloom.metadata(None).await.unwrap();
+            assert_eq!(meta.row_count, 11);
+            assert_eq!(meta.segment_count, 4);
+            let mut bytes = bloom
+                .range_read(0, meta.bloom_filter_size as u32, None)
+                .await
+                .unwrap()
+                .to_vec();
+            bytes.extend_from_slice(&meta.encode_to_vec());
+            index_bytes.insert(format!("bloom/{id}"), bytes);
+        }
+        if let Some(expected) = &expected {
+            assert_eq!(&index_bytes, expected);
+        } else {
+            expected = Some(index_bytes);
+        }
+    }
+}

src/mito2/src/test_util/bench_util.rs

@@ -33,6 +33,8 @@ use store_api::storage::RegionId;
 use table::predicate::Predicate;
 
 use crate::memtable::KeyValues;
+pub use crate::sst::index::bloom_filter::creator::BloomFilterIndexer;
+pub use crate::sst::index::inverted_index::creator::InvertedIndexer;
 use crate::test_util::memtable_util::region_metadata_to_row_schema;
 
 pub struct Host {