feat: add download threshold for faster index scan

Author: darkskygit

src/buffer.rs

@@ -36,15 +36,17 @@ type Fetch = Box<dyn Fn(usize, usize) -> Result<Vec<u8>> + Send + Sync>;
 pub struct LazyBuffer {
     blocks: HashMap<usize, LazyBlock>,
     block_size: usize,
+    download_threshold: usize,
     length: usize,
     fetch: Arc<Fetch>,
 }
 
 impl LazyBuffer {
-    pub fn new(length: usize, block_size: usize, fetch: Fetch) -> Self {
+    pub fn new(length: usize, block_size: usize, download_threshold: usize, fetch: Fetch) -> Self {
         Self {
             blocks: HashMap::new(),
             block_size,
+            download_threshold,
             length,
             fetch: Arc::new(fetch),
         }
@@ -74,7 +76,16 @@ impl LazyBuffer {
         }
         let block_index = offset / self.block_size;
         let block_offset = offset - block_index * self.block_size;
-        self.get_block(block_index)?.read(buf, block_offset)?;
+        if !self.blocks.contains_key(&block_index) && buf.len() <= self.download_threshold {
+            // skip full block fetch if buf is smaller than download_threshold
+            let data = (self.fetch)(offset, buf.len())?;
+            if data.len() != buf.len() {
+                return Err(Error::new(ErrorKind::UnexpectedEof, "read out of bounds"));
+            }
+            buf.copy_from_slice(&data);
+        } else {
+            self.get_block(block_index)?.read(buf, block_offset)?;
+        }
         Ok(())
     }
 }
@@ -92,7 +103,7 @@ mod tests {
         let block_size = 16;
         let total_size = 64;
 
-        let mut lazy_buffer = LazyBuffer::new(total_size, block_size, Box::new(mock_fetch));
+        let mut lazy_buffer = LazyBuffer::new(total_size, block_size, 0, Box::new(mock_fetch));
 
         let mut buf1 = vec![0u8; 16];
         lazy_buffer.read(&mut buf1, 0).unwrap();

src/conn.rs

@@ -53,11 +53,12 @@ impl Connection {
         }
     }
 
-    pub fn new(url: &str, block_size: usize) -> Result<Self, Error> {
+    pub fn new(url: &str, block_size: usize, download_threshold: usize) -> Result<Self, Error> {
         let (rt, length) = Self::init_with_url(url)?;
         let buffer = LazyBuffer::new(
             length,
             block_size,
+            download_threshold,
             Box::new({
                 let url = url.to_string();
                 let rt = rt.clone();

src/vfs.rs

@@ -9,16 +9,8 @@ use std::{
 pub const HTTP_VFS: &str = "http";
 
 pub struct HttpVfs {
-    block_size: usize,
-}
-
-impl Default for HttpVfs {
-    fn default() -> Self {
-        Self {
-            // 4MB block size
-            block_size: 1024 * 4096,
-        }
-    }
+    pub(crate) block_size: usize,
+    pub(crate) download_threshold: usize,
 }
 
 impl Vfs for HttpVfs {
@@ -32,7 +24,11 @@ impl Vfs for HttpVfs {
             ));
         }
 
-        Ok(Connection::new(db, self.block_size)?)
+        Ok(Connection::new(
+            db,
+            self.block_size,
+            self.download_threshold,
+        )?)
     }
 
     fn delete(&self, _db: &str) -> Result<(), Error> {