From ce4fc4529e70cf59dda03594cb78073dde27cb5a Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 23 Oct 2026 15:11:29 +0800 Subject: [PATCH] feat(mito): read columnar series through prepared sources (#9892) * feat(mito): read columnar series through prepared sources Signed-off-by: evenyag * fix(mito): rebuild schema-incompatible columnar sidecars Signed-off-by: evenyag * docs(mito): regenerate series format configuration Signed-off-by: evenyag * fix(mito): batch prepared series reads and preflight inputs Signed-off-by: evenyag --------- Signed-off-by: evenyag --- config/config.md | 2 + config/datanode.example.toml | 3 + config/standalone.example.toml | 3 + src/mito2/AGENTS.md | 4 +- src/mito2/src/config.rs | 19 + src/mito2/src/engine.rs | 15 + src/mito2/src/read/series_candidate.rs | 1 + src/mito2/src/read/series_scan.rs | 113 +- src/mito2/src/read/series_scan/prepared.rs | 627 +++++++++ .../src/read/series_scan/prepared/tests.rs | 1121 +++++++++++++++++ src/mito2/src/read/series_scan/raw.rs | 16 + src/mito2/src/region/opener.rs | 6 +- src/mito2/src/series_index.rs | 14 +- src/mito2/src/series_index/bucket.rs | 68 +- src/mito2/src/series_index/builder.rs | 8 +- src/mito2/src/series_index/catalog.rs | 19 +- src/mito2/src/series_index/columnar.rs | 8 +- src/mito2/src/series_index/columnar/codec.rs | 4 + src/mito2/src/series_index/format.rs | 185 ++- src/mito2/src/series_index/maintenance.rs | 150 ++- src/mito2/src/series_index/samples.rs | 19 +- src/mito2/src/series_index/task.rs | 23 +- src/mito2/src/series_index/version.rs | 7 + src/store-api/src/region_engine.rs | 27 +- tests-integration/tests/http.rs | 1 + 25 files changed, 2362 insertions(+), 101 deletions(-) create mode 100644 src/mito2/src/read/series_scan/prepared.rs create mode 100644 src/mito2/src/read/series_scan/prepared/tests.rs diff --git a/config/config.md b/config/config.md index 6898973b214..930608aa4f5 100644 --- a/config/config.md +++ b/config/config.md @@ -182,6 +182,7 @@ | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_series_index_format` | String | `legacy` | Under development. Select legacy Parquet indexes or paired columnar series and samples. | | `region_engine.mito.experimental_series_index_max_size` | String | `5GiB` | Approximate series and range index size limit in open regions, shared across workers.
Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation
can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots
retained by readers are not counted.
Minimum: 1KiB. Takes effect on restart. | | `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | @@ -658,6 +659,7 @@ | `region_engine.mito.experimental_manifest_keep_removed_file_ttl` | String | `1h` | How long to keep removed files in the `removed_files` field of manifest
after they are removed from manifest.
files will only be removed from `removed_files` field
if both `keep_removed_file_count` and `keep_removed_file_ttl` is reached. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_series_index_format` | String | `legacy` | Under development. Select legacy Parquet indexes or paired columnar series and samples. | | `region_engine.mito.experimental_series_index_max_size` | String | `5GiB` | Approximate series and range index size limit in open regions, shared across workers.
Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation
can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots
retained by readers are not counted.
Minimum: 1KiB. Takes effect on restart. | | `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 40eaff54f3b..2af83991638 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -543,6 +543,9 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Under development. Select legacy Parquet indexes or paired columnar series and samples. +#+ experimental_series_index_format = "legacy" + ## Approximate series and range index size limit in open regions, shared across workers. ## Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation ## can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 9964653564d..4346b4daeb8 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -750,6 +750,9 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Under development. Select legacy Parquet indexes or paired columnar series and samples. +#+ experimental_series_index_format = "legacy" + ## Approximate series and range index size limit in open regions, shared across workers. ## Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation ## can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 421b0df6723..e6058a0c1a2 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -66,14 +66,14 @@ filtered `RecordBatch` stream. and defer new builds when the shared estimate is full. Cleanup still prunes obsolete range entries and expires series indexes to reclaim published usage. Concurrent builds may overshoot; closed-region files are not counted. Series handles and the SST purger own file deletion. -- **Series format PoC**: samples are rebuildable local SST sidecars. Optional writer interfaces (`format.rs`) integrate with reconciliation, catalog publication, and deferred purge. Defaults retain the original Parquet writer and disable samples. The opt-in columnar backend (`series_index/columnar.rs`) implements streaming OpenDAL writes and pinned format reads; its wire contract is in `docs/formats/mito-series.md`. Artifact objects record paths and sizes, including the manifest, so reads need no stat. Codecs currently run inline pending scheduling benchmarks. Prepared sources and query-consumer integration are deferred; ordinary scans only use Parquet index artifacts. +- **Series format PoC**: samples are rebuildable local SST sidecars. Optional writer interfaces (`format.rs`) integrate with reconciliation, catalog publication, and deferred purge. Defaults retain the original Parquet writer and disable samples. The opt-in columnar backend (`series_index/columnar.rs`) implements streaming OpenDAL writes and pinned format reads; its wire contract is in `docs/formats/mito-series.md`. Artifact objects record paths and sizes, including the manifest, so reads need no stat. Codecs currently run inline pending scheduling benchmarks. Columnar `SeriesScan` raw reads use scanner-wide preparation in `read/series_scan/prepared.rs`; ordinary record-batch scans retain their SST path. The paired `experimental_series_index_format` setting selects legacy or columnar artifacts. - **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. - **Columnar samples**: `series_index/samples.rs` implements opt-in per-SST sidecars, including nullable Float64 fields and preserved sequence/delete rows. Its wire contract is `docs/formats/mito-samples.md`; time/ID bounds may overlap across leaves and parts. Both columnar backends share cancellation-safe output ownership - in `series_index/artifact_writer.rs`. Sample reads do not apply query visibility. + in `series_index/artifact_writer.rs`. Format readers do not apply query visibility; the prepared-source adapter applies visibility and the existing merge/dedup pipeline. Manual reconciliation covers all snapshot SSTs regardless of maintenance eligibility, but retains the soft disk limit. Manual builds validate columnar sample and tag schemas before reuse. - **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/config.rs b/src/mito2/src/config.rs index c0fa74281ca..1fce110f7ef 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -54,6 +54,17 @@ const INDEX_CREATE_MEM_THRESHOLD_FACTOR: u64 = 16; /// Fetch option timeout pub(crate) const FETCH_OPTION_TIMEOUT: Duration = Duration::from_secs(3); +/// Paired series-index and sample storage selection for the experimental reader. +#[derive(Debug, Default, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum SeriesIndexFormat { + /// Parquet series indexes and authoritative SST reads. + #[default] + Legacy, + /// Columnar series indexes and per-SST columnar samples. + Columnar, +} + /// Configuration for [MitoEngine](crate::engine::MitoEngine). /// Before using the config, make sure to call `MitoConfig::validate()` to check if the config is valid. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] @@ -92,6 +103,8 @@ pub struct MitoConfig { /// Under development; do not enable. Whether to enable series indexes (default false). /// Indexes are stored on the local filesystem under `{data_home}/series_index`. pub experimental_enable_series_index: bool, + /// Paired artifact format (default legacy). + pub experimental_series_index_format: SeriesIndexFormat, /// Approximate series and range index size limit in open regions (default: 5 GiB). /// Workers share a periodically refreshed estimate and skip maintenance when full. /// In-flight reconciliation can exceed the limit; closed-region files are not counted. @@ -222,6 +235,7 @@ impl Default for MitoConfig { compress_manifest: false, max_background_index_builds: divide_num_cpus(8), experimental_enable_series_index: false, + experimental_series_index_format: SeriesIndexFormat::Legacy, experimental_series_index_max_size: ReadableSize::gb(5), experimental_enable_range_index: false, experimental_series_index_maintenance_interval: @@ -450,6 +464,7 @@ mod tests { ); let mut config: MitoConfig = toml::from_str( "experimental_enable_series_index = true + experimental_series_index_format = 'columnar' experimental_enable_range_index = false experimental_series_index_max_size = '64MiB' experimental_series_index_maintenance_interval = '30s' @@ -458,6 +473,10 @@ mod tests { .unwrap(); config.sanitize("/data").unwrap(); assert!(config.experimental_enable_series_index); + assert_eq!( + config.experimental_series_index_format, + SeriesIndexFormat::Columnar + ); assert_eq!( ReadableSize::mb(64), config.experimental_series_index_max_size diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 47752adbe75..8d930be4df0 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -1062,6 +1062,21 @@ impl EngineInner { // in the query's older SST snapshot. let series_index = region.series_index_store.as_ref().map(|store| { crate::series_index::SeriesIndexReadContext { + format: if region + .series_index_version_control + .formats + .read() + .unwrap() + .index + .as_ref() + .is_some_and(|format| { + format.format() + == crate::series_index::artifact::ArtifactFormat::ColumnarSeries + }) { + crate::config::SeriesIndexFormat::Columnar + } else { + crate::config::SeriesIndexFormat::Legacy + }, store: store.clone(), version: region.series_index_version(), } diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 82210e0ffa7..54b9b8bb5e5 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -834,6 +834,7 @@ mod tests { let (purger, _receiver) = series_index_channel(store.clone()); let handle = SeriesIndexFileHandle::new(metadata.region_id, entry.clone(), purger); let context = SeriesIndexReadContext { + format: crate::config::SeriesIndexFormat::Legacy, store, version: Arc::new(SeriesIndexVersion { series_indexes: [(entry.index_uuid, handle)].into(), diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 971ff93a742..433fac8c7f1 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -14,6 +14,7 @@ //! Per-series scan implementation. +mod prepared; mod raw; use std::collections::HashSet; @@ -187,6 +188,7 @@ impl SeriesScanMode { pub struct SeriesScan { /// Implementation used by this scan. mode: SeriesScanMode, + prepared: Arc, /// Properties of the scanner. properties: ScannerProperties, /// Context of streams. @@ -236,6 +238,7 @@ impl SeriesScan { Self { mode, + prepared: Arc::new(prepared::Preparation::default()), properties, stream_ctx, pruner, @@ -245,6 +248,14 @@ impl SeriesScan { } } + fn uses_columnar(&self) -> bool { + self.stream_ctx + .input + .series_index + .as_ref() + .is_some_and(|context| context.format == crate::config::SeriesIndexFormat::Columnar) + } + fn prepare_raw_scan( &self, label_source: &RawLabelSource, @@ -311,6 +322,9 @@ impl SeriesScan { return Ok(None); } } + if self.uses_columnar() && !is_sparse_metric_metadata(metadata) { + return Ok(None); + } let selected = input.mapper.output_schema(); let mut fields = Vec::new(); let mut specs = Vec::new(); @@ -342,7 +356,9 @@ impl SeriesScan { .build(), )); }; - if !data_type.is_numeric() && !is_native_histogram_value_type(data_type) { + if (self.uses_columnar() && *data_type != ConcreteDataType::float64_datatype()) + || (!data_type.is_numeric() && !is_native_histogram_value_type(data_type)) + { debug!( "Raw series scan preflight rejected field type: region_id={}, column_id={}", metadata.region_id, column.column_id, @@ -370,13 +386,28 @@ impl SeriesScan { ); return Ok(None); } + if self.uses_columnar() { + let ranges = self + .properties + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new)?; + } let schema = Arc::new(SeriesSchema { fields: specs.into(), eval_grid: None, }); let typed_schema = input.mapper.input_arrow_schema(false); let mut predicates = Vec::new(); - if let Some(predicate) = input.predicate_group().predicate_without_region() { + let predicate = if self.uses_columnar() { + input.predicate_group().predicate() + } else { + input.predicate_group().predicate_without_region() + }; + if let Some(predicate) = predicate { if !predicate.dyn_filters().is_empty() { debug!( "Raw series scan preflight rejected dynamic predicate filters: region_id={}", @@ -813,6 +844,9 @@ impl RegionScanner for SeriesScan { } fn raw_series_ordering(&self) -> SeriesOrdering { + if self.uses_columnar() { + return SeriesOrdering::KeyThenTime; + } match self.mode { SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime, SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous, @@ -838,6 +872,35 @@ impl RegionScanner for SeriesScan { let Some(preparation) = self.prepare_raw_scan(label_source)? else { return Ok(None); }; + if self.uses_columnar() { + let partitions = self.properties.num_partitions().max(1); + if partition >= partitions { + return Err(BoxedError::new( + PartitionOutOfRangeSnafu { + given: partition, + all: partitions, + } + .build(), + )); + } + let future = self.prepare_series(ctx, metrics_set, label_source); + let schema = preparation.schema; + let labels = label_source.clone(); + let stream = async_stream::try_stream! { + let source = future.await?.ok_or_else(|| BoxedError::new(InvalidRecordBatchSnafu { reason: "columnar prepared capability disappeared" }.build()))?; + let n = source.num_series() as u128; + let start = (n * partition as u128 / partitions as u128) as usize; + let end = (n * (partition + 1) as u128 / partitions as u128) as usize; + let mut reader = source.read_series(start..end)?; + while let Some(chunk) = reader.try_next().await? { yield chunk; } + }; + return Ok(Some(raw::from_stream( + schema, + labels, + SeriesOrdering::KeyThenTime, + Box::pin(stream), + ))); + } let metrics = new_partition_metrics( &self.stream_ctx, ctx.explain_verbose, @@ -860,7 +923,52 @@ impl RegionScanner for SeriesScan { ))) } + fn prepare_series( + &self, + ctx: &QueryScanContext, + metrics_set: &ExecutionPlanMetricsSet, + label_source: &RawLabelSource, + ) -> futures::future::BoxFuture< + 'static, + Result>, BoxedError>, + > { + if !self.uses_columnar() { + return Box::pin(async { Ok(None) }); + } + let preparation = self + .check_scan_limit() + .map_err(BoxedError::new) + .and_then(|_| self.prepare_raw_scan(label_source)); + let state = self.prepared.clone(); + let stream_ctx = self.stream_ctx.clone(); + let ranges = self + .properties + .partitions + .iter() + .flatten() + .copied() + .collect(); + let labels = label_source.clone(); + let metrics = new_partition_metrics( + &self.stream_ctx, + ctx.explain_verbose, + metrics_set, + 0, + &self.metrics_list, + ); + Box::pin(async move { + let Some(preparation) = preparation? else { + return Ok(None); + }; + state + .prepare(stream_ctx, ranges, labels, preparation, metrics) + .await + .map(Some) + }) + } + fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> { + self.prepared = Arc::new(prepared::Preparation::default()); self.properties.prepare(request); self.check_scan_limit().map_err(BoxedError::new)?; @@ -885,6 +993,7 @@ impl RegionScanner for SeriesScan { } fn reset_state(&mut self) { + self.prepared = Arc::new(prepared::Preparation::default()); self.stream_ctx.input.predicate.clear_dyn_filters(); let num_workers = common_stat::get_total_cpu_cores().max(1); self.pruner = match self.mode { diff --git a/src/mito2/src/read/series_scan/prepared.rs b/src/mito2/src/read/series_scan/prepared.rs new file mode 100644 index 00000000000..8dc002f9b23 --- /dev/null +++ b/src/mito2/src/read/series_scan/prepared.rs @@ -0,0 +1,627 @@ +// 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. + +//! Scanner-wide columnar candidates and independent, lazy sample reads. + +#[cfg(test)] +mod tests; + +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::ops::Range; +use std::sync::Arc; + +use api::v1::SemanticType; +use common_error::ext::BoxedError; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datatypes::arrow::array::{ + Array, ArrayRef, BinaryArray, BooleanArray, DictionaryArray, StringArray, UInt32Array, + UInt64Array, +}; +use datatypes::arrow::buffer::{Buffer, OffsetBuffer}; +use datatypes::arrow::compute::{cast, filter_record_batch}; +use datatypes::arrow::datatypes::UInt32Type; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::series::SeriesOrdering; +use datatypes::value::ValueRef; +use futures::TryStreamExt; +use mito_codec::row_converter::SparsePrimaryKeyCodec; +use snafu::{IntoError, ResultExt}; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, +}; +use store_api::region_engine::{ + PartitionRange, PreparedSeriesSource, RawLabelSource, SendableRawSeriesStream, SeriesInfo, +}; +use store_api::storage::FileId; +use store_api::storage::consts::{ReservedColumnId, SEQUENCE_COLUMN_NAME}; +use table::predicate::Predicate; +use tokio::sync::OnceCell; + +use crate::error::{ + ComputeArrowSnafu, Error, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, +}; +use crate::read::BoxedRecordBatchStream; +use crate::read::scan_region::StreamContext; +use crate::read::scan_util::PartitionMetrics; +use crate::read::seq_scan::SeqScan; +use crate::read::series_scan::{RawScanPreparation, raw}; +use crate::read::stream::ScanBatch; +use crate::series_index::MetricSeriesId; +use crate::series_index::artifact::ArtifactFormat; +use crate::series_index::columnar::ColumnarSeriesFormat; +use crate::series_index::format::{ + self, SamplesFormat, SamplesFormatReader, SamplesReadRequest, SeriesIndexFormat, + SeriesReadRequest, +}; +use crate::series_index::samples::ColumnarSamplesFormat; + +fn invalid(reason: impl Into) -> Error { + InvalidRecordBatchSnafu { + reason: reason.into(), + } + .build() +} + +/// Errors are shared too: a failed execution must not restart discovery in another partition. +#[derive(Default)] +pub(crate) struct Preparation { + source: OnceCell, Arc>>, +} + +impl Preparation { + pub(crate) async fn prepare( + &self, + context: Arc, + ranges: Vec, + labels: RawLabelSource, + raw: RawScanPreparation, + metrics: PartitionMetrics, + ) -> std::result::Result, BoxedError> { + let expected = labels.clone(); + let result = self + .source + .get_or_init(|| async { + ColumnarPreparedSource::prepare(context, ranges, labels, raw, metrics) + .await + .map(Arc::new) + .map_err(Arc::new) + }) + .await; + let source = result.as_ref().map_err(|error| { + BoxedError::new(crate::error::ScanSeriesSnafu.into_error(error.clone())) + })?; + if !same_labels(&source.labels, &expected) { + return Err(BoxedError::new(invalid( + "label binding changed after series preparation", + ))); + } + Ok(source.clone()) + } +} + +fn same_labels(left: &RawLabelSource, right: &RawLabelSource) -> bool { + match (left, right) { + ( + RawLabelSource::PrimaryKey { + metric_name: a, + physical_metadata: am, + logical_tags: at, + }, + RawLabelSource::PrimaryKey { + metric_name: b, + physical_metadata: bm, + logical_tags: bt, + }, + ) => { + a == b + && Arc::ptr_eq(am, bm) + && at.len() == bt.len() + && at + .iter() + .zip(bt.iter()) + .all(|(a, b)| a.column_id == b.column_id && a.name == b.name) + } + _ => false, + } +} + +/// Validates the pinned assignments and coverage without opening artifacts. +pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result<()> { + collect_inputs(context, ranges).map(|_| ()) +} + +fn collect_inputs( + context: &StreamContext, + ranges: &[PartitionRange], +) -> Result<(Vec, HashSet)> { + let input = &context.input; + let metadata = input.region_metadata(); + let index = input + .series_index + .as_ref() + .ok_or_else(|| invalid("missing columnar index context"))?; + let mut files = BTreeMap::>::new(); + for range in ranges { + let range = context + .ranges + .get(range.identifier) + .ok_or_else(|| invalid("invalid prepared scan range"))?; + for row_group in &range.row_group_indices { + if context.is_mem_range_index(*row_group) { + return Err(invalid("columnar prepared reads require flushed memtables")); + } + if !context.is_file_range_index(*row_group) { + return Err(invalid( + "columnar prepared reads do not support extension ranges", + )); + } + files + .entry(row_group.index - input.num_memtables()) + .or_default() + .insert(row_group.row_group_index); + } + } + let mut indexes = HashSet::new(); + let file_indices = files.keys().copied().collect(); + for (file_idx, groups) in files { + let file = &input.files[file_idx]; + if !groups.contains(&-1) + && (groups.len() as u64 != file.meta_ref().num_row_groups + || groups + .iter() + .any(|group| *group < 0 || *group as u64 >= file.meta_ref().num_row_groups)) + { + return Err(invalid( + "columnar samples require complete SST range assignments", + )); + } + if file.meta_ref().region_id != metadata.region_id { + return Err(invalid( + "columnar prepared reads do not support foreign SSTs", + )); + } + let mut covered = false; + for handle in index.version.series_indexes.values() { + if handle.entry().format() == ArtifactFormat::ColumnarSeries + && handle + .entry() + .covers_file(file.meta_ref(), metadata.region_id) + { + indexes.insert(handle.entry().index_uuid); + covered = true; + } + } + if !covered { + return Err(invalid(format!( + "SST {} has no columnar series coverage; run build_series_index()", + file.file_id() + ))); + } + index + .version + .samples + .get(&file.file_id().file_id()) + .filter(|handle| handle.entry().artifact.format == ArtifactFormat::ColumnarSamples) + .ok_or_else(|| { + invalid(format!( + "SST {} has no columnar samples; run build_series_index()", + file.file_id() + )) + })?; + } + Ok((file_indices, indexes)) +} + +struct Candidate { + id: MetricSeriesId, + key: Buffer, + tags: HashMap>, +} + +struct ColumnarPreparedSource { + context: Arc, + labels: RawLabelSource, + raw: Arc, + candidates: Arc>, + samples: Arc>>, + metrics: PartitionMetrics, + // Readers retain this reservation with the candidate set. + reservation: Arc, +} + +impl ColumnarPreparedSource { + async fn prepare( + context: Arc, + ranges: Vec, + labels: RawLabelSource, + raw: RawScanPreparation, + metrics: PartitionMetrics, + ) -> Result { + let input = &context.input; + let metadata = input.region_metadata(); + let index = input + .series_index + .as_ref() + .ok_or_else(|| invalid("missing columnar index context"))?; + let (files, indexes) = collect_inputs(&context, &ranges)?; + let mut samples = Vec::>::new(); + for file_idx in files { + let file = &input.files[file_idx]; + let handle = &index.version.samples[&file.file_id().file_id()]; + let reader = ColumnarSamplesFormat::default() + .reader(index.store.clone(), handle.clone()) + .await?; + format::validate_sample_schema(&reader.schema(), metadata)?; + samples.push(Arc::from(reader)); + } + let reservation = + MemoryConsumer::new("ColumnarPreparedSource").register(&input.scan_memory_pool); + let mut candidates = BTreeMap::::new(); + for id in indexes { + let reader = ColumnarSeriesFormat::default() + .reader( + index.store.clone(), + index.version.series_indexes[&id].clone(), + ) + .await?; + let schema = reader.schema(); + format::validate_series_schema(&schema, metadata)?; + let mut tags = schema + .fields() + .iter() + .skip(3) + .enumerate() + .map(|(i, field)| { + let id = format::field_id(field)?; + Ok((i + 3, id, field.name().clone())) + }) + .collect::>>()?; + tags.sort_unstable_by(|a, b| a.2.cmp(&b.2)); + let mut filters = Vec::new(); + if let Some(predicate) = input.predicate_group().predicate() { + for expr in predicate.exprs() { + if expr + .column_refs() + .iter() + .all(|column| schema.index_of(&column.name).is_ok()) + { + filters.push( + Predicate::to_physical_expr(expr, &schema) + .map_err(|e| invalid(e.to_string()))?, + ); + } + } + } + let mut batches = reader.read(SeriesReadRequest { + projection: schema.fields().iter().map(|f| f.name().clone()).collect(), + ..Default::default() + })?; + while let Some(mut batch) = batches.try_next().await? { + for filter in &filters { + batch = common_recordbatch::filter::batch_filter(&batch, filter) + .map_err(|e| invalid(e.to_string()))?; + } + let table = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("invalid table ID column"))?; + let tsid = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("invalid TSID column"))?; + for row in 0..batch.num_rows() { + let id = MetricSeriesId { + table_id: table.value(row), + tsid: tsid.value(row), + }; + let mut values = HashMap::new(); + for (i, column_id, _) in &tags { + let strings = batch + .column(*i) + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("invalid series tag array"))?; + values.insert( + *column_id, + (!strings.is_null(row)).then(|| strings.value(row).to_owned()), + ); + } + // TODO(yingwen): Separate raw-series identity from label representation. + // The PoC formats assume (table_id, tsid) uniqueness; encode once at + // this adapter boundary until RawSeriesChunk no longer needs sparse keys. + let mut key = Vec::new(); + let codec = SparsePrimaryKeyCodec::schemaless(); + codec + .encode_internal(id.table_id, id.tsid, &mut key) + .map_err(|e| invalid(e.to_string()))?; + codec + .encode_to_vec( + tags.iter().map(|(_, id, _)| { + ( + *id, + values[id] + .as_deref() + .map_or(ValueRef::Null, ValueRef::String), + ) + }), + &mut key, + ) + .map_err(|e| invalid(e.to_string()))?; + if let Some(previous) = candidates.get(&id) { + if previous.key.as_slice() != key.as_slice() { + return Err(invalid( + "conflicting labels for the same table ID and TSID", + )); + } + continue; + } + let bytes = key.len() + + values.values().flatten().map(String::len).sum::() + + values.capacity() * (std::mem::size_of::<(u32, Option)>() + 16) + + std::mem::size_of::() + + 64; + reservation + .try_grow(bytes) + .map_err(|e| invalid(e.to_string()))?; + candidates.insert( + id, + Candidate { + id, + key: Buffer::from(key), + tags: values, + }, + ); + } + } + } + let mut candidates: Vec<_> = candidates.into_values().collect(); + candidates.sort_unstable_by(|a, b| a.key.as_slice().cmp(b.key.as_slice())); + Ok(Self { + context, + labels, + raw: Arc::new(raw), + candidates: Arc::new(candidates), + samples: Arc::new(samples), + metrics, + reservation: Arc::new(reservation), + }) + } +} + +impl PreparedSeriesSource for ColumnarPreparedSource { + fn num_series(&self) -> usize { + self.candidates.len() + } + fn label_source(&self) -> &RawLabelSource { + &self.labels + } + fn series_info(&self, ordinal: usize) -> std::result::Result { + let candidate = self + .candidates + .get(ordinal) + .ok_or_else(|| BoxedError::new(invalid("series ordinal out of bounds")))?; + Ok(SeriesInfo { + primary_key: candidate.key.clone(), + metric: None, + }) + } + fn read_series( + &self, + range: Range, + ) -> std::result::Result { + if range.start > range.end || range.end > self.candidates.len() { + return Err(BoxedError::new(invalid("series range out of bounds"))); + } + let context = self.context.clone(); + let candidates = self.candidates.clone(); + let samples = self.samples.clone(); + let metrics = self.metrics.clone(); + let reservation = self.reservation.clone(); + let stream = async_stream::try_stream! { + let _reservation = reservation; + if range.is_empty() { return; } + let unit = context.input.region_metadata().time_index_type().unit(); + let bounds = context.input.time_range.as_ref(); + let start = bounds.and_then(|r| *r.start()).and_then(|t| t.convert_to_ceil(unit)).map(|t| t.value()); + let end = bounds.and_then(|r| *r.end()).and_then(|t| t.convert_to_ceil(unit)).map(|t| t.value()); + let mut sources = Vec::::new(); + for reader in samples.iter() { + let candidates = candidates.clone(); + let range = range.clone(); + let context = context.clone(); + let reader = reader.clone(); + sources.push(Box::pin(async_stream::try_stream! { + let reservation = MemoryConsumer::new("ColumnarSamplesRead").register(&context.input.scan_memory_pool); + // Account for the ID selection retained by the samples reader, including + // B-tree node overhead, before allocating it. + let selection_bytes = range.len() * (std::mem::size_of::() + 64); + reservation.try_grow(selection_bytes).context(crate::error::MergeCandidateSeriesSnafu)?; + let ids = candidates[range.clone()].iter().map(|c| c.id).collect::>(); + let mut batches = reader.read(SamplesReadRequest { + projection: reader.schema().fields().iter().map(|f| f.name().clone()).collect(), + ids: Some(ids), + start, + end, + ..Default::default() + })?; + while let Some(batch) = batches.try_next().await? { + let tables = batch.column_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + .and_then(|a| a.as_any().downcast_ref::()) + .ok_or_else(|| invalid("missing or invalid sample table ID column"))?; + let tsids = batch.column_by_name(DATA_SCHEMA_TSID_COLUMN_NAME) + .and_then(|a| a.as_any().downcast_ref::()) + .ok_or_else(|| invalid("missing or invalid sample TSID column"))?; + let mut offset = 0; + while offset < batch.num_rows() { + let id = MetricSeriesId { table_id: tables.value(offset), tsid: tsids.value(offset) }; + // Sparse keys begin with the order-preserving table ID/TSID prefix. + let ordinal = candidates[range.clone()].binary_search_by_key(&id, |c| c.id) + .map_err(|_| invalid("sample contains an unselected series"))? + range.start; + let mut end = offset + 1; + while end < batch.num_rows() && tables.value(end) == id.table_id && tsids.value(end) == id.tsid { + end += 1; + } + let flat = sample_to_flat(batch.slice(offset, end - offset), &candidates[ordinal], &context)?; + let flat = filter_visibility(flat, &context)?; + offset = end; + if flat.num_rows() == 0 { continue; } + let bytes = flat.get_array_memory_size(); + reservation.try_grow(bytes).context(crate::error::MergeCandidateSeriesSnafu)?; + yield flat; + reservation.shrink(bytes); + } + } + })); + } + let mut merged = SeqScan::build_flat_reader_from_sources(&context, sources, None, Some(&metrics), false, 1).await?; + while let Some(batch) = merged.try_next().await? { + yield ScanBatch::RecordBatch(batch); + } + }; + Ok(raw::make_stream( + Box::pin(stream), + self.raw.schema.clone(), + self.labels.clone(), + self.raw.fields.clone(), + self.metrics.clone(), + self.raw.predicates.clone(), + SeriesOrdering::KeyThenTime, + )) + } +} + +fn sample_to_flat( + batch: RecordBatch, + candidate: &Candidate, + context: &StreamContext, +) -> Result { + let schema = context.input.mapper.input_arrow_schema(false); + let metadata = context.input.region_metadata(); + let mut arrays = Vec::::new(); + for field in schema.fields() { + let array: ArrayRef = match field.name().as_str() { + "__primary_key" => Arc::new( + DictionaryArray::::try_new( + UInt32Array::from(vec![0; batch.num_rows()]), + Arc::new(BinaryArray::new( + OffsetBuffer::new( + vec![ + 0, + i32::try_from(candidate.key.len()).map_err(|_| { + invalid("series key exceeds Arrow offset limit") + })?, + ] + .into(), + ), + candidate.key.clone(), + None, + )), + ) + .context(ComputeArrowSnafu)?, + ), + "__sequence" | "__op_type" => batch + .column_by_name(field.name()) + .ok_or_else(|| invalid("missing sample internal column"))? + .clone(), + name => { + let column = metadata + .column_by_name(name) + .ok_or_else(|| invalid("unknown flat column"))?; + if column.semantic_type == SemanticType::Tag { + if column.column_id == ReservedColumnId::table_id() { + Arc::new(UInt32Array::from(vec![ + candidate.id.table_id; + batch.num_rows() + ])) + } else if column.column_id == ReservedColumnId::tsid() { + Arc::new(UInt64Array::from(vec![candidate.id.tsid; batch.num_rows()])) + } else { + let value = candidate + .tags + .get(&column.column_id) + .ok_or_else(|| invalid("missing series tag"))?; + Arc::new(StringArray::from(vec![value.as_deref(); batch.num_rows()])) + } + } else { + batch + .column_by_name(name) + .ok_or_else(|| invalid("missing sample field"))? + .clone() + } + } + }; + arrays.push(if array.data_type() == field.data_type() { + array + } else { + cast(&array, field.data_type()).context(ComputeArrowSnafu)? + }); + } + RecordBatch::try_new(schema, arrays).context(NewRecordBatchSnafu) +} + +fn filter_visibility(batch: RecordBatch, context: &StreamContext) -> Result { + let sequence_column = batch + .column_by_name(SEQUENCE_COLUMN_NAME) + .ok_or_else(|| invalid("missing sequence column __sequence"))?; + let sequence = sequence_column + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("sequence column __sequence must be UInt64"))?; + let metadata = context.input.region_metadata(); + let timestamp = &metadata.time_index_column().column_schema; + let time_column = batch + .column_by_name(×tamp.name) + .ok_or_else(|| invalid(format!("missing time index column {}", timestamp.name)))?; + let unit = metadata.time_index_type().unit(); + let (time, actual_unit) = datatypes::timestamp::timestamp_array_to_primitive(time_column) + .ok_or_else(|| { + invalid(format!( + "time index column {} must be a timestamp", + timestamp.name + )) + })?; + if common_time::timestamp::TimeUnit::from(actual_unit) != unit { + return Err(invalid(format!( + "time index column {} has unit {actual_unit:?}, expected {unit:?}", + timestamp.name + ))); + } + let sequence_filter = context + .input + .sequence_range + .map(|range| range.filter(sequence_column)) + .transpose() + .context(ComputeArrowSnafu)?; + let filter = BooleanArray::from( + (0..batch.num_rows()) + .map(|i| { + !sequence.is_null(i) + && !time.is_null(i) + && context + .input + .snapshot_sequence + .is_none_or(|max| sequence.value(i) <= max) + && sequence_filter + .as_ref() + .is_none_or(|filter| filter.value(i)) + && context.input.time_range.as_ref().is_none_or(|range| { + range.contains(&common_time::Timestamp::new(time.value(i), unit)) + }) + }) + .collect::>(), + ); + filter_record_batch(&batch, &filter).context(ComputeArrowSnafu) +} diff --git a/src/mito2/src/read/series_scan/prepared/tests.rs b/src/mito2/src/read/series_scan/prepared/tests.rs new file mode 100644 index 00000000000..a0fbee4c14c --- /dev/null +++ b/src/mito2/src/read/series_scan/prepared/tests.rs @@ -0,0 +1,1121 @@ +// 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::sync::Mutex; +use std::time::Duration; + +use api::v1::helper::row; +use api::v1::region::build_index_request; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, ColumnSchema, Rows}; +use common_error::ext::ErrorExt; +use common_query::series_frame::unpack_raw_series; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ExecutionPlan, metrics::ExecutionPlanMetricsSet}; +use datafusion_expr::{col, lit}; +use datatypes::prelude::ConcreteDataType; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::SampleRef; +use object_store::ObjectStore; +use object_store::layers::mock::{self, MockLayerBuilder, oio}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::region_engine::{ + BoundTag, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, +}; +use store_api::region_request::{ + RegionBuildIndexRequest, RegionDeleteRequest, RegionPutRequest, RegionRequest, +}; +use store_api::storage::{ScanRequest, TimeSeriesDistribution}; +use table::table::scan::RegionScanExec; + +use super::*; +use crate::config::{MitoConfig, SeriesIndexFormat as SelectedFormat}; +use crate::engine::MitoEngine; +use crate::read::flat_projection::FlatProjectionMapper; +use crate::read::scan_region::{PredicateGroup, ScanInput}; +use crate::read::series_scan::SeriesScan; +use crate::region::MitoRegionRef; +use crate::series_index::SeriesIndexReadContext; +use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; +use crate::test_util::{TestEnv, flush_region}; + +async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) { + let mut env = TestEnv::with_prefix("prepared-columnar").await; + let mut metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + metadata + .column_metadatas + .iter_mut() + .find(|c| c.column_id == 2) + .unwrap() + .column_schema + .data_type = ConcreteDataType::float64_datatype(); + let (engine, region) = crate::series_index::prepare_region_with_metadata( + &mut env, + &[1000, 1000, 2000], + MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: SelectedFormat::Columnar, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + Arc::new(metadata), + ) + .await; + (env, engine, region) +} + +fn labels(metadata: store_api::metadata::RegionMetadataRef) -> RawLabelSource { + RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture"), + physical_metadata: metadata, + logical_tags: vec![ + BoundTag { + column_id: 0, + name: Arc::from("tag_0"), + }, + BoundTag { + column_id: 1, + name: Arc::from("tag_1"), + }, + ] + .into(), + } +} + +async fn build(engine: &MitoEngine, region: &MitoRegionRef) { + engine + .handle_request( + region.region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(build_index_request::Options::SeriesIndex(Default::default())), + }), + ) + .await + .unwrap(); +} + +async fn write( + engine: &MitoEngine, + region: &MitoRegionRef, + tsid: u64, + timestamp: i64, + value: Option, +) { + let metadata = region.version().metadata.clone(); + let mut schema = vec![ + ColumnSchema { + column_name: "__primary_key".into(), + datatype: ColumnDataType::Binary as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "ts".into(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ]; + let mut values = vec![ + ValueData::BinaryValue(new_sparse_primary_key(&["a", "x"], &metadata, 10, tsid)), + ValueData::TimestampMillisecondValue(timestamp), + ]; + let request = if let Some(value) = value { + schema.push(ColumnSchema { + column_name: "field_0".into(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }); + values.push(ValueData::F64Value(value)); + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema, + rows: vec![row(values)], + }, + skip_wal: false, + hint: Some(api::v1::WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse as i32, + }), + partition_expr_version: None, + }) + } else { + RegionRequest::Delete(RegionDeleteRequest { + rows: Rows { + schema, + rows: vec![row(values)], + }, + hint: Some(api::v1::WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse as i32, + }), + partition_expr_version: None, + }) + }; + engine + .handle_request(region.region_id, request) + .await + .unwrap(); +} + +struct ReadRecorder { + inner: oio::Reader, + path: String, + reads: Arc>>, +} +impl mock::Read for ReadRecorder { + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + self.reads.lock().unwrap().push((self.path.clone(), range)); + self.inner.read(range).await + } + + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> { + self.reads.lock().unwrap().push((self.path.clone(), range)); + self.inner.open(range).await + } +} + +fn scanner(region: &MitoRegionRef, store: ObjectStore) -> SeriesScan { + scanner_with_format(region, store, SelectedFormat::Columnar) +} + +fn scanner_with_format( + region: &MitoRegionRef, + store: ObjectStore, + format: SelectedFormat, +) -> SeriesScan { + scanner_with_projection( + region, + store, + format, + (0..region.version().metadata.column_metadatas.len()).collect(), + ) +} + +fn scanner_with_projection( + region: &MitoRegionRef, + store: ObjectStore, + format: SelectedFormat, + projection: Vec, +) -> SeriesScan { + SeriesScan::new(scan_input(region, store, format, projection), false) +} + +fn scan_input( + region: &MitoRegionRef, + store: ObjectStore, + format: SelectedFormat, + projection: Vec, +) -> ScanInput { + let version = region.version(); + let mapper = FlatProjectionMapper::new(&version.metadata, projection).unwrap(); + ScanInput::builder(region.access_layer.clone(), mapper) + .with_files( + version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .cloned() + .collect(), + ) + .with_series_index(Some(SeriesIndexReadContext { + format, + store, + version: region.series_index_version(), + })) + .with_predicate( + PredicateGroup::new( + &version.metadata, + &[col("tag_0").eq(lit("a")), col("field_0").gt(lit(0.0))], + ) + .unwrap(), + ) + .build() +} + +fn rows(chunks: &[RawSeriesChunk]) -> Vec<(Vec, i64, f64)> { + chunks + .iter() + .flat_map(|chunk| { + (0..chunk.time.len()).map(move |i| { + let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(i) else { + panic!("expected float"); + }; + ( + chunk.primary_key.to_vec(), + chunk.time.timestamp(i).unwrap(), + value, + ) + }) + }) + .collect() +} + +#[tokio::test] +async fn prepared_sources_share_discovery_and_read_independently() { + let (_env, engine, region) = fixture().await; + // A second series and updates/deletes across SSTs exercise the merge adapter. + write(&engine, ®ion, 1, 1500, Some(2.0)).await; + write(&engine, ®ion, 0, 1000, Some(3.0)).await; + write(&engine, ®ion, 0, 2000, None).await; + flush_region(&engine, region.region_id, None).await; + build(&engine, ®ion).await; + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let store = region.series_index_store.clone().unwrap().layer( + MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| { + Box::new(ReadRecorder { + inner, + path: path.to_owned(), + reads: observed.clone(), + }) + })) + .build() + .unwrap(), + ); + let legacy = scanner_with_format(®ion, store.clone(), SelectedFormat::Legacy); + let legacy_binding = labels(legacy.metadata()); + let oracle = legacy + .scan_series_partition( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + 0, + &legacy_binding, + ) + .unwrap() + .unwrap() + .try_collect::>() + .await + .unwrap(); + reads.lock().unwrap().clear(); + let mut scanner = scanner(®ion, store); + let binding = labels(scanner.metadata()); + let ranges = scanner.properties().partitions[0].clone(); + scanner + .prepare(PrepareRequest::default().with_ranges(vec![ranges, vec![], vec![]])) + .unwrap(); + let ctx = QueryScanContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let (a, b) = tokio::join!( + scanner.prepare_series(&ctx, &metrics, &binding), + scanner.prepare_series(&ctx, &metrics, &binding) + ); + let source = a.unwrap().unwrap(); + assert!(Arc::ptr_eq(&source, &b.unwrap().unwrap())); + assert_eq!(2, source.num_series()); + assert!( + !reads + .lock() + .unwrap() + .iter() + .any(|(p, _)| p.ends_with(".samples")) + ); + let left = source.read_series(0..2).unwrap(); + let right = source.read_series(0..2).unwrap(); + assert!( + !reads + .lock() + .unwrap() + .iter() + .any(|(p, _)| p.ends_with(".samples")) + ); + let (left, right) = tokio::join!(left.try_collect::>(), right.try_collect::>()); + let expected = rows(&left.unwrap()); + assert_eq!(rows(&oracle), expected); + assert_eq!(expected, rows(&right.unwrap())); + assert_eq!( + vec![(1000, 3.0), (1500, 2.0)], + expected + .iter() + .map(|(_, t, v)| (*t, *v)) + .collect::>() + ); + let before = reads + .lock() + .unwrap() + .iter() + .filter(|(p, _)| p.contains("/series/")) + .count(); + let mut partition_rows = Vec::new(); + for p in 0..3 { + let chunks = scanner + .scan_series_partition(&ctx, &metrics, p, &binding) + .unwrap() + .unwrap() + .try_collect::>() + .await + .unwrap(); + partition_rows.extend(rows(&chunks)); + } + assert_eq!(expected, partition_rows); + assert_eq!( + before, + reads + .lock() + .unwrap() + .iter() + .filter(|(p, _)| p.contains("/series/")) + .count() + ); + let survivor = source.read_series(0..2).unwrap(); + drop(source); + drop(scanner); + assert_eq!( + expected, + rows(&survivor.try_collect::>().await.unwrap()) + ); + engine.stop().await.unwrap(); +} + +#[tokio::test] +async fn native_scan_reports_incomplete_coverage_and_uses_prepared_source() { + let (_env, engine, region) = fixture().await; + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }; + // Three files are below the automatic index threshold. + let raw_scanner = engine + .handle_query(region.region_id, request.clone()) + .await + .unwrap(); + let binding = labels(raw_scanner.metadata()); + assert!(raw_scanner.raw_series_schema(&binding).is_err()); + assert!( + raw_scanner + .scan_series_partition( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + 0, + &binding, + ) + .is_err() + ); + assert!( + raw_scanner + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding + ) + .await + .is_err() + ); + build(&engine, ®ion).await; + let raw_scanner = engine + .handle_query(region.region_id, request.clone()) + .await + .unwrap(); + let binding = labels(raw_scanner.metadata()); + let plan = Arc::new(RegionScanExec::new(raw_scanner, request.clone(), None).unwrap()) + .try_to_series_scan(binding.clone()) + .unwrap() + .unwrap(); + let frame = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let chunks = unpack_raw_series(frame, Some(plan.series_schema().clone())) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + vec![1000, 2000], + rows(&chunks).iter().map(|(_, t, _)| *t).collect::>() + ); + write(&engine, ®ion, 0, 3000, Some(4.0)).await; + let raw_scanner = engine + .handle_query(region.region_id, request) + .await + .unwrap(); + let binding = labels(raw_scanner.metadata()); + assert!(raw_scanner.raw_series_schema(&binding).is_err()); + assert!( + raw_scanner + .scan_series_partition( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + 0, + &binding, + ) + .is_err() + ); + assert!( + raw_scanner + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding + ) + .await + .err() + .unwrap() + .output_msg() + .contains("memtables") + ); + engine.stop().await.unwrap(); +} + +#[tokio::test] +async fn manual_build_covers_unknown_sequences_without_a_compaction_window() { + use crate::region::options::CompactionOptions; + use crate::region::version::VersionBuilder; + use crate::series_index::reconcile_series_indexes_with_mode; + use crate::series_index::series_index_channel; + + let (_env, engine, region) = fixture().await; + let original = region.version(); + let mut options = original.options.clone(); + let CompactionOptions::Twcs(twcs) = &mut options.compaction; + twcs.time_window = None; + // Keep the original version pinned while replacing handles to the same SSTs. + let originals: Vec<_> = original + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .map(|f| f.meta_ref().clone()) + .collect(); + let mut unknown = originals.clone(); + for file in &mut unknown { + file.sequence = None; + } + let version = VersionBuilder::from_version(original.clone()) + .options(options) + .compaction_time_window(None) + .remove_files(originals.into_iter()) + .add_files(region.file_purger.clone(), unknown.into_iter()) + .build(); + region.version_control.overwrite_current(Arc::new(version)); + let store = region.series_index_store.clone().unwrap(); + let (purger, _receiver) = series_index_channel(store.clone()); + let reconcile = |manual, capacity| { + reconcile_series_indexes_with_mode( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + capacity, + manual, + ) + }; + assert_eq!(0, reconcile(false, true).await.unwrap().built_series); + assert!( + reconcile(true, false) + .await + .unwrap_err() + .to_string() + .contains("capacity") + ); + assert_eq!(1, reconcile(true, true).await.unwrap().built_series); + assert_eq!(0, reconcile(true, true).await.unwrap().built_series); + let scanner = scanner(®ion, store.clone()); + let binding = labels(scanner.metadata()); + let source = scanner + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding, + ) + .await + .unwrap() + .unwrap(); + let chunks = source + .read_series(0..source.num_series()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + vec![1000, 2000], + rows(&chunks).iter().map(|(_, t, _)| *t).collect::>() + ); + // Missing sample objects must fail even when capacity would otherwise make this a no-op. + let missing = region + .series_index_version() + .samples + .values() + .next() + .unwrap() + .entry() + .artifact + .objects[0] + .path + .clone(); + store.delete(&missing).await.unwrap(); + assert!( + reconcile(true, false) + .await + .unwrap_err() + .to_string() + .contains("capacity") + ); + engine.stop().await.unwrap(); +} + +#[rstest::rstest] +#[case::field(SemanticType::Field)] +#[case::tag(SemanticType::Tag)] +#[tokio::test] +async fn manual_build_repairs_schema_changes(#[case] semantic_type: SemanticType) { + use store_api::metadata::ColumnMetadata; + use store_api::region_request::{AddColumn, AlterKind, RegionAlterRequest}; + + use crate::series_index::{reconcile_series_indexes_with_mode, series_index_channel}; + + let (_env, engine, region) = fixture().await; + let store = region.series_index_store.clone().unwrap(); + build(&engine, ®ion).await; + let original = region.series_index_version(); + let projection = (0..region.version().metadata.column_metadatas.len()).collect::>(); + let query = || { + scanner_with_projection( + ®ion, + store.clone(), + SelectedFormat::Columnar, + projection.clone(), + ) + }; + let scan = query(); + let source = scan + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &labels(scan.metadata()), + ) + .await + .unwrap() + .unwrap(); + let expected = rows( + &source + .read_series(0..source.num_series()) + .unwrap() + .try_collect::>() + .await + .unwrap(), + ); + let data_type = if semantic_type == SemanticType::Field { + ConcreteDataType::float64_datatype() + } else { + ConcreteDataType::string_datatype() + }; + engine + .handle_request( + region.region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: datatypes::schema::ColumnSchema::new( + "added", data_type, true, + ), + semantic_type, + column_id: 4, + }, + location: None, + }], + }, + }), + ) + .await + .unwrap(); + // Even a projection containing only the original fields must reject stale sidecars. + let scan = query(); + let error = scan + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &labels(scan.metadata()), + ) + .await + .err() + .unwrap(); + assert!( + error + .output_msg() + .contains(if semantic_type == SemanticType::Field { + "missing scan fields" + } else { + "missing scan tags" + }), + "{error}" + ); + let (purger, _receiver) = series_index_channel(store.clone()); + let error = reconcile_series_indexes_with_mode( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger, + false, + false, + true, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("capacity")); + assert!(Arc::ptr_eq(&original, ®ion.series_index_version())); + + build(&engine, ®ion).await; + let rebuilt = region.series_index_version(); + for (id, old) in &original.samples { + let new = &rebuilt.samples[id]; + assert_eq!( + semantic_type == SemanticType::Tag, + old.entry().artifact.generation == new.entry().artifact.generation + ); + if semantic_type == SemanticType::Field { + let reader = ColumnarSamplesFormat::default() + .reader(store.clone(), new.clone()) + .await + .unwrap(); + let batches = reader + .read(crate::series_index::format::SamplesReadRequest { + projection: vec!["added".to_string()], + ..Default::default() + }) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.num_rows(), batch.column(0).null_count()); + } + } + } + assert_eq!( + semantic_type == SemanticType::Field, + original.series_indexes.keys().collect::>() + == rebuilt.series_indexes.keys().collect::>() + ); + let scan = query(); + let source = scan + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &labels(scan.metadata()), + ) + .await + .unwrap() + .unwrap(); + let actual = source + .read_series(0..source.num_series()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(expected, rows(&actual)); + + build(&engine, ®ion).await; + let unchanged = region.series_index_version(); + assert!( + rebuilt.series_indexes.keys().collect::>() + == unchanged.series_indexes.keys().collect::>() + ); + for (id, old) in &rebuilt.samples { + assert_eq!(old.entry(), unchanged.samples[id].entry()); + } + engine.stop().await.unwrap(); +} + +struct SamplesReadRecorder { + inner: Arc, + requests: Arc>>, +} + +impl SamplesFormatReader for SamplesReadRecorder { + fn schema(&self) -> datatypes::arrow::datatypes::SchemaRef { + self.inner.schema() + } + + fn read( + &self, + request: SamplesReadRequest, + ) -> Result>> { + self.requests.lock().unwrap().push(request.clone()); + self.inner.read(request) + } +} + +#[tokio::test] +async fn ordinal_range_batches_sample_requests_and_pushes_time_bounds() { + use common_time::{Timestamp, range::TimestampRange}; + + let (_env, engine, region) = fixture().await; + write(&engine, ®ion, 1, 1500, Some(2.0)).await; + write(&engine, ®ion, 2, 1600, Some(3.0)).await; + write(&engine, ®ion, 0, 1000, Some(4.0)).await; + write(&engine, ®ion, 0, 2000, None).await; + flush_region(&engine, region.region_id, None).await; + build(&engine, ®ion).await; + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let store = region.series_index_store.clone().unwrap().layer( + MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| { + Box::new(ReadRecorder { + inner, + path: path.to_owned(), + reads: observed.clone(), + }) + })) + .build() + .unwrap(), + ); + let make_scan = |format| { + let mut input = scan_input( + ®ion, + store.clone(), + format, + (0..region.version().metadata.column_metadatas.len()).collect(), + ); + // Bounds in a finer unit must round up for both inclusive start and exclusive end. + input.time_range = TimestampRange::new( + Timestamp::new_microsecond(1_000_001), + Timestamp::new_microsecond(1_600_001), + ); + SeriesScan::new(input, false) + }; + let legacy = make_scan(SelectedFormat::Legacy); + let ctx = QueryScanContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let oracle = legacy + .scan_series_partition(&ctx, &metrics, 0, &labels(legacy.metadata())) + .unwrap() + .unwrap() + .try_collect::>() + .await + .unwrap(); + let scan = make_scan(SelectedFormat::Columnar); + scan.prepare_series(&ctx, &metrics, &labels(scan.metadata())) + .await + .unwrap() + .unwrap(); + let source = scan.prepared.source.get().unwrap().as_ref().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let source = ColumnarPreparedSource { + context: source.context.clone(), + labels: source.labels.clone(), + raw: source.raw.clone(), + candidates: source.candidates.clone(), + metrics: source.metrics.clone(), + reservation: source.reservation.clone(), + samples: Arc::new( + source + .samples + .iter() + .map(|inner| { + Arc::new(SamplesReadRecorder { + inner: inner.clone(), + requests: requests.clone(), + }) as Arc + }) + .collect(), + ), + }; + let stream = source.read_series(0..source.num_series()).unwrap(); + assert!(requests.lock().unwrap().is_empty()); + let result = stream.try_collect::>().await.unwrap(); + // SST time-range pruning is inexact; apply the exact residual bounds to the oracle. + let expected = rows(&oracle) + .into_iter() + .filter(|(_, time, _)| (1001..1601).contains(time)) + .collect::>(); + assert_eq!(expected, rows(&result)); + assert_eq!( + vec![1500, 1600], + rows(&result).iter().map(|(_, t, _)| *t).collect::>() + ); + { + let requests = requests.lock().unwrap(); + assert_eq!(source.samples.len(), requests.len()); + let ids = source + .candidates + .iter() + .map(|c| c.id) + .collect::>(); + for request in requests.iter() { + assert_eq!(Some(&ids), request.ids.as_ref()); + assert_eq!(Some(1001), request.start); + assert_eq!(Some(1601), request.end); + } + } + // Each selected samples part opens its header once, even when it has multiple series. + let mut headers = HashMap::::new(); + for (path, range) in reads.lock().unwrap().iter() { + if path.ends_with(".samples") && range.offset() == 0 { + *headers.entry(path.clone()).or_default() += 1; + } + } + assert!(!headers.is_empty()); + assert!(headers.values().all(|count| *count == 1)); + requests.lock().unwrap().clear(); + assert!( + source + .read_series(1..1) + .unwrap() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + assert!(requests.lock().unwrap().is_empty()); + assert!(source.read_series(0..source.num_series() + 1).is_err()); + let chunks = source + .read_series(1..3) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(rows(&result), rows(&chunks)); + assert!( + requests + .lock() + .unwrap() + .iter() + .all(|r| r.ids.as_ref().unwrap().len() == 2) + ); + engine.stop().await.unwrap(); +} + +#[rstest::rstest] +#[case(common_time::timestamp::TimeUnit::Second)] +#[case(common_time::timestamp::TimeUnit::Millisecond)] +#[case(common_time::timestamp::TimeUnit::Microsecond)] +#[case(common_time::timestamp::TimeUnit::Nanosecond)] +#[tokio::test] +async fn visibility_uses_named_columns_and_region_time_unit( + #[case] unit: common_time::timestamp::TimeUnit, +) { + use common_time::{Timestamp, range::TimestampRange}; + use datatypes::arrow::array::Int64Array; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use store_api::storage::SequenceRange; + + let (_env, engine, region) = fixture().await; + let mut metadata = (*region.version().metadata).clone(); + let time_name = metadata.time_index_column().column_schema.name.clone(); + metadata + .column_metadatas + .iter_mut() + .find(|c| c.column_schema.name == time_name) + .unwrap() + .column_schema + .data_type = ConcreteDataType::timestamp_datatype(unit); + let mapper = FlatProjectionMapper::new(&Arc::new(metadata), vec![]).unwrap(); + let input = ScanInput::builder(region.access_layer.clone(), mapper) + .with_snapshot_sequence(Some(4)) + .with_sequence_range(SequenceRange::new(Some(1), Some(5))) + .with_time_range(TimestampRange::new( + Timestamp::new(-1, unit), + Timestamp::new(2, unit), + )) + .build(); + let context = StreamContext::seq_scan_ctx(input); + let time_type = DataType::Timestamp(unit.into(), None); + let values: ArrayRef = Arc::new(Int64Array::from(vec![-2, -1, 0, 1, 2])); + let time = cast(&values, &time_type).unwrap(); + // Neither internal column is at its normal flat-format offset. + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(&time_name, time_type, false), + Field::new("other", DataType::Int64, false), + Field::new(SEQUENCE_COLUMN_NAME, DataType::UInt64, false), + ])), + vec![ + time, + values, + Arc::new(UInt64Array::from(vec![1, 2, 3, 4, 5])), + ], + ) + .unwrap(); + let filtered = filter_visibility(batch.clone(), &context).unwrap(); + assert_eq!(3, filtered.num_rows()); + assert_eq!( + &[2, 3, 4], + filtered + .column_by_name(SEQUENCE_COLUMN_NAME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ref() + ); + assert!( + filter_visibility(batch.project(&[0, 1]).unwrap(), &context) + .unwrap_err() + .to_string() + .contains("missing sequence") + ); + assert!( + filter_visibility(batch.project(&[1, 2]).unwrap(), &context) + .unwrap_err() + .to_string() + .contains("missing time index") + ); + let wrong_time = RecordBatch::try_from_iter(vec![ + (time_name.as_str(), batch.column(1).clone()), + (SEQUENCE_COLUMN_NAME, batch.column(2).clone()), + ]) + .unwrap(); + assert!( + filter_visibility(wrong_time, &context) + .unwrap_err() + .to_string() + .contains("must be a timestamp") + ); + let wrong_unit = if unit == common_time::timestamp::TimeUnit::Second { + common_time::timestamp::TimeUnit::Millisecond + } else { + common_time::timestamp::TimeUnit::Second + }; + let wrong_time = RecordBatch::try_from_iter(vec![ + ( + time_name.as_str(), + cast( + batch.column(1), + &DataType::Timestamp(wrong_unit.into(), None), + ) + .unwrap(), + ), + (SEQUENCE_COLUMN_NAME, batch.column(2).clone()), + ]) + .unwrap(); + assert!( + filter_visibility(wrong_time, &context) + .unwrap_err() + .to_string() + .contains("expected") + ); + engine.stop().await.unwrap(); +} + +#[tokio::test] +async fn preflight_rejects_unsupported_assignments_without_artifact_io() { + use crate::read::range::RowGroupIndex; + use crate::series_index::SeriesIndexVersion; + use crate::sst::file::FileHandle; + use crate::sst::file_purger::NoopFilePurger; + + let (_env, engine, region) = fixture().await; + build(&engine, ®ion).await; + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let store = region.series_index_store.clone().unwrap().layer( + MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| { + Box::new(ReadRecorder { + inner, + path: path.to_owned(), + reads: observed.clone(), + }) + })) + .build() + .unwrap(), + ); + let make_context = || { + StreamContext::seq_scan_ctx(scan_input( + ®ion, + store.clone(), + SelectedFormat::Columnar, + (0..region.version().metadata.column_metadatas.len()).collect(), + )) + }; + let mut context = make_context(); + let ranges = context.partition_ranges(); + validate_inputs(&context, &ranges).unwrap(); + context.ranges[ranges[0].identifier].row_group_indices = vec![RowGroupIndex { + index: context.input.num_files(), + row_group_index: -1, + }] + .into(); + assert!( + validate_inputs(&context, &ranges) + .unwrap_err() + .to_string() + .contains("extension ranges") + ); + + let mut context = make_context(); + context.ranges[ranges[0].identifier].row_group_indices = vec![RowGroupIndex { + index: 0, + row_group_index: i64::MAX, + }] + .into(); + assert!( + validate_inputs(&context, &ranges[..1]) + .unwrap_err() + .to_string() + .contains("complete SST") + ); + + let mut context = make_context(); + let mut meta = context.input.files[0].meta_ref().clone(); + meta.region_id = store_api::storage::RegionId::new(123, 456); + context.input.files[0] = FileHandle::new(meta, Arc::new(NoopFilePurger)); + assert!( + validate_inputs(&context, &ranges) + .unwrap_err() + .to_string() + .contains("foreign SST") + ); + + let mut context = make_context(); + let index = context.input.series_index.as_mut().unwrap(); + index.version = Arc::new(SeriesIndexVersion { + series_indexes: index.version.series_indexes.clone(), + ..Default::default() + }); + assert!( + validate_inputs(&context, &ranges) + .unwrap_err() + .to_string() + .contains("no columnar samples") + ); + let scan = SeriesScan::new(context.input, false); + let binding = labels(scan.metadata()); + assert!(scan.raw_series_schema(&binding).is_err()); + assert!( + scan.scan_series_partition( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + 0, + &binding + ) + .is_err() + ); + assert!( + scan.prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding + ) + .await + .is_err() + ); + assert!(reads.lock().unwrap().is_empty()); + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/read/series_scan/raw.rs b/src/mito2/src/read/series_scan/raw.rs index 00d6867921f..414f08628bc 100644 --- a/src/mito2/src/read/series_scan/raw.rs +++ b/src/mito2/src/read/series_scan/raw.rs @@ -287,3 +287,19 @@ pub(crate) fn make_stream( inner: Box::pin(stream), }) } + +/// Wraps deferred preparation while exposing the scan contract before polling. +pub(crate) fn from_stream( + schema: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + inner: BoxStream<'static, Result>, +) -> SendableRawSeriesStream { + Box::pin(RawStream { + schema, + label_decode_spec: label_decode_spec(&label_source), + label_source, + ordering, + inner, + }) +} diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 37c91946db3..76605702cb4 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -436,7 +436,8 @@ impl RegionOpener { Ok(Arc::new(MitoRegion { region_id, version_control, - series_index_version_control: Default::default(), + series_index_version_control: crate::series_index::SeriesIndexVersionControl::default() + .with_format(config.experimental_series_index_format), series_index_store: self.series_index_store.clone(), access_layer: access_layer.clone(), // Region is writable after it is created. @@ -680,7 +681,8 @@ impl RegionOpener { let region = MitoRegion { region_id: self.region_id, version_control: version_control.clone(), - series_index_version_control, + series_index_version_control: series_index_version_control + .with_format(config.experimental_series_index_format), series_index_store: self.series_index_store.clone(), access_layer: access_layer.clone(), // Region is always opened in read only mode. diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index dc8b461afd6..ed72734fe1a 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -21,20 +21,12 @@ mod artifact_writer; mod bucket; mod builder; mod catalog; -#[allow( - dead_code, - reason = "Opt-in format; query consumer integration is deferred." -)] pub(crate) mod columnar; pub(crate) mod format; #[cfg(test)] mod format_tests; mod maintenance; mod purger; -#[allow( - dead_code, - reason = "Samples are an opt-in backend; query integration is deferred." -)] pub(crate) mod samples; mod searcher; mod task; @@ -43,6 +35,11 @@ mod tests; mod version; mod writer; +#[cfg(test)] +pub(crate) use maintenance::reconcile_series_indexes_with_mode; +#[cfg(test)] +pub(crate) use tests::prepare_region_with_metadata; + use std::sync::Arc; use futures::stream::BoxStream; @@ -71,6 +68,7 @@ pub use crate::series_index::writer::{ /// Index storage and pinned catalog snapshot for a query. #[derive(Clone)] pub(crate) struct SeriesIndexReadContext { + pub(crate) format: crate::config::SeriesIndexFormat, pub(crate) store: ObjectStore, pub(crate) version: Arc, } diff --git a/src/mito2/src/series_index/bucket.rs b/src/mito2/src/series_index/bucket.rs index 470caeadd3b..e7395177ce2 100644 --- a/src/mito2/src/series_index/bucket.rs +++ b/src/mito2/src/series_index/bucket.rs @@ -14,7 +14,7 @@ //! Event-time bucket planning and source coverage. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashSet}; use std::time::Duration; use common_time::{TimeToLive, Timestamp}; @@ -32,6 +32,8 @@ pub(crate) struct IndexBucket { pub(crate) start: Timestamp, pub(crate) end: Timestamp, pub(crate) index_ids: SmallVec<[FileId; 2]>, + /// Exact input set when a manual build contains unknown file sequences. + pub(crate) unknown_sequence_sources: HashSet, /// Zero means that merged indexes use incompatible window widths. pub(crate) compaction_window_secs: i64, /// Indexed SST summaries keyed by aligned start. @@ -45,6 +47,11 @@ impl IndexBucket { start: entry.bucket_start, end: entry.bucket_end, index_ids: smallvec![entry.index_uuid], + unknown_sequence_sources: if entry.min_file_sequence == 0 { + entry.source_file_ids.iter().copied().collect() + } else { + HashSet::new() + }, compaction_window_secs: entry.compaction_window_secs, window_sequences: entry.window_sequences.clone(), } @@ -66,6 +73,8 @@ impl IndexBucket { self.window_sequences.clear(); } } + self.unknown_sequence_sources + .extend(other.unknown_sequence_sources); self.index_ids.append(&mut other.index_ids); } @@ -122,21 +131,33 @@ pub(crate) struct SeriesIndexPlan { /// 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. +#[cfg(test)] pub(crate) fn plan_series_indexes( + buckets: Vec, + index_buckets: BTreeMap, + ttl: Option, + now_ms: i64, +) -> SeriesIndexPlan { + plan_series_indexes_with_mode(buckets, index_buckets, ttl, now_ms, false) +} + +pub(crate) fn plan_series_indexes_with_mode( buckets: Vec, mut index_buckets: BTreeMap, ttl: Option, now_ms: i64, + manual: bool, ) -> 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) - }) + !manual + && 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| { @@ -153,17 +174,21 @@ pub(crate) fn plan_series_indexes( if expired(bucket.end) { continue; } - if bucket.has_unknown_sequence { + if bucket.has_unknown_sequence && !manual { continue; } if index_buckets.get(&bucket.start).is_some_and(|indexed| { - indexed.end == bucket.end + (!bucket.has_unknown_sequence + || (manual + && indexed.unknown_sequence_sources + == bucket.files.iter().map(|f| f.file_id().file_id()).collect())) + && 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() { + if let Some(entry) = bucket.to_series_entry(manual) { index_buckets.retain(|_, indexed| { if indexed.start < bucket.end && bucket.start < indexed.end { superseded_index_ids.extend_from_slice(&indexed.index_ids); @@ -317,9 +342,12 @@ fn merge_window_sequences( 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 { + /// Automatic builds defer unknown sequences and buckets with too few files. + fn to_series_entry(&self, manual: bool) -> Option { + if self.files.is_empty() + || (!manual + && (self.has_unknown_sequence || self.files.len() < SERIES_INDEX_TRIGGER_FILES)) + { return None; } let mut source_file_ids = self @@ -332,7 +360,8 @@ impl SeriesBucket { .files .iter() .filter_map(|file| file.meta_ref().sequence.map(|sequence| sequence.get())) - .min()?; + .min() + .unwrap_or(0); Some(SeriesIndexEntry { artifact: None, file_size: 0, @@ -340,7 +369,11 @@ impl SeriesBucket { bucket_start: self.start, bucket_end: self.end, source_file_ids, - min_file_sequence, + min_file_sequence: if self.has_unknown_sequence { + 0 + } else { + min_file_sequence + }, max_file_sequence: self.max_file_sequence, compaction_window_secs: self.compaction_window_secs, window_sequences: self.window_sequences.clone(), @@ -431,8 +464,8 @@ mod tests { ); 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()); + assert!(buckets[0].to_series_entry(false).is_none()); + assert!(buckets[1].to_series_entry(false).is_none()); let mut files = files.to_vec(); files.push(file( None, @@ -442,7 +475,7 @@ mod tests { )); assert!( group_files_into_series_buckets(&files, width, 10)[0] - .to_series_entry() + .to_series_entry(false) .is_none() ); } @@ -483,6 +516,7 @@ mod tests { (70, 100, 20, ids[2]), ] { IndexBucket { + unknown_sequence_sources: HashSet::new(), start: ts(start), end: ts(end), index_ids: smallvec![id], @@ -740,6 +774,7 @@ mod tests { let mut indexes = BTreeMap::new(); for (start, seq) in [(0, 10), (20, 30)] { IndexBucket { + unknown_sequence_sources: HashSet::new(), start: ts(start), end: ts(start + 10), index_ids: smallvec![FileId::random()], @@ -836,6 +871,7 @@ mod tests { fn test_index_map_merge_is_order_independent() { let ts = Timestamp::new_second; let make_index = |start, end, width, windows: &[(i64, u64)]| IndexBucket { + unknown_sequence_sources: HashSet::new(), start: ts(start), end: ts(end), index_ids: smallvec![FileId::random()], diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index c875456a819..bd57d362b1d 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -139,7 +139,13 @@ pub(crate) async fn build_series_index( 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; + return UnexpectedSnafu { + reason: format!( + "series-index source SST {} has no readable input", + file.file_id() + ), + } + .fail(); }; let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; sources.push(Box::pin(try_stream! { diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index 6cb7b33912b..f588be72b9b 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -58,7 +58,8 @@ pub(crate) struct WindowSequence { /// /// A published entry must include every series from every SST in the region /// contained by its bucket and inclusive file-sequence interval. Query planning -/// relies on this complete-coverage contract, not on `source_file_ids`. +/// relies on this complete-coverage contract. Manual builds additionally use exact +/// `source_file_ids` membership for SSTs without a file sequence. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SeriesIndexEntry { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -70,8 +71,8 @@ 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. + /// Source SST IDs retained for diagnostics and exact coverage of unknown-sequence files. Compaction can replace these + /// files without changing indexed data, so known-sequence reuse does not depend on IDs. pub(crate) source_file_ids: Vec, pub(crate) min_file_sequence: u64, pub(crate) max_file_sequence: u64, @@ -104,9 +105,15 @@ impl SeriesIndexEntry { && file.time_range.0 <= file.time_range.1 && self.bucket_start <= file.time_range.0 && file.time_range.1 < self.bucket_end - && file.sequence.is_some_and(|sequence| { - self.min_file_sequence <= sequence.get() && sequence.get() <= self.max_file_sequence - }) + && match file.sequence { + Some(sequence) => { + self.min_file_sequence <= sequence.get() + && sequence.get() <= self.max_file_sequence + } + // Manual builds can cover immutable files without sequence metadata. + // This is exact membership, never a claim about replacement SSTs. + None => self.min_file_sequence == 0 && self.source_file_ids.contains(&file.file_id), + } } } diff --git a/src/mito2/src/series_index/columnar.rs b/src/mito2/src/series_index/columnar.rs index 5ebe8d86fcf..c48ec8e8ee2 100644 --- a/src/mito2/src/series_index/columnar.rs +++ b/src/mito2/src/series_index/columnar.rs @@ -138,11 +138,9 @@ impl Manifest { Field::new(TSID_COLUMN, DataType::UInt64, false), Field::new(ROW_COUNT_COLUMN, DataType::UInt64, false), ]; - fields.extend( - self.tags - .iter() - .map(|tag| Field::new(&tag.name, DataType::Utf8, true)), - ); + fields.extend(self.tags.iter().map(|tag| { + crate::sst::with_field_id(Field::new(&tag.name, DataType::Utf8, true), tag.column_id) + })); Arc::new(Schema::new(fields)) } } diff --git a/src/mito2/src/series_index/columnar/codec.rs b/src/mito2/src/series_index/columnar/codec.rs index 53e0a678432..2d3dafcba5b 100644 --- a/src/mito2/src/series_index/columnar/codec.rs +++ b/src/mito2/src/series_index/columnar/codec.rs @@ -28,6 +28,10 @@ pub(crate) const DICTIONARY: u16 = 2; #[derive(Debug, Clone, Copy)] #[repr(u16)] pub(crate) enum Compression { + #[allow( + dead_code, + reason = "Uncompressed encoding is supported for format tests." + )] None = 0, Zstd = 1, } diff --git a/src/mito2/src/series_index/format.rs b/src/mito2/src/series_index/format.rs index eb552a8eabe..b646f28616d 100644 --- a/src/mito2/src/series_index/format.rs +++ b/src/mito2/src/series_index/format.rs @@ -14,23 +14,27 @@ //! Storage-format boundaries. Encodings stay outside reconciliation and source scheduling. -use std::collections::BTreeSet; +use std::collections::{BTreeSet, HashSet}; use std::fmt::Debug; use std::sync::Arc; +use api::v1::SemanticType; use async_trait::async_trait; use common_recordbatch::filter::SimpleFilterEvaluator; -use datatypes::arrow::datatypes::SchemaRef; +use datatypes::arrow::datatypes::{Field, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::DataType; use futures::stream::BoxStream; use object_store::ObjectStore; -use store_api::metadata::RegionMetadataRef; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::FileId; +use store_api::storage::consts::ReservedColumnId; -use crate::error::Result; +use crate::error::{InvalidRecordBatchSnafu, Result}; use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata, SamplesFileHandle}; use crate::series_index::catalog::SeriesIndexEntry; use crate::series_index::{MetricSeriesId, SeriesIndexFileHandle}; +use crate::sst::PARQUET_FIELD_ID_KEY; /// Writers own partial outputs until finish succeeds; abort removes every partial object. /// A failed finish must clean up its output before returning; implementations also own @@ -43,10 +47,6 @@ pub(crate) trait FormatWriter: Send { } /// A read has independent traversal state and pins its immutable artifact. -#[allow( - dead_code, - reason = "Query adapter will consume the format reader separately." -)] #[derive(Debug, Clone, Default)] pub(crate) struct SeriesReadRequest { pub(crate) projection: Vec, @@ -55,10 +55,6 @@ pub(crate) struct SeriesReadRequest { pub(crate) filters: Vec, } -#[allow( - dead_code, - reason = "Query adapter will consume the format reader separately." -)] pub(crate) trait FormatReader: Send + Sync { fn schema(&self) -> SchemaRef; fn read(&self, request: SeriesReadRequest) -> Result>>; @@ -68,7 +64,6 @@ pub(crate) trait FormatReader: Send + Sync { #[async_trait] pub(crate) trait SeriesIndexFormat: Debug + Send + Sync { fn format(&self) -> ArtifactFormat; - #[allow(dead_code, reason = "Query adapter is deferred.")] async fn reader( &self, _store: ObjectStore, @@ -88,7 +83,6 @@ pub(crate) trait SeriesIndexFormat: Debug + Send + Sync { } /// Selection in the stored timestamp unit. No visibility filtering is implicit. -#[allow(dead_code, reason = "Query adapter is deferred.")] #[derive(Debug, Clone, Default)] pub(crate) struct SamplesReadRequest { pub(crate) projection: Vec, @@ -98,7 +92,6 @@ pub(crate) struct SamplesReadRequest { pub(crate) end: Option, } -#[allow(dead_code, reason = "Query adapter is deferred.")] pub(crate) trait SamplesFormatReader: Send + Sync { fn schema(&self) -> SchemaRef; fn read(&self, request: SamplesReadRequest) -> Result>>; @@ -108,7 +101,6 @@ pub(crate) trait SamplesFormatReader: Send + Sync { #[async_trait] pub(crate) trait SamplesFormat: Debug + Send + Sync { fn format(&self) -> ArtifactFormat; - #[allow(dead_code, reason = "Query adapter is deferred.")] async fn reader( &self, _store: ObjectStore, @@ -137,3 +129,164 @@ pub(crate) struct SeriesFormats { pub(crate) index: Option>, pub(crate) samples: Option>, } + +impl SeriesFormats { + pub(crate) fn from_config(format: crate::config::SeriesIndexFormat) -> Self { + match format { + crate::config::SeriesIndexFormat::Legacy => Self::default(), + crate::config::SeriesIndexFormat::Columnar => Self { + index: Some(Arc::new( + crate::series_index::columnar::ColumnarSeriesFormat::default(), + )), + samples: Some(Arc::new( + crate::series_index::samples::ColumnarSamplesFormat::default(), + )), + }, + } + } +} + +/// Reads a column ID from an artifact schema field. +pub(crate) fn field_id(field: &Field) -> Result { + field + .metadata() + .get(PARQUET_FIELD_ID_KEY) + .and_then(|id| id.parse().ok()) + .ok_or_else(|| invalid(format!("missing column ID for {}", field.name()))) +} + +/// Checks whether samples can be read with the current region schema. +pub(crate) fn validate_sample_schema(schema: &SchemaRef, metadata: &RegionMetadata) -> Result<()> { + let mut ids = HashSet::new(); + for field in schema + .fields() + .iter() + .skip(2) + .filter(|f| f.name() != "__sequence" && f.name() != "__op_type") + { + let id = field_id(field)?; + let column = metadata + .column_by_id(id) + .ok_or_else(|| invalid("sample column absent from scan schema"))?; + if column.semantic_type == SemanticType::Tag + || column.column_schema.name != *field.name() + || column.column_schema.data_type.as_arrow_type() != *field.data_type() + { + return Err(invalid("samples schema changed; rebuild series indexes")); + } + ids.insert(id); + } + if metadata + .column_metadatas + .iter() + .any(|c| c.semantic_type != SemanticType::Tag && !ids.contains(&c.column_id)) + { + return Err(invalid( + "samples are missing scan fields; rebuild series indexes", + )); + } + Ok(()) +} + +/// Checks whether a series index contains exactly the current non-reserved tags. +pub(crate) fn validate_series_schema(schema: &SchemaRef, metadata: &RegionMetadata) -> Result<()> { + let mut ids = HashSet::new(); + for field in schema.fields().iter().skip(3) { + let id = field_id(field)?; + let column = metadata + .column_by_id(id) + .ok_or_else(|| invalid("series tag absent from scan schema"))?; + if column.semantic_type != SemanticType::Tag || column.column_schema.name != *field.name() { + return Err(invalid("series tag schema changed; rebuild series indexes")); + } + ids.insert(id); + } + if metadata + .primary_key_columns() + .any(|c| !ReservedColumnId::is_reserved(c.column_id) && !ids.contains(&c.column_id)) + { + return Err(invalid( + "series index is missing scan tags; rebuild series indexes", + )); + } + Ok(()) +} + +fn invalid(reason: impl Into) -> crate::error::Error { + InvalidRecordBatchSnafu { + reason: reason.into(), + } + .build() +} + +#[cfg(test)] +mod tests { + use datatypes::arrow::datatypes::{DataType, Schema}; + use datatypes::prelude::DataType as _; + use store_api::codec::PrimaryKeyEncoding; + + use super::*; + use crate::test_util::sst_util::sst_region_metadata_with_encoding; + + #[test] + fn schema_compatibility_requires_current_columns() { + let metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + for tags in [false, true] { + let mut fields = vec![ + Field::new("__table_id", DataType::UInt32, false), + Field::new("__tsid", DataType::UInt64, false), + ]; + if tags { + fields.push(Field::new("__row_count", DataType::UInt64, false)); + } + fields.extend( + metadata + .column_metadatas + .iter() + .filter(|c| { + (c.semantic_type == SemanticType::Tag) == tags + && !ReservedColumnId::is_reserved(c.column_id) + }) + .map(|c| { + crate::sst::with_field_id( + Field::new( + &c.column_schema.name, + c.column_schema.data_type.as_arrow_type(), + true, + ), + c.column_id, + ) + }), + ); + let validate = |fields: Vec| { + let schema = Arc::new(Schema::new(fields)); + if tags { + validate_series_schema(&schema, &metadata) + } else { + validate_sample_schema(&schema, &metadata) + } + }; + validate(fields.clone()).unwrap(); + let index = if tags { 3 } else { 2 }; + let mut missing = fields.clone(); + missing.remove(index); + assert!(validate(missing).is_err()); + let mut renamed = fields.clone(); + renamed[index] = renamed[index].clone().with_name("renamed"); + assert!(validate(renamed).is_err()); + let mut replaced = fields.clone(); + replaced[index] = crate::sst::with_field_id(replaced[index].clone(), 100); + assert!(validate(replaced).is_err()); + if !tags { + let mut changed_type = fields.clone(); + changed_type[index] = changed_type[index] + .clone() + .with_data_type(DataType::Float64); + assert!(validate(changed_type).is_err()); + } + let mut reordered = fields.clone(); + reordered[index..].reverse(); + validate(reordered).unwrap(); + } + } +} diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index 1e777a2aa56..d6897ffe900 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -28,14 +28,16 @@ use crate::metrics::{SERIES_INDEX_RECONCILE_ELAPSED, SERIES_INDEX_RECONCILE_TOTA use crate::read::series_candidate::is_sparse_metric_metadata; use crate::region::version::VersionRef; use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::series_index::artifact::ArtifactFormat; use crate::series_index::bucket::{ - group_files_into_series_buckets, plan_series_indexes, rounded_bucket_width, + group_files_into_series_buckets, plan_series_indexes_with_mode, rounded_bucket_width, }; use crate::series_index::builder::{build_range_index, build_series_index}; use crate::series_index::catalog::{ RangeIndexCatalog, SeriesIndexCatalog, delete_catalogs, range_catalog_path, series_catalog_path, store_catalog, }; +use crate::series_index::format; use crate::series_index::purger::IndexFilePurger; use crate::series_index::version::{SeriesIndexFileHandle, SeriesIndexVersion}; @@ -91,6 +93,7 @@ impl ReconcileStats { /// Reconciles indexes for one region snapshot, persists catalogs, then atomically publishes it. #[allow(clippy::too_many_arguments)] +#[cfg(test)] pub(crate) async fn reconcile_series_indexes( worker_id: u32, store: ObjectStore, @@ -100,6 +103,33 @@ pub(crate) async fn reconcile_series_indexes( purger: IndexFilePurger, enable_range_index: bool, allow_builds: bool, +) -> Result { + reconcile_series_indexes_with_mode( + worker_id, + store, + region, + requested_bucket_width, + now_ms, + purger, + enable_range_index, + allow_builds, + false, + ) + .await +} + +/// Manual reconciliation bypasses maintenance eligibility, but not capacity limits. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn reconcile_series_indexes_with_mode( + worker_id: u32, + store: ObjectStore, + region: MitoRegionRef, + requested_bucket_width: Duration, + now_ms: i64, + purger: IndexFilePurger, + enable_range_index: bool, + allow_builds: bool, + manual: bool, ) -> Result { let total_start = Instant::now(); // Use this snapshot throughout reconciliation, even if the region version advances. @@ -123,6 +153,7 @@ pub(crate) async fn reconcile_series_indexes( &mut unpublished, enable_range_index, allow_builds, + manual, ) .await?; SERIES_INDEX_RECONCILE_ELAPSED @@ -184,6 +215,7 @@ async fn build_index_version( unpublished: &mut UnpublishedSeriesFiles, enable_range_index: bool, allow_builds: bool, + manual: bool, ) -> Result<(Option, ReconcileStats)> { let mut stats = ReconcileStats::default(); let files = version @@ -208,6 +240,12 @@ async fn build_index_version( group_files_into_series_buckets(&files, width, (window.as_secs() as i64).max(1)) }) .unwrap_or_default(), + (true, None) if manual => { + let width = i64::try_from(requested_bucket_width.as_secs()) + .unwrap_or(i64::MAX) + .max(1); + group_files_into_series_buckets(&files, width, width) + } (true, None) => { debug!( "Deferring series indexes without compaction window, worker: {worker_id}, region: {}", @@ -240,10 +278,51 @@ async fn build_index_version( .context(crate::error::OpenDalSnafu)?; } } + if manual && reusable && entry.artifact.is_none() { + reusable = store + .exists(&crate::series_index::catalog::series_index_path( + region.region_id, + entry.index_uuid, + )) + .await + .context(crate::error::OpenDalSnafu)?; + } + // Manual builds must repair schema changes even when SST coverage is unchanged. + if manual + && reusable + && selected_format == ArtifactFormat::ColumnarSeries + && let Some(backend) = &formats.index + { + let reader = backend.reader(store.clone(), handle.clone()).await?; + reusable = format::validate_series_schema(&reader.schema(), &version.metadata).is_ok(); + } if !reusable { unusable.insert(entry.index_uuid); } } + let mut reusable_samples = HashSet::new(); + if let Some(backend) = &formats.samples { + for (id, handle) in ¤t.samples { + if !visible.contains(id) || handle.entry().artifact.format != backend.format() { + continue; + } + let mut reusable = true; + for object in &handle.entry().artifact.objects { + reusable &= store + .exists(&object.path) + .await + .context(crate::error::OpenDalSnafu)?; + } + if manual && reusable && backend.format() == ArtifactFormat::ColumnarSamples { + let reader = backend.reader(store.clone(), handle.clone()).await?; + reusable = + format::validate_sample_schema(&reader.schema(), &version.metadata).is_ok(); + } + if reusable { + reusable_samples.insert(*id); + } + } + } let mut index_buckets = current.index_buckets.clone(); for bucket in index_buckets.values_mut() { if bucket.index_ids.iter().any(|id| unusable.contains(id)) { @@ -253,19 +332,40 @@ async fn build_index_version( bucket.window_sequences.clear(); } } - let plan = plan_series_indexes(buckets, index_buckets, version.options.ttl, now_ms); + let plan = + plan_series_indexes_with_mode(buckets, index_buckets, version.options.ttl, now_ms, manual); stats.computed_buckets = plan.computed_buckets; stats.skipped_buckets = plan.skipped_buckets; let mut next = prune_index_version(¤t, &visible, &plan.expired_index_ids, &mut stats); if !allow_builds { + if manual + && files.iter().any(|file| { + !current.series_indexes.values().any(|handle| { + handle.entry().format() == selected_format + && !unusable.contains(&handle.entry().index_uuid) + && handle + .entry() + .covers_file(file.meta_ref(), region.region_id) + }) || (formats.samples.is_some() + && !reusable_samples.contains(&file.file_id().file_id())) + }) + { + return crate::error::InvalidRequestSnafu { + region_id: region.region_id, + reason: "series-index capacity limit prevents complete manual coverage", + } + .fail(); + } if let Some(next) = &mut next { next.index_buckets = plan.index_buckets; } return Ok((next, stats)); } - if plan.builds.is_empty() + if !manual + && plan.builds.is_empty() && next.is_none() && formats.samples.is_none() + && current.samples.is_empty() && (!enable_range_index || current.range_indexes.len() == visible.len()) { return Ok((None, stats)); @@ -282,22 +382,14 @@ async fn build_index_version( index_buckets: Default::default(), }); let index_buckets = plan.index_buckets; + if formats.samples.is_none() { + stats.removed_samples += samples.len(); + samples.clear(); + } if let Some(backend) = &formats.samples { for file in &files { let id = file.file_id().file_id(); - let reusable = if let Some(handle) = samples.get(&id) { - let artifact = &handle.entry().artifact; - let mut exists = artifact.format == backend.format(); - for object in &artifact.objects { - exists &= store - .exists(&object.path) - .await - .context(crate::error::OpenDalSnafu)?; - } - exists - } else { - false - }; + let reusable = reusable_samples.contains(&id); if reusable { continue; } @@ -314,6 +406,12 @@ async fn build_index_version( unpublished.1.push(handle.clone()); samples.insert(id, handle); stats.built_samples += 1; + } else if manual { + return crate::error::InvalidRequestSnafu { + region_id: region.region_id, + reason: format!("SST {id} has no readable input for samples"), + } + .fail(); } } } @@ -347,17 +445,33 @@ async fn build_index_version( } // Cover SSTs outside planned aggregate builds, including skipped buckets. if enable_range_index { - for file in files { + for file in &files { let file_id = file.file_id().file_id(); if range_indexes.contains_key(&file_id) { continue; } - if let Some(entry) = build_range_index(store, region, version, file).await? { + if let Some(entry) = build_range_index(store, region, version, file.clone()).await? { stats.built_range += 1; range_indexes.insert(entry.file_id, entry); } } } + if manual + && files.iter().any(|file| { + !series_indexes.values().any(|handle| { + handle.entry().format() == selected_format + && handle + .entry() + .covers_file(file.meta_ref(), region.region_id) + }) || (formats.samples.is_some() && !samples.contains_key(&file.file_id().file_id())) + }) + { + return crate::error::InvalidRequestSnafu { + region_id: region.region_id, + reason: "manual series-index build did not cover every SST", + } + .fail(); + } stats.removed_series = current .series_indexes .keys() diff --git a/src/mito2/src/series_index/samples.rs b/src/mito2/src/series_index/samples.rs index 23d34f02900..547d77ad0cc 100644 --- a/src/mito2/src/series_index/samples.rs +++ b/src/mito2/src/series_index/samples.rs @@ -198,19 +198,20 @@ impl Manifest { let mut fields = vec![ Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), Field::new(TSID_COLUMN, DataType::UInt64, false), - Field::new( - &self.timestamp.name, - DataType::Timestamp(self.timestamp.unit.arrow(), None), - false, + crate::sst::with_field_id( + Field::new( + &self.timestamp.name, + DataType::Timestamp(self.timestamp.unit.arrow(), None), + false, + ), + self.timestamp.column_id, ), Field::new(SEQUENCE, DataType::UInt64, false), Field::new(OPERATION, DataType::UInt8, false), ]; - fields.extend( - self.fields - .iter() - .map(|f| Field::new(&f.name, DataType::Float64, true)), - ); + fields.extend(self.fields.iter().map(|f| { + crate::sst::with_field_id(Field::new(&f.name, DataType::Float64, true), f.column_id) + })); Arc::new(Schema::new(fields)) } } diff --git a/src/mito2/src/series_index/task.rs b/src/mito2/src/series_index/task.rs index 6ce340827a1..0aa0dbe9c97 100644 --- a/src/mito2/src/series_index/task.rs +++ b/src/mito2/src/series_index/task.rs @@ -32,7 +32,7 @@ use crate::metrics::{ }; use crate::read::series_candidate::is_sparse_metric_metadata; use crate::region::{MitoRegionRef, RegionLeaderState, RegionMapRef, RegionRoleState}; -use crate::series_index::maintenance::{ReconcileStats, reconcile_series_indexes}; +use crate::series_index::maintenance::{ReconcileStats, reconcile_series_indexes_with_mode}; use crate::series_index::purger::{IndexFilePurger, PurgeRequest, run_index_purge_task}; use crate::time_provider::TimeProviderRef; @@ -234,7 +234,7 @@ impl SeriesIndexTask { async fn reconcile_manual(&mut self, region: MitoRegionRef) -> Result { self.validate_manual_region(®ion)?; self.refresh_usage(); - let result = self.reconcile_region(region.clone()).await; + let result = self.reconcile_region_with_mode(region.clone(), true).await; self.refresh_usage(); let result = result?; self.validate_manual_region(®ion)?; @@ -242,12 +242,20 @@ impl SeriesIndexTask { } async fn reconcile_region(&self, region: MitoRegionRef) -> Result { + self.reconcile_region_with_mode(region, false).await + } + + async fn reconcile_region_with_mode( + &self, + region: MitoRegionRef, + manual: bool, + ) -> Result { // Full capacity defers builds, but cleanup must still reclaim published usage. let allow_builds = self.disk_usage.load(Ordering::Relaxed) < self.max_size; if !allow_builds { SERIES_INDEX_CAPACITY_DEFERRED.inc(); } - let result = reconcile_series_indexes( + let result = reconcile_series_indexes_with_mode( self.worker_id, self.store.clone(), region, @@ -256,6 +264,7 @@ impl SeriesIndexTask { self.purger.clone(), self.enable_range_index, allow_builds, + manual, ) .await; if result.is_err() { @@ -474,7 +483,13 @@ mod tests { let (engine, region) = prepare_region(&mut env).await; let (mut task, _purge_receiver) = new_task(region.clone()); task.max_size = 0; - task.reconcile_manual(region.clone()).await.unwrap(); + assert!( + task.reconcile_manual(region.clone()) + .await + .unwrap_err() + .to_string() + .contains("capacity") + ); assert_eq!(0, region.series_index_version().disk_usage()); assert_eq!(0, task.disk_usage.load(Ordering::Relaxed)); diff --git a/src/mito2/src/series_index/version.rs b/src/mito2/src/series_index/version.rs index d5212e85420..9a830d52aab 100644 --- a/src/mito2/src/series_index/version.rs +++ b/src/mito2/src/series_index/version.rs @@ -166,6 +166,13 @@ pub(crate) struct SeriesIndexVersionControl { } impl SeriesIndexVersionControl { + pub(crate) fn with_format(self, format: crate::config::SeriesIndexFormat) -> Self { + *self.formats.write().unwrap() = Arc::new( + crate::series_index::format::SeriesFormats::from_config(format), + ); + self + } + pub(crate) fn current(&self) -> Arc { self.current.read().unwrap().clone() } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 10a8c3d0011..7d63dec1576 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -578,7 +578,7 @@ pub struct SeriesInfo { /// `Some(metric)` matching each chunk. Ordering describes the stream, not ordinal enumeration. /// /// Preparation, identity lookup, and reader creation do not load sample payloads or perform -/// hidden blocking I/O. The selected set is bound to the partition's immutable MVCC view, +/// hidden blocking I/O. The selected set is bound to the scanner's immutable MVCC view, /// predicates, authorization, projection, labels, and existing scan resources. Readers own /// resources needed to continue after this source or scanner is dropped. Dropping one reader /// does not cancel sibling readers; normal ownership cleanup applies without detached work. @@ -688,20 +688,22 @@ pub trait RegionScanner: Debug + DisplayAs + Send { todo!() } - /// Prepares a finite, stable set of series for independent ordinal-range reads. + /// Prepares scanner-wide candidates once for independent ordinal-range reads. /// - /// `None` means unsupported before consumption, so callers may use the legacy record-batch - /// scan. `Some`, including a source with zero series, selects this path; errors propagate - /// without replay. Choose exactly one entrypoint per partition per execution: this prepared - /// source, [`Self::scan_series_partition`], or [`Self::scan_partition`]. In particular, do - /// not call the legacy raw-scan default merely because this returns `None`. - fn prepare_series_partition( + /// Metadata I/O is asynchronous; sample payloads are loaded only by polling readers. + /// The owned future does not borrow the scanner. `None` means unsupported before + /// execution; errors must propagate without replay through a different scan path. + /// Implementations may share preparation with their raw partition adapter. + fn prepare_series( &self, _ctx: &QueryScanContext, _metrics_set: &ExecutionPlanMetricsSet, - _partition: usize, - ) -> Result>, BoxedError> { - Ok(None) + _label_source: &RawLabelSource, + ) -> futures::future::BoxFuture< + 'static, + Result>, BoxedError>, + > { + Box::pin(async { Ok(None) }) } /// Check if there is any predicate exclude region partition exprs that may be executed in this scanner. @@ -1011,7 +1013,8 @@ mod tests { ); assert!( scanner - .prepare_series_partition(&ctx, &metrics, 0) + .prepare_series(&ctx, &metrics, &RawLabelSource::Explicit) + .await .unwrap() .is_none() ); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index bb9216e6734..493fb9530d3 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3434,6 +3434,7 @@ experimental_manifest_keep_removed_file_count = 256 experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_enable_series_index = false +experimental_series_index_format = "legacy" experimental_series_index_max_size = "5GiB" experimental_enable_range_index = false experimental_series_index_maintenance_interval = "5m"