robert3005 opened a new issue, #10692:
URL: https://github.com/apache/arrow-rs/issues/10692
### Describe the bug
Datafusion uses BatchCoalescer and compute::take functions from arrow-rs to
implement hash join. For stringview types this ends up in some cases with
arrays that have a lot of buffers holding very little or repeated data. The
issue occurs if the array on build side doesn't need gc and therefore is
blindly appended to the output of the coalescer, the output of the
BatchCoalescer will never trigger gc base on the buffer count but only ever on
buffer utilisation.
### To Reproduce
use std::sync::Arc;
use std::time::Instant;
use arrow_array::cast::AsArray;
use arrow_array::{Array, ArrayRef, RecordBatch, StringViewArray,
UInt32Array};
use arrow_buffer::{Buffer, ScalarBuffer};
use arrow_schema::{DataType, Field, Schema};
use arrow_select::coalesce::BatchCoalescer;
use arrow_select::take::take;
const ROWS: usize = 8_192;
const STAGES: usize = 12;
/// Every value must exceed 12 bytes, or the views are inlined and there are
no data buffers.
const VALUES: [&str; 5] = [
"UNITED STATES",
"UNITED KINGDOM",
"SAUDI ARABIA_",
"MOZAMBIQUE_XX",
"ARGENTINA_XXX",
];
fn view(len: u32, prefix: &[u8], buffer_index: u32, offset: u32) -> u128 {
let p = u32::from_le_bytes([prefix[0], prefix[1], prefix[2], prefix[3]]);
(len as u128) | ((p as u128) << 32) | ((buffer_index as u128) << 64) |
((offset as u128) << 96)
}
/// `dedupe = true` stores each distinct value once in a tightly sized
buffer -- what a
/// dictionary decoder naturally produces. `false` writes every row's bytes,
like a row-wise
/// decoder, which leaves the buffer looking sparse to the gc test.
fn build(dedupe: bool) -> ArrayRef {
let mut bytes: Vec<u8> = Vec::new();
let mut views = Vec::with_capacity(ROWS);
if dedupe {
let offsets: Vec<u32> = VALUES
.iter()
.map(|v| {
let at = bytes.len() as u32;
bytes.extend_from_slice(v.as_bytes());
at
})
.collect();
for i in 0..ROWS {
let v = VALUES[i % VALUES.len()];
views.push(view(v.len() as u32, v.as_bytes(), 0, offsets[i %
VALUES.len()]));
}
} else {
for i in 0..ROWS {
let v = VALUES[i % VALUES.len()];
let at = bytes.len() as u32;
bytes.extend_from_slice(v.as_bytes());
views.push(view(v.len() as u32, v.as_bytes(), 0, at));
}
}
bytes.shrink_to_fit();
let array = StringViewArray::try_new(ScalarBuffer::from(views),
vec![Buffer::from(bytes)], None)
.expect("valid view array");
Arc::new(array)
}
fn describe(label: &str, array: &dyn Array) {
let v = array.as_string_view();
let used = v.total_buffer_bytes_used();
let capacity: usize = v.data_buffers().iter().map(|b|
b.capacity()).sum();
println!(
" {label:<16} buffers={:<7} views_reference={used:<8}
buffer_capacity={capacity:<9} need_gc={}",
v.data_buffers().len(),
used != 0 && capacity > used * 2,
);
}
/// One pipeline stage: `take` from the input several times and coalesce the
results, as a hash
/// join emits and re-batches its output.
fn stage(schema: &Arc<Schema>, input: &ArrayRef, fan_in: usize) -> ArrayRef {
let mut coalescer = BatchCoalescer::new(Arc::clone(schema), ROWS);
let width = ROWS as u32 * 3 / 4; // does not divide the batch size, so
outputs span sources
let idx = UInt32Array::from_iter_values((0..width).map(|i| i %
input.len() as u32));
for _ in 0..fan_in {
let taken = take(input.as_ref(), &idx, None).unwrap();
coalescer
.push_batch(RecordBatch::try_new(Arc::clone(schema),
vec![taken]).unwrap())
.unwrap();
}
coalescer.finish_buffered_batch().unwrap();
coalescer.next_completed_batch().unwrap().column(0).clone()
}
fn run(label: &str, dedupe: bool) -> std::time::Duration {
let schema = Arc::new(Schema::new(vec![Field::new("s",
DataType::Utf8View, true)]));
let mut array = build(dedupe);
println!("\n=== {label} ===");
describe("source", array.as_ref());
for n in 1..=STAGES {
array = stage(&schema, &array, 2);
if n % 3 == 0 || n <= 2 {
describe(&format!("after stage {n}"), array.as_ref());
}
}
// What a downstream hash join pays per probe batch: `take_byte_view`
clones the whole
// buffer list on every call, and the resulting batch drops it again.
let len = array.len() as u32;
let idx = UInt32Array::from_iter_values((0..len).map(|i| (i * 7919) %
len));
let start = Instant::now();
for _ in 0..400 {
std::hint::black_box(take(array.as_ref(), &idx, None).unwrap());
}
let elapsed = start.elapsed();
println!(" 400x take({len} rows): {elapsed:?}");
elapsed
}
fn main() {
let deduped = run("deduplicated buffer (dictionary-shaped)", true);
let plain = run("per-row buffer (row-decoder-shaped)", false);
println!(
"\ndeduplicated is {:.1}x slower to take from after {STAGES} stages",
deduped.as_secs_f64() / plain.as_secs_f64()
);
}
### Expected behavior
The above code has similar performance whether buffers are deduplicated or
not
### Additional context
This behaviour can be triggered if you have a datafusion join on stringview
array that is well deduplicated
--
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]
