vortex-file 0.87.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

#![expect(clippy::cast_possible_truncation)]
use std::iter;
use std::sync::Arc;
use std::sync::LazyLock;

use bytes::Bytes;
use flatbuffers::FlatBufferBuilder;
use futures::StreamExt;
use futures::TryStreamExt;
use futures::pin_mut;
use rstest::rstest;
use vortex_array::ArrayRef;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::array_session;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::ChunkedArray;
use vortex_array::arrays::ConstantArray;
use vortex_array::arrays::DecimalArray;
use vortex_array::arrays::Dict;
use vortex_array::arrays::ListArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::StructArray;
use vortex_array::arrays::TemporalArray;
use vortex_array::arrays::VarBinArray;
use vortex_array::arrays::VarBinViewArray;
use vortex_array::arrays::dict::DictArraySlotsExt;
use vortex_array::arrays::struct_::StructArrayExt;
use vortex_array::assert_arrays_eq;
use vortex_array::builders::MapBuilder;
use vortex_array::dtype::DType;
use vortex_array::dtype::DecimalDType;
use vortex_array::dtype::MapDType;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::PType;
use vortex_array::dtype::PType::I32;
use vortex_array::dtype::StructFields;
use vortex_array::dtype::i256;
use vortex_array::expr::BoundExpression;
use vortex_array::expr::Expression;
use vortex_array::expr::and;
use vortex_array::expr::cast;
use vortex_array::expr::col;
use vortex_array::expr::eq;
use vortex_array::expr::get_item;
use vortex_array::expr::gt;
use vortex_array::expr::gt_eq;
use vortex_array::expr::lit;
use vortex_array::expr::lt;
use vortex_array::expr::lt_eq;
use vortex_array::expr::or;
use vortex_array::expr::root;
use vortex_array::expr::select;
use vortex_array::extension::datetime::TimeUnit;
use vortex_array::extension::datetime::Timestamp;
use vortex_array::extension::datetime::TimestampOptions;
use vortex_array::field_path;
use vortex_array::scalar::Scalar;
use vortex_array::scalar_fn::ScalarFnVTableExt;
use vortex_array::scalar_fn::fns::pack::Pack;
use vortex_array::scalar_fn::fns::pack::PackOptions;
use vortex_array::stats::PRUNING_STATS;
use vortex_array::stream::ArrayStreamAdapter;
use vortex_array::stream::ArrayStreamExt;
use vortex_array::validity::Validity;
use vortex_btrblocks::BtrBlocksCompressorBuilder;
use vortex_btrblocks::SchemeExt;
use vortex_btrblocks::schemes::string::StringDictScheme;
use vortex_buffer::Buffer;
use vortex_buffer::ByteBuffer;
use vortex_buffer::ByteBufferMut;
use vortex_buffer::buffer;
use vortex_decimal_byte_parts::DecimalByteParts;
use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt;
use vortex_edition::EDITION_DECLARATIONS;
use vortex_edition::EditionSession;
use vortex_edition::EditionSessionExt;
use vortex_edition::declarations::core::CORE_2026_08_3;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_io::session::RuntimeSession;
use vortex_layout::DynLayout;
use vortex_layout::LayoutStrategy;
use vortex_layout::layouts::buffered::BufferedStrategy;
use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy;
use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
use vortex_layout::layouts::struct_::StructStrategy;
use vortex_layout::layouts::table::TableStrategy;
use vortex_layout::layouts::zoned::LegacyStats;
use vortex_layout::layouts::zoned::Zoned;
use vortex_layout::scan::scan_builder::ScanBuilder;
use vortex_layout::scan::split_by::DEFAULT_MAX_SPLIT_ROWS;
use vortex_layout::scan::split_by::SplitBy;
use vortex_layout::session::LayoutSession;
use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
use vortex_session::VortexSession;
use vortex_zigzag::ZigZag;

use crate::MAX_POSTSCRIPT_SIZE;
use crate::OpenOptionsSessionExt;
use crate::V1_FOOTER_FBS_SIZE;
use crate::VERSION;
use crate::VortexFile;
use crate::WriteOptionsSessionExt;
use crate::flatbuffers::footer as fb;
use crate::footer::SegmentSpec;
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
    let session = array_session()
        .with::<LayoutSession>()
        .with::<RuntimeSession>();

    crate::register_default_encodings(&session);
    crate::enable_all_registered_array_encodings(&session);

    session
});

fn strict_sorted(indices: Buffer<u64>) -> StrictSortedBuffer<u64> {
    StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing")
}

fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression {
    expr.bind(file.dtype())
        .and_then(|expr| expr.optimize_recursive())
        .vortex_expect("scan expression should bind")
}

#[tokio::test]
async fn test_eof_values() {
    // this test exists as a reminder to think about whether we should increment the version
    // when we change the footer
    assert_eq!(VERSION, 1);
    assert_eq!(V1_FOOTER_FBS_SIZE, 32);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_read_simple() {
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
    ])
    .into_array();

    let numbers = ChunkedArray::from_iter([
        buffer![1u32, 2, 3, 4].into_array(),
        buffer![5u32, 6, 7, 8].into_array(),
    ])
    .into_array();

    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)]).unwrap();
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let stream = SESSION
        .open_options()
        .open_buffer(buf)
        .unwrap()
        .scan()
        .unwrap()
        .into_array_stream()
        .unwrap();
    pin_mut!(stream);

    let mut row_count = 0;

    while let Some(array) = stream.next().await {
        let array = array.unwrap();
        row_count += array.len();
    }

    assert_eq!(row_count, 8);
}

