kosiew commented on code in PR #19285:
URL: https://github.com/apache/datafusion/pull/19285#discussion_r2659305885
##########
datafusion/physical-plan/src/aggregates/topk/heap.rs:
##########
@@ -161,6 +225,166 @@ where
}
}
+/// An implementation of `ArrowHeap` that deals with string values.
+///
+/// Supports all three UTF-8 string types: `Utf8`, `LargeUtf8`, and `Utf8View`.
+/// String values are compared lexicographically. Null values are not
explicitly handled
+/// and should not appear in the input; the aggregation layer ensures nulls
are managed
+/// appropriately before calling this heap.
+///
+/// Uses string interning to avoid repeated allocations for duplicate strings
within a batch.
+/// The `string_cache` maps string hashes to `Arc<str>` values, amortizing
allocation costs
+/// when the same string appears multiple times (common in trace IDs, user
IDs, etc.).
+pub struct StringHeap {
+ batch: ArrayRef,
+ heap: TopKHeap<Arc<str>>,
+ desc: bool,
+ data_type: DataType,
+ /// Cache of interned strings for the current batch, mapping hash to
`Arc<str>`.
+ /// Cleared on each `set_batch` call to prevent memory leaks from old
batches.
+ string_cache: HashMap<u64, Arc<str>>,
+}
+
+impl StringHeap {
+ pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self {
+ let batch: ArrayRef = Arc::new(StringArray::from(Vec::<&str>::new()));
+ Self {
+ batch,
+ heap: TopKHeap::new(limit, desc),
+ desc,
+ data_type,
+ string_cache: HashMap::new(),
+ }
+ }
+
+ /// Extracts a string value from the current batch at the given row index.
+ ///
+ /// Panics if the row index is out of bounds or if the data type is not
one of
+ /// the supported UTF-8 string types.
+ ///
+ /// Note: Null values should not appear in the input; the aggregation layer
+ /// ensures nulls are filtered before reaching this code.
+ fn value(&self, row_idx: usize) -> &str {
+ extract_string_value(&self.batch, &self.data_type, row_idx)
+ }
+
+ /// Interns a string value, returning a cached `Arc<str>` if available.
+ ///
+ /// This method implements string interning to reduce allocations for
duplicate strings.
+ /// It computes a hash of the input string and checks the cache. If found,
returns the
+ /// cached `Arc<str>`. Otherwise, allocates a new `Arc<str>`, caches it,
and returns it.
+ ///
+ /// This is particularly effective for workloads with repeated string
values like trace IDs.
+ fn intern_string(&mut self, s: &str) -> Arc<str> {
+ // Compute hash of the string
+ let mut hasher = std::collections::hash_map::DefaultHasher::new();
+ s.hash(&mut hasher);
+ let hash = hasher.finish();
+
+ // Check cache and return if found
+ if let Some(cached) = self.string_cache.get(&hash) {
Review Comment:
good idea!
I'll implement it.
--
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]
