mito2/read/
range_cache.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//! Utilities for the partition range scan result cache.
16
17use std::mem;
18use std::sync::Arc;
19
20use async_stream::try_stream;
21use common_telemetry::warn;
22use common_time::Timestamp;
23use common_time::range::TimestampRange;
24use common_time::timestamp::TimeUnit;
25use datafusion_expr::expr::Expr;
26use datafusion_expr::{Between, BinaryExpr, Operator};
27use datatypes::arrow::compute::concat_batches;
28use datatypes::arrow::record_batch::RecordBatch;
29use datatypes::prelude::ConcreteDataType;
30use datatypes::value::scalar_value_to_timestamp;
31use futures::TryStreamExt;
32use snafu::ResultExt;
33use store_api::region_engine::PartitionRange;
34use store_api::storage::{FileId, RegionId, TimeSeriesRowSelector};
35use table::predicate::is_string_timestamp_literal;
36use tokio::sync::{mpsc, oneshot};
37
38use crate::cache::CacheStrategy;
39use crate::error::{ComputeArrowSnafu, Result};
40use crate::read::BoxedRecordBatchStream;
41use crate::read::read_columns::ReadColumns;
42use crate::read::scan_region::StreamContext;
43use crate::read::scan_util::PartitionMetrics;
44use crate::read::series_reader::SeriesRange;
45use crate::region::options::MergeMode;
46use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
47
48const RANGE_CACHE_COMPACT_THRESHOLD_BYTES: usize = 8 * 1024 * 1024;
49
50/// Fingerprint of the scan request fields that affect partition range cache reuse.
51///
52/// It records a normalized view of the projected columns and filters, plus
53/// scan options that can change the returned rows. Schema-dependent metadata
54/// and the partition expression version are included so cached results are not
55/// reused across incompatible schema or partitioning changes.
56#[derive(Debug, Clone, PartialEq, Eq, Hash)]
57pub(crate) struct ScanRequestFingerprint {
58    /// Projection and filters without the time index and partition exprs.
59    inner: Arc<SharedScanRequestFingerprint>,
60    /// Filters with the time index column.
61    time_filters: Option<Arc<Vec<String>>>,
62    series_row_selector: Option<TimeSeriesRowSelector>,
63    append_mode: bool,
64    filter_deleted: bool,
65    merge_mode: MergeMode,
66    stage: RangeScanStage,
67    /// We keep the partition expr version to ensure we won't reuse the fingerprint after we change the partition expr.
68    /// We store the version instead of the whole partition expr or partition expr filters.
69    partition_expr_version: u64,
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
73enum RangeScanStage {
74    Data,
75    CandidateSeries,
76    SeriesData(SeriesRange),
77}
78
79#[derive(Debug)]
80pub(crate) struct ScanRequestFingerprintBuilder {
81    pub(crate) read_columns: ReadColumns,
82    pub(crate) read_column_types: Vec<Option<ConcreteDataType>>,
83    pub(crate) filters: Vec<String>,
84    pub(crate) time_filters: Vec<String>,
85    pub(crate) series_row_selector: Option<TimeSeriesRowSelector>,
86    pub(crate) append_mode: bool,
87    pub(crate) filter_deleted: bool,
88    pub(crate) merge_mode: MergeMode,
89    pub(crate) partition_expr_version: u64,
90}
91
92impl ScanRequestFingerprintBuilder {
93    pub(crate) fn build(self) -> ScanRequestFingerprint {
94        let Self {
95            read_columns,
96            read_column_types,
97            filters,
98            time_filters,
99            series_row_selector,
100            append_mode,
101            filter_deleted,
102            merge_mode,
103            partition_expr_version,
104        } = self;
105
106        ScanRequestFingerprint {
107            inner: Arc::new(SharedScanRequestFingerprint {
108                read_columns,
109                read_column_types,
110                filters,
111            }),
112            time_filters: (!time_filters.is_empty()).then(|| Arc::new(time_filters)),
113            series_row_selector,
114            append_mode,
115            filter_deleted,
116            merge_mode,
117            stage: RangeScanStage::Data,
118            partition_expr_version,
119        }
120    }
121}
122
123/// Non-copiable struct of the fingerprint.
124#[derive(Debug, PartialEq, Eq, Hash)]
125struct SharedScanRequestFingerprint {
126    /// Logical columns of the projection.
127    read_columns: ReadColumns,
128    /// Column types of the projection.
129    /// We keep this to ensure we won't reuse the fingerprint after a schema change.
130    read_column_types: Vec<Option<ConcreteDataType>>,
131    /// Filters without the time index column and region partition exprs.
132    filters: Vec<String>,
133}
134
135impl ScanRequestFingerprint {
136    #[cfg(test)]
137    pub(crate) fn read_columns(&self) -> &ReadColumns {
138        &self.inner.read_columns
139    }
140
141    #[cfg(test)]
142    pub(crate) fn read_column_types(&self) -> &[Option<ConcreteDataType>] {
143        &self.inner.read_column_types
144    }
145
146    #[cfg(test)]
147    pub(crate) fn filters(&self) -> &[String] {
148        &self.inner.filters
149    }
150
151    #[cfg(test)]
152    pub(crate) fn time_filters(&self) -> &[String] {
153        self.time_filters
154            .as_deref()
155            .map(Vec::as_slice)
156            .unwrap_or(&[])
157    }
158
159    pub(crate) fn without_time_filters(&self) -> Self {
160        Self {
161            inner: Arc::clone(&self.inner),
162            time_filters: None,
163            series_row_selector: self.series_row_selector,
164            append_mode: self.append_mode,
165            filter_deleted: self.filter_deleted,
166            merge_mode: self.merge_mode,
167            stage: self.stage,
168            partition_expr_version: self.partition_expr_version,
169        }
170    }
171
172    fn for_candidate_series(&self) -> Self {
173        Self {
174            inner: Arc::clone(&self.inner),
175            time_filters: self.time_filters.clone(),
176            series_row_selector: self.series_row_selector,
177            append_mode: self.append_mode,
178            filter_deleted: self.filter_deleted,
179            merge_mode: self.merge_mode,
180            stage: RangeScanStage::CandidateSeries,
181            partition_expr_version: self.partition_expr_version,
182        }
183    }
184
185    fn for_series_data(&self, range: SeriesRange) -> Self {
186        Self {
187            inner: Arc::clone(&self.inner),
188            time_filters: self.time_filters.clone(),
189            series_row_selector: self.series_row_selector,
190            append_mode: self.append_mode,
191            filter_deleted: self.filter_deleted,
192            merge_mode: self.merge_mode,
193            stage: RangeScanStage::SeriesData(range),
194            partition_expr_version: self.partition_expr_version,
195        }
196    }
197
198    pub(crate) fn estimated_size(&self) -> usize {
199        mem::size_of::<SharedScanRequestFingerprint>()
200            + self.inner.read_columns.estimated_size()
201            + self.inner.read_column_types.capacity() * mem::size_of::<Option<ConcreteDataType>>()
202            + self.inner.filters.capacity() * mem::size_of::<String>()
203            + self
204                .inner
205                .filters
206                .iter()
207                .map(|filter| filter.capacity())
208                .sum::<usize>()
209            + self.time_filters.as_ref().map_or(0, |filters| {
210                mem::size_of::<Vec<String>>()
211                    + filters.capacity() * mem::size_of::<String>()
212                    + filters
213                        .iter()
214                        .map(|filter| filter.capacity())
215                        .sum::<usize>()
216            })
217    }
218}
219
220/// Cache key for range scan outputs.
221#[derive(Debug, Clone, PartialEq, Eq, Hash)]
222pub(crate) struct RangeScanCacheKey {
223    pub(crate) region_id: RegionId,
224    /// Sorted (file_id, row_group_index) pairs that uniquely identify the data this range covers.
225    pub(crate) row_groups: Vec<(FileId, i64)>,
226    pub(crate) scan: ScanRequestFingerprint,
227}
228
229impl RangeScanCacheKey {
230    pub(crate) fn estimated_size(&self) -> usize {
231        mem::size_of::<Self>()
232            + self.row_groups.capacity() * mem::size_of::<(FileId, i64)>()
233            + self.scan.estimated_size()
234    }
235}
236
237/// Cached result for one range scan.
238#[derive(Debug)]
239pub(crate) struct CachedBatchSlice {
240    batch: RecordBatch,
241    slice_lengths: Vec<usize>,
242}
243
244impl CachedBatchSlice {
245    fn metadata_size(&self) -> usize {
246        self.slice_lengths.capacity() * mem::size_of::<usize>()
247    }
248}
249
250pub(crate) struct RangeScanCacheValue {
251    cached_batches: Vec<CachedBatchSlice>,
252    /// Precomputed size of all compacted batches.
253    estimated_batches_size: usize,
254}
255
256impl RangeScanCacheValue {
257    pub(crate) fn new(
258        cached_batches: Vec<CachedBatchSlice>,
259        estimated_batches_size: usize,
260    ) -> Self {
261        Self {
262            cached_batches,
263            estimated_batches_size,
264        }
265    }
266
267    pub(crate) fn estimated_size(&self) -> usize {
268        mem::size_of::<Self>()
269            + self.cached_batches.capacity() * mem::size_of::<CachedBatchSlice>()
270            + self
271                .cached_batches
272                .iter()
273                .map(CachedBatchSlice::metadata_size)
274                .sum::<usize>()
275            + self.estimated_batches_size
276    }
277}
278
279/// Row groups and whether all sources are file-only for a partition range.
280pub(crate) struct PartitionRangeRowGroups {
281    /// Sorted (file_id, row_group_index) pairs.
282    pub(crate) row_groups: Vec<(FileId, i64)>,
283    pub(crate) only_file_sources: bool,
284}
285
286/// Collects (file_id, row_group_index) pairs from a partition range's row group indices.
287pub(crate) fn collect_partition_range_row_groups(
288    stream_ctx: &StreamContext,
289    part_range: &PartitionRange,
290) -> PartitionRangeRowGroups {
291    let range_meta = &stream_ctx.ranges[part_range.identifier];
292    let mut row_groups = Vec::new();
293    let mut only_file_sources = true;
294
295    for index in &range_meta.row_group_indices {
296        if stream_ctx.is_file_range_index(*index) {
297            let file_id = stream_ctx.input.file_from_index(*index).file_id().file_id();
298            row_groups.push((file_id, index.row_group_index));
299        } else {
300            only_file_sources = false;
301        }
302    }
303
304    row_groups.sort_unstable_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()).then(a.1.cmp(&b.1)));
305
306    PartitionRangeRowGroups {
307        row_groups,
308        only_file_sources,
309    }
310}
311
312/// Returns the timestamp range where all time-only predicates are guaranteed true.
313///
314/// Returns `Some(min_to_max)` for empty input (vacuously true everywhere).
315/// Returns `None` if any expression contains an unsupported shape: `OR`, `NOT`,
316/// `IN`, non-literal RHS, unsupported operator, column-name mismatch, an `=`
317/// literal that cannot be represented exactly in the column unit, or overflow
318/// during bound adjustment.
319///
320/// This is intentionally stricter than `extract_time_range_from_expr` in
321/// `table::predicate`: lower bounds round up and upper bounds round down. If a
322/// partition's file-time range is contained by the returned range, every row in
323/// that partition satisfies the original time predicates.
324///
325/// `IsNull`/`IsNotNull` on the time index are not routed into `time_filters`
326/// today. If that changes, handle them here before stripping time filters from
327/// the cache key.
328pub(crate) fn implied_time_range_from_exprs(
329    ts_col_name: &str,
330    ts_col_unit: TimeUnit,
331    exprs: &[&Expr],
332) -> Option<TimestampRange> {
333    let mut acc = TimestampRange::min_to_max();
334    for expr in exprs {
335        let r = implied_time_range_from_expr(ts_col_name, ts_col_unit, expr)?;
336        acc = acc.and(&r);
337    }
338    Some(acc)
339}
340
341fn implied_time_range_from_expr(
342    ts_col_name: &str,
343    ts_col_unit: TimeUnit,
344    expr: &Expr,
345) -> Option<TimestampRange> {
346    match expr {
347        Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
348            Operator::And => {
349                let l = implied_time_range_from_expr(ts_col_name, ts_col_unit, left)?;
350                let r = implied_time_range_from_expr(ts_col_name, ts_col_unit, right)?;
351                Some(l.and(&r))
352            }
353            Operator::Eq | Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => {
354                implied_from_cmp(ts_col_name, ts_col_unit, left, *op, right)
355            }
356            // `OR` would require a strict intersection over a union of half-planes
357            // (not the loose-span union provided by `TimestampRange::or`), so we
358            // refuse it. Any other operator is unsupported.
359            _ => None,
360        },
361        Expr::Between(Between {
362            expr,
363            negated,
364            low,
365            high,
366        }) => {
367            if *negated {
368                return None;
369            }
370            implied_from_between(ts_col_name, ts_col_unit, expr, low, high)
371        }
372        // Includes `IsNull`, `IsNotNull`, `Not`, `InList`, function calls, etc.
373        _ => None,
374    }
375}
376
377fn match_ts_column_literal<'a>(
378    ts_col_name: &str,
379    left: &'a Expr,
380    right: &'a Expr,
381) -> Option<(Timestamp, bool)> {
382    let (col, scalar, reverse) = match (left, right) {
383        (Expr::Column(c), Expr::Literal(s, _)) => (c, s, false),
384        (Expr::Literal(s, _), Expr::Column(c)) => (c, s, true),
385        _ => return None,
386    };
387    if col.name != ts_col_name {
388        return None;
389    }
390    // Reject string literals: their conversion needs a timezone we do not have,
391    // and the existing extractor in `table::predicate` rejects them too.
392    if is_string_timestamp_literal(scalar) {
393        return None;
394    }
395    scalar_value_to_timestamp(scalar, None).map(|t| (t, reverse))
396}
397
398fn implied_from_cmp(
399    ts_col_name: &str,
400    ts_col_unit: TimeUnit,
401    left: &Expr,
402    op: Operator,
403    right: &Expr,
404) -> Option<TimestampRange> {
405    let (ts, reverse) = match_ts_column_literal(ts_col_name, left, right)?;
406    // Normalize to "column OP literal".
407    let op = if reverse {
408        match op {
409            Operator::Lt => Operator::Gt,
410            Operator::LtEq => Operator::GtEq,
411            Operator::Gt => Operator::Lt,
412            Operator::GtEq => Operator::LtEq,
413            Operator::Eq => Operator::Eq,
414            _ => return None,
415        }
416    } else {
417        op
418    };
419
420    match op {
421        Operator::GtEq => {
422            // ts >= L. Round the lower bound up in the column unit.
423            let b = ts.convert_to_ceil(ts_col_unit)?;
424            Some(TimestampRange::from_start(b))
425        }
426        Operator::Gt => {
427            // ts > L. floor(L) + 1 is the tight lower bound in the column unit.
428            let v = ts.convert_to(ts_col_unit)?.value().checked_add(1)?;
429            Some(TimestampRange::from_start(Timestamp::new(v, ts_col_unit)))
430        }
431        Operator::LtEq => {
432            // ts <= L. Round the upper bound down in the column unit.
433            let b = ts.convert_to(ts_col_unit)?;
434            Some(TimestampRange::until_end(b, true))
435        }
436        Operator::Lt => {
437            // ts < L. `ts < ceil(L)` is the tight bound: equal to `ts < L` when
438            // L is exactly representable, and `ts <= floor(L)` otherwise.
439            let b = ts.convert_to_ceil(ts_col_unit)?;
440            Some(TimestampRange::until_end(b, false))
441        }
442        Operator::Eq => {
443            // ts = L. Only provable when L is exactly representable.
444            let f = ts.convert_to(ts_col_unit)?;
445            let c = ts.convert_to_ceil(ts_col_unit)?;
446            if f.value() != c.value() {
447                return None;
448            }
449            Some(TimestampRange::single(f))
450        }
451        _ => None,
452    }
453}
454
455fn implied_from_between(
456    ts_col_name: &str,
457    ts_col_unit: TimeUnit,
458    expr: &Expr,
459    low: &Expr,
460    high: &Expr,
461) -> Option<TimestampRange> {
462    let Expr::Column(c) = expr else {
463        return None;
464    };
465    if c.name != ts_col_name {
466        return None;
467    }
468    let (low_s, high_s) = match (low, high) {
469        (Expr::Literal(l, _), Expr::Literal(h, _)) => (l, h),
470        _ => return None,
471    };
472    if is_string_timestamp_literal(low_s) || is_string_timestamp_literal(high_s) {
473        return None;
474    }
475    let low_ts = scalar_value_to_timestamp(low_s, None)?;
476    let high_ts = scalar_value_to_timestamp(high_s, None)?;
477    // BETWEEN low AND high is equivalent to ts >= low AND ts <= high.
478    let lo = low_ts.convert_to_ceil(ts_col_unit)?;
479    let hi = high_ts.convert_to(ts_col_unit)?;
480    Some(TimestampRange::new_inclusive(Some(lo), Some(hi)))
481}
482
483/// Builds a cache key for the given partition range if it is eligible for caching.
484pub(crate) fn build_range_cache_key(
485    stream_ctx: &StreamContext,
486    part_range: &PartitionRange,
487) -> Option<RangeScanCacheKey> {
488    build_range_cache_key_inner(stream_ctx, part_range, None)
489}
490
491/// Builds a cache key for a candidate-series partition-range result.
492pub(crate) fn build_candidate_range_cache_key(
493    stream_ctx: &StreamContext,
494    part_range: &PartitionRange,
495) -> Option<RangeScanCacheKey> {
496    build_range_cache_key_inner(
497        stream_ctx,
498        part_range,
499        Some(RangeScanStage::CandidateSeries),
500    )
501}
502
503/// Builds a cache key for a two-phase series-data partition-range result.
504#[allow(dead_code)]
505pub(crate) fn build_series_range_cache_key(
506    stream_ctx: &StreamContext,
507    part_range: &PartitionRange,
508    range: SeriesRange,
509) -> Option<RangeScanCacheKey> {
510    build_range_cache_key_inner(
511        stream_ctx,
512        part_range,
513        Some(RangeScanStage::SeriesData(range)),
514    )
515}
516
517fn build_range_cache_key_inner(
518    stream_ctx: &StreamContext,
519    part_range: &PartitionRange,
520    stage: Option<RangeScanStage>,
521) -> Option<RangeScanCacheKey> {
522    if !stream_ctx.input.cache_strategy.has_range_result_cache() {
523        return None;
524    }
525
526    let fingerprint = stream_ctx.scan_fingerprint.as_ref()?;
527
528    // Dyn filters can change at runtime, so we can't cache when they're present.
529    let has_dyn_filters = stream_ctx
530        .input
531        .predicate_group()
532        .predicate_without_region()
533        .is_some_and(|p| !p.dyn_filters().is_empty());
534    if has_dyn_filters {
535        return None;
536    }
537
538    let rg = collect_partition_range_row_groups(stream_ctx, part_range);
539    if !rg.only_file_sources || rg.row_groups.is_empty() {
540        return None;
541    }
542
543    // If the implied range covers this partition's `FileTimeRange`, drop
544    // time-only predicates from the cache key so that queries with different
545    // but equally-covering time bounds share an entry. `None` means some
546    // time-only predicate had an unsupported shape (e.g. `OR`), so we keep
547    // them in the key.
548    let range_meta = &stream_ctx.ranges[part_range.identifier];
549    let (file_min, file_max) = range_meta.time_range;
550    let covers = match &stream_ctx.scan_implied_time_range {
551        // An empty implied range can never cover a non-empty file range, so
552        // short-circuit. We also skip the unit asserts because
553        // `TimestampRange::empty()` uses `Timestamp::default()` (millisecond),
554        // which would falsely trip the asserts for non-ms time index columns.
555        Some(implied) if !implied.is_empty() => {
556            // The `contains` check is sound only when `file_min`/`file_max`
557            // share the implied range's unit (the time index column's unit).
558            // Mito stores time index values in that unit; assert to catch any
559            // future drift.
560            if let Some(ts) = implied.start().as_ref().or(implied.end().as_ref()) {
561                assert_eq!(file_min.unit(), ts.unit());
562                assert_eq!(file_max.unit(), ts.unit());
563            }
564            implied.contains(&file_min) && implied.contains(&file_max)
565        }
566        _ => false,
567    };
568    let scan = if covers {
569        fingerprint.without_time_filters()
570    } else {
571        fingerprint.clone()
572    };
573    let scan = match stage {
574        Some(RangeScanStage::CandidateSeries) => scan.for_candidate_series(),
575        Some(RangeScanStage::SeriesData(assignment)) => scan.for_series_data(assignment),
576        Some(RangeScanStage::Data) | None => scan,
577    };
578
579    Some(RangeScanCacheKey {
580        region_id: stream_ctx.input.region_metadata().region_id,
581        row_groups: rg.row_groups,
582        scan,
583    })
584}
585
586/// Returns a stream that replays cached record batches.
587pub(crate) fn cached_flat_range_stream(value: Arc<RangeScanCacheValue>) -> BoxedRecordBatchStream {
588    Box::pin(try_stream! {
589        for cached_batch in &value.cached_batches {
590            let mut offset = 0;
591            for &len in &cached_batch.slice_lengths {
592                yield cached_batch.batch.slice(offset, len);
593                offset += len;
594            }
595        }
596    })
597}
598
599enum CacheConcatCommand {
600    Compact(Vec<RecordBatch>),
601    Finish {
602        pending: Vec<RecordBatch>,
603        key: RangeScanCacheKey,
604        cache_strategy: CacheStrategy,
605        part_metrics: PartitionMetrics,
606        result_tx: Option<oneshot::Sender<Result<Arc<RangeScanCacheValue>>>>,
607    },
608}
609
610#[derive(Default)]
611struct CacheConcatState {
612    cached_batches: Vec<CachedBatchSlice>,
613    estimated_size: usize,
614}
615
616impl CacheConcatState {
617    async fn compact(
618        &mut self,
619        batches: Vec<RecordBatch>,
620        limiter: &crate::cache::RangeResultMemoryLimiter,
621    ) -> Result<()> {
622        if batches.is_empty() {
623            return Ok(());
624        }
625
626        let input_size = batches
627            .iter()
628            .map(RecordBatch::get_array_memory_size)
629            .sum::<usize>();
630        let _permit = limiter.acquire(input_size).await?;
631
632        let compacted = compact_record_batches(batches)?;
633        self.estimated_size += compacted.batch.get_array_memory_size();
634        self.cached_batches.push(compacted);
635        Ok(())
636    }
637
638    fn finish(self) -> RangeScanCacheValue {
639        RangeScanCacheValue::new(self.cached_batches, self.estimated_size)
640    }
641}
642
643fn compact_record_batches(batches: Vec<RecordBatch>) -> Result<CachedBatchSlice> {
644    debug_assert!(!batches.is_empty());
645
646    let slice_lengths = batches.iter().map(RecordBatch::num_rows).collect();
647    build_cached_batch_slice(batches, slice_lengths)
648}
649
650fn build_cached_batch_slice(
651    batches: Vec<RecordBatch>,
652    slice_lengths: Vec<usize>,
653) -> Result<CachedBatchSlice> {
654    let batch = if batches.len() == 1 {
655        batches.into_iter().next().unwrap()
656    } else {
657        let schema = batches[0].schema();
658        concat_batches(&schema, &batches).context(ComputeArrowSnafu)?
659    };
660
661    Ok(CachedBatchSlice {
662        batch,
663        slice_lengths,
664    })
665}
666
667async fn run_cache_concat_task(
668    mut rx: mpsc::UnboundedReceiver<CacheConcatCommand>,
669    limiter: Arc<crate::cache::RangeResultMemoryLimiter>,
670    skip_threshold_bytes: usize,
671) {
672    let mut state = CacheConcatState::default();
673
674    while let Some(cmd) = rx.recv().await {
675        match cmd {
676            CacheConcatCommand::Compact(batches) => {
677                if let Err(err) = state.compact(batches, &limiter).await {
678                    warn!(err; "Failed to compact range cache batches");
679                    return;
680                }
681                // Close the channel to stop further work as soon as the cached
682                // size exceeds the configured cache budget.
683                if state.estimated_size > skip_threshold_bytes {
684                    return;
685                }
686            }
687            CacheConcatCommand::Finish {
688                pending,
689                key,
690                cache_strategy,
691                part_metrics,
692                result_tx,
693            } => {
694                let compact_result = state
695                    .compact(pending, &limiter)
696                    .await
697                    .map(|()| state.finish());
698                let result = match compact_result {
699                    Ok(v) => {
700                        let value = Arc::new(v);
701                        part_metrics
702                            .inc_range_cache_size(key.estimated_size() + value.estimated_size());
703                        cache_strategy.put_range_result(key, value.clone());
704
705                        Ok(value)
706                    }
707                    Err(e) => {
708                        warn!(e; "Failed to finalize range cache batches");
709
710                        Err(e)
711                    }
712                };
713
714                if let Some(tx) = result_tx {
715                    let _ = tx.send(result);
716                }
717
718                break;
719            }
720        }
721    }
722}
723
724struct CacheBatchBuffer {
725    buffered_batches: Vec<RecordBatch>,
726    buffered_rows: usize,
727    buffered_size: usize,
728    sender: Option<mpsc::UnboundedSender<CacheConcatCommand>>,
729}
730
731impl CacheBatchBuffer {
732    fn new(cache_strategy: &CacheStrategy) -> Self {
733        let sender = cache_strategy.range_result_memory_limiter().map(|limiter| {
734            let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0);
735            let (tx, rx) = mpsc::unbounded_channel();
736            common_runtime::spawn_query(run_cache_concat_task(
737                rx,
738                limiter.clone(),
739                skip_threshold_bytes,
740            ));
741            tx
742        });
743
744        Self {
745            buffered_batches: Vec::new(),
746            buffered_rows: 0,
747            buffered_size: 0,
748            sender,
749        }
750    }
751
752    fn push(&mut self, batch: RecordBatch) -> Result<()> {
753        if self.sender.is_none() {
754            return Ok(());
755        }
756
757        self.buffered_rows += batch.num_rows();
758        self.buffered_size += batch.get_array_memory_size();
759        self.buffered_batches.push(batch);
760
761        if self.buffered_batches.len() > 1
762            && (self.buffered_rows > DEFAULT_READ_BATCH_SIZE
763                || self.buffered_size > RANGE_CACHE_COMPACT_THRESHOLD_BYTES)
764        {
765            self.notify_compact();
766        }
767
768        Ok(())
769    }
770
771    fn notify_compact(&mut self) {
772        if self.buffered_batches.is_empty() || self.sender.is_none() {
773            return;
774        }
775
776        let batches = mem::take(&mut self.buffered_batches);
777        self.buffered_rows = 0;
778        self.buffered_size = 0;
779
780        let Some(sender) = &self.sender else {
781            return;
782        };
783        if sender.send(CacheConcatCommand::Compact(batches)).is_err() {
784            self.sender = None;
785        }
786    }
787
788    fn finish(
789        mut self,
790        key: RangeScanCacheKey,
791        cache_strategy: CacheStrategy,
792        part_metrics: PartitionMetrics,
793        result_tx: Option<oneshot::Sender<Result<Arc<RangeScanCacheValue>>>>,
794    ) {
795        let Some(sender) = self.sender.take() else {
796            return;
797        };
798
799        if sender
800            .send(CacheConcatCommand::Finish {
801                pending: mem::take(&mut self.buffered_batches),
802                key,
803                cache_strategy,
804                part_metrics,
805                result_tx,
806            })
807            .is_err()
808        {
809            self.sender = None;
810        }
811    }
812}
813
814/// Wraps a stream to cache its output for future range cache hits.
815pub(crate) fn cache_flat_range_stream(
816    mut stream: BoxedRecordBatchStream,
817    cache_strategy: CacheStrategy,
818    key: RangeScanCacheKey,
819    part_metrics: PartitionMetrics,
820) -> BoxedRecordBatchStream {
821    Box::pin(try_stream! {
822        let mut buffer = CacheBatchBuffer::new(&cache_strategy);
823        while let Some(batch) = stream.try_next().await? {
824            buffer.push(batch.clone())?;
825            yield batch;
826        }
827
828        buffer.finish(key, cache_strategy, part_metrics, None);
829    })
830}
831
832/// Creates a `cache_flat_range_stream` with dummy internals for benchmarking.
833///
834/// This avoids exposing `RangeScanCacheKey`, `ScanRequestFingerprint`, and
835/// `PartitionMetrics` publicly.
836#[cfg(feature = "test")]
837pub fn bench_cache_flat_range_stream(
838    stream: BoxedRecordBatchStream,
839    cache_size_bytes: u64,
840    region_id: RegionId,
841) -> BoxedRecordBatchStream {
842    use std::time::Instant;
843
844    use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
845
846    use crate::region::options::MergeMode;
847
848    let cache_manager = Arc::new(
849        crate::cache::CacheManager::builder()
850            .range_result_cache_size(cache_size_bytes)
851            .build(),
852    );
853    let cache_strategy = CacheStrategy::EnableAll(cache_manager);
854
855    let fingerprint = ScanRequestFingerprintBuilder {
856        read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()),
857        read_column_types: vec![],
858        filters: vec![],
859        time_filters: vec![],
860        series_row_selector: None,
861        append_mode: false,
862        filter_deleted: false,
863        merge_mode: MergeMode::LastRow,
864        partition_expr_version: 0,
865    }
866    .build();
867
868    let key = RangeScanCacheKey {
869        region_id,
870        row_groups: vec![],
871        scan: fingerprint,
872    };
873
874    let metrics_set = ExecutionPlanMetricsSet::new();
875    let part_metrics =
876        PartitionMetrics::new(region_id, 0, "bench", Instant::now(), false, &metrics_set);
877
878    cache_flat_range_stream(stream, cache_strategy, key, part_metrics)
879}
880
881#[cfg(test)]
882mod tests {
883    use std::sync::Arc;
884    use std::time::Instant;
885
886    use common_time::Timestamp;
887    use common_time::range::TimestampRange;
888    use common_time::timestamp::TimeUnit;
889    use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
890    use datafusion_common::ScalarValue;
891    use datafusion_expr::{Expr, col, lit};
892    use smallvec::smallvec;
893    use store_api::storage::{FileId, RegionId};
894
895    use super::*;
896    use crate::cache::CacheManager;
897    use crate::read::flat_projection::FlatProjectionMapper;
898    use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex};
899    use crate::read::scan_region::{PredicateGroup, ScanInput};
900    use crate::sst::file::FileTimeRange;
901    use crate::test_util::memtable_util::metadata_with_primary_key;
902    use crate::test_util::scheduler_util::SchedulerEnv;
903    use crate::test_util::sst_util::sst_file_handle_with_file_id;
904
905    fn test_cache_strategy() -> CacheStrategy {
906        CacheStrategy::EnableAll(Arc::new(
907            CacheManager::builder()
908                .range_result_cache_size(1024 * 1024)
909                .build(),
910        ))
911    }
912
913    fn test_scan_fingerprint(
914        filters: Vec<String>,
915        time_filters: Vec<String>,
916        series_row_selector: Option<TimeSeriesRowSelector>,
917        filter_deleted: bool,
918        partition_expr_version: u64,
919    ) -> ScanRequestFingerprint {
920        let read_columns = ReadColumns::from_deduped_column_ids([1, 2]);
921        ScanRequestFingerprintBuilder {
922            read_columns,
923            read_column_types: vec![None, None],
924            filters,
925            time_filters,
926            series_row_selector,
927            append_mode: false,
928            filter_deleted,
929            merge_mode: MergeMode::LastRow,
930            partition_expr_version,
931        }
932        .build()
933    }
934
935    fn test_cache_context(strategy: &CacheStrategy) -> (RangeScanCacheKey, PartitionMetrics) {
936        let region_id = RegionId::new(1, 1);
937        let key = RangeScanCacheKey {
938            region_id,
939            row_groups: vec![],
940            scan: test_scan_fingerprint(vec![], vec![], None, false, 0),
941        };
942
943        let metrics_set = ExecutionPlanMetricsSet::new();
944        let part_metrics =
945            PartitionMetrics::new(region_id, 0, "test", Instant::now(), false, &metrics_set);
946
947        assert!(strategy.get_range_result(&key).is_none());
948        (key, part_metrics)
949    }
950
951    async fn finish_cache_batch_buffer(
952        buffer: CacheBatchBuffer,
953        key: RangeScanCacheKey,
954        cache_strategy: CacheStrategy,
955        part_metrics: PartitionMetrics,
956    ) -> Result<Arc<RangeScanCacheValue>> {
957        let (tx, rx) = oneshot::channel();
958        common_telemetry::info!("finish start");
959        buffer.finish(key, cache_strategy, part_metrics, Some(tx));
960        common_telemetry::info!("finish end");
961        rx.await.context(crate::error::RecvSnafu)?
962    }
963
964    async fn new_stream_context(
965        filters: Vec<Expr>,
966        query_time_range: Option<TimestampRange>,
967        partition_time_range: FileTimeRange,
968    ) -> (StreamContext, PartitionRange) {
969        let env = SchedulerEnv::new().await;
970        let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false));
971        let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap();
972        let predicate = PredicateGroup::new(metadata.as_ref(), &filters).unwrap();
973        let file_id = FileId::random();
974        let file = sst_file_handle_with_file_id(
975            file_id,
976            partition_time_range.0.value(),
977            partition_time_range.1.value(),
978        );
979        let input = ScanInput::new(env.access_layer.clone(), mapper)
980            .with_predicate(predicate)
981            .with_time_range(query_time_range)
982            .with_files(vec![file])
983            .with_cache(test_cache_strategy());
984        let range_meta = RangeMeta {
985            time_range: partition_time_range,
986            indices: smallvec![SourceIndex {
987                index: 0,
988                num_row_groups: 1,
989            }],
990            row_group_indices: smallvec![RowGroupIndex {
991                index: 0,
992                row_group_index: 0,
993            }],
994            num_rows: 10,
995        };
996        let partition_range = range_meta.new_partition_range(0);
997        let (scan_fingerprint, scan_implied_time_range) =
998            match crate::read::scan_region::build_scan_fingerprint(&input) {
999                Some(b) => (Some(b.fingerprint), b.implied_time_range),
1000                None => (None, None),
1001            };
1002        let stream_ctx = StreamContext {
1003            input,
1004            ranges: vec![range_meta],
1005            scan_fingerprint,
1006            scan_implied_time_range,
1007            query_start: Instant::now(),
1008        };
1009
1010        (stream_ctx, partition_range)
1011    }
1012
1013    /// Helper to create a timestamp millisecond literal.
1014    fn ts_lit(val: i64) -> Expr {
1015        lit(ScalarValue::TimestampMillisecond(Some(val), None))
1016    }
1017
1018    fn normalized_exprs(exprs: impl IntoIterator<Item = Expr>) -> Vec<String> {
1019        let mut exprs = exprs
1020            .into_iter()
1021            .map(|expr| expr.to_string())
1022            .collect::<Vec<_>>();
1023        exprs.sort_unstable();
1024        exprs
1025    }
1026
1027    async fn assert_range_cache_filters(
1028        filters: Vec<Expr>,
1029        query_time_range: Option<TimestampRange>,
1030        partition_time_range: FileTimeRange,
1031        expected_filters: Vec<Expr>,
1032        expected_time_filters: Vec<Expr>,
1033    ) {
1034        let (stream_ctx, part_range) =
1035            new_stream_context(filters, query_time_range, partition_time_range).await;
1036
1037        let key = build_range_cache_key(&stream_ctx, &part_range).unwrap();
1038
1039        assert_eq!(
1040            key.scan.filters(),
1041            normalized_exprs(expected_filters).as_slice()
1042        );
1043        assert_eq!(
1044            key.scan.time_filters(),
1045            normalized_exprs(expected_time_filters).as_slice()
1046        );
1047    }
1048
1049    #[tokio::test]
1050    async fn range_cache_time_filter_key_cases() {
1051        let partition = (
1052            Timestamp::new_millisecond(1000),
1053            Timestamp::new_millisecond(2000),
1054        );
1055
1056        struct Case {
1057            filters: Vec<Expr>,
1058            query_time_range: Option<TimestampRange>,
1059            expected_filters: Vec<Expr>,
1060            expected_time_filters: Vec<Expr>,
1061        }
1062
1063        // Time filters are stripped only when their implied range fully covers
1064        // the partition's file-time range. `is_not_null(ts)` stays in regular
1065        // filters because it is not routed into `time_filters`.
1066        for case in [
1067            Case {
1068                filters: vec![
1069                    col("ts").gt_eq(ts_lit(1000)),
1070                    col("ts").lt(ts_lit(2001)),
1071                    col("ts").is_not_null(),
1072                    col("k0").eq(lit("foo")),
1073                ],
1074                query_time_range: TimestampRange::with_unit(1000, 2002, TimeUnit::Millisecond),
1075                expected_filters: vec![col("k0").eq(lit("foo")), col("ts").is_not_null()],
1076                expected_time_filters: vec![],
1077            },
1078            Case {
1079                filters: vec![
1080                    col("ts").gt_eq(ts_lit(500)),
1081                    col("ts").lt(ts_lit(3000)),
1082                    col("k0").eq(lit("foo")),
1083                ],
1084                query_time_range: TimestampRange::with_unit(500, 3000, TimeUnit::Millisecond),
1085                expected_filters: vec![col("k0").eq(lit("foo"))],
1086                expected_time_filters: vec![],
1087            },
1088            Case {
1089                filters: vec![
1090                    col("ts").gt_eq(ts_lit(1000)),
1091                    col("ts").lt_eq(ts_lit(2000)),
1092                    col("k0").eq(lit("foo")),
1093                ],
1094                query_time_range: TimestampRange::with_unit(1000, 2001, TimeUnit::Millisecond),
1095                expected_filters: vec![col("k0").eq(lit("foo"))],
1096                expected_time_filters: vec![],
1097            },
1098            Case {
1099                filters: vec![
1100                    col("ts").between(ts_lit(1000), ts_lit(2000)),
1101                    col("k0").eq(lit("foo")),
1102                ],
1103                query_time_range: TimestampRange::with_unit(1000, 2001, TimeUnit::Millisecond),
1104                expected_filters: vec![col("k0").eq(lit("foo"))],
1105                expected_time_filters: vec![],
1106            },
1107            Case {
1108                filters: vec![col("ts").gt_eq(ts_lit(1200)), col("k0").eq(lit("foo"))],
1109                query_time_range: TimestampRange::with_unit(1200, 2001, TimeUnit::Millisecond),
1110                expected_filters: vec![col("k0").eq(lit("foo"))],
1111                expected_time_filters: vec![col("ts").gt_eq(ts_lit(1200))],
1112            },
1113            Case {
1114                filters: vec![
1115                    col("ts").gt_eq(ts_lit(1500)),
1116                    col("ts").is_not_null(),
1117                    col("k0").eq(lit("foo")),
1118                ],
1119                query_time_range: None,
1120                expected_filters: vec![col("k0").eq(lit("foo")), col("ts").is_not_null()],
1121                expected_time_filters: vec![col("ts").gt_eq(ts_lit(1500))],
1122            },
1123        ] {
1124            assert_range_cache_filters(
1125                case.filters,
1126                case.query_time_range,
1127                partition,
1128                case.expected_filters,
1129                case.expected_time_filters,
1130            )
1131            .await;
1132        }
1133    }
1134
1135    #[tokio::test]
1136    async fn two_distinct_queries_share_cache_key_when_both_cover() {
1137        let partition_range = (
1138            Timestamp::new_millisecond(1000),
1139            Timestamp::new_millisecond(2000),
1140        );
1141
1142        let (ctx_a, part_a) = new_stream_context(
1143            vec![
1144                col("ts").gt_eq(ts_lit(500)),
1145                col("ts").lt(ts_lit(3000)),
1146                col("k0").eq(lit("foo")),
1147            ],
1148            TimestampRange::with_unit(500, 3000, TimeUnit::Millisecond),
1149            partition_range,
1150        )
1151        .await;
1152        let (ctx_b, part_b) = new_stream_context(
1153            vec![
1154                col("ts").gt_eq(ts_lit(100)),
1155                col("ts").lt(ts_lit(5000)),
1156                col("k0").eq(lit("foo")),
1157            ],
1158            TimestampRange::with_unit(100, 5000, TimeUnit::Millisecond),
1159            partition_range,
1160        )
1161        .await;
1162
1163        let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap();
1164        let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap();
1165        assert_eq!(key_a.scan, key_b.scan);
1166        assert!(key_a.scan.time_filters().is_empty());
1167    }
1168
1169    #[tokio::test]
1170    async fn scan_stages_have_separate_cache_keys() {
1171        let partition_range = (
1172            Timestamp::new_millisecond(1000),
1173            Timestamp::new_millisecond(2000),
1174        );
1175        let (ctx, part_range) =
1176            new_stream_context(vec![col("k0").eq(lit("foo"))], None, partition_range).await;
1177
1178        let data_key = build_range_cache_key(&ctx, &part_range).unwrap();
1179        let candidate_key = build_candidate_range_cache_key(&ctx, &part_range).unwrap();
1180        let range_0 = SeriesRange::new(0, 2).unwrap();
1181        let range_1 = SeriesRange::new(1, 2).unwrap();
1182        let series_key_0 = build_series_range_cache_key(&ctx, &part_range, range_0).unwrap();
1183        let series_key_1 = build_series_range_cache_key(&ctx, &part_range, range_1).unwrap();
1184
1185        assert_ne!(data_key.scan, candidate_key.scan);
1186        assert_ne!(data_key.scan, series_key_0.scan);
1187        assert_ne!(candidate_key.scan, series_key_0.scan);
1188        assert_ne!(series_key_0.scan, series_key_1.scan);
1189        assert_eq!(data_key.row_groups, candidate_key.row_groups);
1190        assert_eq!(data_key.row_groups, series_key_0.row_groups);
1191    }
1192
1193    #[tokio::test]
1194    async fn disables_optimization_on_or_clause() {
1195        let partition_range = (
1196            Timestamp::new_millisecond(1000),
1197            Timestamp::new_millisecond(2000),
1198        );
1199
1200        let or_a = col("ts").gt_eq(ts_lit(1000)).or(col("ts").lt(ts_lit(500)));
1201        let or_b = col("ts").gt_eq(ts_lit(900)).or(col("ts").lt(ts_lit(400)));
1202
1203        let (ctx_a, part_a) = new_stream_context(
1204            vec![or_a.clone(), col("k0").eq(lit("foo"))],
1205            None,
1206            partition_range,
1207        )
1208        .await;
1209        let (ctx_b, part_b) = new_stream_context(
1210            vec![or_b.clone(), col("k0").eq(lit("foo"))],
1211            None,
1212            partition_range,
1213        )
1214        .await;
1215
1216        assert!(ctx_a.scan_implied_time_range.is_none());
1217        let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap();
1218        let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap();
1219        assert_ne!(key_a.scan, key_b.scan);
1220        assert_eq!(
1221            key_a.scan.time_filters(),
1222            normalized_exprs([or_a]).as_slice()
1223        );
1224    }
1225
1226    #[tokio::test]
1227    async fn empty_implied_range_does_not_panic_on_non_ms_file_range() {
1228        // Contradictory time predicates make the implied range empty. The
1229        // empty range's sentinel timestamps use `Timestamp::default()` (ms),
1230        // so without the `is_empty()` short-circuit the unit asserts would
1231        // panic against a non-ms `range_meta.time_range`.
1232        let partition = (
1233            Timestamp::new_millisecond(1000),
1234            Timestamp::new_millisecond(2000),
1235        );
1236
1237        let (mut ctx, part_range) = new_stream_context(
1238            vec![col("ts").gt_eq(ts_lit(1500)), col("k0").eq(lit("foo"))],
1239            TimestampRange::with_unit(1500, 3000, TimeUnit::Millisecond),
1240            partition,
1241        )
1242        .await;
1243
1244        ctx.scan_implied_time_range = Some(TimestampRange::empty());
1245        ctx.ranges[0].time_range = (
1246            Timestamp::new(1_000_000_000, TimeUnit::Nanosecond),
1247            Timestamp::new(2_000_000_000, TimeUnit::Nanosecond),
1248        );
1249
1250        let key = build_range_cache_key(&ctx, &part_range).unwrap();
1251        // Empty implied range cannot cover, so time filters stay in the key.
1252        assert!(!key.scan.time_filters().is_empty());
1253    }
1254
1255    fn ms_ts(v: i64) -> Timestamp {
1256        Timestamp::new_millisecond(v)
1257    }
1258
1259    fn implied_ms(expr: Expr) -> Option<TimestampRange> {
1260        implied_time_range_from_exprs("ts", TimeUnit::Millisecond, &[&expr])
1261    }
1262
1263    #[test]
1264    fn implied_time_range_supported_exprs() {
1265        for (expr, expected) in [
1266            (
1267                col("ts").gt_eq(ts_lit(1000)),
1268                Some(TimestampRange::from_start(ms_ts(1000))),
1269            ),
1270            (
1271                col("ts").gt(ts_lit(1000)),
1272                Some(TimestampRange::from_start(ms_ts(1001))),
1273            ),
1274            (
1275                col("ts").lt_eq(ts_lit(2000)),
1276                Some(TimestampRange::until_end(ms_ts(2000), true)),
1277            ),
1278            (
1279                col("ts").lt(ts_lit(2000)),
1280                Some(TimestampRange::until_end(ms_ts(2000), false)),
1281            ),
1282            (
1283                col("ts").eq(ts_lit(1500)),
1284                Some(TimestampRange::single(ms_ts(1500))),
1285            ),
1286            (
1287                ts_lit(1000).lt_eq(col("ts")),
1288                Some(TimestampRange::from_start(ms_ts(1000))),
1289            ),
1290            (
1291                col("ts").between(ts_lit(1000), ts_lit(2000)),
1292                Some(TimestampRange::new_inclusive(
1293                    Some(ms_ts(1000)),
1294                    Some(ms_ts(2000)),
1295                )),
1296            ),
1297            (
1298                col("ts")
1299                    .gt_eq(ts_lit(1000))
1300                    .and(col("ts").lt(ts_lit(2000))),
1301                TimestampRange::with_unit(1000, 2000, TimeUnit::Millisecond),
1302            ),
1303            (
1304                col("ts")
1305                    .gt_eq(ts_lit(1000))
1306                    .and(col("ts").lt(ts_lit(5000)))
1307                    .and(col("ts").lt_eq(ts_lit(3000))),
1308                TimestampRange::with_unit(1000, 3001, TimeUnit::Millisecond),
1309            ),
1310        ] {
1311            assert_eq!(implied_ms(expr), expected);
1312        }
1313
1314        assert_eq!(
1315            implied_time_range_from_exprs("ts", TimeUnit::Millisecond, &[]),
1316            Some(TimestampRange::min_to_max())
1317        );
1318    }
1319
1320    #[test]
1321    fn implied_time_range_unsupported_exprs() {
1322        let not_between = Expr::Between(Between {
1323            expr: Box::new(col("ts")),
1324            negated: true,
1325            low: Box::new(ts_lit(1000)),
1326            high: Box::new(ts_lit(2000)),
1327        });
1328
1329        for expr in [
1330            not_between,
1331            col("ts").gt_eq(ts_lit(1000)).or(col("ts").lt(ts_lit(500))),
1332            Expr::Not(Box::new(col("ts").gt_eq(ts_lit(1000)))),
1333            col("ts").in_list(vec![ts_lit(1000), ts_lit(2000)], false),
1334            col("ts").gt_eq(col("other")),
1335            col("other_ts").gt_eq(ts_lit(1000)),
1336        ] {
1337            assert!(implied_ms(expr).is_none());
1338        }
1339    }
1340
1341    #[test]
1342    fn implied_time_range_unit_conversion() {
1343        let second_1 = lit(ScalarValue::TimestampSecond(Some(1), None));
1344        let ns_1500 = lit(ScalarValue::TimestampNanosecond(Some(1_500_000_000), None));
1345        let ns_1500_5 = lit(ScalarValue::TimestampNanosecond(Some(1_500_500_000), None));
1346
1347        for (expr, expected) in [
1348            (
1349                col("ts").gt_eq(second_1.clone()),
1350                Some(TimestampRange::from_start(ms_ts(1000))),
1351            ),
1352            (
1353                col("ts").lt_eq(second_1),
1354                Some(TimestampRange::until_end(ms_ts(1000), true)),
1355            ),
1356            (
1357                col("ts").eq(ns_1500),
1358                Some(TimestampRange::single(ms_ts(1500))),
1359            ),
1360            (col("ts").eq(ns_1500_5.clone()), None),
1361            (
1362                col("ts").gt_eq(ns_1500_5.clone()),
1363                Some(TimestampRange::from_start(ms_ts(1501))),
1364            ),
1365            (
1366                col("ts").lt_eq(ns_1500_5.clone()),
1367                Some(TimestampRange::until_end(ms_ts(1500), true)),
1368            ),
1369            (
1370                col("ts").gt(ns_1500_5.clone()),
1371                Some(TimestampRange::from_start(ms_ts(1501))),
1372            ),
1373            (
1374                col("ts").lt(ns_1500_5),
1375                Some(TimestampRange::until_end(ms_ts(1501), false)),
1376            ),
1377        ] {
1378            assert_eq!(implied_ms(expr), expected);
1379        }
1380    }
1381
1382    #[test]
1383    fn normalizes_and_clears_time_filters() {
1384        let normalized =
1385            test_scan_fingerprint(vec!["k0 = 'foo'".to_string()], vec![], None, true, 0);
1386
1387        assert!(normalized.time_filters().is_empty());
1388
1389        let fingerprint = test_scan_fingerprint(
1390            vec!["k0 = 'foo'".to_string()],
1391            vec!["ts >= 1000".to_string()],
1392            Some(TimeSeriesRowSelector::LastRow),
1393            true,
1394            7,
1395        );
1396
1397        let reset = fingerprint.without_time_filters();
1398
1399        assert_eq!(reset.read_columns(), fingerprint.read_columns());
1400        assert_eq!(reset.read_column_types(), fingerprint.read_column_types());
1401        assert_eq!(reset.filters(), fingerprint.filters());
1402        assert!(reset.time_filters().is_empty());
1403        assert_eq!(reset.series_row_selector, fingerprint.series_row_selector);
1404        assert_eq!(reset.append_mode, fingerprint.append_mode);
1405        assert_eq!(reset.filter_deleted, fingerprint.filter_deleted);
1406        assert_eq!(reset.merge_mode, fingerprint.merge_mode);
1407        assert_eq!(
1408            reset.partition_expr_version,
1409            fingerprint.partition_expr_version
1410        );
1411    }
1412
1413    fn test_schema() -> Arc<datatypes::arrow::datatypes::Schema> {
1414        use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
1415
1416        Arc::new(Schema::new(vec![Field::new(
1417            "value",
1418            ArrowDataType::Int64,
1419            false,
1420        )]))
1421    }
1422
1423    fn make_batch(values: &[i64]) -> RecordBatch {
1424        use datatypes::arrow::array::Int64Array;
1425
1426        RecordBatch::try_new(
1427            test_schema(),
1428            vec![Arc::new(Int64Array::from(values.to_vec()))],
1429        )
1430        .unwrap()
1431    }
1432
1433    fn make_large_binary_batch(rows: usize, bytes_per_row: usize) -> RecordBatch {
1434        use datatypes::arrow::array::BinaryArray;
1435        use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
1436
1437        let schema = Arc::new(Schema::new(vec![Field::new(
1438            "value",
1439            ArrowDataType::Binary,
1440            false,
1441        )]));
1442        let payload = vec![b'x'; bytes_per_row];
1443        let values = (0..rows).map(|_| payload.as_slice()).collect::<Vec<_>>();
1444
1445        RecordBatch::try_new(schema, vec![Arc::new(BinaryArray::from_vec(values))]).unwrap()
1446    }
1447
1448    #[test]
1449    fn compact_record_batches_keeps_original_boundaries() {
1450        let batches = vec![make_batch(&[1, 2]), make_batch(&[3]), make_batch(&[4, 5])];
1451
1452        let compacted = compact_record_batches(batches).unwrap();
1453
1454        assert_eq!(compacted.batch.num_rows(), 5);
1455        assert_eq!(compacted.slice_lengths, vec![2, 1, 2]);
1456    }
1457
1458    #[tokio::test]
1459    async fn cached_flat_range_stream_replays_original_batches() {
1460        let value = Arc::new(RangeScanCacheValue::new(
1461            vec![CachedBatchSlice {
1462                batch: make_batch(&[1, 2, 3]),
1463                slice_lengths: vec![2, 1],
1464            }],
1465            make_batch(&[1, 2, 3]).get_array_memory_size(),
1466        ));
1467
1468        let replayed = cached_flat_range_stream(value)
1469            .try_collect::<Vec<_>>()
1470            .await
1471            .unwrap();
1472
1473        assert_eq!(replayed.len(), 2);
1474        assert_eq!(replayed[0].num_rows(), 2);
1475        assert_eq!(replayed[1].num_rows(), 1);
1476    }
1477
1478    #[tokio::test]
1479    async fn cache_batch_buffer_finishes_pending_batches() {
1480        let strategy = test_cache_strategy();
1481        let batch = make_batch(&[1, 2, 3]);
1482        let expected_size = batch.get_array_memory_size();
1483        let (key, part_metrics) = test_cache_context(&strategy);
1484
1485        let mut buffer = CacheBatchBuffer::new(&strategy);
1486        buffer.push(batch).unwrap();
1487
1488        let value = finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics)
1489            .await
1490            .unwrap();
1491        assert_eq!(value.cached_batches.len(), 1);
1492        assert_eq!(value.cached_batches[0].slice_lengths, vec![3]);
1493        assert_eq!(value.estimated_batches_size, expected_size);
1494        assert!(Arc::ptr_eq(
1495            &value,
1496            &strategy.get_range_result(&key).unwrap()
1497        ));
1498    }
1499
1500    #[tokio::test]
1501    async fn cache_batch_buffer_compacts_when_rows_exceed_default_batch_size() {
1502        let strategy = test_cache_strategy();
1503        let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]);
1504        let (key, part_metrics) = test_cache_context(&strategy);
1505
1506        let mut buffer = CacheBatchBuffer::new(&strategy);
1507        buffer.push(batch.clone()).unwrap();
1508        buffer.push(batch).unwrap();
1509
1510        assert_eq!(buffer.buffered_rows, 0);
1511        assert!(buffer.buffered_batches.is_empty());
1512
1513        let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics)
1514            .await
1515            .unwrap();
1516        assert_eq!(value.cached_batches.len(), 1);
1517        assert_eq!(
1518            value.cached_batches[0].slice_lengths,
1519            vec![
1520                DEFAULT_READ_BATCH_SIZE / 2 + 1,
1521                DEFAULT_READ_BATCH_SIZE / 2 + 1
1522            ]
1523        );
1524    }
1525
1526    #[tokio::test]
1527    async fn cache_batch_buffer_compacts_when_buffered_size_exceeds_threshold() {
1528        let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE, 4096);
1529        let strategy = CacheStrategy::EnableAll(Arc::new(
1530            CacheManager::builder()
1531                .range_result_cache_size((large_batch.get_array_memory_size() * 3) as u64)
1532                .build(),
1533        ));
1534        let (key, part_metrics) = test_cache_context(&strategy);
1535
1536        let mut buffer = CacheBatchBuffer::new(&strategy);
1537        buffer.push(large_batch.clone()).unwrap();
1538
1539        assert_eq!(buffer.buffered_rows, large_batch.num_rows());
1540        assert_eq!(buffer.buffered_batches.len(), 1);
1541
1542        buffer.push(large_batch.clone()).unwrap();
1543
1544        assert_eq!(buffer.buffered_rows, 0);
1545        assert!(buffer.buffered_batches.is_empty());
1546
1547        let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics)
1548            .await
1549            .unwrap();
1550        assert_eq!(value.cached_batches.len(), 1);
1551        assert_eq!(
1552            value.cached_batches[0].slice_lengths,
1553            vec![large_batch.num_rows(), large_batch.num_rows()]
1554        );
1555    }
1556
1557    #[tokio::test]
1558    async fn cache_batch_buffer_skips_cache_when_compacted_size_exceeds_limit() {
1559        let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE / 2 + 1, 4096);
1560        // Budget only fits two large batches.
1561        let budget = (large_batch.get_array_memory_size() as u64) * 2 + 1;
1562        let strategy = CacheStrategy::EnableAll(Arc::new(
1563            CacheManager::builder()
1564                .range_result_cache_size(budget)
1565                .build(),
1566        ));
1567        let (key, part_metrics) = test_cache_context(&strategy);
1568
1569        let mut buffer = CacheBatchBuffer::new(&strategy);
1570        for _ in 0..4 {
1571            buffer.push(large_batch.clone()).unwrap();
1572        }
1573        assert!(
1574            finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics)
1575                .await
1576                .is_err()
1577        );
1578        assert!(strategy.get_range_result(&key).is_none());
1579    }
1580}
mito2/read/range_cache.rs

mito2/read/
range_cache.rs