/// Wide decimals split into multi-part arrays only when the writer allows the v2 format. The
/// default writer allows every registered encoding once editions are disabled; a strategy built
/// from the session keeps the enabled editions' restrictions regardless.
#[rstest]
#[case::default_writer(false, false)]
#[case::custom_layout(true, false)]
#[case::explicit_compressor(true, true)]
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn decimal_writer_uses_default_or_session_restrictions(
    #[values(false, true)] use_i256: bool,
    #[case] custom_strategy: bool,
    #[case] explicit_compressor: bool,
) -> VortexResult<()> {
    let session = array_session()
        .with::<EditionSession>()
        .with::<LayoutSession>()
        .with::<RuntimeSession>();
    crate::register_default_encodings(&session);
    for declaration in EDITION_DECLARATIONS {
        session.register_edition(declaration)?;
    }
    session.enable_edition(CORE_2026_08_3)?;

    let array = if use_i256 {
        DecimalArray::new(
            (0..1024u128)
                .map(|i| i256::from_parts(i * 17, 1i128 << 70))
                .collect::<Buffer<i256>>(),
            DecimalDType::new(76, 2),
            Validity::NonNullable,
        )
    } else {
        DecimalArray::new(
            (0..1024i128)
                .map(|i| (1i128 << 70) + i * 17)
                .collect::<Buffer<i128>>(),
            DecimalDType::new(38, 2),
            Validity::NonNullable,
        )
    }
    .into_array();
    let strategy = crate::strategy::WriteStrategyBuilder::from_session(&session)
        .with_row_block_size(256)
        .with_data_block_target_bytes(None);
    let strategy = if explicit_compressor {
        strategy.with_btrblocks_builder(BtrBlocksCompressorBuilder::from_session(&session))
    } else {
        strategy
    }
    .build();

    for disable_editions in [false, true] {
        let mut options = session.write_options();
        if custom_strategy {
            options = options.with_strategy(Arc::clone(&strategy));
        }
        let options = if disable_editions {
            options.disable_editions()
        } else {
            options
        };
        let mut buffer = ByteBufferMut::empty();
        options
            .write(&mut buffer, array.clone().to_array_stream())
            .await?;
        let actual = session
            .open_options()
            .open_buffer(buffer)?
            .scan()?
            .into_array_stream()?
            .read_all()
            .await?;
        let uses_v2 = actual.depth_first_traversal().any(|array| {
            array
                .as_opt::<DecimalByteParts>()
                .is_some_and(|parts| !parts.lower_parts().is_empty())
        });
        assert_eq!(uses_v2, disable_editions && !custom_strategy);
        assert_arrays_eq!(array, actual, &mut session.create_execution_ctx());
    }
    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_round_trip_many_types() {
    let strings = VarBinArray::from(vec!["ab", "foo", "bar"]).into_array();

    let numbers = buffer![1u32, 2, 3].into_array();

    let decimal_2 = DecimalArray::new(
        buffer![100i8, 10i8, 2i8],
        DecimalDType::new(2, 1),
        Validity::from_iter([false, true, false]),
    )
    .into_array();

    let decimal_4 = DecimalArray::new(
        buffer![100i16, 10i16, 2i16],
        DecimalDType::new(4, 2),
        Validity::from_iter([false, true, false]),
    )
    .into_array();

    let decimal_9 = DecimalArray::new(
        buffer![100i32, 10i32, 2i32],
        DecimalDType::new(9, 2),
        Validity::from_iter([false, true, false]),
    )
    .into_array();

    let decimal_17 = DecimalArray::new(
        buffer![100i64, 10i64, 20234i64],
        DecimalDType::new(17, 2),
        Validity::from_iter([false, true, false]),
    )
    .into_array();

    let decimal_35 = DecimalArray::new(
        buffer![100i128, 139348340i128, 23943942i128],
        DecimalDType::new(35, 2),
        Validity::from_iter([true, false, false]),
    )
    .into_array();

    let st = StructArray::from_fields(&[
        ("strings", strings),
        ("numbers", numbers),
        ("decimal_2", decimal_2),
        ("decimal_4", decimal_4),
        ("decimal_9", decimal_9),
        ("decimal_17", decimal_17),
        ("decimal_35", decimal_35),
    ])
    .unwrap();
    let dtype = st.dtype().clone();
    let mut buf = ByteBufferMut::empty();

    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let chunks: Vec<_> = SESSION
        .open_options()
        .open_buffer(buf)
        .unwrap()
        .scan()
        .unwrap()
        .into_array_stream()
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    let read = ChunkedArray::try_new(chunks, dtype).unwrap();

    assert_eq!(read.len(), 3);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_read_simple_with_spawn() {
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
    ])
    .into_array();

    let numbers = ChunkedArray::from_iter([
        buffer![1u32, 2, 3, 4].into_array(),
        buffer![5u32, 6, 7, 8].into_array(),
    ])
    .into_array();

    let lists = ChunkedArray::from_iter([
        ListArray::from_iter_slow::<i32, _>(
            vec![vec![11, 12], vec![21, 22], vec![31, 32], vec![41, 42]],
            Arc::new(I32.into()),
        )
        .unwrap()
        .into_array(),
        ListArray::from_iter_slow::<i64, _>(
            vec![vec![51, 52], vec![61, 62], vec![71, 72], vec![81, 82]],
            Arc::new(I32.into()),
        )
        .unwrap()
        .into_array(),
    ])
    .into_array();

    let st =
        StructArray::from_fields(&[("strings", strings), ("numbers", numbers), ("lists", lists)])
            .unwrap();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    assert!(!buf.is_empty());
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_read_projection() {
    let mut ctx = SESSION.create_execution_ctx();
    let strings_expected = ["ab", "foo", "bar", "baz", "ab", "foo", "bar", "baz"];
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(strings_expected[..4].to_vec()).into_array(),
        VarBinArray::from(strings_expected[4..].to_vec()).into_array(),
    ])
    .into_array();
    let strings_dtype = strings.dtype().clone();

    let numbers_expected = [1u32, 2, 3, 4, 5, 6, 7, 8];
    let numbers = ChunkedArray::from_iter([
        Buffer::copy_from(&numbers_expected[..4]).into_array(),
        Buffer::copy_from(&numbers_expected[4..]).into_array(),
    ])
    .into_array();
    let numbers_dtype = numbers.dtype().clone();

    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)]).unwrap();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();
    let array = file
        .scan()
        .unwrap()
        .with_projection(bind_scan_expr(&file, select(["strings"], root())))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap();

    assert_eq!(
        array.dtype(),
        &DType::Struct(
            StructFields::new(["strings"].into(), vec![strings_dtype]),
            Nullability::NonNullable,
        )
    );

    let actual = array
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(0)
        .clone();
    let expected = VarBinArray::from(strings_expected.to_vec()).into_array();
    assert_arrays_eq!(actual, expected, &mut ctx);

    let array = file
        .scan()
        .unwrap()
        .with_projection(bind_scan_expr(&file, select(["numbers"], root())))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap();

    assert_eq!(
        array.dtype(),
        &DType::Struct(
            StructFields::new(["numbers"].into(), vec![numbers_dtype]),
            Nullability::NonNullable,
        )
    );

    let actual = array
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(0)
        .clone();
    let expected = Buffer::copy_from(numbers_expected).into_array();
    assert_arrays_eq!(actual, expected, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn unequal_batches() {
    let mut ctx = SESSION.create_execution_ctx();
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(vec!["ab", "foo", "bar", "bob"]).into_array(),
        VarBinArray::from(vec!["baz", "ab", "foo", "bar", "baz", "alice"]).into_array(),
    ])
    .into_array();

    let numbers = ChunkedArray::from_iter([
        buffer![1u32, 2, 3, 4, 5].into_array(),
        buffer![6u32, 7, 8, 9, 10].into_array(),
    ])
    .into_array();

    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)]).unwrap();
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let stream = SESSION
        .open_options()
        .open_buffer(buf)
        .unwrap()
        .scan()
        .unwrap()
        .into_array_stream()
        .unwrap();
    pin_mut!(stream);

    let mut item_count = 0;

    while let Some(array) = stream.next().await {
        let array = array.unwrap();
        item_count += array.len();

        let numbers = array
            .execute::<StructArray>(&mut ctx)
            .unwrap()
            .unmasked_field_by_name("numbers")
            .unwrap()
            .clone()
            .execute::<PrimitiveArray>(&mut ctx)
            .unwrap();
        assert_eq!(numbers.ptype(), PType::U32);
    }
    assert_eq!(item_count, 10);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn write_chunked() {
    let strings = VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array();
    let string_dtype = strings.dtype().clone();
    let strings_chunked = ChunkedArray::try_new(iter::repeat_n(strings, 4), string_dtype)
        .unwrap()
        .into_array();
    let numbers = buffer![1u32, 2, 3, 4].into_array();
    let numbers_dtype = numbers.dtype().clone();
    let numbers_chunked = ChunkedArray::try_new(iter::repeat_n(numbers, 4), numbers_dtype)
        .unwrap()
        .into_array();
    let st = StructArray::try_new(
        ["strings", "numbers"].into(),
        vec![strings_chunked, numbers_chunked],
        16,
        Validity::NonNullable,
    )
    .unwrap()
    .into_array();
    let st_dtype = st.dtype().clone();

    let chunked_st = ChunkedArray::try_new(iter::repeat_n(st, 3), st_dtype)
        .unwrap()
        .into_array();
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, chunked_st.into_array().to_array_stream())
        .await
        .unwrap();

    let stream = SESSION
        .open_options()
        .open_buffer(buf)
        .unwrap()
        .scan()
        .unwrap()
        .into_array_stream()
        .unwrap();
    pin_mut!(stream);

    let mut array_len: usize = 0;
    while let Some(array) = stream.next().await {
        array_len += array.unwrap().len();
    }
    assert_eq!(array_len, 48);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_empty_varbin_array_roundtrip() {
    let empty = VarBinArray::from(Vec::<&str>::new()).into_array();

    let st = StructArray::from_fields(&[("a", empty)]).unwrap();
    let dtype = st.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    let result = file
        .scan()
        .unwrap()
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap();

    assert_eq!(result.len(), 0);
    assert_eq!(result.dtype(), &dtype);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn issue_5385_filter_casted_column() {
    let array = StructArray::try_from_iter([("x", buffer![1u8, 2, 3, 4, 5])])
        .unwrap()
        .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, array.to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();
    let result = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            eq(
                cast(
                    get_item("x", root()),
                    DType::Primitive(PType::U16, Nullability::NonNullable),
                ),
                lit(1u16),
            ),
        ))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap();

    assert_arrays_eq!(
        result,
        StructArray::try_from_iter([("x", buffer![1u8])]).unwrap(),
        &mut SESSION.create_execution_ctx()
    );
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn filter_string() {
    let mut ctx = SESSION.create_execution_ctx();
    let names_orig = VarBinArray::from_iter(
        vec![Some("Joseph"), None, Some("Angela"), Some("Mikhail"), None],
        DType::Utf8(Nullability::Nullable),
    )
    .into_array();
    let ages_orig =
        PrimitiveArray::from_option_iter([Some(25), Some(31), None, Some(57), None]).into_array();
    let st = StructArray::try_new(
        ["name", "age"].into(),
        vec![names_orig, ages_orig],
        5,
        Validity::NonNullable,
    )
    .unwrap()
    .into_array();
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();
    let result: Vec<_> = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            eq(get_item("name", root()), lit("Joseph")),
        ))
        .into_array_stream()
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(result.len(), 1);
    let names_actual = result[0]
        .clone()
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(0)
        .clone();
    let names_expected =
        VarBinArray::from_iter(vec![Some("Joseph")], DType::Utf8(Nullability::Nullable))
            .into_array();
    assert_arrays_eq!(names_actual, names_expected, &mut ctx);

    let ages_actual = result[0]
        .clone()
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(1)
        .clone();
    let ages_expected = PrimitiveArray::from_option_iter([Some(25i32)]).into_array();
    assert_arrays_eq!(ages_actual, ages_expected, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn filter_or() {
    let mut ctx = SESSION.create_execution_ctx();
    let names = VarBinArray::from_iter(
        vec![Some("Joseph"), None, Some("Angela"), Some("Mikhail"), None],
        DType::Utf8(Nullability::Nullable),
    );
    let ages = PrimitiveArray::from_option_iter([Some(25), Some(31), None, Some(57), None]);
    let st = StructArray::try_new(
        ["name", "age"].into(),
        vec![names.into_array(), ages.into_array()],
        5,
        Validity::NonNullable,
    )
    .unwrap()
    .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();
    let result: Vec<_> = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            or(
                eq(get_item("name", root()), lit("Angela")),
                and(
                    gt_eq(get_item("age", root()), lit(20)),
                    lt_eq(get_item("age", root()), lit(30)),
                ),
            ),
        ))
        .into_array_stream()
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(result.len(), 1);
    let names_actual = result[0]
        .clone()
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(0)
        .clone();
    let names_expected = VarBinArray::from_iter(
        vec![Some("Joseph"), Some("Angela")],
        DType::Utf8(Nullability::Nullable),
    )
    .into_array();
    assert_arrays_eq!(names_actual, names_expected, &mut ctx);

    let ages_actual = result[0]
        .clone()
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(1)
        .clone();
    let ages_expected = PrimitiveArray::from_option_iter([Some(25i32), None]).into_array();
    assert_arrays_eq!(ages_actual, ages_expected, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn filter_and() {
    let mut ctx = SESSION.create_execution_ctx();
    let names = VarBinArray::from_iter(
        vec![Some("Joseph"), None, Some("Angela"), Some("Mikhail"), None],
        DType::Utf8(Nullability::Nullable),
    );
    let ages = PrimitiveArray::from_option_iter([Some(25), Some(31), None, Some(57), None]);
    let st = StructArray::try_new(
        ["name", "age"].into(),
        vec![names.into_array(), ages.into_array()],
        5,
        Validity::NonNullable,
    )
    .unwrap()
    .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();
    let result: Vec<_> = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            and(
                gt(get_item("age", root()), lit(21)),
                lt_eq(get_item("age", root()), lit(33)),
            ),
        ))
        .into_array_stream()
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(result.len(), 1);
    let names_actual = result[0]
        .clone()
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(0)
        .clone();
    let names_expected = VarBinArray::from_iter(
        vec![Some("Joseph"), None],
        DType::Utf8(Nullability::Nullable),
    )
    .into_array();
    assert_arrays_eq!(names_actual, names_expected, &mut ctx);

    let ages_actual = result[0]
        .clone()
        .execute::<StructArray>(&mut ctx)
        .unwrap()
        .unmasked_field(1)
        .clone();
    let ages_expected = PrimitiveArray::from_option_iter([Some(25i32), Some(31i32)]).into_array();
    assert_arrays_eq!(ages_actual, ages_expected, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_with_indices_simple() {
    let mut ctx = SESSION.create_execution_ctx();
    let expected_numbers_split: Vec<Buffer<i16>> = (0..5).map(|_| (0_i16..100).collect()).collect();
    let expected_array = StructArray::from_fields(&[(
        "numbers",
        ChunkedArray::from_iter(
            expected_numbers_split
                .iter()
                .cloned()
                .map(IntoArray::into_array),
        )
        .into_array(),
    )])
    .unwrap();
    let expected_numbers: Vec<i16> = expected_numbers_split.into_iter().flatten().collect();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, expected_array.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    // test no indices
    let actual_kept_array = file
        .scan()
        .unwrap()
        .with_row_indices(strict_sorted(Buffer::<u64>::empty()))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    assert_eq!(actual_kept_array.len(), 0);

    // test a few indices
    let kept_indices = [0_u64, 3, 99, 100, 101, 399, 400, 401, 499];

    let actual_kept_array = file
        .scan()
        .unwrap()
        .with_row_indices(strict_sorted(Buffer::from_iter(kept_indices)))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();
    let actual_kept_numbers_array = actual_kept_array
        .unmasked_field(0)
        .clone()
        .execute::<PrimitiveArray>(&mut ctx)
        .unwrap();

    let expected_kept_numbers: Vec<i16> = kept_indices
        .iter()
        .map(|&x| expected_numbers[x as usize])
        .collect();
    let expected_array = Buffer::copy_from(&expected_kept_numbers).into_array();
    assert_arrays_eq!(actual_kept_numbers_array, expected_array, &mut ctx);

    // test all indices
    let actual_array = file
        .scan()
        .unwrap()
        .with_row_indices(strict_sorted((0u64..500).collect::<Buffer<_>>()))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();
    let actual_numbers_array = actual_array.unmasked_field(0).clone();
    let expected_array = Buffer::copy_from(&expected_numbers).into_array();
    assert_arrays_eq!(actual_numbers_array, expected_array, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_with_indices_on_two_columns() {
    let mut ctx = SESSION.create_execution_ctx();
    let strings_expected = ["ab", "foo", "bar", "baz", "ab", "foo", "bar", "baz"];
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(strings_expected[..4].to_vec()).into_array(),
        VarBinArray::from(strings_expected[4..].to_vec()).into_array(),
    ])
    .into_array();

    let numbers_expected = [1u32, 2, 3, 4, 5, 6, 7, 8];
    let numbers = ChunkedArray::from_iter([
        Buffer::copy_from(&numbers_expected[..4]).into_array(),
        Buffer::copy_from(&numbers_expected[4..]).into_array(),
    ])
    .into_array();

    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)]).unwrap();
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    let kept_indices = [0_u64, 3, 7];
    let array = file
        .scan()
        .unwrap()
        .with_row_indices(strict_sorted(Buffer::from_iter(kept_indices)))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    let strings_actual = array.unmasked_field(0).clone();
    let strings_expected_vec: Vec<&str> = kept_indices
        .iter()
        .map(|&x| strings_expected[x as usize])
        .collect();
    let strings_expected_array = VarBinArray::from(strings_expected_vec).into_array();
    assert_arrays_eq!(strings_actual, strings_expected_array, &mut ctx);

    let numbers_actual = array.unmasked_field(1).clone();
    let numbers_expected_vec: Vec<u32> = kept_indices
        .iter()
        .map(|&x| numbers_expected[x as usize])
        .collect();
    let numbers_expected_array = Buffer::copy_from(&numbers_expected_vec).into_array();
    assert_arrays_eq!(numbers_actual, numbers_expected_array, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_with_indices_and_with_row_filter_simple() {
    let mut ctx = SESSION.create_execution_ctx();
    let expected_numbers_split: Vec<Buffer<i16>> = (0..5).map(|_| (0_i16..100).collect()).collect();
    let expected_array = StructArray::from_fields(&[(
        "numbers",
        ChunkedArray::from_iter(
            expected_numbers_split
                .iter()
                .cloned()
                .map(IntoArray::into_array),
        )
        .into_array(),
    )])
    .unwrap();
    let expected_numbers: Vec<i16> = expected_numbers_split.into_iter().flatten().collect();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, expected_array.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    let actual_kept_array = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            gt(get_item("numbers", root()), lit(50_i16)),
        ))
        .with_row_indices(strict_sorted(Buffer::empty()))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    assert_eq!(actual_kept_array.len(), 0);

    // test a few indices
    let kept_indices = [0u64, 3, 99, 100, 101, 399, 400, 401, 499];

    let actual_kept_array = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            gt(get_item("numbers", root()), lit(50_i16)),
        ))
        .with_row_indices(strict_sorted(Buffer::from_iter(kept_indices)))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    let actual_kept_numbers_array = actual_kept_array
        .unmasked_field(0)
        .clone()
        .execute::<PrimitiveArray>(&mut ctx)
        .unwrap();

    let expected_kept_numbers: Buffer<i16> = kept_indices
        .iter()
        .map(|&x| expected_numbers[x as usize])
        .filter(|&x| x > 50)
        .collect();
    let expected_array = expected_kept_numbers.into_array();
    assert_arrays_eq!(actual_kept_numbers_array, expected_array, &mut ctx);

    // test all indices
    let actual_array = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            gt(get_item("numbers", root()), lit(50_i16)),
        ))
        .with_row_indices(strict_sorted((0..500).collect::<Buffer<_>>()))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    let actual_numbers_array = actual_array.unmasked_field(0).clone();
    let expected_filtered: Buffer<i16> = expected_numbers
        .iter()
        .filter(|&&x| x > 50)
        .cloned()
        .collect();
    let expected_numbers_array = expected_filtered.into_array();
    assert_arrays_eq!(actual_numbers_array, expected_numbers_array, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn filter_string_chunked() {
    let mut ctx = SESSION.create_execution_ctx();
    let name_chunk1 =
        VarBinViewArray::from_iter_nullable_str([Some("Joseph"), Some("James"), Some("Angela")])
            .into_array();
    let age_chunk1 = PrimitiveArray::from_option_iter([Some(25_i32), Some(31), None]).into_array();
    let name_chunk2 = VarBinViewArray::from_iter_nullable_str([
        Some("Pharrell".to_owned()),
        Some("Khalil".to_owned()),
        Some("Mikhail".to_owned()),
        None,
    ])
    .into_array();
    let age_chunk2 =
        PrimitiveArray::from_option_iter([Some(57_i32), Some(18), None, Some(32)]).into_array();

    let chunk1 = StructArray::from_fields(&[("name", name_chunk1), ("age", age_chunk1)])
        .unwrap()
        .into_array();
    let chunk2 = StructArray::from_fields(&[("name", name_chunk2), ("age", age_chunk2)])
        .unwrap()
        .into_array();
    let dtype = chunk1.dtype().clone();

    let array = ChunkedArray::try_new(vec![chunk1, chunk2], dtype)
        .unwrap()
        .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, array.to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    let actual_array = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            eq(get_item("name", root()), lit("Joseph")),
        ))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    assert_eq!(actual_array.len(), 1);
    let names_actual = actual_array.unmasked_field(0).clone();
    let names_expected =
        VarBinArray::from_iter(vec![Some("Joseph")], DType::Utf8(Nullability::Nullable))
            .into_array();
    assert_arrays_eq!(names_actual, names_expected, &mut ctx);

    let ages_actual = actual_array.unmasked_field(1).clone();
    let ages_expected = PrimitiveArray::from_option_iter([Some(25i32)]).into_array();
    assert_arrays_eq!(ages_actual, ages_expected, &mut ctx);
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_pruning_with_or() {
    let mut ctx = SESSION.create_execution_ctx();
    let letter_chunk1 = VarBinViewArray::from_iter_nullable_str([
        Some("A".to_owned()),
        Some("B".to_owned()),
        Some("D".to_owned()),
    ])
    .into_array();
    let number_chunk1 =
        PrimitiveArray::from_option_iter([Some(25_i32), Some(31), None]).into_array();
    let letter_chunk2 = VarBinViewArray::from_iter_nullable_str([
        Some("G".to_owned()),
        Some("I".to_owned()),
        Some("J".to_owned()),
        None,
    ])
    .into_array();
    let number_chunk2 =
        PrimitiveArray::from_option_iter([Some(4_i32), Some(18), None, Some(21)]).into_array();
    let letter_chunk3 = VarBinViewArray::from_iter_nullable_str([
        Some("L".to_owned()),
        None,
        Some("O".to_owned()),
        Some("P".to_owned()),
    ])
    .into_array();
    let number_chunk3 =
        PrimitiveArray::from_option_iter([Some(10_i32), Some(15), None, Some(22)]).into_array();
    let letter_chunk4 = VarBinViewArray::from_iter_nullable_str([
        Some("X".to_owned()),
        Some("Y".to_owned()),
        Some("Z".to_owned()),
    ])
    .into_array();
    let number_chunk4 =
        PrimitiveArray::from_option_iter([Some(66_i32), Some(77), Some(88)]).into_array();

    let chunk1 = StructArray::from_fields(&[("letter", letter_chunk1), ("number", number_chunk1)])
        .unwrap()
        .into_array();
    let chunk2 = StructArray::from_fields(&[("letter", letter_chunk2), ("number", number_chunk2)])
        .unwrap()
        .into_array();
    let chunk3 = StructArray::from_fields(&[("letter", letter_chunk3), ("number", number_chunk3)])
        .unwrap()
        .into_array();
    let chunk4 = StructArray::from_fields(&[("letter", letter_chunk4), ("number", number_chunk4)])
        .unwrap()
        .into_array();
    let dtype = chunk1.dtype().clone();

    let array = ChunkedArray::try_new(vec![chunk1, chunk2, chunk3, chunk4], dtype)
        .unwrap()
        .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, array.to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    let actual_array = file
        .scan()
        .unwrap()
        .with_filter(bind_scan_expr(
            &file,
            or(
                lt_eq(get_item("letter", root()), lit("J")),
                lt(get_item("number", root()), lit(25)),
            ),
        ))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    assert_eq!(actual_array.len(), 10);
    let letters_actual = actual_array.unmasked_field(0).clone();
    let letters_expected = VarBinViewArray::from_iter_nullable_str([
        Some("A".to_owned()),
        Some("B".to_owned()),
        Some("D".to_owned()),
        Some("G".to_owned()),
        Some("I".to_owned()),
        Some("J".to_owned()),
        None,
        Some("L".to_owned()),
        None,
        Some("P".to_owned()),
    ])
    .into_array();
    assert_arrays_eq!(letters_actual, letters_expected, &mut ctx);

    let numbers_actual = actual_array.unmasked_field(1).clone();
    let numbers_expected = PrimitiveArray::from_option_iter([
        Some(25_i32),
        Some(31),
        None,
        Some(4),
        Some(18),
        None,
        Some(21),
        Some(10),
        Some(15),
        Some(22),
    ])
    .into_array();
    assert_arrays_eq!(numbers_actual, numbers_expected, &mut ctx);
}

