From d7571c127876b7a6c4e8ab65431278eb729865a2 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 3 Sep 2026 08:43:43 +0000 Subject: [PATCH] feat(mito2): add SST range index searcher (#9003) * feat(mito2): add SST range index searcher Signed-off-by: evenyag * refactor(mito2): reuse parquet index reader Signed-off-by: evenyag * refactor(mito2): simplify range index pruning Signed-off-by: evenyag * test(mito2): cover missing range index series Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/src/series_index/searcher.rs | 213 +-------- src/mito2/src/sst/parquet.rs | 1 + src/mito2/src/sst/parquet/index_reader.rs | 190 ++++++++ src/mito2/src/sst/range_index.rs | 2 + src/mito2/src/sst/range_index/searcher.rs | 517 ++++++++++++++++++++++ 5 files changed, 731 insertions(+), 192 deletions(-) create mode 100644 src/mito2/src/sst/parquet/index_reader.rs create mode 100644 src/mito2/src/sst/range_index/searcher.rs diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs index 0f8499b230..e923de9f64 100644 --- a/src/mito2/src/series_index/searcher.rs +++ b/src/mito2/src/series_index/searcher.rs @@ -12,77 +12,33 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; -use std::ops::Range; -use std::sync::Arc; - use api::v1::SemanticType; use async_stream::try_stream; -use bytes::Bytes; use common_recordbatch::filter::SimpleFilterEvaluator; use common_time::range::TimestampRange; -use datafusion_common::pruning::PruningStatistics; -use datafusion_common::{Column, ScalarValue}; use datafusion_expr::{Expr, col, lit}; -use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt32Array, UInt64Array}; +use datatypes::arrow::array::{ArrayRef, UInt32Array, UInt64Array}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use futures::TryStreamExt; use object_store::ObjectStore; -use parquet::DecodeResult; -use parquet::arrow::ProjectionMask; -use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; -use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; -use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::RegionMetadataRef; use table::predicate::Predicate; use crate::error::{ - InvalidMetaSnafu, InvalidRecordBatchSnafu, OpenDalSnafu, ReadParquetSnafu, RecordBatchSnafu, - Result, UnexpectedSnafu, + InvalidMetaSnafu, InvalidRecordBatchSnafu, RecordBatchSnafu, Result, UnexpectedSnafu, }; use crate::series_index::{ MAX_TS_COLUMN, METRIC_SERIES_ID_BATCH_SIZE, MIN_TS_COLUMN, MetricSeriesId, MetricSeriesIdStream, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, series_index_schema, }; -use crate::sst::parquet::format::{column_null_counts, column_values_by_type}; -use crate::sst::parquet::helper::fetch_byte_ranges; -use crate::sst::parquet::metadata::MetadataLoader; +use crate::sst::parquet::index_reader::ParquetIndexReader; use crate::sst::parquet::prefilter::simple_tag_filters; -use crate::sst::parquet::reader::MetadataCacheMetrics; - -#[derive(Clone)] -struct SeriesIndexRangeFetcher { - object_store: ObjectStore, -} - -impl SeriesIndexRangeFetcher { - async fn fetch(&self, path: &str, ranges: &[Range]) -> Result> { - fetch_byte_ranges(path, self.object_store.clone(), ranges) - .await - .context(OpenDalSnafu) - } -} - -#[derive(Clone)] -struct SeriesIndexMetadataProvider { - object_store: ObjectStore, -} - -impl SeriesIndexMetadataProvider { - async fn load(&self, path: &str) -> Result> { - let mut metrics = MetadataCacheMetrics::default(); - MetadataLoader::new(self.object_store.clone(), path, 0) - .load(&mut metrics) - .await - .map(Arc::new) - } -} /// Searches a series-index file for metric series matching query predicates. pub struct SeriesIndexSearcher { - range_fetcher: SeriesIndexRangeFetcher, - metadata_provider: SeriesIndexMetadataProvider, + object_store: ObjectStore, filters: Vec<(Expr, SimpleFilterEvaluator)>, empty_time_range: bool, } @@ -108,10 +64,7 @@ impl SeriesIndexSearcher { } Ok(Self { - range_fetcher: SeriesIndexRangeFetcher { - object_store: object_store.clone(), - }, - metadata_provider: SeriesIndexMetadataProvider { object_store }, + object_store, filters, empty_time_range, }) @@ -123,56 +76,21 @@ impl SeriesIndexSearcher { return Ok(Box::pin(futures::stream::empty())); } - let parquet_metadata = self.metadata_provider.load(path).await?; - let arrow_metadata = - ArrowReaderMetadata::try_new(parquet_metadata, ArrowReaderOptions::new()) - .with_context(|_| ReadParquetSnafu { - path: path.to_string(), - })?; - validate_index_schema(arrow_metadata.schema())?; + let reader = ParquetIndexReader::open(self.object_store.clone(), path).await?; + validate_index_schema(reader.schema())?; // An older index file may not contain tags added by schema evolution. // Ignore filters on those tags to preserve a conservative candidate set. - let (pruning_predicate, filters) = self.filters_for_schema(arrow_metadata.schema()); - let row_groups = row_groups_to_read( - arrow_metadata.metadata().row_groups(), - arrow_metadata.schema().clone(), - &pruning_predicate, - ); - let projection = projection_mask( - arrow_metadata.parquet_schema(), - arrow_metadata.schema(), - &filters, - )?; - let mut decoder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata) - .with_row_groups(row_groups) - .with_projection(projection) - .build() - .with_context(|_| ReadParquetSnafu { - path: path.to_string(), - })?; - let path = path.to_string(); - let range_fetcher = self.range_fetcher.clone(); + let (pruning_predicate, filters) = self.filters_for_schema(reader.schema()); + let mut projection_columns = Vec::with_capacity(filters.len() + 2); + projection_columns.extend([TABLE_ID_COLUMN, TSID_COLUMN]); + projection_columns.extend(filters.iter().map(SimpleFilterEvaluator::column_name)); + let mut batches = reader.read(&pruning_predicate, &projection_columns)?; Ok(Box::pin(try_stream! { let mut last_series = None; let mut output = Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE); - loop { - let batch = match decoder - .try_decode() - .with_context(|_| ReadParquetSnafu { path: path.clone() })? - { - DecodeResult::NeedsData(ranges) => { - let data = range_fetcher.fetch(&path, &ranges).await?; - decoder - .push_ranges(ranges, data) - .with_context(|_| ReadParquetSnafu { path: path.clone() })?; - continue; - } - DecodeResult::Data(batch) => batch, - DecodeResult::Finished => break, - }; - + while let Some(batch) = batches.try_next().await? { let mut mask = BooleanBuffer::new_set(batch.num_rows()); for filter in &filters { let column = column(&batch, filter.column_name())?; @@ -307,36 +225,6 @@ fn validate_index_schema(schema: &SchemaRef) -> Result<()> { Ok(()) } -fn projection_mask( - parquet_schema: &parquet::schema::types::SchemaDescriptor, - arrow_schema: &SchemaRef, - filters: &[SimpleFilterEvaluator], -) -> Result { - let mut indices = HashSet::new(); - for name in [TABLE_ID_COLUMN, TSID_COLUMN] { - let index = arrow_schema - .index_of(name) - .ok() - .with_context(|| InvalidRecordBatchSnafu { - reason: format!("series index is missing internal column {name}"), - })?; - indices.insert(index); - } - for filter in filters { - let index = arrow_schema - .index_of(filter.column_name()) - .ok() - .with_context(|| InvalidRecordBatchSnafu { - reason: format!( - "series index is missing predicate column {}", - filter.column_name() - ), - })?; - indices.insert(index); - } - Ok(ProjectionMask::roots(parquet_schema, indices)) -} - fn column<'a>( batch: &'a datatypes::arrow::record_batch::RecordBatch, name: &str, @@ -351,62 +239,10 @@ fn column<'a>( Ok(batch.column(index)) } -struct SeriesIndexPruningStats<'a> { - row_groups: &'a [RowGroupMetaData], - schema: SchemaRef, -} - -impl PruningStatistics for SeriesIndexPruningStats<'_> { - fn min_values(&self, column: &Column) -> Option { - self.column_values(column, true) - } - - fn max_values(&self, column: &Column) -> Option { - self.column_values(column, false) - } - - fn num_containers(&self) -> usize { - self.row_groups.len() - } - - fn null_counts(&self, column: &Column) -> Option { - let column_index = self.schema.index_of(&column.name).ok()?; - column_null_counts(self.row_groups, column_index) - } - - fn row_counts(&self, _column: &Column) -> Option { - None - } - - fn contained(&self, _column: &Column, _values: &HashSet) -> Option { - None - } -} - -impl SeriesIndexPruningStats<'_> { - fn column_values(&self, column: &Column, is_min: bool) -> Option { - let column_index = self.schema.index_of(&column.name).ok()?; - let data_type = self.schema.field(column_index).data_type(); - column_values_by_type(self.row_groups, data_type, column_index, is_min) - } -} - -fn row_groups_to_read( - row_groups: &[RowGroupMetaData], - schema: SchemaRef, - predicate: &Predicate, -) -> Vec { - let stats = SeriesIndexPruningStats { row_groups, schema }; - predicate - .prune_with_stats(&stats, &stats.schema) - .into_iter() - .enumerate() - .filter_map(|(row_group, keep)| keep.then_some(row_group)) - .collect() -} - #[cfg(test)] mod tests { + use std::sync::Arc; + use datafusion_expr::{col, lit}; use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt8Array}; use datatypes::arrow::datatypes::{Field, Schema}; @@ -705,18 +541,11 @@ mod tests { None, ) .unwrap(); - let parquet_metadata = searcher.metadata_provider.load(path).await.unwrap(); - let arrow_metadata = - ArrowReaderMetadata::try_new(parquet_metadata, ArrowReaderOptions::new()).unwrap(); - let (pruning_predicate, _) = searcher.filters_for_schema(arrow_metadata.schema()); - assert_eq!( - row_groups_to_read( - arrow_metadata.metadata().row_groups(), - arrow_metadata.schema().clone(), - &pruning_predicate, - ), - vec![1] - ); + let reader = ParquetIndexReader::open(object_store.clone(), path) + .await + .unwrap(); + let (pruning_predicate, _) = searcher.filters_for_schema(reader.schema()); + assert_eq!(reader.row_groups_to_read(&pruning_predicate), vec![1]); let empty = SeriesIndexSearcher::try_new( metadata, diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 0b563aea9a..3d90a36092 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -30,6 +30,7 @@ pub mod file_range; pub mod flat_format; pub mod format; pub(crate) mod helper; +pub(crate) mod index_reader; pub(crate) mod index_writer; pub(crate) mod json_align; pub mod metadata; diff --git a/src/mito2/src/sst/parquet/index_reader.rs b/src/mito2/src/sst/parquet/index_reader.rs new file mode 100644 index 0000000000..bd894f207a --- /dev/null +++ b/src/mito2/src/sst/parquet/index_reader.rs @@ -0,0 +1,190 @@ +// 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. + +//! Reader for standalone Parquet index files. + +use std::collections::HashSet; +use std::sync::Arc; + +use datafusion_common::pruning::PruningStatistics; +use datafusion_common::{Column, ScalarValue}; +use datatypes::arrow::array::{ArrayRef, BooleanArray}; +use datatypes::arrow::datatypes::SchemaRef; +use datatypes::arrow::record_batch::RecordBatch; +use futures::StreamExt; +use futures::stream::BoxStream; +use object_store::ObjectStore; +use parquet::DecodeResult; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +use parquet::file::metadata::RowGroupMetaData; +use snafu::{OptionExt, ResultExt}; +use table::predicate::Predicate; + +use crate::error::{InvalidRecordBatchSnafu, OpenDalSnafu, ReadParquetSnafu, Result}; +use crate::sst::parquet::format::{column_null_counts, column_values_by_type}; +use crate::sst::parquet::helper::fetch_byte_ranges; +use crate::sst::parquet::metadata::MetadataLoader; +use crate::sst::parquet::reader::MetadataCacheMetrics; + +/// Reads a standalone index file stored in Parquet format. +pub(crate) struct ParquetIndexReader { + object_store: ObjectStore, + path: String, + arrow_metadata: ArrowReaderMetadata, +} + +impl ParquetIndexReader { + /// Opens `path` and loads its Parquet metadata. + pub(crate) async fn open(object_store: ObjectStore, path: &str) -> Result { + let mut metrics = MetadataCacheMetrics::default(); + let parquet_metadata = MetadataLoader::new(object_store.clone(), path, 0) + .load(&mut metrics) + .await?; + let arrow_metadata = + ArrowReaderMetadata::try_new(Arc::new(parquet_metadata), ArrowReaderOptions::new()) + .with_context(|_| ReadParquetSnafu { + path: path.to_string(), + })?; + + Ok(Self { + object_store, + path: path.to_string(), + arrow_metadata, + }) + } + + /// Returns the Arrow schema of the index file. + pub(crate) fn schema(&self) -> &SchemaRef { + self.arrow_metadata.schema() + } + + /// Returns row groups that may match `predicate`. + pub(crate) fn row_groups_to_read(&self, predicate: &Predicate) -> Vec { + let stats = IndexRowGroupPruningStats { + row_groups: self.arrow_metadata.metadata().row_groups(), + schema: self.arrow_metadata.schema(), + }; + predicate + .prune_with_stats(&stats, stats.schema) + .into_iter() + .enumerate() + .filter_map(|(row_group, keep)| keep.then_some(row_group)) + .collect() + } + + /// Returns a stream of projected batches from row groups matching `predicate`. + pub(crate) fn read( + &self, + predicate: &Predicate, + projection_columns: &[&str], + ) -> Result>> { + let projection = self.projection_mask(projection_columns)?; + let row_groups = self.row_groups_to_read(predicate); + if row_groups.is_empty() { + return Ok(futures::stream::empty().boxed()); + } + + let mut decoder = ParquetPushDecoderBuilder::new_with_metadata(self.arrow_metadata.clone()) + .with_row_groups(row_groups) + .with_projection(projection) + .build() + .with_context(|_| ReadParquetSnafu { + path: self.path.clone(), + })?; + let path = self.path.clone(); + let object_store = self.object_store.clone(); + + Ok(async_stream::try_stream! { + loop { + match decoder + .try_decode() + .with_context(|_| ReadParquetSnafu { path: path.clone() })? + { + DecodeResult::NeedsData(ranges) => { + let data = fetch_byte_ranges(&path, object_store.clone(), &ranges) + .await + .context(OpenDalSnafu)?; + decoder + .push_ranges(ranges, data) + .with_context(|_| ReadParquetSnafu { path: path.clone() })?; + } + DecodeResult::Data(batch) => yield batch, + DecodeResult::Finished => break, + } + } + } + .boxed()) + } + + fn projection_mask(&self, projection_columns: &[&str]) -> Result { + let mut indices = HashSet::with_capacity(projection_columns.len()); + for name in projection_columns { + let index = self + .arrow_metadata + .schema() + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("Parquet index is missing projected column {name}"), + })?; + indices.insert(index); + } + Ok(ProjectionMask::roots( + self.arrow_metadata.parquet_schema(), + indices, + )) + } +} + +struct IndexRowGroupPruningStats<'a> { + row_groups: &'a [RowGroupMetaData], + schema: &'a SchemaRef, +} + +impl PruningStatistics for IndexRowGroupPruningStats<'_> { + fn min_values(&self, column: &Column) -> Option { + self.column_values(column, true) + } + + fn max_values(&self, column: &Column) -> Option { + self.column_values(column, false) + } + + fn num_containers(&self) -> usize { + self.row_groups.len() + } + + fn null_counts(&self, column: &Column) -> Option { + let column_index = self.schema.index_of(&column.name).ok()?; + column_null_counts(self.row_groups, column_index) + } + + fn row_counts(&self, _column: &Column) -> Option { + None + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + None + } +} + +impl IndexRowGroupPruningStats<'_> { + fn column_values(&self, column: &Column, is_min: bool) -> Option { + let column_index = self.schema.index_of(&column.name).ok()?; + let data_type = self.schema.field(column_index).data_type(); + column_values_by_type(self.row_groups, data_type, column_index, is_min) + } +} diff --git a/src/mito2/src/sst/range_index.rs b/src/mito2/src/sst/range_index.rs index 2fa8142a23..fa766acacb 100644 --- a/src/mito2/src/sst/range_index.rs +++ b/src/mito2/src/sst/range_index.rs @@ -14,8 +14,10 @@ //! Per-SST series row-range index. +mod searcher; mod writer; +pub use searcher::SstRangeIndexSearcher; use store_api::metric_engine_consts::{ DATA_SCHEMA_TABLE_ID_COLUMN_NAME as TABLE_ID_COLUMN, DATA_SCHEMA_TSID_COLUMN_NAME as TSID_COLUMN, diff --git a/src/mito2/src/sst/range_index/searcher.rs b/src/mito2/src/sst/range_index/searcher.rs new file mode 100644 index 0000000000..b86c68c887 --- /dev/null +++ b/src/mito2/src/sst/range_index/searcher.rs @@ -0,0 +1,517 @@ +// 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::cmp::Ordering; +use std::ops::Range; + +use datafusion_expr::{col, lit}; +use datatypes::arrow::array::{Int64Array, UInt32Array, UInt64Array}; +use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use futures::TryStreamExt; +use object_store::ObjectStore; +use snafu::{OptionExt, ensure}; +use table::predicate::Predicate; + +use crate::error::{InvalidRecordBatchSnafu, Result, UnexpectedSnafu}; +use crate::series_index::MetricSeriesId; +use crate::sst::parquet::index_reader::ParquetIndexReader; +use crate::sst::range_index::{ + END_COLUMN, ROW_GROUP_ID_COLUMN, START_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, +}; + +/// Searches per-SST range-index files for the rows of candidate metric series. +pub struct SstRangeIndexSearcher { + reader: ParquetIndexReader, +} + +impl SstRangeIndexSearcher { + /// Opens the range-index file at `path` and loads its Parquet metadata. + pub async fn open(object_store: ObjectStore, path: &str) -> Result { + let reader = ParquetIndexReader::open(object_store, path).await?; + validate_index_schema(reader.schema())?; + Ok(Self { reader }) + } + + /// Returns the row ranges for `series` in one source SST row group. + /// + /// `series` is one batch emitted by a + /// [`MetricSeriesIdStream`](crate::series_index::MetricSeriesIdStream). The + /// returned half-open ranges are relative to the start of `row_group_id`, + /// sorted, non-overlapping, and coalesced when adjacent. The number of + /// returned ranges may be less than the number of input series if some + /// series don't exist in the row group. + pub async fn search( + &self, + row_group_id: u32, + series: &[MetricSeriesId], + ) -> Result>> { + if series.is_empty() { + return Ok(Vec::new()); + } + + validate_sorted_series(series)?; + let predicate = search_predicate(row_group_id, series)?; + let mut batches = self.reader.read( + &predicate, + &[ + ROW_GROUP_ID_COLUMN, + TABLE_ID_COLUMN, + TSID_COLUMN, + START_COLUMN, + END_COLUMN, + ], + )?; + let mut merge = RangeMergeState::new(row_group_id, series); + + while let Some(batch) = batches.try_next().await? { + if merge.append_batch(&batch)? { + break; + } + } + + Ok(merge.finish()) + } +} + +fn validate_sorted_series(series: &[MetricSeriesId]) -> Result<()> { + if let Some(pair) = series.windows(2).find(|pair| pair[0] > pair[1]) { + return InvalidRecordBatchSnafu { + reason: format!( + "range index search series are not sorted: {:?} appears before {:?}", + pair[0], pair[1] + ), + } + .fail(); + } + Ok(()) +} + +fn search_predicate(row_group_id: u32, series: &[MetricSeriesId]) -> Result { + let min_table_id = series + .first() + .context(UnexpectedSnafu { + reason: "cannot build a range-index predicate for an empty series set", + })? + .table_id; + let max_table_id = series + .last() + .context(UnexpectedSnafu { + reason: "cannot build a range-index predicate for an empty series set", + })? + .table_id; + + Ok(Predicate::new(vec![ + col(ROW_GROUP_ID_COLUMN).eq(lit(row_group_id)), + col(TABLE_ID_COLUMN).gt_eq(lit(min_table_id)), + col(TABLE_ID_COLUMN).lt_eq(lit(max_table_id)), + ])) +} + +fn validate_index_schema(schema: &SchemaRef) -> Result<()> { + for (name, data_type) in [ + (ROW_GROUP_ID_COLUMN, DataType::UInt32), + (TABLE_ID_COLUMN, DataType::UInt32), + (TSID_COLUMN, DataType::UInt64), + (START_COLUMN, DataType::Int64), + (END_COLUMN, DataType::Int64), + ] { + let field = schema + .field_with_name(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("range index is missing column {name}"), + })?; + ensure!( + field.data_type() == &data_type && !field.is_nullable(), + InvalidRecordBatchSnafu { + reason: format!( + "range index column {name} must be non-nullable {data_type:?}, got {:?}", + field.data_type() + ), + } + ); + } + Ok(()) +} + +struct RangeMergeState<'a> { + /// Source SST row group whose ranges are being searched. + row_group_id: u32, + /// Sorted metric series to match against the range index. + series: &'a [MetricSeriesId], + /// Cursor to the next series to match. + series_index: usize, + /// Last range-index key read, used to validate ordering across batches. + last_index_key: Option<(u32, MetricSeriesId)>, + /// Matching row ranges, sorted and coalesced when adjacent. + ranges: Vec>, +} + +impl<'a> RangeMergeState<'a> { + fn new(row_group_id: u32, series: &'a [MetricSeriesId]) -> Self { + Self { + row_group_id, + series, + series_index: 0, + last_index_key: None, + ranges: Vec::new(), + } + } + + /// Appends matches from `batch` and returns whether the merge is complete. + fn append_batch( + &mut self, + batch: &datatypes::arrow::record_batch::RecordBatch, + ) -> Result { + let row_group_ids = typed_column::(batch, ROW_GROUP_ID_COLUMN, "UInt32")?; + let table_ids = typed_column::(batch, TABLE_ID_COLUMN, "UInt32")?; + let tsids = typed_column::(batch, TSID_COLUMN, "UInt64")?; + let starts = typed_column::(batch, START_COLUMN, "Int64")?; + let ends = typed_column::(batch, END_COLUMN, "Int64")?; + + for row in 0..batch.num_rows() { + let index_series = MetricSeriesId { + table_id: table_ids.value(row), + tsid: tsids.value(row), + }; + let index_key = (row_group_ids.value(row), index_series); + ensure!( + self.last_index_key.is_none_or(|last| last < index_key), + InvalidRecordBatchSnafu { + reason: format!( + "range index rows are not strictly sorted: {index_key:?} follows {:?}", + self.last_index_key + ), + } + ); + self.last_index_key = Some(index_key); + + match index_key.0.cmp(&self.row_group_id) { + Ordering::Less => continue, + Ordering::Greater => return Ok(true), + Ordering::Equal => {} + } + + while self.series_index < self.series.len() + && self.series[self.series_index] < index_series + { + self.advance_series(); + } + if self.series_index == self.series.len() { + return Ok(true); + } + + match self.series[self.series_index].cmp(&index_series) { + Ordering::Less => { + return UnexpectedSnafu { + reason: "range-index merge cursor did not advance past a smaller series", + } + .fail(); + } + Ordering::Greater => continue, + Ordering::Equal => { + self.append_range(starts.value(row), ends.value(row), row)?; + self.advance_series(); + if self.series_index == self.series.len() { + return Ok(true); + } + } + } + } + Ok(false) + } + + fn advance_series(&mut self) { + let current = self.series[self.series_index]; + while self.series_index < self.series.len() && self.series[self.series_index] == current { + self.series_index += 1; + } + } + + fn append_range(&mut self, start: i64, end: i64, row: usize) -> Result<()> { + let start = usize::try_from(start).map_err(|_| { + InvalidRecordBatchSnafu { + reason: format!("range index contains negative start offset at row {row}"), + } + .build() + })?; + let end = usize::try_from(end).map_err(|_| { + InvalidRecordBatchSnafu { + reason: format!("range index contains negative end offset at row {row}"), + } + .build() + })?; + ensure!( + start < end, + InvalidRecordBatchSnafu { + reason: format!("range index contains invalid range {start}..{end} at row {row}"), + } + ); + + if let Some(last) = self.ranges.last_mut() { + ensure!( + start >= last.end, + InvalidRecordBatchSnafu { + reason: format!( + "range index contains overlapping or unsorted range {start}..{end} after {}..{}", + last.start, last.end + ), + } + ); + if start == last.end { + last.end = end; + return Ok(()); + } + } + self.ranges.push(start..end); + Ok(()) + } + + fn finish(self) -> Vec> { + self.ranges + } +} + +fn typed_column<'a, T: 'static>( + batch: &'a datatypes::arrow::record_batch::RecordBatch, + name: &str, + data_type: &str, +) -> Result<&'a T> { + let index = batch + .schema() + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("range index batch is missing column {name}"), + })?; + batch + .column(index) + .as_any() + .downcast_ref::() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("range index column {name} is not {data_type}"), + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::arrow::array::{ArrayRef, BinaryArray}; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + use object_store::services::Memory; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::RegionMetadataRef; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + + use super::*; + use crate::sst::range_index::{ + SstRangeIndexWriter, SstRangeIndexWriterOptions, range_index_schema, + }; + use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; + + fn object_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap().finish() + } + + fn series(table_id: u32, tsid: u64) -> MetricSeriesId { + MetricSeriesId { table_id, tsid } + } + + fn primary_key_batch(metadata: &RegionMetadataRef, ids: &[(u32, u64)]) -> RecordBatch { + let primary_keys = ids + .iter() + .map(|(table_id, tsid)| new_sparse_primary_key(&["a", "x"], metadata, *table_id, *tsid)) + .collect::>(); + let schema = Arc::new(Schema::new(vec![Field::new( + PRIMARY_KEY_COLUMN_NAME, + DataType::Binary, + false, + )])); + RecordBatch::try_new( + schema, + vec![Arc::new(BinaryArray::from_iter_values( + primary_keys.iter().map(Vec::as_slice), + ))], + ) + .unwrap() + } + + async fn write_index(store: &ObjectStore, path: &str) { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mut writer = SstRangeIndexWriter::try_new( + metadata.clone(), + store.clone(), + path, + SstRangeIndexWriterOptions { + index_row_group_size: 2, + }, + ) + .await + .unwrap(); + writer + .write( + 0, + &primary_key_batch( + &metadata, + &[(1, 10), (1, 10), (1, 20), (2, 10), (2, 20), (2, 20)], + ), + ) + .await + .unwrap(); + writer + .write(1, &primary_key_batch(&metadata, &[(2, 20), (2, 20)])) + .await + .unwrap(); + writer.finish().await.unwrap(); + } + + #[tokio::test] + async fn search_filters_exact_series_pairs_and_coalesces_ranges() { + let store = object_store(); + let path = "range-search.parquet"; + write_index(&store, path).await; + let searcher = SstRangeIndexSearcher::open(store, path).await.unwrap(); + + let ranges = searcher + .search(0, &[series(1, 10), series(2, 20)]) + .await + .unwrap(); + assert_eq!(ranges, vec![0..2, 4..6]); + + let ranges = searcher + .search(0, &[series(1, 10), series(1, 20)]) + .await + .unwrap(); + assert_eq!(ranges, vec![0..3]); + + let ranges = searcher + .search(0, &[series(1, 15), series(2, 20)]) + .await + .unwrap(); + assert_eq!(ranges, vec![4..6]); + + let ranges = searcher + .search(0, &[series(1, 10), series(2, 30)]) + .await + .unwrap(); + assert_eq!(ranges, vec![0..2]); + + let ranges = searcher + .search(1, &[series(2, 20), series(2, 20)]) + .await + .unwrap(); + assert_eq!(ranges, vec![0..2]); + + assert!( + searcher + .search(1, &[series(1, 10)]) + .await + .unwrap() + .is_empty() + ); + + assert!(searcher.search(0, &[]).await.unwrap().is_empty()); + + let error = searcher + .search(0, &[series(2, 20), series(1, 10)]) + .await + .unwrap_err(); + assert!(error.to_string().contains("not sorted"), "{error}"); + } + + #[tokio::test] + async fn opening_a_missing_index_fails() { + assert!( + SstRangeIndexSearcher::open(object_store(), "does-not-exist.parquet") + .await + .is_err() + ); + } + + #[tokio::test] + async fn pruning_uses_the_source_row_group_and_table_id_range() { + let store = object_store(); + let path = "range-pruning.parquet"; + write_index(&store, path).await; + let reader = ParquetIndexReader::open(store, path).await.unwrap(); + let predicate = search_predicate(0, &[series(1, 999), series(2, 999)]).unwrap(); + + assert_eq!(reader.row_groups_to_read(&predicate), vec![0, 1]); + } + + #[test] + fn validates_schema_and_range_offsets() { + let nullable_schema = Arc::new(Schema::new(vec![ + Field::new(ROW_GROUP_ID_COLUMN, DataType::UInt32, false), + Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), + Field::new(TSID_COLUMN, DataType::UInt64, false), + Field::new(START_COLUMN, DataType::Int64, true), + Field::new(END_COLUMN, DataType::Int64, false), + ])); + assert!(validate_index_schema(&nullable_schema).is_err()); + + let batch = RecordBatch::try_new( + range_index_schema(), + vec![ + Arc::new(UInt32Array::from(vec![0])) as ArrayRef, + Arc::new(UInt32Array::from(vec![1])), + Arc::new(UInt64Array::from(vec![10])), + Arc::new(Int64Array::from(vec![-1])), + Arc::new(Int64Array::from(vec![2])), + ], + ) + .unwrap(); + let selected = [series(1, 10)]; + let mut merge = RangeMergeState::new(0, &selected); + assert!(merge.append_batch(&batch).is_err()); + + let unsorted_batch = RecordBatch::try_new( + range_index_schema(), + vec![ + Arc::new(UInt32Array::from(vec![0, 0])) as ArrayRef, + Arc::new(UInt32Array::from(vec![1, 1])), + Arc::new(UInt64Array::from(vec![20, 10])), + Arc::new(Int64Array::from(vec![0, 1])), + Arc::new(Int64Array::from(vec![1, 2])), + ], + ) + .unwrap(); + let selected = [series(1, 20), series(1, 30)]; + let mut merge = RangeMergeState::new(0, &selected); + assert!(merge.append_batch(&unsorted_batch).is_err()); + + let make_batch = |tsid, start, end| { + RecordBatch::try_new( + range_index_schema(), + vec![ + Arc::new(UInt32Array::from(vec![0])) as ArrayRef, + Arc::new(UInt32Array::from(vec![1])), + Arc::new(UInt64Array::from(vec![tsid])), + Arc::new(Int64Array::from(vec![start])), + Arc::new(Int64Array::from(vec![end])), + ], + ) + .unwrap() + }; + let selected = [series(1, 10), series(1, 20)]; + let mut merge = RangeMergeState::new(0, &selected); + assert!(!merge.append_batch(&make_batch(10, 0, 1)).unwrap()); + assert!(merge.append_batch(&make_batch(20, 1, 2)).unwrap()); + assert_eq!(merge.finish(), vec![0..2]); + } +}