// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use super::*;
use crate::writer::{MosaicWriter, OutputFile, WriterOptions};
use arrow_array::*;
use arrow_schema::{DataType, Field, Schema, TimeUnit};
use std::sync::Arc;

fn columns_to_arrow_schema(columns: &[(String, DataType, bool)]) -> Schema {
    Schema::new(
        columns
            .iter()
            .map(|(name, dt, nullable)| Field::new(name, dt.clone(), *nullable))
            .collect::<Vec<_>>(),
    )
}

fn legacy_timestamp_nanos_fields() -> arrow_schema::Fields {
    vec![
        Field::new("millis", DataType::Int64, false),
        Field::new("nanos_of_milli", DataType::Int32, false),
    ]
    .into()
}

fn legacy_timestamp_nanos_type() -> DataType {
    DataType::Struct(legacy_timestamp_nanos_fields())
}

struct ByteArrayInputFile {
    data: Vec<u8>,
}

impl ByteArrayInputFile {
    fn new(data: Vec<u8>) -> Self {
        Self { data }
    }
}

impl InputFile for ByteArrayInputFile {
    fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
        let start = offset as usize;
        let end = start + buf.len();
        if end > self.data.len() {
            return Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "read past end",
            ));
        }
        buf.copy_from_slice(&self.data[start..end]);
        Ok(())
    }
}

struct MemOutputFile {
    buf: Vec<u8>,
}

impl MemOutputFile {
    fn new() -> Self {
        Self { buf: Vec::new() }
    }
}

impl OutputFile for MemOutputFile {
    fn write(&mut self, data: &[u8]) -> io::Result<()> {
        self.buf.extend_from_slice(data);
        Ok(())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
    fn pos(&self) -> u64 {
        self.buf.len() as u64
    }
}

fn values_to_batch(rows: &[Vec<Value>], columns: &[(String, DataType, bool)]) -> RecordBatch {
    let num_cols = columns.len();
    let mut fields = Vec::with_capacity(num_cols);
    let mut arrays: Vec<Arc<dyn Array>> = Vec::with_capacity(num_cols);

    for (c, (name, dt, nullable)) in columns.iter().enumerate() {
        fields.push(Field::new(name, dt.clone(), *nullable));
        arrays.push(build_array_from_values(rows, c, dt));
    }

    RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap()
}

fn build_array_from_values(rows: &[Vec<Value>], col: usize, dt: &DataType) -> Arc<dyn Array> {
    match dt {
        DataType::Boolean => {
            let vals: Vec<Option<bool>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Boolean(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(BooleanArray::from(vals))
        }
        DataType::Int8 => {
            let vals: Vec<Option<i8>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::TinyInt(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Int8Array::from(vals))
        }
        DataType::Int16 => {
            let vals: Vec<Option<i16>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::SmallInt(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Int16Array::from(vals))
        }
        DataType::Int32 => {
            let vals: Vec<Option<i32>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Integer(v) | Value::Date(v) | Value::Time(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Int32Array::from(vals))
        }
        DataType::Int64 => {
            let vals: Vec<Option<i64>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::BigInt(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Int64Array::from(vals))
        }
        DataType::Float32 => {
            let vals: Vec<Option<f32>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Float(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Float32Array::from(vals))
        }
        DataType::Float64 => {
            let vals: Vec<Option<f64>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Double(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Float64Array::from(vals))
        }
        DataType::Date32 => {
            let vals: Vec<Option<i32>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Date(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Date32Array::from(vals))
        }
        DataType::Time32(_) => {
            let vals: Vec<Option<i32>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Time(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(Time32MillisecondArray::from(vals))
        }
        DataType::Utf8 => {
            let vals: Vec<Option<String>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::String(b) => Some(String::from_utf8(b.clone()).unwrap()),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(StringArray::from(vals))
        }
        DataType::Binary => {
            let vals: Vec<Option<Vec<u8>>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::Bytes(b) => Some(b.clone()),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            let refs: Vec<Option<&[u8]>> = vals.iter().map(|v| v.as_deref()).collect();
            Arc::new(BinaryArray::from(refs))
        }
        DataType::Decimal128(p, s) => {
            let vals: Vec<Option<i128>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::DecimalCompact(v) => Some(*v as i128),
                    Value::DecimalLarge(bytes) => Some(biginteger_bytes_to_i128(bytes)),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            Arc::new(
                Decimal128Array::from(vals)
                    .with_precision_and_scale(*p, *s)
                    .unwrap(),
            )
        }
        DataType::Timestamp(TimeUnit::Millisecond, tz) => {
            let vals: Vec<Option<i64>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::TimestampMillis(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            let arr = TimestampMillisecondArray::from(vals);
            Arc::new(if let Some(tz) = tz {
                arr.with_timezone(tz.clone())
            } else {
                arr
            })
        }
        DataType::Timestamp(TimeUnit::Microsecond, tz) => {
            let vals: Vec<Option<i64>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::TimestampMicros(v) => Some(*v),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            let arr = TimestampMicrosecondArray::from(vals);
            Arc::new(if let Some(tz) = tz {
                arr.with_timezone(tz.clone())
            } else {
                arr
            })
        }
        DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
            let vals: Vec<Option<i64>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::TimestampNanos {
                        millis,
                        nanos_of_milli,
                    } => Some(crate::types::millis_nanos_to_ns(*millis, *nanos_of_milli).unwrap()),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            let arr = TimestampNanosecondArray::from(vals);
            Arc::new(if let Some(tz) = tz {
                arr.with_timezone(tz.clone())
            } else {
                arr
            })
        }
        DataType::Struct(struct_fields)
            if crate::types::is_timestamp_nanos_struct(struct_fields) =>
        {
            let millis_vals: Vec<Option<i64>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::TimestampNanos { millis, .. } => Some(*millis),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            let nanos_vals: Vec<Option<i32>> = rows
                .iter()
                .map(|row| match &row[col] {
                    Value::TimestampNanos { nanos_of_milli, .. } => Some(*nanos_of_milli),
                    Value::Null => None,
                    _ => panic!("type mismatch"),
                })
                .collect();
            let nulls: Option<arrow_buffer::NullBuffer> = {
                let any_null = rows.iter().any(|row| row[col].is_null());
                if any_null {
                    let bools: Vec<bool> = rows.iter().map(|row| !row[col].is_null()).collect();
                    Some(arrow_buffer::NullBuffer::from(bools))
                } else {
                    None
                }
            };
            Arc::new(StructArray::new(
                struct_fields.clone(),
                vec![
                    Arc::new(Int64Array::from(
                        millis_vals
                            .into_iter()
                            .map(|v| v.unwrap_or(0))
                            .collect::<Vec<_>>(),
                    )),
                    Arc::new(Int32Array::from(
                        nanos_vals
                            .into_iter()
                            .map(|v| v.unwrap_or(0))
                            .collect::<Vec<_>>(),
                    )),
                ],
                nulls,
            ))
        }
        _ => panic!("unsupported type in test helper: {:?}", dt),
    }
}

fn biginteger_bytes_to_i128(bytes: &[u8]) -> i128 {
    let negative = bytes[0] & 0x80 != 0;
    let pad = if negative { 0xFF } else { 0x00 };
    let mut buf = [pad; 16];
    let start = 16 - bytes.len();
    buf[start..].copy_from_slice(bytes);
    i128::from_be_bytes(buf)
}

fn write_values(
    writer: &mut MosaicWriter<MemOutputFile>,
    columns: &[(String, DataType, bool)],
    rows: &[Vec<Value>],
) {
    if rows.is_empty() {
        return;
    }
    for row in rows {
        let batch = values_to_batch(std::slice::from_ref(row), columns);
        writer.write_batch(&batch).unwrap();
    }
}

fn batch_col_bool<'a>(batch: &'a RecordBatch, name: &str) -> &'a BooleanArray {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<BooleanArray>()
        .unwrap()
}

fn batch_col_i8<'a>(batch: &'a RecordBatch, name: &str) -> &'a Int8Array {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<Int8Array>()
        .unwrap()
}

fn batch_col_i16<'a>(batch: &'a RecordBatch, name: &str) -> &'a Int16Array {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<Int16Array>()
        .unwrap()
}

fn batch_col_i32<'a>(batch: &'a RecordBatch, name: &str) -> &'a Int32Array {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<Int32Array>()
        .unwrap()
}

fn batch_col_i64<'a>(batch: &'a RecordBatch, name: &str) -> &'a Int64Array {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<Int64Array>()
        .unwrap()
}

fn batch_col_f32<'a>(batch: &'a RecordBatch, name: &str) -> &'a Float32Array {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<Float32Array>()
        .unwrap()
}

fn batch_col_f64<'a>(batch: &'a RecordBatch, name: &str) -> &'a Float64Array {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap()
}

fn batch_col_string<'a>(batch: &'a RecordBatch, name: &str) -> &'a StringArray {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<StringArray>()
        .unwrap()
}

fn batch_col_binary<'a>(batch: &'a RecordBatch, name: &str) -> &'a BinaryArray {
    let idx = batch.schema().index_of(name).unwrap();
    batch
        .column(idx)
        .as_any()
        .downcast_ref::<BinaryArray>()
        .unwrap()
}

#[test]
fn test_roundtrip_basic() {
    let columns = vec![
        ("age".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows_to_write: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Integer(20 + (i % 50)),
                Value::String(format!("user_{}", i).into_bytes()),
                Value::Double(i as f64 * 1.5),
            ]
        })
        .collect();

    write_values(&mut writer, &columns, &rows_to_write);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    assert_eq!(reader.schema().columns.len(), 3);
    assert_eq!(reader.num_row_groups(), 1);

    let mut rg = reader.row_group_reader(0).unwrap();
    assert_eq!(rg.num_rows(), 100);

    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);
    assert_eq!(batch.num_columns(), 3);

    let ages = batch_col_i32(&batch, "age");
    let names = batch_col_string(&batch, "name");
    let scores = batch_col_f64(&batch, "score");

    for i in 0..100usize {
        assert_eq!(ages.value(i), 20 + (i as i32 % 50));
        assert_eq!(names.value(i), format!("user_{}", i));
        assert!((scores.value(i) - i as f64 * 1.5).abs() < 1e-10);
    }
}

#[test]
fn test_roundtrip_with_nulls() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("value".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows = vec![
        vec![
            Value::Integer(1),
            Value::String(b"hello".to_vec()),
            Value::Double(1.0),
        ],
        vec![Value::Integer(2), Value::Null, Value::Double(2.0)],
        vec![
            Value::Integer(3),
            Value::String(b"world".to_vec()),
            Value::Null,
        ],
        vec![Value::Integer(4), Value::Null, Value::Null],
    ];
    write_values(&mut writer, &columns, &rows);

    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 4);

    let ids = batch_col_i32(&batch, "id");
    let names = batch_col_string(&batch, "name");
    let vals = batch_col_f64(&batch, "value");

    assert_eq!(ids.value(0), 1);
    assert_eq!(ids.value(1), 2);
    assert_eq!(ids.value(2), 3);
    assert_eq!(ids.value(3), 4);

    assert!(!names.is_null(0));
    assert_eq!(names.value(0), "hello");
    assert!(names.is_null(1));
    assert!(!names.is_null(2));
    assert_eq!(names.value(2), "world");
    assert!(names.is_null(3));

    assert!(!vals.is_null(0));
    assert!((vals.value(0) - 1.0).abs() < 1e-10);
    assert!(!vals.is_null(1));
    assert!((vals.value(1) - 2.0).abs() < 1e-10);
    assert!(vals.is_null(2));
    assert!(vals.is_null(3));
}

