From 7aacc947b30deb79aff8fd18bfe8a9eef942bedf Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Wed, 16 Sep 2026 06:15:52 +0000 Subject: [PATCH] perf(mito2): lazily extract sparse primary key index values (#9176) Signed-off-by: Lei, HUANG --- src/index/src/bloom_filter/creator.rs | 65 +++++ src/mito-codec/src/index.rs | 45 ++- src/mito-codec/src/row_converter/sparse.rs | 3 + .../src/row_converter/sparse/checked.rs | 244 ++++++++++++++++ src/mito2/Cargo.toml | 5 + src/mito2/benches/bench_index_update.rs | 276 ++++++++++++++++++ src/mito2/src/sst/index.rs | 63 +--- .../src/sst/index/bloom_filter/creator.rs | 79 +++-- .../src/sst/index/inverted_index/creator.rs | 79 +++-- src/mito2/src/sst/index/primary_key.rs | 108 +++++++ src/mito2/src/sst/index/sparse_test.rs | 198 +++++++++++++ src/mito2/src/test_util/bench_util.rs | 2 + 12 files changed, 1020 insertions(+), 147 deletions(-) create mode 100644 src/mito-codec/src/row_converter/sparse/checked.rs create mode 100644 src/mito2/benches/bench_index_update.rs create mode 100644 src/mito2/src/sst/index/primary_key.rs create mode 100644 src/mito2/src/sst/index/sparse_test.rs diff --git a/src/index/src/bloom_filter/creator.rs b/src/index/src/bloom_filter/creator.rs index 5a9b5a0b36a..8bbd4cf1c13 100644 --- a/src/index/src/bloom_filter/creator.rs +++ b/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) -> 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()); diff --git a/src/mito-codec/src/index.rs b/src/mito-codec/src/index.rs index 539a3adeddb..e0f2c819afe 100644 --- a/src/mito-codec/src/index.rs +++ b/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, + ) -> Result> { + 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. /// diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index c70293886f8..b2b3e0e3778 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/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. diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs new file mode 100644 index 00000000000..63666f91c4c --- /dev/null +++ b/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 { + 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> { + 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::().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 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()); + } +} diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 100a6fb1023..ccc3fd6f176 100644 --- a/src/mito2/Cargo.toml +++ b/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"] diff --git a/src/mito2/benches/bench_index_update.rs b/src/mito2/benches/bench_index_update.rs new file mode 100644 index 00000000000..eba89cf60ff --- /dev/null +++ b/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, 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::::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 = 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); diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 268b8387397..c8fcd6fbc7e 100644 --- a/src/mito2/src/sst/index.rs +++ b/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> { - let primary_key_index = primary_key_column_index(batch.num_columns()); - let pk_dict_array = batch - .column(primary_key_index) - .as_any() - .downcast_ref::() - .context(InvalidRecordBatchSnafu { - reason: "Primary key column is not a dictionary array", - })?; - let pk_values_array = pk_dict_array - .values() - .as_any() - .downcast_ref::() - .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 = None; - - let pk_indices = keys.values(); - for ¤t_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; diff --git a/src/mito2/src/sst/index/bloom_filter/creator.rs b/src/mito2/src/sst/index/bloom_filter/creator.rs index 7828cdb9fdd..cfdd751bf7f 100644 --- a/src/mito2/src/sst/index/bloom_filter/creator.rs +++ b/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, /// 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> = 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(()) } diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 9cad6fa84ee..0243016fc7a 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/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, + /// 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> = 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)> { diff --git a/src/mito2/src/sst/index/primary_key.rs b/src/mito2/src/sst/index/primary_key.rs new file mode 100644 index 00000000000..4724ae44e6f --- /dev/null +++ b/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 { + let pk = batch + .column(primary_key_column_index(batch.num_columns())) + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "Primary key column is not a dictionary array", + })?; + let values = pk.values().as_any().downcast_ref::().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 { + let &key = self.keys.first()?; + let count = self + .keys + .iter() + .take_while(|&¤t| 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::::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![ + (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() + ); + } +} diff --git a/src/mito2/src/sst/index/sparse_test.rs b/src/mito2/src/sst/index/sparse_test.rs new file mode 100644 index 00000000000..15d26bd1bc7 --- /dev/null +++ b/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::::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 = 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::>(); + 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); + } + } +} diff --git a/src/mito2/src/test_util/bench_util.rs b/src/mito2/src/test_util/bench_util.rs index 8f182e41578..1fa1a53a439 100644 --- a/src/mito2/src/test_util/bench_util.rs +++ b/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 {