persist key/value type metadata and add MassMap::cast

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "massmap"
-version = "0.1.3"
+version = "0.1.4"
 edition = "2024"
 authors = ["SF-Zhou <sfzhou.scut@gmail.com>"]
 homepage = "https://github.com/SF-Zhou/massmap"

README.md

@@ -67,7 +67,7 @@ massmap convert -i examples/demo.json -o examples/demo.massmap --bucket-count 32
 #> {
 #>   "header": {
 #>     "meta_offset": 486,
-#>     "meta_length": 192
+#>     "meta_length": 239
 #>   },
 #>   "meta": {
 #>     "entry_count": 47,
@@ -82,7 +82,9 @@ massmap convert -i examples/demo.json -o examples/demo.massmap --bucket-count 32
 #>       "parameters": {
 #>         "seed": 0
 #>       }
-#>     }
+#>     },
+#>     "key_type": "alloc::string::String",
+#>     "value_type": "serde_json::value::Value"
 #>   }
 #> }
 
@@ -91,7 +93,7 @@ massmap info examples/demo.massmap -k 1999
 #> {
 #>   "header": {
 #>     "meta_offset": 486,
-#>     "meta_length": 192
+#>     "meta_length": 239
 #>   },
 #>   "meta": {
 #>     "entry_count": 47,
@@ -106,15 +108,17 @@ massmap info examples/demo.massmap -k 1999
 #>       "parameters": {
 #>         "seed": 0
 #>       }
-#>     }
+#>     },
+#>     "key_type": "alloc::string::String",
+#>     "value_type": "serde_json::value::Value"
 #>   }
 #> }
-#> 1999: Some(Number(7229))
+#> Get 1999: Some(Number(7229))
 
 # 4. view the raw bytes of the massmap file
 hexdump -C examples/demo.massmap
 #> 00000000  4d 41 53 53 4d 41 50 21  00 00 00 00 00 00 01 e6  |MASSMAP!........|
-#> 00000010  00 00 00 00 00 00 00 c0  92 92 a4 31 39 39 34 cd  |...........1994.|
+#> 00000010  00 00 00 00 00 00 00 ef  92 92 a4 31 39 39 34 cd  |...........1994.|
 #> 00000020  0f f1 92 a4 32 30 32 30  ce 00 01 18 94 91 92 a4  |....2020........|
 #> 00000030  32 30 32 34 ce 00 01 76  05 92 92 a4 31 39 38 39  |2024...v....1989|
 #> 00000040  cd 06 00 92 a4 32 30 30  39 cd 66 44 92 92 a4 31  |.....2009.fD...1|
@@ -143,20 +147,22 @@ hexdump -C examples/demo.massmap
 #> 000001b0  e9 91 92 a4 31 39 38 33  cd 02 4c 91 92 a4 31 39  |....1983..L...19|
 #> 000001c0  38 32 cd 02 15 92 92 a4  31 39 39 30 cd 06 7f 92  |82......1990....|
 #> 000001d0  a4 32 30 31 39 ce 00 01  11 be 91 92 a4 32 30 32  |.2019........202|
