mito2/sst/parquet/
push_decoder.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//! Push decoder stream implementation for SST parquet files.
16
17use std::io;
18use std::ops::Range;
19use std::sync::{Arc, Mutex};
20use std::time::Instant;
21
22use async_trait::async_trait;
23use common_base::range_read::{Metadata, RangeReader};
24use tokio::sync::OnceCell;
25
26use bytes::Bytes;
27use common_recordbatch::QueryMemoryTracker;
28use datatypes::arrow::record_batch::RecordBatch;
29use futures::StreamExt;
30use futures::stream::BoxStream;
31use object_store::ObjectStore;
32use parquet::DecodeResult;
33use parquet::arrow::ProjectionMask;
34use parquet::arrow::arrow_reader::{ArrowReaderMetadata, RowSelection};
35use parquet::arrow::push_decoder::ParquetPushDecoderBuilder;
36use snafu::{IntoError, OptionExt, ResultExt, ensure};
37
38use crate::cache::CacheStrategy;
39use crate::cache::file_cache::{FileType, IndexKey};
40use crate::error::{
41    Error, LoadParquetPageSnafu, OpenDalSnafu, ReadParquetSnafu, Result, UnexpectedSnafu,
42};
43use crate::metrics::{READ_STAGE_ELAPSED, READ_STAGE_FETCH_PAGES};
44use crate::sst::file::RegionFileId;
45use crate::sst::parquet::helper::fetch_byte_ranges;
46use crate::sst::parquet::row_group::{ParquetFetchMetrics, compute_total_range_size};
47
48/// Fetches parquet byte ranges through Greptime's cache hierarchy.
49///
50/// The push decoder decides which ranges are required for decoding, while this
51/// fetcher keeps cache lookup, local write-cache reads, and remote I/O explicit
52/// in Greptime code.
53pub struct SstParquetRangeFetcher {
54    /// Region file ID for cache key.
55    region_file_id: RegionFileId,
56    /// Path to the parquet file in object storage.
57    file_path: String,
58    /// Object store for reading data.
59    object_store: ObjectStore,
60    /// Cache strategy for reading pages.
61    cache_strategy: CacheStrategy,
62    /// Immutable SST size, resolved once when no manifest hint is available.
63    file_size: OnceCell<u64>,
64    file_size_hint: u64,
65    /// Optional metrics for tracking fetch operations.
66    fetch_metrics: Option<ParquetFetchMetrics>,
67    /// Shared quota for decoding this fetcher's row group.
68    memory_tracker: Option<QueryMemoryTracker>,
69}
70
71impl SstParquetRangeFetcher {
72    /// Creates a new [SstParquetRangeFetcher].
73    pub fn new(
74        region_file_id: RegionFileId,
75        file_path: String,
76        object_store: ObjectStore,
77        cache_strategy: CacheStrategy,
78        fetch_metrics: Option<ParquetFetchMetrics>,
79    ) -> Self {
80        Self {
81            region_file_id,
82            file_path,
83            object_store,
84            cache_strategy,
85            file_size: OnceCell::new(),
86            file_size_hint: 0,
87            fetch_metrics,
88            memory_tracker: None,
89        }
90    }
91
92    /// Sets the shared quota for this decoder's projected row group.
93    pub(crate) fn with_memory_tracker(mut self, tracker: Option<QueryMemoryTracker>) -> Self {
94        self.memory_tracker = tracker;
95        self
96    }
97
98    /// Supplies the immutable SST size from its manifest entry.
99    pub(crate) fn with_file_size(mut self, size: u64) -> Self {
100        self.file_size_hint = size;
101        self
102    }
103
104    async fn resolve_file_size(&self, ranges: &[Range<u64>]) -> Result<u64> {
105        self.file_size
106            .get_or_try_init(|| async {
107                if self.file_size_hint != 0 {
108                    return Ok(self.file_size_hint);
109                }
110                if let Some(cache) = self.cache_strategy.parquet_page_cache()
111                    && let Some(size) = cache.known_file_size(self.region_file_id, ranges).await
112                {
113                    return Ok(size);
114                }
115                if let Some(cache) = self.cache_strategy.write_cache() {
116                    let key = IndexKey::new(
117                        self.region_file_id.region_id(),
118                        self.region_file_id.file_id(),
119                        FileType::Parquet,
120                    );
121                    let files = cache.file_cache();
122                    if files.get(key).await.is_some() {
123                        // The write-cache index uses u32 sizes, so stat the actual
124                        // file instead. Eviction or local I/O failure falls back below.
125                        if let Ok(meta) =
126                            files.local_store().stat(&files.cache_file_path(key)).await
127                        {
128                            return Ok(meta.content_length());
129                        }
130                    }
131                }
132                Ok(self
133                    .object_store
134                    .stat(&self.file_path)
135                    .await
136                    .context(OpenDalSnafu)?
137                    .content_length())
138            })
139            .await
140            .copied()
141    }
142
143    /// Fetches decoder ranges through fixed-size pages and the existing local/remote hierarchy.
144    async fn fetch_bytes_with_cache(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
145        let start = self.fetch_metrics.as_ref().map(|_| Instant::now());
146        let _timer = READ_STAGE_FETCH_PAGES.start_timer();
147        ensure!(
148            ranges.iter().all(|r| r.start <= r.end),
149            UnexpectedSnafu {
150                reason: "Reversed Parquet byte range"
151            }
152        );
153        if ranges.iter().all(Range::is_empty) {
154            return Ok(vec![Bytes::new(); ranges.len()]);
155        }
156        let reader = ParquetRangeReader {
157            fetcher: self,
158            local: self.fetch_metrics.as_ref().map(|_| Mutex::default()),
159            remote: self.fetch_metrics.as_ref().map(|_| Mutex::default()),
160        };
161        let output = if let Some(cache) = self.cache_strategy.parquet_page_cache() {
162            let size = self.resolve_file_size(ranges).await?;
163            let (output, stats) = cache
164                .fetch(self.region_file_id, size, ranges, &reader)
165                .await
166                .map_err(|error| {
167                    if let Some(source) = crate::cache::find_shared_page_source::<Error>(&error) {
168                        LoadParquetPageSnafu.into_error(source)
169                    } else {
170                        ReadParquetSnafu {
171                            path: &self.file_path,
172                        }
173                        .into_error(parquet::errors::ParquetError::External(Box::new(error)))
174                    }
175                })?;
176            if let Some(metrics) = &self.fetch_metrics {
177                let mut data = metrics.data.lock().unwrap();
178                data.page_cache_hit += usize::from(stats.cached_bytes > 0);
179                data.pages_to_fetch_mem += stats.cached_fragments;
180                data.page_size_to_fetch_mem += stats.cached_bytes;
181            }
182            output
183        } else {
184            reader.read_ranges(ranges).await?
185        };
186        if let (Some(metrics), Some(local), Some(remote), Some(start)) =
187            (&self.fetch_metrics, &reader.local, &reader.remote, start)
188        {
189            let mut data = metrics.data.lock().unwrap();
190            let local = local.lock().unwrap();
191            let remote = remote.lock().unwrap();
192            data.write_cache_hit += usize::from(local.ranges != 0);
193            data.cache_miss += usize::from(remote.ranges != 0);
194            data.pages_to_fetch_write_cache += local.ranges;
195            data.page_size_to_fetch_write_cache += local.bytes;
196            data.write_cache_fetch_elapsed += local.elapsed();
197            data.pages_to_fetch_store += remote.ranges;
198            data.page_size_to_fetch_store += remote.bytes;
199            data.store_fetch_elapsed += remote.elapsed();
200            data.page_size_needed += output.iter().map(|b| b.len() as u64).sum::<u64>();
201            data.total_fetch_elapsed += start.elapsed();
202        }
203        Ok(output)
204    }
205}
206
207/// Aggregates one fetch's physical reads without turning request hit counters
208/// into per-I/O counters. Elapsed time spans concurrent reads rather than summing them.
209#[derive(Default)]
210struct StorageReadMetrics {
211    ranges: usize,
212    bytes: u64,
213    interval: Option<(Instant, Instant)>,
214}
215
216impl StorageReadMetrics {
217    fn record(&mut self, ranges: usize, bytes: u64, start: Instant) {
218        self.ranges += ranges;
219        self.bytes += bytes;
220        let end = Instant::now();
221        self.interval = Some(match self.interval {
222            Some((first, last)) => (first.min(start), last.max(end)),
223            None => (start, end),
224        });
225    }
226
227    fn elapsed(&self) -> std::time::Duration {
228        self.interval
229            .map(|(start, end)| end - start)
230            .unwrap_or_default()
231    }
232}
233
234/// Reads one immutable SST through the local write cache or object store.
235/// Both cached and uncached reads use the existing range merging.
236struct ParquetRangeReader<'a> {
237    fetcher: &'a SstParquetRangeFetcher,
238    local: Option<Mutex<StorageReadMetrics>>,
239    remote: Option<Mutex<StorageReadMetrics>>,
240}
241
242impl ParquetRangeReader<'_> {
243    async fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
244        let fetcher = self.fetcher;
245        let key = IndexKey::new(
246            fetcher.region_file_id.region_id(),
247            fetcher.region_file_id.file_id(),
248            FileType::Parquet,
249        );
250        let bytes = if fetcher.fetch_metrics.is_none() {
251            0
252        } else {
253            compute_total_range_size(ranges).1
254        };
255        let start = self.local.as_ref().map(|_| Instant::now());
256        if let Some(cache) = fetcher.cache_strategy.write_cache()
257            && let Some(data) = cache.file_cache().read_ranges(key, ranges).await
258        {
259            if let (Some(metrics), Some(start)) = (&self.local, start) {
260                metrics.lock().unwrap().record(ranges.len(), bytes, start);
261            }
262            return Ok(data);
263        }
264        let _timer = READ_STAGE_ELAPSED
265            .with_label_values(&["cache_miss_read"])
266            .start_timer();
267        let start = self.remote.as_ref().map(|_| Instant::now());
268        let data = fetch_byte_ranges(&fetcher.file_path, fetcher.object_store.clone(), ranges)
269            .await
270            .context(OpenDalSnafu)?;
271        if let (Some(metrics), Some(start)) = (&self.remote, start) {
272            metrics.lock().unwrap().record(ranges.len(), bytes, start);
273        }
274        Ok(data)
275    }
276}
277
278#[async_trait]
279impl RangeReader for ParquetRangeReader<'_> {
280    async fn metadata(&self) -> io::Result<Metadata> {
281        self.fetcher
282            .resolve_file_size(&[])
283            .await
284            .map(|content_length| Metadata { content_length })
285            .map_err(|error| io::Error::other(Arc::new(error)))
286    }
287
288    async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
289        self.read_ranges(ranges)
290            .await
291            .map_err(|error| io::Error::other(Arc::new(error)))
292    }
293
294    async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
295        let mut data = self
296            .read_ranges(std::slice::from_ref(&range))
297            .await
298            .map_err(|error| io::Error::other(Arc::new(error)))?;
299        if data.len() != 1 {
300            return Err(io::Error::new(
301                io::ErrorKind::InvalidData,
302                "Invalid Parquet read response count",
303            ));
304        }
305        data.pop().ok_or_else(|| {
306            io::Error::new(
307                io::ErrorKind::UnexpectedEof,
308                "Missing Parquet read response",
309            )
310        })
311    }
312}
313
314/// Builds a parquet record batch stream driven directly by [ParquetPushDecoderBuilder].
315pub fn build_sst_parquet_record_batch_stream(
316    arrow_metadata: ArrowReaderMetadata,
317    row_group_idx: usize,
318    row_selection: Option<RowSelection>,
319    projection: ProjectionMask,
320    fetcher: SstParquetRangeFetcher,
321    file_path: String,
322    batch_size: usize,
323) -> Result<BoxStream<'static, Result<RecordBatch>>> {
324    let uncompressed_bytes = if fetcher.memory_tracker.is_some() {
325        let row_group = arrow_metadata.metadata().row_group(row_group_idx);
326        row_group
327            .columns()
328            .iter()
329            .enumerate()
330            .filter(|(index, _)| projection.leaf_included(*index))
331            .try_fold(0usize, |total, (_, column)| {
332                usize::try_from(column.uncompressed_size())
333                    .ok()
334                    .and_then(|size| total.checked_add(size))
335            })
336            .context(UnexpectedSnafu {
337                reason: "Invalid projected row-group uncompressed size",
338            })?
339    } else {
340        0
341    };
342    let mut builder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata)
343        .with_row_groups(vec![row_group_idx])
344        .with_projection(projection)
345        .with_batch_size(batch_size);
346
347    if let Some(selection) = row_selection {
348        builder = builder.with_row_selection(selection);
349    }
350
351    let mut decoder = builder
352        .build()
353        .context(ReadParquetSnafu { path: &file_path })?;
354
355    Ok(async_stream::try_stream! {
356        let _reservation = if let Some(tracker) = &fetcher.memory_tracker {
357            Some(tracker.reserve(uncompressed_bytes)
358                .context(crate::error::RecordBatchSnafu)?)
359        } else {
360            None
361        };
362        loop {
363            match decoder.try_decode().context(ReadParquetSnafu { path: &file_path })? {
364                DecodeResult::NeedsData(ranges) => {
365                    let data = fetcher.fetch_bytes_with_cache(&ranges).await?;
366                    decoder.push_ranges(ranges, data)
367                        .context(ReadParquetSnafu { path: &file_path })?;
368                }
369                DecodeResult::Data(batch) => yield batch,
370                DecodeResult::Finished => break,
371            }
372        }
373    }
374    .boxed())
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380
381    #[tokio::test]
382    async fn test_decoder_memory_projection_lifetime_and_error() {
383        use std::sync::Arc;
384
385        use datatypes::arrow::array::{Int64Array, StringArray};
386        use datatypes::arrow::datatypes::{DataType, Field, Schema};
387        use object_store::services::Memory;
388        use parquet::arrow::ArrowWriter;
389        use parquet::arrow::arrow_reader::ArrowReaderOptions;
390        use parquet::file::properties::WriterProperties;
391
392        use crate::sst::parquet::metadata::MetadataLoader;
393        use crate::sst::parquet::reader::MetadataCacheMetrics;
394
395        let schema = Arc::new(Schema::new(vec![
396            Field::new("value", DataType::Int64, false),
397            Field::new("payload", DataType::Utf8, false),
398        ]));
399        let batch = RecordBatch::try_new(
400            schema.clone(),
401            vec![
402                Arc::new(Int64Array::from_iter_values(0..8)),
403                Arc::new(StringArray::from_iter_values(
404                    (0..8).map(|i| format!("{i}:{}", "x".repeat(512))),
405                )),
406            ],
407        )
408        .unwrap();
409        let mut bytes = Vec::new();
410        let props = WriterProperties::builder()
411            .set_max_row_group_row_count(Some(4))
412            .build();
413        let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(props)).unwrap();
414        writer.write(&batch).unwrap();
415        writer.close().unwrap();
416        let store = ObjectStore::new(Memory::default()).unwrap();
417        let file_size = bytes.len() as u64;
418        store.write("memory.parquet", bytes).await.unwrap();
419        let metadata = MetadataLoader::new(store.clone(), "memory.parquet", file_size)
420            .load(&mut MetadataCacheMetrics::default())
421            .await
422            .unwrap();
423        let sizes: Vec<_> = metadata
424            .row_groups()
425            .iter()
426            .map(|rg| rg.column(0).uncompressed_size() as usize)
427            .collect();
428        let aligned = |bytes: usize| bytes.div_ceil(1024) * 1024;
429        let limit = aligned(*sizes.iter().max().unwrap());
430        assert!(metadata.row_group(0).total_byte_size() as usize > limit);
431        let projection = ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0]);
432        let metadata =
433            ArrowReaderMetadata::try_new(Arc::new(metadata), ArrowReaderOptions::new()).unwrap();
434        let tracker = QueryMemoryTracker::builder(limit).build();
435        let cache = CacheStrategy::EnableAll(Arc::new(
436            crate::cache::CacheManager::builder()
437                .page_cache_size(1024 * 1024)
438                .build(),
439        ));
440        let region_file_id = RegionFileId::new(
441            store_api::storage::RegionId::new(1, 1),
442            store_api::storage::FileId::random(),
443        );
444        let build = |row_group_idx, path: &str, selection| {
445            let fetcher = SstParquetRangeFetcher::new(
446                region_file_id,
447                path.into(),
448                store.clone(),
449                if path == "memory.parquet" {
450                    cache.clone()
451                } else {
452                    CacheStrategy::Disabled
453                },
454                None,
455            )
456            .with_memory_tracker(Some(tracker.clone()));
457            build_sst_parquet_record_batch_stream(
458                metadata.clone(),
459                row_group_idx,
460                selection,
461                projection.clone(),
462                fetcher,
463                path.into(),
464                1,
465            )
466            .unwrap()
467        };
468        // Construction (including an unpolled stream's drop) must not reserve memory.
469        drop(build(0, "memory.parquet", None));
470        assert_eq!(0, tracker.current());
471        for (row_group_idx, size) in sizes.iter().enumerate() {
472            let mut first = build(row_group_idx, "memory.parquet", None);
473            assert_eq!(1, first.next().await.unwrap().unwrap().num_rows());
474            assert_eq!(aligned(*size), tracker.current());
475            let mut competing = build(row_group_idx, "memory.parquet", None);
476            assert!(competing.next().await.unwrap().is_err());
477            assert_eq!(aligned(*size), tracker.current());
478            // Reserve before decoding, even when the selection would require no page reads.
479            let mut skipped = build(
480                row_group_idx,
481                "memory.parquet",
482                Some(RowSelection::from(vec![
483                    parquet::arrow::arrow_reader::RowSelector::skip(4),
484                ])),
485            );
486            assert!(skipped.next().await.unwrap().is_err());
487            assert_eq!(aligned(*size), tracker.current());
488            while first.next().await.transpose().unwrap().is_some() {}
489            // Completed streams may remain alive without holding their quota.
490            assert_eq!(0, tracker.current());
491        }
492        // The second read must acquire quota even when all pages are cached.
493        store.delete("memory.parquet").await.unwrap();
494        let mut cancelled = build(0, "memory.parquet", None);
495        cancelled.next().await.unwrap().unwrap();
496        drop(cancelled);
497        assert_eq!(0, tracker.current());
498        let mut failed = build(0, "missing.parquet", None);
499        assert!(failed.next().await.unwrap().is_err());
500        drop(failed);
501        assert_eq!(0, tracker.current());
502        let mut skipped = build(
503            0,
504            "memory.parquet",
505            Some(RowSelection::from(vec![
506                parquet::arrow::arrow_reader::RowSelector::skip(4),
507            ])),
508        );
509        assert!(skipped.next().await.is_none());
510        assert_eq!(0, tracker.current());
511    }
512    #[tokio::test]
513    async fn large_local_file_size_is_lossless_and_missing_file_falls_back() {
514        use crate::cache::{CacheManager, file_cache::IndexValue};
515        use crate::test_util::TestEnv;
516        use common_base::readable_size::ReadableSize;
517        use common_test_util::temp_dir::create_temp_dir;
518        use object_store::services::{Fs, Memory};
519        use std::io::{Seek, SeekFrom, Write};
520        use store_api::storage::{FileId, RegionId};
521
522        let env = TestEnv::new().await;
523        let dir = create_temp_dir("large-page-cache");
524        let local = ObjectStore::new(Fs::default().root(dir.path().to_str().unwrap())).unwrap();
525        let write_cache = env.create_write_cache(local, ReadableSize::gb(8)).await;
526        let files = write_cache.file_cache();
527        let id = RegionFileId::new(RegionId::new(1, 1), FileId::random());
528        let key = IndexKey::new(id.region_id(), id.file_id(), FileType::Parquet);
529        let path = dir.path().join(files.cache_file_path(key));
530        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
531        let mut file = std::fs::File::create(&path).unwrap();
532        let size = 5_u64 * 1024 * 1024 * 1024;
533        let offset = 4_u64 * 1024 * 1024 * 1024 + 17;
534        // A sparse file reproduces the u32 overflow without writing 5 GiB.
535        file.set_len(size).unwrap();
536        file.seek(SeekFrom::Start(offset)).unwrap();
537        file.write_all(b"page").unwrap();
538        drop(file);
539        files
540            .put(
541                key,
542                IndexValue {
543                    file_size: size as u32,
544                },
545            )
546            .await;
547        let remote = ObjectStore::new(Memory::default()).unwrap();
548        let strategy = CacheStrategy::EnableAll(Arc::new(
549            CacheManager::builder()
550                .page_cache_size(1024 * 1024)
551                .write_cache(Some(write_cache))
552                .build(),
553        ));
554        let fetcher = SstParquetRangeFetcher::new(
555            id,
556            "remote.parquet".into(),
557            remote.clone(),
558            strategy.clone(),
559            None,
560        );
561        let range = offset..offset + 4;
562        let ranges = std::slice::from_ref(&range);
563        assert_eq!(fetcher.resolve_file_size(ranges).await.unwrap(), size);
564        assert_eq!(
565            fetcher.fetch_bytes_with_cache(ranges).await.unwrap(),
566            vec![Bytes::from_static(b"page")]
567        );
568
569        // Keep the stale index entry but remove its file; a fresh fetcher must stat remote.
570        std::fs::remove_file(&path).unwrap();
571        remote
572            .write("remote.parquet", b"remote".as_slice())
573            .await
574            .unwrap();
575        let fetcher =
576            SstParquetRangeFetcher::new(id, "remote.parquet".into(), remote, strategy, None);
577        assert_eq!(fetcher.resolve_file_size(&[]).await.unwrap(), 6);
578        assert_eq!(
579            fetcher.fetch_bytes_with_cache(&[0..3, 3..6]).await.unwrap(),
580            vec![Bytes::from_static(b"rem"), Bytes::from_static(b"ote")]
581        );
582    }
583
584    #[tokio::test]
585    async fn test_page_fetch_partial_hits_preserve_request_metrics() {
586        use crate::cache::CacheManager;
587        use crate::cache::parquet_page::PAGE_SIZE;
588        use object_store::services::Memory;
589        use std::sync::Arc;
590        use store_api::storage::{FileId, RegionId};
591
592        let store = ObjectStore::new(Memory::default()).unwrap();
593        let size = PAGE_SIZE * 2 + 10;
594        store
595            .write(
596                "bytes.parquet",
597                (0..size).map(|n| n as u8).collect::<Vec<_>>(),
598            )
599            .await
600            .unwrap();
601        let manager = Arc::new(CacheManager::builder().page_cache_size(1024 * 1024).build());
602        let id = RegionFileId::new(RegionId::new(1, 1), FileId::random());
603        let metrics = ParquetFetchMetrics::default();
604        let fetcher = SstParquetRangeFetcher::new(
605            id,
606            "bytes.parquet".into(),
607            store.clone(),
608            CacheStrategy::EnableAll(manager),
609            Some(metrics.clone()),
610        )
611        .with_file_size(size);
612        assert_eq!(
613            fetcher.fetch_bytes_with_cache(&[1..5, 1..1]).await.unwrap()[0],
614            Bytes::from_static(&[1, 2, 3, 4])
615        );
616        assert_eq!(
617            fetcher.fetch_bytes_with_cache(&[3..7, 1..1]).await.unwrap()[0],
618            Bytes::from_static(&[3, 4, 5, 6])
619        );
620        assert_eq!(
621            fetcher
622                .fetch_bytes_with_cache(&[PAGE_SIZE - 2..PAGE_SIZE + 2, 1..1])
623                .await
624                .unwrap()[0],
625            Bytes::from_static(&[254, 255, 0, 1])
626        );
627        {
628            let data = metrics.data.lock().unwrap();
629            assert_eq!(data.page_cache_hit, 2);
630            assert_eq!(data.cache_miss, 2);
631            assert_eq!(data.write_cache_hit, 0);
632            assert_eq!(data.page_size_needed, 12);
633            assert_eq!(data.page_size_to_fetch_mem, 6);
634            assert_eq!(data.page_size_to_fetch_store, PAGE_SIZE * 2);
635        }
636        // Completed pages serve another fetch without contacting the deleted backing object.
637        store.delete("bytes.parquet").await.unwrap();
638        let result = fetcher
639            .fetch_bytes_with_cache(&[PAGE_SIZE + 1..PAGE_SIZE + 3, size..size])
640            .await
641            .unwrap();
642        assert_eq!(&result[0][..], &[1, 2]);
643        assert!(result[1].is_empty());
644        let disabled = SstParquetRangeFetcher::new(
645            id,
646            "bytes.parquet".into(),
647            store,
648            CacheStrategy::Disabled,
649            None,
650        );
651        assert!(
652            disabled
653                .fetch_bytes_with_cache(&[1..5, 1..1])
654                .await
655                .is_err()
656        );
657    }
658
659    #[tokio::test]
660    async fn test_page_cache_coalesces_storage_reads() {
661        use std::sync::atomic::{AtomicUsize, Ordering};
662
663        use object_store::layers::mock;
664        use object_store::services::Memory;
665        use store_api::storage::{FileId, RegionId};
666
667        use crate::cache::CacheManager;
668        use crate::cache::parquet_page::PAGE_SIZE;
669
670        struct CountingReader {
671            inner: mock::oio::Reader,
672            reads: Arc<AtomicUsize>,
673        }
674
675        impl mock::oio::Read for CountingReader {
676            async fn open(
677                &self,
678                range: mock::BytesRange,
679            ) -> object_store::Result<(mock::RpRead, Box<dyn mock::oio::ReadStreamDyn>)>
680            {
681                self.reads.fetch_add(1, Ordering::Relaxed);
682                self.inner.open(range).await
683            }
684
685            async fn read(
686                &self,
687                range: mock::BytesRange,
688            ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> {
689                self.reads.fetch_add(1, Ordering::Relaxed);
690                self.inner.read(range).await
691            }
692        }
693
694        let reads = Arc::new(AtomicUsize::new(0));
695        let layer = mock::MockLayerBuilder::default()
696            .reader_factory(Arc::new({
697                let reads = reads.clone();
698                move |_, _, inner| {
699                    Box::new(CountingReader {
700                        inner,
701                        reads: reads.clone(),
702                    })
703                }
704            }))
705            .build()
706            .unwrap();
707        let store = ObjectStore::new(Memory::default()).unwrap().layer(layer);
708        let bytes = Bytes::from(
709            (0..PAGE_SIZE * 128 + 13)
710                .map(|n| (n * 31 + n / PAGE_SIZE) as u8)
711                .collect::<Vec<_>>(),
712        );
713        store.write("bytes.parquet", bytes.clone()).await.unwrap();
714        let manager = Arc::new(
715            CacheManager::builder()
716                .page_cache_size(16 * 1024 * 1024)
717                .build(),
718        );
719        let fetcher = SstParquetRangeFetcher::new(
720            RegionFileId::new(RegionId::new(1, 1), FileId::random()),
721            "bytes.parquet".into(),
722            store,
723            CacheStrategy::EnableAll(manager),
724            None,
725        )
726        .with_file_size(bytes.len() as u64);
727
728        // A cold request spanning 65 adjacent cache pages opens one storage read.
729        let contiguous = 3..PAGE_SIZE * 65 - 3;
730        for _ in 0..2 {
731            assert_eq!(
732                fetcher
733                    .fetch_bytes_with_cache(std::slice::from_ref(&contiguous))
734                    .await
735                    .unwrap(),
736                vec![bytes.slice(contiguous.start as usize..contiguous.end as usize)]
737            );
738            assert_eq!(reads.load(Ordering::Relaxed), 1);
739        }
740
741        // Two distant cold pages stay separate; cached and duplicate ranges keep
742        // their requested order without adding storage reads.
743        let scattered = [
744            PAGE_SIZE * 120 + 3..PAGE_SIZE * 120 + 17,
745            PAGE_SIZE * 80 + 7..PAGE_SIZE * 80 + 11,
746            1..5,
747            PAGE_SIZE * 120 + 4..PAGE_SIZE * 120 + 8,
748        ];
749        let expected: Vec<_> = scattered
750            .iter()
751            .map(|range| bytes.slice(range.start as usize..range.end as usize))
752            .collect();
753        for _ in 0..2 {
754            assert_eq!(
755                fetcher.fetch_bytes_with_cache(&scattered).await.unwrap(),
756                expected
757            );
758            assert_eq!(reads.load(Ordering::Relaxed), 3);
759        }
760    }
761
762    #[tokio::test]
763    async fn test_page_cache_preserves_storage_error_semantics() {
764        use std::sync::atomic::{AtomicBool, Ordering};
765
766        use common_error::ext::ErrorExt;
767        use object_store::layers::mock;
768        use object_store::services::Memory;
769        use store_api::storage::{FileId, RegionId};
770
771        use crate::cache::CacheManager;
772
773        struct FailingReader {
774            inner: mock::oio::Reader,
775            fail: Arc<AtomicBool>,
776            kind: object_store::ErrorKind,
777            temporary: bool,
778        }
779
780        impl FailingReader {
781            fn error(&self) -> object_store::Error {
782                let error = object_store::Error::new(self.kind, "page read failed");
783                if self.temporary {
784                    error.set_temporary()
785                } else {
786                    error
787                }
788            }
789        }
790
791        impl mock::oio::Read for FailingReader {
792            async fn open(
793                &self,
794                range: mock::BytesRange,
795            ) -> object_store::Result<(mock::RpRead, Box<dyn mock::oio::ReadStreamDyn>)>
796            {
797                if self.fail.load(Ordering::Relaxed) {
798                    return Err(self.error());
799                }
800                self.inner.open(range).await
801            }
802
803            async fn read(
804                &self,
805                range: mock::BytesRange,
806            ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> {
807                if self.fail.load(Ordering::Relaxed) {
808                    return Err(self.error());
809                }
810                self.inner.read(range).await
811            }
812        }
813
814        for (kind, temporary) in [
815            (object_store::ErrorKind::PermissionDenied, false),
816            (object_store::ErrorKind::NotFound, false),
817            (object_store::ErrorKind::Unexpected, true),
818        ] {
819            let mut baseline = None;
820            for cache_enabled in [false, true] {
821                let fail = Arc::new(AtomicBool::new(true));
822                let layer = mock::MockLayerBuilder::default()
823                    .reader_factory(Arc::new({
824                        let fail = fail.clone();
825                        move |_, _, inner| {
826                            Box::new(FailingReader {
827                                inner,
828                                fail: fail.clone(),
829                                kind,
830                                temporary,
831                            })
832                        }
833                    }))
834                    .build()
835                    .unwrap();
836                let store = ObjectStore::new(Memory::default()).unwrap().layer(layer);
837                store.write("bytes.parquet", vec![7; 100]).await.unwrap();
838                let strategy = if cache_enabled {
839                    CacheStrategy::EnableAll(Arc::new(
840                        CacheManager::builder().page_cache_size(1024 * 1024).build(),
841                    ))
842                } else {
843                    CacheStrategy::Disabled
844                };
845                let fetcher = SstParquetRangeFetcher::new(
846                    RegionFileId::new(RegionId::new(1, 1), FileId::random()),
847                    "bytes.parquet".into(),
848                    store,
849                    strategy,
850                    None,
851                )
852                .with_file_size(100);
853                let range = 0..4;
854                let ranges = std::slice::from_ref(&range);
855                let error = fetcher.fetch_bytes_with_cache(ranges).await.unwrap_err();
856                let semantics = (
857                    error.status_code(),
858                    error.retry_hint(),
859                    error.is_object_not_found(),
860                );
861                if cache_enabled {
862                    assert!(matches!(error, Error::LoadParquetPage { .. }));
863                    assert_eq!(Some(semantics), baseline);
864                } else {
865                    assert!(matches!(error, Error::OpenDal { .. }));
866                    baseline = Some(semantics);
867                }
868                // Failed shared loads must not poison the cache.
869                fail.store(false, Ordering::Relaxed);
870                assert_eq!(
871                    fetcher.fetch_bytes_with_cache(ranges).await.unwrap(),
872                    vec![Bytes::from_static(&[7; 4])]
873                );
874            }
875        }
876    }
877
878    #[tokio::test]
879    async fn test_cached_and_disabled_fetch_return_identical_ranges() {
880        use crate::cache::CacheManager;
881        use crate::cache::parquet_page::PAGE_SIZE;
882        use object_store::services::Memory;
883        use std::sync::Arc;
884        use store_api::storage::{FileId, RegionId};
885
886        let store = ObjectStore::new(Memory::default()).unwrap();
887        let size = PAGE_SIZE + 10;
888        store
889            .write(
890                "bytes.parquet",
891                (0..size).map(|n| n as u8).collect::<Vec<_>>(),
892            )
893            .await
894            .unwrap();
895        let id = RegionFileId::new(RegionId::new(1, 1), FileId::random());
896        let manager = Arc::new(CacheManager::builder().page_cache_size(1024 * 1024).build());
897        let cached = SstParquetRangeFetcher::new(
898            id,
899            "bytes.parquet".into(),
900            store.clone(),
901            CacheStrategy::EnableAll(manager.clone()),
902            None,
903        );
904        let disabled = SstParquetRangeFetcher::new(
905            id,
906            "bytes.parquet".into(),
907            store.clone(),
908            CacheStrategy::Disabled,
909            None,
910        );
911        let compact = SstParquetRangeFetcher::new(
912            id,
913            "bytes.parquet".into(),
914            store,
915            CacheStrategy::Compaction(manager),
916            None,
917        );
918        let ranges = vec![size - 3..size, 1..6, PAGE_SIZE - 2..PAGE_SIZE + 3, 4..8];
919        let expected = disabled.fetch_bytes_with_cache(&ranges).await.unwrap();
920        assert_eq!(
921            cached.fetch_bytes_with_cache(&ranges).await.unwrap(),
922            expected
923        );
924        assert_eq!(
925            cached.fetch_bytes_with_cache(&ranges).await.unwrap(),
926            expected
927        );
928        assert_eq!(
929            compact.fetch_bytes_with_cache(&ranges).await.unwrap(),
930            expected
931        );
932    }
933}
mito2/sst/parquet/push_decoder.rs

mito2/sst/parquet/
push_decoder.rs
