kumarUjjawal commented on code in PR #24476:
URL: https://github.com/apache/datafusion/pull/24476#discussion_r3832250229
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs:
##########
@@ -546,6 +587,52 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
}
}
+/// Compact view buffers after `take` and subset nested dictionary values the
same way.
+fn gc_taken_values(array: ArrayRef) -> ArrayRef {
Review Comment:
`spill/mod.rs:423` already has a `gc_array` for this it recurses into nested
children and gates `gc()` behind a size threshold, which matters since
`gc()`
never no-ops. As written this copies the string data on every `take_n` even
when nothing is reclaimable.
Worth reusing?
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs:
##########
@@ -546,6 +587,52 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
}
}
+/// Compact view buffers after `take` and subset nested dictionary values the
same way.
+fn gc_taken_values(array: ArrayRef) -> ArrayRef {
+ match array.data_type() {
+ DataType::Utf8View => Arc::new(array.as_string_view().gc()) as
ArrayRef,
+ DataType::BinaryView => Arc::new(array.as_binary_view().gc()) as
ArrayRef,
+ DataType::Dictionary(_, _) => {
+ let arr = array.as_ref();
+ downcast_dictionary_array!(
+ arr => gc_dictionary(arr),
+ _ => array
+ )
+ }
+ _ => array,
+ }
+}
+
+/// Keep only values referenced by keys then gc those values
+fn gc_dictionary<K: ArrowDictionaryKeyType>(array: &DictionaryArray<K>) ->
ArrayRef {
Review Comment:
Is this reachable? It needs `Dictionary(K1, Dictionary(K2, V))` and I
couldn't
find anything that produces one. If there is a path, it'd be worth a test
since
nothing exercises this today
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs:
##########
@@ -515,28 +528,56 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
new_to_old.push(old);
}
- self.value_dedup = HashTable::new();
- self.value_dedup_size = 0;
- self.null_inner_slot = None;
-
- self.hash_values(&all_inner_values);
+ // Subset leftover values and gc view/nested buffers so dropped
strings are freed
+ let leftover_non_null_len = new_to_old.len() - null_old_slot.is_some()
as usize;
+ let leftover_values = if leftover_non_null_len == 0 {
+ None
+ } else {
+ let leftover_indices = Int64Array::from_iter(
+ new_to_old[..leftover_non_null_len]
+ .iter()
+ .map(|&i| i as i64),
+ );
+ let taken = take(&*all_inner_values, &leftover_indices, None)
+ .expect("take leftover values");
+ Some(gc_taken_values(taken))
+ };
+ // Drop fully-emitted slots from value_dedup and remap leftover slots
in place
+ let old_hashes = std::mem::take(&mut self.slot_hash);
+ for &old in &emit_new_to_old {
+ if old_to_new[old] != usize::MAX {
+ continue;
+ }
+ let hash = old_hashes[old];
+ if let Ok(entry) = self.value_dedup.find_entry(hash,
|&(entry_hash, slot)| {
+ entry_hash == hash && slot == old
+ }) {
+ entry.remove();
+ }
+ }
+ self.null_inner_slot = None;
+ self.slot_hash.resize(leftover_non_null_len, 0);
Review Comment:
Am I right that this can alias? `*slot = new_slot` runs while later
iterations still match on `slot == old_slot`, so two entries sharing a hash can
end up swapped, silently grouping rows under the wrong key. Rare, but
`value_dedup` does support collisions.
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs:
##########
@@ -515,28 +528,56 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
new_to_old.push(old);
}
- self.value_dedup = HashTable::new();
- self.value_dedup_size = 0;
- self.null_inner_slot = None;
-
- self.hash_values(&all_inner_values);
+ // Subset leftover values and gc view/nested buffers so dropped
strings are freed
+ let leftover_non_null_len = new_to_old.len() - null_old_slot.is_some()
as usize;
Review Comment:
The leftover `take` + `gc` looks redundant here, `append_val` already
copies into the builder's own storage (`bytes_view.rs:221`), so the rebuilt
`inner` was never holding the original allocation. That makes this an extra
full copy per `take_n`.
--
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]
