kumarUjjawal commented on code in PR #23565:
URL: https://github.com/apache/datafusion/pull/23565#discussion_r4118969964
##########
datafusion/physical-plan/src/spill/mod.rs:
##########
@@ -800,29 +804,133 @@ pub(crate) fn gc_view_arrays(batch: &RecordBatch) ->
Result<RecordBatch> {
}
}
+/// Compacts the data buffers of a view array, or returns `None` when the
+/// array is too small for compaction to be worth it.
+///
+/// Arrow's `gc()` copies the bytes of every view separately, so repeated
+/// values each get their own copy. For a dictionary-encoded Parquet column
+/// this inflates the spilled data by the average repeat count: 1M rows of
+/// 1000 distinct 64 byte values spill as 64 MB
+/// (<https://github.com/apache/datafusion/issues/23564>).
+/// [`gc_dedup_view_array`] copies each distinct value once instead, and falls
+/// back to `gc()` when the values turn out to be distinct.
+fn gc_view_array<T: ByteViewType>(
+ array: &GenericByteViewArray<T>,
+) -> Option<GenericByteViewArray<T>> {
+ if !should_gc_view_array(array) {
+ return None;
+ }
+ Some(gc_dedup_view_array(array).unwrap_or_else(|| array.gc()))
+}
+
+/// Number of non-inline values [`gc_dedup_view_array`] deduplicates before
+/// it checks whether the array has enough repeats to be worth it.
+const DEDUP_SAMPLE_VALUES: usize = 256;
+
+/// Like `gc()`, but copies each distinct value once, and zeroes null views.
+///
+/// Returns `None` if the first [`DEDUP_SAMPLE_VALUES`] non-inline values
+/// hold almost no repeats, since hashing every value then costs more than
+/// deduplication saves.
+fn gc_dedup_view_array<T: ByteViewType>(
+ array: &GenericByteViewArray<T>,
+) -> Option<GenericByteViewArray<T>> {
+ let buffers = array.data_buffers();
+ let bytes_of = |view: &ByteView| {
+ let start = view.offset as usize;
+ &buffers[view.buffer_index as usize][start..start + view.length as
usize]
+ };
+ let hasher = DefaultHashBuilder::default();
+
+ // (input view, output view) of the first occurrence of each distinct value
+ let mut distinct: HashTable<(u128, u128)> = HashTable::new();
+ let mut completed: Vec<Buffer> = vec![];
+ let mut data: Vec<u8> = vec![];
+ let mut non_inline = 0;
+ let mut views = Vec::with_capacity(array.len());
+
+ for (i, &raw) in array.views().iter().enumerate() {
+ if array.is_null(i) {
+ views.push(0);
+ continue;
+ }
+ if (raw as u32) <= MAX_INLINE_VIEW_LEN {
+ views.push(raw);
+ continue;
+ }
+
+ let view = ByteView::from(raw);
+ let bytes = bytes_of(&view);
+ let hash = hasher.hash_one(bytes);
+ let found = distinct.find(hash, |(first, _)| {
+ *first == raw || bytes_of(&ByteView::from(*first)) == bytes
+ });
+ let new_view = match found {
+ Some((_, new_view)) => *new_view,
+ None => {
+ // A view offset is a `u32`, and a buffer may not exceed
`i32::MAX`
+ if data.len() + bytes.len() > i32::MAX as usize {
+ completed.push(Buffer::from_vec(std::mem::take(&mut
data)));
+ }
+ let new_view = ByteView {
+ buffer_index: completed.len() as u32,
+ offset: data.len() as u32,
+ ..view
+ }
+ .as_u128();
+ data.extend_from_slice(bytes);
+ distinct.insert_unique(hash, (raw, new_view), |(first, _)| {
+ hasher.hash_one(bytes_of(&ByteView::from(*first)))
+ });
+ new_view
+ }
+ };
+ views.push(new_view);
+
+ non_inline += 1;
+ if non_inline == DEDUP_SAMPLE_VALUES
Review Comment:
When labels cycle through 1,000 values, the first 256 are distinct, so we
fall back to gc() despite many repeats later. Could we sample across the batch
and add a test for this ordering?
--
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]
