| 758 | } |
| 759 | |
| 760 | fn sample(&mut self, batch: &RecordBatch) -> Result<RecordBatch> { |
| 761 | let range = self.upper_bound - self.lower_bound; |
| 762 | if range <= 0.0 { |
| 763 | return Ok(RecordBatch::new_empty(batch.schema())); |
| 764 | } |
| 765 | |
| 766 | // Select rows where random value falls in [lower, upper) |
| 767 | let indices: Vec<u32> = (0..batch.num_rows()) |
| 768 | .filter(|_| { |
| 769 | let r: f64 = self.rng.random(); |
| 770 | r >= self.lower_bound && r < self.upper_bound |
| 771 | }) |
| 772 | .map(|i| i as u32) |
| 773 | .collect(); |
| 774 | |
| 775 | if indices.is_empty() { |
| 776 | return Ok(RecordBatch::new_empty(batch.schema())); |
| 777 | } |
| 778 | |
| 779 | compute::take_record_batch(batch, &UInt32Array::from(indices)) |
| 780 | .map_err(DataFusionError::from) |
| 781 | } |
| 782 | } |
| 783 | |
| 784 | /// Stream adapter that applies sampling to each batch. |