#[tokio::test]
async fn test_repeated_projection() {
    let mut ctx = SESSION.create_execution_ctx();
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
    ])
    .into_array();

    let single_column_array = StructArray::from_fields(&[("strings", strings.clone())])
        .unwrap()
        .into_array();

    let expected = StructArray::from_fields(&[("strings", strings.clone()), ("strings", strings)])
        .unwrap()
        .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, single_column_array.into_array().to_array_stream())
        .await
        .unwrap();

    let file = SESSION.open_options().open_buffer(buf).unwrap();

    let actual = file
        .scan()
        .unwrap()
        .with_projection(bind_scan_expr(
            &file,
            select(["strings", "strings"], root()),
        ))
        .into_array_stream()
        .unwrap()
        .read_all()
        .await
        .unwrap()
        .execute::<StructArray>(&mut ctx)
        .unwrap();

    assert_arrays_eq!(actual, expected, &mut ctx);
}

async fn chunked_file() -> VortexResult<VortexFile> {
    let array = ChunkedArray::from_iter([
        buffer![0, 1, 2].into_array(),
        buffer![3, 4, 5].into_array(),
        buffer![6, 7, 8].into_array(),
    ])
    .into_array();

    let mut writer = vec![];
    SESSION
        .write_options()
        .write(&mut writer, array.to_array_stream())
        .await?;
    let buffer: Bytes = writer.into();
    SESSION.open_options().open_buffer(buffer)
}

