| 181 | } |
| 182 | |
| 183 | fn execute( |
| 184 | &self, |
| 185 | partition: usize, |
| 186 | context: Arc<TaskContext>, |
| 187 | ) -> Result<SendableRecordBatchStream> { |
| 188 | trace!( |
| 189 | "Start AsyncFuncExpr::execute for partition {} of context session_id {} and task_id {:?}", |
| 190 | partition, |
| 191 | context.session_id(), |
| 192 | context.task_id() |
| 193 | ); |
| 194 | |
| 195 | // first execute the input stream |
| 196 | let input_stream = self.input.execute(partition, Arc::clone(&context))?; |
| 197 | |
| 198 | // TODO: Track `elapsed_compute` in `BaselineMetrics` |
| 199 | // Issue: <https://github.com/apache/datafusion/issues/19658> |
| 200 | let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); |
| 201 | |
| 202 | // now, for each record batch, evaluate the async expressions and add the columns to the result |
| 203 | let async_exprs_captured = Arc::new(self.async_exprs.clone()); |
| 204 | let schema_captured = self.schema(); |
| 205 | let config_options_ref = Arc::clone(context.session_config().options()); |
| 206 | |
| 207 | let coalesced_input_stream = CoalesceInputStream { |
| 208 | input_stream, |
| 209 | batch_coalescer: LimitedBatchCoalescer::new( |
| 210 | Arc::clone(&self.input.schema()), |
| 211 | config_options_ref.execution.batch_size, |
| 212 | None, |
| 213 | ), |
| 214 | }; |
| 215 | |
| 216 | let stream_with_async_functions = coalesced_input_stream.then(move |batch| { |
| 217 | // need to clone *again* to capture the async_exprs and schema in the |
| 218 | // stream and satisfy lifetime requirements. |
| 219 | let async_exprs_captured = Arc::clone(&async_exprs_captured); |
| 220 | let schema_captured = Arc::clone(&schema_captured); |
| 221 | let config_options = Arc::clone(&config_options_ref); |
| 222 | let baseline_metrics_captured = baseline_metrics.clone(); |
| 223 | |
| 224 | async move { |
| 225 | let batch = batch?; |
| 226 | // append the result of evaluating the async expressions to the output |
| 227 | let mut output_arrays = batch.columns().to_vec(); |
| 228 | for async_expr in async_exprs_captured.iter() { |
| 229 | let output = async_expr |
| 230 | .invoke_with_args(&batch, Arc::clone(&config_options)) |
| 231 | .await?; |
| 232 | output_arrays.push(output.to_array(batch.num_rows())?); |
| 233 | } |
| 234 | let batch = RecordBatch::try_new(schema_captured, output_arrays)?; |
| 235 | |
| 236 | Ok(batch.record_output(&baseline_metrics_captured)) |
| 237 | } |
| 238 | }); |
| 239 | |
| 240 | // Adapt the stream with the output schema |