martin-g commented on code in PR #25040:
URL: https://github.com/apache/datafusion/pull/25040#discussion_r3956245380
##########
datafusion/core/src/datasource/memory_test.rs:
##########
@@ -501,4 +503,359 @@ mod tests {
);
Ok(())
}
+
+ /// A schema of one non-nullable Int32 column called "a".
+ fn one_column_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
+ }
+
+ /// A batch of the rows 1, 2 and 3 in the column "a".
+ fn one_column_batch(schema: &SchemaRef) -> Result<RecordBatch> {
+ Ok(RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+ )?)
+ }
+
+ /// A `MemTable` of one partition holding the rows 1, 2 and 3.
+ fn one_column_table() -> Result<MemTable> {
+ let schema = one_column_schema();
+ MemTable::try_new(Arc::clone(&schema),
vec![vec![one_column_batch(&schema)?]])
+ }
+
+ /// A `MemTable` that holds no partition. `MemTable::try_new` rejects an
+ /// empty partition list, so a caller reaches this state through the public
+ /// `batches` field.
+ fn zero_partition_table(schema: SchemaRef) -> Result<MemTable> {
+ let mut table = MemTable::try_new(schema, vec![vec![]])?;
+ table.batches.clear();
+ Ok(table)
+ }
+
+ /// Run a DELETE or an UPDATE plan and return the count that it emits.
+ async fn run_dml(
+ plan: Arc<dyn ExecutionPlan>,
+ session_ctx: &SessionContext,
+ ) -> Result<u64> {
+ Ok(extract_count(collect(plan, session_ctx.task_ctx()).await?))
+ }
+
+ /// Read one partition of a `MemTable` as a plain vector of batches.
+ async fn read_partition(table: &MemTable, partition: usize) ->
Vec<RecordBatch> {
+ table.batches[partition].read().await.clone()
+ }
+
+ /// The values of the first column of `batch`, which must hold no null.
+ fn column_values(batch: &RecordBatch) -> Vec<i32> {
+ batch
+ .column(0)
+ .as_primitive::<Int32Type>()
+ .iter()
+ .map(|value| value.expect("expected non null"))
+ .collect()
+ }
+
+ // A DELETE on a table without a partition affects no row
+ #[tokio::test]
+ async fn test_delete_from_zero_partition() -> Result<()> {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+ let table = zero_partition_table(one_column_schema())?;
+
+ let plan = table.delete_from(&state, vec![col("a").gt(lit(1))]).await?;
+ assert_eq!(run_dml(plan, &session_ctx).await?, 0);
+ Ok(())
+ }
+
+ // An UPDATE on a table without a partition affects no row
+ #[tokio::test]
+ async fn test_update_zero_partition() -> Result<()> {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+ let table = zero_partition_table(one_column_schema())?;
+
+ let plan = table
+ .update(&state, vec![("a".to_string(), lit(7))], vec![])
+ .await?;
+ assert_eq!(run_dml(plan, &session_ctx).await?, 0);
+ Ok(())
+ }
+
+ // A DELETE skips a batch of no row and drops it from the partition
+ #[tokio::test]
+ async fn test_delete_from_empty_batch() -> Result<()> {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+ let schema = one_column_schema();
+ let table = MemTable::try_new(
+ Arc::clone(&schema),
+ vec![vec![
+ RecordBatch::new_empty(Arc::clone(&schema)),
+ one_column_batch(&schema)?,
+ ]],
+ )?;
+
+ let plan = table.delete_from(&state, vec![col("a").gt(lit(1))]).await?;
+ assert_eq!(run_dml(plan, &session_ctx).await?, 2);
+
+ // The empty batch is gone and the row 1 remains
+ let partition = read_partition(&table, 0).await;
+ assert_eq!(partition.len(), 1);
+ assert_eq!(column_values(&partition[0]), vec![1]);
+ Ok(())
+ }
+
+ // An UPDATE skips a batch of no row and drops it from the partition
+ #[tokio::test]
+ async fn test_update_empty_batch() -> Result<()> {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+ let schema = one_column_schema();
+ let table = MemTable::try_new(
+ Arc::clone(&schema),
+ vec![vec![
+ RecordBatch::new_empty(Arc::clone(&schema)),
+ one_column_batch(&schema)?,
+ ]],
+ )?;
+
+ let plan = table
+ .update(
+ &state,
+ vec![("a".to_string(), lit(7))],
+ vec![col("a").gt(lit(1))],
+ )
+ .await?;
+ assert_eq!(run_dml(plan, &session_ctx).await?, 2);
+
+ // The empty batch is gone and the rows 2 and 3 now hold 7
+ let partition = read_partition(&table, 0).await;
+ assert_eq!(partition.len(), 1);
+ assert_eq!(column_values(&partition[0]), vec![1, 7, 7]);
+ Ok(())
+ }
+
+ // The DELETE plan has one partition and rejects a request for another
+ #[tokio::test]
+ async fn test_delete_exec_rejects_other_partition() -> Result<()> {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+ let table = one_column_table()?;
+
+ let plan = table.delete_from(&state, vec![]).await?;
+ let Err(err) = plan.execute(1, session_ctx.task_ctx()) else {
+ panic!("expected an error for partition 1");
+ };
+ assert_contains!(
+ err.strip_backtrace(),
+ "MemDeleteExec has one partition, but partition 1 was requested"
+ );
+
+ // The failed request leaves the rows alone
+ assert_eq!(read_partition(&table, 0).await[0].num_rows(), 3);
+ Ok(())
+ }
+
+ // The UPDATE plan has one partition and rejects a request for another
+ #[tokio::test]
+ async fn test_update_exec_rejects_other_partition() -> Result<()> {
+ let session_ctx = SessionContext::new();
+ let state = session_ctx.state();
+ let table = one_column_table()?;
+
+ let plan = table
+ .update(&state, vec![("a".to_string(), lit(7))], vec![])
+ .await?;
+ let Err(err) = plan.execute(1, session_ctx.task_ctx()) else {
+ panic!("expected an error for partition 1");
+ };
+ assert_contains!(
+ err.strip_backtrace(),
+ "MemUpdateExec has one partition, but partition 1 was requested"
+ );
+
+ // The failed request leaves the rows alone
+ assert_eq!(
+ column_values(&read_partition(&table, 0).await[0]),
+ vec![1, 2, 3]
+ );
+ Ok(())
+ }
+
+ // A DELETE whose `WHERE` clause names an unknown column fails while the
+ // plan is built, before any row changes
+ #[tokio::test]
Review Comment:
SqlLogicTests are preferred than unit tests when possible because they don't
require compilation.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
