xet-data 1.5.2

Data processing pipeline for chunking, deduplication, and file reconstruction; used in the Hugging Face Xet client tools. Intended to be used through the API in the hf-xet package.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use std::fs::File;
use std::io::{BufRead, BufReader, Read, Write};
use std::path::PathBuf;

use clap::Parser;
use regex::Regex;
use xet_core_structures::merklehash::{MerkleHash, compute_data_hash, file_hash, xorb_hash};
use xet_core_structures::metadata_shard::chunk_verification::range_hash_from_chunks;

#[derive(Debug, Copy, Clone)]
enum HashType {
    Chunk,
    Xorb,
    File,
    Range,
}

impl std::str::FromStr for HashType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "chunk" => Ok(HashType::Chunk),
            "xorb" => Ok(HashType::Xorb),
            "file" => Ok(HashType::File),
            "range" => Ok(HashType::Range),
            _ => Err(format!("Invalid hash type: {s}")),
        }
    }
}

#[derive(Debug, Parser)]
#[command(
    version,
    about,
    long_about = "Example of computing different hash functions for different inputs"
)]
struct HashArgs {
    /// Hash type to compute, choices are: chunk, xorb, file, range
    #[arg(short, long)]
    hash_type: HashType,
    /// Output file or uses stdout if not specified, where to write the hash
    #[arg(short, long)]
    output: Option<PathBuf>,
    /// Input file or uses stdin if not specified. For File, Xorb and Range hash types, expects input format to be a
    /// newline-separated list of chunk hashes and lengths in the format of `<hash> <length>`
    #[arg(short, long)]
    input: Option<PathBuf>,
}

fn main() {
    let args = HashArgs::parse();

    let mut input: Box<dyn BufRead> = if let Some(path) = args.input {
        Box::new(BufReader::new(File::open(path).unwrap()))
    } else {
        Box::new(std::io::stdin().lock())
    };

    let mut output: Box<dyn Write> = if let Some(path) = args.output {
        Box::new(File::create(path).unwrap())
    } else {
        Box::new(std::io::stdout())
    };

    if matches!(args.hash_type, HashType::Chunk) {
        let mut buf = vec![];
        input.read_to_end(&mut buf).unwrap();
        let hash = compute_data_hash(&buf);
        output.write_all(format!("{hash}").as_bytes()).unwrap();
        return;
    }

    let chunks_list = read_input_as_chunks_list(&mut input);
    let hash = match args.hash_type {
        HashType::Chunk => unreachable!("already handled"),
        HashType::Xorb => xorb_hash(&chunks_list),
        HashType::File => file_hash(&chunks_list),
        HashType::Range => {
            let hashes_only = chunks_list.into_iter().map(|(hash, _len)| hash).collect::<Vec<_>>();
            range_hash_from_chunks(&hashes_only)
        },
    };

    output.write_all(format!("{hash}").as_bytes()).unwrap();
}

fn read_input_as_chunks_list(input: &mut impl BufRead) -> Vec<(MerkleHash, u64)> {
    let line_regex = Regex::new(r"^([0-9a-fA-F]+)\s+(\d+)$").unwrap();

    let mut res = vec![];
    for line_result in input.lines() {
        let line: String = line_result.expect("Failed to read line");
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let caps = line_regex.captures(line).expect("Failed to parse line");
        let hash = MerkleHash::from_hex(caps.get(1).expect("Failed to find hash in line ").as_str())
            .expect("Failed to parse hash");
        let length: u64 = caps
            .get(2)
            .expect("Failed to find length in line")
            .as_str()
            .parse()
            .expect("Failed to parse length");
        res.push((hash, length));
    }

    res
}