#[tokio::test]
async fn basic_file_roundtrip() -> VortexResult<()> {
    let vxf = chunked_file().await?;
    let result = vxf.scan()?.into_array_stream()?.read_all().await?;

    let expected = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8].into_array();
    assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());

    Ok(())
}

#[tokio::test]
async fn file_excluding_dtype() -> VortexResult<()> {
    let array = ChunkedArray::from_iter([
        buffer![0, 1, 2].into_array(),
        buffer![3, 4, 5].into_array(),
        buffer![6, 7, 8].into_array(),
    ])
    .into_array();
    let dtype = array.dtype().clone();

    let mut writer = vec![];
    SESSION
        .write_options()
        .exclude_dtype()
        .write(&mut writer, array.to_array_stream())
        .await?;
    let buffer: Bytes = writer.into();

    // Fail to open without DType.
    let vxf = SESSION.open_options().open_buffer(buffer.clone());
    assert!(vxf.is_err(), "Opening without DType should fail");

    let vxf = SESSION
        .open_options()
        .with_dtype(dtype.clone())
        .open_buffer(buffer)?;
    assert_eq!(vxf.dtype(), &dtype);
    assert_eq!(vxf.row_count(), 9);

    Ok(())
}

#[tokio::test]
async fn file_take() -> VortexResult<()> {
    let vxf = chunked_file().await?;
    let result = vxf
        .scan()?
        .with_row_indices(StrictSortedBuffer::try_new(buffer![0, 1, 8])?)
        .into_array_stream()?
        .read_all()
        .await?;

    let expected = buffer![0i32, 1, 8].into_array();
    assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());

    Ok(())
}

#[tokio::test]
#[should_panic(
    expected = "FileStatsAccumulator temporarily does not support nullable top-level structs"
)]
async fn write_nullable_top_level_struct() {
    let ages = PrimitiveArray::from_option_iter([Some(25), Some(31), None, Some(57), None]);

    let array = StructArray::try_new(
        ["age"].into(),
        vec![ages.into_array()],
        5,
        Validity::AllValid,
    )
    .unwrap()
    .into_array();

    let mut writer = vec![];
    SESSION
        .write_options()
        .write(&mut writer, array.to_array_stream())
        .await
        .unwrap();
}

async fn round_trip(
    array: &ArrayRef,
    f: impl FnOnce(ScanBuilder<ArrayRef>) -> VortexResult<ScanBuilder<ArrayRef>>,
) -> VortexResult<ArrayRef> {
    let mut writer = vec![];
    SESSION
        .write_options()
        .write(&mut writer, array.to_array_stream())
        .await?;
    let buffer: Bytes = writer.into();

    let vxf = SESSION
        .open_options()
        .with_dtype(array.dtype().clone())
        .open_buffer(buffer)?;

    assert_eq!(vxf.dtype(), array.dtype());
    assert_eq!(vxf.row_count(), array.len() as u64);

    f(vxf.scan()?)?.into_array_stream()?.read_all().await
}

#[tokio::test]
async fn write_nullable_nested_struct() -> VortexResult<()> {
    let nested_dtype = DType::struct_(
        [(
            "nested_field",
            DType::Primitive(PType::F16, Nullability::Nullable),
        )],
        Nullability::Nullable,
    );

    let struct_ = ConstantArray::new(Scalar::null(nested_dtype.clone()), 3).into_array();

    let array = StructArray::try_new(
        ["struct"].into(),
        vec![struct_.into_array()],
        3,
        Validity::NonNullable,
    )?
    .into_array();

    let mut ctx = SESSION.create_execution_ctx();
    let result = round_trip(&array, Ok)
        .await?
        .execute::<StructArray>(&mut ctx)?;

    assert_eq!(result.len(), 3);
    assert_eq!(result.struct_fields().nfields(), 1);
    assert!(result.all_valid(&mut ctx)?);

    let nested_struct = result
        .unmasked_field_by_name("struct")?
        .clone()
        .execute::<StructArray>(&mut ctx)?;
    assert_eq!(nested_struct.dtype(), &nested_dtype);
    assert_eq!(nested_struct.len(), 3);
    assert!(nested_struct.all_invalid(&mut ctx)?);

    Ok(())
}

#[tokio::test]
async fn scan_empty_fields() -> VortexResult<()> {
    let array = (0..10000).collect::<PrimitiveArray>();
    let projection = Pack
        .new_expr(
            PackOptions {
                names: Default::default(),
                nullability: Nullability::Nullable,
            },
            [],
        )
        .bind(array.dtype())?
        .optimize_recursive()?;

    let result = round_trip(&array.clone().into_array(), |scan| {
        Ok(scan.with_projection(projection))
    })
    .await?;

    assert_eq!(result.len(), array.len());

    Ok(())
}

#[tokio::test]
async fn test_into_tokio_array_stream() -> VortexResult<()> {
    let strings = ChunkedArray::from_iter([
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
        VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
    ])
    .into_array();

    let numbers = ChunkedArray::from_iter([
        buffer![1u32, 2, 3, 4].into_array(),
        buffer![5u32, 6, 7, 8].into_array(),
    ])
    .into_array();

    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)])?;
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await?;

    let file = SESSION.open_options().open_buffer(buf)?;
    let stream = file.scan().unwrap().into_array_stream()?;
    let array = stream.read_all().await?;

    assert_eq!(array.len(), 8);

    Ok(())
}

#[tokio::test]
async fn test_array_stream_no_double_dict_encode() -> VortexResult<()> {
    let num_vals = 2048;
    let mut values = Vec::<i64>::with_capacity(num_vals);
    values.extend(iter::repeat_n(0, num_vals / 2));
    values.extend(iter::repeat_n(1, num_vals / 2));

    let array = PrimitiveArray::from_iter(values).into_array();
    let mut buf = Vec::new();
    SESSION
        .write_options()
        .write(&mut buf, array.to_array_stream())
        .await?;
    let file = SESSION.open_options().open_buffer(buf)?;
    let read_array = file.scan()?.into_array_stream()?.read_all().await?;

    let dict = read_array
        .as_opt::<Dict>()
        .expect("expected root to be dictionary");
    assert!(
        !dict.codes().is::<Dict>(),
        "dictionary codes should not be dictionary encoded"
    );
    Ok(())
}

#[tokio::test]
async fn test_writer_basic_push() -> VortexResult<()> {
    let strings = VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array();
    let numbers = buffer![1u32, 2, 3, 4].into_array();
    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)])?.into_array();
    let dtype = st.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone());

    writer.push(st.clone()).await?;
    let summary = writer.finish().await?;

    assert_eq!(summary.row_count(), 4);

    let file = SESSION.open_options().open_buffer(buf)?;
    let result = file.scan()?.into_array_stream()?.read_all().await?;

    assert_eq!(result.len(), 4);
    assert_eq!(result.dtype(), &dtype);

    Ok(())
}

#[tokio::test]
async fn test_writer_multiple_pushes() -> VortexResult<()> {
    let mut ctx = SESSION.create_execution_ctx();
    let chunk1 =
        StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3].into_array())])?.into_array();
    let chunk2 =
        StructArray::from_fields(&[("numbers", buffer![4u32, 5, 6].into_array())])?.into_array();
    let chunk3 =
        StructArray::from_fields(&[("numbers", buffer![7u32, 8, 9].into_array())])?.into_array();

    let dtype = chunk1.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone());

    writer.push(chunk1).await?;
    writer.push(chunk2).await?;
    writer.push(chunk3).await?;

    let summary = writer.finish().await?;
    assert_eq!(summary.row_count(), 9);

    let file = SESSION.open_options().open_buffer(buf)?;
    let result = file.scan()?.into_array_stream()?.read_all().await?;

    assert_eq!(result.len(), 9);
    let numbers = result
        .execute::<StructArray>(&mut ctx)?
        .unmasked_field_by_name("numbers")?
        .clone();
    let expected = buffer![1u32, 2, 3, 4, 5, 6, 7, 8, 9].into_array();
    assert_arrays_eq!(numbers, expected, &mut ctx);

    Ok(())
}

#[tokio::test]
async fn test_writer_push_stream() -> VortexResult<()> {
    let mut ctx = SESSION.create_execution_ctx();
    let chunk1 =
        StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3].into_array())])?.into_array();
    let chunk2 =
        StructArray::from_fields(&[("numbers", buffer![4u32, 5, 6].into_array())])?.into_array();

    let dtype = chunk1.dtype().clone();

    let stream = futures::stream::iter(vec![Ok(chunk1), Ok(chunk2)]);
    let sendable_stream = ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype.clone(), stream));

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone());

    writer.push_stream(sendable_stream).await?;

    let summary = writer.finish().await?;
    assert_eq!(summary.row_count(), 6);

    let file = SESSION.open_options().open_buffer(buf)?;
    let result = file.scan()?.into_array_stream()?.read_all().await?;

    assert_eq!(result.len(), 6);
    let numbers = result
        .execute::<StructArray>(&mut ctx)?
        .unmasked_field_by_name("numbers")?
        .clone();
    let expected = buffer![1u32, 2, 3, 4, 5, 6].into_array();
    assert_arrays_eq!(numbers, expected, &mut ctx);

    Ok(())
}

