zhuqi-lucas commented on code in PR #10901:
URL: https://github.com/apache/arrow-rs/pull/10901#discussion_r4119104400
##########
parquet/src/arrow/array_reader/cached_array_reader.rs:
##########
@@ -168,22 +172,33 @@ impl CachedArrayReader {
/// Remove batches from cache that have been completely consumed
/// This is only called for Consumer role readers
- fn cleanup_consumed_batches(&self) {
+ fn cleanup_consumed_batches(&mut self) {
let current_batch_id =
self.get_batch_id_from_position(self.outer_position);
// Remove batches that are at least one batch behind the current
position
// This ensures we don't remove batches that might still be needed for
the current batch
// We can safely remove batch_id if current_batch_id > batch_id + 1
- if current_batch_id.val > 1 {
- let mut cache = self.shared_cache.write().unwrap();
- for batch_id_to_remove in 0..(current_batch_id.val - 1) {
- cache.remove(
- self.column_idx,
- BatchID {
- val: batch_id_to_remove,
- },
- );
- }
+ if current_batch_id.val <= 1 {
+ return;
+ }
+ let end = current_batch_id.val - 1;
Review Comment:
I left this one as is, because it would change behaviour rather than
simplify. With `end = current_batch_id.val` the loop removes up to `current -
1`; the `-1` keeps it at `current - 2`, and that boundary is what `main` does
today:
```rust
if current_batch_id.val > 1 {
for batch_id_to_remove in 0..(current_batch_id.val - 1) {
```
This PR only makes that same range incremental, so widening it by one batch
is a separate claim — that no consumer can still ask for `current - 1`. That
may well be true (the `local_cache.retain` just above keeps `>= current`, which
points the same way), but it deserves its own PR and its own reasoning rather
than riding along with a cleanup change. Happy to open one if you think the
wider range is right.
##########
parquet/src/arrow/array_reader/cached_array_reader.rs:
##########
@@ -168,23 +172,34 @@ impl CachedArrayReader {
/// Remove batches from cache that have been completely consumed
/// This is only called for Consumer role readers
- fn cleanup_consumed_batches(&self) {
+ fn cleanup_consumed_batches(&mut self) {
let current_batch_id =
self.get_batch_id_from_position(self.outer_position);
// Remove batches that are at least one batch behind the current
position
// This ensures we don't remove batches that might still be needed for
the current batch
// We can safely remove batch_id if current_batch_id > batch_id + 1
- if current_batch_id.val > 1 {
- let mut cache = self.shared_cache.write().unwrap();
- for batch_id_to_remove in 0..(current_batch_id.val - 1) {
- cache.remove(
- self.column_idx,
- BatchID {
- val: batch_id_to_remove,
- },
- );
- }
+ if current_batch_id.val <= 1 {
+ return;
+ }
+ let end = current_batch_id.val - 1;
+ // Everything below `cleaned_up_to` was removed by an earlier call.
+ // Rescanning from 0 each time made this quadratic in the number of
+ // batches in the row group, and took the shared cache's write lock on
+ // every `consume_batch` even when there was nothing left to remove.
+ if end <= self.cleaned_up_to {
+ return;
+ }
+ let mut cache = self.shared_cache.write().unwrap();
+ for batch_id_to_remove in self.cleaned_up_to..end {
+ cache.remove(
+ self.column_idx,
+ BatchID {
+ val: batch_id_to_remove,
+ },
+ );
}
+ drop(cache);
Review Comment:
This one is from the earlier round and the code has since moved —
`cleaned_up_to` is now set *before* the write lock is taken:
```rust
let start = self.cleaned_up_to;
self.cleaned_up_to = end;
let mut cache = self.shared_cache.write().unwrap();
```
The watermark is per-reader state, not shared, so it does not need the lock;
taking the lock last also keeps the critical section down to the removals
themselves.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
