getChan commented on code in PR #13906:
URL: https://github.com/apache/datafusion/pull/13906#discussion_r1897906605
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -643,3 +643,44 @@ impl RecordBatchStore {
+ self.batches_size
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::Int32Array;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::record_batch::RecordBatch;
+ use arrow_array::Float64Array;
+
+ /// This test ensures the size calculation is correct for RecordBatches
with multiple columns.
+ #[test]
+ fn test_record_batch_store_size() {
+ // given
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("ints", DataType::Int32, true),
+ Field::new("float64", DataType::Float64, false),
+ ]));
+ let mut record_batch_store =
RecordBatchStore::new(Arc::clone(&schema));
+ let int_array =
+ Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4),
Some(5)]); // 5 * 4 = 20
+ let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
// 5 * 8 = 40
+
+ let record_batch_entry = RecordBatchEntry {
+ id: 0,
+ batch: RecordBatch::try_new(
+ schema,
+ vec![Arc::new(int_array), Arc::new(float64_array)],
+ )
+ .unwrap(),
+ uses: 1,
+ };
+
+ // when insert record batch entry
+ record_batch_store.insert(record_batch_entry);
+ assert_eq!(record_batch_store.batches_size, 60);
Review Comment:
Before, it was 252.
--
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
