SubhamSinghal commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4501543072
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -8062,55 +7891,48 @@ mod tests {
Arc::new(Int32Array::from(vals)),
],
)?)?;
- assert_eq!(state.states.len(), PARTITIONS);
-
- let key_bytes: usize = state.states.keys().map(|pk|
pk.capacity()).sum();
- let scratch_bytes = state.partition_groups.capacity()
- * (size_of::<Vec<u8>>() + size_of::<Vec<u32>>())
- + state.ob_runs.capacity() * (size_of::<Vec<u8>>() +
size_of::<Vec<u32>>());
+ assert_eq!(state.core.partitions.len(), PARTITIONS);
+
+ // The store holds a copy of every key too, in the gathered `pk`
+ // column, and so does the partition-key scratch buffer, which still
+ // holds this batch's encoding. Both are taken out: either alone is
+ // about `key_bytes`, so leaving one in would cover the keys whether or
+ // not `size()` charges them.
+ let key_bytes: usize =
state.core.partitions.keys().map(Vec::len).sum();
assert!(key_bytes >= PARTITIONS * KEY_WIDTH, "key bytes {key_bytes}");
- assert!(scratch_bytes > 0, "scratch tables never allocated");
-
- // Reconstruct the total from its parts. Both terms above have to
- // appear for this to balance, so dropping either from `size()`
- // fails here rather than being absorbed by the slack in some
- // other term.
- let expected = size_of::<PartitionedTopKDenseRank>()
- + state.row_converter.size()
- + state.partition_converter.size()
- + state.scratch_rows.size()
- + state.partition_scratch_rows.size()
- + key_bytes
- + state.states.values().map(|s| s.size()).sum::<usize>()
- + state.states.capacity()
- * (size_of::<Vec<u8>>() + size_of::<DenseRankPartitionState>())
- + scratch_bytes
- + state.store.size();
- assert_eq!(state.size(), expected);
+ let outside_copies = state.size()
+ - state.core.store.size()
+ - state.core.partition_scratch_rows.size();
+ assert!(
Review Comment:
Addressed in
https://github.com/apache/datafusion/commit/5f96ec422cf781200fb32420e09c6aaa6ed4fed8
--
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]
