mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-05 21:18:57 +00:00
feat(mito2): add SST range index searcher (#9003)
* feat(mito2): add SST range index searcher Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(mito2): reuse parquet index reader Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(mito2): simplify range index pruning Signed-off-by: evenyag <realevenyag@gmail.com> * test(mito2): cover missing range index series Signed-off-by: evenyag <realevenyag@gmail.com> --------- Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
Yingwen
@@ -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<u64>]) -> Result<Vec<Bytes>> {
|
||||
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<Arc<ParquetMetaData>> {
|
||||
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<ProjectionMask> {
|
||||
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<ArrayRef> {
|
||||
self.column_values(column, true)
|
||||
}
|
||||
|
||||
fn max_values(&self, column: &Column) -> Option<ArrayRef> {
|
||||
self.column_values(column, false)
|
||||
}
|
||||
|
||||
fn num_containers(&self) -> usize {
|
||||
self.row_groups.len()
|
||||
}
|
||||
|
||||
fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
|
||||
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<ArrayRef> {
|
||||
None
|
||||
}
|
||||
|
||||
fn contained(&self, _column: &Column, _values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl SeriesIndexPruningStats<'_> {
|
||||
fn column_values(&self, column: &Column, is_min: bool) -> Option<ArrayRef> {
|
||||
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<usize> {
|
||||
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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<Self> {
|
||||
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<usize> {
|
||||
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<BoxStream<'static, Result<RecordBatch>>> {
|
||||
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<ProjectionMask> {
|
||||
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<ArrayRef> {
|
||||
self.column_values(column, true)
|
||||
}
|
||||
|
||||
fn max_values(&self, column: &Column) -> Option<ArrayRef> {
|
||||
self.column_values(column, false)
|
||||
}
|
||||
|
||||
fn num_containers(&self) -> usize {
|
||||
self.row_groups.len()
|
||||
}
|
||||
|
||||
fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
|
||||
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<ArrayRef> {
|
||||
None
|
||||
}
|
||||
|
||||
fn contained(&self, _column: &Column, _values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl IndexRowGroupPruningStats<'_> {
|
||||
fn column_values(&self, column: &Column, is_min: bool) -> Option<ArrayRef> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Self> {
|
||||
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<Vec<Range<usize>>> {
|
||||
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<Predicate> {
|
||||
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<Range<usize>>,
|
||||
}
|
||||
|
||||
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<bool> {
|
||||
let row_group_ids = typed_column::<UInt32Array>(batch, ROW_GROUP_ID_COLUMN, "UInt32")?;
|
||||
let table_ids = typed_column::<UInt32Array>(batch, TABLE_ID_COLUMN, "UInt32")?;
|
||||
let tsids = typed_column::<UInt64Array>(batch, TSID_COLUMN, "UInt64")?;
|
||||
let starts = typed_column::<Int64Array>(batch, START_COLUMN, "Int64")?;
|
||||
let ends = typed_column::<Int64Array>(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<Range<usize>> {
|
||||
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::<T>()
|
||||
.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::<Vec<_>>();
|
||||
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]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user