Vortex can't combine two files it has previous written, when one of the blocks in the rewritten file contains a sliced sparse array of UUIDs alongside a dense one. #10417

use std::process::ExitCode;
use std::sync::Arc;

use arrow::array::{FixedSizeBinaryArray, RecordBatch};
use arrow::datatypes::{DataType, Field, Schema};
use futures::StreamExt;
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use vortex::VortexSessionDefault;
use vortex::array::arrays::ChunkedArray;
use vortex::array::{ArrayRef, IntoArray};
use vortex::arrow::ArrowSessionExt;
use vortex::buffer::ByteBuffer;
use vortex::error::VortexResult;
use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt};
use vortex::io::session::RuntimeSessionExt;
use vortex::session::VortexSession;

const DENSE_ROWS: usize = 10_000;
const DENSE_SEED: u64 = 1;
const SPARSE_ROWS: usize = 8_000;
const SPARSE_UUID_EVERY: usize = 20;
const SPARSE_SEED: u64 = 2;

fn uuid_batch(rows: usize, seed: u64, has_uuid: impl Fn(usize) -> bool) -> RecordBatch {
    let mut rng = SmallRng::seed_from_u64(seed);
    let uuids = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
        (0..rows)
            .map(|row| has_uuid(row).then(|| rng.random::<[u8; 16]>()))
            .collect::<Vec<_>>()
            .into_iter(),
        16,
    )
    .unwrap();
    let schema = Schema::new(vec![
        Field::new("id", DataType::FixedSizeBinary(16), true)
            .with_extension_type(arrow_schema::extension::Uuid),
    ]);
    RecordBatch::try_new(Arc::new(schema), vec![Arc::new(uuids)]).unwrap()
}

async fn write(session: &VortexSession, chunks: Vec<ArrayRef>) -> VortexResult<Vec<u8>> {
    let dtype = chunks[0].dtype().clone();
    let array = ChunkedArray::try_new(chunks, dtype)?.into_array();
    let mut bytes = Vec::new();
    session
        .write_options()
        .write(&mut bytes, array.to_array_stream())
        .await?;
    Ok(bytes)
}

async fn read(session: &VortexSession, bytes: Vec<u8>) -> VortexResult<Vec<ArrayRef>> {
    let file = session
        .open_options()
        .open_buffer(ByteBuffer::from(bytes))?;
    let mut stream = file.scan()?.into_array_stream()?;
    let mut chunks = Vec::new();
    while let Some(chunk) = stream.next().await {
        chunks.push(chunk?);
    }
    Ok(chunks)
}

async fn write_vortex_file(session: &VortexSession, batch: RecordBatch) -> VortexResult<Vec<u8>> {
    let schema = batch.schema();
    let array = session.arrow().from_arrow_record_batch(batch, &schema)?;
    write(session, vec![array]).await
}

#[tokio::main]
async fn main() -> ExitCode {
    let session = VortexSession::default().with_tokio();

    let dense = uuid_batch(DENSE_ROWS, DENSE_SEED, |_| true);
    let sparse = uuid_batch(SPARSE_ROWS, SPARSE_SEED, |row| row % SPARSE_UUID_EVERY == 0);

    let mut chunks = Vec::new();
    for (name, batch) in [("dense", dense), ("sparse", sparse)] {
        let file = write_vortex_file(&session, batch).await.unwrap();
        println!("wrote {name} file ({} bytes)", file.len());
        chunks.extend(read(&session, file).await.unwrap());
    }

    println!(
        "rewriting {} chunks read back from both files into one file",
        chunks.len()
    );
    match write(&session, chunks).await {
        Ok(file) => {
            println!(
                "rewrite succeeded ({} bytes): bug not reproduced",
                file.len()
            );
            ExitCode::SUCCESS
        }
        Err(err) => {
            eprintln!("rewrite failed: {err}");
            ExitCode::FAILURE
        }
    }
}