diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index e68ca58609..20ca9b53ca 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Incremental writer for aggregate series index files | +| `series_index` | `src/mito2/src/series_index/` | Incremental writer and predicate searcher for aggregate series index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 2933327a38..73591d4ffa 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -30,7 +30,6 @@ use datatypes::arrow::compute::SortOptions; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; -use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec}; use snafu::{OptionExt, ResultExt, ensure}; @@ -51,6 +50,7 @@ use crate::read::range_cache::{ }; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges}; +use crate::series_index::{METRIC_SERIES_ID_BATCH_SIZE, MetricSeriesId, MetricSeriesIdStream}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::format::PrimaryKeyArray; use crate::sst::parquet::prefilter::{ @@ -59,17 +59,6 @@ use crate::sst::parquet::prefilter::{ use crate::sst::parquet::reader::ReaderMetrics; use crate::sst::parquet::row_group::ParquetFetchMetrics; -const CANDIDATE_SERIES_BATCH_SIZE: usize = 500; - -/// Identifies one series in a physical metric region. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct MetricSeriesId { - pub(crate) table_id: u32, - pub(crate) tsid: u64, -} - -pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result>>; - /// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan). pub(crate) struct SeriesCandidateScanner { stream_ctx: Arc, @@ -537,7 +526,7 @@ fn decode_metric_series( let codec = SparsePrimaryKeyCodec::new(&metadata); Ok(Box::pin(try_stream! { let mut last_series = None; - let mut output = Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE); + let mut output = Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE); while let Some(batch) = input.try_next().await? { let array = batch .column(0) @@ -556,10 +545,10 @@ fn decode_metric_series( } last_series = Some(series); output.push(series); - if output.len() == CANDIDATE_SERIES_BATCH_SIZE { + if output.len() == METRIC_SERIES_ID_BATCH_SIZE { yield std::mem::replace( &mut output, - Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE), + Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE), ); } } diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 0a3e008f27..15e4792432 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -40,7 +40,8 @@ use crate::read::scan_util::{ should_split_flat_batches_for_merge, }; use crate::read::seq_scan::SeqScan; -use crate::read::series_candidate::{MetricSeriesId, validate_metric_metadata}; +use crate::read::series_candidate::validate_metric_metadata; +use crate::series_index::MetricSeriesId; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::prefilter::prefilter_flat_batch_by_primary_key; diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 21ced00cf6..e35769bec4 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -12,10 +12,34 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Series index writer. +//! Series index writer and searcher. +mod searcher; mod writer; +use futures::stream::BoxStream; +pub use searcher::SeriesIndexSearcher; pub use writer::{ SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema, }; + +use crate::error::Result; + +pub(crate) const MIN_TS_COLUMN: &str = "__series_min_ts"; +pub(crate) const MAX_TS_COLUMN: &str = "__series_max_ts"; +pub(crate) const ROW_COUNT_COLUMN: &str = "__series_row_count"; +pub(crate) const TABLE_ID_COLUMN: &str = "__table_id"; +pub(crate) const TSID_COLUMN: &str = "__tsid"; +pub(crate) const METRIC_SERIES_ID_BATCH_SIZE: usize = 500; + +/// Identifies one series in a physical metric region. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MetricSeriesId { + /// Logical table ID inside the physical metric region. + pub table_id: u32, + /// Time-series ID inside the logical table. + pub tsid: u64, +} + +/// Stream of bounded batches of matching metric-series IDs. +pub type MetricSeriesIdStream = BoxStream<'static, Result>>; diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs new file mode 100644 index 0000000000..a303bcaebd --- /dev/null +++ b/src/mito2/src/series_index/searcher.rs @@ -0,0 +1,734 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::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::buffer::BooleanBuffer; +use datatypes::arrow::datatypes::{DataType, SchemaRef}; +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, +}; +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::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, + filters: Vec<(Expr, SimpleFilterEvaluator)>, + empty_time_range: bool, +} + +impl SeriesIndexSearcher { + /// Creates a searcher reusable across series-index files of `metadata`. + pub fn try_new( + metadata: RegionMetadataRef, + object_store: ObjectStore, + predicate: Option<&Predicate>, + time_range: Option, + ) -> Result { + // Keep search-time metadata validation identical to the writer. + series_index_schema(&metadata)?; + + let mut filters = simple_tag_filters(&metadata, None, predicate); + let (empty_time_range, time_exprs) = time_range_filters(&metadata, time_range)?; + for expr in time_exprs { + let filter = SimpleFilterEvaluator::try_new(&expr).context(UnexpectedSnafu { + reason: "failed to build an internal series-index time filter", + })?; + filters.push((expr, filter)); + } + + Ok(Self { + range_fetcher: SeriesIndexRangeFetcher { + object_store: object_store.clone(), + }, + metadata_provider: SeriesIndexMetadataProvider { object_store }, + filters, + empty_time_range, + }) + } + + /// Searches `path` and returns sorted batches of matching metric-series IDs. + pub async fn search(&self, path: &str) -> Result { + if self.empty_time_range { + 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())?; + + // 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(); + + 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, + }; + + let mut mask = BooleanBuffer::new_set(batch.num_rows()); + for filter in &filters { + let column = column(&batch, filter.column_name())?; + let evaluated = filter.evaluate_array(column).context(RecordBatchSnafu)?; + mask = &mask & &evaluated; + } + + let table_ids = column(&batch, TABLE_ID_COLUMN)? + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "series index __table_id is not UInt32", + })?; + let tsids = column(&batch, TSID_COLUMN)? + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "series index __tsid is not UInt64", + })?; + + for (row, matched) in mask.iter().enumerate() { + if !matched { + continue; + } + let series = MetricSeriesId { + table_id: table_ids.value(row), + tsid: tsids.value(row), + }; + if last_series == Some(series) { + continue; + } + last_series = Some(series); + output.push(series); + if output.len() == METRIC_SERIES_ID_BATCH_SIZE { + yield std::mem::replace( + &mut output, + Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE), + ); + } + } + } + if !output.is_empty() { + yield output; + } + })) + } + + fn filters_for_schema(&self, schema: &SchemaRef) -> (Predicate, Vec) { + let (exprs, filters): (Vec<_>, Vec<_>) = self + .filters + .iter() + .filter(|(_, filter)| schema.field_with_name(filter.column_name()).is_ok()) + .cloned() + .unzip(); + (Predicate::new(exprs), filters) + } +} + +fn time_range_filters( + metadata: &RegionMetadataRef, + time_range: Option, +) -> Result<(bool, Vec)> { + let Some(time_range) = time_range else { + return Ok((false, Vec::new())); + }; + if time_range.is_empty() { + return Ok((true, Vec::new())); + } + + let time_index = metadata.time_index_column(); + ensure!( + time_index.semantic_type == SemanticType::Timestamp, + InvalidMetaSnafu { + reason: "series index metadata has no timestamp time index", + } + ); + let timestamp_type = + time_index + .column_schema + .data_type + .as_timestamp() + .context(InvalidMetaSnafu { + reason: "series index time index is not a timestamp", + })?; + let unit = timestamp_type.unit(); + let mut exprs = Vec::with_capacity(2); + // A series overlaps [start, end) only if its maximum is at least start. + // Round start up so a series ending before an unaligned start is pruned. + if let Some(start) = time_range + .start() + .and_then(|start| start.convert_to_ceil(unit)) + { + exprs.push(col(MAX_TS_COLUMN).gt_eq(lit(start.value()))); + } + // A series overlaps [start, end) only if its minimum is less than end. + // Round the exclusive end up to avoid pruning the containing unit interval. + if let Some(end) = time_range.end().and_then(|end| end.convert_to_ceil(unit)) { + exprs.push(col(MIN_TS_COLUMN).lt(lit(end.value()))); + } + Ok((false, exprs)) +} + +fn validate_index_schema(schema: &SchemaRef) -> Result<()> { + for (name, data_type) in [ + (MIN_TS_COLUMN, DataType::Int64), + (MAX_TS_COLUMN, DataType::Int64), + (ROW_COUNT_COLUMN, DataType::UInt64), + (TABLE_ID_COLUMN, DataType::UInt32), + (TSID_COLUMN, DataType::UInt64), + ] { + let field = schema + .field_with_name(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index is missing internal column {name}"), + })?; + ensure!( + field.data_type() == &data_type && !field.is_nullable(), + InvalidRecordBatchSnafu { + reason: format!( + "series index internal column {name} must be non-nullable {data_type:?}, got {:?}", + field.data_type() + ), + } + ); + } + 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, +) -> Result<&'a ArrayRef> { + let index = batch + .schema() + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index batch is missing column {name}"), + })?; + 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 datafusion_expr::{col, lit}; + use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt8Array}; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use futures::TryStreamExt; + use object_store::services::Memory; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + + use super::*; + use crate::series_index::{SeriesIndexWriter, SeriesIndexWriterOptions}; + 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 flat_batch(primary_keys: &[Vec], timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("__primary_key", DataType::Binary, false), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(BinaryArray::from_iter_values( + primary_keys.iter().map(Vec::as_slice), + )), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + async fn write_index( + metadata: RegionMetadataRef, + object_store: ObjectStore, + path: &str, + rows: &[(u32, u64, &str, &str, i64)], + row_group_size: usize, + ) { + let primary_keys = rows + .iter() + .map(|(table_id, tsid, tag_0, tag_1, _)| { + new_sparse_primary_key(&[*tag_0, *tag_1], &metadata, *table_id, *tsid) + }) + .collect::>(); + let timestamps = rows.iter().map(|row| row.4).collect::>(); + let mut writer = SeriesIndexWriter::try_new( + metadata, + object_store, + path, + SeriesIndexWriterOptions { row_group_size }, + ) + .await + .unwrap(); + writer + .write(&flat_batch(&primary_keys, ×tamps)) + .await + .unwrap(); + writer.finish().await.unwrap(); + } + + async fn collect_ids(stream: MetricSeriesIdStream) -> Vec { + stream + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect() + } + + #[tokio::test] + async fn search_applies_candidate_tag_filters_and_time_overlap() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "search.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "a", "y", 30), + ], + 2, + ) + .await; + + // The field filter is not available in the series index and is ignored, + // matching candidate-primary-key filter behavior. + let predicate = Predicate::new(vec![ + col("tag_0").eq(lit("a")), + col("field_0").gt(lit(0_u64)), + ]); + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(31), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + Some(&predicate), + Some(time_range), + ) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + + // Both bounds fall between millisecond ticks. Rounding the inclusive + // start and exclusive end upward leaves only the 30 ms series. + let time_range = TimestampRange::new( + common_time::Timestamp::new_microsecond(20_001), + common_time::Timestamp::new_microsecond(30_001), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + None, + Some(time_range), + ) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + + // The stored maximum equal to the query start intersects, while the + // stored minimum equal to the exclusive query end does not. + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(30), + ) + .unwrap(); + let searcher = + SeriesIndexSearcher::try_new(metadata, object_store, None, Some(time_range)).unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 20 + }] + ); + } + + #[tokio::test] + async fn search_skips_filters_for_columns_missing_from_older_index() { + let old_metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "schema-evolution.parquet"; + write_index( + old_metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "a", "y", 30), + ], + 2, + ) + .await; + + let mut builder = RegionMetadataBuilder::from_existing(old_metadata.as_ref().clone()); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag_2", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 4, + }); + let mut primary_key = old_metadata.primary_key.clone(); + primary_key.push(4); + builder.primary_key(primary_key); + let current_metadata = Arc::new(builder.build().unwrap()); + + let predicate = + Predicate::new(vec![col("tag_0").eq(lit("a")), col("tag_2").eq(lit("new"))]); + let searcher = + SeriesIndexSearcher::try_new(current_metadata, object_store, Some(&predicate), None) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![ + MetricSeriesId { + table_id: 1, + tsid: 10, + }, + MetricSeriesId { + table_id: 1, + tsid: 30, + }, + ] + ); + } + + #[tokio::test] + async fn search_streams_fixed_size_batches() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let rows = (0..501_u64) + .map(|tsid| (1, tsid, "a", "x", tsid as i64)) + .collect::>(); + write_index( + metadata.clone(), + object_store.clone(), + "batching.parquet", + &rows, + 100, + ) + .await; + + let searcher = SeriesIndexSearcher::try_new(metadata, object_store, None, None).unwrap(); + let batches = searcher + .search("batching.parquet") + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches.iter().map(Vec::len).collect::>(), [500, 1]); + assert_eq!( + batches[0][0], + MetricSeriesId { + table_id: 1, + tsid: 0 + } + ); + assert_eq!( + batches[1][0], + MetricSeriesId { + table_id: 1, + tsid: 500 + } + ); + } + + #[tokio::test] + async fn search_prunes_row_groups_and_empty_ranges() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "pruning.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 0, "a", "x", 0), + (1, 1, "b", "x", 1), + (1, 2, "m", "x", 2), + (1, 3, "m", "x", 3), + (1, 4, "y", "x", 4), + (1, 5, "z", "x", 5), + ], + 2, + ) + .await; + + let predicate = Predicate::new(vec![col("tag_0").eq(lit("m"))]); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + Some(&predicate), + 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 empty = SeriesIndexSearcher::try_new( + metadata, + object_store, + None, + Some(TimestampRange::empty()), + ) + .unwrap(); + assert!( + empty + .search("does-not-need-to-exist.parquet") + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + } +} diff --git a/src/mito2/src/series_index/writer.rs b/src/mito2/src/series_index/writer.rs index ce5b7405f2..9ebba289fb 100644 --- a/src/mito2/src/series_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -45,15 +45,13 @@ use crate::error::{ DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, }; +use crate::series_index::{ + MAX_TS_COLUMN, MIN_TS_COLUMN, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, +}; use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; -const MIN_TS_COLUMN: &str = "__series_min_ts"; -const MAX_TS_COLUMN: &str = "__series_max_ts"; -const ROW_COUNT_COLUMN: &str = "__series_row_count"; -const TABLE_ID_COLUMN: &str = "__table_id"; -const TSID_COLUMN: &str = "__tsid"; const WRITE_BATCH_SIZE: usize = 1024; type ParquetWriter = AsyncArrowWriter; diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index f3f0a34a68..4ccdcb4108 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -36,7 +36,7 @@ use datafusion_common::ScalarValue; use datatypes::arrow::array::{ ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, UInt64Array, }; -use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; use datatypes::types::json_type::JsonNativeType; @@ -136,12 +136,22 @@ pub(crate) fn column_values( column_index: usize, is_min: bool, ) -> Option { - let null_scalar: ScalarValue = column - .column_schema - .data_type - .as_arrow_type() - .try_into() - .ok()?; + column_values_by_type( + row_groups, + &column.column_schema.data_type.as_arrow_type(), + column_index, + is_min, + ) +} + +/// Returns min/max values of a parquet column with the given Arrow data type. +pub(crate) fn column_values_by_type( + row_groups: &[impl Borrow], + data_type: &ArrowDataType, + column_index: usize, + is_min: bool, +) -> Option { + let null_scalar: ScalarValue = data_type.try_into().ok()?; let scalar_values = row_groups .iter() .map(|meta| { @@ -152,16 +162,22 @@ pub(crate) fn column_values( } else { *s.max_opt()? }))), - Statistics::Int32(s) => Some(ScalarValue::Int32(Some(if is_min { - *s.min_opt()? - } else { - *s.max_opt()? - }))), - Statistics::Int64(s) => Some(ScalarValue::Int64(Some(if is_min { - *s.min_opt()? - } else { - *s.max_opt()? - }))), + Statistics::Int32(s) => { + let value = if is_min { *s.min_opt()? } else { *s.max_opt()? }; + if data_type == &ArrowDataType::UInt32 { + Some(ScalarValue::UInt32(Some(value as u32))) + } else { + Some(ScalarValue::Int32(Some(value))) + } + } + Statistics::Int64(s) => { + let value = if is_min { *s.min_opt()? } else { *s.max_opt()? }; + if data_type == &ArrowDataType::UInt64 { + Some(ScalarValue::UInt64(Some(value as u64))) + } else { + Some(ScalarValue::Int64(Some(value))) + } + } Statistics::Int96(_) => None, Statistics::Float(s) => Some(ScalarValue::Float32(Some(if is_min { *s.min_opt()? diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index fcaaae4a54..3af3b20d9d 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; +use datafusion_expr::Expr; use datatypes::arrow::array::{Array, BinaryArray, BooleanArray, BooleanBufferBuilder}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::SchemaRef; @@ -272,15 +273,9 @@ pub(crate) fn build_primary_key_filter( expected_metadata: Option<&RegionMetadata>, predicate: Option<&Predicate>, ) -> Option { - let filters = predicate + let filters = simple_tag_filters(sst_metadata, expected_metadata, predicate) .into_iter() - .flat_map(|predicate| predicate.exprs()) - .filter_map(|expr| SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr)) - .filter_map(|filter_ctx| { - (filter_ctx.semantic_type() == SemanticType::Tag) - .then(|| filter_ctx.filter().as_filter().cloned()) - .flatten() - }) + .map(|(_, filter)| filter) .collect::>(); if filters.is_empty() { return None; @@ -291,6 +286,33 @@ pub(crate) fn build_primary_key_filter( Some(CachedPrimaryKeyFilter::new(filter)) } +/// Extracts simple tag filters that can be applied to encoded primary keys or series indexes. +pub(crate) fn simple_tag_filters( + sst_metadata: &RegionMetadataRef, + expected_metadata: Option<&RegionMetadata>, + predicate: Option<&Predicate>, +) -> Vec<(Expr, SimpleFilterEvaluator)> { + predicate + .into_iter() + .flat_map(|predicate| predicate.exprs()) + .filter_map(|expr| { + SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr) + .map(|filter_ctx| (expr, filter_ctx)) + }) + .filter_map(|(expr, filter_ctx)| { + (filter_ctx.semantic_type() == SemanticType::Tag) + .then(|| { + filter_ctx + .filter() + .as_filter() + .cloned() + .map(|filter| (expr.clone(), filter)) + }) + .flatten() + }) + .collect() +} + /// How the parquet reader should apply each predicate. /// /// The reader runs in two phases. Predicates routed into `prefilter_builder`