mito2/cache/
parquet_page.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//! Fixed byte blocks shared across row groups of an immutable Parquet file.
16
17use std::io;
18use std::mem;
19use std::ops::Range;
20use std::sync::Arc;
21use std::time::Duration;
22
23use bytes::Bytes;
24use common_base::range_read::RangeReader;
25use common_cache::paged::{PageFetchStats, PagedCache};
26
27use crate::metrics;
28use crate::sst::file::RegionFileId;
29
30/// Default byte blocks, independent of the Parquet format's encoded data-page boundaries.
31pub(crate) const PAGE_SIZE: u64 = 64 * 1024;
32
33/// Internal tuning knob, in bytes; read only when constructing a cache.
34const PAGE_SIZE_ENV: &str = "GREPTIME_PARQUET_PAGE_SIZE";
35/// Expire pages one hour after insertion; physical reclamation uses cache maintenance.
36const PAGE_TTL: Duration = Duration::from_secs(3600);
37
38fn parse_page_size(value: Option<&str>) -> u64 {
39    match value {
40        None => PAGE_SIZE,
41        Some(value) => match value.parse::<u64>() {
42            Ok(size) if (1..=MAX_PAGE_SIZE).contains(&size) => size,
43            _ => {
44                common_telemetry::warn!(
45                    "Invalid {PAGE_SIZE_ENV}: expected bytes in 1..={MAX_PAGE_SIZE}; using {PAGE_SIZE}"
46                );
47                PAGE_SIZE
48            }
49        },
50    }
51}
52
53/// Largest block whose bytes and entry overhead fit the cache's u32 weight.
54pub(crate) const MAX_PAGE_SIZE: u64 = u32::MAX as u64
55    - (mem::size_of::<(RegionFileId, u64)>() + mem::size_of::<Bytes>() + mem::size_of::<u64>())
56        as u64;
57
58/// Caches immutable Parquet file bytes in fixed-size blocks shared across row groups.
59pub(crate) struct ParquetPageCache {
60    page_size: u64,
61    cache: PagedCache<RegionFileId>,
62}
63
64impl ParquetPageCache {
65    /// Constructs the production cache, resolving the internal block-size override once.
66    pub(crate) fn from_env(capacity: u64) -> Arc<Self> {
67        let value = std::env::var(PAGE_SIZE_ENV).ok();
68        Self::new(capacity, parse_page_size(value.as_deref()))
69    }
70
71    /// Creates a cache with a total byte capacity and a byte-block size.
72    pub(crate) fn new(capacity: u64, page_size: u64) -> Arc<Self> {
73        Self::with_ttl(capacity, page_size, PAGE_TTL)
74    }
75
76    fn with_ttl(capacity: u64, page_size: u64, ttl: Duration) -> Arc<Self> {
77        Arc::new(Self {
78            page_size,
79            cache: PagedCache::new_with_ttl(
80                "parquet_pages",
81                capacity,
82                page_size,
83                |_: &(RegionFileId, u64), bytes: &Bytes| {
84                    (mem::size_of::<(RegionFileId, u64)>()
85                        + mem::size_of::<Bytes>()
86                        + mem::size_of::<u64>()
87                        + bytes.len()) as u32
88                },
89                metrics::cache_page_observer("page", "page", |hits, misses| {
90                    (usize::from(hits > 0), usize::from(misses > 0))
91                }),
92                Some(ttl),
93            ),
94        })
95    }
96
97    /// Recovers the immutable file size from a requested block already in memory.
98    pub(crate) async fn known_file_size(
99        &self,
100        file: RegionFileId,
101        ranges: &[Range<u64>],
102    ) -> Option<u64> {
103        self.cache.known_file_size(&file, ranges).await
104    }
105
106    /// Fetches ranges in input order, sharing missing blocks and returning [`PageFetchStats`].
107    /// The reader must address the immutable file identified by `file`.
108    pub(crate) async fn fetch<R: RangeReader + ?Sized>(
109        &self,
110        file: RegionFileId,
111        file_size: u64,
112        ranges: &[Range<u64>],
113        reader: &R,
114    ) -> Result<(Vec<Bytes>, PageFetchStats), Arc<io::Error>> {
115        // Also protect internal callers that supply an explicit block size.
116        if self.page_size == 0 || self.page_size > MAX_PAGE_SIZE {
117            return Err(Arc::new(io::Error::new(
118                io::ErrorKind::InvalidInput,
119                format!(
120                    "Parquet cache page size {} is outside 1..={MAX_PAGE_SIZE} bytes",
121                    self.page_size
122                ),
123            )));
124        }
125        let result = self.cache.fetch(file, file_size, ranges, reader).await?;
126        metrics::cache_page_fetch_metrics("page", &result.1);
127        Ok(result)
128    }
129
130    /// Checks whether a completed block is available without recording a lookup event.
131    #[cfg(test)]
132    pub(crate) async fn contains(&self, file: RegionFileId, page_id: u64) -> bool {
133        self.cache.contains(&file, page_id).await
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use async_trait::async_trait;
141    use common_base::range_read::Metadata;
142    use std::sync::Mutex;
143    use store_api::storage::{FileId, RegionId};
144
145    #[derive(Default)]
146    struct Reader(Mutex<Vec<Range<u64>>>);
147
148    #[async_trait]
149    impl RangeReader for Reader {
150        async fn metadata(&self) -> io::Result<Metadata> {
151            Ok(Metadata {
152                content_length: PAGE_SIZE * 2 + 10,
153            })
154        }
155        async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
156            self.0.lock().unwrap().push(range.clone());
157            tokio::task::yield_now().await;
158            Ok(Bytes::from(range.map(|n| n as u8).collect::<Vec<_>>()))
159        }
160    }
161
162    #[tokio::test]
163    async fn test_aligns_and_reuses_partial_overlaps_without_merging_reads() {
164        let cache = ParquetPageCache::new(1024 * 1024, PAGE_SIZE);
165        let file = RegionFileId::new(RegionId::new(1, 1), FileId::random());
166        let size = PAGE_SIZE * 2 + 10;
167        let reader = Reader::default();
168        let (result, stats) = cache
169            .fetch(
170                file,
171                size,
172                &[1..4, PAGE_SIZE - 2..PAGE_SIZE + 2, size - 3..size],
173                &reader,
174            )
175            .await
176            .unwrap();
177        assert_eq!(&result[0][..], &[1, 2, 3]);
178        assert_eq!(&result[1][..], &[254, 255, 0, 1]);
179        assert_eq!(&result[2][..], &[7, 8, 9]);
180        assert_eq!(
181            (
182                stats.requested_bytes,
183                stats.unique_requested_bytes,
184                stats.aligned_bytes
185            ),
186            (10, 10, size)
187        );
188        assert_eq!(
189            reader.0.lock().unwrap().as_slice(),
190            &[0..PAGE_SIZE, PAGE_SIZE..PAGE_SIZE * 2, PAGE_SIZE * 2..size]
191        );
192        let (_, stats) = cache
193            .fetch(file, size, &[2..6, 4..8], &reader)
194            .await
195            .unwrap();
196        assert_eq!(
197            (stats.cache_hit, stats.cache_miss, stats.cached_bytes),
198            (1, 0, 8)
199        );
200        assert_eq!(reader.0.lock().unwrap().len(), 3);
201    }
202
203    #[tokio::test]
204    async fn file_identity_includes_region_and_concurrent_reads_share() {
205        let cache = ParquetPageCache::new(1024 * 1024, PAGE_SIZE);
206        let file_id = FileId::random();
207        let first = RegionFileId::new(RegionId::new(1, 1), file_id);
208        let second = RegionFileId::new(RegionId::new(2, 1), file_id);
209        let reader = Reader::default();
210        let ranges = [0..4, 2..6];
211        let (a, b) = tokio::join!(
212            cache.fetch(first, 100, &ranges, &reader),
213            cache.fetch(first, 100, &ranges, &reader)
214        );
215        assert_eq!(a.unwrap().0, b.unwrap().0);
216        assert_eq!(reader.0.lock().unwrap().len(), 1);
217        cache.fetch(second, 100, &ranges, &reader).await.unwrap();
218        assert_eq!(reader.0.lock().unwrap().len(), 2);
219        assert!(cache.contains(first, 0).await);
220        assert!(cache.contains(second, 0).await);
221    }
222
223    #[tokio::test]
224    async fn configurable_blocks_preserve_cross_block_and_tail_reads() {
225        let file = RegionFileId::new(RegionId::new(1, 1), FileId::random());
226        let size = PAGE_SIZE * 2 + 10;
227        // Include a non-power-of-two size: the layout only requires a nonzero divisor.
228        for page_size in [1, 37, 4096, PAGE_SIZE, PAGE_SIZE * 4] {
229            let cache = ParquetPageCache::new(16 * 1024 * 1024, page_size);
230            let reader = Reader::default();
231            let boundary = page_size.min(size - 1);
232            let ranges = [
233                0..3,
234                boundary - 1..boundary + 1,
235                PAGE_SIZE - 2..PAGE_SIZE + 2,
236                size - 3..size,
237                size..size,
238            ];
239            let expected: Vec<_> = ranges
240                .iter()
241                .map(|r| Bytes::from(r.clone().map(|n| n as u8).collect::<Vec<_>>()))
242                .collect();
243            let (actual, _) = cache.fetch(file, size, &ranges, &reader).await.unwrap();
244            assert_eq!(actual, expected);
245            let reads = reader.0.lock().unwrap().len();
246            assert!(reader.0.lock().unwrap().iter().all(|r| r.end <= size));
247            let (cached, stats) = cache.fetch(file, size, &ranges, &reader).await.unwrap();
248            assert_eq!(cached, expected);
249            assert_eq!(stats.cache_miss, 0);
250            assert_eq!(reader.0.lock().unwrap().len(), reads);
251        }
252    }
253
254    #[tokio::test]
255    async fn builder_uses_fixed_block_size() {
256        use crate::cache::{CacheManager, CacheStrategy};
257
258        let strategy = CacheStrategy::EnableAll(Arc::new(
259            CacheManager::builder().page_cache_size(1024 * 1024).build(),
260        ));
261        let cache = strategy.parquet_page_cache().unwrap();
262        let reader = Reader::default();
263        let file = RegionFileId::new(RegionId::new(1, 1), FileId::random());
264        cache
265            .fetch(file, PAGE_SIZE * 2, std::slice::from_ref(&(0..1)), &reader)
266            .await
267            .unwrap();
268        assert_eq!(
269            reader.0.lock().unwrap().as_slice(),
270            std::slice::from_ref(&(0..PAGE_SIZE))
271        );
272        let disabled = CacheStrategy::EnableAll(Arc::new(CacheManager::builder().build()));
273        assert!(disabled.parquet_page_cache().is_none());
274    }
275
276    #[test]
277    fn block_size_override_uses_valid_bytes_or_default() {
278        for (value, expected) in [
279            (None, PAGE_SIZE),
280            (Some(""), PAGE_SIZE),
281            (Some("64KiB"), PAGE_SIZE),
282            (Some("0"), PAGE_SIZE),
283            (Some("-1"), PAGE_SIZE),
284            (Some("37"), 37),
285            (Some("131072"), 131072),
286        ] {
287            assert_eq!(parse_page_size(value), expected);
288        }
289        assert_eq!(
290            parse_page_size(Some(&MAX_PAGE_SIZE.to_string())),
291            MAX_PAGE_SIZE
292        );
293        assert_eq!(
294            parse_page_size(Some(&(MAX_PAGE_SIZE + 1).to_string())),
295            PAGE_SIZE
296        );
297    }
298
299    #[tokio::test]
300    async fn expired_pages_reload_without_file_invalidation() {
301        let cache = ParquetPageCache::with_ttl(1024 * 1024, PAGE_SIZE, Duration::from_millis(50));
302        let retained = ParquetPageCache::new(1024 * 1024, PAGE_SIZE);
303        let file = RegionFileId::new(RegionId::new(1, 1), FileId::random());
304        let reader = Reader::default();
305        let ranges = [0..4, 4..8];
306        let expected = cache.fetch(file, 100, &ranges, &reader).await.unwrap().0;
307        retained.fetch(file, 100, &ranges, &reader).await.unwrap();
308        cache.cache.run_pending_tasks().await;
309        retained.cache.run_pending_tasks().await;
310        assert!(retained.contains(file, 0).await);
311        tokio::time::timeout(Duration::from_secs(5), async {
312            while cache.contains(file, 0).await {
313                tokio::time::sleep(Duration::from_millis(5)).await;
314            }
315        })
316        .await
317        .unwrap();
318        assert!(retained.contains(file, 0).await);
319        let (actual, stats) = cache.fetch(file, 100, &ranges, &reader).await.unwrap();
320        assert_eq!(actual, expected);
321        assert_eq!(stats.cache_miss, 1);
322        assert_eq!(reader.0.lock().unwrap().len(), 3);
323    }
324
325    #[tokio::test]
326    async fn invalid_block_size_is_rejected_before_io() {
327        let file = RegionFileId::new(RegionId::new(1, 1), FileId::random());
328        for page_size in [0, MAX_PAGE_SIZE + 1, u64::MAX] {
329            let cache = ParquetPageCache::new(1024, page_size);
330            let reader = Reader::default();
331            let error = cache
332                .fetch(file, 100, std::slice::from_ref(&(0..4)), &reader)
333                .await
334                .unwrap_err();
335            assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
336            assert!(reader.0.lock().unwrap().is_empty());
337        }
338    }
339
340    #[tokio::test]
341    async fn oversized_page_is_returned_without_retention() {
342        let cache = ParquetPageCache::new(1, PAGE_SIZE);
343        let file = RegionFileId::new(RegionId::new(1, 1), FileId::random());
344        let reader = Reader::default();
345        let result = cache
346            .fetch(file, 100, &[0..4, 4..8], &reader)
347            .await
348            .unwrap();
349        assert_eq!(&result.0[0][..], &[0, 1, 2, 3]);
350        cache.cache.run_pending_tasks().await;
351        assert!(!cache.contains(file, 0).await);
352    }
353}
mito2/cache/parquet_page.rs

mito2/cache/
parquet_page.rs