#[tokio::test]
async fn test_writer_bytes_written() -> VortexResult<()> {
    let array = StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3, 4, 5].into_array())])?
        .into_array();
    let dtype = array.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype);

    assert_eq!(writer.bytes_written(), 0);

    writer.push(array.clone()).await?;
    writer.push(array).await?;

    let bytes_after_push = writer.bytes_written();
    assert!(
        bytes_after_push > 0,
        "Bytes should have been written after pushing twice"
    );

    let summary = writer.finish().await?;
    assert_eq!(summary.row_count(), 10);

    assert!(!buf.is_empty(), "Buffer should contain data");

    Ok(())
}

#[rstest]
#[case::table_one_leaf(true, 1, false, 32)]
#[case::table_two_shared_leaves(true, 2, false, 64)]
#[case::table_field_override(true, 2, true, 64)]
#[case::struct_default(false, 1, false, 32)]
#[tokio::test]
async fn test_writer_buffered_bytes(
    #[case] use_table_strategy: bool,
    #[case] leaf_count: usize,
    #[case] field_override: bool,
    #[case] expected_buffered_bytes: u64,
) -> VortexResult<()> {
    const BUFFER_SIZE: u64 = 16;

    let fields = [
        ("a", buffer![1u32, 2, 3, 4].into_array()),
        ("b", buffer![5u32, 6, 7, 8].into_array()),
    ];
    let array = StructArray::from_fields(&fields[..leaf_count])?.into_array();

    let new_leaf = || -> Arc<dyn LayoutStrategy> {
        Arc::new(BufferedStrategy::new(
            ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()),
            BUFFER_SIZE,
        ))
    };
    let validity: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
    let strategy: Arc<dyn LayoutStrategy> = if use_table_strategy {
        let mut table = TableStrategy::new(validity, new_leaf());
        if field_override {
            table = table.with_field_writer(field_path!(b), new_leaf());
        }
        Arc::new(table)
    } else {
        Arc::new(StructStrategy::new(validity, new_leaf()))
    };

    let mut buf = ByteBufferMut::empty();
    let options = SESSION.write_options().with_strategy(Arc::clone(&strategy));
    let buffered_bytes = options.buffered_bytes_tracker();
    let mut writer = options.writer(&mut buf, array.dtype().clone());

    assert_eq!(writer.buffered_bytes(), 0);

    // The third push forces two chunks through the capacity-one input channel while keeping the
    // writer open. Each physical leaf retains two BUFFER_SIZE chunks while peeking for more input.
    writer.push(array.clone()).await?;
    writer.push(array.clone()).await?;
    writer.push(array).await?;

    assert_eq!(writer.buffered_bytes(), expected_buffered_bytes);

    let summary = writer.finish().await?;
    assert_eq!(summary.row_count(), 12);
    assert_eq!(buffered_bytes.buffered_bytes(), 0);

    Ok(())
}

#[tokio::test]
async fn test_buffered_bytes_are_writer_scoped() -> VortexResult<()> {
    const BUFFER_SIZE: u64 = 16;

    let array =
        StructArray::from_fields(&[("a", buffer![1u32, 2, 3, 4].into_array())])?.into_array();
    let leaf = Arc::new(BufferedStrategy::new(
        ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()),
        BUFFER_SIZE,
    ));
    let strategy: Arc<dyn LayoutStrategy> = Arc::new(TableStrategy::new(
        Arc::new(FlatLayoutStrategy::default()),
        leaf,
    ));

    let mut first_buf = ByteBufferMut::empty();
    let mut first = SESSION
        .write_options()
        .with_strategy(Arc::clone(&strategy))
        .writer(&mut first_buf, array.dtype().clone());
    let mut second_buf = ByteBufferMut::empty();
    let mut second = SESSION
        .write_options()
        .with_strategy(strategy)
        .writer(&mut second_buf, array.dtype().clone());

    first.push(array.clone()).await?;
    first.push(array.clone()).await?;
    first.push(array.clone()).await?;
    second.push(array.clone()).await?;
    second.push(array.clone()).await?;
    second.push(array).await?;

    assert_eq!(first.buffered_bytes(), 2 * BUFFER_SIZE);
    assert_eq!(second.buffered_bytes(), 2 * BUFFER_SIZE);

    first.finish().await?;
    second.finish().await?;

    Ok(())
}

#[tokio::test]
async fn test_encoding_registered_after_write_options() -> VortexResult<()> {
    // A session that does not know about ZigZag yet.
    let session = array_session()
        .with::<EditionSession>()
        .with::<LayoutSession>()
        .with::<RuntimeSession>();

    // Configure the options before the encoding is registered; `write` is what snapshots the
    // session's encodings, so registering in between must still be honoured.
    let options = session
        .write_options()
        .with_strategy(Arc::new(FlatLayoutStrategy::default()));
    vortex_zigzag::initialize(&session);
    crate::enable_all_registered_array_encodings(&session);

    let array = ZigZag::try_new(buffer![1u32, 2, 3, 4].into_array())?.into_array();
    let dtype = array.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    options.write(&mut buf, array.to_array_stream()).await?;

    let chunks: Vec<_> = session
        .open_options()
        .open_buffer(buf)?
        .scan()?
        .into_array_stream()?
        .try_collect()
        .await?;
    let read = ChunkedArray::try_new(chunks, dtype)?.into_array();
    let mut ctx = session.create_execution_ctx();
    assert_arrays_eq!(read, buffer![-1i32, 1, -2, 2].into_array(), &mut ctx);

    Ok(())
}

#[tokio::test]
async fn test_writer_empty_chunks() -> VortexResult<()> {
    let mut ctx = SESSION.create_execution_ctx();
    let empty = StructArray::from_fields(&[(
        "numbers",
        PrimitiveArray::new::<u32>(buffer![], Validity::NonNullable).into_array(),
    )])?
    .into_array();
    let non_empty =
        StructArray::from_fields(&[("numbers", buffer![1u32, 2].into_array())])?.into_array();

    let dtype = empty.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone());

    writer.push(empty.clone()).await?;
    writer.push(non_empty).await?;
    writer.push(empty).await?;

    let summary = writer.finish().await?;
    assert_eq!(summary.row_count(), 2);

    let file = SESSION.open_options().open_buffer(buf)?;
    let result = file.scan()?.into_array_stream()?.read_all().await?;

    assert_eq!(result.len(), 2);
    let numbers = result
        .execute::<StructArray>(&mut ctx)?
        .unmasked_field_by_name("numbers")?
        .clone();
    let expected = buffer![1u32, 2].into_array();
    assert_arrays_eq!(numbers, expected, &mut ctx);

    Ok(())
}

#[tokio::test]
async fn test_writer_mixed_push_and_stream() -> VortexResult<()> {
    let mut ctx = SESSION.create_execution_ctx();
    let chunk1 =
        StructArray::from_fields(&[("numbers", buffer![1u32, 2].into_array())])?.into_array();
    let chunk2 =
        StructArray::from_fields(&[("numbers", buffer![3u32, 4].into_array())])?.into_array();
    let chunk3 =
        StructArray::from_fields(&[("numbers", buffer![5u32, 6].into_array())])?.into_array();

    let dtype = chunk1.dtype().clone();

    let stream = futures::stream::iter(vec![Ok(chunk2.clone())]);
    let sendable_stream = ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype.clone(), stream));

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone());

    writer.push(chunk1).await?;
    writer.push_stream(sendable_stream).await?;
    writer.push(chunk3).await?;

    let summary = writer.finish().await?;
    assert_eq!(summary.row_count(), 6);

    let file = SESSION.open_options().open_buffer(buf)?;
    let result = file.scan()?.into_array_stream()?.read_all().await?;

    assert_eq!(result.len(), 6);
    let numbers = result
        .execute::<StructArray>(&mut ctx)?
        .unmasked_field_by_name("numbers")?
        .clone();
    let expected = buffer![1u32, 2, 3, 4, 5, 6].into_array();
    assert_arrays_eq!(numbers, expected, &mut ctx);

    Ok(())
}

#[tokio::test]
async fn test_writer_with_complex_types() -> VortexResult<()> {
    let mut ctx = SESSION.create_execution_ctx();
    let strings = VarBinArray::from(vec!["hello", "world", "test"]).into_array();
    let numbers = buffer![100i32, 200, 300].into_array();
    let lists = ListArray::from_iter_slow::<i32, _>(
        vec![vec![1, 2], vec![3, 4, 5], vec![6]],
        Arc::new(I32.into()),
    )?;

    let chunk = StructArray::from_fields(&[
        ("strings", strings),
        ("numbers", numbers),
        ("lists", lists.into_array()),
    ])?
    .into_array();

    let dtype = chunk.dtype().clone();

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone());

    writer.push(chunk).await?;
    let footer = writer.finish().await?;

    assert_eq!(footer.row_count(), 3);

    let file = SESSION.open_options().open_buffer(buf)?;
    let result = file.scan()?.into_array_stream()?.read_all().await?;

    assert_eq!(result.len(), 3);
    assert_eq!(result.dtype(), &dtype);

    let strings_field = result
        .execute::<StructArray>(&mut ctx)?
        .unmasked_field_by_name("strings")
        .cloned()?;
    let strings_view = strings_field.execute::<VarBinViewArray>(&mut ctx)?;
    let mask = strings_view
        .validity()?
        .execute_mask(strings_view.len(), &mut ctx)?;
    let strings = (0..strings_view.len())
        .map(|i| {
            mask.value(i)
                .then(|| unsafe { String::from_utf8_unchecked(strings_view.bytes_at(i).to_vec()) })
        })
        .collect::<Vec<_>>();
    assert_eq!(
        strings,
        vec![
            Some("hello".to_string()),
            Some("world".to_string()),
            Some("test".to_string())
        ]
    );

    Ok(())
}

/// Write `array` with list decomposition forced on (through the full compress/zone pipeline) and
/// read the whole thing back.
async fn write_read_roundtrip(array: ArrayRef) -> VortexResult<ArrayRef> {
    write_read_roundtrip_with_layout(array, true).await
}

async fn write_read_roundtrip_with_layout(
    array: ArrayRef,
    use_list_layout: bool,
) -> VortexResult<ArrayRef> {
    let strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION)
        .with_list_layout()
        .build();
    let mut buf = ByteBufferMut::empty();
    if use_list_layout {
        SESSION
            .write_options()
            .with_strategy(strategy)
            .write(&mut buf, array.to_array_stream())
            .await?;
    } else {
        SESSION
            .write_options()
            .write(&mut buf, array.to_array_stream())
            .await?;
    }
    SESSION
        .open_options()
        .open_buffer(buf)?
        .scan()?
        .into_array_stream()?
        .read_all()
        .await
}

/// A `list<list<i32>>` column round-trips through the `TableStrategy` dispatcher, exercising list
/// decomposition recursing into itself (the outer list's `elements` are themselves lists).
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn nested_list_of_list_roundtrip() -> VortexResult<()> {
    let inner = ListArray::try_new(
        buffer![1i32, 2, 3, 4, 5, 6].into_array(),
        buffer![0u32, 2, 5, 5, 6].into_array(),
        Validity::NonNullable,
    )?
    .into_array();
    let outer = ListArray::try_new(
        inner,
        buffer![0u32, 2, 4].into_array(),
        Validity::NonNullable,
    )?
    .into_array();
    let st = StructArray::from_fields(&[("nested", outer)])?.into_array();

    let result = write_read_roundtrip(st.clone()).await?;
    assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx());
    Ok(())
}

type MapEntryFixture<'a> = (i32, Option<&'a str>);
type MapRowFixture<'a> = Option<Vec<MapEntryFixture<'a>>>;

fn map_array_from_rows(rows: &[MapRowFixture<'_>], keys_sorted: bool) -> VortexResult<ArrayRef> {
    let map_dtype = MapDType::try_new(
        DType::Primitive(I32, Nullability::NonNullable),
        DType::Utf8(Nullability::Nullable),
        keys_sorted,
    )?;
    let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable);
    let mut builder = MapBuilder::<u64, u64>::with_capacity_in(
        map_dtype,
        Nullability::Nullable,
        rows.len(),
        vortex_buffer::BufferAllocatorRef::static_ref(),
    );

    for row in rows {
        let scalar = match row {
            Some(entries) => {
                let entries = entries
                    .iter()
                    .map(|(key, value)| {
                        let key = Scalar::primitive(*key, Nullability::NonNullable);
                        let value = value.map_or_else(
                            || Scalar::null(DType::Utf8(Nullability::Nullable)),
                            |value| Scalar::utf8(value, Nullability::Nullable),
                        );
                        (key, value)
                    })
                    .collect::<Vec<_>>();
                Scalar::try_map(dtype.clone(), entries)?
            }
            None => Scalar::null(dtype.clone()),
        };
        builder.append_value(scalar.as_map())?;
    }

    Ok(builder.finish_into_map().into_array())
}

/// A struct containing a Map column crosses both the default flat writer and the list layout
/// strategy without changing map nullability, empty rows, duplicate keys, or scalar values.
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn struct_with_map_column_roundtrip() -> VortexResult<()> {
    for use_list_layout in [false, true] {
        let maps = map_array_from_rows(
            &[
                Some(vec![(1, Some("one")), (2, None)]),
                None,
                Some(vec![]),
                Some(vec![(1, Some("dup-old")), (1, Some("dup-new"))]),
            ],
            false,
        )?;
        let st = StructArray::from_fields(&[
            ("id", buffer![10i32, 20, 30, 40].into_array()),
            ("attrs", maps),
        ])?
        .into_array();

        let result = write_read_roundtrip_with_layout(st.clone(), use_list_layout).await?;
        assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx());
    }

    Ok(())
}

/// A `struct<{ items: list<struct<{a,b}>>? }>` column round-trips, exercising list decomposition
/// recursing into struct decomposition (list `elements` are structs) plus a nullable list validity
/// child.
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn nested_struct_list_struct_roundtrip() -> VortexResult<()> {
    let inner_struct = StructArray::from_fields(&[
        ("a", buffer![1i32, 2, 3, 4, 5].into_array()),
        ("b", buffer![10i32, 20, 30, 40, 50].into_array()),
    ])?
    .into_array();
    let items = ListArray::try_new(
        inner_struct,
        buffer![0u32, 2, 5, 5].into_array(),
        Validity::Array(BoolArray::from_iter([true, false, true]).into_array()),
    )?
    .into_array();
    let st = StructArray::from_fields(&[("items", items)])?.into_array();

    let result = write_read_roundtrip(st.clone()).await?;
    assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx());
    Ok(())
}

#[tokio::test]
async fn test_writer_with_statistics() -> VortexResult<()> {
    let array = StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3, 4, 5].into_array())])?
        .into_array();

    let mut buf = ByteBufferMut::empty();
    let mut writer = SESSION
        .write_options()
        .with_file_statistics(PRUNING_STATS.to_vec())
        .writer(&mut buf, array.dtype().clone());

    writer.push(array).await?;
    let summary = writer.finish().await?;

    assert!(summary.footer().statistics().is_some());
    assert_eq!(summary.row_count(), 5);

    Ok(())
}

#[tokio::test]
async fn test_file_metadata_roundtrip() -> VortexResult<()> {
    let array =
        StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3].into_array())])?.into_array();
    let small = ByteBuffer::copy_from(b"{\"source\":\"test\"}");
    let large = ByteBuffer::copy_from(vec![7u8; usize::from(MAX_POSTSCRIPT_SIZE) + 1024]);
    let empty = ByteBuffer::empty_aligned(vortex_buffer::Alignment::new(16));
    let aligned = ByteBuffer::copy_from_aligned(b"aligned", vortex_buffer::Alignment::new(64));

    let mut buf = ByteBufferMut::empty();
    let summary = SESSION
        .write_options()
        .with_metadata_segment("json", ByteBuffer::copy_from(b"old"))
        .with_metadata_segment("json", small.clone()) // last write wins
        .with_metadata_segment("large", large.clone())
        .with_metadata_segment("empty", empty)
        .with_metadata_segment("aligned", aligned.clone())
        .write(&mut buf, array.to_array_stream())
        .await?;

    assert_eq!(summary.footer().metadata_segments().count(), 4);
    // The footer holds only locators, so its size does not grow with the large value.
    assert!(summary.footer().approx_byte_size().unwrap() < large.len());
    for (_key, locator) in summary.footer().metadata_segments() {
        assert!(locator.alignment.is_offset_aligned(locator.offset as usize));
    }

    let bytes = ByteBuffer::from(buf);

    let default = SESSION.open_options().open_buffer(bytes.clone())?;
    assert_eq!(default.row_count(), 3);
    assert_eq!(default.metadata_segments().count(), 0);
    assert!(default.metadata_segment("json").is_none());

    let file = SESSION
        .open_options()
        .include_metadata()
        .open_buffer(bytes.clone())?;
    assert_eq!(file.metadata_segments().count(), 4);
    assert_eq!(
        file.metadata_segment("json").map(ByteBuffer::as_slice),
        Some(small.as_slice())
    );
    assert_eq!(
        file.metadata_segment("large").map(ByteBuffer::as_slice),
        Some(large.as_slice())
    );
    assert!(
        file.metadata_segment("empty")
            .vortex_expect("empty")
            .is_empty()
    );
    let resolved_aligned = file.metadata_segment("aligned").vortex_expect("aligned");
    assert_eq!(resolved_aligned.as_slice(), aligned.as_slice());
    assert!(resolved_aligned.is_aligned(vortex_buffer::Alignment::new(64)));
    assert!(file.metadata_segment("missing").is_none());

    // Resolved values are copied out, not sliced from the file buffer.
    let file_range = {
        let s = bytes.as_ptr() as usize;
        s..s + bytes.len()
    };
    let resolved = file.metadata_segment("json").vortex_expect("json");
    assert!(!file_range.contains(&(resolved.as_ptr() as usize)));

    Ok(())
}

fn with_invalid_metadata_alignment(bytes: &ByteBuffer, exponent: u8) -> ByteBuffer {
    let eof_offset = bytes.len() - crate::EOF_SIZE;
    let postscript_len =
        u16::from_le_bytes(bytes[eof_offset + 2..eof_offset + 4].try_into().unwrap()) as usize;
    let postscript_offset = eof_offset - postscript_len;
    let old = flatbuffers::root::<fb::Postscript>(&bytes[postscript_offset..eof_offset]).unwrap();

    let copy_segment = |segment: fb::PostscriptSegment<'_>| {
        (
            segment.offset(),
            segment.length(),
            segment.alignment_exponent(),
        )
    };
    let dtype = old.dtype().map(copy_segment);
    let layout = copy_segment(old.layout().unwrap());
    let statistics = old.statistics().map(copy_segment);
    let footer = copy_segment(old.footer().unwrap());
    let metadata = old.metadata().unwrap().get(0);
    let metadata_key = metadata.key().to_string();
    let metadata_segment = copy_segment(metadata.segment());

    fn create_segment<'a>(
        fbb: &mut FlatBufferBuilder<'a>,
        (offset, length, alignment_exponent): (u64, u32, u8),
    ) -> flatbuffers::WIPOffset<fb::PostscriptSegment<'a>> {
        fb::PostscriptSegment::create(
            fbb,
            &fb::PostscriptSegmentArgs {
                offset,
                length,
                alignment_exponent,
                _compression: None,
                _encryption: None,
            },
        )
    }

    let mut fbb = FlatBufferBuilder::new();
    let dtype = dtype.map(|segment| create_segment(&mut fbb, segment));
    let layout = create_segment(&mut fbb, layout);
    let statistics = statistics.map(|segment| create_segment(&mut fbb, segment));
    let footer = create_segment(&mut fbb, footer);
    let key = fbb.create_string(&metadata_key);
    let invalid_segment =
        create_segment(&mut fbb, (metadata_segment.0, metadata_segment.1, exponent));
    let metadata = fb::PostscriptMetadata::create(
        &mut fbb,
        &fb::PostscriptMetadataArgs {
            key: Some(key),
            segment: Some(invalid_segment),
        },
    );
    let metadata = fbb.create_vector(&[metadata]);
    let postscript = fb::Postscript::create(
        &mut fbb,
        &fb::PostscriptArgs {
            dtype,
            layout: Some(layout),
            statistics,
            footer: Some(footer),
            metadata: Some(metadata),
        },
    );
    fbb.finish_minimal(postscript);
    let postscript = fbb.finished_data();

    let mut corrupted =
        ByteBufferMut::with_capacity(postscript_offset + postscript.len() + crate::EOF_SIZE);
    corrupted.extend_from_slice(&bytes[..postscript_offset]);
    corrupted.extend_from_slice(postscript);
    corrupted.extend_from_slice(&VERSION.to_le_bytes());
    corrupted.extend_from_slice(&(postscript.len() as u16).to_le_bytes());
    corrupted.extend_from_slice(&crate::MAGIC_BYTES);
    corrupted.freeze()
}

