2010YOUY01 commented on code in PR #13377:
URL: https://github.com/apache/datafusion/pull/13377#discussion_r1841773403
##########
datafusion/physical-plan/src/spill.rs:
##########
@@ -175,11 +247,103 @@ mod tests {
)?;
let file = BufReader::new(File::open(spill_file.path())?);
- let reader = arrow::ipc::reader::FileReader::try_new(file, None)?;
+ let reader = FileReader::try_new(file, None)?;
assert_eq!(reader.num_batches(), 4);
assert_eq!(reader.schema(), schema);
Ok(())
}
+
+ #[test]
+ fn test_get_record_batch_memory_size() {
+ // Create a simple record batch with two columns
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("ints", DataType::Int32, true),
+ Field::new("float64", DataType::Float64, false),
+ ]));
+
+ let int_array =
+ Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4),
Some(5)]);
+ let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
+
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::new(int_array), Arc::new(float64_array)],
+ )
+ .unwrap();
+
+ let size = get_record_batch_memory_size(&batch);
+ assert_eq!(size, 60);
+ }
+
+ #[test]
+ fn test_get_record_batch_memory_size_empty() {
+ // Test with empty record batch
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "ints",
+ DataType::Int32,
+ false,
+ )]));
+
+ let int_array: Int32Array = Int32Array::from(vec![] as Vec<i32>);
+ let batch = RecordBatch::try_new(schema,
vec![Arc::new(int_array)]).unwrap();
+
+ let size = get_record_batch_memory_size(&batch);
+ assert_eq!(size, 0, "Empty batch should have 0 memory size");
+ }
+
+ #[test]
+ fn test_get_record_batch_memory_size_shared_buffer() {
+ // Test with slices that share the same underlying buffer
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("slice1", DataType::Int32, false),
+ Field::new("slice2", DataType::Int32, false),
+ ]));
+
+ let original = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let slice1 = original.slice(0, 3);
+ let slice2 = original.slice(2, 3);
+
+ let batch =
+ RecordBatch::try_new(schema, vec![Arc::new(slice1),
Arc::new(slice2)])
+ .unwrap();
+
+ let size = get_record_batch_memory_size(&batch);
+ // The size should only count the shared buffer once
+ assert_eq!(size, 20);
+ }
+
+ #[test]
+ fn test_get_record_batch_memory_size_nested_array() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new(
+ "nested_int",
+ DataType::List(Arc::new(Field::new("item", DataType::Int32,
true))),
+ false,
+ ),
+ Field::new(
+ "nested_int2",
+ DataType::List(Arc::new(Field::new("item", DataType::Int32,
true))),
+ false,
+ ),
+ ]));
+
+ let int_list_array = ListArray::from_iter_primitive::<Int32Type, _,
_>(vec![
+ Some(vec![Some(1), Some(2), Some(3)]),
+ ]);
+
+ let int_list_array2 = ListArray::from_iter_primitive::<Int32Type, _,
_>(vec![
+ Some(vec![Some(4), Some(5), Some(6)]),
+ ]);
+
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::new(int_list_array), Arc::new(int_list_array2)],
+ )
+ .unwrap();
+
+ let size = get_record_batch_memory_size(&batch);
+ assert_eq!(size, 8320);
+ }
}
Review Comment:
Agreed.
Also I believe there are some tools can automatically do similar checks
(mutate code and make sure some test case will fail, if don't then there is
some issue with test coverage), like https://mutants.rs/
We can investigate how to integrate them into the project 😄
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
