diff --git a/AGENTS.md b/AGENTS.md index 819679f304..e8feb4abe4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +69,11 @@ a module map, read/write paths, change-coupling points, and gotchas: - [`tests/compatibility/AGENTS.md`](tests/compatibility/AGENTS.md) - [`tests/perf/AGENTS.md`](tests/perf/AGENTS.md) +## Rust imports + +Prefer crate-rooted imports (`use crate::...`) over `use super::...` in +production code. Relative imports using `super` are allowed in tests. + ## Read before changing code - [`.agents/architecture-invariants.md`](.agents/architecture-invariants.md) — diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index b22ee758bb..faa581331b 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, catalogs, immutable snapshots, file lifecycle, and background maintenance | +| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning and SST builders, catalogs, immutable snapshots, file lifecycle, and background maintenance | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | @@ -62,6 +62,8 @@ filtered `RecordBatch` stream. - **Manifest format** (`manifest/action.rs`): affects crash recovery and follower replay. Keep it backward compatible. - **SST/Parquet layout** (`sst/`): readers must stay compatible with existing files. +- **Series-index coverage** (`series_index/catalog.rs`): `SeriesIndexEntry` stores + compaction-window width and SST summaries keyed by aligned start in both catalogs and Parquet footers. - **Request types** (`request.rs`): usually tied to proto definitions consumed by `datanode`. - **WAL/memtable encoding** (`wal/`, `memtable/`): breaks replay if changed incompatibly. diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index f08178456a..8173050f2a 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -17,10 +17,18 @@ // These components are consumed by the upcoming query and maintenance integration. #[allow(dead_code)] mod catalog; +// Consumed by the follow-up background-maintenance integration. +#[allow(dead_code)] +mod bucket; +// Consumed by the follow-up background-maintenance integration. +#[allow(dead_code)] +mod builder; #[allow(dead_code)] mod purger; mod searcher; mod task; +#[cfg(test)] +mod tests; #[allow(dead_code)] mod version; mod writer; diff --git a/src/mito2/src/series_index/bucket.rs b/src/mito2/src/series_index/bucket.rs new file mode 100644 index 0000000000..7fecb25344 --- /dev/null +++ b/src/mito2/src/series_index/bucket.rs @@ -0,0 +1,877 @@ +// 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. + +//! Event-time bucket planning and source coverage. + +use std::collections::BTreeMap; +use std::time::Duration; + +use common_time::{TimeToLive, Timestamp}; +use smallvec::{SmallVec, smallvec}; +use store_api::storage::FileId; + +use crate::series_index::catalog::{SeriesIndexEntry, WindowSequence}; +use crate::sst::file::FileHandle; + +const SERIES_INDEX_TRIGGER_FILES: usize = 4; + +/// Index files sharing a non-overlapping, half-open time bucket. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct IndexBucket { + pub(crate) start: Timestamp, + pub(crate) end: Timestamp, + pub(crate) index_ids: SmallVec<[FileId; 2]>, + /// Zero means that merged indexes use incompatible window widths. + pub(crate) compaction_window_secs: i64, + /// Indexed SST summaries keyed by aligned start. + /// See [`SeriesIndexEntry::window_sequences`] for the layout and sequence assumption. + pub(crate) window_sequences: BTreeMap, +} + +impl IndexBucket { + pub(crate) fn from_entry(entry: &SeriesIndexEntry) -> Self { + Self { + start: entry.bucket_start, + end: entry.bucket_end, + index_ids: smallvec![entry.index_uuid], + compaction_window_secs: entry.compaction_window_secs, + window_sequences: entry.window_sequences.clone(), + } + } + + fn merge(&mut self, mut other: Self) { + self.start = self.start.min(other.start); + self.end = self.end.max(other.end); + if self.index_ids.is_empty() { + self.compaction_window_secs = other.compaction_window_secs; + self.window_sequences = std::mem::take(&mut other.window_sequences); + } else if !other.index_ids.is_empty() { + if self.compaction_window_secs == other.compaction_window_secs { + merge_window_sequences(&mut self.window_sequences, other.window_sequences); + } else { + // Do not compare maps with different window boundaries, including + // a previous incompatible merge, against current SST coverage. + self.compaction_window_secs = 0; + self.window_sequences.clear(); + } + } + self.index_ids.append(&mut other.index_ids); + } + + /// Inserts a bucket, consuming overlapping entries and expanding their interval. + /// Each map key equals the stored bucket start; adjacent intervals remain separate. + pub(crate) fn insert_into(mut self, buckets: &mut BTreeMap) { + if let Some((&previous_start, previous)) = buckets.range(..self.start).next_back() + && previous.end > self.start + { + self.start = previous_start; + } + // Expanding the end may expose more overlaps, so look up the next entry again. + while let Some((&next_start, _)) = buckets.range(self.start..self.end).next() { + if let Some(other) = buckets.remove(&next_start) { + self.merge(other); + } + } + buckets.insert(self.start, self); + } +} + +/// SSTs grouped into a half-open time interval for an aggregate series-index build. +/// +/// Reconciliation may expand the interval through existing index coverage. A changed +/// bucket is rebuilt from all its SSTs. Unknown source sequences prevent builds, +/// since their coverage cannot be determined safely. +#[derive(Debug, Clone)] +pub(crate) struct SeriesBucket { + pub(crate) start: Timestamp, + pub(crate) end: Timestamp, + pub(crate) files: SmallVec<[FileHandle; 2]>, + pub(crate) has_unknown_sequence: bool, + /// Maximum known source sequence, or zero when all sequences are unknown. + pub(crate) max_file_sequence: u64, + pub(crate) compaction_window_secs: i64, + /// SST summaries keyed by compaction-window-aligned start. + /// Each SST contributes one entry; equal starts merge by maximum end and sequence. + /// Ranges may overlap. See [`WindowSequence`] for the sequence assumption. + /// Missing SST sequences use zero as a placeholder and set `has_unknown_sequence`, + /// preventing index builds and reuse. + pub(crate) window_sequences: BTreeMap, +} + +/// Next bucket coverage and the work required to publish it, computed without I/O. +pub(crate) struct SeriesIndexPlan { + pub(crate) index_buckets: BTreeMap, + pub(crate) builds: Vec<(SeriesBucket, SeriesIndexEntry)>, + pub(crate) expired_index_ids: Vec, + /// Retire only after replacement builds and catalog publication succeed. + pub(crate) superseded_index_ids: Vec, + pub(crate) computed_buckets: usize, + pub(crate) skipped_buckets: usize, +} + +/// Plans whole-bucket replacements when source SST summaries change. The returned bucket +/// map is publishable only after every planned build and the catalog writes succeed. +pub(crate) fn plan_series_indexes( + buckets: Vec, + mut index_buckets: BTreeMap, + ttl: Option, + now_ms: i64, +) -> SeriesIndexPlan { + // Reconciliation changes geometry, not established coverage. In particular, + // a deferred bridge must not change the indexed snapshot. + let buckets = reconcile_series_buckets(buckets, &index_buckets); + let computed_buckets = buckets.len(); + let expired = |end| { + ttl.is_some_and(|ttl| { + ttl.is_expired(&end, &Timestamp::new_millisecond(now_ms)) + .unwrap_or(false) + }) + }; + let mut expired_index_ids = Vec::new(); + index_buckets.retain(|_, bucket| { + if expired(bucket.end) { + expired_index_ids.extend_from_slice(&bucket.index_ids); + false + } else { + true + } + }); + let mut builds = Vec::new(); + let mut superseded_index_ids = Vec::new(); + for bucket in buckets { + if expired(bucket.end) { + continue; + } + if bucket.has_unknown_sequence { + continue; + } + if index_buckets.get(&bucket.start).is_some_and(|indexed| { + indexed.end == bucket.end + && indexed.compaction_window_secs == bucket.compaction_window_secs + && indexed.window_sequences == bucket.window_sequences + }) { + continue; + } + if let Some(entry) = bucket.to_series_entry() { + index_buckets.retain(|_, indexed| { + if indexed.start < bucket.end && bucket.start < indexed.end { + superseded_index_ids.extend_from_slice(&indexed.index_ids); + false + } else { + true + } + }); + IndexBucket::from_entry(&entry).insert_into(&mut index_buckets); + builds.push((bucket, entry)); + } + } + index_buckets.retain(|_, bucket| !bucket.index_ids.is_empty()); + SeriesIndexPlan { + index_buckets, + skipped_buckets: computed_buckets - builds.len(), + builds, + expired_index_ids, + superseded_index_ids, + computed_buckets, + } +} + +pub(crate) fn rounded_bucket_width( + requested: Duration, + compaction_window: Duration, +) -> Option { + let window_secs = i64::try_from(compaction_window.as_secs()).ok()?.max(1); + let requested_secs = i64::try_from(requested.as_secs()) + .unwrap_or(i64::MAX) + .max(1); + let multiples = requested_secs / window_secs + i64::from(requested_secs % window_secs != 0); + multiples.checked_mul(window_secs) +} + +/// Groups SSTs across levels into sorted, disjoint time buckets without planning builds. +/// +/// Both widths must be positive, and `width_secs` must be a multiple of the compaction +/// window width. Inclusive SST ranges are rounded outward to aligned, +/// half-open intervals in seconds. Overlapping intervals merge; adjacent ones stay +/// separate. Each merged bucket tracks its maximum sequence and any unknown sequence. +pub(crate) fn group_files_into_series_buckets( + files: &[FileHandle], + width_secs: i64, + compaction_window_secs: i64, +) -> Vec { + let mut spans = files + .iter() + .map(|file| { + let start = file.time_range().0.split().0; + let end = file.time_range().1.split().0; + let sequence = file + .meta_ref() + .sequence + .map_or(0, |sequence| sequence.get()); + let first_window = start.div_euclid(compaction_window_secs); + let last_window = end.div_euclid(compaction_window_secs); + let window_sequences = BTreeMap::from([( + first_window.saturating_mul(compaction_window_secs), + WindowSequence { + start: first_window.saturating_mul(compaction_window_secs), + end: last_window + .saturating_add(1) + .saturating_mul(compaction_window_secs), + max_sequence: sequence, + }, + )]); + SeriesBucket { + start: Timestamp::new_second( + start.div_euclid(width_secs).saturating_mul(width_secs), + ), + end: Timestamp::new_second( + end.div_euclid(width_secs) + .saturating_add(1) + .saturating_mul(width_secs), + ), + files: smallvec![file.clone()], + has_unknown_sequence: file.meta_ref().sequence.is_none(), + max_file_sequence: sequence, + compaction_window_secs, + window_sequences, + } + }) + .collect::>(); + spans.sort_unstable_by(|a, b| a.start.cmp(&b.start).then_with(|| b.end.cmp(&a.end))); + group_series_buckets(spans) +} + +/// Expands SST buckets through existing index coverage before grouping build inputs. +/// +/// `buckets` must be sorted by start and disjoint. `index_buckets` must contain +/// disjoint intervals keyed by their starts. Expansion preserves start ordering, +/// but may introduce overlaps, which are merged in the returned disjoint buckets. +/// Established index coverage is borrowed so deferred builds cannot modify it. +fn reconcile_series_buckets( + mut buckets: Vec, + index_buckets: &BTreeMap, +) -> Vec { + for bucket in &mut buckets { + if let Some((&start, previous)) = index_buckets.range(..=bucket.start).next_back() + && previous.end > bucket.start + { + bucket.start = start; + bucket.end = bucket.end.max(previous.end); + } + // Disjoint index intervals make the last overlapping interval's end the + // furthest boundary. Expanding to it cannot expose another index interval. + if let Some((_, last)) = index_buckets.range(bucket.start..bucket.end).next_back() { + bucket.end = bucket.end.max(last.end); + } + } + group_series_buckets(buckets) +} + +/// Merges overlapping spans sorted by nondecreasing start; equal starts are allowed. +/// Returns sorted, disjoint buckets. Adjacent half-open intervals remain separate. +/// All spans must use the same compaction-window width for their SST summaries. +fn group_series_buckets(spans: Vec) -> Vec { + let mut buckets: Vec = Vec::new(); + for mut span in spans { + if let Some(last) = buckets.last_mut() + && span.start < last.end + { + last.end = last.end.max(span.end); + last.files.append(&mut span.files); + last.has_unknown_sequence |= span.has_unknown_sequence; + last.max_file_sequence = last.max_file_sequence.max(span.max_file_sequence); + merge_window_sequences(&mut last.window_sequences, span.window_sequences); + } else { + buckets.push(span); + } + } + buckets +} + +/// Merges summaries sharing a start, relying on [`WindowSequence`]'s sequence assumption. +fn merge_window_sequences( + target: &mut BTreeMap, + source: BTreeMap, +) { + for (start, summary) in source { + target + .entry(start) + .and_modify(|current| { + current.end = current.end.max(summary.end); + current.max_sequence = current.max_sequence.max(summary.max_sequence); + }) + .or_insert(summary); + } +} + +impl SeriesBucket { + /// Creates entry metadata with a fresh UUID and sorted source file IDs. + /// Returns `None` for unknown sequences or too few files. + fn to_series_entry(&self) -> Option { + if self.has_unknown_sequence || self.files.len() < SERIES_INDEX_TRIGGER_FILES { + return None; + } + let mut source_file_ids = self + .files + .iter() + .map(|file| file.file_id().file_id()) + .collect::>(); + source_file_ids.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes())); + let min_file_sequence = self + .files + .iter() + .filter_map(|file| file.meta_ref().sequence.map(|sequence| sequence.get())) + .min()?; + Some(SeriesIndexEntry { + index_uuid: FileId::random(), + bucket_start: self.start, + bucket_end: self.end, + source_file_ids, + min_file_sequence, + max_file_sequence: self.max_file_sequence, + compaction_window_secs: self.compaction_window_secs, + window_sequences: self.window_sequences.clone(), + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::num::NonZeroU64; + + use common_time::timestamp::TimeUnit; + + use super::*; + use crate::sst::file::FileMeta; + use crate::test_util::new_noop_file_purger; + + fn coverage(intervals: &[(i64, i64, u64)]) -> BTreeMap { + intervals + .iter() + .map(|&(start, end, max_sequence)| { + ( + start, + WindowSequence { + start, + end, + max_sequence, + }, + ) + }) + .collect() + } + + fn file(sequence: Option, level: u8, start: Timestamp, end: Timestamp) -> FileHandle { + FileHandle::new( + FileMeta { + file_id: FileId::random(), + sequence: sequence.and_then(NonZeroU64::new), + level, + time_range: (start, end), + ..Default::default() + }, + new_noop_file_purger(), + ) + } + + #[test] + fn test_second_resolution_buckets_merge_spans_across_levels() { + let width = rounded_bucket_width(Duration::from_secs(11), Duration::from_secs(10)).unwrap(); + let files = [ + file( + Some(1), + 0, + Timestamp::new_millisecond(-1), + Timestamp::new_millisecond(19999), + ), + file( + Some(2), + 1, + Timestamp::new_microsecond(19000000), + Timestamp::new_microsecond(39000000), + ), + file( + Some(3), + 2, + Timestamp::new_nanosecond(39000000000), + Timestamp::new_nanosecond(40000000000), + ), + file( + Some(4), + 1, + Timestamp::new_second(60), + Timestamp::new_second(61), + ), + ]; + let buckets = group_files_into_series_buckets(&files, width, 10); + let spans = buckets + .iter() + .map(|b| (b.start, b.end, b.files.len())) + .collect::>(); + assert_eq!( + vec![ + (Timestamp::new_second(-20), Timestamp::new_second(60), 3), + (Timestamp::new_second(60), Timestamp::new_second(80), 1), + ], + spans + ); + assert_eq!(3, buckets[0].max_file_sequence); + assert_eq!(4, buckets[1].max_file_sequence); + assert!(buckets[0].to_series_entry().is_none()); + assert!(buckets[1].to_series_entry().is_none()); + let mut files = files.to_vec(); + files.push(file( + None, + 0, + Timestamp::new_second(0), + Timestamp::new_second(1), + )); + assert!( + group_files_into_series_buckets(&files, width, 10)[0] + .to_series_entry() + .is_none() + ); + } + + #[test] + fn test_seconds_do_not_require_millisecond_conversion() { + let start = Timestamp::new_second(i64::MAX / 1000 + 100); + assert!(start.convert_to(TimeUnit::Millisecond).is_none()); + let buckets = group_files_into_series_buckets(&[file(Some(1), 0, start, start)], 1, 1); + assert_eq!(start, buckets[0].start); + assert_eq!(Timestamp::new_second(start.value() + 1), buckets[0].end); + + let width = rounded_bucket_width(Duration::ZERO, Duration::from_millis(100)).unwrap(); + let buckets = group_files_into_series_buckets( + &[file( + Some(1), + 0, + Timestamp::new_nanosecond(-1), + Timestamp::new_microsecond(1), + )], + width, + 1, + ); + assert_eq!( + (Timestamp::new_second(-1), Timestamp::new_second(1)), + (buckets[0].start, buckets[0].end) + ); + } + + #[test] + fn test_reconcile_bridge_expands_through_index_and_sst_buckets() { + let ts = Timestamp::new_second; + let mut indexes = BTreeMap::new(); + let ids = [FileId::random(), FileId::random(), FileId::random()]; + for (start, end, max_sequence, id) in [ + (0, 20, 10, ids[0]), + (30, 60, 30, ids[1]), + (70, 100, 20, ids[2]), + ] { + IndexBucket { + start: ts(start), + end: ts(end), + index_ids: smallvec![id], + compaction_window_secs: 10, + window_sequences: coverage(&[(start, end, max_sequence)]), + } + .insert_into(&mut indexes); + } + let files = [ + file(Some(15), 1, ts(10), ts(30)), + file(Some(31), 0, ts(50), ts(70)), + file(Some(32), 0, ts(90), ts(95)), + file(Some(32), 1, ts(90), ts(95)), + file(Some(33), 0, ts(90), ts(95)), + file(Some(34), 0, ts(100), ts(105)), + ]; + let planned = group_files_into_series_buckets(&files, 10, 10); + assert_eq!(4, planned.len()); + let plan = plan_series_indexes(planned, indexes, None, 0); + assert_eq!((2, 1), (plan.computed_buckets, plan.skipped_buckets)); + let [(bucket, entry)] = plan.builds.as_slice() else { + panic!("expected one replacement build"); + }; + assert_eq!((ts(0), ts(100)), (bucket.start, bucket.end)); + // Include sequence 15 and both files at 32, but exclude the adjacent bucket. + let expected = files[..5] + .iter() + .map(|file| file.file_id().file_id()) + .collect::>(); + assert_eq!( + expected, + bucket + .files + .iter() + .map(|file| file.file_id().file_id()) + .collect() + ); + assert_eq!(expected, entry.source_file_ids.iter().copied().collect()); + assert_eq!((15, 33), (entry.min_file_sequence, entry.max_file_sequence)); + assert!(plan.expired_index_ids.is_empty()); + assert_eq!(1, plan.index_buckets.len()); + let merged = &plan.index_buckets[&ts(0)]; + assert_eq!((ts(0), ts(100)), (merged.start, merged.end)); + assert_eq!(33, bucket.max_file_sequence); + assert_eq!(bucket.window_sequences, merged.window_sequences); + assert_eq!(ids.as_slice(), plan.superseded_index_ids.as_slice()); + assert_eq!([entry.index_uuid].as_slice(), merged.index_ids.as_slice()); + } + + #[test] + fn test_plan_reuses_replaced_sources_and_rebuilds_changed_bucket() { + let ts = Timestamp::new_second; + let make_file = |sequence| file(Some(sequence), 0, ts(1), ts(2)); + let original = (1..=4).map(make_file).collect::>(); + let initial = plan_series_indexes( + group_files_into_series_buckets(&original, 10, 10), + BTreeMap::new(), + None, + 0, + ); + let first_id = initial.builds[0].1.index_uuid; + + // New file IDs with already indexed sequences do not invalidate the index. + let mut files = (1..=4).map(make_file).collect::>(); + let replaced = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + initial.index_buckets.clone(), + None, + 0, + ); + assert!(replaced.builds.is_empty()); + assert!(replaced.expired_index_ids.is_empty()); + assert_eq!(initial.index_buckets, replaced.index_buckets); + + files.push(make_file(5)); + let ready = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + replaced.index_buckets, + None, + 0, + ); + let [(bucket, entry)] = ready.builds.as_slice() else { + panic!("a changed window must rebuild the whole bucket"); + }; + let expected = files[..] + .iter() + .map(|file| file.file_id().file_id()) + .collect::>(); + assert_eq!( + expected, + bucket + .files + .iter() + .map(|file| file.file_id().file_id()) + .collect() + ); + assert_eq!(expected, entry.source_file_ids.iter().copied().collect()); + assert_eq!((1, 5), (entry.min_file_sequence, entry.max_file_sequence)); + assert_eq!( + [entry.index_uuid].as_slice(), + ready.index_buckets[&ts(0)].index_ids.as_slice() + ); + assert_eq!(vec![first_id], ready.superseded_index_ids); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + ready.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(ready.index_buckets, repeated.index_buckets); + + // TTL retires the replacement; the previous index is already superseded. + let expired = plan_series_indexes( + Vec::new(), + ready.index_buckets, + Some(TimeToLive::Duration(Duration::from_secs(10))), + 21_000, + ); + assert!(expired.builds.is_empty()); + assert!(expired.index_buckets.is_empty()); + assert_eq!( + [entry.index_uuid].as_slice(), + expired.expired_index_ids.as_slice() + ); + } + + #[test] + fn test_sst_summaries_round_ranges_outward() { + let ts = Timestamp::new_millisecond; + let files = [ + file(Some(10), 0, ts(-1), ts(19_999)), + file(Some(30), 1, ts(10_000), ts(20_000)), + ]; + let buckets = group_files_into_series_buckets(&files, 100, 10); + assert_eq!(1, buckets.len()); + assert_eq!( + coverage(&[(-10, 20, 10), (10, 30, 30)]), + buckets[0].window_sequences + ); + } + + #[rstest::rstest] + #[case(32)] + #[case(33)] + #[case(100_000)] + fn test_wide_sst_reuses_unchanged_inputs_and_rebuilds_after_splitting(#[case] windows: i64) { + let ts = Timestamp::new_second; + let end = windows * 10; + let files = (1..=4) + .map(|seq| file(Some(seq), 0, ts(0), ts(end - 1))) + .collect::>(); + let initial = plan_series_indexes( + group_files_into_series_buckets(&files, end, 10), + BTreeMap::new(), + None, + 0, + ); + assert_eq!(1, initial.builds.len()); + assert_eq!( + coverage(&[(0, end, 4)]), + initial.builds[0].1.window_sequences + ); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, end, 10), + initial.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(initial.index_buckets, repeated.index_buckets); + + // Splitting changes the summaries and conservatively triggers one rebuild. + let split = (1..=4) + .rev() + .flat_map(|seq| { + [ + file(Some(seq), 1, ts(end / 2), ts(end - 1)), + file(Some(seq), 1, ts(0), ts(end / 2 - 1)), + ] + }) + .collect::>(); + let replaced = plan_series_indexes( + group_files_into_series_buckets(&split, end, 10), + repeated.index_buckets, + None, + 0, + ); + assert_eq!(1, replaced.builds.len()); + assert_eq!(1, replaced.superseded_index_ids.len()); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&split, end, 10), + replaced.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(replaced.index_buckets, repeated.index_buckets); + } + + #[rstest::rstest] + #[case::existing_start(0)] + #[case::new_start(20)] + fn test_new_data_changes_sst_summary(#[case] start: i64) { + let ts = Timestamp::new_second; + let mut files = vec![ + file(Some(1), 0, ts(1), ts(29)), + file(Some(10), 0, ts(1), ts(29)), + file(Some(30), 0, ts(1), ts(9)), + file(Some(20), 0, ts(41), ts(49)), + ]; + let initial = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + BTreeMap::new(), + None, + 0, + ); + assert_eq!(1, initial.builds.len()); + // The maximum end survives even when the highest sequence is in a shorter SST. + assert_eq!( + coverage(&[(0, 30, 30), (40, 50, 20)]), + initial.builds[0].1.window_sequences + ); + // New data has a sequence greater than those in the indexed snapshot. + files.push(file(Some(31), 0, ts(start + 1), ts(start + 2))); + let plan = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + initial.index_buckets, + None, + 0, + ); + let [(bucket, entry)] = plan.builds.as_slice() else { + panic!("new data must rebuild the bucket"); + }; + assert_eq!(files.len(), bucket.files.len()); + assert_eq!(31, entry.max_file_sequence); + let expected = if start == 0 { + coverage(&[(0, 30, 31), (40, 50, 20)]) + } else { + coverage(&[(0, 30, 30), (20, 30, 31), (40, 50, 20)]) + }; + assert_eq!(expected, entry.window_sequences); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + plan.index_buckets, + None, + 0, + ); + assert!(repeated.builds.is_empty()); + } + + #[test] + fn test_deferred_bridge_preserves_established_coverage() { + let ts = Timestamp::new_second; + let mut indexes = BTreeMap::new(); + for (start, seq) in [(0, 10), (20, 30)] { + IndexBucket { + start: ts(start), + end: ts(start + 10), + index_ids: smallvec![FileId::random()], + compaction_window_secs: 10, + window_sequences: coverage(&[(start, start + 10, seq)]), + } + .insert_into(&mut indexes); + } + let mut files = vec![ + file(Some(13), 0, ts(1), ts(21)), + file(Some(30), 0, ts(21), ts(22)), + ]; + let deferred = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + indexes.clone(), + None, + 0, + ); + assert!(deferred.builds.is_empty()); + assert!(deferred.superseded_index_ids.is_empty()); + assert_eq!(indexes, deferred.index_buckets); + files.extend((11..=12).map(|seq| file(Some(seq), 0, ts(1), ts(2)))); + let plan = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + deferred.index_buckets, + None, + 0, + ); + assert_eq!(1, plan.builds.len()); + assert_eq!(4, plan.builds[0].0.files.len()); + assert_eq!(2, plan.superseded_index_ids.len()); + assert_eq!( + coverage(&[(0, 30, 13), (20, 30, 30)]), + plan.builds[0].1.window_sequences + ); + + files.push(file(None, 0, ts(1), ts(2))); + let unknown = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + indexes.clone(), + None, + 0, + ); + assert!(unknown.builds.is_empty()); + assert_eq!(indexes, unknown.index_buckets); + } + + #[rstest::rstest] + #[case::removed_window(100, 10, false)] + #[case::added_window(100, 10, true)] + #[case::window_width(100, 20, false)] + #[case::bucket_width(200, 10, false)] + fn test_coverage_shape_changes_rebuild( + #[case] bucket_width: i64, + #[case] window_width: i64, + #[case] add_window: bool, + ) { + let ts = Timestamp::new_second; + let mut files = (27..=30) + .map(|seq| file(Some(seq), 0, ts(11), ts(12))) + .collect::>(); + let mut original = files.clone(); + original.push(file(Some(10), 0, ts(1), ts(2))); + let initial = plan_series_indexes( + group_files_into_series_buckets(&original, 100, 10), + BTreeMap::new(), + None, + 0, + ); + if add_window { + files = original; + files.push(file(Some(15), 0, ts(21), ts(22))); + } else if bucket_width != 100 || window_width != 10 { + files = original; + } + let plan = plan_series_indexes( + group_files_into_series_buckets(&files, bucket_width, window_width), + initial.index_buckets, + None, + 0, + ); + assert_eq!(1, plan.builds.len()); + assert_eq!(1, plan.superseded_index_ids.len()); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, bucket_width, window_width), + plan.index_buckets, + None, + 0, + ); + assert!(repeated.builds.is_empty()); + } + + #[test] + fn test_index_map_merge_is_order_independent() { + let ts = Timestamp::new_second; + let make_index = |start, end, width, windows: &[(i64, u64)]| IndexBucket { + start: ts(start), + end: ts(end), + index_ids: smallvec![FileId::random()], + compaction_window_secs: width, + window_sequences: windows + .iter() + .map(|&(start, max_sequence)| { + ( + start, + WindowSequence { + start, + end: start + width, + max_sequence, + }, + ) + }) + .collect(), + }; + let indexes = [ + make_index(0, 20, 10, &[(0, 10), (10, 20)]), + make_index(10, 30, 10, &[(10, 30), (20, 15)]), + make_index(20, 40, 10, &[(20, 25), (30, 5)]), + ]; + for order in [[0, 1, 2], [2, 1, 0], [1, 0, 2], [0, 2, 1]] { + let mut map = BTreeMap::new(); + for i in order { + indexes[i].clone().insert_into(&mut map); + } + assert_eq!(1, map.len()); + assert_eq!(10, map[&ts(0)].compaction_window_secs); + assert_eq!( + coverage(&[(0, 10, 10), (10, 20, 30), (20, 30, 25), (30, 40, 5)]), + map[&ts(0)].window_sequences + ); + make_index(10, 30, 20, &[(0, 30)]).insert_into(&mut map); + indexes[0].clone().insert_into(&mut map); + assert_eq!(0, map[&ts(0)].compaction_window_secs); + assert!(map[&ts(0)].window_sequences.is_empty()); + } + } +} diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs new file mode 100644 index 0000000000..8d361be364 --- /dev/null +++ b/src/mito2/src/series_index/builder.rs @@ -0,0 +1,742 @@ +// 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. + +//! Independently builds range and aggregate series indexes from SST readers. + +use std::sync::Arc; + +use async_stream::try_stream; +use common_telemetry::warn; +use futures::TryStreamExt; +use object_store::ObjectStore; +use snafu::{OptionExt, ensure}; +use store_api::storage::FileId; + +use crate::error::{Result, UnexpectedSnafu}; +use crate::read::BoxedRecordBatchStream; +use crate::read::compat::FlatCompatBatch; +use crate::read::flat_merge::FlatMergeReader; +use crate::read::flat_projection::FlatProjectionMapper; +use crate::read::prune::FlatPruneReader; +use crate::read::read_columns::ReadColumns; +use crate::region::MitoRegionRef; +use crate::region::version::VersionRef; +use crate::series_index::bucket::SeriesBucket; +use crate::series_index::catalog::{ + SeriesIndexEntry, range_index_path, series_index_path, series_metadata, +}; +use crate::series_index::purger::{IndexFilePurger, IndexFileType, file_operation}; +use crate::series_index::version::SeriesIndexFileHandle; +use crate::series_index::{SeriesIndexWriter, SeriesIndexWriterOptions}; +use crate::sst::file::FileHandle; +use crate::sst::parquet::reader::{FlatRowGroupReader, ReaderMetrics}; +use crate::sst::parquet::row_group::ParquetFetchMetrics; +use crate::sst::range_index::{SstRangeIndexWriter, SstRangeIndexWriterOptions}; + +async fn reader_input( + region: &MitoRegionRef, + file: FileHandle, +) -> Result< + Option<( + Arc, + crate::sst::parquet::row_selection::RowGroupSelection, + )>, +> { + Ok(region + .access_layer + .read_sst(file) + .projection(Some(ReadColumns::new([]))) + .build_reader_input(&mut ReaderMetrics::default()) + .await? + .map(|(context, selection)| (Arc::new(context), selection))) +} + +/// Builds one range index, or returns `None` when the SST has no readable input. +pub(crate) async fn build_range_index( + store: &ObjectStore, + region: &MitoRegionRef, + version: &VersionRef, + file: FileHandle, +) -> Result> { + let file_id = file.file_id().file_id(); + let Some((context, mut selection)) = reader_input(region, file).await? else { + return Ok(None); + }; + let mapper = FlatProjectionMapper::new(&version.metadata, [])?; + let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; + let path = range_index_path(region.region_id, file_id); + let mut writer = SstRangeIndexWriter::try_new( + version.metadata.clone(), + store.clone(), + &path, + SstRangeIndexWriterOptions::default(), + ) + .await?; + let result: Result<()> = async { + let fetch_metrics = ParquetFetchMetrics::default(); + while let Some((row_group_id, row_selection)) = selection.pop_first() { + let parquet_reader = context + .reader_builder() + .build(context.build_context( + row_group_id, + Some(row_selection), + Some(&fetch_metrics), + )) + .await?; + let mut reader = FlatPruneReader::new_with_row_group_reader( + context.clone(), + FlatRowGroupReader::new(context.clone(), parquet_reader), + context.pre_filter_mode().skip_fields(), + ); + while let Some(batch) = reader.next_batch().await? { + let batch = match &compat { + Some(compat) => compat.compat(batch)?, + None => batch, + }; + writer.write(row_group_id as u32, &batch).await?; + } + } + Ok(()) + } + .await; + if let Err(error) = result { + if let Err(cleanup_error) = writer.abort().await { + warn!(cleanup_error; "Failed to abort range-index build"); + } + return Err(error); + } + writer.finish().await?; + file_operation(IndexFileType::Range, "build", "success"); + Ok(Some(file_id)) +} + +/// Builds only the series index. Callers build needed range indexes separately. +pub(crate) async fn build_series_index( + store: &ObjectStore, + region: &MitoRegionRef, + version: &VersionRef, + bucket: &SeriesBucket, + entry: &SeriesIndexEntry, + purger: &IndexFilePurger, +) -> Result { + let mut sources = Vec::::new(); + let mapper = FlatProjectionMapper::new(&version.metadata, [])?; + let schema = mapper.input_arrow_schema(false); + for file in &bucket.files { + let Some((context, mut selection)) = reader_input(region, file.clone()).await? else { + continue; + }; + let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; + sources.push(Box::pin(try_stream! { + let fetch_metrics = ParquetFetchMetrics::default(); + while let Some((row_group_id, row_selection)) = selection.pop_first() { + let parquet_reader = context.reader_builder().build(context.build_context( + row_group_id, + Some(row_selection), + Some(&fetch_metrics), + )).await?; + let mut reader = FlatPruneReader::new_with_row_group_reader( + context.clone(), + FlatRowGroupReader::new(context.clone(), parquet_reader), + context.pre_filter_mode().skip_fields(), + ); + while let Some(batch) = reader.next_batch().await? { + yield match &compat { + Some(compat) => compat.compat(batch)?, + None => batch, + }; + } + } + })); + } + ensure!( + !sources.is_empty(), + UnexpectedSnafu { + reason: "series-index bucket has no readable SST", + } + ); + let mut visible: BoxedRecordBatchStream = if sources.len() == 1 { + sources.pop().context(UnexpectedSnafu { + reason: "series-index source disappeared", + })? + } else { + Box::pin( + FlatMergeReader::new(schema, sources, 8192, None) + .await? + .into_stream(), + ) + }; + // TODO(yingwen): Deduplicate update-mode rows before series indexes are used by queries. + let path = series_index_path(region.region_id, entry.index_uuid); + let mut writer = SeriesIndexWriter::try_new( + version.metadata.clone(), + store.clone(), + &path, + SeriesIndexWriterOptions::default(), + Some(series_metadata(entry)?), + ) + .await?; + let result: Result<()> = async { + while let Some(batch) = visible.try_next().await? { + writer.write(&batch).await?; + } + Ok(()) + } + .await; + if let Err(error) = result { + if let Err(cleanup_error) = writer.abort().await { + warn!(cleanup_error; "Failed to abort series-index build"); + } + return Err(error); + } + writer.finish().await?; + file_operation(IndexFileType::Series, "build", "success"); + Ok(SeriesIndexFileHandle::new( + region.region_id, + entry.clone(), + purger.clone(), + )) +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, HashMap}; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, Ordering}; + + use datatypes::data_type::ConcreteDataType; + use object_store::layers::mock::{self, MockLayerBuilder, oio}; + use object_store::services::Memory; + use store_api::region_engine::RegionEngine; + + use super::*; + use crate::series_index::bucket::{group_files_into_series_buckets, plan_series_indexes}; + use crate::series_index::catalog::{ + SeriesIndexCatalog, load_version_control, series_catalog_path, store_catalog, + }; + use crate::series_index::purger::series_index_channel; + use crate::series_index::tests::prepare_region; + use crate::test_util::TestEnv; + + fn build_input(version: &VersionRef) -> (SeriesBucket, SeriesIndexEntry) { + let files = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .cloned() + .collect::>(); + let mut plan = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + BTreeMap::new(), + None, + 0, + ); + assert_eq!(1, plan.builds.len()); + plan.builds.pop().unwrap() + } + + fn metadata_with_seconds(version: &VersionRef) -> store_api::metadata::RegionMetadataRef { + let mut metadata = (*version.metadata).clone(); + let time_index = metadata.time_index_column_pos(); + metadata.column_metadatas[time_index] + .column_schema + .data_type = ConcreteDataType::timestamp_second_datatype(); + Arc::new(metadata) + } + + #[rstest::rstest] + #[case::range_then_series(true, false)] + #[case::series_only(false, false)] + #[case::series_failure(true, true)] + #[tokio::test] + async fn test_build_indexes_without_publication( + #[case] build_ranges: bool, + #[case] fail_series: bool, + ) { + let mut env = TestEnv::with_prefix("series-builder").await; + let (engine, region) = prepare_region(&mut env).await; + let version = region.version(); + let (bucket, entry) = build_input(&version); + let files = &bucket.files; + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, mut receiver) = series_index_channel(store.clone()); + let mut range_bytes = HashMap::new(); + if build_ranges { + for file in files { + let id = build_range_index(&store, ®ion, &version, file.clone()) + .await + .unwrap() + .unwrap(); + range_bytes.insert( + id, + store + .read(&range_index_path(region.region_id, id)) + .await + .unwrap() + .to_bytes(), + ); + } + } + // Both builders use the captured version, even if live metadata changes. + region + .version_control + .alter_metadata(metadata_with_seconds(&version)); + let layer = writer_layer(&WriterStates::default(), move |path| { + assert!(path.contains("/series/"), "series stage opened {path}"); + if fail_series { + WriterFailure::Finish + } else { + WriterFailure::None + } + }); + let result = build_series_index( + &store.clone().layer(layer), + ®ion, + &version, + &bucket, + &entry, + &purger, + ) + .await; + if fail_series { + assert!(result.is_err()); + } else { + let handle = result.unwrap(); + assert_eq!(handle.entry(), &entry); + assert!( + store + .exists(&series_index_path(region.region_id, entry.index_uuid)) + .await + .unwrap() + ); + store_catalog( + &store, + &series_catalog_path(region.region_id), + &SeriesIndexCatalog { + indexes: vec![handle.entry().clone()], + }, + ) + .await + .unwrap(); + let recovered = load_version_control(&store, region.region_id, &purger) + .await + .current(); + let repeated = plan_series_indexes( + group_files_into_series_buckets(files, 100, 10), + recovered.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(recovered.index_buckets, repeated.index_buckets); + } + // Series success or failure leaves completed range indexes unchanged. + for file in files { + let id = file.file_id().file_id(); + let path = range_index_path(region.region_id, id); + if let Some(bytes) = range_bytes.get(&id) { + assert_eq!(*bytes, store.read(&path).await.unwrap().to_bytes()); + } else { + assert!(!store.exists(&path).await.unwrap()); + } + } + assert!(region.series_index_version().range_indexes.is_empty()); + assert!(region.series_index_version().series_indexes.is_empty()); + assert!(receiver.try_recv().is_err()); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_build_indexes_after_time_index_widening() { + use api::v1::helper::row; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; + use datatypes::arrow::array::{TimestampMicrosecondArray, UInt64Array}; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use store_api::region_request::{ + AlterKind, ModifyColumnType, RegionAlterRequest, RegionPutRequest, RegionRequest, + }; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + + use crate::test_util::sst_util::new_sparse_primary_key; + use crate::test_util::{CreateRequestBuilder, flush_region, rows_schema}; + + let mut env = TestEnv::with_prefix("series-builder-widen").await; + let (engine, region) = prepare_region(&mut env).await; + let metadata = region.version().metadata.clone(); + engine + .handle_request( + region.region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnType { + column_name: metadata.time_index_column().column_schema.name.clone(), + target_type: ConcreteDataType::timestamp_microsecond_datatype(), + }], + }, + }), + ) + .await + .unwrap(); + + let mut request = CreateRequestBuilder::new().build(); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + let full_schema = rows_schema(&request); + let mut pk = full_schema[0].clone(); + pk.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + pk.datatype = ColumnDataType::Binary.into(); + pk.semantic_type = SemanticType::Tag.into(); + let mut ts = full_schema[5].clone(); + ts.datatype = ColumnDataType::TimestampMicrosecond.into(); + engine + .handle_request( + region.region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows: Rows { + schema: vec![pk, ts, full_schema[4].clone()], + rows: [500_500, 2_500_500, 3_500_500] + .into_iter() + .map(|ts| { + row(vec![ + ValueData::BinaryValue(new_sparse_primary_key( + &["a", "x"], + &metadata, + 10, + 0, + )), + ValueData::TimestampMicrosecondValue(ts), + ValueData::U64Value(1), + ]) + }) + .collect(), + }, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + partition_expr_version: None, + }), + ) + .await + .unwrap(); + flush_region(&engine, region.region_id, None).await; + let version = region.version(); + let (bucket, entry) = build_input(&version); + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, _receiver) = series_index_channel(store.clone()); + for file in &bucket.files { + build_range_index(&store, ®ion, &version, file.clone()) + .await + .unwrap() + .unwrap(); + } + let _handle = build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + .await + .unwrap(); + let bytes = store + .read(&series_index_path(region.region_id, entry.index_uuid)) + .await + .unwrap() + .to_bytes(); + let batches = ParquetRecordBatchReaderBuilder::try_new(bytes) + .unwrap() + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(1, batches.len()); + let batch = &batches[0]; + assert_eq!(1, batch.num_rows()); + for (column, expected) in [(0, 500_500), (1, 4_000_000)] { + assert_eq!( + expected, + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + } + assert_eq!( + 7, + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_range_failure_preserves_completed_indexes_and_stops_series_stage() { + let mut env = TestEnv::with_prefix("range-builder-stage-failure").await; + let (engine, region) = prepare_region(&mut env).await; + let version = region.version(); + let (bucket, entry) = build_input(&version); + let files = &bucket.files; + let failed_path = range_index_path(region.region_id, files[1].file_id().file_id()); + let states = WriterStates::default(); + let layer = writer_layer(&states, move |path| { + if path == failed_path { + WriterFailure::Finish + } else { + WriterFailure::None + } + }); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let (purger, _receiver) = series_index_channel(store.clone()); + let mut completed = Vec::new(); + let result: Result> = async { + for file in files { + let Some(id) = build_range_index(&store, ®ion, &version, file.clone()).await? + else { + // Defer series construction if a needed range is not ready. + return Ok(None); + }; + completed.push(id); + } + build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + .await + .map(Some) + } + .await; + assert!(format!("{:?}", result.unwrap_err()).contains("injected index finish failure")); + assert_eq!(vec![files[0].file_id().file_id()], completed); + { + let states = states.lock().unwrap(); + assert_eq!(2, states.len()); + assert!(states.keys().all(|path| !path.contains("/series/"))); + assert_eq!( + 1, + states[&range_index_path(region.region_id, completed[0])].closed + ); + assert_eq!( + 1, + states[&range_index_path(region.region_id, files[1].file_id().file_id())].aborted + ); + } + for (i, file) in files.iter().enumerate() { + assert_eq!( + i == 0, + store + .exists(&range_index_path( + region.region_id, + file.file_id().file_id() + )) + .await + .unwrap() + ); + } + engine.stop().await.unwrap(); + } + + #[derive(Default)] + struct WriterState { + closed: usize, + aborted: usize, + } + + type WriterStates = Arc>>; + + enum WriterFailure { + None, + Finish, + Abort, + } + + fn writer_layer( + states: &WriterStates, + failure: impl Fn(&str) -> WriterFailure + Send + Sync + 'static, + ) -> mock::MockLayer { + let states = states.clone(); + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| -> oio::Writer { + states + .lock() + .unwrap() + .insert(path.to_string(), WriterState::default()); + Box::new(RecordingWriter { + inner, + path: path.to_string(), + states: states.clone(), + failure: failure(path), + }) + })) + .build() + .unwrap() + } + + struct RecordingWriter { + inner: oio::Writer, + path: String, + states: WriterStates, + failure: WriterFailure, + } + + impl mock::Write for RecordingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.inner.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + if matches!(self.failure, WriterFailure::Finish) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected index finish failure", + )); + } + let metadata = self.inner.close().await?; + self.states + .lock() + .unwrap() + .get_mut(&self.path) + .unwrap() + .closed += 1; + Ok(metadata) + } + + async fn abort(&mut self) -> mock::Result<()> { + self.states + .lock() + .unwrap() + .get_mut(&self.path) + .unwrap() + .aborted += 1; + if matches!(self.failure, WriterFailure::Abort) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected abort failure", + )); + } + self.inner.abort().await + } + } + + struct FailingSstReader { + inner: oio::Reader, + fail: Arc, + } + + impl mock::Read for FailingSstReader { + async fn read( + &self, + range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, mock::Buffer)> { + if self.fail.load(Ordering::Relaxed) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected SST read failure", + )); + } + self.inner.read(range).await + } + + async fn open( + &self, + range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, Box)> { + if self.fail.load(Ordering::Relaxed) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected SST read failure", + )); + } + self.inner.open(range).await + } + } + + #[rstest::rstest] + #[case::range(false, false, false)] + #[case::range_abort_failure(false, false, true)] + #[case::series_setup(true, true, false)] + #[case::series_read(true, false, false)] + #[case::series_abort_failure(true, false, true)] + #[tokio::test] + async fn test_read_failure_aborts_unfinished_outputs( + #[case] series: bool, + #[case] fail_before_open: bool, + #[case] fail_abort: bool, + ) { + let fail = Arc::new(AtomicBool::new(false)); + let read_fail = fail.clone(); + let read_layer = MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| -> oio::Reader { + if path.ends_with(".parquet") { + Box::new(FailingSstReader { + inner, + fail: read_fail.clone(), + }) + } else { + inner + } + })) + .build() + .unwrap(); + let mut env = TestEnv::with_prefix("series-builder-read-failure") + .await + .with_mock_layer(read_layer); + let (engine, region) = prepare_region(&mut env).await; + let version = region.version(); + let (mut bucket, mut entry) = build_input(&version); + // A single source is first polled after opening the series writer. + bucket.files.truncate(1); + entry.source_file_ids = vec![bucket.files[0].file_id().file_id()]; + let states = WriterStates::default(); + let write_fail = fail.clone(); + let write_layer = writer_layer(&states, move |_| { + write_fail.store(true, Ordering::Relaxed); + if fail_abort { + WriterFailure::Abort + } else { + WriterFailure::None + } + }); + let store = ObjectStore::new(Memory::default()) + .unwrap() + .layer(write_layer); + fail.store(fail_before_open, Ordering::Relaxed); + let result = if !series { + build_range_index(&store, ®ion, &version, bucket.files[0].clone()) + .await + .map(|_| ()) + } else { + let (purger, _receiver) = series_index_channel(store.clone()); + build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + .await + .map(|_| ()) + }; + fail.store(false, Ordering::Relaxed); + let error = result.unwrap_err(); + assert!( + format!("{error:?}").contains("injected SST read failure"), + "{error:?}" + ); + { + let states = states.lock().unwrap(); + assert_eq!(usize::from(!fail_before_open), states.len()); + for (path, state) in states.iter() { + assert_eq!(0, state.closed, "{path}"); + assert_eq!(1, state.aborted, "{path}"); + } + } + assert!(store.list("/").await.unwrap().is_empty()); + assert!(region.series_index_version().series_indexes.is_empty()); + engine.stop().await.unwrap(); + } +} diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index a59c1926fc..56c61da54f 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -14,6 +14,8 @@ //! Index catalog persistence, coverage metadata, and file paths. +use std::collections::BTreeMap; + use common_telemetry::warn; use common_time::Timestamp; use object_store::{ErrorKind, ObjectStore}; @@ -28,11 +30,29 @@ use crate::series_index::purger::IndexFilePurger; use crate::series_index::version::{ SeriesIndexFileHandle, SeriesIndexVersion, SeriesIndexVersionControl, }; +pub(crate) use crate::sst::range_index::range_index_path; + const SERIES_DIR: &str = "series"; const RANGE_CATALOG: &str = "range-index.json"; const SERIES_CATALOG: &str = "series-index.json"; const SERIES_METADATA_KEY: &str = "greptime.series_index"; +/// Summary of SSTs sharing a compaction-window-aligned start. +/// +/// New data changes must have sequences greater than those already indexed. A new +/// start adds a map entry; new data at an existing start raises its maximum sequence. +/// The maximum end and sequence may come from different files, so this summary +/// does not establish uniform sequence coverage throughout the interval. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct WindowSequence { + /// Inclusive start in epoch seconds, equal to the key in `window_sequences`. + pub(crate) start: i64, + /// Exclusive interval end in epoch seconds. + pub(crate) end: i64, + /// Maximum sequence among the SSTs sharing this aligned start. + pub(crate) max_sequence: u64, +} + /// Self-describing coverage stored in a series-index Parquet footer. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SeriesIndexEntry { @@ -41,9 +61,19 @@ pub(crate) struct SeriesIndexEntry { pub(crate) bucket_start: Timestamp, /// Exclusive bucket end. pub(crate) bucket_end: Timestamp, + /// Source SST IDs retained for debugging only. Compaction can replace these + /// files without changing indexed data, so IDs must not determine index reuse. pub(crate) source_file_ids: Vec, pub(crate) min_file_sequence: u64, pub(crate) max_file_sequence: u64, + /// Width used to align the half-open compaction windows, in seconds. + pub(crate) compaction_window_secs: i64, + /// Source SST summaries keyed by aligned start; intervals may overlap. + /// Each file contributes one summary regardless of its span. Equal starts merge + /// by taking the maximum end and sequence. See [`WindowSequence`] for the + /// sequence assumption used to detect new data. Compaction changing summary + /// boundaries may conservatively trigger a rebuild. + pub(crate) window_sequences: BTreeMap, } #[derive(Debug, Default, Serialize, Deserialize)] @@ -135,9 +165,9 @@ pub(crate) async fn load_version_control( .await .unwrap_or_default(); // TODO: Handle catalog entries whose index files are missing from storage. - let version = SeriesIndexVersion { - range_indexes: range.indexes.into_iter().collect(), - series_indexes: series + let version = SeriesIndexVersion::new( + range.indexes.into_iter().collect(), + series .indexes .into_iter() .map(|entry| { @@ -147,7 +177,7 @@ pub(crate) async fn load_version_control( ) }) .collect(), - }; + ); let control = SeriesIndexVersionControl::default(); control.publish(std::sync::Arc::new(version)); control @@ -155,6 +185,7 @@ pub(crate) async fn load_version_control( #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::Arc; use common_time::Timestamp; @@ -164,8 +195,9 @@ mod tests { use store_api::storage::{FileId, RegionId}; use crate::series_index::catalog::{ - SeriesIndexCatalog, SeriesIndexEntry, load_catalog, load_version_control, - series_catalog_path, series_metadata, store_catalog, + RangeIndexCatalog, SeriesIndexCatalog, SeriesIndexEntry, WindowSequence, load_catalog, + load_version_control, range_catalog_path, series_catalog_path, series_metadata, + store_catalog, }; use crate::series_index::purger::series_index_channel; @@ -194,31 +226,54 @@ mod tests { } #[tokio::test] - async fn test_load_catalog_returns_none_on_error() { + async fn test_load_catalog_defaults_on_missing_invalid_or_unreadable_catalog() { let store = ObjectStore::new(Memory::default()).unwrap(); - let path = series_catalog_path(RegionId::new(1, 1)); - // Missing catalog. + let region_id = RegionId::new(1, 1); + let (purger, _receiver) = series_index_channel(store.clone()); assert!( - load_catalog::(&store, &path) - .await - .is_none() - ); - store.write(&path, "invalid").await.unwrap(); - assert!( - load_catalog::(&store, &path) + load_catalog::(&store, &series_catalog_path(region_id)) .await .is_none() ); + let control = load_version_control(&store, region_id, &purger).await; + assert!(control.current().range_indexes.is_empty()); + assert!(control.current().series_indexes.is_empty()); + let file_id = FileId::random(); + store + .write( + &range_catalog_path(region_id), + serde_json::to_vec(&RangeIndexCatalog { + indexes: vec![file_id], + }) + .unwrap(), + ) + .await + .unwrap(); + store + .write(&series_catalog_path(region_id), "invalid") + .await + .unwrap(); + let control = load_version_control(&store, region_id, &purger).await; + assert!(control.current().range_indexes.contains(&file_id)); + assert!(control.current().series_indexes.is_empty()); let layer = MockLayerBuilder::default() .reader_factory(Arc::new(|_, _, _| Box::new(FailingCatalogReader))) .build() .unwrap(); - let store = store.layer(layer); assert!( - load_catalog::(&store, &path) + load_catalog::(&store, &series_catalog_path(region_id)) .await .is_none() ); + let store = store.layer(layer); + assert!( + load_catalog::(&store, &range_catalog_path(region_id)) + .await + .is_none() + ); + let control = load_version_control(&store, region_id, &purger).await; + assert!(control.current().range_indexes.is_empty()); + assert!(control.current().series_indexes.is_empty()); } #[tokio::test] @@ -232,6 +287,25 @@ mod tests { source_file_ids: vec![FileId::random()], min_file_sequence: 1, max_file_sequence: 2, + compaction_window_secs: 10, + window_sequences: BTreeMap::from([ + ( + 0, + WindowSequence { + start: 0, + end: 20, + max_sequence: 1, + }, + ), + ( + 20, + WindowSequence { + start: 20, + end: 100, + max_sequence: 2, + }, + ), + ]), }; store_catalog( &store, @@ -248,6 +322,13 @@ mod tests { .current(); assert!(current.range_indexes.is_empty()); assert_eq!(&entry, current.series_indexes[&entry.index_uuid].entry()); + assert_eq!(current.index_buckets.len(), 1); + let bucket = ¤t.index_buckets[&entry.bucket_start]; + assert_eq!(bucket.start, entry.bucket_start); + assert_eq!(bucket.end, entry.bucket_end); + assert_eq!(bucket.index_ids.as_slice(), &[entry.index_uuid]); + assert_eq!(bucket.compaction_window_secs, entry.compaction_window_secs); + assert_eq!(bucket.window_sequences, entry.window_sequences); let metadata = series_metadata(&entry).unwrap(); let decoded: SeriesIndexEntry = diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs new file mode 100644 index 0000000000..fa65dadcdd --- /dev/null +++ b/src/mito2/src/series_index/tests.rs @@ -0,0 +1,114 @@ +// 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. + +//! Shared fixtures for series-index construction tests. + +use std::sync::Arc; + +use api::v1::helper::row; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{RegionPutRequest, RegionRequest}; +use store_api::storage::RegionId; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::region::MitoRegionRef; +use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; +use crate::test_util::{CreateRequestBuilder, TestEnv, flush_region, rows_schema}; + +/// Builds real sparse SSTs; background maintenance is disabled so tests control publication. +pub(super) async fn prepare_region(env: &mut TestEnv) -> (MitoEngine, MitoRegionRef) { + prepare_region_with_timestamps(env, &[1000, 2000, 3000, 4000]).await +} + +async fn prepare_region_with_timestamps( + env: &mut TestEnv, + timestamps: &[i64], +) -> (MitoEngine, MitoRegionRef) { + let engine = env.create_engine(MitoConfig::default()).await; + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new().build(); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + request + .options + .insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); + request + .options + .insert("memtable.type".to_string(), "bulk".to_string()); + request + .options + .insert("sst_format".to_string(), "flat".to_string()); + request + .options + .insert("compaction.type".to_string(), "twcs".to_string()); + request.options.insert( + "compaction.twcs.time_window".to_string(), + "100s".to_string(), + ); + // Keep the source SSTs stable while tests reconcile multiple buckets. + request.options.insert( + "compaction.twcs.trigger_file_num".to_string(), + "100".to_string(), + ); + let full_schema = rows_schema(&request); + let mut pk_column = full_schema[0].clone(); + pk_column.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + pk_column.datatype = ColumnDataType::Binary.into(); + pk_column.semantic_type = SemanticType::Tag.into(); + let schema = vec![pk_column, full_schema[5].clone(), full_schema[4].clone()]; + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + for &ts in timestamps { + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows: Rows { + schema: schema.clone(), + rows: vec![row(vec![ + ValueData::BinaryValue(new_sparse_primary_key( + &["a", "x"], + &metadata, + 10, + 0, + )), + ValueData::TimestampMillisecondValue(ts), + ValueData::U64Value(1), + ])], + }, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + partition_expr_version: None, + }), + ) + .await + .unwrap(); + flush_region(&engine, region_id, None).await; + } + let region = engine.get_region(region_id).unwrap(); + (engine, region) +} diff --git a/src/mito2/src/series_index/version.rs b/src/mito2/src/series_index/version.rs index 7db23da2f9..4db883e762 100644 --- a/src/mito2/src/series_index/version.rs +++ b/src/mito2/src/series_index/version.rs @@ -14,13 +14,15 @@ //! Immutable index snapshots and aggregate series-file handles. -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::{self, Debug, Formatter}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; +use common_time::Timestamp; use store_api::storage::{FileId, RegionId}; +use crate::series_index::bucket::IndexBucket; use crate::series_index::catalog::SeriesIndexEntry; use crate::series_index::purger::{IndexFilePurger, PurgeRequest}; use crate::sst::file::RegionFileId; @@ -85,11 +87,30 @@ impl Drop for SeriesIndexFileHandleInner { /// Immutable series-index snapshot for one region. #[derive(Debug, Default)] pub(crate) struct SeriesIndexVersion { + /// Range indexes for visible SSTs; reconciliation removes IDs absent from its SST snapshot. + /// Physical deletion is independently handled by the SST file purger. pub(crate) range_indexes: HashSet, pub(crate) series_indexes: HashMap, + pub(crate) index_buckets: BTreeMap, } impl SeriesIndexVersion { + /// Restores bucket lookup from immutable index coverage stored in the catalog. + pub(crate) fn new( + range_indexes: HashSet, + series_indexes: HashMap, + ) -> Self { + let mut index_buckets = BTreeMap::new(); + for handle in series_indexes.values() { + IndexBucket::from_entry(handle.entry()).insert_into(&mut index_buckets); + } + Self { + range_indexes, + series_indexes, + index_buckets, + } + } + fn mark_all_deleted(&self) { self.series_indexes .values() diff --git a/src/mito2/src/sst/range_index.rs b/src/mito2/src/sst/range_index.rs index c4e576b6cd..6d9eadc553 100644 --- a/src/mito2/src/sst/range_index.rs +++ b/src/mito2/src/sst/range_index.rs @@ -28,8 +28,6 @@ pub use writer::{ }; pub use crate::sst::range_index::deleter::RangeIndexDeleter; -// Used by the upcoming query and index-building integration. -#[allow(unused_imports)] pub(crate) use crate::sst::range_index::deleter::range_index_path; const ROW_GROUP_ID_COLUMN: &str = "row_group_id";