This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11252-b94434d8c0880243fdd80c813a7ee6eb4a29ce72 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 103ebe7548db35c37e8ff35fcb02438bcd182fa7 Author: Andrew Lamb <[email protected]> AuthorDate: Mon Oct 5 01:46:48 2026 +0000 test: Assert CachedArrayReader cleanup watermark in consumer tests (#11252) # Which issue does this PR close? - Follow-up to #10901 from @zhuqi-lucas # Rationale for this change The existing tests wouldn't fail if we messed up the accounting of the newly added `cleaned_up_to` field. That would be nice to catch # What changes are included in this PR? Added some new tests that `cleaned_up_to` is kept properly updated # Are these changes tested? Yes, they are tests. # Are there any user-facing changes? No. Co-authored-by: Claude Fable 5.1 <[email protected]> --- .../src/arrow/array_reader/cached_array_reader.rs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/parquet/src/arrow/array_reader/cached_array_reader.rs b/parquet/src/arrow/array_reader/cached_array_reader.rs index 980d67f31b..43f8e18f2a 100644 --- a/parquet/src/arrow/array_reader/cached_array_reader.rs +++ b/parquet/src/arrow/array_reader/cached_array_reader.rs @@ -614,6 +614,8 @@ mod tests { // After first consume_batch, batch 0 should still be in cache // (current_batch_id = 3/3 = 1, cleanup only happens if current_batch_id > 1) assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_some()); + // Nothing was removed, so the watermark has not moved + assert_eq!(consumer_reader.cleaned_up_to, 0); // Read second batch (positions 3-5, batch 1) let read2 = consumer_reader.read_records(3).unwrap(); @@ -626,6 +628,8 @@ mod tests { // (current_batch_id = 6/3 = 2, cleanup removes batches 0..(2-1) = 0..1, so removes batch 0) assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_none()); assert!(cache.read().unwrap().get(0, BatchID { val: 1 }).is_some()); + // Watermark records that batch 0 is gone + assert_eq!(consumer_reader.cleaned_up_to, 1); // Read third batch (positions 6-8, batch 2) let read3 = consumer_reader.read_records(3).unwrap(); @@ -639,6 +643,19 @@ mod tests { assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_none()); assert!(cache.read().unwrap().get(0, BatchID { val: 1 }).is_none()); assert!(cache.read().unwrap().get(0, BatchID { val: 2 }).is_some()); + // Watermark advanced by exactly one batch: only batch 1 was newly removed + assert_eq!(consumer_reader.cleaned_up_to, 2); + + // Read one more record (position 9, still batch 3) and consume again. + // current_batch_id is unchanged, so cleanup has nothing new to remove: + // the watermark must stay put and batch 2 must remain cached. + let read4 = consumer_reader.read_records(1).unwrap(); + assert_eq!(read4, 1); + assert_eq!(consumer_reader.outer_position, 10); + let array4 = consumer_reader.consume_batch().unwrap(); + assert_eq!(array4.len(), 1); + assert_eq!(consumer_reader.cleaned_up_to, 2); + assert!(cache.read().unwrap().get(0, BatchID { val: 2 }).is_some()); } #[test] @@ -683,6 +700,8 @@ mod tests { ); assert_eq!(consumer.skip_records(9).unwrap(), 9); assert_eq!(consumer.read_records(3).unwrap(), 3); + // Skipping and reading do not touch the shared cache; only consume does + assert_eq!(consumer.cleaned_up_to, 0); let array = consumer.consume_batch().unwrap(); assert_eq!(array.len(), 3); @@ -698,6 +717,8 @@ mod tests { ); } assert!(cache.read().unwrap().get(0, BatchID { val: 3 }).is_some()); + // Watermark jumped over the skipped batches in one step + assert_eq!(consumer.cleaned_up_to, 3); } #[test]
