From d6147dd14d526df1cbe7b1a55a3033045725269c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 8 Sep 2026 12:01:10 +0100 Subject: [PATCH 1/2] Assert array buffer sizes in Python IO doctests Signed-off-by: Robert Kruszewski --- vortex-python/src/io.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index b5f2d2bcfbd..bd8e0941821 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -302,18 +302,16 @@ impl PyVortexWriteOptions { /// Let's model some stock ticker data. As you may know, the stock market always (noisly) goes /// up: /// - /// >>> import os /// >>> import random /// >>> sprl = vx.array([random.randint(i, i + 10) for i in range(100_000)]) /// - /// If we naively wrote 4-bytes for each of these integers to a file we'd have 400,000 bytes! - /// Let's see how small this is when we write with the default Vortex write options (which are - /// also used by :func:`vortex.io.write`): + /// If we naively wrote 8-bytes for each of these integers to a file we'd have 800,000 bytes! + /// Let's see how small the array buffers are when we write with the default Vortex write + /// options (which are also used by :func:`vortex.io.write`): /// /// >>> vx.io.VortexWriteOptions.default().write(sprl, "chonky.vortex") - /// >>> import os - /// >>> os.path.getsize('chonky.vortex') - /// 215788 + /// >>> vx.open("chonky.vortex").scan().read_all().nbytes + /// 213248 /// /// Wow, Vortex manages to use about two bytes per integer! So advanced. So tiny. /// @@ -322,8 +320,8 @@ impl PyVortexWriteOptions { /// We sure can. /// /// >>> vx.io.VortexWriteOptions.compact().write(sprl, "tiny.vortex") - /// >>> os.path.getsize('tiny.vortex') - /// 54992 + /// >>> vx.open("tiny.vortex").scan().read_all().nbytes + /// 52564 /// /// Random numbers are not (usually) composed of random bytes! #[staticmethod] From e8208a8c813016b0a9397362a05ba2588871fc5f Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 8 Sep 2026 12:47:38 +0100 Subject: [PATCH 2/2] Test stock ticker file sizes in vortex-file Signed-off-by: Robert Kruszewski --- Cargo.lock | 2 ++ vortex-file/Cargo.toml | 2 ++ vortex-file/src/tests.rs | 59 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 159b3b020d4..8791225e055 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11050,7 +11050,9 @@ dependencies = [ "object_store", "parking_lot", "pin-project-lite", + "rand 0.10.2", "rstest", + "tempfile", "tokio", "tracing", "url", diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 626a3fabffe..79eb280f05f 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -63,7 +63,9 @@ vortex-zstd = { workspace = true, optional = true } [dev-dependencies] allocator-api2 = { workspace = true } divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-io = { workspace = true, features = ["tokio"] } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..6274ac1f0aa 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::cast_possible_truncation)] +use std::fs; use std::iter; use std::sync::Arc; use std::sync::LazyLock; @@ -11,7 +12,11 @@ use flatbuffers::FlatBufferBuilder; use futures::StreamExt; use futures::TryStreamExt; use futures::pin_mut; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; use rstest::rstest; +use tempfile::tempdir; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -76,7 +81,10 @@ use vortex_edition::EditionSession; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_flatbuffers::footer as fb; +use vortex_io::VortexWrite; +use vortex_io::runtime::tokio::TokioRuntime; use vortex_io::session::RuntimeSession; +use vortex_io::std_file::FileWrite; use vortex_layout::DynLayout; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::buffered::BufferedStrategy; @@ -129,6 +137,57 @@ async fn test_eof_values() { assert_eq!(V1_FOOTER_FBS_SIZE, 32); } +// Optional encodings affect both compression choices and the registry stored in the footer. +#[rstest] +#[case::default( + BtrBlocksCompressorBuilder::default(), + match (cfg!(feature = "zstd"), cfg!(feature = "unstable_encodings")) { + (false, false) => 215_876, + (true, false) => 215_900, + (false, true) => 70_100, + (true, true) => 70_164, + } +)] +#[cfg_attr( + feature = "zstd", + case::compact( + BtrBlocksCompressorBuilder::default().with_compact(), + if cfg!(feature = "unstable_encodings") { 55_248 } else { 55_032 } + ) +)] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_stock_ticker_file_size( + #[case] compressor: BtrBlocksCompressorBuilder, + #[case] expected_size: u64, +) -> VortexResult<()> { + // Same stock-ticker distribution as the Python IO doctest, with a fixed Rust RNG seed. + let mut rng = StdRng::seed_from_u64(0); + let array = PrimitiveArray::from_iter((0..100_000i64).map(|i| rng.random_range(i..=i + 10))) + .into_array(); + let directory = tempdir()?; + let path = directory.path().join("stock_ticker.vortex"); + let mut writer = FileWrite::create(&path, TokioRuntime::current()).await?; + SESSION + .write_options() + .with_strategy( + crate::strategy::WriteStrategyBuilder::default() + .with_btrblocks_builder(compressor) + .build(), + ) + .write(&mut writer, array.clone().to_array_stream()) + .await?; + writer.shutdown().await?; + + let file = SESSION + .open_options() + .open_buffer(ByteBuffer::from(fs::read(&path)?))?; + let actual = file.scan()?.into_array_stream()?.read_all().await?; + assert_arrays_eq!(actual, array, &mut SESSION.create_execution_ctx()); + assert_eq!(fs::metadata(path)?.len(), expected_size); + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_read_simple() {