#[test]
fn test_roundtrip_with_zstd() {
    let columns = vec![
        ("a".to_string(), DataType::Int64, true),
        ("b".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            zstd_level: 3,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..1000)
        .map(|i| vec![Value::BigInt(i), Value::BigInt(i * 2)])
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 1000);

    let col_a = batch_col_i64(&batch, "a");
    let col_b = batch_col_i64(&batch, "b");

    for i in 0..1000usize {
        assert_eq!(col_a.value(i), i as i64);
        assert_eq!(col_b.value(i), i as i64 * 2);
    }
}

#[test]
fn test_roundtrip_all_types() {
    let columns = vec![
        ("f_boolean".to_string(), DataType::Boolean, true),
        ("f_tinyint".to_string(), DataType::Int8, true),
        ("f_smallint".to_string(), DataType::Int16, true),
        ("f_int".to_string(), DataType::Int32, true),
        ("f_bigint".to_string(), DataType::Int64, true),
        ("f_float".to_string(), DataType::Float32, true),
        ("f_double".to_string(), DataType::Float64, true),
        ("f_string".to_string(), DataType::Utf8, true),
        ("f_bytes".to_string(), DataType::Binary, true),
        (
            "f_decimal_compact".to_string(),
            DataType::Decimal128(10, 2),
            true,
        ),
        ("f_date".to_string(), DataType::Date32, true),
        (
            "f_timestamp".to_string(),
            DataType::Timestamp(TimeUnit::Millisecond, None),
            true,
        ),
        (
            "f_timestamp_high".to_string(),
            DataType::Struct(
                vec![
                    arrow_schema::Field::new("millis", DataType::Int64, false),
                    arrow_schema::Field::new("nanos_of_milli", DataType::Int32, false),
                ]
                .into(),
            ),
            true,
        ),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            num_buckets: DEFAULT_NUM_BUCKETS,
            ..Default::default()
        },
    )
    .unwrap();

    let rows = vec![vec![
        Value::Boolean(true),
        Value::TinyInt(42),
        Value::SmallInt(1234),
        Value::Integer(999999),
        Value::BigInt(123456789012345),
        Value::Float(1.25),
        Value::Double(9.876543210),
        Value::String("hello world".as_bytes().to_vec()),
        Value::Bytes(vec![1, 2, 3, 4, 5]),
        Value::DecimalCompact(1234567),
        Value::Date(19000),
        Value::TimestampMillis(1700000000000),
        Value::TimestampNanos {
            millis: 1700000000000,
            nanos_of_milli: 123456,
        },
    ]];
    write_values(&mut writer, &columns, &rows);

    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 1);

    assert!(batch_col_bool(&batch, "f_boolean").value(0));
    assert_eq!(batch_col_i8(&batch, "f_tinyint").value(0), 42);
    assert_eq!(batch_col_i16(&batch, "f_smallint").value(0), 1234);
    assert_eq!(batch_col_i32(&batch, "f_int").value(0), 999999);
    assert_eq!(
        batch_col_i64(&batch, "f_bigint").value(0),
        123456789012345i64
    );
    assert!((batch_col_f32(&batch, "f_float").value(0) - 1.25).abs() < 0.001);
    assert!((batch_col_f64(&batch, "f_double").value(0) - 9.876543210).abs() < 1e-9);
    assert_eq!(batch_col_string(&batch, "f_string").value(0), "hello world");
    assert_eq!(
        batch_col_binary(&batch, "f_bytes").value(0),
        &[1u8, 2, 3, 4, 5]
    );

    let col = |name: &str| -> usize { batch.schema().index_of(name).unwrap() };

    let c_dec = batch
        .column(col("f_decimal_compact"))
        .as_any()
        .downcast_ref::<Decimal128Array>()
        .unwrap();
    assert_eq!(c_dec.value(0), 1234567i128);

    let c_date = batch
        .column(col("f_date"))
        .as_any()
        .downcast_ref::<Date32Array>()
        .unwrap();
    assert_eq!(c_date.value(0), 19000);

    let c_ts = batch
        .column(col("f_timestamp"))
        .as_any()
        .downcast_ref::<TimestampMillisecondArray>()
        .unwrap();
    assert_eq!(c_ts.value(0), 1700000000000i64);

    let c_ts_high = batch
        .column(col("f_timestamp_high"))
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .unwrap();
    assert_eq!(
        c_ts_high.value(0),
        crate::types::millis_nanos_to_ns(1700000000000, 123456).unwrap()
    );
}

fn write_and_read(
    columns: Vec<(String, DataType, bool)>,
    rows: &[Vec<Value>],
) -> (MosaicReader<ByteArrayInputFile>, Vec<u8>) {
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            num_buckets: DEFAULT_NUM_BUCKETS,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data.clone()), len).unwrap();
    (reader, data)
}

/// Mirrors Paimon's testSchemaEvolutionTypeWidening.
/// Writes with narrow types (INT, FLOAT, TINYINT) and verifies the reader
/// produces the exact values that a higher layer can widen to BIGINT/DOUBLE/INT.
#[test]
fn test_evolution_type_widening_readback() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float32, true),
        ("amount".to_string(), DataType::Int8, true),
    ];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("name_{}", i).into_bytes()),
                Value::Float(i as f32 * 0.5),
                Value::TinyInt((i % 127) as i8),
            ]
        })
        .collect();

    let (reader, _) = write_and_read(columns, &rows);

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);

    let ids = batch_col_i32(&batch, "id");
    let names = batch_col_string(&batch, "name");
    let scores = batch_col_f32(&batch, "score");
    let amounts = batch_col_i8(&batch, "amount");

    for i in 0..100usize {
        assert_eq!(ids.value(i) as i64, i as i64);
        assert_eq!(names.value(i), format!("name_{}", i));
        assert!((scores.value(i) as f64 - i as f64 * 0.5).abs() < 1e-5);
        assert_eq!(amounts.value(i) as i32, i as i32 % 127);
    }
}

/// Mirrors Paimon's testSchemaEvolutionDroppedColumn.
/// Writes 3 columns, then reads back and accesses only column "c" by name,
/// ignoring "a" and "b". This is the pattern used by the adapter when the read
/// schema has fewer columns than the write schema.
#[test]
fn test_evolution_dropped_column_readback() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..50)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("v{}", i).into_bytes()),
                Value::Double(i as f64 * 1.1),
            ]
        })
        .collect();

    let (reader, _) = write_and_read(columns, &rows);

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 50);

    let cc = batch_col_f64(&batch, "c");
    for i in 0..50usize {
        assert!((cc.value(i) - i as f64 * 1.1).abs() < 1e-10);
    }
}

/// Mirrors Paimon's testSchemaEvolutionConstAndDictCast.
/// Writes data that triggers const encoding (all-same value) and dict encoding
/// (few distinct values), then reads back and verifies all values are correct.
/// The adapter layer would then widen these to BIGINT/INT/DOUBLE.
#[test]
fn test_evolution_const_and_dict_encoding_readback() {
    let columns = vec![
        ("const_col".to_string(), DataType::Int32, true),
        ("dict_col".to_string(), DataType::Int16, true),
        ("plain_col".to_string(), DataType::Float32, true),
    ];
    let rows: Vec<Vec<Value>> = (0..200)
        .map(|i| {
            vec![
                Value::Integer(42),
                Value::SmallInt((i % 5) as i16),
                Value::Float(i as f32),
            ]
        })
        .collect();

    let (reader, _) = write_and_read(columns, &rows);

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 200);

    let const_col = batch_col_i32(&batch, "const_col");
    let dict_col = batch_col_i16(&batch, "dict_col");
    let plain_col = batch_col_f32(&batch, "plain_col");

    for i in 0..200usize {
        assert_eq!(const_col.value(i), 42);
        assert_eq!(dict_col.value(i), (i % 5) as i16);
        assert!((plain_col.value(i) - i as f32).abs() < 1e-5);
    }
}

#[test]
fn test_stats_roundtrip() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["id".to_string(), "score".to_string()],
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Integer(i * 2),
                Value::String(format!("row_{}", i).into_bytes()),
                Value::Double(i as f64 * 0.5),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let stats = reader.row_group_stats(0).unwrap();
    assert_eq!(stats.len(), 2);

    // Column 0 (id): min=0, max=198
    assert_eq!(stats[0].column_index, 0);
    assert_eq!(stats[0].null_count, 0);
    assert!(matches!(&stats[0].min, Some(Value::Integer(0))));
    assert!(matches!(&stats[0].max, Some(Value::Integer(198))));

    // Column 2 (score): min=0.0, max=49.5
    assert_eq!(stats[1].column_index, 2);
    assert_eq!(stats[1].null_count, 0);
    match &stats[1].min {
        Some(Value::Double(v)) => assert!(v.abs() < 1e-10),
        other => panic!("expected Double(0.0), got {:?}", other),
    }
    match &stats[1].max {
        Some(Value::Double(v)) => assert!((*v - 49.5).abs() < 1e-10),
        other => panic!("expected Double(49.5), got {:?}", other),
    }
}

#[test]
fn test_stats_with_nulls() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["a".to_string(), "b".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows = vec![
        vec![Value::Integer(10), Value::Null],
        vec![Value::Null, Value::Null],
        vec![Value::Integer(5), Value::BigInt(100)],
        vec![Value::Integer(20), Value::BigInt(50)],
    ];
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let stats = reader.row_group_stats(0).unwrap();
    assert_eq!(stats.len(), 2);

    assert_eq!(stats[0].column_index, 0);
    assert_eq!(stats[0].null_count, 1);
    assert!(matches!(&stats[0].min, Some(Value::Integer(5))));
    assert!(matches!(&stats[0].max, Some(Value::Integer(20))));

    assert_eq!(stats[1].column_index, 1);
    assert_eq!(stats[1].null_count, 2);
    assert!(matches!(&stats[1].min, Some(Value::BigInt(50))));
    assert!(matches!(&stats[1].max, Some(Value::BigInt(100))));
}

#[test]
fn test_stats_all_null_column() {
    let columns = vec![("x".to_string(), DataType::Int32, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["x".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..10).map(|_| vec![Value::Null]).collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let stats = reader.row_group_stats(0).unwrap();
    assert_eq!(stats.len(), 1);
    assert_eq!(stats[0].null_count, 10);
    assert!(stats[0].min.is_none());
    assert!(stats[0].max.is_none());
}

#[test]
fn test_no_stats_minimal_overhead() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Int32, true),
    ];

    let rows: Vec<Vec<Value>> = (0..10)
        .map(|i| vec![Value::Integer(i), Value::Integer(i * 2)])
        .collect();

    // Write without stats
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec![],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let no_stats_size = writer.output().buf.len();

    // Write with stats on column "a"
    let out2 = MemOutputFile::new();
    let mut writer2 = MosaicWriter::new(
        out2,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["a".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer2, &columns, &rows);
    writer2.close().unwrap();
    let with_stats_size = writer2.output().buf.len();

    // No-stats file should have exactly 1 byte overhead per row group (varint 0)
    // With stats: varint(1) + varint(col_idx=0) + varint(null_count=0) + has_non_null(1) + min(4) + max(4) = ~12 bytes
    let overhead = with_stats_size - no_stats_size;
    assert!(
        overhead < 20,
        "stats overhead too large: {} bytes",
        overhead
    );
    assert!(
        overhead >= 10,
        "stats overhead suspiciously small: {} bytes",
        overhead
    );

    // Verify no-stats reader still works
    let no_stats_data = writer.output().buf.clone();
    let no_stats_len = no_stats_data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(no_stats_data), no_stats_len).unwrap();
    let stats = reader.row_group_stats(0).unwrap();
    assert!(stats.is_empty());
}

