metesynnada commented on code in PR #5520:
URL: https://github.com/apache/arrow-datafusion/pull/5520#discussion_r1133821465
##########
datafusion/core/src/datasource/memory.rs:
##########
@@ -388,4 +465,115 @@ mod tests {
Ok(())
}
+
+ #[tokio::test]
+ async fn test_insert_into_single_partition() -> Result<()> {
+ // Create a new session context
+ let session_ctx = SessionContext::new();
+ // Create a new schema with one field called "a" of type Int32
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, false)]));
+
+ // Create a new batch of data to insert into the table
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(Int32Array::from_slice([1, 2, 3]))],
+ )?;
+ // Create a new table with one partition that contains the batch of
data
+ let initial_table = Arc::new(MemTable::try_new(
+ schema.clone(),
+ vec![vec![batch.clone()]],
+ )?);
+ // Convert the table into a provider so that it can be used in a query
+ let provider = provider_as_source(Arc::new(MemTable::try_new(
+ schema.clone(),
+ vec![vec![batch.clone()]],
+ )?));
+ // Create a table scan logical plan to read from the table
+ let table_scan =
+ Arc::new(LogicalPlanBuilder::scan("source", provider,
None)?.build()?);
+ // Insert the data from the provider into the table
+ initial_table
+ .insert_into(&session_ctx.state(), &table_scan)
+ .await?;
+ // Ensure that the table now contains two batches of data in the same
partition
+ assert_eq!(initial_table.batches.read().await.get(0).unwrap().len(),
2);
+
+ // Create a new provider with 2 partitions
+ let multi_partition_provider =
provider_as_source(Arc::new(MemTable::try_new(
+ schema.clone(),
+ vec![vec![batch.clone()], vec![batch]],
+ )?));
+ // Create a new table scan logical plan to read from the provider
+ let table_scan = Arc::new(
+ LogicalPlanBuilder::scan("source", multi_partition_provider, None)?
+ .build()?,
+ );
Review Comment:
Addressed 👍🏻
##########
datafusion/core/src/execution/context.rs:
##########
@@ -2714,6 +2728,46 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn sql_table_insert() -> Result<()> {
Review Comment:
Addressed 👍🏻
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
