| 608 | } |
| 609 | |
| 610 | fn execute( |
| 611 | &self, |
| 612 | _partition: usize, |
| 613 | _context: Arc<datafusion_execution::TaskContext>, |
| 614 | ) -> Result<datafusion_execution::SendableRecordBatchStream> { |
| 615 | // Create a single batch with the count |
| 616 | let count_array = UInt64Array::from(vec![self.rows_affected]); |
| 617 | let batch = ArrowRecordBatch::try_new( |
| 618 | Arc::clone(&self.schema), |
| 619 | vec![Arc::new(count_array) as ArrayRef], |
| 620 | )?; |
| 621 | |
| 622 | // Create a stream that yields just this one batch |
| 623 | let stream = futures::stream::iter(vec![Ok(batch)]); |
| 624 | Ok(Box::pin(RecordBatchStreamAdapter::new( |
| 625 | Arc::clone(&self.schema), |
| 626 | stream, |
| 627 | ))) |
| 628 | } |
| 629 | } |