InProgressByteViewArray::append_views_and_update_buffer_index/GenericByteViewBuilder::append_array should deduplicate buffers #10692

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()
    );
}