feat: allow lazily chunking unsorted iteration

benches/my_bench.rs

@@ -95,7 +95,7 @@ fn main() {
 
         // Open finished file
         let all_items = if unsorted_read {
-            UnsortedShardReader::<T1>::open(tmp.path())?.collect::<Result<_, _>>()?
+            UnsortedShardReader::<T1>::open(tmp.path()).collect::<Result<_, _>>()?
         } else {
             let reader = ShardReader::<T1>::open(tmp.path())?;
 

src/lib.rs

@@ -1853,10 +1853,20 @@ mod shard_tests {
             }
 
             // Check the unsorted read
+            assert_eq!(n_items, UnsortedShardReader::<T1>::len(&[tmp.path()])?);
             let unsorted_reader = UnsortedShardReader::<T1>::open(tmp.path());
             let all_items_res: Result<Vec<_>, Error> = unsorted_reader.collect();
             let all_items = all_items_res?;
             assert!(set_compare(&true_items, &all_items));
+
+            // Check that the skip_lazy feature produces the expected results.
+            let mut unsorted_reader_skip = UnsortedShardReader::<T1>::open(tmp.path());
+            let to_skip = (disk_chunk_size * 3) + 1;
+            let skipped = unsorted_reader_skip.skip_lazy(to_skip)?;
+            assert_eq!(to_skip, skipped);
+            let all_items_res_skip: Result<Vec<_>, Error> = unsorted_reader_skip.collect();
+            let all_items_skip = all_items_res_skip?;
+            assert_eq!(&all_items[to_skip..], &all_items_skip);
         }
         Ok(())
     }
@@ -2117,5 +2127,11 @@ mod shard_tests {
         let shard_files = Vec::<PathBuf>::new();
         let reader = UnsortedShardReader::<u8>::open_set(&shard_files);
         assert_eq!(reader.count(), 0);
+
+        // Test that skipping an empty set works correctly.
+        let mut reader = UnsortedShardReader::<u8>::open_set(&shard_files);
+        let skipped = reader.skip_lazy(10).unwrap();
+        assert_eq!(0, skipped);
+        assert_eq!(reader.count(), 0);
     }
 }

src/unsorted.rs

@@ -69,13 +69,28 @@ where
         }
     }
 
+    /// Compute the total number of elements in this set of shard files.
+    pub fn len<P: AsRef<Path>>(shard_files: &[P]) -> Result<usize, Error>
+    where
+        <S as SortKey<T>>::Key: 'static,
+    {
+        // Create a set reader, and consume all the files, just getting counts.
+        let files_reader = Self::open_set(shard_files);
+        let mut count = 0;
+        for file in files_reader.shard_reader_iter {
+            let file = file?;
+            count += file.len();
+        }
+        Ok(count)
+    }
+
     /// Skip the first count items that would be returned by iteration.
     /// This avoids reading anything into memory besides the indexes for each
     /// file, and the first chunk that wouldn't be skipped.
     ///
     /// Returns the number of items skipped. This will only ever be less than
     /// count if we exhausted all shard files.
-    pub fn skip(&mut self, count: usize) -> Result<usize, Error> {
+    pub fn skip_lazy(&mut self, count: usize) -> Result<usize, Error> {
         let mut skipped = 0;
         loop {
             let Some(mut file_reader) = self.take_active_reader().transpose()? else {
@@ -97,6 +112,7 @@ where
         Ok(skipped)
     }
 
+    /// Take the current active shard reader, or advance to the next.
     fn take_active_reader(&mut self) -> Option<Result<UnsortedShardFileReader<T, S>, Error>> {
         let reader = if let Some(reader) = self.active_shard_reader.take() {
             reader
@@ -160,6 +176,7 @@ struct UnsortedShardFileReader<T, S = DefaultSort>
 where
     S: SortKey<T>,
 {
+    count: usize,
     file_index_iter: Box<dyn Iterator<Item = KeylessShardRecord>>,
     shard_iter: Option<UnsortedShardIter<T>>,
     phantom: PhantomData<S>,
@@ -178,6 +195,7 @@ where
         <S as SortKey<T>>::Key: 'static,
     {
         let reader = ShardReaderSingle::<T, S>::open(path)?;
+        let count = reader.len();
         let mut file_index_iter = reader.index.into_iter().map(|r| KeylessShardRecord {
             offset: r.offset,
             len_bytes: r.len_bytes,
@@ -188,12 +206,18 @@ where
             return Ok(None);
         };
         Ok(Some(Self {
+            count,
             shard_iter: Some(UnsortedShardIter::new(reader.file, first_index_entry)?),
             file_index_iter: Box::new(file_index_iter),
             phantom: Default::default(),
         }))
     }
 
+    /// Return the total number of items in this shard file.
+    pub fn len(&self) -> usize {
+        self.count
+    }
+
     /// Get the next item out of this file.
     pub fn next(&mut self) -> Result<Option<T>, Error> {
         loop {
@@ -240,6 +264,11 @@ where
                     skipped,
                     exhausted: false,
                 });
+            } else {
+                // Advance to the next shard, if there is one.
+                if let Some(next_index_record) = self.file_index_iter.next() {
+                    self.shard_iter = Some(shard_iter.reset(next_index_record)?);
+                }
             }
         }
     }