Merge pull request #319 from meilisearch/dup-sort-custom-comparator

Add Dup Sort Custom Comparator support

Author: Kerollmops

examples/custom-dupsort-comparator.rs

@@ -0,0 +1,62 @@
+use std::cmp::Ordering;
+use std::error::Error;
+use std::fs;
+use std::path::Path;
+
+use byteorder::BigEndian;
+use heed::{DatabaseFlags, EnvOpenOptions};
+use heed_traits::Comparator;
+use heed_types::{Str, U128};
+
+enum DescendingIntCmp {}
+
+impl Comparator for DescendingIntCmp {
+    fn compare(a: &[u8], b: &[u8]) -> Ordering {
+        a.cmp(b).reverse()
+    }
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let env_path = Path::new("target").join("custom-dupsort-cmp.mdb");
+
+    let _ = fs::remove_dir_all(&env_path);
+
+    fs::create_dir_all(&env_path)?;
+    let env = unsafe {
+        EnvOpenOptions::new()
+            .map_size(10 * 1024 * 1024) // 10MB
+            .max_dbs(3)
+            .open(env_path)?
+    };
+
+    let mut wtxn = env.write_txn()?;
+    let db = env
+        .database_options()
+        .types::<Str, U128<BigEndian>>()
+        .flags(DatabaseFlags::DUP_SORT)
+        .dup_sort_comparator::<DescendingIntCmp>()
+        .create(&mut wtxn)?;
+    wtxn.commit()?;
+
+    let mut wtxn = env.write_txn()?;
+
+    // We fill our database with entries.
+    db.put(&mut wtxn, "1", &1)?;
+    db.put(&mut wtxn, "1", &2)?;
+    db.put(&mut wtxn, "1", &3)?;
+    db.put(&mut wtxn, "2", &4)?;
+    db.put(&mut wtxn, "1", &5)?;
+    db.put(&mut wtxn, "0", &0)?;
+
+    // We check that the keys are in lexicographic and values in descending order.
+    let mut iter = db.iter(&wtxn)?;
+    assert_eq!(iter.next().transpose()?, Some(("0", 0)));
+    assert_eq!(iter.next().transpose()?, Some(("1", 5)));
+    assert_eq!(iter.next().transpose()?, Some(("1", 3)));
+    assert_eq!(iter.next().transpose()?, Some(("1", 2)));
+    assert_eq!(iter.next().transpose()?, Some(("1", 1)));
+    assert_eq!(iter.next().transpose()?, Some(("2", 4)));
+    drop(iter);
+
+    Ok(())
+}

heed/Cargo.toml

@@ -131,6 +131,10 @@ path = "../examples/cursor-append.rs"
 name = "custom-comparator"
 path = "../examples/custom-comparator.rs"
 
+[[example]]
+name = "custom-dupsort-comparator"
+path = "../examples/custom-dupsort-comparator.rs"
+
 [[example]]
 name = "multi-env"
 path = "../examples/multi-env.rs"

heed/src/cookbook.rs

@@ -5,6 +5,8 @@
 //! - [Create Custom and Prefix Codecs](#create-custom-and-prefix-codecs)
 //! - [Change the Environment Size Dynamically](#change-the-environment-size-dynamically)
 //! - [Advanced Multithreaded Access of Entries](#advanced-multithreaded-access-of-entries)
+//! - [Use Custom Key Comparator](#use-custom-key-comparator)
+//! - [Use Custom Dupsort Comparator](#use-custom-dupsort-comparator)
 //!
 //! # Decode Values on Demand
 //!
@@ -441,9 +443,146 @@
 //! unsafe impl Sync for ImmutableMap<'_> {}
 //! ```
 //!
