Repository: dathere/qsv (language: rust)

Issue:
`extdedup` memory-mapped on-disk hash table
Hi! For inspiration and learning purposes I checked out `extdedup` and its description states:

> Remove duplicate rows from an arbitrarily large CSV/text file using a memory-mapped, [on-disk hash table](https://crates.io/crates/odht). Unlike the dedup command, this command does not load the entire file into memory nor does it sort the deduped file.

But looking at [the implementation](https://github.com/dathere/qsv/blob/4a8eb2bf0a273ac55728d560671f8652e432518c/src/odhtcache.rs#L92-L95) I can't seem to find where it performs `mmap` for the backing memory. It uses a `HashTableOwned` which according to the docs:

> A [HashTableOwned](https://docs.rs/odht/latest/odht/struct.HashTableOwned.html) keeps the underlying data **on the heap** and can resize itself on demand.

I also looked through the source for `odht` and it too doesn't seem perform any mmap unless I've missed something. If I understand correctly, it is supposed to use `mmap`-backed memory like this with the `HashTable`-type rather than `HashTableOwned`-type:

```rust
use memmap2::MmapOptions;
use odht::{bytes_needed, Config, FxHashFn, HashTable};

struct Conf;

impl Config for Conf {
    type Key = u64;
    type Value = u32;
    type EncodedKey = [u8; 8];
    type EncodedValue = [u8; 4];
    type H = FxHashFn;

    #[inline]
    fn encode_key(k: &Self::Key) -> Self::EncodedKey {
        k.to_le_bytes()
    }
    #[inline]
    fn encode_value(v: &Self::Value) -> Self::EncodedValue {
        v.to_le_bytes()
    }
    #[inline]
    fn decode_key(k: &Self::EncodedKey) -> Self::Key {
        u64::from_le_bytes(*k)
    }
    #[inline]
    fn decode_value(v: &Self::EncodedValue) -> Self::Value {
        u32::from_le_bytes(*v)
    }
}

fn main() {
    let max_cap = 1024;
    let load_factor = 95;
    let len = bytes_needed::<Conf>(max_cap, load_factor);

    let mut mmap = MmapOptions::new()
        .len(len)
        .map_anon()
        .expect("Failed to mmap");

    let mut table = HashTable::<Conf, &mut [u8]>::init_in_place(&mut mmap, max_cap, load_factor)
        .expect("Failed to create table");
    table.insert(&1, &2);
    table.insert(&3, &4);

    assert_eq!(table.len(), 2);
    assert_eq!(table.get(&1), Some(2u32));
    assert_eq!(table.get(&3), Some(4u32));
}
```


Discussion / hints:
Hi @Svenskunganka ,
Thanks for the detailed investigation.
And you're right, I thought the odht crate implemented it, and since it's used by the compiler team, I just went with it giving the name of the crate.

It seems the odht crate is designed to work with any mem-mapping implementation by setting up `odht::Config` and passing the mapped memory to `HashTable::from_raw_bytes()`.

It seems odht focused solely on a platform-independent hash table implementation and binary format, leaving the choice of how to actually persist the odht to the crate consumer.

If you're up to it, can you actually implement the memory mapping?

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index c6dedff74..e3607c656 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5215,6 +5215,7 @@ dependencies = [
  "jsonschema",
  "localzone",
  "log",
+ "memmap2",
  "mimalloc",
  "minijinja",
  "minijinja-contrib",
diff --git a/Cargo.toml b/Cargo.toml
index 4aa7bd52d..90227c8d6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -143,6 +143,7 @@ jsonschema = { version = "0.28", features = [
 ], default-features = false }
 localzone = { version = "0.3", features = ["auto_validation"] }
 log = "0.4"
+memmap2 = "0.9"
 mimalloc = { version = "0.1", default-features = false, optional = true }
 minijinja = { version = "2", features = [
     "json",
@@ -178,6 +179,7 @@ polars = { version = "0.45", features = [
     # "cloud",
     "coalesce",
     "cross_join",
+    "cse",
     "csv",
     "decompress",
     "diagonal_concat",
diff --git a/src/cmd/extdedup.rs b/src/cmd/extdedup.rs
index 3305ac39e..708dba9e4 100644
--- a/src/cmd/extdedup.rs
+++ b/src/cmd/extdedup.rs
@@ -39,6 +39,8 @@ extdedup options:
                                If more than 50, this is the memory in MB to allocate, capped
                                at 90 percent of total memory.
                                [default: 10]
+    --temp-dir <arg>           Directory to store temporary hash table file.
+                               If not specified, defaults to operating system temp directory.
 
 Common options:
                                CSV MODE ONLY:
@@ -56,6 +58,7 @@ Common options:
 use std::{
     fs,
     io::{self, stdin, stdout, BufRead, Write},
+    path::PathBuf,
 };
 
 use indicatif::HumanCount;
@@ -81,6 +84,7 @@ struct Args {
     flag_dupes_output:   Option<String>,
     flag_human_readable: bool,
     flag_memory_limit:   Option<u64>,
+    flag_temp_dir:       Option<String>,
     flag_quiet:          bool,
 }
 
@@ -138,7 +142,8 @@ fn dedup_csv(args: Args, mem_limited_buffer: u64) -> Result<u64, crate::clitypes
         dupewtr.write_byte_record(&dupe_headers)?;
     }
 
-    let mut dedup_cache = odhtcache::ExtDedupCache::new(mem_limited_buffer);
+    let temp_dir = args.flag_temp_dir.map(PathBuf::from);
+    let mut dedup_cache = odhtcache::ExtDedupCache::new(mem_limited_buffer, temp_dir);
     let mut dupes_count = 0_u64;
     let sel = rconfig.selection(&headers)?;
 
@@ -236,7 +241,8 @@ fn dedup_lines(args: Args, mem_limited_buffer: u64) -> Result<u64, crate::clityp
             fs::File::create("nul")?,
         )
     };
-    let mut dedup_cache = odhtcache::ExtDedupCache::new(mem_limited_buffer);
+    let temp_dir = args.flag_temp_dir.map(PathBuf::from);
+    let mut dedup_cache = odhtcache::ExtDedupCache::new(mem_limited_buffer, temp_dir);
     let mut dupes_count = 0_u64;
     let mut line_work = String::with_capacity(1024);
     for (row_idx, line) in input_reader.lines().enumerate() {
diff --git a/src/odhtcache.rs b/src/odhtcache.rs
index cc6c57c5f..ff3b375e3 100644
--- a/src/odhtcache.rs
+++ b/src/odhtcache.rs
@@ -1,11 +1,14 @@
-// blatantly copied from https://github.com/race604/dedup/blob/master/src/cache.rs
-use std::collections::HashSet;
+// inspired by https://github.com/race604/dedup/blob/master/src/cache.rs
+use std::{collections::HashSet, path::PathBuf};
 
 use log::debug;
-use odht::{Config, FxHashFn, HashTableOwned};
+use memmap2::MmapMut;
+use odht::{bytes_needed, Config, FxHashFn, HashTableOwned};
+use tempfile::NamedTempFile;
 
 struct ExtDedupConfig;
 
+const ODHT_CAPACITY: usize = 10_000_000; // 10 million initial capacity
 const CHUNK_SIZE: usize = 127;
 
 impl Config for ExtDedupConfig {
@@ -41,10 +44,13 @@ pub struct ExtDedupCache {
     disk:       Option<HashTableOwned<ExtDedupConfig>>,
     memo_limit: u64,
     memo_size:  u64,
+    temp_file:  Option<NamedTempFile>,
+    mmap:       Option<MmapMut>,
+    temp_dir:   PathBuf,
 }
 
 impl ExtDedupCache {
-    pub fn new(memo_limit: u64) -> Self {
+    pub fn new(memo_limit: u64, temp_dir: Option<PathBuf>) -> Self {
         Self {
             memo:       HashSet::new(),
             disk:       None,
@@ -54,9 +60,47 @@ impl ExtDedupCache {
                 memo_limit
             },
             memo_size:  0,
+            temp_file:  None,
+            mmap:       None,
+            temp_dir:   temp_dir.unwrap_or_else(std::env::temp_dir),
         }
     }
 
+    fn create_mmap(&mut self) -> std::io::Result<()> {
+        let temp_file = tempfile::Builder::new()
+            .prefix("qsv-extdedup-")
+            .suffix(".tmp")
+            .tempfile_in(&self.temp_dir)?;
+
+        // Calculate required space for the hash table
+        let load_factor = 95;
+        let required_bytes = bytes_needed::<ExtDedupConfig>(ODHT_CAPACITY, load_factor);
+
+        // Ensure file is large enough
+        temp_file.as_file().set_len(required_bytes as u64)?;
+
+        let mut mmap = unsafe { MmapMut::map_mut(temp_file.as_file())? };
+
+        // Create a properly initialized table
+        let table = HashTableOwned::<ExtDedupConfig>::with_capacity(ODHT_CAPACITY, load_factor);
+
+        // Copy the initialized bytes to mmap
+        let raw_bytes = table.raw_bytes();
+        if mmap.len() >= raw_bytes.len() {
+            mmap[..raw_bytes.len()].copy_from_slice(raw_bytes);
+            mmap.flush()?;
+        } else {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "Mmap size too small for ODHT table",
+            ));
+        }
+
+        self.mmap = Some(mmap);
+        self.temp_file = Some(temp_file);
+        Ok(())
+    }
+
     #[inline]
     pub fn insert(&mut self, item: &str) -> bool {
         if self.memo_size >= self.memo_limit {
@@ -89,13 +133,32 @@ impl ExtDedupCache {
     }
 
     fn insert_on_disk(&mut self, item: &str) -> bool {
-        let disk = self.disk.get_or_insert_with(|| {
+        if self.disk.is_none() {
             debug!("Create new disk cache");
-            HashTableOwned::<ExtDedupConfig>::with_capacity(1_000_000, 95)
-        });
+            match self.create_mmap() {
+                Ok(()) => {
+                    if let Some(mmap) = &mut self.mmap {
+                        // Create the table from the properly initialized mmap
+                        self.disk = Some(unsafe {
+                            HashTableOwned::<ExtDedupConfig>::from_raw_bytes_unchecked(mmap)
+                        });
+                    }
+                },
+                Err(e) => {
+                    debug!("Failed to create memory map: {}", e);
+                    // Fallback to regular HashTableOwned if mmap fails
+                    self.disk = Some(HashTableOwned::<ExtDedupConfig>::with_capacity(
+                        1_000_000, 95,
+                    ));
+                },
+            }
+        }
+
         let mut res = false;
-        for key in ExtDedupCache::item_to_keys(item) {
-            res = disk.insert(&key, &true).is_none() || res;
+        if let Some(disk) = &mut self.disk {
+            for key in ExtDedupCache::item_to_keys(item) {
+                res = disk.insert(&key, &true).is_none() || res;
+            }
         }
         res
     }
@@ -115,7 +178,7 @@ impl ExtDedupCache {
     }
 
     fn dump_to_disk(&mut self) {
-        // debug!("Memory cache is full, dump to disk");
+        debug!("Memory cache is full, dump to disk");
         let keys = self.memo.drain().collect::<Vec<_>>();
         for key in keys {
             self.insert_on_disk(&key);
@@ -124,6 +187,14 @@ impl ExtDedupCache {
     }
 }
 
+impl Drop for ExtDedupCache {
+    fn drop(&mut self) {
+        // Explicitly drop mmap first
+        self.mmap.take();
+        // temp_file will be automatically deleted when dropped
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use rand::{distributions::Alphanumeric, thread_rng, Rng};
@@ -132,7 +203,7 @@ mod tests {
 
     #[test]
     fn test_basic_cache() {
-        let mut cache = ExtDedupCache::new(0);
+        let mut cache = ExtDedupCache::new(0, None);
         assert!(cache.insert("hello"));
         assert!(cache.insert("world"));
 
@@ -143,13 +214,13 @@ mod tests {
 
     #[test]
     fn test_limit_memory() {
-        let mut cache = ExtDedupCache::new(1024);
+        let mut cache = ExtDedupCache::new(1024, None);
         for _ in 0..100 {
             cache.insert(&rand_string(32));
         }
         assert!(cache.memo.len() < 100);
         assert!(cache.disk.is_some());
-        assert!(cache.disk.unwrap().len() > 0);
+        assert!(cache.disk.as_ref().unwrap().len() > 0);
     }
 
     fn rand_string(len: usize) -> String {


Test patch: