| 485 | } |
| 486 | |
| 487 | fn next_batch(&mut self, batch_size: usize) -> Option<Vec<usize>> { |
| 488 | if batch_size == 0 { |
| 489 | return None; |
| 490 | } |
| 491 | |
| 492 | match &mut self.state { |
| 493 | RowSelectionState::All { next, total_rows } => { |
| 494 | if *next >= *total_rows { |
| 495 | return None; |
| 496 | } |
| 497 | |
| 498 | let end = (*next + batch_size).min(*total_rows); |
| 499 | let batch: Vec<usize> = (*next..end).collect(); |
| 500 | *next = end; |
| 501 | Some(batch) |
| 502 | } |
| 503 | RowSelectionState::Ranges { |
| 504 | total_rows, |
| 505 | ranges, |
| 506 | range_idx, |
| 507 | next_in_range, |
| 508 | } => { |
| 509 | if *range_idx >= ranges.len() || *total_rows == 0 { |
| 510 | return None; |
| 511 | } |
| 512 | |
| 513 | let mut batch = Vec::with_capacity(batch_size); |
| 514 | while batch.len() < batch_size && *range_idx < ranges.len() { |
| 515 | let range = &ranges[*range_idx]; |
| 516 | if *next_in_range > range.to() { |
| 517 | *range_idx += 1; |
| 518 | if *range_idx < ranges.len() { |
| 519 | *next_in_range = ranges[*range_idx].from(); |
| 520 | } |
| 521 | continue; |
| 522 | } |
| 523 | |
| 524 | batch.push(*next_in_range as usize); |
| 525 | *next_in_range += 1; |
| 526 | } |
| 527 | |
| 528 | if batch.is_empty() { |
| 529 | None |
| 530 | } else { |
| 531 | Some(batch) |
| 532 | } |
| 533 | } |
| 534 | } |
| 535 | } |
| 536 | } |
| 537 | |
| 538 | fn decode_delta_varints(bytes: &[u8]) -> crate::Result<Vec<i64>> { |