diskann_disk/storage/

cached_reader.rs

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */
use std::io::{Read, Seek};

use diskann::{ANNError, ANNResult};
use diskann_providers::storage::StorageReadProvider;
use tracing::info;

/// Sequential cached reads with a generic storage provider with read access.
pub struct CachedReader<Storage>
where
    Storage: StorageReadProvider,
{
    /// File reader
    reader: Storage::Reader,

    /// # bytes to cache in one shot read
    cache_size: u64,

    /// Underlying buf for cache
    cache_buf: Vec<u8>,

    /// Offset into cache_buf for cur_pos
    cur_off: u64,

    /// File size
    size: u64,
}

impl<Storage> CachedReader<Storage>
where
    Storage: StorageReadProvider,
{
    pub fn new(
        filename: &str,
        cache_size: u64,
        storage_provider: &Storage,
    ) -> std::io::Result<Self> {
        info!("Opening: {}", filename);
        let mut reader = storage_provider.open_reader(filename)?;
        let size = storage_provider.get_length(filename)?;

        let cache_size = cache_size.min(size);
        let mut cache_buf = vec![0; cache_size as usize];
        reader.read_exact(&mut cache_buf)?;
        info!(
            "Opened: {}, size: {}, cache_size: {}",
            filename, size, cache_size
        );

        Ok(Self {
            reader,
            cache_size,
            cache_buf,
            cur_off: 0,
            size,
        })
    }

    pub fn get_file_size(&self) -> u64 {
        self.size
    }

    pub fn read(&mut self, read_buf: &mut [u8]) -> ANNResult<()> {
        let n_bytes = read_buf.len() as u64;
        if n_bytes <= (self.cache_size - self.cur_off) {
            // case 1: cache contains all data
            read_buf.copy_from_slice(
                &self.cache_buf
                    [(self.cur_off as usize)..(self.cur_off as usize + n_bytes as usize)],
            );
            self.cur_off += n_bytes;
        } else {
            // case 2: cache contains some data
            let cached_bytes = self.cache_size - self.cur_off;
            if n_bytes - cached_bytes > self.size - self.reader.stream_position()? {
                return Err(ANNError::log_index_error(format_args!(
                    "Reading beyond end of file, n_bytes: {} cached_bytes: {} fsize: {} current pos: {}",
                    n_bytes,
                    cached_bytes,
                    self.size,
                    self.reader.stream_position()?
                )));
            }

            read_buf[..cached_bytes as usize]
                .copy_from_slice(&self.cache_buf[self.cur_off as usize..]);
            // go to disk and fetch more data
            self.reader
                .read_exact(&mut read_buf[cached_bytes as usize..])?;
            // reset cur off
            self.cur_off = self.cache_size;

            let size_left = self.size - self.reader.stream_position()?;
            if size_left >= self.cache_size {
                self.reader.read_exact(&mut self.cache_buf)?;
                self.cur_off = 0;
            }
            // note that if size_left < cache_size, then cur_off = cache_size,
            // so subsequent reads will all be directly from file
        }
        Ok(())
    }

    pub fn read_u32(&mut self) -> ANNResult<u32> {
        let mut bytes = [0u8; 4];
        self.read(&mut bytes)?;
        Ok(u32::from_le_bytes(bytes))
    }
}

#[cfg(test)]
mod cached_reader_test {

    use diskann_providers::storage::{StorageWriteProvider, VirtualStorageProvider};
    use vfs::MemoryFS;

    use super::*;

    #[test]
    fn cached_reader_works() {
        let file_name = "/cached_reader_works_test2.bin";
        //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8]
        let data: [u8; 72] = [
            2, 0, 1, 2, 8, 0, 1, 3, 0x00, 0x01, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00,
            0x40, 0x40, 0x00, 0x00, 0x80, 0x40, 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40,
            0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, 0x00, 0x00, 0x10, 0x41, 0x00, 0x00,
            0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, 0x00, 0x00, 0x50, 0x41,
            0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x11, 0x80, 0x41,
        ];
        let storage_provider = VirtualStorageProvider::new_memory();
        {
            let mut writer = storage_provider.create_for_write(file_name).unwrap();
            writer.write_all(&data).unwrap();
        }

        let mut reader =
            CachedReader::<VirtualStorageProvider<MemoryFS>>::new(file_name, 8, &storage_provider)
                .unwrap();
        assert_eq!(reader.get_file_size(), 72);
        assert_eq!(reader.cache_size, 8);

        let mut all_from_cache_buf = vec![0; 4];
        reader.read(all_from_cache_buf.as_mut_slice()).unwrap();
        assert_eq!(all_from_cache_buf, [2, 0, 1, 2]);
        assert_eq!(reader.cur_off, 4);

        let mut partial_from_cache_buf = vec![0; 6];
        reader.read(partial_from_cache_buf.as_mut_slice()).unwrap();
        assert_eq!(partial_from_cache_buf, [8, 0, 1, 3, 0x00, 0x01]);
        assert_eq!(reader.cur_off, 0);

        let mut over_cache_size_buf = vec![0; 60];
        reader.read(over_cache_size_buf.as_mut_slice()).unwrap();
        assert_eq!(
            over_cache_size_buf,
            [
                0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40,
                0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00,
                0x00, 0x41, 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41,
                0x00, 0x00, 0x40, 0x41, 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00,
                0x70, 0x41, 0x00, 0x11
            ]
        );

        let mut remaining_less_than_cache_size_buf = vec![0; 2];
        reader
            .read(remaining_less_than_cache_size_buf.as_mut_slice())
            .unwrap();
        assert_eq!(remaining_less_than_cache_size_buf, [0x80, 0x41]);
        assert_eq!(reader.cur_off, reader.cache_size);

        storage_provider
            .delete(file_name)
            .expect("Failed to delete file");
    }

    #[test]
    #[should_panic(expected = "Reading beyond end of file")]
    fn failed_for_reading_beyond_end_of_file() {
        let file_name = "/failed_for_reading_beyond_end_of_file_test_2.bin";
        //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8]
        let data: [u8; 72] = [
            2, 0, 1, 2, 8, 0, 1, 3, 0x00, 0x01, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00,
            0x40, 0x40, 0x00, 0x00, 0x80, 0x40, 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40,
            0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, 0x00, 0x00, 0x10, 0x41, 0x00, 0x00,
            0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, 0x00, 0x00, 0x50, 0x41,
            0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x11, 0x80, 0x41,
        ];
        let storage_provider = VirtualStorageProvider::new_memory();
        {
            let mut writer = storage_provider.create_for_write(file_name).unwrap();
            writer.write_all(&data).unwrap();
        }

        let mut reader =
            CachedReader::<VirtualStorageProvider<MemoryFS>>::new(file_name, 8, &storage_provider)
                .unwrap();
        storage_provider
            .delete(file_name)
            .expect("Failed to delete file");

        let mut over_size_buf = vec![0; 73];
        reader.read(over_size_buf.as_mut_slice()).unwrap();
    }
}