Attempt to produce the next sliced batch from the current batch. Returns `Some(batch)` if a slice was produced, `None` if the current batch is exhausted and we need to poll upstream for more data.
(&mut self)
| 661 | /// Returns `Some(batch)` if a slice was produced, `None` if the current batch |
| 662 | /// is exhausted and we need to poll upstream for more data. |
| 663 | fn next_sliced_batch(&mut self) -> Option<Result<RecordBatch>> { |
| 664 | let batch = self.current_batch.take()?; |
| 665 | |
| 666 | // Assert slice boundary safety - offset should never exceed batch size |
| 667 | debug_assert!( |
| 668 | self.offset <= batch.num_rows(), |
| 669 | "Offset {} exceeds batch size {}", |
| 670 | self.offset, |
| 671 | batch.num_rows() |
| 672 | ); |
| 673 | |
| 674 | let remaining = batch.num_rows() - self.offset; |
| 675 | let to_take = remaining.min(self.batch_size); |
| 676 | let out = batch.slice(self.offset, to_take); |
| 677 | |
| 678 | self.metrics.batches_split.add(1); |
| 679 | self.offset += to_take; |
| 680 | if self.offset < batch.num_rows() { |
| 681 | // More data remains in this batch, store it back |
| 682 | self.current_batch = Some(batch); |
| 683 | } else { |
| 684 | // Batch is exhausted, reset offset |
| 685 | // Note: current_batch is already None since we took it at the start |
| 686 | self.offset = 0; |
| 687 | } |
| 688 | Some(Ok(out)) |
| 689 | } |
| 690 | |
| 691 | /// Poll the upstream input for the next batch. |
| 692 | /// |