Storage: allow lookup of multiple embeddings at the same time

Author: danieldk

Cargo.toml

@@ -39,9 +39,10 @@ default = ["memmap"]
 
 [dev-dependencies]
 approx = "0.3"
-maplit = "1"
-lazy_static = "1"
 criterion = "0.3"
+lazy_static = "1"
+maplit = "1"
+tempfile = "3"
 
 [[bench]]
 name = "array"

src/chunks/storage/array.rs

@@ -3,7 +3,7 @@ use std::io::{Read, Seek, SeekFrom, Write};
 use std::mem::size_of;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-use ndarray::{Array2, ArrayView2, ArrayViewMut2, CowArray, Ix1};
+use ndarray::{Array2, ArrayView2, ArrayViewMut2, Axis, CowArray, Ix1};
 
 use super::{Storage, StorageView, StorageViewMut};
 use crate::chunks::io::{ChunkIdentifier, ReadChunk, TypeId, WriteChunk};
@@ -24,7 +24,7 @@ mod mmap {
     use byteorder::ByteOrder;
     use byteorder::{LittleEndian, ReadBytesExt};
     use memmap::{Mmap, MmapOptions};
-    use ndarray::{ArrayView2, CowArray, Ix1};
+    use ndarray::{Array2, ArrayView2, Axis, CowArray, Ix1};
     use ndarray::{Dimension, Ix2};
 
     #[cfg(target_endian = "little")]
@@ -65,6 +65,27 @@ mod mmap {
             CowArray::from(embedding)
         }
 
+        #[allow(clippy::let_and_return)]
+        fn embeddings(&self, indices: &[usize]) -> Array2<f32> {
+            #[allow(clippy::cast_ptr_alignment,unused_mut)]
+            let embeddings =
+            // Alignment is ok, padding guarantees that the pointer is at
+            // a multiple of 4.
+		unsafe { ArrayView2::from_shape_ptr(self.shape, self.map.as_ptr() as *const f32) };
+
+            #[allow(unused_mut)]
+            let mut selected_embeddings = embeddings.select(Axis(0), indices);
+
+            #[cfg(target_endian = "big")]
+            LittleEndian::from_slice_f32(
+                selected_embeddings
+                    .as_slice_mut()
+                    .expect("Cannot borrow matrix as mutable slice"),
+            );
+
+            selected_embeddings
+        }
+
         fn shape(&self) -> (usize, usize) {
             self.shape.into_pattern()
         }
@@ -245,6 +266,10 @@ impl Storage for NdArray {
         CowArray::from(self.inner.row(idx))
     }
 
+    fn embeddings(&self, indices: &[usize]) -> Array2<f32> {
+        self.inner.select(Axis(0), indices)
+    }
+
     fn shape(&self) -> (usize, usize) {
         self.inner.dim()
     }
@@ -321,7 +346,7 @@ mod tests {
     use ndarray::Array2;
 
     use crate::chunks::io::{ReadChunk, WriteChunk};
-    use crate::chunks::storage::{NdArray, StorageView};
+    use crate::chunks::storage::{NdArray, Storage, StorageView};
 
     const N_ROWS: usize = 100;
     const N_COLS: usize = 100;
@@ -342,6 +367,19 @@ mod tests {
         read.read_u64::<LittleEndian>().unwrap()
     }
 
+    #[test]
+    fn embeddings_returns_expected_embeddings() {
+        const CHECK_INDICES: &[usize] = &[0, 50, 99, 0];
+
+        let check_arr = test_ndarray();
+
+        let embeddings = check_arr.embeddings(CHECK_INDICES);
+
+        for (embedding, &idx) in embeddings.outer_iter().zip(CHECK_INDICES) {
+            assert_eq!(embedding, check_arr.embedding(idx));
+        }
+    }
+
     #[test]
     fn ndarray_correct_chunk_size() {
         let check_arr = test_ndarray();
@@ -365,4 +403,36 @@ mod tests {
         let arr = NdArray::read_chunk(&mut cursor).unwrap();
         assert_eq!(arr.view(), check_arr.view());
     }
+
+    #[cfg(feature = "memmap")]
+    mod mmap {
+        use std::io::{BufReader, BufWriter, Seek, SeekFrom};
+
+        use tempfile::tempfile;
+
+        use super::test_ndarray;
+        use crate::chunks::io::{MmapChunk, WriteChunk};
+        use crate::chunks::storage::{MmapArray, Storage};
+
+        fn test_mmap_array() -> MmapArray {
+            let array = test_ndarray();
+            let mut tmp = tempfile().unwrap();
+            array.write_chunk(&mut BufWriter::new(&mut tmp)).unwrap();
+            tmp.seek(SeekFrom::Start(0)).unwrap();
+            MmapArray::mmap_chunk(&mut BufReader::new(tmp)).unwrap()
+        }
+
+        #[test]
+        fn embeddings_returns_expected_embeddings() {
+            const CHECK_INDICES: &[usize] = &[0, 50, 99, 0];
+
+            let check_arr = test_mmap_array();
+
+            let embeddings = check_arr.embeddings(CHECK_INDICES);
+
+            for (embedding, &idx) in embeddings.outer_iter().zip(CHECK_INDICES) {
+                assert_eq!(embedding, check_arr.embedding(idx));
+            }
+        }
+    }
 }

src/chunks/storage/mod.rs

@@ -1,6 +1,6 @@
 //! Embedding matrix representations.
 
-use ndarray::{ArrayView2, ArrayViewMut2, CowArray, Ix1};
+use ndarray::{Array2, ArrayView2, ArrayViewMut2, CowArray, Ix1};
 
 mod array;
 #[cfg(feature = "memmap")]
@@ -23,6 +23,9 @@ pub use self::wrappers::{StorageViewWrap, StorageWrap};
 pub trait Storage {
     fn embedding(&self, idx: usize) -> CowArray<f32, Ix1>;
 
+    /// Retrieve multiple embeddings.
+    fn embeddings(&self, indices: &[usize]) -> Array2<f32>;
+
     fn shape(&self) -> (usize, usize);
 }
 

src/chunks/storage/quantized.rs

@@ -3,7 +3,7 @@ use std::io::{Read, Seek, SeekFrom, Write};
 use std::mem::size_of;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-use ndarray::{Array, Array1, Array2, ArrayView1, ArrayView2, CowArray, IntoDimension, Ix1};
+use ndarray::{Array, Array1, Array2, ArrayView1, ArrayView2, Axis, CowArray, IntoDimension, Ix1};
 use rand::{RngCore, SeedableRng};
 use rand_xorshift::XorShiftRng;
 use reductive::pq::{QuantizeVector, ReconstructVector, TrainPQ, PQ};
@@ -257,6 +257,21 @@ impl Storage for QuantizedArray {
         CowArray::from(reconstructed)
     }
 
+    fn embeddings(&self, indices: &[usize]) -> Array2<f32> {
+        let quantized_select = self.quantized_embeddings.select(Axis(0), indices);
+        let mut reconstructed = self.quantizer.reconstruct_batch(quantized_select);
+
+        if let Some(ref norms) = self.norms {
+            let norms_select = norms.select(Axis(0), indices);
+            reconstructed *= &norms_select
+                .view()
+                .into_shape((norms_select.len(), 1))
+                .unwrap();
+        }
+
+        reconstructed
+    }
+
     fn shape(&self) -> (usize, usize) {
         (
             self.quantized_embeddings.nrows(),
@@ -451,7 +466,7 @@ mod mmap {
     use std::io::{BufReader, Seek, SeekFrom, Write};
 
     use memmap::{Mmap, MmapOptions};
-    use ndarray::{Array1, ArrayView2, CowArray, Ix1};
+    use ndarray::{Array1, Array2, ArrayView2, Axis, CowArray, Ix1};
     use reductive::pq::{QuantizeVector, ReconstructVector, PQ};
 
     use super::{PQRead, QuantizedArray, Storage};
@@ -530,6 +545,23 @@ mod mmap {
             CowArray::from(reconstructed)
         }
 
+        fn embeddings(&self, indices: &[usize]) -> Array2<f32> {
+            let quantized = unsafe { self.quantized_embeddings() };
+
+            let quantized_select = quantized.select(Axis(0), indices);
+            let mut reconstructed = self.quantizer.reconstruct_batch(quantized_select);
+
+            if let Some(ref norms) = self.norms {
+                let norms_select = norms.select(Axis(0), indices);
+                reconstructed *= &norms_select
+                    .view()
+                    .into_shape((norms_select.len(), 1))
+                    .unwrap();
+            }
+
+            reconstructed
+        }
+
         fn shape(&self) -> (usize, usize) {
             (
                 self.quantized_embeddings.len() / self.quantizer.quantized_len(),
@@ -654,6 +686,19 @@ mod tests {
         }
     }
 
+    #[test]
+    fn embeddings_returns_expected_embeddings() {
+        const CHECK_INDICES: &[usize] = &[0, 50, 99, 0];
+
+        let check_arr = test_quantized_array(true);
+
+        let embeddings = check_arr.embeddings(CHECK_INDICES);
+
+        for (embedding, &idx) in embeddings.outer_iter().zip(CHECK_INDICES) {
+            assert_eq!(embedding, check_arr.embedding(idx));
+        }
+    }
+
     #[test]
     fn quantized_array_correct_chunk_size() {
         let check_arr = test_quantized_array(false);
@@ -700,64 +745,96 @@ mod tests {
         storage_eq(&quantized, &reconstructed);
     }
 
-    #[test]
     #[cfg(feature = "memmap")]
-    fn mmap_quantized_array() {
-        use crate::chunks::io::MmapChunk;
-        use crate::chunks::storage::MmapQuantizedArray;
-        use std::fs::File;
-        use std::io::BufReader;
-
-        let mut storage_read =
-            BufReader::new(File::open("testdata/quantized_storage.bin").unwrap());
-        let check_arr = QuantizedArray::read_chunk(&mut storage_read).unwrap();
-
-        // Memory map matrix.
-        storage_read.seek(SeekFrom::Start(0)).unwrap();
-        let arr = MmapQuantizedArray::mmap_chunk(&mut storage_read).unwrap();
-
-        // Check
-        storage_eq(&arr, &check_arr);
-    }
+    mod mmap {
+        use std::io::{BufReader, BufWriter, Cursor, Seek, SeekFrom};
+
+        use tempfile::tempfile;
+
+        use super::{storage_eq, test_quantized_array};
+        use crate::chunks::io::{MmapChunk, ReadChunk, WriteChunk};
+        use crate::chunks::storage::quantized::Reconstruct;
+        use crate::chunks::storage::{MmapQuantizedArray, QuantizedArray, Storage};
+
+        fn test_mmap_quantized_array(norms: bool) -> MmapQuantizedArray {
+            let quantized = test_quantized_array(norms);
+            let mut tmp = tempfile().unwrap();
+            quantized
+                .write_chunk(&mut BufWriter::new(&mut tmp))
+                .unwrap();
+            tmp.seek(SeekFrom::Start(0)).unwrap();
+            MmapQuantizedArray::mmap_chunk(&mut BufReader::new(tmp)).unwrap()
+        }
 
-    #[test]
-    #[cfg(feature = "memmap")]
-    fn reconstruct_mmap_quantized_array() {
-        use std::fs::File;
-        use std::io::BufReader;
+        #[test]
+        fn embeddings_returns_expected_embeddings() {
+            const CHECK_INDICES: &[usize] = &[0, 50, 99, 0];
 
-        use crate::chunks::io::MmapChunk;
-        use crate::chunks::storage::MmapQuantizedArray;
+            let check_arr = test_mmap_quantized_array(true);
 
-        let mut storage_read =
-            BufReader::new(File::open("testdata/quantized_storage.bin").unwrap());
-        let quantized = MmapQuantizedArray::mmap_chunk(&mut storage_read).unwrap();
-        let reconstructed = quantized.reconstruct();
-        storage_eq(&quantized, &reconstructed);
-    }
+            let embeddings = check_arr.embeddings(CHECK_INDICES);
 
-    #[test]
-    #[cfg(feature = "memmap")]
-    fn write_mmap_quantized_array() {
-        use crate::chunks::io::MmapChunk;
-        use crate::chunks::storage::MmapQuantizedArray;
-        use std::fs::File;
-        use std::io::BufReader;
-
-        // Memory map matrix.
-        let mut storage_read =
-            BufReader::new(File::open("testdata/quantized_storage.bin").unwrap());
-        let check_arr = MmapQuantizedArray::mmap_chunk(&mut storage_read).unwrap();
-
-        // Write matrix
-        let mut cursor = Cursor::new(Vec::new());
-        check_arr.write_chunk(&mut cursor).unwrap();
+            for (embedding, &idx) in embeddings.outer_iter().zip(CHECK_INDICES) {
+                assert_eq!(embedding, check_arr.embedding(idx));
+            }
+        }
 
-        // Read using non-mmap'ed reader.
-        cursor.seek(SeekFrom::Start(0)).unwrap();
-        let arr = QuantizedArray::read_chunk(&mut cursor).unwrap();
+        #[test]
+        fn reconstruct_mmap_quantized_array() {
+            use std::fs::File;
+            use std::io::BufReader;
+
+            use crate::chunks::io::MmapChunk;
+            use crate::chunks::storage::MmapQuantizedArray;
 
-        // Check
-        storage_eq(&arr, &check_arr);
+            let mut storage_read =
+                BufReader::new(File::open("testdata/quantized_storage.bin").unwrap());
+            let quantized = MmapQuantizedArray::mmap_chunk(&mut storage_read).unwrap();
+            let reconstructed = quantized.reconstruct();
+            storage_eq(&quantized, &reconstructed);
+        }
+
+        #[test]
+        fn mmap_quantized_array() {
+            use crate::chunks::io::MmapChunk;
+            use crate::chunks::storage::MmapQuantizedArray;
+            use std::fs::File;
+            use std::io::BufReader;
+
+            let mut storage_read =
+                BufReader::new(File::open("testdata/quantized_storage.bin").unwrap());
+            let check_arr = QuantizedArray::read_chunk(&mut storage_read).unwrap();
+
+            // Memory map matrix.
+            storage_read.seek(SeekFrom::Start(0)).unwrap();
+            let arr = MmapQuantizedArray::mmap_chunk(&mut storage_read).unwrap();
+
+            // Check
+            storage_eq(&arr, &check_arr);
+        }
+
+        #[test]
+        fn write_mmap_quantized_array() {
+            use crate::chunks::io::MmapChunk;
+            use crate::chunks::storage::MmapQuantizedArray;
+            use std::fs::File;
+            use std::io::BufReader;
+
+            // Memory map matrix.
+            let mut storage_read =
+                BufReader::new(File::open("testdata/quantized_storage.bin").unwrap());
+            let check_arr = MmapQuantizedArray::mmap_chunk(&mut storage_read).unwrap();
+
+            // Write matrix
+            let mut cursor = Cursor::new(Vec::new());
+            check_arr.write_chunk(&mut cursor).unwrap();
+
+            // Read using non-mmap'ed reader.
+            cursor.seek(SeekFrom::Start(0)).unwrap();
+            let arr = QuantizedArray::read_chunk(&mut cursor).unwrap();
+
+            // Check
+            storage_eq(&arr, &check_arr);
+        }
     }
 }