(
self,
ctx: &dyn CliSessionContext,
print_options: &PrintOptions,
)
| 257 | } |
| 258 | |
| 259 | async fn execute( |
| 260 | self, |
| 261 | ctx: &dyn CliSessionContext, |
| 262 | print_options: &PrintOptions, |
| 263 | ) -> Result<()> { |
| 264 | let now = Instant::now(); |
| 265 | let (df, adjusted) = self |
| 266 | .create_and_execute_logical_plan(ctx, print_options) |
| 267 | .await?; |
| 268 | let physical_plan = df.create_physical_plan().await?; |
| 269 | let task_ctx = ctx.task_ctx(); |
| 270 | let options = task_ctx.session_config().options(); |
| 271 | |
| 272 | // Track memory usage for the query result if it's bounded |
| 273 | let reservation = |
| 274 | MemoryConsumer::new("DataFusion-Cli").register(task_ctx.memory_pool()); |
| 275 | |
| 276 | if physical_plan.boundedness().is_unbounded() { |
| 277 | if physical_plan.pipeline_behavior() == EmissionType::Final { |
| 278 | return plan_err!( |
| 279 | "The given query can generate a valid result only once \ |
| 280 | the source finishes, but the source is unbounded" |
| 281 | ); |
| 282 | } |
| 283 | // As the input stream comes, we can generate results. |
| 284 | // However, memory safety is not guaranteed. |
| 285 | let stream = execute_stream(physical_plan, task_ctx.clone())?; |
| 286 | print_options |
| 287 | .print_stream(stream, now, &options.format) |
| 288 | .await?; |
| 289 | } else { |
| 290 | // Bounded stream; collected results size is limited by the maxrows option |
| 291 | let schema = physical_plan.schema(); |
| 292 | let mut stream = execute_stream(physical_plan, task_ctx.clone())?; |
| 293 | let mut results = vec![]; |
| 294 | let mut row_count = 0_usize; |
| 295 | let max_rows = match print_options.maxrows { |
| 296 | MaxRows::Unlimited => usize::MAX, |
| 297 | MaxRows::Limited(n) => n, |
| 298 | }; |
| 299 | while let Some(batch) = stream.next().await { |
| 300 | let batch = batch?; |
| 301 | let curr_num_rows = batch.num_rows(); |
| 302 | // Stop collecting results if the number of rows exceeds the limit |
| 303 | // results batch should include the last batch that exceeds the limit |
| 304 | if row_count < max_rows.saturating_add(curr_num_rows) { |
| 305 | // Try to grow the reservation to accommodate the batch in memory |
| 306 | reservation.try_grow(get_record_batch_memory_size(&batch))?; |
| 307 | results.push(batch); |
| 308 | } |
| 309 | row_count += curr_num_rows; |
| 310 | } |
| 311 | adjusted.into_inner().print_batches( |
| 312 | schema, |
| 313 | &results, |
| 314 | now, |
| 315 | row_count, |
| 316 | &options.format, |
no test coverage detected