#[tokio::test]
async fn test_file_metadata_malformed_alignment_returns_error_on_default_open() -> VortexResult<()>
{
    let mut output = ByteBufferMut::empty();
    SESSION
        .write_options()
        .with_metadata_segment("key", ByteBuffer::copy_from(b"value"))
        .write(&mut output, buffer![1u32].into_array().to_array_stream())
        .await?;
    let corrupted = with_invalid_metadata_alignment(&ByteBuffer::from(output), 64);

    for include_metadata in [false, true] {
        let result = SESSION
            .open_options()
            .with_include_metadata(include_metadata)
            .open_buffer(corrupted.clone());
        let error = match result {
            Ok(_) => panic!("invalid alignment exponent must fail open"),
            Err(error) => error,
        };
        assert!(error.to_string().contains("Alignment exponent"));
    }

    Ok(())
}

#[tokio::test]
async fn timestamp_unit_mismatch() -> Result<(), Box<dyn std::error::Error>> {
    // Write file with MILLISECONDS timestamps
    let ts_array = PrimitiveArray::from_iter(vec![1704067200000i64, 1704153600000, 1704240000000])
        .into_array();
    let temporal = TemporalArray::new_timestamp(ts_array, TimeUnit::Milliseconds, None);

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, temporal.into_array().to_array_stream())
        .await?;

    // Read with SECONDS filter scalar
    let filter_expr = gt(
        root(),
        lit(Scalar::extension::<Timestamp>(
            TimestampOptions {
                unit: TimeUnit::Seconds,
                tz: None,
            },
            Scalar::from(1704153600i64),
        )),
    );

    let file = SESSION.open_options().open_buffer(buf)?;
    let filter = filter_expr.bind(file.dtype())?.optimize_recursive()?;
    let mut stream = file.scan()?.with_filter(filter).into_array_stream()?;
    let result = stream.try_next().await;

    assert!(result.is_err());

    Ok(())
}

/// Regression test: filtering a milliseconds timestamp column with a seconds scalar should
/// always error, regardless of how the internal children of `DateTimePartsArray` are encoded.
///
/// The compressor's built-in constant detection encodes the seconds/subseconds children
/// (`[0, 0, 0]`) as `ConstantArray`s. The scanner should still detect the time unit
/// mismatch and error, not silently return wrong results.
#[tokio::test]
async fn timestamp_unit_mismatch_errors_with_constant_children()
-> Result<(), Box<dyn std::error::Error>> {
    let compressor = vortex_btrblocks::BtrBlocksCompressor::from_session(&SESSION);

    // Write file with MILLISECONDS timestamps using this compressor.
    let ts_array = PrimitiveArray::from_iter(vec![1704067200000i64, 1704153600000, 1704240000000])
        .into_array();
    let temporal = TemporalArray::new_timestamp(ts_array, TimeUnit::Milliseconds, None);

    let strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION)
        .with_compressor(compressor)
        .build();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .with_strategy(strategy)
        .write(&mut buf, temporal.into_array().to_array_stream())
        .await?;

    // Read with SECONDS filter scalar — should error due to time unit mismatch.
    let filter_expr = gt(
        root(),
        lit(Scalar::extension::<Timestamp>(
            TimestampOptions {
                unit: TimeUnit::Seconds,
                tz: None,
            },
            Scalar::from(1704153600i64),
        )),
    );

    let file = SESSION.open_options().open_buffer(buf)?;
    let filter = filter_expr.bind(file.dtype())?.optimize_recursive()?;
    let stream = file.scan()?.with_filter(filter).into_array_stream()?;
    let results = stream.try_collect::<Vec<_>>().await;

    assert!(
        results.is_err(),
        "Expected error from timestamp unit mismatch (ms vs s), but got {} results. \
         This indicates the scanner silently applied the filter incorrectly when \
         DateTimePartsArray children use ConstantArray encoding.",
        results?.len()
    );

    Ok(())
}

/// Collect all segment byte offsets reachable from a layout node.
fn collect_segment_offsets(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) -> Vec<u64> {
    let mut result = Vec::new();
    collect_segment_offsets_inner(layout, segment_specs, &mut result);
    result
}

fn collect_segment_offsets_inner(
    layout: &dyn DynLayout,
    segment_specs: &[SegmentSpec],
    result: &mut Vec<u64>,
) {
    for seg_id in layout.segment_ids() {
        result.push(segment_specs[*seg_id as usize].offset);
    }
    for child in layout.children().unwrap() {
        collect_segment_offsets_inner(child.as_ref(), segment_specs, result);
    }
}

/// Assert that all offsets in `before` are less than all offsets in `after`.
fn assert_offsets_ordered(before: &[u64], after: &[u64], context: &str) {
    if let (Some(&max_before), Some(&min_after)) = (before.iter().max(), after.iter().min()) {
        assert!(
            max_before < min_after,
            "{context}: expected all 'before' offsets < all 'after' offsets, \
             but max before = {max_before} >= min after = {min_after}"
        );
    }
}

/// Whether any node in the layout tree is a dict layout.
fn layout_has_dict(layout: &dyn DynLayout) -> bool {
    layout.encoding_id().as_ref() == "vortex.dict"
        || layout
            .children()
            .unwrap()
            .iter()
            .any(|child| layout_has_dict(child.as_ref()))
}

/// Rows in the [`large_flat_file`] fixture; spans the sub-split threshold.
const FLAT_N_ROWS: u64 = 250_000;

