| 757 | } |
| 758 | |
| 759 | async fn write_all( |
| 760 | &self, |
| 761 | mut data: SendableRecordBatchStream, |
| 762 | _context: &Arc<TaskContext>, |
| 763 | ) -> Result<u64> { |
| 764 | let num_partitions = self.batches.len(); |
| 765 | |
| 766 | // buffer up the data round robin style into num_partitions |
| 767 | |
| 768 | let mut new_batches = vec![vec![]; num_partitions]; |
| 769 | let mut i = 0; |
| 770 | let mut row_count = 0; |
| 771 | while let Some(batch) = data.next().await.transpose()? { |
| 772 | row_count += batch.num_rows(); |
| 773 | new_batches[i].push(batch); |
| 774 | i = (i + 1) % num_partitions; |
| 775 | } |
| 776 | |
| 777 | // write the outputs into the batches |
| 778 | for (target, mut batches) in self.batches.iter().zip(new_batches) { |
| 779 | // Append all the new batches in one go to minimize locking overhead |
| 780 | target.write().await.append(&mut batches); |
| 781 | } |
| 782 | |
| 783 | Ok(row_count as u64) |
| 784 | } |
| 785 | } |
| 786 | |
| 787 | #[cfg(test)] |