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::ops::Range;
18
19use bytes::{Bytes, BytesMut};
20use common_recordbatch::QueryMemoryTracker;
21use datatypes::arrow::record_batch::RecordBatch;
22use futures::StreamExt;
23use futures::stream::BoxStream;
24use object_store::ObjectStore;
25use parquet::DecodeResult;
26use parquet::arrow::ProjectionMask;
27use parquet::arrow::arrow_reader::{ArrowReaderMetadata, RowSelection};
28use parquet::arrow::push_decoder::ParquetPushDecoderBuilder;
29use snafu::{OptionExt, ResultExt, ensure};
30
31use crate::cache::file_cache::{FileType, IndexKey};
32use crate::cache::{CacheStrategy, PageRangePart};
33use crate::error::{OpenDalSnafu, ReadParquetSnafu, Result, UnexpectedSnafu};
34use crate::metrics::{READ_STAGE_ELAPSED, READ_STAGE_FETCH_PAGES};
35use crate::sst::file::RegionFileId;
36use crate::sst::parquet::helper::fetch_byte_ranges;
37use crate::sst::parquet::row_group::{ParquetFetchMetrics, compute_total_range_size};
38
39/// Fetches parquet byte ranges through Greptime's cache hierarchy.
40///
41/// The push decoder decides which ranges are required for decoding, while this
42/// fetcher keeps cache lookup, local write-cache reads, and remote I/O explicit
43/// in Greptime code.
44pub struct SstParquetRangeFetcher {
45    /// Region file ID for cache key.
46    region_file_id: RegionFileId,
47    /// Path to the parquet file in object storage.
48    file_path: String,
49    /// Object store for reading data.
50    object_store: ObjectStore,
51    /// Cache strategy for reading pages.
52    cache_strategy: CacheStrategy,
53    /// Row group index for cache key.
54    row_group_idx: usize,
55    /// Optional metrics for tracking fetch operations.
56    fetch_metrics: Option<ParquetFetchMetrics>,
57    /// Shared quota for decoding this fetcher's row group.
58    memory_tracker: Option<QueryMemoryTracker>,
59}
60
61impl SstParquetRangeFetcher {
62    /// Creates a new [SstParquetRangeFetcher].
63    pub fn new(
64        region_file_id: RegionFileId,
65        file_path: String,
66        object_store: ObjectStore,
67        cache_strategy: CacheStrategy,
68        row_group_idx: usize,
69        fetch_metrics: Option<ParquetFetchMetrics>,
70    ) -> Self {
71        Self {
72            region_file_id,
73            file_path,
74            object_store,
75            cache_strategy,
76            row_group_idx,
77            fetch_metrics,
78            memory_tracker: None,
79        }
80    }
81
82    /// Sets the shared quota for this decoder's projected row group.
83    pub(crate) fn with_memory_tracker(mut self, tracker: Option<QueryMemoryTracker>) -> Self {
84        self.memory_tracker = tracker;
85        self
86    }
87
88    /// Fetches byte ranges from page cache, write cache, or object store.
89    async fn fetch_bytes_with_cache(&self, ranges: Vec<Range<u64>>) -> Result<Vec<Bytes>> {
90        let fetch_start = self
91            .fetch_metrics
92            .as_ref()
93            .map(|_| std::time::Instant::now());
94        let _timer = READ_STAGE_FETCH_PAGES.start_timer();
95
96        let mut page_lookup = self.cache_strategy.get_page_ranges(
97            self.region_file_id.file_id(),
98            self.row_group_idx,
99            &ranges,
100        );
101        if let Some(lookup) = &page_lookup
102            && lookup.cached_bytes > 0
103            && let Some(metrics) = &self.fetch_metrics
104        {
105            let mut metrics_data = metrics.data.lock().unwrap();
106            metrics_data.page_cache_hit += 1;
107            metrics_data.pages_to_fetch_mem += lookup.cached_range_count;
108            metrics_data.page_size_to_fetch_mem += lookup.cached_bytes;
109            metrics_data.page_size_needed += lookup.cached_bytes;
110        }
111
112        // Fast path: all requested ranges can be assembled from cached fragments.
113        if page_lookup
114            .as_ref()
115            .map(|lookup| lookup.is_fully_cached())
116            .unwrap_or(false)
117        {
118            let lookup = page_lookup.take().unwrap();
119            if let Some(metrics) = &self.fetch_metrics
120                && let Some(start) = fetch_start
121            {
122                metrics.data.lock().unwrap().total_fetch_elapsed += start.elapsed();
123            }
124            return assemble_ranges(&ranges, lookup.cached_parts, &[]);
125        }
126
127        let missing_ranges = page_lookup
128            .as_ref()
129            .map(|lookup| lookup.missing_ranges.clone())
130            .unwrap_or_else(|| ranges.clone());
131
132        // Calculate total range size for metrics.
133        let (_, unaligned_size) = compute_total_range_size(&missing_ranges);
134
135        // Check write cache.
136        let key = IndexKey::new(
137            self.region_file_id.region_id(),
138            self.region_file_id.file_id(),
139            FileType::Parquet,
140        );
141        let fetch_write_cache_start = self
142            .fetch_metrics
143            .as_ref()
144            .map(|_| std::time::Instant::now());
145        let write_cache_result = match self.cache_strategy.write_cache() {
146            Some(cache) => cache.file_cache().read_ranges(key, &missing_ranges).await,
147            None => None,
148        };
149
150        let fetched_pages = match write_cache_result {
151            Some(data) => {
152                if let Some(metrics) = &self.fetch_metrics {
153                    let elapsed = fetch_write_cache_start
154                        .map(|start| start.elapsed())
155                        .unwrap_or_default();
156                    let range_size_needed: u64 =
157                        missing_ranges.iter().map(|r| r.end - r.start).sum();
158                    let mut metrics_data = metrics.data.lock().unwrap();
159                    metrics_data.write_cache_fetch_elapsed += elapsed;
160                    metrics_data.write_cache_hit += 1;
161                    metrics_data.pages_to_fetch_write_cache += missing_ranges.len();
162                    metrics_data.page_size_to_fetch_write_cache += unaligned_size;
163                    metrics_data.page_size_needed += range_size_needed;
164                }
165                data
166            }
167            None => {
168                // Fetch data from object store.
169                let _timer = READ_STAGE_ELAPSED
170                    .with_label_values(&["cache_miss_read"])
171                    .start_timer();
172
173                let start = self
174                    .fetch_metrics
175                    .as_ref()
176                    .map(|_| std::time::Instant::now());
177                let data =
178                    fetch_byte_ranges(&self.file_path, self.object_store.clone(), &missing_ranges)
179                        .await
180                        .context(OpenDalSnafu)?;
181
182                if let Some(metrics) = &self.fetch_metrics {
183                    let elapsed = start.map(|start| start.elapsed()).unwrap_or_default();
184                    let range_size_needed: u64 =
185                        missing_ranges.iter().map(|r| r.end - r.start).sum();
186                    let mut metrics_data = metrics.data.lock().unwrap();
187                    metrics_data.store_fetch_elapsed += elapsed;
188                    metrics_data.cache_miss += 1;
189                    metrics_data.pages_to_fetch_store += missing_ranges.len();
190                    metrics_data.page_size_to_fetch_store += unaligned_size;
191                    metrics_data.page_size_needed += range_size_needed;
192                }
193                data
194            }
195        };
196        ensure!(
197            fetched_pages.len() == missing_ranges.len(),
198            UnexpectedSnafu {
199                reason: format!(
200                    "Invalid parquet range fetch: {} missing ranges but {} fetched byte ranges",
201                    missing_ranges.len(),
202                    fetched_pages.len()
203                ),
204            }
205        );
206
207        self.cache_strategy.put_page_ranges(
208            self.region_file_id.file_id(),
209            self.row_group_idx,
210            &missing_ranges,
211            &fetched_pages,
212        );
213
214        if let (Some(metrics), Some(start)) = (&self.fetch_metrics, fetch_start) {
215            metrics.data.lock().unwrap().total_fetch_elapsed += start.elapsed();
216        }
217
218        if let Some(lookup) = page_lookup {
219            let fetched_parts = missing_ranges
220                .into_iter()
221                .zip(fetched_pages)
222                .map(|(range, bytes)| PageRangePart { range, bytes })
223                .collect::<Vec<_>>();
224            return assemble_ranges(&ranges, lookup.cached_parts, &fetched_parts);
225        }
226
227        Ok(fetched_pages)
228    }
229}
230
231fn assemble_ranges(
232    ranges: &[Range<u64>],
233    cached_parts: Vec<Vec<PageRangePart>>,
234    fetched_parts: &[PageRangePart],
235) -> Result<Vec<Bytes>> {
236    ensure!(
237        ranges.len() == cached_parts.len(),
238        UnexpectedSnafu {
239            reason: format!(
240                "Invalid parquet range assembly: {} requested ranges but {} cached part groups",
241                ranges.len(),
242                cached_parts.len()
243            ),
244        }
245    );
246
247    ranges
248        .iter()
249        .zip(cached_parts)
250        .map(|(range, mut parts)| {
251            parts.extend(
252                fetched_parts
253                    .iter()
254                    .filter_map(|part| overlapping_part(range, part)),
255            );
256            assemble_range(range, parts)
257        })
258        .collect()
259}
260
261fn overlapping_part(range: &Range<u64>, part: &PageRangePart) -> Option<PageRangePart> {
262    let start = range.start.max(part.range.start);
263    let end = range.end.min(part.range.end);
264    if start >= end {
265        return None;
266    }
267
268    let slice_start = (start - part.range.start) as usize;
269    let slice_end = (end - part.range.start) as usize;
270    Some(PageRangePart {
271        range: start..end,
272        bytes: part.bytes.slice(slice_start..slice_end),
273    })
274}
275
276fn assemble_range(range: &Range<u64>, mut parts: Vec<PageRangePart>) -> Result<Bytes> {
277    if range.start >= range.end {
278        return Ok(Bytes::new());
279    }
280
281    parts.sort_unstable_by_key(|part| part.range.start);
282    if parts.len() == 1 && parts[0].range == *range {
283        return Ok(parts.pop().unwrap().bytes);
284    }
285
286    let mut cursor = range.start;
287    let mut output = BytesMut::with_capacity((range.end - range.start) as usize);
288    for part in parts {
289        ensure!(
290            part.range.start <= cursor,
291            UnexpectedSnafu {
292                reason: format!(
293                    "Missing cached parquet bytes for range {}..{}, next part starts at {}",
294                    range.start, range.end, part.range.start
295                ),
296            }
297        );
298        if part.range.end <= cursor {
299            continue;
300        }
301
302        let slice_start = (cursor - part.range.start) as usize;
303        let slice_end = (part.range.end.min(range.end) - part.range.start) as usize;
304        output.extend_from_slice(&part.bytes.slice(slice_start..slice_end));
305        cursor = part.range.end.min(range.end);
306        if cursor >= range.end {
307            break;
308        }
309    }
310
311    ensure!(
312        cursor == range.end,
313        UnexpectedSnafu {
314            reason: format!(
315                "Missing cached parquet bytes for range {}..{}, assembled through {}",
316                range.start, range.end, cursor
317            ),
318        }
319    );
320
321    Ok(output.freeze())
322}
323
324/// Builds a parquet record batch stream driven directly by [ParquetPushDecoderBuilder].
325pub fn build_sst_parquet_record_batch_stream(
326    arrow_metadata: ArrowReaderMetadata,
327    row_group_idx: usize,
328    row_selection: Option<RowSelection>,
329    projection: ProjectionMask,
330    fetcher: SstParquetRangeFetcher,
331    file_path: String,
332    batch_size: usize,
333) -> Result<BoxStream<'static, Result<RecordBatch>>> {
334    let uncompressed_bytes = if fetcher.memory_tracker.is_some() {
335        let row_group = arrow_metadata.metadata().row_group(row_group_idx);
336        row_group
337            .columns()
338            .iter()
339            .enumerate()
340            .filter(|(index, _)| projection.leaf_included(*index))
341            .try_fold(0usize, |total, (_, column)| {
342                usize::try_from(column.uncompressed_size())
343                    .ok()
344                    .and_then(|size| total.checked_add(size))
345            })
346            .context(UnexpectedSnafu {
347                reason: "Invalid projected row-group uncompressed size",
348            })?
349    } else {
350        0
351    };
352    let mut builder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata)
353        .with_row_groups(vec![row_group_idx])
354        .with_projection(projection)
355        .with_batch_size(batch_size);
356
357    if let Some(selection) = row_selection {
358        builder = builder.with_row_selection(selection);
359    }
360
361    let mut decoder = builder
362        .build()
363        .context(ReadParquetSnafu { path: &file_path })?;
364
365    Ok(async_stream::try_stream! {
366        let _reservation = if let Some(tracker) = &fetcher.memory_tracker {
367            Some(tracker.reserve(uncompressed_bytes)
368                .context(crate::error::RecordBatchSnafu)?)
369        } else {
370            None
371        };
372        loop {
373            match decoder.try_decode().context(ReadParquetSnafu { path: &file_path })? {
374                DecodeResult::NeedsData(ranges) => {
375                    let data = fetcher.fetch_bytes_with_cache(ranges.clone()).await?;
376                    decoder.push_ranges(ranges, data)
377                        .context(ReadParquetSnafu { path: &file_path })?;
378                }
379                DecodeResult::Data(batch) => yield batch,
380                DecodeResult::Finished => break,
381            }
382        }
383    }
384    .boxed())
385}
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390
391    #[tokio::test]
392    async fn test_decoder_memory_projection_lifetime_and_error() {
393        use std::sync::Arc;
394
395        use datatypes::arrow::array::{Int64Array, StringArray};
396        use datatypes::arrow::datatypes::{DataType, Field, Schema};
397        use object_store::services::Memory;
398        use parquet::arrow::ArrowWriter;
399        use parquet::arrow::arrow_reader::ArrowReaderOptions;
400        use parquet::file::properties::WriterProperties;
401
402        use crate::sst::parquet::metadata::MetadataLoader;
403        use crate::sst::parquet::reader::MetadataCacheMetrics;
404
405        let schema = Arc::new(Schema::new(vec![
406            Field::new("value", DataType::Int64, false),
407            Field::new("payload", DataType::Utf8, false),
408        ]));
409        let batch = RecordBatch::try_new(
410            schema.clone(),
411            vec![
412                Arc::new(Int64Array::from_iter_values(0..8)),
413                Arc::new(StringArray::from_iter_values(
414                    (0..8).map(|i| format!("{i}:{}", "x".repeat(512))),
415                )),
416            ],
417        )
418        .unwrap();
419        let mut bytes = Vec::new();
420        let props = WriterProperties::builder()
421            .set_max_row_group_row_count(Some(4))
422            .build();
423        let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(props)).unwrap();
424        writer.write(&batch).unwrap();
425        writer.close().unwrap();
426        let store = ObjectStore::new(Memory::default()).unwrap();
427        let file_size = bytes.len() as u64;
428        store.write("memory.parquet", bytes).await.unwrap();
429        let metadata = MetadataLoader::new(store.clone(), "memory.parquet", file_size)
430            .load(&mut MetadataCacheMetrics::default())
431            .await
432            .unwrap();
433        let sizes: Vec<_> = metadata
434            .row_groups()
435            .iter()
436            .map(|rg| rg.column(0).uncompressed_size() as usize)
437            .collect();
438        let aligned = |bytes: usize| bytes.div_ceil(1024) * 1024;
439        let limit = aligned(*sizes.iter().max().unwrap());
440        assert!(metadata.row_group(0).total_byte_size() as usize > limit);
441        let projection = ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0]);
442        let metadata =
443            ArrowReaderMetadata::try_new(Arc::new(metadata), ArrowReaderOptions::new()).unwrap();
444        let tracker = QueryMemoryTracker::builder(limit).build();
445        let cache = CacheStrategy::EnableAll(Arc::new(
446            crate::cache::CacheManager::builder()
447                .page_cache_size(1024 * 1024)
448                .build(),
449        ));
450        let region_file_id = RegionFileId::new(
451            store_api::storage::RegionId::new(1, 1),
452            store_api::storage::FileId::random(),
453        );
454        let build = |row_group_idx, path: &str, selection| {
455            let fetcher = SstParquetRangeFetcher::new(
456                region_file_id,
457                path.into(),
458                store.clone(),
459                if path == "memory.parquet" {
460                    cache.clone()
461                } else {
462                    CacheStrategy::Disabled
463                },
464                row_group_idx,
465                None,
466            )
467            .with_memory_tracker(Some(tracker.clone()));
468            build_sst_parquet_record_batch_stream(
469                metadata.clone(),
470                row_group_idx,
471                selection,
472                projection.clone(),
473                fetcher,
474                path.into(),
475                1,
476            )
477            .unwrap()
478        };
479        // Construction (including an unpolled stream's drop) must not reserve memory.
480        drop(build(0, "memory.parquet", None));
481        assert_eq!(0, tracker.current());
482        for (row_group_idx, size) in sizes.iter().enumerate() {
483            let mut first = build(row_group_idx, "memory.parquet", None);
484            assert_eq!(1, first.next().await.unwrap().unwrap().num_rows());
485            assert_eq!(aligned(*size), tracker.current());
486            let mut competing = build(row_group_idx, "memory.parquet", None);
487            assert!(competing.next().await.unwrap().is_err());
488            assert_eq!(aligned(*size), tracker.current());
489            // Reserve before decoding, even when the selection would require no page reads.
490            let mut skipped = build(
491                row_group_idx,
492                "memory.parquet",
493                Some(RowSelection::from(vec![
494                    parquet::arrow::arrow_reader::RowSelector::skip(4),
495                ])),
496            );
497            assert!(skipped.next().await.unwrap().is_err());
498            assert_eq!(aligned(*size), tracker.current());
499            while first.next().await.transpose().unwrap().is_some() {}
500            // Completed streams may remain alive without holding their quota.
501            assert_eq!(0, tracker.current());
502        }
503        // The second read must acquire quota even when all pages are cached.
504        store.delete("memory.parquet").await.unwrap();
505        let mut cancelled = build(0, "memory.parquet", None);
506        cancelled.next().await.unwrap().unwrap();
507        drop(cancelled);
508        assert_eq!(0, tracker.current());
509        let mut failed = build(0, "missing.parquet", None);
510        assert!(failed.next().await.unwrap().is_err());
511        drop(failed);
512        assert_eq!(0, tracker.current());
513        let mut skipped = build(
514            0,
515            "memory.parquet",
516            Some(RowSelection::from(vec![
517                parquet::arrow::arrow_reader::RowSelector::skip(4),
518            ])),
519        );
520        assert!(skipped.next().await.is_none());
521        assert_eq!(0, tracker.current());
522    }
523
524    #[test]
525    fn test_assemble_range_from_cached_subrange_and_fetched_tail() {
526        let cached_parts = vec![vec![PageRangePart {
527            range: 400..500,
528            bytes: Bytes::from(vec![1; 100]),
529        }]];
530        let fetched_parts = vec![PageRangePart {
531            range: 500..600,
532            bytes: Bytes::from(vec![2; 100]),
533        }];
534
535        let requested = 400..600;
536        let output = assemble_ranges(
537            std::slice::from_ref(&requested),
538            cached_parts,
539            &fetched_parts,
540        )
541        .unwrap();
542        assert_eq!(1, output.len());
543        assert_eq!(vec![1; 100].as_slice(), &output[0][..100]);
544        assert_eq!(vec![2; 100].as_slice(), &output[0][100..]);
545    }
546
547    #[test]
548    fn test_assemble_range_returns_single_covering_part_without_copy() {
549        let bytes = Bytes::from_static(b"abcdef");
550        let cached_parts = vec![vec![PageRangePart {
551            range: 10..16,
552            bytes: bytes.clone(),
553        }]];
554
555        let requested = 10..16;
556        let output = assemble_ranges(std::slice::from_ref(&requested), cached_parts, &[]).unwrap();
557        assert_eq!(bytes, output[0]);
558    }
559
560    #[test]
561    fn test_assemble_range_clamps_overlapping_part_to_requested_end() {
562        let parts = vec![PageRangePart {
563            range: 0..10,
564            bytes: Bytes::from_static(b"0123456789"),
565        }];
566
567        let output = assemble_range(&(2..5), parts).unwrap();
568        assert_eq!(Bytes::from_static(b"234"), output);
569    }
570}
mito2/sst/parquet/push_decoder.rs

mito2/sst/parquet/
push_decoder.rs