#[test]
fn test_stats_string_column() {
    let columns = vec![("s".to_string(), DataType::Utf8, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["s".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows = vec![
        vec![Value::String(b"banana".to_vec())],
        vec![Value::String(b"apple".to_vec())],
        vec![Value::String(b"cherry".to_vec())],
    ];
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let stats = reader.row_group_stats(0).unwrap();
    assert_eq!(stats.len(), 1);
    assert!(matches!(&stats[0].min, Some(Value::String(b)) if b == b"apple"));
    assert!(matches!(&stats[0].max, Some(Value::String(b)) if b == b"cherry"));
}

#[test]
fn test_cursor_api_basic() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("val".to_string(), DataType::Utf8, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..50)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("v{}", i).into_bytes()),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 50);

    let ids = batch_col_i32(&batch, "id");
    let vals = batch_col_string(&batch, "val");

    for i in 0..50usize {
        assert_eq!(ids.value(i), i as i32);
        assert_eq!(vals.value(i), format!("v{}", i));
    }
}

#[test]
fn test_cursor_api_overwrites_previous_row() {
    let columns = vec![("x".to_string(), DataType::Int32, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    let rows = vec![vec![Value::Integer(100)], vec![Value::Integer(200)]];
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 2);

    let xs = batch_col_i32(&batch, "x");
    assert_eq!(xs.value(0), 100);
    assert_eq!(xs.value(1), 200);
}

#[test]
fn test_projection_subset() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
        ("d".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 4,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..30)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("s{}", i).into_bytes()),
                Value::Double(i as f64),
                Value::BigInt(i as i64 * 100),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let col_a = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "a")
        .unwrap();
    let col_c = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c")
        .unwrap();

    let mut rg = reader
        .row_group_reader_projected(0, &[col_a, col_c])
        .unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 30);
    assert_eq!(batch.num_columns(), 2);

    let ca = batch_col_i32(&batch, "a");
    let cc = batch_col_f64(&batch, "c");
    for i in 0..30usize {
        assert_eq!(ca.value(i), i as i32);
        assert!((cc.value(i) - i as f64).abs() < 1e-10);
    }
}

#[test]
fn test_projection_empty_columns_preserves_row_count() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows = vec![
        vec![Value::Integer(1), Value::String(b"a".to_vec())],
        vec![Value::Integer(2), Value::String(b"b".to_vec())],
        vec![Value::Integer(3), Value::String(b"c".to_vec())],
    ];
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let mut rg = reader.row_group_reader_projected(0, &[]).unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_columns(), 0);
    assert_eq!(batch.num_rows(), rows.len());
    assert!(batch.schema().fields().is_empty());
}

#[test]
fn test_projection_single_column() {
    let columns = vec![
        ("x".to_string(), DataType::Int32, true),
        ("y".to_string(), DataType::Int32, true),
        ("z".to_string(), DataType::Int32, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 3,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..20)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::Integer(i * 10),
                Value::Integer(i * 100),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let col_y = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "y")
        .unwrap();

    let mut rg = reader.row_group_reader_projected(0, &[col_y]).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 20);
    assert_eq!(batch.num_columns(), 1);

    let cy = batch_col_i32(&batch, "y");
    for i in 0..20usize {
        assert_eq!(cy.value(i), i as i32 * 10);
    }
}

#[test]
fn test_projection_preserves_requested_order() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..10)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("s{}", i).into_bytes()),
                Value::Double(i as f64 * 0.5),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let mut reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    reader.project(&["c", "a", "b"]).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_columns(), 3);
    let schema = batch.schema();
    let fields: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
    assert_eq!(fields, vec!["c", "a", "b"]);

    let col_c = batch
        .column(0)
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();
    let col_a = batch
        .column(1)
        .as_any()
        .downcast_ref::<Int32Array>()
        .unwrap();
    let col_b = batch
        .column(2)
        .as_any()
        .downcast_ref::<StringArray>()
        .unwrap();
    for i in 0..10usize {
        assert_eq!(col_a.value(i), i as i32);
        assert_eq!(col_b.value(i), format!("s{}", i));
        assert!((col_c.value(i) - i as f64 * 0.5).abs() < 1e-10);
    }
}

#[test]
fn test_projection_empty_via_project_method() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..5)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("v{}", i).into_bytes()),
                Value::Double(i as f64),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let mut reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    reader.project(&[]).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_columns(), 0);
    assert_eq!(batch.num_rows(), 5);
}

#[test]
fn test_multiple_row_groups() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("data".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            row_group_max_size: 200,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let total_rows = 500;
    let rows: Vec<Vec<Value>> = (0..total_rows)
        .map(|i| vec![Value::Integer(i), Value::BigInt(i as i64 * 3)])
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1, "expected multiple row groups");

    let mut offset = 0usize;
    for rg_idx in 0..reader.num_row_groups() {
        let mut rg = reader.row_group_reader(rg_idx).unwrap();
        let batch = rg.read_columns().unwrap();
        let ids = batch_col_i32(&batch, "id");
        let data = batch_col_i64(&batch, "data");
        for i in 0..batch.num_rows() {
            assert_eq!(ids.value(i), (offset + i) as i32);
            assert_eq!(data.value(i), (offset + i) as i64 * 3);
        }
        offset += batch.num_rows();
    }
    assert_eq!(offset, total_rows as usize);
}

#[test]
fn test_multiple_row_groups_with_projection() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            row_group_max_size: 100,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let total_rows = 200;
    let rows: Vec<Vec<Value>> = (0..total_rows)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("r{}", i).into_bytes()),
                Value::Double(i as f64),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1);

    let col_c = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c")
        .unwrap();
    let mut count = 0usize;
    for rg_idx in 0..reader.num_row_groups() {
        let mut rg = reader.row_group_reader_projected(rg_idx, &[col_c]).unwrap();
        let batch = rg.read_columns().unwrap();
        let cc = batch_col_f64(&batch, "c");
        for i in 0..batch.num_rows() {
            assert!((cc.value(i) - (count + i) as f64).abs() < 1e-10);
        }
        count += batch.num_rows();
    }
    assert_eq!(count, total_rows as usize);
}

#[test]
fn test_single_row() {
    let columns = vec![
        ("k".to_string(), DataType::Int32, true),
        ("v".to_string(), DataType::Utf8, true),
    ];
    let (reader, _) = write_and_read(
        columns,
        &[vec![Value::Integer(42), Value::String(b"only".to_vec())]],
    );
    let mut rg = reader.row_group_reader(0).unwrap();
    assert_eq!(rg.num_rows(), 1);
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 1);
    assert_eq!(batch_col_i32(&batch, "k").value(0), 42);
    assert_eq!(batch_col_string(&batch, "v").value(0), "only");
}

#[test]
fn test_zero_rows() {
    let columns = vec![("a".to_string(), DataType::Int32, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    assert_eq!(reader.num_row_groups(), 0);
}

#[test]
fn test_many_columns_few_buckets() {
    let columns: Vec<(String, DataType, bool)> = (0..50)
        .map(|i| (format!("col_{:03}", i), DataType::Int32, true))
        .collect();
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 3,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..20)
        .map(|r| (0..50).map(|c| Value::Integer(r * 50 + c)).collect())
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    assert_eq!(reader.schema().columns.len(), 50);

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 20);
    assert_eq!(batch.num_columns(), 50);

    for r in 0..20usize {
        for c in 0..50usize {
            let col = batch
                .column(c)
                .as_any()
                .downcast_ref::<Int32Array>()
                .unwrap();
            assert_eq!(col.value(r), (r as i32) * 50 + c as i32);
        }
    }
}

#[test]
fn test_single_bucket() {
    let columns = vec![
        ("x".to_string(), DataType::Int32, true),
        ("y".to_string(), DataType::Float64, true),
        ("z".to_string(), DataType::Utf8, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    let rows: Vec<Vec<Value>> = (0..30)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::Double(i as f64),
                Value::String(format!("s{}", i).into_bytes()),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 30);

    let xs = batch_col_i32(&batch, "x");
    for i in 0..30usize {
        assert_eq!(xs.value(i), i as i32);
    }
}

#[test]
fn test_const_encoding_with_nulls() {
    let columns = vec![("v".to_string(), DataType::Int32, true)];
    let mut rows = Vec::new();
    for i in 0..20 {
        if i % 3 == 0 {
            rows.push(vec![Value::Null]);
        } else {
            rows.push(vec![Value::Integer(77)]);
        }
    }
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 20);

    let vs = batch_col_i32(&batch, "v");
    for i in 0..20usize {
        if i % 3 == 0 {
            assert!(vs.is_null(i));
        } else {
            assert_eq!(vs.value(i), 77);
        }
    }
}

#[test]
fn test_dict_encoding_with_nulls() {
    let columns = vec![("v".to_string(), DataType::Int32, true)];
    let mut rows = Vec::new();
    for i in 0..100 {
        if i % 5 == 0 {
            rows.push(vec![Value::Null]);
        } else {
            rows.push(vec![Value::Integer(i % 4)]);
        }
    }
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);

    let vs = batch_col_i32(&batch, "v");
    for i in 0..100usize {
        if i % 5 == 0 {
            assert!(vs.is_null(i));
        } else {
            assert_eq!(vs.value(i), i as i32 % 4);
        }
    }
}

#[test]
fn test_const_string_encoding() {
    let columns = vec![("s".to_string(), DataType::Utf8, true)];
    let rows: Vec<Vec<Value>> = (0..50)
        .map(|_| vec![Value::String(b"constant_value".to_vec())])
        .collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 50);

    let ss = batch_col_string(&batch, "s");
    for i in 0..50usize {
        assert_eq!(ss.value(i), "constant_value");
    }
}

#[test]
fn test_dict_string_encoding() {
    let dict_values = ["apple", "banana", "cherry"];
    let columns = vec![("fruit".to_string(), DataType::Utf8, true)];
    let rows: Vec<Vec<Value>> = (0..90)
        .map(|i| vec![Value::String(dict_values[i % 3].as_bytes().to_vec())])
        .collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 90);

    let fruits = batch_col_string(&batch, "fruit");
    for i in 0..90usize {
        assert_eq!(fruits.value(i), dict_values[i % 3]);
    }
}

#[test]
fn test_dict_string_encoding_with_nulls() {
    let dict_values = ["foo", "bar"];
    let columns = vec![("s".to_string(), DataType::Utf8, true)];
    let mut rows = Vec::new();
    for i in 0..60 {
        if i % 4 == 0 {
            rows.push(vec![Value::Null]);
        } else {
            rows.push(vec![Value::String(dict_values[i % 2].as_bytes().to_vec())]);
        }
    }
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 60);

    let ss = batch_col_string(&batch, "s");
    for i in 0..60usize {
        if i % 4 == 0 {
            assert!(ss.is_null(i));
        } else {
            assert_eq!(ss.value(i), dict_values[i % 2]);
        }
    }
}

#[test]
fn test_varchar_char_roundtrip() {
    let columns = vec![
        ("vc".to_string(), DataType::Utf8, true),
        ("ch".to_string(), DataType::Utf8, true),
    ];
    let rows = vec![
        vec![
            Value::String(b"hello".to_vec()),
            Value::String(b"world12345".to_vec()),
        ],
        vec![Value::String(b"test".to_vec()), Value::Null],
        vec![Value::Null, Value::String(b"abcdefghij".to_vec())],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 3);

    let vc = batch_col_string(&batch, "vc");
    let ch = batch_col_string(&batch, "ch");

    assert_eq!(vc.value(0), "hello");
    assert_eq!(ch.value(0), "world12345");
    assert_eq!(vc.value(1), "test");
    assert!(ch.is_null(1));
    assert!(vc.is_null(2));
    assert_eq!(ch.value(2), "abcdefghij");
}

#[test]
fn test_binary_varbinary_roundtrip() {
    let columns = vec![
        ("bin".to_string(), DataType::Binary, true),
        ("vbin".to_string(), DataType::Binary, true),
    ];
    let rows = vec![
        vec![
            Value::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]),
            Value::Bytes(vec![1, 2, 3]),
        ],
        vec![Value::Null, Value::Bytes(vec![0xFF; 100])],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 2);

    let bin = batch_col_binary(&batch, "bin");
    let vbin = batch_col_binary(&batch, "vbin");

    assert_eq!(bin.value(0), &[0xDE, 0xAD, 0xBE, 0xEF]);
    assert_eq!(vbin.value(0), &[1, 2, 3]);
    assert!(bin.is_null(1));
    assert_eq!(vbin.value(1).len(), 100);
    assert!(vbin.value(1).iter().all(|&x| x == 0xFF));
}

#[test]
fn test_time_roundtrip() {
    let columns = vec![(
        "t".to_string(),
        DataType::Time32(TimeUnit::Millisecond),
        true,
    )];
    let rows = vec![
        vec![Value::Time(0)],
        vec![Value::Time(43200000)],
        vec![Value::Null],
        vec![Value::Time(86399999)],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 4);

    let ts = batch
        .column(0)
        .as_any()
        .downcast_ref::<Time32MillisecondArray>()
        .unwrap();
    assert_eq!(ts.value(0), 0);
    assert_eq!(ts.value(1), 43200000);
    assert!(ts.is_null(2));
    assert_eq!(ts.value(3), 86399999);
}

#[test]
fn test_timestamp_ltz_roundtrip() {
    let columns = vec![
        (
            "ts_millis".to_string(),
            DataType::Timestamp(TimeUnit::Millisecond, Some("Asia/Shanghai".into())),
            true,
        ),
        (
            "ts_micros".to_string(),
            DataType::Timestamp(TimeUnit::Microsecond, Some("America/New_York".into())),
            true,
        ),
    ];
    let rows = vec![
        vec![
            Value::TimestampMillis(1700000000000),
            Value::TimestampMicros(1700000000000000),
        ],
        vec![Value::Null, Value::Null],
        vec![Value::TimestampMillis(0), Value::TimestampMicros(0)],
    ];
    let (reader, _) = write_and_read(columns, &rows);

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 3);

    let millis_idx = batch.schema().index_of("ts_millis").unwrap();
    let micros_idx = batch.schema().index_of("ts_micros").unwrap();

    let ts0 = batch
        .column(millis_idx)
        .as_any()
        .downcast_ref::<TimestampMillisecondArray>()
        .unwrap();
    assert_eq!(ts0.value(0), 1700000000000i64);
    assert!(ts0.is_null(1));
    assert_eq!(ts0.value(2), 0);
    assert_eq!(ts0.timezone().unwrap(), "Asia/Shanghai");

    let ts1 = batch
        .column(micros_idx)
        .as_any()
        .downcast_ref::<TimestampMicrosecondArray>()
        .unwrap();
    assert_eq!(ts1.value(0), 1700000000000000i64);
    assert!(ts1.is_null(1));
    assert_eq!(ts1.value(2), 0);
    assert_eq!(ts1.timezone().unwrap(), "America/New_York");
}

#[test]
fn test_timestamp_nanos_roundtrip() {
    let columns = vec![(
        "ts_nanos".to_string(),
        DataType::Timestamp(TimeUnit::Nanosecond, None),
        true,
    )];
    let rows = vec![
        vec![Value::TimestampNanos {
            millis: 1700000000000,
            nanos_of_milli: 123456,
        }],
        vec![Value::Null],
        vec![Value::TimestampNanos {
            millis: -1,
            nanos_of_milli: 999999,
        }],
        vec![Value::TimestampNanos {
            millis: 0,
            nanos_of_milli: 0,
        }],
    ];
    let (reader, _) = write_and_read(columns, &rows);

    assert_eq!(
        reader.schema().columns[0].data_type,
        DataType::Timestamp(TimeUnit::Nanosecond, None)
    );

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    let ts = batch
        .column(0)
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .unwrap();
    assert_eq!(
        ts.value(0),
        crate::types::millis_nanos_to_ns(1700000000000, 123456).unwrap()
    );
    assert!(ts.is_null(1));
    assert_eq!(ts.value(2), -1);
    assert_eq!(ts.value(3), 0);
}

#[test]
fn test_timestamp_ltz_nanos_roundtrip() {
    let columns = vec![(
        "ts_ltz_nanos".to_string(),
        DataType::Timestamp(TimeUnit::Nanosecond, Some("Asia/Shanghai".into())),
        true,
    )];
    let rows = vec![
        vec![Value::TimestampNanos {
            millis: 1700000000000,
            nanos_of_milli: 42,
        }],
        vec![Value::Null],
    ];
    let (reader, _) = write_and_read(columns, &rows);

    assert_eq!(
        reader.schema().columns[0].data_type,
        DataType::Timestamp(TimeUnit::Nanosecond, Some("Asia/Shanghai".into()))
    );

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    let ts = batch
        .column(0)
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .unwrap();
    assert_eq!(ts.timezone().unwrap(), "Asia/Shanghai");
    assert_eq!(
        ts.value(0),
        crate::types::millis_nanos_to_ns(1700000000000, 42).unwrap()
    );
    assert!(ts.is_null(1));
}

#[test]
fn test_legacy_timestamp_nanos_struct_writes_as_arrow_nanos() {
    let columns = vec![("ts".to_string(), legacy_timestamp_nanos_type(), true)];
    let rows = vec![vec![Value::TimestampNanos {
        millis: 1700000000000,
        nanos_of_milli: 123456,
    }]];
    let (reader, _) = write_and_read(columns, &rows);

    assert_eq!(
        reader.schema().columns[0].data_type,
        DataType::Timestamp(TimeUnit::Nanosecond, None)
    );

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    let ts = batch
        .column(0)
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .unwrap();
    assert_eq!(
        ts.value(0),
        crate::types::millis_nanos_to_ns(1700000000000, 123456).unwrap()
    );
}

#[test]
fn test_legacy_timestamp_nanos_struct_rejects_child_null() {
    let legacy_fields: arrow_schema::Fields = vec![
        Field::new("millis", DataType::Int64, true),
        Field::new("nanos_of_milli", DataType::Int32, true),
    ]
    .into();
    let legacy_dt = DataType::Struct(legacy_fields.clone());
    let schema = Schema::new(vec![Field::new("ts", legacy_dt.clone(), true)]);
    let struct_array = StructArray::new(
        legacy_fields,
        vec![
            Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef,
            Arc::new(Int32Array::from(vec![None])) as ArrayRef,
        ],
        Some(arrow_buffer::NullBuffer::from(vec![true])),
    );
    let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap();
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &Schema::new(vec![Field::new("ts", legacy_dt, true)]),
        WriterOptions::default(),
    )
    .unwrap();

    let err = writer.write_batch(&batch).unwrap_err();
    assert!(err.to_string().contains("null child"));
}

#[test]
fn test_legacy_timestamp_nanos_struct_rejects_invalid_nanos() {
    let legacy_dt = legacy_timestamp_nanos_type();
    let schema = Schema::new(vec![Field::new("ts", legacy_dt.clone(), true)]);
    let fields = legacy_timestamp_nanos_fields();
    let struct_array = StructArray::new(
        fields,
        vec![
            Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef,
            Arc::new(Int32Array::from(vec![Some(1_000_000)])) as ArrayRef,
        ],
        Some(arrow_buffer::NullBuffer::from(vec![true])),
    );
    let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap();
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &Schema::new(vec![Field::new("ts", legacy_dt, true)]),
        WriterOptions::default(),
    )
    .unwrap();

    let err = writer.write_batch(&batch).unwrap_err();
    assert!(err.to_string().contains("invalid nanos_of_milli"));
}

#[test]
fn test_legacy_timestamp_nanos_struct_rejects_overflow() {
    let legacy_dt = legacy_timestamp_nanos_type();
    let schema = Schema::new(vec![Field::new("ts", legacy_dt.clone(), true)]);
    let fields = legacy_timestamp_nanos_fields();
    let struct_array = StructArray::new(
        fields,
        vec![
            Arc::new(Int64Array::from(vec![Some(i64::MAX)])) as ArrayRef,
            Arc::new(Int32Array::from(vec![Some(0)])) as ArrayRef,
        ],
        Some(arrow_buffer::NullBuffer::from(vec![true])),
    );
    let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap();
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &Schema::new(vec![Field::new("ts", legacy_dt, true)]),
        WriterOptions::default(),
    )
    .unwrap();

    let err = writer.write_batch(&batch).unwrap_err();
    assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    assert!(err.to_string().contains("overflow"));
}

#[test]
fn test_timestamp_nanos_stats_min_max() {
    let columns = vec![(
        "ts".to_string(),
        DataType::Timestamp(TimeUnit::Nanosecond, None),
        true,
    )];
    let rows = vec![
        vec![Value::TimestampNanos {
            millis: -1,
            nanos_of_milli: 999999,
        }],
        vec![Value::Null],
        vec![Value::TimestampNanos {
            millis: 0,
            nanos_of_milli: 42,
        }],
        vec![Value::TimestampNanos {
            millis: -1,
            nanos_of_milli: 999998,
        }],
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            stats_columns: vec!["ts".to_string()],
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();

    let stats = writer.row_group_stats(0);
    assert_eq!(stats.len(), 1);
    assert_eq!(stats[0].null_count, 1);
    assert!(matches!(
        stats[0].min,
        Some(Value::TimestampNanos {
            millis: -1,
            nanos_of_milli: 999998
        })
    ));
    assert!(matches!(
        stats[0].max,
        Some(Value::TimestampNanos {
            millis: 0,
            nanos_of_milli: 42
        })
    ));
}

#[test]
fn test_timestamp_micros_roundtrip() {
    let columns = vec![
        (
            "ts_millis".to_string(),
            DataType::Timestamp(TimeUnit::Millisecond, None),
            true,
        ),
        (
            "ts_micros".to_string(),
            DataType::Timestamp(TimeUnit::Microsecond, None),
            true,
        ),
        (
            "ts_nanos".to_string(),
            DataType::Struct(
                vec![
                    arrow_schema::Field::new("millis", DataType::Int64, false),
                    arrow_schema::Field::new("nanos_of_milli", DataType::Int32, false),
                ]
                .into(),
            ),
            true,
        ),
    ];
    let rows = vec![
        vec![
            Value::TimestampMillis(1700000000000),
            Value::TimestampMicros(1_700_000_000_000_000),
            Value::TimestampNanos {
                millis: 1700000000000,
                nanos_of_milli: 123456,
            },
        ],
        vec![Value::Null, Value::Null, Value::Null],
        vec![
            Value::TimestampMillis(0),
            Value::TimestampMicros(0),
            Value::TimestampNanos {
                millis: 0,
                nanos_of_milli: 0,
            },
        ],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 3);

    let millis_pos = batch.schema().index_of("ts_millis").unwrap();
    let micros_pos = batch.schema().index_of("ts_micros").unwrap();
    let nanos_pos = batch.schema().index_of("ts_nanos").unwrap();

    let ts_millis = batch
        .column(millis_pos)
        .as_any()
        .downcast_ref::<TimestampMillisecondArray>()
        .unwrap();
    let ts_micros = batch
        .column(micros_pos)
        .as_any()
        .downcast_ref::<TimestampMicrosecondArray>()
        .unwrap();
    let ts_nanos = batch
        .column(nanos_pos)
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .unwrap();

    assert_eq!(ts_millis.value(0), 1700000000000i64);
    assert_eq!(ts_micros.value(0), 1_700_000_000_000_000i64);
    assert_eq!(
        ts_nanos.value(0),
        crate::types::millis_nanos_to_ns(1700000000000, 123456).unwrap()
    );

    assert!(ts_millis.is_null(1));
    assert!(ts_micros.is_null(1));
    assert!(ts_nanos.is_null(1));

    assert_eq!(ts_millis.value(2), 0);
    assert_eq!(ts_micros.value(2), 0);
    assert_eq!(ts_nanos.value(2), 0);
}

#[test]
fn test_decimal_large_roundtrip() {
    let columns = vec![("d".to_string(), DataType::Decimal128(30, 5), true)];
    let large_bytes = vec![
        0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0x01, 0x02, 0x03, 0x04,
    ];
    let rows = vec![
        vec![Value::DecimalLarge(large_bytes.clone())],
        vec![Value::Null],
        vec![Value::DecimalLarge(vec![0xFF])],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 3);

    let ds = batch
        .column(0)
        .as_any()
        .downcast_ref::<Decimal128Array>()
        .unwrap();
    assert_eq!(ds.value(0), biginteger_bytes_to_i128(&large_bytes));
    assert!(ds.is_null(1));
    assert_eq!(ds.value(2), biginteger_bytes_to_i128(&[0xFF]));
}

#[test]
fn test_writer_close_idempotent() {
    let columns = vec![("a".to_string(), DataType::Int32, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &[vec![Value::Integer(1)]]);
    writer.close().unwrap();
    let size_after_first_close = writer.output().buf.len();

    writer.close().unwrap();
    let size_after_second_close = writer.output().buf.len();
    assert_eq!(size_after_first_close, size_after_second_close);
}

#[test]
fn test_all_null_rows() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..30)
        .map(|_| vec![Value::Null, Value::Null, Value::Null])
        .collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 30);

    for c in 0..batch.num_columns() {
        for r in 0..30usize {
            assert!(batch.column(c).is_null(r));
        }
    }
}

#[test]
fn test_mixed_encodings_multi_column() {
    let columns = vec![
        ("all_null".to_string(), DataType::Int32, true),
        ("const_int".to_string(), DataType::Int64, true),
        ("dict_str".to_string(), DataType::Utf8, true),
        ("plain_dbl".to_string(), DataType::Float64, true),
    ];
    let dict_vals = ["alpha", "beta", "gamma"];
    let mut rows = Vec::new();
    for i in 0..100 {
        rows.push(vec![
            Value::Null,
            Value::BigInt(999),
            Value::String(dict_vals[i % 3].as_bytes().to_vec()),
            Value::Double(i as f64 * 0.1),
        ]);
    }
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);

    let col0 = batch_col_i32(&batch, "all_null");
    let col1 = batch_col_i64(&batch, "const_int");
    let col2 = batch_col_string(&batch, "dict_str");
    let col3 = batch_col_f64(&batch, "plain_dbl");

    for i in 0..100usize {
        assert!(col0.is_null(i));
        assert_eq!(col1.value(i), 999);
        assert_eq!(col2.value(i), dict_vals[i % 3]);
        assert!((col3.value(i) - i as f64 * 0.1).abs() < 1e-10);
    }
}

#[test]
fn test_cursor_api_with_projection() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();
    let rows: Vec<Vec<Value>> = (0..40)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("x{}", i).into_bytes()),
                Value::BigInt(i as i64 * 7),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let col_b = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "b")
        .unwrap();

    let mut rg = reader.row_group_reader_projected(0, &[col_b]).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 40);

    let bs = batch_col_string(&batch, "b");
    for i in 0..40usize {
        assert_eq!(bs.value(i), format!("x{}", i));
    }
}

#[test]
fn test_boolean_all_true() {
    let columns = vec![("b".to_string(), DataType::Boolean, true)];
    let rows: Vec<Vec<Value>> = (0..20).map(|_| vec![Value::Boolean(true)]).collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 20);

    let bs = batch_col_bool(&batch, "b");
    for i in 0..20usize {
        assert!(bs.value(i));
    }
}

#[test]
fn test_boolean_alternating() {
    let columns = vec![("b".to_string(), DataType::Boolean, true)];
    let rows: Vec<Vec<Value>> = (0..40).map(|i| vec![Value::Boolean(i % 2 == 0)]).collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 40);

    let bs = batch_col_bool(&batch, "b");
    for i in 0..40usize {
        assert_eq!(bs.value(i), i % 2 == 0);
    }
}

#[test]
fn test_zstd_with_multiple_row_groups() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            zstd_level: 1,
            row_group_max_size: 200,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let total = 300;
    let rows: Vec<Vec<Value>> = (0..total)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("name_{}", i).into_bytes()),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1);

    let mut count = 0usize;
    for rg_idx in 0..reader.num_row_groups() {
        let mut rg = reader.row_group_reader(rg_idx).unwrap();
        let batch = rg.read_columns().unwrap();
        let ids = batch_col_i32(&batch, "id");
        let names = batch_col_string(&batch, "name");
        for i in 0..batch.num_rows() {
            assert_eq!(ids.value(i), (count + i) as i32);
            assert_eq!(names.value(i), format!("name_{}", count + i));
        }
        count += batch.num_rows();
    }
    assert_eq!(count, total as usize);
}

#[test]
fn test_stats_across_multiple_row_groups() {
    let columns = vec![("v".to_string(), DataType::Int32, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            row_group_max_size: 100,
            stats_columns: vec!["v".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..500).map(|i| vec![Value::Integer(i)]).collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1);

    let mut prev_max = i32::MIN;
    for rg_idx in 0..reader.num_row_groups() {
        let stats = reader.row_group_stats(rg_idx).unwrap();
        assert_eq!(stats.len(), 1);
        assert_eq!(stats[0].null_count, 0);
        if let (Some(Value::Integer(min)), Some(Value::Integer(max))) =
            (&stats[0].min, &stats[0].max)
        {
            assert!(*min >= prev_max || rg_idx == 0);
            assert!(*max >= *min);
            prev_max = *max;
        } else {
            panic!("expected Integer stats");
        }
    }
}

#[test]
fn test_schema_order_preserved_after_roundtrip() {
    let columns = vec![
        ("name".to_string(), DataType::Utf8, true),
        ("age".to_string(), DataType::Int32, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["age".to_string()],
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..50)
        .map(|i| {
            vec![
                Value::String(format!("user_{}", i).into_bytes()),
                Value::Integer(20 + i),
                Value::Double(i as f64 * 0.5),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    // Internal storage is name-sorted
    assert_eq!(reader.schema().columns[0].name, "age");
    assert_eq!(reader.schema().columns[1].name, "name");
    assert_eq!(reader.schema().columns[2].name, "score");
    // original_order preserves input order: name=1, age=0, score=2
    assert_eq!(reader.schema().original_order, vec![1, 0, 2]);

    let stats = reader.row_group_stats(0).unwrap();
    assert_eq!(stats.len(), 1);
    assert_eq!(stats[0].column_index, 0);
    assert!(matches!(&stats[0].min, Some(Value::Integer(20))));
    assert!(matches!(&stats[0].max, Some(Value::Integer(69))));

    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    // Output batch is in original input order (name, age, score)
    assert_eq!(batch.schema().field(0).name(), "name");
    assert_eq!(batch.schema().field(1).name(), "age");
    assert_eq!(batch.schema().field(2).name(), "score");
    assert_eq!(batch_col_string(&batch, "name").value(0), "user_0");
    assert_eq!(batch_col_i32(&batch, "age").value(0), 20);
}

#[test]
fn test_write_batch_non_nullable_rejects_nulls() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, false),
        ("b".to_string(), DataType::Int32, false),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let bad_batch = RecordBatch::try_new(
        Arc::new(Schema::new(vec![
            Field::new("a", DataType::Int32, true),
            Field::new("b", DataType::Int32, true),
        ])),
        vec![
            Arc::new(Int32Array::from(vec![None, Some(1)])),
            Arc::new(Int32Array::from(vec![Some(2), Some(3)])),
        ],
    )
    .unwrap();
    assert!(writer.write_batch(&bad_batch).is_err());

    let good_rows = vec![vec![Value::Integer(222), Value::Integer(333)]];
    write_values(&mut writer, &columns, &good_rows);

    writer.close().unwrap();
    let data = &writer.output().buf;

    let input = ByteArrayInputFile::new(data.clone());
    let reader = MosaicReader::new(input, data.len() as u64).unwrap();
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 1);
    assert_eq!(batch_col_i32(&batch, "a").value(0), 222);
    assert_eq!(batch_col_i32(&batch, "b").value(0), 333);
}

fn write_and_read_paged(
    columns: Vec<(String, DataType, bool)>,
    rows: &[Vec<Value>],
) -> (MosaicReader<ByteArrayInputFile>, Vec<u8>) {
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data.clone()), len).unwrap();
    (reader, data)
}

#[test]
fn test_paged_roundtrip_basic() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..200)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("user_{}", i).into_bytes()),
                Value::Double(i as f64 * 1.5),
            ]
        })
        .collect();

    let (reader, _) = write_and_read_paged(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 200);

    let ids = batch_col_i32(&batch, "id");
    let names = batch_col_string(&batch, "name");
    let scores = batch_col_f64(&batch, "score");

    for i in 0..200usize {
        assert_eq!(ids.value(i), i as i32);
        assert_eq!(names.value(i), format!("user_{}", i));
        assert!((scores.value(i) - i as f64 * 1.5).abs() < 1e-10);
    }
}

#[test]
fn test_paged_projection() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
        ("d".to_string(), DataType::Int64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("s{}", i).into_bytes()),
                Value::Double(i as f64),
                Value::BigInt(i as i64 * 100),
            ]
        })
        .collect();

    let (reader, _) = write_and_read_paged(columns, &rows);
    let col_a = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "a")
        .unwrap();
    let col_c = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c")
        .unwrap();

    let mut rg = reader
        .row_group_reader_projected(0, &[col_a, col_c])
        .unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);
    assert_eq!(batch.num_columns(), 2);

    let ca = batch_col_i32(&batch, "a");
    let cc = batch_col_f64(&batch, "c");
    for i in 0..100usize {
        assert_eq!(ca.value(i), i as i32);
        assert!((cc.value(i) - i as f64).abs() < 1e-10);
    }
}

#[test]
fn test_paged_mixed_encodings() {
    let columns = vec![
        ("all_null".to_string(), DataType::Int32, true),
        ("const_int".to_string(), DataType::Int64, true),
        ("dict_str".to_string(), DataType::Utf8, true),
        ("plain_dbl".to_string(), DataType::Float64, true),
    ];
    let dict_vals = ["alpha", "beta", "gamma"];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Null,
                Value::BigInt(999),
                Value::String(dict_vals[i % 3].as_bytes().to_vec()),
                Value::Double(i as f64 * 0.1),
            ]
        })
        .collect();

    let (reader, _) = write_and_read_paged(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);

    let col0 = batch_col_i32(&batch, "all_null");
    let col1 = batch_col_i64(&batch, "const_int");
    let col2 = batch_col_string(&batch, "dict_str");
    let col3 = batch_col_f64(&batch, "plain_dbl");

    for i in 0..100usize {
        assert!(col0.is_null(i));
        assert_eq!(col1.value(i), 999);
        assert_eq!(col2.value(i), dict_vals[i % 3]);
        assert!((col3.value(i) - i as f64 * 0.1).abs() < 1e-10);
    }
}

#[test]
fn test_paged_const_with_nulls() {
    let columns = vec![("v".to_string(), DataType::Int32, true)];
    let rows: Vec<Vec<Value>> = (0..50)
        .map(|i| {
            if i % 3 == 0 {
                vec![Value::Null]
            } else {
                vec![Value::Integer(42)]
            }
        })
        .collect();

    let (reader, _) = write_and_read_paged(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 50);

    let vs = batch_col_i32(&batch, "v");
    for i in 0..50usize {
        if i % 3 == 0 {
            assert!(vs.is_null(i));
        } else {
            assert_eq!(vs.value(i), 42);
        }
    }
}

#[test]
fn test_paged_dict_with_nulls() {
    let columns = vec![("v".to_string(), DataType::Int32, true)];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            if i % 5 == 0 {
                vec![Value::Null]
            } else {
                vec![Value::Integer(i % 3)]
            }
        })
        .collect();

    let (reader, _) = write_and_read_paged(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);

    let vs = batch_col_i32(&batch, "v");
    for i in 0..100usize {
        if i % 5 == 0 {
            assert!(vs.is_null(i));
        } else {
            assert_eq!(vs.value(i), i as i32 % 3);
        }
    }
}

#[test]
fn test_paged_multiple_row_groups() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("data".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            row_group_max_size: 200,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let total_rows = 500;
    let rows: Vec<Vec<Value>> = (0..total_rows)
        .map(|i| vec![Value::Integer(i), Value::BigInt(i as i64 * 3)])
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1);

    let mut offset = 0usize;
    for rg_idx in 0..reader.num_row_groups() {
        let mut rg = reader.row_group_reader(rg_idx).unwrap();
        let batch = rg.read_columns().unwrap();
        let ids = batch_col_i32(&batch, "id");
        let data = batch_col_i64(&batch, "data");
        for i in 0..batch.num_rows() {
            assert_eq!(ids.value(i), (offset + i) as i32);
            assert_eq!(data.value(i), (offset + i) as i64 * 3);
        }
        offset += batch.num_rows();
    }
    assert_eq!(offset, total_rows as usize);
}

#[test]
fn test_paged_string_dict_encoding() {
    let dict_values = ["apple", "banana", "cherry"];
    let columns = vec![("fruit".to_string(), DataType::Utf8, true)];
    let rows: Vec<Vec<Value>> = (0..90)
        .map(|i| vec![Value::String(dict_values[i % 3].as_bytes().to_vec())])
        .collect();
    let (reader, _) = write_and_read_paged(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 90);

    let fruits = batch_col_string(&batch, "fruit");
    for i in 0..90usize {
        assert_eq!(fruits.value(i), dict_values[i % 3]);
    }
}

#[test]
fn test_read_columns_basic() {
    let columns = vec![
        ("age".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Integer(20 + (i % 50)),
                Value::String(format!("user_{}", i).into_bytes()),
                Value::Double(i as f64 * 1.5),
            ]
        })
        .collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_rows(), 100);
    assert_eq!(batch.num_columns(), 3);

    let ages = batch_col_i32(&batch, "age");
    let names = batch_col_string(&batch, "name");
    let scores = batch_col_f64(&batch, "score");

    assert_eq!(ages.null_count(), 0);

    for i in 0..100 {
        assert_eq!(ages.value(i), 20 + (i as i32 % 50));
        assert_eq!(names.value(i), format!("user_{}", i));
        assert!((scores.value(i) - i as f64 * 1.5).abs() < 1e-10);
    }
}

#[test]
fn test_read_columns_with_nulls() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("value".to_string(), DataType::Float64, true),
    ];
    let rows = vec![
        vec![
            Value::Integer(1),
            Value::String(b"hello".to_vec()),
            Value::Double(1.0),
        ],
        vec![Value::Integer(2), Value::Null, Value::Double(2.0)],
        vec![
            Value::Integer(3),
            Value::String(b"world".to_vec()),
            Value::Null,
        ],
        vec![Value::Integer(4), Value::Null, Value::Null],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    let ids = batch_col_i32(&batch, "id");
    let names = batch_col_string(&batch, "name");
    let values = batch_col_f64(&batch, "value");

    assert_eq!(ids.value(0), 1);
    assert_eq!(ids.value(1), 2);
    assert_eq!(ids.value(2), 3);
    assert_eq!(ids.value(3), 4);

    assert!(!names.is_null(0));
    assert_eq!(names.value(0), "hello");
    assert!(names.is_null(1));
    assert!(!names.is_null(2));
    assert_eq!(names.value(2), "world");
    assert!(names.is_null(3));

    assert!(!values.is_null(0));
    assert!(!values.is_null(1));
    assert!(values.is_null(2));
    assert!(values.is_null(3));
}

#[test]
fn test_read_columns_with_projection() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
        ("d".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            num_buckets: 4,
            ..Default::default()
        },
    )
    .unwrap();
    let rows: Vec<Vec<Value>> = (0..30)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("s{}", i).into_bytes()),
                Value::Double(i as f64),
                Value::BigInt(i as i64 * 100),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let col_a = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "a")
        .unwrap();
    let col_c = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c")
        .unwrap();

    let mut rg = reader
        .row_group_reader_projected(0, &[col_a, col_c])
        .unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_columns(), 2);
    assert!(batch.schema().index_of("a").is_ok());
    assert!(batch.schema().index_of("c").is_ok());

    let ca = batch_col_i32(&batch, "a");
    let cc = batch_col_f64(&batch, "c");
    for i in 0..30 {
        assert_eq!(ca.value(i), i as i32);
        assert!((cc.value(i) - i as f64).abs() < 1e-10);
    }
}

#[test]
fn test_read_columns_with_projection_monolithic_bucket() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Utf8, true),
        ("c".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..12)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("b{}", i).into_bytes()),
                Value::Double(i as f64 * 2.5),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let col_b = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "b")
        .unwrap();

    let mut rg = reader.row_group_reader_projected(0, &[col_b]).unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_rows(), 12);
    assert_eq!(batch.num_columns(), 1);
    assert!(batch.schema().index_of("a").is_err());
    assert!(batch.schema().index_of("c").is_err());

    let b = batch_col_string(&batch, "b");
    for i in 0..12usize {
        assert_eq!(b.value(i), format!("b{}", i));
    }
}

#[test]
fn test_read_columns_paged() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..200)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::String(format!("user_{}", i).into_bytes()),
                Value::Double(i as f64 * 1.5),
            ]
        })
        .collect();
    let (reader, _) = write_and_read_paged(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    let ids = batch_col_i32(&batch, "id");
    let names = batch_col_string(&batch, "name");
    let scores = batch_col_f64(&batch, "score");

    for i in 0..200 {
        assert_eq!(ids.value(i), i as i32);
        assert_eq!(names.value(i), format!("user_{}", i));
        assert!((scores.value(i) - i as f64 * 1.5).abs() < 1e-10);
    }
}

#[test]
fn test_read_columns_mixed_encodings() {
    let columns = vec![
        ("all_null".to_string(), DataType::Int32, true),
        ("const_int".to_string(), DataType::Int64, true),
        ("dict_str".to_string(), DataType::Utf8, true),
        ("plain_dbl".to_string(), DataType::Float64, true),
    ];
    let dict_vals = ["alpha", "beta", "gamma"];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Null,
                Value::BigInt(999),
                Value::String(dict_vals[i % 3].as_bytes().to_vec()),
                Value::Double(i as f64 * 0.1),
            ]
        })
        .collect();
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    let col0 = batch_col_i32(&batch, "all_null");
    assert_eq!(col0.len(), 100);
    for i in 0..100 {
        assert!(col0.is_null(i));
    }

    let col1 = batch_col_i64(&batch, "const_int");
    for i in 0..100 {
        assert_eq!(col1.value(i), 999);
    }

    let col2 = batch_col_string(&batch, "dict_str");
    for i in 0..100 {
        assert_eq!(col2.value(i), dict_vals[i % 3]);
    }

    let col3 = batch_col_f64(&batch, "plain_dbl");
    for i in 0..100 {
        assert!((col3.value(i) - i as f64 * 0.1).abs() < 1e-10);
    }
}

#[test]
fn test_read_columns_binary_offsets() {
    let columns = vec![("s".to_string(), DataType::Utf8, true)];
    let rows: Vec<Vec<Value>> = vec![
        vec![Value::String(b"hello".to_vec())],
        vec![Value::String(b"world!".to_vec())],
        vec![Value::String(b"".to_vec())],
        vec![Value::String(b"test".to_vec())],
    ];
    let (reader, _) = write_and_read(columns, &rows);
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();

    let col = batch_col_string(&batch, "s");
    assert_eq!(col.value(0), "hello");
    assert_eq!(col.value(1), "world!");
    assert_eq!(col.value(2), "");
    assert_eq!(col.value(3), "test");
}

#[test]
fn test_read_columns_multiple_row_groups() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("data".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            row_group_max_size: 200,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let total_rows = 500;
    let rows: Vec<Vec<Value>> = (0..total_rows)
        .map(|i| vec![Value::Integer(i), Value::BigInt(i as i64 * 3)])
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1);

    let mut all_id = Vec::new();
    let mut all_data = Vec::new();
    for rg_idx in 0..reader.num_row_groups() {
        let mut rg = reader.row_group_reader(rg_idx).unwrap();
        let batch = rg.read_columns().unwrap();

        let ids = batch_col_i32(&batch, "id");
        let data = batch_col_i64(&batch, "data");

        for i in 0..ids.len() {
            all_id.push(ids.value(i));
            all_data.push(data.value(i));
        }
    }

    assert_eq!(all_id.len(), total_rows as usize);
    assert_eq!(all_data.len(), total_rows as usize);
    for i in 0..total_rows as usize {
        assert_eq!(all_id[i], i as i32);
        assert_eq!(all_data[i], i as i64 * 3);
    }
}

// ====================== read_ranges coalescing tests ======================

#[test]
fn test_read_ranges_coalesces_adjacent() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct CountingInputFile {
        data: Vec<u8>,
        read_count: AtomicUsize,
    }

    impl InputFile for CountingInputFile {
        fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
            self.read_count.fetch_add(1, Ordering::Relaxed);
            let start = offset as usize;
            buf.copy_from_slice(&self.data[start..start + buf.len()]);
            Ok(())
        }
    }

    let data = vec![0u8; 4096];
    let input = CountingInputFile {
        data,
        read_count: AtomicUsize::new(0),
    };

    // 4 adjacent ranges (gap = 0) should merge into 1 read_at call
    let ranges = vec![(0u64, 100usize), (100, 200), (300, 100), (400, 200)];
    let results = input.read_ranges(&ranges).unwrap();
    assert_eq!(results.len(), 4);
    assert_eq!(results[0].len(), 100);
    assert_eq!(results[1].len(), 200);
    assert_eq!(results[2].len(), 100);
    assert_eq!(results[3].len(), 200);
    assert_eq!(input.read_count.load(Ordering::Relaxed), 1);
}

#[test]
fn test_read_ranges_shared_reuses_coalesced_buffer() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct CountingInputFile {
        data: Vec<u8>,
        read_count: AtomicUsize,
    }

    impl InputFile for CountingInputFile {
        fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
            self.read_count.fetch_add(1, Ordering::Relaxed);
            let start = offset as usize;
            buf.copy_from_slice(&self.data[start..start + buf.len()]);
            Ok(())
        }
    }

    let data: Vec<u8> = (0..512).map(|i| (i % 256) as u8).collect();
    let input = CountingInputFile {
        data: data.clone(),
        read_count: AtomicUsize::new(0),
    };

    let ranges = vec![(0u64, 64usize), (64, 64), (200, 32)];
    let results = input.read_ranges_shared(&ranges).unwrap();

    assert_eq!(input.read_count.load(Ordering::Relaxed), 1);
    assert_eq!(results[0].as_slice(), &data[0..64]);
    assert_eq!(results[1].as_slice(), &data[64..128]);
    assert_eq!(results[2].as_slice(), &data[200..232]);
    assert!(Arc::ptr_eq(&results[0].data, &results[1].data));
    assert!(Arc::ptr_eq(&results[0].data, &results[2].data));
}

#[test]
fn test_read_range_buffer_new_validates_bounds() {
    let data = Arc::new(vec![1, 2, 3, 4]);

    let buffer = ReadRangeBuffer::new(data.clone(), 1, 2).unwrap();
    assert_eq!(buffer.as_slice(), &[2, 3]);

    match ReadRangeBuffer::new(data, 3, 2) {
        Ok(_) => panic!("expected out-of-bounds range to fail"),
        Err(err) => assert_eq!(err.kind(), io::ErrorKind::InvalidInput),
    }
}

#[test]
fn test_read_ranges_splits_large_gap() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct CountingInputFile {
        data: Vec<u8>,
        read_count: AtomicUsize,
    }

    impl InputFile for CountingInputFile {
        fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
            self.read_count.fetch_add(1, Ordering::Relaxed);
            let start = offset as usize;
            buf.copy_from_slice(&self.data[start..start + buf.len()]);
            Ok(())
        }
    }

    // Gap > 1MB between two ranges — should NOT merge
    let size = 2 * 1024 * 1024 + 200;
    let data = vec![42u8; size];
    let input = CountingInputFile {
        data,
        read_count: AtomicUsize::new(0),
    };

    let gap = COALESCE_GAP as usize + 1;
    let ranges = vec![(0u64, 100usize), ((100 + gap) as u64, 100)];
    let results = input.read_ranges(&ranges).unwrap();
    assert_eq!(results.len(), 2);
    assert_eq!(results[0], vec![42u8; 100]);
    assert_eq!(results[1], vec![42u8; 100]);
    assert_eq!(input.read_count.load(Ordering::Relaxed), 2);
}

#[test]
fn test_read_ranges_coalesces_small_gap() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct CountingInputFile {
        data: Vec<u8>,
        read_count: AtomicUsize,
    }

    impl InputFile for CountingInputFile {
        fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
            self.read_count.fetch_add(1, Ordering::Relaxed);
            let start = offset as usize;
            buf.copy_from_slice(&self.data[start..start + buf.len()]);
            Ok(())
        }
    }

    // Gap = 512KB (< 1MB) — should merge into 1 read
    let gap: usize = 512 * 1024;
    let total = 100 + gap + 100;
    let data: Vec<u8> = (0..total).map(|i| (i % 256) as u8).collect();
    let input = CountingInputFile {
        data: data.clone(),
        read_count: AtomicUsize::new(0),
    };

    let ranges = vec![(0u64, 100usize), ((100 + gap) as u64, 100)];
    let results = input.read_ranges(&ranges).unwrap();
    assert_eq!(results[0], &data[0..100]);
    assert_eq!(results[1], &data[100 + gap..100 + gap + 100]);
    assert_eq!(input.read_count.load(Ordering::Relaxed), 1);
}

#[test]
fn test_read_ranges_out_of_order() {
    let data: Vec<u8> = (0..200).map(|i| i as u8).collect();
    let input = ByteArrayInputFile::new(data.clone());

    // Ranges not in file order — results should still match original order
    let ranges = vec![(150u64, 50usize), (0, 50), (50, 50)];
    let results = input.read_ranges(&ranges).unwrap();
    assert_eq!(results[0], &data[150..200]);
    assert_eq!(results[1], &data[0..50]);
    assert_eq!(results[2], &data[50..100]);
}

#[test]
fn test_read_ranges_empty() {
    let input = ByteArrayInputFile::new(vec![0u8; 100]);
    let results = input.read_ranges(&[]).unwrap();
    assert!(results.is_empty());
}

// ====================== Paged bucket multi-bucket tests ======================

#[test]
fn test_paged_multi_bucket_projection() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Int64, true),
        ("c".to_string(), DataType::Float64, true),
        ("d".to_string(), DataType::Utf8, true),
        ("e".to_string(), DataType::Int32, true),
        ("f".to_string(), DataType::Int64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..200)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::BigInt(i as i64 * 2),
                Value::Double(i as f64 * 0.5),
                Value::String(format!("v{}", i).into_bytes()),
                Value::Integer(i * 10),
                Value::BigInt(i as i64 * 100),
            ]
        })
        .collect();

    // Use 3 buckets to spread columns across multiple paged buckets
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            num_buckets: 3,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    // Project only "a" and "f" — should hit different buckets
    let col_a = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "a")
        .unwrap();
    let col_f = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "f")
        .unwrap();

    let mut rg = reader
        .row_group_reader_projected(0, &[col_a, col_f])
        .unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 200);
    assert_eq!(batch.num_columns(), 2);

    let ca = batch_col_i32(&batch, "a");
    let cf = batch_col_i64(&batch, "f");
    for i in 0..200usize {
        assert_eq!(ca.value(i), i as i32);
        assert_eq!(cf.value(i), i as i64 * 100);
    }
}

#[test]
fn test_paged_wide_table_sparse_projection() {
    // 20 columns, project only 2 — tests that skipped columns aren't read
    let num_cols = 20;
    let columns: Vec<(String, DataType, bool)> = (0..num_cols)
        .map(|i| (format!("col_{}", i), DataType::Int64, true))
        .collect();
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|r| {
            (0..num_cols)
                .map(|c| Value::BigInt((r * num_cols + c) as i64))
                .collect()
        })
        .collect();

    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            num_buckets: 4,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    // Project col_3 and col_17
    let idx3 = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "col_3")
        .unwrap();
    let idx17 = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "col_17")
        .unwrap();

    let mut rg = reader
        .row_group_reader_projected(0, &[idx3, idx17])
        .unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);
    assert_eq!(batch.num_columns(), 2);

    let c3 = batch_col_i64(&batch, "col_3");
    let c17 = batch_col_i64(&batch, "col_17");
    for r in 0..100usize {
        assert_eq!(c3.value(r), (r * num_cols + 3) as i64);
        assert_eq!(c17.value(r), (r * num_cols + 17) as i64);
    }
}

#[test]
fn test_mixed_paged_and_monolithic_buckets() {
    // Create a scenario where some buckets are paged and others are monolithic.
    // Use page_size_threshold high enough that small buckets stay monolithic.
    let columns = vec![
        ("tiny".to_string(), DataType::Int32, true),
        ("big_a".to_string(), DataType::Int64, true),
        ("big_b".to_string(), DataType::Float64, true),
        ("big_c".to_string(), DataType::Utf8, true),
    ];
    let rows: Vec<Vec<Value>> = (0..500)
        .map(|i| {
            vec![
                Value::Integer(i),
                Value::BigInt(i as i64 * 7),
                Value::Double(i as f64 / 3.0),
                Value::String(format!("str_{:05}", i).into_bytes()),
            ]
        })
        .collect();

    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 4096,
            num_buckets: 4,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    // Read all columns
    let mut rg = reader.row_group_reader(0).unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 500);

    let col_tiny = batch_col_i32(&batch, "tiny");
    let col_a = batch_col_i64(&batch, "big_a");
    let col_b = batch_col_f64(&batch, "big_b");
    let col_c = batch_col_string(&batch, "big_c");
    for i in 0..500usize {
        assert_eq!(col_tiny.value(i), i as i32);
        assert_eq!(col_a.value(i), i as i64 * 7);
        assert!((col_b.value(i) - i as f64 / 3.0).abs() < 1e-10);
        assert_eq!(col_c.value(i), format!("str_{:05}", i));
    }

    // Project only "tiny" and "big_c"
    let idx_tiny = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "tiny")
        .unwrap();
    let idx_c = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "big_c")
        .unwrap();
    let mut rg2 = reader
        .row_group_reader_projected(0, &[idx_tiny, idx_c])
        .unwrap();
    let batch2 = rg2.read_columns().unwrap();
    assert_eq!(batch2.num_rows(), 500);
    assert_eq!(batch2.num_columns(), 2);

    let col_tiny2 = batch_col_i32(&batch2, "tiny");
    let col_c2 = batch_col_string(&batch2, "big_c");
    for i in 0..500usize {
        assert_eq!(col_tiny2.value(i), i as i32);
        assert_eq!(col_c2.value(i), format!("str_{:05}", i));
    }
}

#[test]
fn test_paged_all_null_columns_in_projection() {
    let columns = vec![
        ("real".to_string(), DataType::Int32, true),
        ("null_col".to_string(), DataType::Int64, true),
        ("also_real".to_string(), DataType::Float64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| vec![Value::Integer(i), Value::Null, Value::Double(i as f64)])
        .collect();

    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    // Project including the ALL_NULL column
    let idx_null = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "null_col")
        .unwrap();
    let idx_real = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "also_real")
        .unwrap();
    let mut rg = reader
        .row_group_reader_projected(0, &[idx_null, idx_real])
        .unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);

    let null_col = batch_col_i64(&batch, "null_col");
    let real_col = batch_col_f64(&batch, "also_real");
    for i in 0..100usize {
        assert!(null_col.is_null(i));
        assert!((real_col.value(i) - i as f64).abs() < 1e-10);
    }
}

#[test]
fn test_paged_all_null_only_projection() {
    let columns = vec![
        ("real".to_string(), DataType::Int32, true),
        ("null_col".to_string(), DataType::Int64, true),
    ];
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| vec![Value::Integer(i), Value::Null])
        .collect();

    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    let idx_null = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "null_col")
        .unwrap();
    let mut rg = reader.row_group_reader_projected(0, &[idx_null]).unwrap();
    let batch = rg.read_columns().unwrap();

    assert_eq!(batch.num_rows(), 100);
    assert_eq!(batch.num_columns(), 1);

    let null_col = batch_col_i64(&batch, "null_col");
    for i in 0..100usize {
        assert!(null_col.is_null(i));
    }
}

#[test]
fn test_paged_adjacent_columns_coalesced_read() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct CountingInputFile {
        data: Vec<u8>,
        read_count: AtomicUsize,
    }

    impl InputFile for CountingInputFile {
        fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
            self.read_count.fetch_add(1, Ordering::Relaxed);
            let start = offset as usize;
            let end = start + buf.len();
            if end > self.data.len() {
                return Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "read past end",
                ));
            }
            buf.copy_from_slice(&self.data[start..end]);
            Ok(())
        }
    }

    // 10 columns in 1 bucket, project col 2,3,4 (adjacent)
    let columns: Vec<(String, DataType, bool)> = (0..10)
        .map(|i| (format!("c{}", i), DataType::Int64, true))
        .collect();
    let rows: Vec<Vec<Value>> = (0..100)
        .map(|r| {
            (0..10)
                .map(|c| Value::BigInt((r * 10 + c) as i64))
                .collect()
        })
        .collect();

    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_ZSTD,
            page_size_threshold: 1,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;

    let input = CountingInputFile {
        data: data.clone(),
        read_count: AtomicUsize::new(0),
    };
    let reader = MosaicReader::new(input, len).unwrap();
    let open_reads = reader.input().read_count.load(Ordering::Relaxed);

    let idx2 = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c2")
        .unwrap();
    let idx3 = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c3")
        .unwrap();
    let idx4 = reader
        .schema()
        .columns
        .iter()
        .position(|c| c.name == "c4")
        .unwrap();

    let mut rg = reader
        .row_group_reader_projected(0, &[idx2, idx3, idx4])
        .unwrap();
    let batch = rg.read_columns().unwrap();
    assert_eq!(batch.num_rows(), 100);
    assert_eq!(batch.num_columns(), 3);

    let c2 = batch_col_i64(&batch, "c2");
    let c3 = batch_col_i64(&batch, "c3");
    let c4 = batch_col_i64(&batch, "c4");
    for r in 0..100usize {
        assert_eq!(c2.value(r), (r * 10 + 2) as i64);
        assert_eq!(c3.value(r), (r * 10 + 3) as i64);
        assert_eq!(c4.value(r), (r * 10 + 4) as i64);
    }

    // After open: 1 tail-prefetch IO covers footer + schema + index
    // projected read should do: round1 (directory) + round2 (3 adjacent slots coalesced)
    // = 2 read_ranges calls, each coalesced into 1 read_at
    let projection_reads = reader.input().read_count.load(Ordering::Relaxed) - open_reads;
    assert_eq!(
        projection_reads, 2,
        "expected 2 read_at calls (dir + coalesced slots)"
    );
}

#[test]
fn test_file_open_single_io() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    struct CountingInputFile {
        data: Vec<u8>,
        read_count: AtomicUsize,
    }

    impl InputFile for CountingInputFile {
        fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
            self.read_count.fetch_add(1, Ordering::Relaxed);
            let start = offset as usize;
            let end = start + buf.len();
            if end > self.data.len() {
                return Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "read past end",
                ));
            }
            buf.copy_from_slice(&self.data[start..end]);
            Ok(())
        }
    }

    // Write a file with modest metadata (fits within 256KB tail)
    let columns: Vec<(String, DataType, bool)> = vec![
        ("a".into(), DataType::Int64, false),
        ("b".into(), DataType::Int64, false),
        ("c".into(), DataType::Utf8, true),
    ];
    let rows: Vec<Vec<Value>> = (0..10)
        .map(|r| {
            vec![
                Value::BigInt(r as i64),
                Value::BigInt(r as i64 * 2),
                Value::String(format!("row{}", r).into_bytes()),
            ]
        })
        .collect();

    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions::default(),
    )
    .unwrap();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let file_data = writer.output().buf.clone();
    let file_len = file_data.len() as u64;

    let input = CountingInputFile {
        data: file_data,
        read_count: AtomicUsize::new(0),
    };

    // Open the file — should require only 1 read_at (tail prefetch covers all metadata)
    let reader = MosaicReader::new(input, file_len).unwrap();
    assert_eq!(
        reader.input().read_count.load(Ordering::Relaxed),
        1,
        "file open should use only 1 IO when metadata fits in tail prefetch"
    );
}

#[test]
fn test_writer_stats_basic() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("name".to_string(), DataType::Utf8, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["id".to_string(), "score".to_string()],
            num_buckets: 2,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..100)
        .map(|i| {
            vec![
                Value::Integer(i * 2),
                Value::String(format!("row_{}", i).into_bytes()),
                Value::Double(i as f64 * 0.5),
            ]
        })
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();

    assert_eq!(writer.num_row_groups(), 1);
    let stats = writer.row_group_stats(0);
    assert_eq!(stats.len(), 2);

    assert_eq!(stats[0].column_index, 0);
    assert_eq!(stats[0].null_count, 0);
    assert!(matches!(&stats[0].min, Some(Value::Integer(0))));
    assert!(matches!(&stats[0].max, Some(Value::Integer(198))));

    assert_eq!(stats[1].column_index, 2);
    assert_eq!(stats[1].null_count, 0);
    match &stats[1].min {
        Some(Value::Double(v)) => assert!(v.abs() < 1e-10),
        other => panic!("expected Double(0.0), got {:?}", other),
    }
    match &stats[1].max {
        Some(Value::Double(v)) => assert!((*v - 49.5).abs() < 1e-10),
        other => panic!("expected Double(49.5), got {:?}", other),
    }
}

#[test]
fn test_writer_stats_with_nulls() {
    let columns = vec![
        ("a".to_string(), DataType::Int32, true),
        ("b".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["a".to_string(), "b".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows = vec![
        vec![Value::Integer(10), Value::Null],
        vec![Value::Null, Value::Null],
        vec![Value::Integer(5), Value::BigInt(100)],
        vec![Value::Integer(20), Value::BigInt(50)],
    ];
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();

    assert_eq!(writer.num_row_groups(), 1);
    let stats = writer.row_group_stats(0);
    assert_eq!(stats.len(), 2);

    assert_eq!(stats[0].column_index, 0);
    assert_eq!(stats[0].null_count, 1);
    assert!(matches!(&stats[0].min, Some(Value::Integer(5))));
    assert!(matches!(&stats[0].max, Some(Value::Integer(20))));

    assert_eq!(stats[1].column_index, 1);
    assert_eq!(stats[1].null_count, 2);
    assert!(matches!(&stats[1].min, Some(Value::BigInt(50))));
    assert!(matches!(&stats[1].max, Some(Value::BigInt(100))));
}

#[test]
fn test_writer_stats_all_null() {
    let columns = vec![("x".to_string(), DataType::Int32, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["x".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..10).map(|_| vec![Value::Null]).collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();

    assert_eq!(writer.num_row_groups(), 1);
    let stats = writer.row_group_stats(0);
    assert_eq!(stats.len(), 1);
    assert_eq!(stats[0].null_count, 10);
    assert!(stats[0].min.is_none());
    assert!(stats[0].max.is_none());
}

#[test]
fn test_writer_stats_matches_reader_stats() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, true),
        ("score".to_string(), DataType::Float64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            stats_columns: vec!["id".to_string(), "score".to_string()],
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let rows: Vec<Vec<Value>> = (0..50)
        .map(|i| vec![Value::Integer(i), Value::Double(i as f64 * 2.0)])
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();

    let writer_stats = writer.row_group_stats(0);

    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let reader_stats = reader.row_group_stats(0).unwrap();

    assert_eq!(writer_stats.len(), reader_stats.len());
    for (ws, rs) in writer_stats.iter().zip(reader_stats.iter()) {
        assert_eq!(ws.column_index, rs.column_index);
        assert_eq!(ws.null_count, rs.null_count);
        assert_eq!(format!("{:?}", ws.min), format!("{:?}", rs.min));
        assert_eq!(format!("{:?}", ws.max), format!("{:?}", rs.max));
    }
}

#[test]
fn test_row_group_num_rows() {
    let columns = vec![
        ("id".to_string(), DataType::Int32, false),
        ("val".to_string(), DataType::Int64, true),
    ];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            compression: COMPRESSION_NONE,
            row_group_max_size: 200,
            num_buckets: 1,
            ..Default::default()
        },
    )
    .unwrap();

    let total_rows = 500;
    let rows: Vec<Vec<Value>> = (0..total_rows)
        .map(|i| vec![Value::Integer(i), Value::BigInt(i as i64 * 2)])
        .collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert!(reader.num_row_groups() > 1);

    let mut total = 0usize;
    for rg_idx in 0..reader.num_row_groups() {
        let num_rows = reader.row_group_num_rows(rg_idx).unwrap();
        assert!(num_rows > 0);
        let mut rg = reader.row_group_reader(rg_idx).unwrap();
        let batch = rg.read_columns().unwrap();
        assert_eq!(num_rows, batch.num_rows());
        total += num_rows;
    }
    assert_eq!(total, total_rows as usize);
}

#[test]
fn test_stats_empty_string_min() {
    let columns = vec![("s".to_string(), DataType::Utf8, true)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions {
            stats_columns: vec!["s".to_string()],
            ..Default::default()
        },
    )
    .unwrap();
    let rows: Vec<Vec<Value>> = vec![
        vec![Value::String(b"".to_vec())],
        vec![Value::String(b"b".to_vec())],
    ];
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();

    let stats = writer.row_group_stats(0);
    assert_eq!(stats.len(), 1);
    assert!(stats[0].min.is_some());
    assert!(stats[0].max.is_some());
    assert_eq!(stats[0].min.as_ref().unwrap().to_be_bytes(), b"");
    assert_eq!(stats[0].max.as_ref().unwrap().to_be_bytes(), b"b");
    assert_eq!(stats[0].null_count, 0);

    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();
    let reader_stats = reader.row_group_stats(0).unwrap();
    assert_eq!(reader_stats.len(), 1);
    assert!(reader_stats[0].min.is_some());
    assert!(reader_stats[0].max.is_some());
}

#[test]
fn test_row_group_num_rows_out_of_range() {
    let columns = vec![("x".to_string(), DataType::Int32, false)];
    let out = MemOutputFile::new();
    let mut writer = MosaicWriter::new(
        out,
        &columns_to_arrow_schema(&columns),
        WriterOptions::default(),
    )
    .unwrap();
    let rows: Vec<Vec<Value>> = (0..10).map(|i| vec![Value::Integer(i)]).collect();
    write_values(&mut writer, &columns, &rows);
    writer.close().unwrap();
    let data = writer.output().buf.clone();
    let len = data.len() as u64;
    let reader = MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap();

    assert_eq!(reader.num_row_groups(), 1);
    assert_eq!(reader.row_group_num_rows(0).unwrap(), 10);
    assert!(reader.row_group_num_rows(1).is_err());
    assert!(reader.row_group_num_rows(999).is_err());
}
