| 526 | } |
| 527 | |
| 528 | async fn experiment( |
| 529 | schema: SchemaRef, |
| 530 | initial_data: Vec<Vec<RecordBatch>>, |
| 531 | inserted_data: Vec<Vec<RecordBatch>>, |
| 532 | ) -> Result<Vec<Vec<RecordBatch>>> { |
| 533 | let expected_count: u64 = inserted_data |
| 534 | .iter() |
| 535 | .flat_map(|batches| batches.iter().map(|batch| batch.num_rows() as u64)) |
| 536 | .sum(); |
| 537 | |
| 538 | // Create a new session context |
| 539 | let session_ctx = SessionContext::new(); |
| 540 | // Create and register the initial table with the provided schema and data |
| 541 | let initial_table = Arc::new(CacheTable::try_new(schema.clone(), initial_data)?); |
| 542 | session_ctx.register_table("t", initial_table.clone())?; |
| 543 | // Create and register the source table with the provided schema and inserted data |
| 544 | let source_table = Arc::new(CacheTable::try_new(schema.clone(), inserted_data)?); |
| 545 | session_ctx.register_table("source", source_table.clone())?; |
| 546 | // Convert the source table into a provider so that it can be used in a query |
| 547 | let source = provider_as_source(source_table); |
| 548 | // Create a table scan logical plan to read from the source table |
| 549 | let scan_plan = LogicalPlanBuilder::scan("source", source, None)?.build()?; |
| 550 | // Create an insert plan to insert the source data into the initial table |
| 551 | let insert_into_table = |
| 552 | LogicalPlanBuilder::insert_into(scan_plan, "t", &schema, InsertOp::Append)?.build()?; |
| 553 | // Create a physical plan from the insert plan |
| 554 | let plan = session_ctx |
| 555 | .state() |
| 556 | .create_physical_plan(&insert_into_table) |
| 557 | .await?; |
| 558 | |
| 559 | // Execute the physical plan and collect the results |
| 560 | let res = collect(plan, session_ctx.task_ctx()).await?; |
| 561 | assert_eq!(extract_count(res), expected_count); |
| 562 | |
| 563 | // Read the data from the initial table and store it in a vector of partitions |
| 564 | let mut partitions = vec![]; |
| 565 | for partition in initial_table.batches.iter() { |
| 566 | let part = partition.read().await.clone(); |
| 567 | partitions.push(part); |
| 568 | } |
| 569 | Ok(partitions) |
| 570 | } |
| 571 | |
| 572 | /// Returns the value of results. For example, returns 6 given the following |
| 573 | /// |