Test IO file sizes in Rust and array sizes in Python doctests

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"] }

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() {

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]