+//! # Use Custom Key Comparator
+//!
+//! LMDB keys are sorted in lexicographic order by default. To change this behavior
+//! you can implement a custom [`Comparator`] and provide it when creating the database.
+//!
+//! Under the hood this translates into a [`mdb_set_compare`] call.
+//!
+//! ```
+//! use std::cmp::Ordering;
+//! use std::error::Error;
+//! use std::str;
+//!
+//! use heed::EnvOpenOptions;
+//! use heed_traits::Comparator;
+//! use heed_types::{Str, Unit};
+//!
+//! enum StringAsIntCmp {}
+//!
+//! // This function takes two strings which represent integers,
+//! // parses them into i32s and compare the parsed value.
+//! // Therefore "-1000" < "-100" must be true even without '0' padding.
+//! impl Comparator for StringAsIntCmp {
+//!     fn compare(a: &[u8], b: &[u8]) -> Ordering {
+//!         let a: i32 = str::from_utf8(a).unwrap().parse().unwrap();
+//!         let b: i32 = str::from_utf8(b).unwrap().parse().unwrap();
+//!         a.cmp(&b)
+//!     }
+//! }
+//!
+//! fn main() -> Result<(), Box<dyn Error>> {
+//!     let path = tempfile::tempdir()?;
+//!
+//!     let env = unsafe {
+//!         EnvOpenOptions::new()
+//!             .map_size(10 * 1024 * 1024) // 10MB
+//!             .max_dbs(3)
+//!             .open(path)?
+//!     };
+//!
+//!     let mut wtxn = env.write_txn()?;
+//!     let db = env
+//!         .database_options()
+//!         .types::<Str, Unit>()
+//!         .key_comparator::<StringAsIntCmp>()
+//!         .create(&mut wtxn)?;
+//!     wtxn.commit()?;
+//!
+//!     let mut wtxn = env.write_txn()?;
+//!
+//!     // We fill our database with entries.
+//!     db.put(&mut wtxn, "-100000", &())?;
+//!     db.put(&mut wtxn, "-10000", &())?;
+//!     db.put(&mut wtxn, "-1000", &())?;
+//!     db.put(&mut wtxn, "-100", &())?;
+//!     db.put(&mut wtxn, "100", &())?;
+//!
+//!     // We check that the key are in the right order ("-100" < "-1000" < "-10000"...)
+//!     let mut iter = db.iter(&wtxn)?;
+//!     assert_eq!(iter.next().transpose()?, Some(("-100000", ())));
+//!     assert_eq!(iter.next().transpose()?, Some(("-10000", ())));
+//!     assert_eq!(iter.next().transpose()?, Some(("-1000", ())));
+//!     assert_eq!(iter.next().transpose()?, Some(("-100", ())));
+//!     assert_eq!(iter.next().transpose()?, Some(("100", ())));
+//!     drop(iter);
+//!
+//!     Ok(())
+//! }
+//! ```
+//!
+//! # Use Custom Dupsort Comparator
+//!
+//! When using DUPSORT LMDB sorts values of the same key in lexicographic order by default.
+//! To change this behavior you can implement a custom [`Comparator`] and provide it when
+//! creating the database.
+//!
+//! Under the hood this translates into a [`mdb_set_dupsort`] call.
+//!
+//! ```
+//! use std::cmp::Ordering;
+//! use std::error::Error;
+//!
+//! use byteorder::BigEndian;
+//! use heed::{DatabaseFlags, EnvOpenOptions};
+//! use heed_traits::Comparator;
+//! use heed_types::{Str, U128};
+//!
+//! enum DescendingIntCmp {}
+//!
+//! impl Comparator for DescendingIntCmp {
+//!     fn compare(a: &[u8], b: &[u8]) -> Ordering {
+//!         a.cmp(&b).reverse()
+//!     }
+//! }
+//!
+//! fn main() -> Result<(), Box<dyn Error>> {
+//!     let path = tempfile::tempdir()?;
+//!
+//!     let env = unsafe {
+//!         EnvOpenOptions::new()
+//!             .map_size(10 * 1024 * 1024) // 10MB
+//!             .max_dbs(3)
+//!             .open(path)?
+//!     };
+//!
+//!     let mut wtxn = env.write_txn()?;
+//!     let db = env
+//!         .database_options()
+//!         .types::<Str, U128<BigEndian>>()
+//!         .flags(DatabaseFlags::DUP_SORT)
+//!         .dup_sort_comparator::<DescendingIntCmp>()
+//!         .create(&mut wtxn)?;
+//!     wtxn.commit()?;
+//!
+//!     let mut wtxn = env.write_txn()?;
+//!
+//!     // We fill our database with entries.
+//!     db.put(&mut wtxn, "1", &1)?;
+//!     db.put(&mut wtxn, "1", &2)?;
+//!     db.put(&mut wtxn, "1", &3)?;
+//!     db.put(&mut wtxn, "2", &4)?;
+//!     db.put(&mut wtxn, "1", &5)?;
+//!     db.put(&mut wtxn, "0", &0)?;
+//!
+//!     // We check that the keys are in lexicographic and values in descending order.
+//!     let mut iter = db.iter(&wtxn)?;
+//!     assert_eq!(iter.next().transpose()?, Some(("0", 0)));
+//!     assert_eq!(iter.next().transpose()?, Some(("1", 5)));
+//!     assert_eq!(iter.next().transpose()?, Some(("1", 3)));
+//!     assert_eq!(iter.next().transpose()?, Some(("1", 2)));
+//!     assert_eq!(iter.next().transpose()?, Some(("1", 1)));
+//!     assert_eq!(iter.next().transpose()?, Some(("2", 4)));
+//!     drop(iter);
+//!
+//!     Ok(())
+//! }
+//! ```
+//!
 
 // To let cargo generate doc links
 #![allow(unused_imports)]
 