/// A single flat (unchunked) [`FLAT_N_ROWS`]-row layout with alternating-sign values, so filters
/// select rows on both sides of any split boundary. Returns the opened file and original array.
async fn large_flat_file() -> VortexResult<(VortexFile, ArrayRef)> {
    let values =
        Buffer::from_iter((0..FLAT_N_ROWS as i32).map(|i| if i % 2 == 0 { i } else { -i }))
            .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .with_strategy(Arc::new(FlatLayoutStrategy::default()))
        .write(&mut buf, values.to_array_stream())
        .await?;

    Ok((SESSION.open_options().open_buffer(buf)?, values))
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> {
    // A single flat (unchunked) 250k-row layout spans the 100k sub-split threshold, so the scan
    // must decode it as multiple row-range splits.
    let mut ctx = SESSION.create_execution_ctx();
    let (file, values) = large_flat_file().await?;

    // Sub-division caps each split at DEFAULT_MAX_SPLIT_ROWS while tiling the file exactly.
    let splits = file.splits()?;
    assert!(splits.len() > 1, "expected sub-divided splits: {splits:?}");
    assert!(
        splits
            .iter()
            .all(|r| r.end - r.start <= DEFAULT_MAX_SPLIT_ROWS)
    );
    assert_eq!(splits.first().map(|r| r.start), Some(0));
    assert_eq!(splits.last().map(|r| r.end), Some(FLAT_N_ROWS));
    assert!(splits.windows(2).all(|w| w[0].end == w[1].start));

    // A full scan across the sub-splits returns the original rows.
    let result = file.scan()?.into_array_stream()?.read_all().await?;
    assert_arrays_eq!(result, values, &mut ctx);

    // A filtered scan crossing sub-split boundaries selects exactly the matching rows.
    let result = file
        .scan()?
        .with_filter(bind_scan_expr(&file, gt(root(), lit(0i32))))
        .into_array_stream()?
        .read_all()
        .await?;
    let expected =
        Buffer::from_iter((0..FLAT_N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array();
    assert_arrays_eq!(result, expected, &mut ctx);

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_layout_split_keeps_large_chunk_whole() -> VortexResult<()> {
    // The same over-wide single chunk as above, scanned with `SplitBy::Layout`: the scan follows
    // the layout's chunk boundaries exactly, so the file decodes as one batch.
    let mut ctx = SESSION.create_execution_ctx();
    let (file, values) = large_flat_file().await?;

    let mut chunks: Vec<ArrayRef> = file
        .scan()?
        .with_split_by(SplitBy::Layout)
        .into_array_stream()?
        .try_collect()
        .await?;
    assert_eq!(chunks.len(), 1, "expected a single un-split chunk");
    assert_arrays_eq!(chunks.remove(0), values, &mut ctx);

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_sub_splitting_max_rows_caps_scan_batches() -> VortexResult<()> {
    // A custom `max_rows` tiles the 250k-row chunk into five evenly sized batches.
    const MAX_ROWS: u64 = 50_000;
    let (file, _) = large_flat_file().await?;

    let chunks: Vec<ArrayRef> = file
        .scan()?
        .with_split_by(SplitBy::LayoutSubSplitting { max_rows: MAX_ROWS })
        .into_array_stream()?
        .try_collect()
        .await?;
    assert_eq!(chunks.len(), (FLAT_N_ROWS / MAX_ROWS) as usize);
    assert!(chunks.iter().all(|c| c.len() as u64 <= MAX_ROWS));

    Ok(())
}

#[rstest]
#[case::unaligned(33_333)]
#[case::exceeds_file(300_000)]
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_flat_chunk_scan_with_row_count_splits(
    #[case] rows_per_split: usize,
) -> VortexResult<()> {
    // Fixed-size splits ignore chunk boundaries entirely, so scans must produce identical
    // results whether the split size straddles the chunk arbitrarily or exceeds the file's
    // row count (a single split).
    let mut ctx = SESSION.create_execution_ctx();
    let (file, values) = large_flat_file().await?;

    let result = file
        .scan()?
        .with_split_by(SplitBy::RowCount(rows_per_split))
        .into_array_stream()?
        .read_all()
        .await?;
    assert_arrays_eq!(result, values, &mut ctx);

    let result = file
        .scan()?
        .with_split_by(SplitBy::RowCount(rows_per_split))
        .with_filter(bind_scan_expr(&file, gt(root(), lit(0i32))))
        .into_array_stream()?
        .read_all()
        .await?;
    let expected =
        Buffer::from_iter((0..FLAT_N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array();
    assert_arrays_eq!(result, expected, &mut ctx);

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_string_chunks_stay_fine_grained_under_split_cap() -> VortexResult<()> {
    // Default writing targets ~1MiB uncompressed blocks, so ~120-byte strings chunk at a few
    // thousand rows (~8k with today's defaults). These natural boundaries sit far below the
    // sub-split cap, and SplitBy::LayoutSubSplitting must pass them through untouched.
    let mut ctx = SESSION.create_execution_ctx();
    const N_ROWS: usize = 40_000;
    let strings = VarBinArray::from_iter(
        (0..N_ROWS).map(|i| Some(format!("{i:0>120}"))),
        DType::Utf8(Nullability::Nullable),
    )
    .into_array();
    let st = StructArray::from_fields(&[("s", strings)])?.into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.to_array_stream())
        .await?;

    let file = SESSION.open_options().open_buffer(buf)?;

    let splits = file.splits()?;
    assert!(
        splits.len() > 1,
        "expected multiple natural chunks: {splits:?}"
    );
    assert!(
        splits
            .iter()
            .all(|r| r.end - r.start < DEFAULT_MAX_SPLIT_ROWS / 4),
        "string chunks should stay fine-grained, nowhere near the split cap: {splits:?}"
    );
    assert_eq!(splits.first().map(|r| r.start), Some(0));
    assert_eq!(splits.last().map(|r| r.end), Some(N_ROWS as u64));
    assert!(splits.windows(2).all(|w| w[0].end == w[1].start));

    let result = file.scan()?.into_array_stream()?.read_all().await?;
    assert_arrays_eq!(result, st, &mut ctx);

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> {
    // Create low-cardinality strings to trigger dict encoding, plus an integer column.
    let n = 100_000;
    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
    let strings = VarBinArray::from(values).into_array();
    let numbers = PrimitiveArray::from_iter(0..n as i32).into_array();

    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)])?;

    let mut buf = ByteBufferMut::empty();
    let summary = SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await?;

    let footer = summary.footer();
    let segment_specs = footer.segment_map();
    let root = footer.layout();

    // Walk the layout tree and find all dict layouts.
    // Verify codes segments come before values segments in byte order within each run.
    fn check_dict_ordering(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) {
        if layout.encoding_id().as_ref() == "vortex.dict" {
            // child 0 = values, child 1 = codes
            let values_offsets =
                collect_segment_offsets(layout.slot(0).unwrap().unwrap().as_ref(), segment_specs);
            let codes_offsets =
                collect_segment_offsets(layout.slot(1).unwrap().unwrap().as_ref(), segment_specs);

            assert_offsets_ordered(
                &codes_offsets,
                &values_offsets,
                "dict: codes should come before values",
            );
        }

        for child in layout.children().unwrap() {
            check_dict_ordering(child.as_ref(), segment_specs);
        }
    }

    check_dict_ordering(root.as_ref(), segment_specs);

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn dict_probe_honours_configured_compressor() -> VortexResult<()> {
    // Low-cardinality strings so the default cascade picks a dictionary.
    let n = 32_768;
    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
    let strings = VarBinArray::from(values).into_array();

    let mut buf = ByteBufferMut::empty();
    let summary = SESSION
        .write_options()
        .with_strategy(crate::strategy::WriteStrategyBuilder::from_session(&SESSION).build())
        .write(&mut buf, strings.clone().to_array_stream())
        .await?;
    assert!(
        layout_has_dict(summary.footer().layout().as_ref()),
        "default builder should produce a dict layout for low-cardinality strings"
    );

    let no_string_dict =
        BtrBlocksCompressorBuilder::from_session(&SESSION).exclude_schemes([StringDictScheme.id()]);
    let mut buf = ByteBufferMut::empty();
    let summary = SESSION
        .write_options()
        .with_strategy(
            crate::strategy::WriteStrategyBuilder::from_session(&SESSION)
                .with_btrblocks_builder(no_string_dict)
                .build(),
        )
        .write(&mut buf, strings.to_array_stream())
        .await?;
    assert!(
        !layout_has_dict(summary.footer().layout().as_ref()),
        "excluding StringDict from the configured compressor should disable the dict layout"
    );

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn probe_compressor_override_is_independent() -> VortexResult<()> {
    // Low-cardinality strings the default cascade would dict-encode.
    let n = 32_768;
    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
    let strings = VarBinArray::from(values).into_array();

    let probe_without_dict = BtrBlocksCompressorBuilder::from_session(&SESSION)
        .exclude_schemes([StringDictScheme.id()])
        .build();

    let mut buf = ByteBufferMut::empty();
    let summary = SESSION
        .write_options()
        .with_strategy(
            crate::strategy::WriteStrategyBuilder::from_session(&SESSION)
                .with_probe_compressor(probe_without_dict)
                .build(),
        )
        .write(&mut buf, strings.to_array_stream())
        .await?;
    assert!(
        !layout_has_dict(summary.footer().layout().as_ref()),
        "probe override should disable the dict layout independently of the data/stats compressor"
    );

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
    // Create a multi-column struct with enough rows to produce zone maps.
    let n = 100_000;
    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
    let strings = VarBinArray::from(values).into_array();
    let numbers = PrimitiveArray::from_iter(0..n as i32).into_array();
    let floats = PrimitiveArray::from_iter((0..n).map(|i| i as f64 * 0.1)).into_array();

    let st = StructArray::from_fields(&[
        ("strings", strings),
        ("numbers", numbers),
        ("floats", floats),
    ])?;

    let mut buf = ByteBufferMut::empty();
    let summary = SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await?;

    let footer = summary.footer();
    let segment_specs = footer.segment_map();
    let root = footer.layout();

    // Find all zoned layouts and verify data segments come before zone map segments.
    fn check_zoned_ordering(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) {
        if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
            // child 0 = data, child 1 = zones
            let data_offsets =
                collect_segment_offsets(layout.slot(0).unwrap().unwrap().as_ref(), segment_specs);
            let zones_offsets =
                collect_segment_offsets(layout.slot(1).unwrap().unwrap().as_ref(), segment_specs);

            assert_offsets_ordered(
                &data_offsets,
                &zones_offsets,
                "zoned: data should come before zones",
            );
        }

        for child in layout.children().unwrap() {
            check_zoned_ordering(child.as_ref(), segment_specs);
        }
    }

    check_zoned_ordering(root.as_ref(), segment_specs);

    // Additionally: all zone map segments across all columns should appear after
    // all data segments across all columns.
    let mut all_data_offsets = Vec::new();
    let mut all_zones_offsets = Vec::new();

    fn collect_all_zoned(
        layout: &dyn DynLayout,
        segment_specs: &[SegmentSpec],
        all_data: &mut Vec<u64>,
        all_zones: &mut Vec<u64>,
    ) {
        if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
            // child 0 = data, child 1 = zones
            all_data.extend(collect_segment_offsets(
                layout.slot(0).unwrap().unwrap().as_ref(),
                segment_specs,
            ));
            all_zones.extend(collect_segment_offsets(
                layout.slot(1).unwrap().unwrap().as_ref(),
                segment_specs,
            ));
            return;
        }
        for child in layout.children().unwrap() {
            collect_all_zoned(child.as_ref(), segment_specs, all_data, all_zones);
        }
    }

    collect_all_zoned(
        root.as_ref(),
        segment_specs,
        &mut all_data_offsets,
        &mut all_zones_offsets,
    );

    assert_offsets_ordered(
        &all_data_offsets,
        &all_zones_offsets,
        "global: all data segments should come before all zone map segments",
    );

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn test_can_prune_composite_predicates() -> VortexResult<()> {
    // Regression test for `can_prune` after `ScalarFnConstantRule` was removed
    // (#7575): composite falsification trees no longer constant-fold during
    // execution, so `can_prune` must read the one-row evaluated result instead
    // of requiring a `Columnar::Constant`. `Eq` is affected too: its
    // falsification is internally `or(min > lit, lit > max)`.
    let st = StructArray::from_fields(&[
        ("age", buffer![15i32, 18, 22, 25].into_array()),
        ("price", buffer![120i32, 130, 140, 150].into_array()),
    ])?;
    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, st.into_array().to_array_stream())
        .await?;
    let file = SESSION.open_options().open_buffer(buf)?;

    // Bare comparisons: falsified directly by min/max stats.
    assert!(file.can_prune(&gt(col("age"), lit(30)))?);
    assert!(file.can_prune(&lt(col("price"), lit(100)))?);

    // Composite predicates whose falsifications are boolean trees.
    assert!(file.can_prune(&and(gt(col("age"), lit(30)), lt(col("price"), lit(100))))?);
    assert!(file.can_prune(&or(gt(col("age"), lit(30)), lt(col("age"), lit(10))))?);
    assert!(file.can_prune(&eq(col("age"), lit(5)))?);

    // Non-falsifiable controls: rows may match, so pruning must refuse.
    assert!(!file.can_prune(&gt(col("age"), lit(20)))?);
    assert!(!file.can_prune(&eq(col("age"), lit(18)))?);
    assert!(!file.can_prune(&and(gt(col("age"), lit(20)), gt(col("price"), lit(100))))?);

    Ok(())
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> {
    use vortex_buffer::ByteBuffer;

    let mut ctx = SESSION.create_execution_ctx();

    let empty: Vec<u8> = vec![];
    let chunk0: Vec<Vec<u8>> = vec![
        vec![0x1d, 0x00],
        empty.clone(),
        vec![0x1d, 0x10, 0x9d, 0x08],
        empty.clone(),
        empty.clone(),
        empty.clone(),
        empty.clone(),
        empty.clone(),
        empty.clone(),
    ];
    let chunk1: Vec<Vec<u8>> = vec![
        empty.clone(),
        empty.clone(),
        vec![0x40],
        empty.clone(),
        empty.clone(),
        empty.clone(),
        empty.clone(),
        empty.clone(),
        vec![0x24],
        vec![0x43, 0xff],
    ];
    let mut big = vec![0xffu8; 112];
    big[89] = 0x03;
    let mut chunk2: Vec<Vec<u8>> = vec![empty.clone(); 10];
    chunk2[8] = big;

    let bin = DType::Binary(Nullability::NonNullable);
    let mk_struct = |vals: Vec<Vec<u8>>| -> VortexResult<ArrayRef> {
        let yyw = VarBinArray::from_vec(vals, bin.clone()).into_array();
        Ok(StructArray::from_fields(&[("yyw", yyw)])?.into_array())
    };
    let array =
        ChunkedArray::from_iter([mk_struct(chunk0)?, mk_struct(chunk1)?, mk_struct(chunk2)?])
            .into_array();

    let mut buf = ByteBufferMut::empty();
    SESSION
        .write_options()
        .write(&mut buf, array.to_array_stream())
        .await?;

    let mut literal = vec![0x6fu8; 5];
    literal.extend(iter::repeat_n(0xffu8, 57));
    literal.push(0x98);
    assert_eq!(literal.len(), 63);

    let filter = gt(
        get_item("yyw", root()),
        lit(Scalar::binary(
            ByteBuffer::from(literal),
            Nullability::NonNullable,
        )),
    );

    let file = SESSION.open_options().open_buffer(buf)?;
    let filter = filter.bind(file.dtype())?.optimize_recursive()?;
    let result = file
        .scan()?
        .with_filter(filter)
        .into_array_stream()?
        .read_all()
        .await?
        .execute::<StructArray>(&mut ctx)?;

    assert_eq!(result.len(), 1);
    Ok(())
}