-#> 000001e0  32 ce 00 01 4e c2 92 95  2f 20 1b 92 00 20 92 a8  |2...N.../ ... ..|
-#> 000001f0  66 6f 6c 64 68 61 73 68  81 a4 73 65 65 64 00 dc  |foldhash..seed..|
-#> 00000200  00 20 93 18 15 02 93 2d  0c 01 93 39 13 02 93 4c  |. .....-...9...L|
-#> 00000210  13 02 93 5f 25 04 93 cc  84 0a 01 93 cc 8e 1c 03  |..._%...........|
-#> 00000220  93 cc aa 13 02 93 cc bd  0a 01 93 cc c7 0a 01 93  |................|
-#> 00000230  00 00 00 93 00 00 00 93  cc d1 15 02 93 cc e6 13  |................|
-#> 00000240  02 93 00 00 00 93 cc f9  25 04 93 cd 01 1e 0a 01  |........%.......|
-#> 00000250  93 cd 01 28 13 02 93 cd  01 3b 0a 01 93 cd 01 45  |...(.....;.....E|
-#> 00000260  13 02 93 00 00 00 93 cd  01 58 13 02 93 00 00 00  |.........X......|
-#> 00000270  93 cd 01 6b 0c 01 93 cd  01 77 0a 01 93 cd 01 81  |...k.....w......|
-#> 00000280  13 02 93 cd 01 94 13 02  93 cd 01 a7 0a 01 93 cd  |................|
-#> 00000290  01 b1 0a 01 93 cd 01 bb  0a 01 93 cd 01 c5 15 02  |................|
-#> 000002a0  93 cd 01 da 0c 01                                 |......|
-#> 000002a6
+#> 000001e0  32 ce 00 01 4e c2 92 97  2f 20 1b 92 00 20 92 a8  |2...N.../ ... ..|
+#> 000001f0  66 6f 6c 64 68 61 73 68  81 a4 73 65 65 64 00 b5  |foldhash..seed..|
+#> 00000200  61 6c 6c 6f 63 3a 3a 73  74 72 69 6e 67 3a 3a 53  |alloc::string::S|
+#> 00000210  74 72 69 6e 67 b8 73 65  72 64 65 5f 6a 73 6f 6e  |tring.serde_json|
+#> 00000220  3a 3a 76 61 6c 75 65 3a  3a 56 61 6c 75 65 dc 00  |::value::Value..|
+#> 00000230  20 93 18 15 02 93 2d 0c  01 93 39 13 02 93 4c 13  | .....-...9...L.|
+#> 00000240  02 93 5f 25 04 93 cc 84  0a 01 93 cc 8e 1c 03 93  |.._%............|
+#> 00000250  cc aa 13 02 93 cc bd 0a  01 93 cc c7 0a 01 93 00  |................|
+#> 00000260  00 00 93 00 00 00 93 cc  d1 15 02 93 cc e6 13 02  |................|
+#> 00000270  93 00 00 00 93 cc f9 25  04 93 cd 01 1e 0a 01 93  |.......%........|
+#> 00000280  cd 01 28 13 02 93 cd 01  3b 0a 01 93 cd 01 45 13  |..(.....;.....E.|
+#> 00000290  02 93 00 00 00 93 cd 01  58 13 02 93 00 00 00 93  |........X.......|
+#> 000002a0  cd 01 6b 0c 01 93 cd 01  77 0a 01 93 cd 01 81 13  |..k.....w.......|
+#> 000002b0  02 93 cd 01 94 13 02 93  cd 01 a7 0a 01 93 cd 01  |................|
+#> 000002c0  b1 0a 01 93 cd 01 bb 0a  01 93 cd 01 c5 15 02 93  |................|
+#> 000002d0  cd 01 da 0c 01                                    |.....|
 ```
 
 ## Configuration

examples/massmap.rs

@@ -2,9 +2,11 @@ use clap::{Parser, Subcommand};
 use foldhash::fast::FixedState;
 use massmap::{
     MassMap, MassMapBuilder, MassMapDefaultHashLoader, MassMapHashConfig, MassMapHashLoader,
-    MassMapMerger,
+    MassMapInner, MassMapMerger, MassMapReader,
 };
+use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use std::fmt::Display;
 use std::fs::File;
 use std::io::{BufReader, Error, ErrorKind, Result};
 use std::path::{Path, PathBuf};
@@ -116,26 +118,27 @@ impl MassMapHashLoader for MassMapTolerableHashLoader {
     }
 }
 
-fn run_info(args: InfoArgs) -> Result<()> {
-    let file = File::open(&args.input)?;
-
-    let map = MassMap::<String, serde_json::Value, _, MassMapTolerableHashLoader>::load(file)?;
-
-    let json = serde_json::to_string_pretty(&map.info())
-        .map_err(|e| Error::other(format!("Failed to format JSON: {e}")))?;
-    println!("{}", json);
-
-    if let Some(key) = args.key {
-        println!("{}: {:?}", key, map.get(&key)?);
+fn do_query<K, R>(
+    map: MassMap<K, serde_json::Value, R, MassMapTolerableHashLoader>,
+    key: Option<K>,
+    bucket: Option<u64>,
+) -> Result<()>
+where
+    K: Serialize + for<'de> Deserialize<'de> + Display + std::hash::Hash + Eq,
+    R: MassMapReader,
+{
+    if let Some(key) = key {
+        println!("Get {}: {:?}", key, map.get(&key)?);
     }
 
-    if let Some(bucket_index) = args.bucket {
-        if bucket_index >= map.meta.bucket_count {
+    if let Some(bucket_index) = bucket {
+        if bucket_index as usize >= map.bucket_count() {
             return Err(Error::new(
                 ErrorKind::InvalidInput,
                 format!(
                     "Bucket index {} out of range >= {}",
-                    bucket_index, map.meta.bucket_count
+                    bucket_index,
+                    map.bucket_count()
                 ),
             ));
         }
@@ -148,6 +151,56 @@ fn run_info(args: InfoArgs) -> Result<()> {
     Ok(())
 }
 
+fn run_info(args: InfoArgs) -> Result<()> {
+    let file = File::open(&args.input)?;
+
+    let map = MassMapInner::<_, MassMapTolerableHashLoader>::load(file)?;
+
+    let json = serde_json::to_string_pretty(&map.info())
+        .map_err(|e| Error::other(format!("Failed to format JSON: {e}")))?;
+    println!("{}", json);
+
+    match map.meta.key_type.as_str() {
+        "u8" => do_query(
+            map.cast::<u8, _>(),
+            args.key.map(|x| x.parse().unwrap()),
+            args.bucket,
+        )?,
+        "u16" => do_query(
+            map.cast::<u16, _>(),
+            args.key.map(|x| x.parse().unwrap()),
+            args.bucket,
+        )?,
+        "u32" => do_query(
+            map.cast::<u32, _>(),
+            args.key.map(|x| x.parse().unwrap()),
+            args.bucket,
+        )?,
+        "u64" => do_query(
+            map.cast::<u64, _>(),
+            args.key.map(|x| x.parse().unwrap()),
+            args.bucket,
+        )?,
+        "u128" => do_query(
+            map.cast::<u128, _>(),
+            args.key.map(|x| x.parse().unwrap()),
+            args.bucket,
+        )?,
+        _ if map.meta.key_type == std::any::type_name::<String>() => {
+            do_query(map.cast::<String, _>(), args.key, args.bucket)?
+        }
+        _ => {
+            assert!(
+                args.key.is_none() && args.bucket.is_none(),
+                "Unsupported key type: {}",
+                map.meta.key_type
+            );
+        }
+    }
+
+    Ok(())
+}
+
 fn run_convert(args: ConvertArgs) -> Result<()> {
     let entries = load_entries_from_json(&args.input)?;
     let writer = File::create(&args.output)?;

src/builder.rs

@@ -212,6 +212,8 @@ impl<H: MassMapHashLoader> MassMapBuilder<H> {
             bucket_count: self.bucket_count,
             occupied_bucket_count,
             occupied_bucket_range,
+            key_type: std::any::type_name::<K>().to_string(),
+            value_type: std::any::type_name::<V>().to_string(),
         };
 
         let meta_offset = offset.load(Ordering::Relaxed) + buf_writer.buffer().len() as u64;
@@ -292,39 +294,39 @@ impl MassMapMerger {
             ));
         }
 
-        maps.sort_by_key(|m| m.meta.occupied_bucket_range.start);
+        maps.sort_by_key(|m| m.meta().occupied_bucket_range.start);
 
         let mut entry_count = 0;
         let mut bucket_metas =
-            vec![MassMapBucketMeta::default(); maps[0].meta.bucket_count as usize];
-        let hash_config = maps[0].meta.hash_config.clone();
+            vec![MassMapBucketMeta::default(); maps[0].meta().bucket_count as usize];
+        let hash_config = maps[0].meta().hash_config.clone();
         let mut occupied_bucket_count = 0;
         let mut occupied_bucket_range = 0..0;
         let mut global_offset = 0u64;
 
         for map in &maps {
-            if map.meta.hash_config != hash_config {
+            if map.meta().hash_config != hash_config {
                 return Err(Error::new(
                     ErrorKind::InvalidData,
                     "Incompatible hash configurations between massmaps",
                 ));
             }
-            if map.meta.bucket_count != bucket_metas.len() as u64 {
+            if map.meta().bucket_count != bucket_metas.len() as u64 {
                 return Err(Error::new(
                     ErrorKind::InvalidData,
                     "Incompatible bucket counts between massmaps",
                 ));
             }
 
-            if map.meta.entry_count == 0 {
+            if map.meta().entry_count == 0 {
                 continue;
             }
 
-            occupied_bucket_count += map.meta.occupied_bucket_count;
+            occupied_bucket_count += map.meta().occupied_bucket_count;
             if occupied_bucket_range.is_empty() {
-                occupied_bucket_range = map.meta.occupied_bucket_range.clone();
-            } else if occupied_bucket_range.end <= map.meta.occupied_bucket_range.start {
-                occupied_bucket_range.end = map.meta.occupied_bucket_range.end;
+                occupied_bucket_range = map.meta().occupied_bucket_range.clone();
+            } else if occupied_bucket_range.end <= map.meta().occupied_bucket_range.start {
+                occupied_bucket_range.end = map.meta().occupied_bucket_range.end;
             } else {
                 return Err(Error::new(
                     ErrorKind::InvalidData,
@@ -333,24 +335,24 @@ impl MassMapMerger {
             }
 
             // update bucket metas.
-            for idx in map.meta.occupied_bucket_range.clone() {
+            for idx in map.meta().occupied_bucket_range.clone() {
                 let bucket_meta = &mut bucket_metas[idx as usize];
-                *bucket_meta = map.bucket_metas[idx as usize];
+                *bucket_meta = map.bucket_metas()[idx as usize];
                 if bucket_meta.count > 0 {
                     bucket_meta.offset += global_offset;
                 }
             }
-            entry_count += map.meta.entry_count;
+            entry_count += map.meta().entry_count;
 
             // copy buckets from reader to writer directly.
             let mut current_offset = MassMapHeader::SIZE as u64;
-            let finished_offset = map.header.meta_offset;
+            let finished_offset = map.header().meta_offset;
             while current_offset < finished_offset {
                 let chunk = std::cmp::min(
                     finished_offset - current_offset,
                     self.writer_buffer_size as u64,
                 );
-                map.reader.read_exact_at(current_offset, chunk, |data| {
+                map.reader().read_exact_at(current_offset, chunk, |data| {
                     writer.write_all_at(data, global_offset + MassMapHeader::SIZE as u64)?;
                     Ok(())
                 })?;
@@ -365,6 +367,8 @@ impl MassMapMerger {
             bucket_count: bucket_metas.len() as u64,
             occupied_bucket_count,
             occupied_bucket_range,
+            key_type: std::any::type_name::<K>().to_string(),
+            value_type: std::any::type_name::<V>().to_string(),
         };
 
         let meta_offset = global_offset + MassMapHeader::SIZE as u64;
@@ -585,9 +589,9 @@ mod tests {
             threads.push(std::thread::spawn(move || {
                 let entries = (0..N).filter(|v| (v % M) / (M / P) == i).map(|v| (v, v));
                 let map = create_simple_map(entries, M, M);
-                assert_eq!(map.meta.occupied_bucket_count, M / P);
-                assert_eq!(map.meta.entry_count, N / P);
-                assert_eq!(map.meta.occupied_bucket_range.start, (M / P) * i);
+                assert_eq!(map.meta().occupied_bucket_count, M / P);
+                assert_eq!(map.meta().entry_count, N / P);
+                assert_eq!(map.meta().occupied_bucket_range.start, (M / P) * i);
 
                 for item in map.iter() {
                     let (k, v) = item.unwrap();

src/lib.rs

@@ -42,7 +42,7 @@ mod writer;
 pub use writer::MassMapWriter;
 
 mod massmap;
-pub use massmap::{MassMap, MassMapIter};
+pub use massmap::{MassMap, MassMapInner, MassMapIter};
 
 mod builder;
 pub use builder::{MassMapBuilder, MassMapMerger};