-use crate::envs::EnvOpenOptions;
-use crate::{BytesDecode, BytesEncode, Database};
+use crate::mdb::ffi::{mdb_set_compare, mdb_set_dupsort};
+use crate::{BytesDecode, BytesEncode, Comparator, Database, EnvOpenOptions};

heed/src/databases/database.rs

@@ -55,9 +55,9 @@ use crate::*;
 /// # Ok(()) }
 /// ```
 #[derive(Debug)]
-pub struct DatabaseOpenOptions<'e, 'n, T, KC, DC, C = DefaultComparator> {
+pub struct DatabaseOpenOptions<'e, 'n, T, KC, DC, C = DefaultComparator, CDUP = DefaultComparator> {
     env: &'e Env<T>,
-    types: marker::PhantomData<(KC, DC, C)>,
+    types: marker::PhantomData<(KC, DC, C, CDUP)>,
     name: Option<&'n str>,
     flags: AllDatabaseFlags,
 }
@@ -74,7 +74,7 @@ impl<'e, T> DatabaseOpenOptions<'e, 'static, T, Unspecified, Unspecified> {
     }
 }
 
-impl<'e, 'n, T, KC, DC, C> DatabaseOpenOptions<'e, 'n, T, KC, DC, C> {
+impl<'e, 'n, T, KC, DC, C, CDUP> DatabaseOpenOptions<'e, 'n, T, KC, DC, C, CDUP> {
     /// Change the type of the database.
     ///
     /// The default types are [`Unspecified`] and require a call to [`Database::remap_types`]
@@ -91,7 +91,19 @@ impl<'e, 'n, T, KC, DC, C> DatabaseOpenOptions<'e, 'n, T, KC, DC, C> {
     /// Change the customized key compare function of the database.
     ///
     /// By default no customized compare function will be set when opening a database.
-    pub fn key_comparator<NC>(self) -> DatabaseOpenOptions<'e, 'n, T, KC, DC, NC> {
+    pub fn key_comparator<NC>(self) -> DatabaseOpenOptions<'e, 'n, T, KC, DC, NC, CDUP> {
+        DatabaseOpenOptions {
+            env: self.env,
+            types: Default::default(),
+            name: self.name,
+            flags: self.flags,
+        }
+    }
+
+    /// Change the customized dup sort compare function of the database.
+    ///
+    /// By default no customized compare function will be set when opening a database.
+    pub fn dup_sort_comparator<NCDUP>(self) -> DatabaseOpenOptions<'e, 'n, T, KC, DC, C, NCDUP> {
         DatabaseOpenOptions {
             env: self.env,
             types: Default::default(),
@@ -132,15 +144,16 @@ impl<'e, 'n, T, KC, DC, C> DatabaseOpenOptions<'e, 'n, T, KC, DC, C> {
     ///
     /// If not done, you might raise `Io(Os { code: 22, kind: InvalidInput, message: "Invalid argument" })`
     /// known as `EINVAL`.
-    pub fn open(&self, rtxn: &RoTxn) -> Result<Option<Database<KC, DC, C>>>
+    pub fn open(&self, rtxn: &RoTxn) -> Result<Option<Database<KC, DC, C, CDUP>>>
     where
         KC: 'static,
         DC: 'static,
         C: Comparator + 'static,
+        CDUP: Comparator + 'static,
     {
         assert_eq_env_txn!(self.env, rtxn);
 
-        match self.env.raw_init_database::<C>(rtxn.txn_ptr(), self.name, self.flags) {
+        match self.env.raw_init_database::<C, CDUP>(rtxn.txn_ptr(), self.name, self.flags) {
             Ok(dbi) => Ok(Some(Database::new(self.env.env_mut_ptr().as_ptr() as _, dbi))),
             Err(Error::Mdb(e)) if e.not_found() => Ok(None),
             Err(e) => Err(e),
@@ -156,29 +169,30 @@ impl<'e, 'n, T, KC, DC, C> DatabaseOpenOptions<'e, 'n, T, KC, DC, C> {
     /// LMDB has an important restriction on the unnamed database when named ones are opened.
     /// The names of the named databases are stored as keys in the unnamed one and are immutable,
     /// and these keys can only be read and not written.
-    pub fn create(&self, wtxn: &mut RwTxn) -> Result<Database<KC, DC, C>>
+    pub fn create(&self, wtxn: &mut RwTxn) -> Result<Database<KC, DC, C, CDUP>>
     where
         KC: 'static,
         DC: 'static,
         C: Comparator + 'static,
+        CDUP: Comparator + 'static,
     {
         assert_eq_env_txn!(self.env, wtxn);
 
         let flags = self.flags | AllDatabaseFlags::CREATE;
-        match self.env.raw_init_database::<C>(wtxn.txn_ptr(), self.name, flags) {
+        match self.env.raw_init_database::<C, CDUP>(wtxn.txn_ptr(), self.name, flags) {
             Ok(dbi) => Ok(Database::new(self.env.env_mut_ptr().as_ptr() as _, dbi)),
             Err(e) => Err(e),
         }
     }
 }
 
-impl<T, KC, DC, C> Clone for DatabaseOpenOptions<'_, '_, T, KC, DC, C> {
+impl<T, KC, DC, C, CDUP> Clone for DatabaseOpenOptions<'_, '_, T, KC, DC, C, CDUP> {
     fn clone(&self) -> Self {
         *self
     }
 }
 
-impl<T, KC, DC, C> Copy for DatabaseOpenOptions<'_, '_, T, KC, DC, C> {}
+impl<T, KC, DC, C, CDUP> Copy for DatabaseOpenOptions<'_, '_, T, KC, DC, C, CDUP> {}
 
 /// A typed database that accepts only the types it was created with.
 ///
@@ -292,14 +306,14 @@ impl<T, KC, DC, C> Copy for DatabaseOpenOptions<'_, '_, T, KC, DC, C> {}
 /// wtxn.commit()?;
 /// # Ok(()) }
 /// ```
-pub struct Database<KC, DC, C = DefaultComparator> {
+pub struct Database<KC, DC, C = DefaultComparator, CDUP = DefaultComparator> {
     pub(crate) env_ident: usize,
     pub(crate) dbi: ffi::MDB_dbi,
-    marker: marker::PhantomData<(KC, DC, C)>,
+    marker: marker::PhantomData<(KC, DC, C, CDUP)>,
 }
 
-impl<KC, DC, C> Database<KC, DC, C> {
-    pub(crate) fn new(env_ident: usize, dbi: ffi::MDB_dbi) -> Database<KC, DC, C> {
+impl<KC, DC, C, CDUP> Database<KC, DC, C, CDUP> {
+    pub(crate) fn new(env_ident: usize, dbi: ffi::MDB_dbi) -> Database<KC, DC, C, CDUP> {
         Database { env_ident, dbi, marker: std::marker::PhantomData }
     }
 
@@ -2693,20 +2707,21 @@ impl<KC, DC, C> Database<KC, DC, C> {
     }
 }
 
-impl<KC, DC, C> Clone for Database<KC, DC, C> {
-    fn clone(&self) -> Database<KC, DC, C> {
+impl<KC, DC, C, CDUP> Clone for Database<KC, DC, C, CDUP> {
+    fn clone(&self) -> Database<KC, DC, C, CDUP> {
         *self
     }
 }
 
-impl<KC, DC, C> Copy for Database<KC, DC, C> {}
+impl<KC, DC, C, CDUP> Copy for Database<KC, DC, C, CDUP> {}
 
-impl<KC, DC, C> fmt::Debug for Database<KC, DC, C> {
+impl<KC, DC, C, CDUP> fmt::Debug for Database<KC, DC, C, CDUP> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_struct("Database")
             .field("key_codec", &any::type_name::<KC>())
             .field("data_codec", &any::type_name::<DC>())
-            .field("comparator", &any::type_name::<C>())
+            .field("key_comparator", &any::type_name::<C>())
+            .field("dup_sort_comparator", &any::type_name::<CDUP>())
             .finish()
     }
 }