Method generate_next_batch

datafusion/core/tests/execution/coop.rs:93–112

Generate the next RecordBatch.

(&mut self)

    /// Generate the next RecordBatch.
    fn generate_next_batch(&mut self) -> datafusion_common::Result<Option<RecordBatch>> {
        self.poll_count += 1;

        let mut builder = Int64Array::builder(self.batch_size);
        for _ in 0..self.batch_size {
            match self.value_range.next() {
                None => break,
                Some(v) => builder.append_value(v),
            }
        }
        let array = builder.finish();

        if array.is_empty() {
            return Ok(None);
        }

        let batch =
            RecordBatch::try_new(Arc::clone(&self.schema), vec![Arc::new(array)])?;
        Ok(Some(batch))
    }

    fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> {
        let mut new = self.clone();