Skip to main content

mito2/read/
series_reader.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Reads selected metric series from partition ranges.
16
17use std::collections::HashSet;
18use std::sync::Arc;
19use std::time::Instant;
20
21use async_stream::try_stream;
22use futures::TryStreamExt;
23use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec};
24use snafu::ResultExt;
25use store_api::region_engine::PartitionRange;
26use store_api::storage::SequenceRange;
27use tokio::sync::Semaphore;
28
29#[cfg(feature = "enterprise")]
30use crate::error::InvalidRequestSnafu;
31use crate::error::{JoinSnafu, Result, UnexpectedSnafu};
32use crate::read::BoxedRecordBatchStream;
33use crate::read::pruner::PartitionPruner;
34use crate::read::range_cache::{CachedRangeReader, build_series_range_cache_key};
35use crate::read::scan_region::StreamContext;
36use crate::read::scan_util::{
37    PartitionMetrics, SplitRecordBatchStream, compute_average_batch_size,
38    compute_parallel_channel_size, filter_flat_batch_by_sequence, new_filter_metrics,
39    scan_flat_mem_ranges, should_split_flat_batches_for_merge,
40};
41use crate::read::seq_scan::SeqScan;
42use crate::read::series_candidate::validate_metric_metadata;
43use crate::series_index::MetricSeriesId;
44use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
45use crate::sst::parquet::flat_format::primary_key_column_index;
46use crate::sst::parquet::prefilter::prefilter_flat_batch_by_primary_key;
47use crate::sst::parquet::reader::ReaderMetrics;
48use crate::sst::parquet::row_group::ParquetFetchMetrics;
49
50const TSID_DOMAIN_END: u128 = 1u128 << u64::BITS;
51
52/// A stable partition of the TSID integer domain.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
54pub(crate) struct SeriesRange {
55    start: u128,
56    end: u128,
57}
58
59impl SeriesRange {
60    pub(crate) fn new(partition: usize, partitions: usize) -> Option<Self> {
61        if partitions == 0 || partition >= partitions {
62            return None;
63        }
64
65        let partitions = partitions as u128;
66        let partition = partition as u128;
67        let boundary = |partition: u128| {
68            let numerator = partition * TSID_DOMAIN_END;
69            numerator.div_ceil(partitions)
70        };
71        Some(Self {
72            start: boundary(partition),
73            end: boundary(partition + 1),
74        })
75    }
76
77    fn partition_for(tsid: u64, partitions: usize) -> usize {
78        ((tsid as u128 * partitions as u128) >> u64::BITS) as usize
79    }
80}
81
82/// All series assigned to one data-reader partition.
83#[derive(Debug)]
84pub(crate) struct AssignedSeriesBatch {
85    range: SeriesRange,
86    series: Vec<MetricSeriesId>,
87    enable_range_cache: bool,
88}
89
90impl AssignedSeriesBatch {
91    fn new(range: SeriesRange, series: Vec<MetricSeriesId>, enable_range_cache: bool) -> Self {
92        Self {
93            range,
94            series,
95            enable_range_cache,
96        }
97    }
98
99    pub(crate) fn range(&self) -> SeriesRange {
100        self.range
101    }
102
103    pub(crate) fn series(&self) -> &[MetricSeriesId] {
104        &self.series
105    }
106
107    pub(crate) fn enable_range_cache(&self) -> bool {
108        self.enable_range_cache
109    }
110}
111
112/// Collects candidate batches and assigns every TSID by its integer range.
113pub(crate) struct SeriesBatchCollector {
114    assignments: Vec<Vec<MetricSeriesId>>,
115    num_series: usize,
116}
117
118impl SeriesBatchCollector {
119    pub(crate) fn new(partitions: usize) -> Option<Self> {
120        (partitions > 0).then(|| Self {
121            assignments: (0..partitions).map(|_| Vec::new()).collect(),
122            num_series: 0,
123        })
124    }
125
126    pub(crate) fn push(&mut self, batch: Vec<MetricSeriesId>) {
127        self.num_series += batch.len();
128        let partitions = self.assignments.len();
129        for series in batch {
130            let partition = SeriesRange::partition_for(series.tsid, partitions);
131            self.assignments[partition].push(series);
132        }
133    }
134
135    pub(crate) fn len(&self) -> usize {
136        self.num_series
137    }
138
139    pub(crate) fn finish(self, enable_range_cache: bool) -> Vec<AssignedSeriesBatch> {
140        let partitions = self.assignments.len();
141        self.assignments
142            .into_iter()
143            .enumerate()
144            .map(|(partition, series)| {
145                AssignedSeriesBatch::new(
146                    SeriesRange::new(partition, partitions).unwrap(),
147                    series,
148                    enable_range_cache,
149                )
150            })
151            .collect()
152    }
153}
154
155/// Immutable allow-list for one partition's metric series.
156#[derive(Clone, Debug)]
157struct MetricSeriesFilter {
158    range: SeriesRange,
159    series: Arc<HashSet<MetricSeriesId>>,
160    sorted_series: Arc<Vec<MetricSeriesId>>,
161    enable_range_cache: bool,
162}
163
164impl MetricSeriesFilter {
165    fn new(assigned: &AssignedSeriesBatch) -> Self {
166        let mut sorted_series = assigned.series().to_vec();
167        sorted_series.sort_unstable();
168        sorted_series.dedup();
169        let series = sorted_series.iter().copied().collect();
170        Self {
171            range: assigned.range(),
172            series: Arc::new(series),
173            sorted_series: Arc::new(sorted_series),
174            enable_range_cache: assigned.enable_range_cache(),
175        }
176    }
177
178    fn primary_key_filter(&self, codec: SparsePrimaryKeyCodec) -> Box<dyn PrimaryKeyFilter> {
179        Box::new(MetricSeriesPrimaryKeyFilter {
180            codec,
181            series: self.series.clone(),
182            last_primary_key: Vec::new(),
183            last_match: None,
184        })
185    }
186
187    fn overlaps_encoded_bounds(
188        &self,
189        codec: &SparsePrimaryKeyCodec,
190        encoded_min: &[u8],
191        encoded_max: &[u8],
192    ) -> Option<bool> {
193        let (min_table_id, min_tsid) = codec.decode_ids(encoded_min).ok()?;
194        let (max_table_id, max_tsid) = codec.decode_ids(encoded_max).ok()?;
195        let min = MetricSeriesId {
196            table_id: min_table_id,
197            tsid: min_tsid,
198        };
199        let max = MetricSeriesId {
200            table_id: max_table_id,
201            tsid: max_tsid,
202        };
203        if min > max {
204            return None;
205        }
206        if min_table_id != max_table_id {
207            return Some(true);
208        }
209
210        Some(u128::from(min_tsid) < self.range.end && u128::from(max_tsid) >= self.range.start)
211    }
212}
213
214struct MetricSeriesPrimaryKeyFilter {
215    codec: SparsePrimaryKeyCodec,
216    series: Arc<HashSet<MetricSeriesId>>,
217    last_primary_key: Vec<u8>,
218    last_match: Option<bool>,
219}
220
221impl PrimaryKeyFilter for MetricSeriesPrimaryKeyFilter {
222    fn matches(&mut self, primary_key: &[u8]) -> mito_codec::error::Result<bool> {
223        if let Some(last_match) = self.last_match
224            && self.last_primary_key == primary_key
225        {
226            return Ok(last_match);
227        }
228
229        let (table_id, tsid) = self.codec.decode_ids(primary_key)?;
230        let matched = self.series.contains(&MetricSeriesId { table_id, tsid });
231        self.last_primary_key.clear();
232        self.last_primary_key.extend_from_slice(primary_key);
233        self.last_match = Some(matched);
234        Ok(matched)
235    }
236}
237
238fn filter_flat_stream_by_series(
239    mut input: BoxedRecordBatchStream,
240    codec: SparsePrimaryKeyCodec,
241    filter: MetricSeriesFilter,
242) -> BoxedRecordBatchStream {
243    Box::pin(try_stream! {
244        let mut primary_key_filter = filter.primary_key_filter(codec);
245        while let Some(batch) = input.try_next().await? {
246            let pk_idx = primary_key_column_index(batch.num_columns());
247            if let Some(batch) = prefilter_flat_batch_by_primary_key(
248                batch,
249                pk_idx,
250                primary_key_filter.as_mut(),
251            )? {
252                yield batch;
253            }
254        }
255    })
256}
257
258/// Reads all collected metric series assigned to one partition.
259pub(crate) struct SeriesReader {
260    stream_ctx: Arc<StreamContext>,
261    partition_ranges: Vec<PartitionRange>,
262    filter: MetricSeriesFilter,
263    codec: SparsePrimaryKeyCodec,
264    partition_pruner: Arc<PartitionPruner>,
265    range_semaphore: Arc<Semaphore>,
266    part_metrics: PartitionMetrics,
267}
268
269impl SeriesReader {
270    /// Creates a reader for the series assigned to one data partition.
271    ///
272    /// `partition_pruner` must come from the candidate scanner's pruner, which
273    /// has predicate prefiltering disabled. The file read path applies precise
274    /// filters with tag filtering skipped because candidate discovery has already
275    /// enforced tag predicates through the exact assigned-series set.
276    ///
277    /// A single `range_semaphore` covers both the range-build phase and the final
278    /// merge.
279    pub(crate) fn try_new(
280        stream_ctx: Arc<StreamContext>,
281        partition_ranges: Vec<PartitionRange>,
282        assigned_series: AssignedSeriesBatch,
283        partition_pruner: Arc<PartitionPruner>,
284        range_semaphore: Arc<Semaphore>,
285        part_metrics: PartitionMetrics,
286    ) -> Result<Self> {
287        validate_metric_metadata(&stream_ctx)?;
288        #[cfg(feature = "enterprise")]
289        snafu::ensure!(
290            stream_ctx.input.extension_ranges().is_empty(),
291            InvalidRequestSnafu {
292                region_id: stream_ctx.input.region_metadata().region_id,
293                reason: "series reader does not support extension ranges",
294            }
295        );
296
297        let filter = MetricSeriesFilter::new(&assigned_series);
298        let codec = SparsePrimaryKeyCodec::new(stream_ctx.input.region_metadata());
299        Ok(Self {
300            stream_ctx,
301            partition_ranges,
302            filter,
303            codec,
304            partition_pruner,
305            range_semaphore,
306            part_metrics,
307        })
308    }
309
310    pub(crate) async fn build_stream(&self) -> Result<BoxedRecordBatchStream> {
311        if self.partition_ranges.is_empty() || self.filter.series.is_empty() {
312            return Ok(Box::pin(futures::stream::empty()));
313        }
314
315        let mut tasks = Vec::with_capacity(self.partition_ranges.len());
316        for part_range in self.partition_ranges.iter().copied() {
317            let stream_ctx = self.stream_ctx.clone();
318            let filter = self.filter.clone();
319            let codec = self.codec.clone();
320            let partition_pruner = self.partition_pruner.clone();
321            let range_semaphore = self.range_semaphore.clone();
322            let part_metrics = self.part_metrics.clone();
323            tasks.push(common_runtime::spawn_query(async move {
324                let _permit = range_semaphore.acquire().await.map_err(|error| {
325                    UnexpectedSnafu {
326                        reason: format!("failed to acquire series range permit: {error}"),
327                    }
328                    .build()
329                })?;
330                build_series_partition_range(
331                    stream_ctx,
332                    part_range,
333                    filter,
334                    codec,
335                    partition_pruner,
336                    part_metrics,
337                )
338                .await
339            }));
340        }
341
342        let mut range_streams = Vec::with_capacity(tasks.len());
343        let mut estimated_batch_sizes = Vec::with_capacity(tasks.len());
344        for task in tasks {
345            let (stream, estimated_batch_size) = task.await.context(JoinSnafu)??;
346            range_streams.push(stream);
347            estimated_batch_sizes.push(estimated_batch_size);
348        }
349
350        // Every range task above has finished, so all build permits are released
351        // and the final merge can reuse the same semaphore.
352        let estimated_batch_size = compute_average_batch_size(estimated_batch_sizes);
353        SeqScan::build_flat_reader_from_sources(
354            &self.stream_ctx,
355            range_streams,
356            Some(self.range_semaphore.clone()),
357            Some(&self.part_metrics),
358            true,
359            compute_parallel_channel_size(estimated_batch_size),
360        )
361        .await
362    }
363}
364
365async fn build_series_partition_range(
366    stream_ctx: Arc<StreamContext>,
367    part_range: PartitionRange,
368    filter: MetricSeriesFilter,
369    codec: SparsePrimaryKeyCodec,
370    partition_pruner: Arc<PartitionPruner>,
371    part_metrics: PartitionMetrics,
372) -> Result<(BoxedRecordBatchStream, usize)> {
373    let range = filter.range;
374    let cache_key = filter
375        .enable_range_cache
376        .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range))
377        .flatten();
378    let ctx = stream_ctx.clone();
379    let metrics = part_metrics.clone();
380    CachedRangeReader {
381        stream_ctx: &stream_ctx,
382        part_range: &part_range,
383        key: cache_key,
384        metrics: &part_metrics,
385    }
386    .build(move || async move {
387        build_uncached_series_partition_range(
388            ctx,
389            part_range,
390            filter,
391            codec,
392            partition_pruner,
393            metrics,
394        )
395        .await
396    })
397    .await
398}
399
400async fn build_uncached_series_partition_range(
401    stream_ctx: Arc<StreamContext>,
402    part_range: PartitionRange,
403    filter: MetricSeriesFilter,
404    codec: SparsePrimaryKeyCodec,
405    partition_pruner: Arc<PartitionPruner>,
406    part_metrics: PartitionMetrics,
407) -> Result<(BoxedRecordBatchStream, usize)> {
408    let range_meta = &stream_ctx.ranges[part_range.identifier];
409    let split_batch_size = should_split_flat_batches_for_merge(&stream_ctx, range_meta);
410    let mut sources = Vec::with_capacity(range_meta.row_group_indices.len());
411
412    for index in range_meta.row_group_indices.iter().copied() {
413        if stream_ctx.is_mem_range_index(index) {
414            let stream = Box::pin(scan_flat_mem_ranges(
415                stream_ctx.clone(),
416                part_metrics.clone(),
417                index,
418                range_meta.time_range,
419            ));
420            sources.push(filter_flat_stream_by_series(
421                stream,
422                codec.clone(),
423                filter.clone(),
424            ));
425            continue;
426        }
427
428        if stream_ctx.is_file_range_index(index) {
429            let file = stream_ctx.input.file_from_index(index);
430            if matches!(
431                file.primary_key_range(stream_ctx.input.primary_key_mapper())
432                    .and_then(|(min, max)| { filter.overlaps_encoded_bounds(&codec, &min, &max) }),
433                Some(false)
434            ) {
435                continue;
436            }
437
438            if partition_pruner.try_skip_manifest_pruned_file_range(index, &part_metrics) {
439                continue;
440            }
441            let mut reader_metrics = ReaderMetrics {
442                filter_metrics: new_filter_metrics(part_metrics.explain_verbose()),
443                ..Default::default()
444            };
445            let file_ranges = partition_pruner
446                .build_file_ranges(index, &part_metrics, &mut reader_metrics)
447                .await?;
448            part_metrics.inc_num_file_ranges(file_ranges.len());
449            part_metrics.merge_reader_metrics(&reader_metrics, None);
450            let ranges = file_ranges
451                .iter()
452                .filter(|file_range| {
453                    !matches!(
454                        file_range.primary_key_range().and_then(|(min, max)| {
455                            filter.overlaps_encoded_bounds(&codec, min, max)
456                        }),
457                        Some(false)
458                    )
459                })
460                .cloned()
461                .collect::<smallvec::SmallVec<[_; 2]>>();
462            if ranges.is_empty() {
463                continue;
464            }
465
466            let stream = scan_series_file_ranges(
467                part_metrics.clone(),
468                ranges,
469                filter.clone(),
470                codec.clone(),
471                stream_ctx.input.sequence_range,
472                stream_ctx.input.region_metadata().region_id,
473            );
474            sources.push(Box::pin(stream) as BoxedRecordBatchStream);
475            continue;
476        }
477
478        return UnexpectedSnafu {
479            reason: format!(
480                "series reader received unsupported range index {}",
481                index.index
482            ),
483        }
484        .fail();
485    }
486
487    if split_batch_size.is_some() {
488        sources = sources
489            .into_iter()
490            .map(|stream| Box::pin(SplitRecordBatchStream::new(stream)) as BoxedRecordBatchStream)
491            .collect();
492    }
493    let estimated_batch_size = split_batch_size.unwrap_or(DEFAULT_READ_BATCH_SIZE);
494    let stream = SeqScan::build_flat_reader_from_sources(
495        &stream_ctx,
496        sources,
497        None,
498        Some(&part_metrics),
499        false,
500        compute_parallel_channel_size(estimated_batch_size),
501    )
502    .await?;
503    Ok((stream, estimated_batch_size))
504}
505
506fn scan_series_file_ranges(
507    part_metrics: PartitionMetrics,
508    ranges: smallvec::SmallVec<[crate::sst::parquet::file_range::FileRange; 2]>,
509    filter: MetricSeriesFilter,
510    codec: SparsePrimaryKeyCodec,
511    sequence_range: Option<SequenceRange>,
512    region_id: store_api::storage::RegionId,
513) -> impl futures::Stream<Item = Result<datatypes::arrow::record_batch::RecordBatch>> {
514    try_stream! {
515        let fetch_metrics = part_metrics
516            .explain_verbose()
517            .then(|| Arc::new(ParquetFetchMetrics::default()));
518        let mut reader_metrics = ReaderMetrics {
519            fetch_metrics: fetch_metrics.clone(),
520            ..Default::default()
521        };
522        let mut primary_key_filter = filter.primary_key_filter(codec);
523
524        for range in ranges {
525            let build_start = Instant::now();
526            let searcher = range.range_index_searcher().await?;
527            let reader = if let Some(searcher) = searcher {
528                let row_group_id = u32::try_from(range.row_group_index()).map_err(|_| UnexpectedSnafu {
529                    reason: format!("row group index exceeds u32: {}", range.row_group_index()),
530                }.build())?;
531                let selected = searcher.search(row_group_id, &filter.sorted_series).await?;
532                range.reader_by_row_ranges(selected, fetch_metrics.as_deref()).await?
533            } else {
534                range.reader_by_primary_key(primary_key_filter.as_mut(), fetch_metrics.as_deref()).await?
535            };
536            let build_cost = build_start.elapsed();
537            reader_metrics.build_cost += build_cost;
538            part_metrics.inc_build_reader_cost(build_cost);
539            let Some(mut reader) = reader else {
540                continue;
541            };
542
543            let scan_start = Instant::now();
544            let file_sequence_trusted = range
545                .file_handle()
546                .is_effective_target_sequence_trusted(region_id);
547            while let Some(record_batch) = reader.next_batch().await? {
548                reader_metrics.num_record_batches += 1;
549                reader_metrics.num_batches += 1;
550                reader_metrics.num_rows += record_batch.num_rows();
551
552                let Some(record_batch) = filter_flat_batch_by_sequence(
553                    record_batch,
554                    sequence_range,
555                    file_sequence_trusted,
556                )? else {
557                    continue;
558                };
559
560                let num_rows_before_filter = record_batch.num_rows();
561                let Some(record_batch) = range.precise_filter_flat(
562                    record_batch,
563                    range.pre_filter_mode().skip_fields(),
564                    true,
565                )? else {
566                    reader_metrics.filter_metrics.rows_precise_filtered +=
567                        num_rows_before_filter;
568                    continue;
569                };
570                reader_metrics.filter_metrics.rows_precise_filtered +=
571                    num_rows_before_filter - record_batch.num_rows();
572
573                let record_batch = if let Some(mapper) = range.compaction_projection_mapper() {
574                    mapper.project(record_batch)?
575                } else {
576                    record_batch
577                };
578                if let Some(compat) = range.compat_batch() {
579                    yield compat.compat(record_batch)?;
580                } else {
581                    yield record_batch;
582                }
583            }
584            reader_metrics.scan_cost += scan_start.elapsed();
585        }
586
587        reader_metrics.observe_rows("series_data");
588        reader_metrics.filter_metrics.observe();
589        part_metrics.merge_reader_metrics(&reader_metrics, None);
590    }
591}
592
593#[cfg(test)]
594mod tests {
595    use store_api::codec::PrimaryKeyEncoding;
596
597    use super::*;
598    use crate::error::DecodeSnafu;
599    use crate::test_util::sst_util::sst_region_metadata_with_encoding;
600
601    fn series(table_id: u32, tsid: u64) -> MetricSeriesId {
602        MetricSeriesId { table_id, tsid }
603    }
604
605    fn assigned_batch(
606        partitions: usize,
607        partition: usize,
608        series: Vec<MetricSeriesId>,
609    ) -> AssignedSeriesBatch {
610        let mut collector = SeriesBatchCollector::new(partitions).unwrap();
611        collector.push(series);
612        collector.finish(true).remove(partition)
613    }
614
615    #[test]
616    fn series_range_assignment_is_stable_across_batch_boundaries() {
617        let input = vec![
618            series(1, 0),
619            series(2, 1u64 << 62),
620            series(1, 1u64 << 63),
621            series(2, 3u64 << 62),
622            series(1, u64::MAX),
623            series(3, 0),
624        ];
625        let mut first = SeriesBatchCollector::new(4).unwrap();
626        first.push(input.clone());
627        let first = first.finish(true);
628
629        let mut second = SeriesBatchCollector::new(4).unwrap();
630        for chunk in input.chunks(2) {
631            second.push(chunk.to_vec());
632        }
633        let second = second.finish(true);
634
635        assert_eq!(
636            first
637                .iter()
638                .map(AssignedSeriesBatch::series)
639                .collect::<Vec<_>>(),
640            second
641                .iter()
642                .map(AssignedSeriesBatch::series)
643                .collect::<Vec<_>>()
644        );
645        assert_eq!(&[series(1, 0), series(3, 0)], first[0].series());
646        assert_eq!(&[series(2, 1u64 << 62)], first[1].series());
647        assert_eq!(
648            input.len(),
649            first
650                .iter()
651                .map(|batch| batch.series().len())
652                .sum::<usize>()
653        );
654    }
655
656    #[test]
657    fn collector_tracks_size_and_cache_policy() {
658        let mut collector = SeriesBatchCollector::new(2).unwrap();
659        collector.push(vec![series(1, 1), series(1, u64::MAX)]);
660        collector.push(vec![series(2, 2)]);
661        assert_eq!(3, collector.len());
662
663        let assignments = collector.finish(false);
664        assert_eq!(
665            3,
666            assignments
667                .iter()
668                .map(|batch| batch.series().len())
669                .sum::<usize>()
670        );
671        assert!(assignments.iter().all(|batch| !batch.enable_range_cache()));
672    }
673
674    #[test]
675    fn series_ranges_cover_the_tsid_domain() {
676        let ranges = (0..3)
677            .map(|partition| SeriesRange::new(partition, 3).unwrap())
678            .collect::<Vec<_>>();
679        assert_eq!(0, ranges[0].start);
680        assert_eq!(TSID_DOMAIN_END, ranges[2].end);
681        assert_eq!(ranges[0].end, ranges[1].start);
682        assert_eq!(ranges[1].end, ranges[2].start);
683        assert_eq!(0, SeriesRange::partition_for(0, 3));
684        assert_eq!(2, SeriesRange::partition_for(u64::MAX, 3));
685
686        for (partition, range) in ranges.iter().enumerate() {
687            assert_eq!(partition, SeriesRange::partition_for(range.start as u64, 3));
688            if range.end < TSID_DOMAIN_END {
689                assert_eq!(
690                    partition + 1,
691                    SeriesRange::partition_for(range.end as u64, 3)
692                );
693            }
694        }
695    }
696
697    #[test]
698    fn metric_series_filter_matches_encoded_primary_key() {
699        let metadata = Arc::new(sst_region_metadata_with_encoding(
700            PrimaryKeyEncoding::Sparse,
701        ));
702        let assigned = assigned_batch(1, 0, vec![series(1, 10), series(2, 20)]);
703        let filter = MetricSeriesFilter::new(&assigned);
704        let codec = SparsePrimaryKeyCodec::new(&metadata);
705        let mut primary_key_filter = filter.primary_key_filter(codec);
706
707        for selected in [series(1, 10), series(2, 20)] {
708            assert!(
709                primary_key_filter
710                    .matches(&encode_series(&metadata, selected.table_id, selected.tsid))
711                    .context(DecodeSnafu)
712                    .unwrap()
713            );
714        }
715        for unselected in [series(2, 10), series(1, 20)] {
716            assert!(
717                !primary_key_filter
718                    .matches(&encode_series(
719                        &metadata,
720                        unselected.table_id,
721                        unselected.tsid,
722                    ))
723                    .context(DecodeSnafu)
724                    .unwrap()
725            );
726        }
727    }
728
729    fn encode_series(
730        metadata: &store_api::metadata::RegionMetadataRef,
731        table_id: u32,
732        tsid: u64,
733    ) -> Vec<u8> {
734        let codec = SparsePrimaryKeyCodec::new(metadata);
735        let mut primary_key = Vec::new();
736        codec
737            .encode_internal(table_id, tsid, &mut primary_key)
738            .unwrap();
739        primary_key
740    }
741
742    #[test]
743    fn series_range_overlaps_single_table_primary_key_bounds() {
744        let metadata = Arc::new(sst_region_metadata_with_encoding(
745            PrimaryKeyEncoding::Sparse,
746        ));
747        let assigned = assigned_batch(2, 0, vec![series(1, 10), series(2, 20)]);
748        let filter = MetricSeriesFilter::new(&assigned);
749        let codec = SparsePrimaryKeyCodec::new(&metadata);
750
751        // The full assigned range overlaps even though the actual candidate bounds do not.
752        assert_eq!(
753            Some(true),
754            filter.overlaps_encoded_bounds(
755                &codec,
756                &encode_series(&metadata, 1, 100),
757                &encode_series(&metadata, 1, 200),
758            )
759        );
760        assert_eq!(
761            Some(false),
762            filter.overlaps_encoded_bounds(
763                &codec,
764                &encode_series(&metadata, 1, 1u64 << 63),
765                &encode_series(&metadata, 1, u64::MAX),
766            )
767        );
768        assert_eq!(
769            Some(true),
770            filter.overlaps_encoded_bounds(
771                &codec,
772                &encode_series(&metadata, 3, 10),
773                &encode_series(&metadata, 3, 20),
774            )
775        );
776    }
777
778    #[test]
779    fn series_range_keeps_multiple_table_primary_key_bounds() {
780        let metadata = Arc::new(sst_region_metadata_with_encoding(
781            PrimaryKeyEncoding::Sparse,
782        ));
783        let assigned = assigned_batch(2, 0, vec![series(1, 10), series(2, 20)]);
784        let filter = MetricSeriesFilter::new(&assigned);
785        let codec = SparsePrimaryKeyCodec::new(&metadata);
786
787        assert_eq!(
788            Some(true),
789            filter.overlaps_encoded_bounds(
790                &codec,
791                &encode_series(&metadata, 1, 1u64 << 63),
792                &encode_series(&metadata, 2, u64::MAX),
793            )
794        );
795    }
796
797    #[test]
798    fn series_statistics_keep_invalid_or_inverted_bounds() {
799        let metadata = Arc::new(sst_region_metadata_with_encoding(
800            PrimaryKeyEncoding::Sparse,
801        ));
802        let assigned = assigned_batch(2, 0, vec![series(1, 10)]);
803        let filter = MetricSeriesFilter::new(&assigned);
804        let codec = SparsePrimaryKeyCodec::new(&metadata);
805
806        assert_eq!(
807            None,
808            filter.overlaps_encoded_bounds(&codec, b"invalid", b"bounds")
809        );
810        assert_eq!(
811            None,
812            filter.overlaps_encoded_bounds(
813                &codec,
814                &encode_series(&metadata, 2, 0),
815                &encode_series(&metadata, 1, 0),
816            )
817        );
818    }
819}