feat(parquet): add retention filter to discard unrequested prefetch data

HippoBaro · HippoBaro · commit 58ef433c9ef4 · 2026-04-14T23:59:33.000-04:00
Per-range release guarantees that every consumed row group is freed. But
when the IO layer prefetches aggressively (eg. streaming entire files or
over-reading across row group boundaries) data for row groups the
decoder will never process enters `PushBuffers` and is never consumed,
so `release_range` is never called for it.

Add `RetentionSet`, a sorted set of byte ranges derived from column
chunk metadata for the queued row groups. Incoming buffers are filtered
at push time: only portions overlapping a retained range are stored
(zero-copy via `Bytes::slice`); everything else is discarded before
reaching `PushBuffers`.

Together with the per-range release in the previous commit, this closes
the loop on memory management: the IO layer is free to push data in any
shape — coalesced, prefetched, uniform-sized, or even the entire file —
without knowledge of Parquet layout. The decoder admits only what it
will consume, and releases it at row-group boundaries.

Signed-off-by: Hippolyte Barraud &lt;hippolyte.barraud@datadoghq.com&gt;
diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs
@@ -28,6 +28,7 @@ use crate::arrow::arrow_reader::{
 use crate::errors::ParquetError;
 use crate::file::metadata::ParquetMetaData;
 use crate::util::push_buffers::PushBuffers;
+use crate::util::retention::RetentionSet;
 use arrow_array::RecordBatch;
 use bytes::Bytes;
 use reader_builder::RowGroupReaderBuilder;
@@ -185,6 +186,7 @@ impl ParquetPushDecoderBuilder {
         // Prepare to build RowGroup readers
         let file_len = 0; // not used in push decoder
         let buffers = PushBuffers::new(file_len);
+        let retention = RetentionSet::from_row_groups(&parquet_metadata, &row_groups);
         let row_group_reader_builder = RowGroupReaderBuilder::new(
             batch_size,
             projection,
@@ -197,6 +199,7 @@ impl ParquetPushDecoderBuilder {
             max_predicate_cache_size,
             buffers,
             row_selection_policy,
+            Some(retention),
         );
 
         // Initialize the decoder with the configured options
diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
@@ -34,6 +34,7 @@ use crate::errors::ParquetError;
 use crate::file::metadata::ParquetMetaData;
 use crate::file::page_index::offset_index::OffsetIndexMetaData;
 use crate::util::push_buffers::PushBuffers;
+use crate::util::retention::RetentionSet;
 use bytes::Bytes;
 use data::DataRequest;
 use filter::AdvanceResult;
@@ -168,6 +169,10 @@ pub(crate) struct RowGroupReaderBuilder {
 
     /// The underlying data store
     buffers: PushBuffers,
+
+    /// Optional retention filter.  When present, incoming `push_data` buffers
+    /// are trimmed to only keep byte ranges the decoder will eventually need.
+    retention: Option<RetentionSet>,
 }
 
 impl RowGroupReaderBuilder {
@@ -185,6 +190,7 @@ impl RowGroupReaderBuilder {
         max_predicate_cache_size: usize,
         buffers: PushBuffers,
         row_selection_policy: RowSelectionPolicy,
+        retention: Option<RetentionSet>,
     ) -> Self {
         Self {
             batch_size,
@@ -199,12 +205,23 @@ impl RowGroupReaderBuilder {
             row_selection_policy,
             state: Some(RowGroupDecoderState::Finished),
             buffers,
+            retention,
         }
     }
 
-    /// Push new data buffers that can be used to satisfy pending requests
+    /// Push new data buffers that can be used to satisfy pending requests.
+    ///
+    /// When a [`RetentionSet`] is configured, incoming buffers are filtered so
+    /// that only byte ranges the decoder will eventually need are stored.
+    /// Portions outside the retention set are silently discarded.
     pub fn push_data(&mut self, ranges: Vec<Range<u64>>, buffers: Vec<Bytes>) {
-        self.buffers.push_ranges(ranges, buffers);
+        let (ranges, buffers) = match &self.retention {
+            Some(retention) => retention.filter(ranges, buffers),
+            None => (ranges, buffers),
+        };
+        if !ranges.is_empty() {
+            self.buffers.push_ranges(ranges, buffers);
+        }
     }
 
     /// Returns the total number of buffered bytes available
diff --git a/parquet/src/util/mod.rs b/parquet/src/util/mod.rs
@@ -21,6 +21,7 @@ mod bit_pack;
 pub(crate) mod interner;
 
 pub(crate) mod push_buffers;
+pub(crate) mod retention;
 #[cfg(any(test, feature = "test_common"))]
 pub(crate) mod test_common;
 pub mod utf8;
diff --git a/parquet/src/util/retention.rs b/parquet/src/util/retention.rs
@@ -0,0 +1,240 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use bytes::Bytes;
+use std::ops::Range;
+
+use crate::file::metadata::ParquetMetaData;
+
+/// A sorted, non-overlapping set of byte ranges that the decoder expects to
+/// consume.
+///
+/// When attached to a [`RowGroupReaderBuilder`], incoming buffers are filtered
+/// against this set: only the portions that overlap a retained range are
+/// stored. Everything else is silently discarded.
+///
+/// This prevents speculatively prefetched data for row groups the decoder will
+/// never process from accumulating in memory.
+///
+/// [`RowGroupReaderBuilder`]: crate::arrow::push_decoder::reader_builder::RowGroupReaderBuilder
+#[derive(Debug, Clone)]
+pub(crate) struct RetentionSet {
+    /// Sorted, non-overlapping, merged ranges.
+    ranges: Vec<Range<u64>>,
+}
+
+impl RetentionSet {
+    /// Build a retention set from the column chunk byte ranges of the given
+    /// row groups.
+    ///
+    /// All column chunks (regardless of projection) for each queued row group
+    /// are included — this is a conservative superset of what the decoder will
+    /// actually read.
+    pub fn from_row_groups(metadata: &ParquetMetaData, row_groups: &[usize]) -> Self {
+        let total_cols: usize = row_groups
+            .iter()
+            .map(|&rg| metadata.row_group(rg).columns().len())
+            .sum();
+        let mut ranges: Vec<Range<u64>> = Vec::with_capacity(total_cols);
+        for &rg_idx in row_groups {
+            let rg = metadata.row_group(rg_idx);
+            for col in rg.columns() {
+                let (start, len) = col.byte_range();
+                ranges.push(start..start + len);
+            }
+        }
+        ranges.sort_unstable_by_key(|r| r.start);
+        let mut merged: Vec<Range<u64>> = Vec::with_capacity(ranges.len());
+        for range in ranges {
+            if let Some(last) = merged.last_mut() {
+                if range.start <= last.end {
+                    last.end = last.end.max(range.end);
+                    continue;
+                }
+            }
+            merged.push(range);
+        }
+        Self { ranges: merged }
+    }
+
+    /// Filter incoming ranges and buffers, keeping only the portions that
+    /// overlap the retention set.
+    ///
+    /// Each retained portion is a zero-copy [`Bytes::slice`] of the original
+    /// buffer. Portions that fall entirely outside the retention set are
+    /// dropped.
+    pub fn filter(
+        &self,
+        ranges: Vec<Range<u64>>,
+        buffers: Vec<Bytes>,
+    ) -> (Vec<Range<u64>>, Vec<Bytes>) {
+        let mut out_ranges = Vec::new();
+        let mut out_buffers = Vec::new();
+
+        for (range, buffer) in ranges.into_iter().zip(buffers) {
+            // Find the first retention range that could overlap: the first
+            // whose end is past range.start.
+            let start_idx = self.ranges.partition_point(|r| r.end <= range.start);
+
+            for ret in &self.ranges[start_idx..] {
+                if ret.start >= range.end {
+                    break;
+                }
+                let overlap_start = range.start.max(ret.start);
+                let overlap_end = range.end.min(ret.end);
+                let buf_offset = (overlap_start - range.start) as usize;
+                let buf_len = (overlap_end - overlap_start) as usize;
+                out_ranges.push(overlap_start..overlap_end);
+                out_buffers.push(buffer.slice(buf_offset..buf_offset + buf_len));
+            }
+        }
+
+        (out_ranges, out_buffers)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #![allow(clippy::single_range_in_vec_init)]
+    use super::*;
+
+    fn make_retention(ranges: &[Range<u64>]) -> RetentionSet {
+        let mut sorted: Vec<Range<u64>> = ranges.to_vec();
+        sorted.sort_unstable_by_key(|r| r.start);
+        let mut merged: Vec<Range<u64>> = Vec::new();
+        for range in sorted {
+            if let Some(last) = merged.last_mut() {
+                if range.start <= last.end {
+                    last.end = last.end.max(range.end);
+                    continue;
+                }
+            }
+            merged.push(range);
+        }
+        RetentionSet { ranges: merged }
+    }
+
+    #[test]
+    fn exact_match() {
+        let ret = make_retention(&[10..20]);
+        let buf = Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+        let (ranges, buffers) = ret.filter(vec![10..20], vec![buf]);
+        assert_eq!(ranges, vec![10..20]);
+        assert_eq!(buffers.len(), 1);
+        assert_eq!(&*buffers[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+    }
+
+    #[test]
+    fn no_overlap() {
+        let ret = make_retention(&[10..20]);
+        let buf = Bytes::from_static(&[1, 2, 3]);
+        let (ranges, buffers) = ret.filter(vec![0..3], vec![buf]);
+        assert!(ranges.is_empty());
+        assert!(buffers.is_empty());
+    }
+
+    #[test]
+    fn partial_overlap_left() {
+        let ret = make_retention(&[10..20]);
+        let buf = Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+        // Buffer covers 5..15, retention is 10..20 → keep 10..15
+        let (ranges, buffers) = ret.filter(vec![5..15], vec![buf]);
+        assert_eq!(ranges, vec![10..15]);
+        assert_eq!(&*buffers[0], &[6, 7, 8, 9, 10]);
+    }
+
+    #[test]
+    fn partial_overlap_right() {
+        let ret = make_retention(&[10..20]);
+        let buf = Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+        // Buffer covers 15..25, retention is 10..20 → keep 15..20
+        let (ranges, buffers) = ret.filter(vec![15..25], vec![buf]);
+        assert_eq!(ranges, vec![15..20]);
+        assert_eq!(&*buffers[0], &[1, 2, 3, 4, 5]);
+    }
+
+    #[test]
+    fn buffer_spans_gap_between_retention_ranges() {
+        // Retention: [10..20) and [30..40). Buffer covers 5..45.
+        let ret = make_retention(&[10..20, 30..40]);
+        let data: Vec<u8> = (0..40).collect();
+        let buf = Bytes::from(data);
+        let (ranges, buffers) = ret.filter(vec![5..45], vec![buf]);
+        assert_eq!(ranges, vec![10..20, 30..40]);
+        assert_eq!(buffers.len(), 2);
+        // First slice: bytes at offset 5..15 in the buffer (values 5..15)
+        assert_eq!(&*buffers[0], &[5, 6, 7, 8, 9, 10, 11, 12, 13, 14]);
+        // Second slice: bytes at offset 25..35 in the buffer (values 25..35)
+        assert_eq!(&*buffers[1], &[25, 26, 27, 28, 29, 30, 31, 32, 33, 34]);
+    }
+
+    #[test]
+    fn superset_buffer_trimmed() {
+        let ret = make_retention(&[10..20]);
+        let data: Vec<u8> = (0..50).collect();
+        let buf = Bytes::from(data);
+        let (ranges, buffers) = ret.filter(vec![0..50], vec![buf]);
+        assert_eq!(ranges, vec![10..20]);
+        assert_eq!(&*buffers[0], &[10, 11, 12, 13, 14, 15, 16, 17, 18, 19]);
+    }
+
+    #[test]
+    fn empty_retention_discards_everything() {
+        let ret = RetentionSet { ranges: Vec::new() };
+        let buf = Bytes::from_static(&[1, 2, 3]);
+        let (ranges, buffers) = ret.filter(vec![0..3], vec![buf]);
+        assert!(ranges.is_empty());
+        assert!(buffers.is_empty());
+    }
+
+    #[test]
+    fn multiple_input_buffers() {
+        let ret = make_retention(&[10..20, 30..40]);
+        let buf1 = Bytes::from_static(&[1, 2, 3, 4, 5]);
+        let buf2 = Bytes::from_static(&[1, 2, 3, 4, 5]);
+        let buf3 = Bytes::from_static(&[1, 2, 3, 4, 5]);
+        let (ranges, buffers) = ret.filter(vec![0..5, 10..15, 35..40], vec![buf1, buf2, buf3]);
+        // First buffer: no overlap. Second: exact. Third: exact.
+        assert_eq!(ranges, vec![10..15, 35..40]);
+        assert_eq!(buffers.len(), 2);
+    }
+
+    #[test]
+    fn zero_copy_slicing() {
+        let ret = make_retention(&[10..20]);
+        let data: Vec<u8> = (0..30).collect();
+        let buf = Bytes::from(data);
+        let original_ptr = buf.as_ptr();
+        let (_, buffers) = ret.filter(vec![0..30], vec![buf]);
+        // The output slice should point into the same allocation,
+        // offset by 10 bytes.
+        assert_eq!(buffers[0].as_ptr(), unsafe { original_ptr.add(10) },);
+    }
+
+    #[test]
+    fn adjacent_retention_ranges_are_merged() {
+        // Two abutting ranges should merge into one.
+        let ret = make_retention(&[10..20, 20..30]);
+        assert_eq!(ret.ranges, vec![10..30]);
+        let data: Vec<u8> = (0..40).collect();
+        let buf = Bytes::from(data);
+        let (ranges, buffers) = ret.filter(vec![0..40], vec![buf]);
+        // Should produce a single slice, not two.
+        assert_eq!(ranges, vec![10..30]);
+        assert_eq!(buffers.len(), 1);
+    }
+}
