fn criterion_benchmark()

in arrow-ipc/benches/ipc_reader.rs [31:172]

• 122 lines of code
• 17 McCabe index (conditional complexity)

fn criterion_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("arrow_ipc_reader");

    group.bench_function("StreamReader/read_10", |b| {
        let buffer = ipc_stream(IpcWriteOptions::default());
        b.iter(move || {
            let projection = None;
            let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap();
            for _ in 0..10 {
                reader.next().unwrap().unwrap();
            }
            assert!(reader.next().is_none());
        })
    });

    group.bench_function("StreamReader/no_validation/read_10", |b| {
        let buffer = ipc_stream(IpcWriteOptions::default());
        b.iter(move || {
            let projection = None;
            let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap();
            unsafe {
                // safety: we created a valid IPC file
                reader = reader.with_skip_validation(true);
            }
            for _ in 0..10 {
                reader.next().unwrap().unwrap();
            }
            assert!(reader.next().is_none());
        })
    });

    group.bench_function("StreamReader/read_10/zstd", |b| {
        let buffer = ipc_stream(
            IpcWriteOptions::default()
                .try_with_compression(Some(CompressionType::ZSTD))
                .unwrap(),
        );
        b.iter(move || {
            let projection = None;
            let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap();
            for _ in 0..10 {
                reader.next().unwrap().unwrap();
            }
            assert!(reader.next().is_none());
        })
    });

    group.bench_function("StreamReader/no_validation/read_10/zstd", |b| {
        let buffer = ipc_stream(
            IpcWriteOptions::default()
                .try_with_compression(Some(CompressionType::ZSTD))
                .unwrap(),
        );
        b.iter(move || {
            let projection = None;
            let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap();
            unsafe {
                // safety: we created a valid IPC file
                reader = reader.with_skip_validation(true);
            }
            for _ in 0..10 {
                reader.next().unwrap().unwrap();
            }
            assert!(reader.next().is_none());
        })
    });

    // --- Create IPC File ---
    group.bench_function("FileReader/read_10", |b| {
        let buffer = ipc_file();
        b.iter(move || {
            let projection = None;
            let cursor = Cursor::new(buffer.as_slice());
            let mut reader = FileReader::try_new(cursor, projection).unwrap();
            for _ in 0..10 {
                reader.next().unwrap().unwrap();
            }
            assert!(reader.next().is_none());
        })
    });

    group.bench_function("FileReader/no_validation/read_10", |b| {
        let buffer = ipc_file();
        b.iter(move || {
            let projection = None;
            let cursor = Cursor::new(buffer.as_slice());
            let mut reader = FileReader::try_new(cursor, projection).unwrap();
            unsafe {
                // safety: we created a valid IPC file
                reader = reader.with_skip_validation(true);
            }
            for _ in 0..10 {
                reader.next().unwrap().unwrap();
            }
            assert!(reader.next().is_none());
        })
    });

    // write to an actual file
    let dir = tempdir().unwrap();
    let path = dir.path().join("test.arrow");
    let mut file = std::fs::File::create(&path).unwrap();
    file.write_all(&ipc_file()).unwrap();
    drop(file);

    group.bench_function("FileReader/read_10/mmap", |b| {
        let path = &path;
        b.iter(move || {
            let ipc_file = std::fs::File::open(path).expect("failed to open file");
            let mmap = unsafe { memmap2::Mmap::map(&ipc_file).expect("failed to mmap file") };

            // Convert the mmap region to an Arrow `Buffer` to back the arrow arrays.
            let bytes = bytes::Bytes::from_owner(mmap);
            let buffer = Buffer::from(bytes);
            let decoder = IPCBufferDecoder::new(buffer);
            assert_eq!(decoder.num_batches(), 10);

            for i in 0..decoder.num_batches() {
                decoder.get_batch(i);
            }
        })
    });

    group.bench_function("FileReader/no_validation/read_10/mmap", |b| {
        let path = &path;
        b.iter(move || {
            let ipc_file = std::fs::File::open(path).expect("failed to open file");
            let mmap = unsafe { memmap2::Mmap::map(&ipc_file).expect("failed to mmap file") };

            // Convert the mmap region to an Arrow `Buffer` to back the arrow arrays.
            let bytes = bytes::Bytes::from_owner(mmap);
            let buffer = Buffer::from(bytes);
            let decoder = IPCBufferDecoder::new(buffer);
            let decoder = unsafe { decoder.with_skip_validation(true) };
            assert_eq!(decoder.num_batches(), 10);

            for i in 0..decoder.num_batches() {
                decoder.get_batch(i);
            }
        })
    });
}