| 520 | } |
| 521 | |
| 522 | fn compact_retained_state(&mut self, emit_groups: usize) -> Result<()> { |
| 523 | // EmitTo::First is used to recover from memory pressure. Simply |
| 524 | // removing emitted entries in place is not enough because mixed batches |
| 525 | // would continue to pin their original Array arrays, even if only a few |
| 526 | // retained rows remain. |
| 527 | // |
| 528 | // Rebuild the retained state from scratch so fully emitted batches are |
| 529 | // dropped, mixed batches are compacted to arrays containing only the |
| 530 | // surviving rows, and retained metadata is right-sized. |
| 531 | let emit_groups = emit_groups as u32; |
| 532 | let old_batches = take(&mut self.batches); |
| 533 | let old_batch_entries = take(&mut self.batch_entries); |
| 534 | |
| 535 | let mut batches = Vec::new(); |
| 536 | let mut batch_entries = Vec::new(); |
| 537 | |
| 538 | for (batch, entries) in old_batches.into_iter().zip(old_batch_entries) { |
| 539 | let retained_len = entries.iter().filter(|(g, _)| *g >= emit_groups).count(); |
| 540 | |
| 541 | if retained_len == 0 { |
| 542 | continue; |
| 543 | } |
| 544 | |
| 545 | if retained_len == entries.len() { |
| 546 | // Nothing was emitted from this batch, so we keep the existing |
| 547 | // array and only renumber the remaining group IDs so that they |
| 548 | // start from 0. |
| 549 | let mut retained_entries = entries; |
| 550 | for (g, _) in &mut retained_entries { |
| 551 | *g -= emit_groups; |
| 552 | } |
| 553 | retained_entries.shrink_to_fit(); |
| 554 | batches.push(batch); |
| 555 | batch_entries.push(retained_entries); |
| 556 | continue; |
| 557 | } |
| 558 | |
| 559 | let mut retained_entries = Vec::with_capacity(retained_len); |
| 560 | let mut retained_rows = Vec::with_capacity(retained_len); |
| 561 | |
| 562 | for (g, r) in entries { |
| 563 | if g >= emit_groups { |
| 564 | // Compute the new `(group_idx, row_idx)` pair for a |
| 565 | // retained row. `group_idx` is renumbered to start from |
| 566 | // 0, and `row_idx` points into the new dense batch we are |
| 567 | // building. |
| 568 | retained_entries.push((g - emit_groups, retained_rows.len() as u32)); |
| 569 | retained_rows.push(r); |
| 570 | } |
| 571 | } |
| 572 | |
| 573 | debug_assert_eq!(retained_entries.len(), retained_len); |
| 574 | debug_assert_eq!(retained_rows.len(), retained_len); |
| 575 | |
| 576 | let batch = if retained_len == batch.len() { |
| 577 | batch |
| 578 | } else { |
| 579 | // Compact mixed batches so retained rows no longer pin the |