feat: range aware file write

wasm/hf_xet_wasm/src/wasm_file_cleaner.rs

@@ -176,7 +176,8 @@ impl SingleFileCleaner {
         let metadata_ext = FileMetadataExt::new(sha256);
 
         // Now finish the deduplication process.
-        let (file_hash, remaining_file_data, deduplication_metrics) = self.dedup_manager.finalize(Some(metadata_ext));
+        let (file_hash, _chunk_hashes, remaining_file_data, deduplication_metrics) =
+            self.dedup_manager.finalize(Some(metadata_ext));
 
         // Let's check some things that should be invariants
         {

xet_client/src/cas_client/interface.rs

@@ -1,5 +1,5 @@
 use bytes::Bytes;
-use xet_core_structures::merklehash::MerkleHash;
+use xet_core_structures::merklehash::{ChunkHashList, MerkleHash};
 use xet_core_structures::metadata_shard::file_structs::MDBFileInfo;
 use xet_core_structures::xorb_object::SerializedXorbObject;
 
@@ -70,4 +70,8 @@ pub trait Client: Send + Sync {
         progress_callback: Option<ProgressCallback>,
         upload_permit: ConnectionPermit,
     ) -> Result<u64>;
+
+    /// Retrieve the chunk hashes and sizes for a file stored in CAS.
+    /// Returns a list of (chunk_hash, chunk_uncompressed_size) pairs.
+    async fn get_file_chunk_hashes(&self, file_id: &MerkleHash) -> Result<ChunkHashList>;
 }

xet_client/src/cas_client/remote_client.rs

@@ -8,8 +8,9 @@ use http::HeaderValue;
 use http::header::{CONTENT_LENGTH, HeaderMap, RANGE};
 use reqwest::{Body, Response, StatusCode, Url};
 use reqwest_middleware::ClientWithMiddleware;
+use serde::Deserialize;
 use tracing::{event, info, instrument};
-use xet_core_structures::merklehash::MerkleHash;
+use xet_core_structures::merklehash::{ChunkHashList, MerkleHash};
 use xet_core_structures::metadata_shard::file_structs::{FileDataSequenceEntry, FileDataSequenceHeader, MDBFileInfo};
 use xet_core_structures::xorb_object::SerializedXorbObject;
 use xet_runtime::core::xet_config;
@@ -735,6 +736,45 @@ impl Client for RemoteClient {
 
         Ok(n_upload_bytes)
     }
+
+    #[instrument(skip_all, name = "RemoteClient::get_file_chunk_hashes", fields(file.hash = file_id.hex()))]
+    async fn get_file_chunk_hashes(&self, file_id: &MerkleHash) -> Result<ChunkHashList> {
+        let url = Url::parse(&format!("{}/v2/file-chunk-hashes/{}", self.endpoint, file_id.hex()))?;
+
+        let api_tag = "cas::get_file_chunk_hashes";
+        let client = self.authenticated_http_client.clone();
+
+        let response: FileChunkHashesResponse = RetryWrapper::new(api_tag)
+            .run_and_extract_json(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
+            .await?;
+
+        let chunks = response.chunks.into_iter().map(|entry| (entry.hash, entry.size)).collect();
+
+        Ok(chunks)
+    }
+}
+
+/// Response from `GET /v2/file-chunk-hashes/{file_id}`.
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct FileChunkHashesResponse {
+    chunks: Vec<ChunkHashEntry>,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct ChunkHashEntry {
+    #[serde(deserialize_with = "deserialize_merkle_hash")]
+    hash: MerkleHash,
+    size: u64,
+}
+
+fn deserialize_merkle_hash<'de, D>(deserializer: D) -> std::result::Result<MerkleHash, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    let s = String::deserialize(deserializer)?;
+    MerkleHash::from_hex(&s).map_err(serde::de::Error::custom)
 }
 
 #[cfg(test)]

xet_client/src/cas_client/simulation/local_client.rs

@@ -15,7 +15,7 @@ use rand::Rng;
 use tempfile::TempDir;
 use tokio::time::{Duration, Instant};
 use tracing::{error, info, warn};
-use xet_core_structures::merklehash::MerkleHash;
+use xet_core_structures::merklehash::{ChunkHashList, MerkleHash};
 use xet_core_structures::metadata_shard::file_structs::MDBFileInfo;
 use xet_core_structures::metadata_shard::shard_file_reconstructor::FileReconstructor;
 use xet_core_structures::metadata_shard::shard_in_memory::MDBInMemoryShard;
@@ -1033,6 +1033,25 @@ impl Client for LocalClient {
         // Should not reach here, but return error if we do.
         Err(ClientError::PresignedUrlExpirationError)
     }
+
+    async fn get_file_chunk_hashes(&self, file_id: &MerkleHash) -> Result<ChunkHashList> {
+        self.apply_api_delay().await;
+
+        let Some((file_info, _)) = self.shard_manager.get_file_reconstruction_info(file_id).await? else {
+            return Err(ClientError::FileNotFound(*file_id));
+        };
+
+        let mut result = Vec::new();
+        for segment in &file_info.segments {
+            let xorb_obj = self.xorb_footer(&segment.xorb_hash).await?;
+            let pairs = xorb_obj
+                .chunk_hash_sizes(segment.chunk_index_start, segment.chunk_index_end)
+                .map_err(|err| ClientError::Other(format!("chunk_hash_sizes error: {err}")))?;
+            result.extend(pairs);
+        }
+
+        Ok(result)
+    }
 }
 
 fn map_heed_db_error(e: heed::Error) -> ClientError {

xet_client/src/cas_client/simulation/local_server/server.rs

@@ -485,6 +485,13 @@ impl Client for LocalTestServer {
             .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit)
             .await
     }
+
+    async fn get_file_chunk_hashes(
+        &self,
+        file_id: &xet_core_structures::merklehash::MerkleHash,
+    ) -> Result<xet_core_structures::merklehash::ChunkHashList> {
+        self.client.get_file_chunk_hashes(file_id).await
+    }
 }
 
 #[cfg(test)]

xet_client/src/cas_client/simulation/local_server/simulation_control_client.rs

@@ -5,7 +5,7 @@ use std::time::Duration;
 use async_trait::async_trait;
 use bytes::Bytes;
 use http::header::HeaderMap;
-use xet_core_structures::merklehash::MerkleHash;
+use xet_core_structures::merklehash::{ChunkHashList, MerkleHash};
 use xet_core_structures::xorb_object::XorbObject;
 
 use super::simulation_types::{
@@ -152,6 +152,10 @@ impl Client for SimulationControlClient {
             .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit)
             .await
     }
+
+    async fn get_file_chunk_hashes(&self, file_id: &MerkleHash) -> Result<ChunkHashList> {
+        self.remote_client.get_file_chunk_hashes(file_id).await
+    }
 }
 
 #[async_trait]

xet_client/src/cas_client/simulation/memory_client.rs

@@ -11,7 +11,7 @@ use tokio::sync::RwLock;
 use tokio::time::{Duration, Instant};
 use tracing::{error, info};
 use xet_core_structures::MerkleHashMap;
-use xet_core_structures::merklehash::MerkleHash;
+use xet_core_structures::merklehash::{ChunkHashList, MerkleHash};
 use xet_core_structures::metadata_shard::file_structs::MDBFileInfo;
 use xet_core_structures::metadata_shard::shard_in_memory::MDBInMemoryShard;
 use xet_core_structures::metadata_shard::streaming_shard::MDBMinimalShard;
@@ -897,6 +897,38 @@ impl Client for MemoryClient {
         }
         Ok((Bytes::from(all_decompressed), all_chunk_indices))
     }
+
+    async fn get_file_chunk_hashes(&self, file_id: &MerkleHash) -> Result<ChunkHashList> {
+        self.apply_api_delay().await;
+
+        let file_info = {
+            let shard = self.shard.read().await;
+            shard
+                .get_file_reconstruction_info(file_id)
+                .ok_or(ClientError::FileNotFound(*file_id))?
+        };
+
+        let xorbs = self.xorbs.read().await;
+        let mut result = Vec::new();
+
+        for segment in &file_info.segments {
+            let storage = xorbs
+                .get(&segment.xorb_hash)
+                .ok_or(ClientError::XORBNotFound(segment.xorb_hash))?;
+
+            let xorb_obj = match storage {
+                XorbStorage::Materialized(entry) => std::borrow::Cow::Borrowed(&entry.xorb_object),
+                XorbStorage::Random(xorb) => std::borrow::Cow::Owned(xorb.get_xorb_object()),
+            };
+
+            let pairs = xorb_obj
+                .chunk_hash_sizes(segment.chunk_index_start, segment.chunk_index_end)
+                .map_err(|err| ClientError::Other(format!("chunk_hash_sizes error: {err}")))?;
+            result.extend(pairs);
+        }
+
+        Ok(result)
+    }
 }
 
 fn generate_fetch_url(hash: &MerkleHash, byte_range: &FileRange, timestamp: Instant) -> String {

xet_client/src/cas_client/simulation/simulation_client.rs

@@ -186,4 +186,11 @@ impl Client for RemoteSimulationClient {
             .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit)
             .await
     }
+
+    async fn get_file_chunk_hashes(
+        &self,
+        file_id: &xet_core_structures::merklehash::MerkleHash,
+    ) -> Result<xet_core_structures::merklehash::ChunkHashList> {
+        self.inner.get_file_chunk_hashes(file_id).await
+    }
 }

xet_client/src/cas_client/simulation/simulation_server.rs

@@ -500,6 +500,13 @@ impl Client for LocalTestServer {
             .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit)
             .await
     }
+
+    async fn get_file_chunk_hashes(
+        &self,
+        file_id: &xet_core_structures::merklehash::MerkleHash,
+    ) -> Result<xet_core_structures::merklehash::ChunkHashList> {
+        self.client.get_file_chunk_hashes(file_id).await
+    }
 }
 
 #[async_trait]

xet_core_structures/src/merklehash/mod.rs

@@ -48,6 +48,9 @@ pub mod data_hash;
 pub use data_hash::*;
 pub type MerkleHash = DataHash;
 
+/// List of (chunk_hash, chunk_uncompressed_size) pairs for a file or xorb range.
+pub type ChunkHashList = Vec<(MerkleHash, u64)>;
+
 mod aggregated_hashes;
 pub mod passthrough_hasher;
 pub mod passthrough_hashmap;

xet_core_structures/src/xorb_object/xorb_object_format.rs

@@ -14,7 +14,7 @@ use super::constants::{TARGET_CHUNK_SIZE, XORB_BLOCK_SIZE};
 use super::xorb_chunk_format::{deserialize_chunk, deserialize_chunk_header, serialize_chunk, write_chunk_header};
 use super::{CompressionScheme, RawXorbData, XorbChunkHeader};
 use crate::error::{CoreError, Validate};
-use crate::merklehash::{DataHash, MerkleHash};
+use crate::merklehash::{ChunkHashList, DataHash, MerkleHash};
 use crate::metadata_shard::chunk_verification::range_hash_from_chunks;
 use crate::serialization_utils::*;
 
@@ -1241,6 +1241,21 @@ impl XorbObject {
         Ok(incl_end - before_start)
     }
 
+    /// Returns (chunk_hash, uncompressed_size) pairs for chunks in [start, end).
+    pub fn chunk_hash_sizes(&self, start: u32, end: u32) -> Result<ChunkHashList, CoreError> {
+        self.validate_xorb_object_info()?;
+        if end > self.info.num_chunks || start > end {
+            return Err(CoreError::InvalidArguments);
+        }
+        (start..end)
+            .map(|i| {
+                let hash = self.info.chunk_hashes[i as usize];
+                let size = self.uncompressed_chunk_length(i)? as u64;
+                Ok((hash, size))
+            })
+            .collect()
+    }
+
     /// Helper method to verify that info object is complete
     fn validate_xorb_object_info(&self) -> Result<(), CoreError> {
         if self.info.num_chunks == 0 {

xet_data/src/deduplication/file_deduplication.rs

@@ -2,7 +2,7 @@ use std::result::Result;
 
 use more_asserts::{debug_assert_le, debug_assert_lt};
 use xet_core_structures::MerkleHashMap;
-use xet_core_structures::merklehash::{MerkleHash, file_hash};
+use xet_core_structures::merklehash::{ChunkHashList, MerkleHash, file_hash};
 use xet_core_structures::metadata_shard::file_structs::{
     FileDataSequenceEntry, FileDataSequenceHeader, FileMetadataExt, FileVerificationEntry, MDBFileInfo,
 };
@@ -32,7 +32,7 @@ pub struct FileDeduper<DataInterfaceType: DeduplicationDataInterface> {
     new_data_hash_lookup: MerkleHashMap<usize>,
 
     /// The current chunk hashes for this file.
-    chunk_hashes: Vec<(MerkleHash, u64)>,
+    chunk_hashes: ChunkHashList,
 
     /// The current file data entries.
     file_info: Vec<FileDataSequenceEntry>,
@@ -373,8 +373,11 @@ impl<DataInterfaceType: DeduplicationDataInterface> FileDeduper<DataInterfaceTyp
     /// and remaining data.  Also returns the aggregated deduplication metrics and the list of xorb hashes that were
     /// registered as part of this run.
     ///
-    /// Returns (file hash, data aggregation, deduplication metrics)
-    pub fn finalize(self, metadata_ext: Option<FileMetadataExt>) -> (MerkleHash, DataAggregator, DeduplicationMetrics) {
+    /// Returns (file hash, chunk_hashes, data aggregation, deduplication metrics)
+    pub fn finalize(
+        self,
+        metadata_ext: Option<FileMetadataExt>,
+    ) -> (MerkleHash, ChunkHashList, DataAggregator, DeduplicationMetrics) {
         let file_hash = file_hash(&self.chunk_hashes);
 
         let metadata = FileDataSequenceHeader::new(file_hash, self.file_info.len(), true, metadata_ext.is_some());
@@ -408,6 +411,6 @@ impl<DataInterfaceType: DeduplicationDataInterface> FileDeduper<DataInterfaceTyp
 
         let remaining_data = DataAggregator::new(self.new_data, fi, self.internally_referencing_entries, self.file_id);
 
-        (file_hash, remaining_data, self.deduplication_metrics)
+        (file_hash, self.chunk_hashes, remaining_data, self.deduplication_metrics)
     }
 }

xet_data/src/processing/bin/example.rs

@@ -104,7 +104,7 @@ async fn clean(mut reader: impl Read, mut writer: impl Write, size: u64) -> Resu
 
     debug_assert_eq!(size_read, size);
 
-    let (file_info, _) = handle.finish().await?;
+    let (file_info, _chunk_hashes, _metrics) = handle.finish().await?;
 
     translator.finalize().await?;
 

xet_data/src/processing/data_client.rs

@@ -46,7 +46,8 @@ pub async fn clean_bytes(
 ) -> Result<(XetFileInfo, DeduplicationMetrics)> {
     let (_id, mut handle) = processor.start_clean(None, bytes.len() as u64, sha256_policy)?;
     handle.add_data(&bytes).await?;
-    handle.finish().await
+    let (info, _chunk_hashes, metrics) = handle.finish().await?;
+    Ok((info, metrics))
 }
 
 #[instrument(skip_all, name = "clean_file", fields(file.name = tracing::field::Empty, file.len = tracing::field::Empty))]
@@ -75,7 +76,8 @@ pub async fn clean_file(
         handle.add_data(&buffer[0..bytes]).await?;
     }
 
-    handle.finish().await
+    let (info, _chunk_hashes, metrics) = handle.finish().await?;
+    Ok((info, metrics))
 }
 
 /// Computes the xet hash for a single file without uploading.

xet_data/src/processing/file_cleaner.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use tracing::{Instrument, debug_span, info, instrument};
+use xet_core_structures::merklehash::ChunkHashList;
 use xet_core_structures::metadata_shard::Sha256;
 use xet_core_structures::metadata_shard::file_structs::FileMetadataExt;
 use xet_runtime::core::{XetRuntime, xet_config};
@@ -194,7 +195,7 @@ impl SingleFileCleaner {
 
     /// Return the representation of the file after clean as a pointer file instance.
     #[instrument(skip_all, name = "FileCleaner::finish", fields(file_name=self.file_name.as_ref().map(|s|s.to_string())))]
-    pub async fn finish(mut self) -> Result<(XetFileInfo, DeduplicationMetrics)> {
+    pub async fn finish(mut self) -> Result<(XetFileInfo, ChunkHashList, DeduplicationMetrics)> {
         // Chunk the rest of the data.
         if let Some(chunk) = self.chunker.finish() {
             let data = Arc::new([chunk]);
@@ -209,7 +210,7 @@ impl SingleFileCleaner {
         };
         let metadata_ext = sha256.map(FileMetadataExt::new);
 
-        let (file_hash, remaining_file_data, deduplication_metrics) =
+        let (file_hash, chunk_hashes, remaining_file_data, deduplication_metrics) =
             self.dedup_manager_fut.await?.finalize(metadata_ext);
 
         let file_info = XetFileInfo {
@@ -246,6 +247,6 @@ impl SingleFileCleaner {
             end_processing_ts = Utc::now().to_rfc3339(),
         );
 
-        Ok((file_info, deduplication_metrics))
+        Ok((file_info, chunk_hashes, deduplication_metrics))
     }
 }

xet_data/src/processing/file_download_session.rs

@@ -323,7 +323,7 @@ mod tests {
             .start_clean(Some("test".into()), data.len() as u64, Sha256Policy::Compute)
             .unwrap();
         cleaner.add_data(data).await.unwrap();
-        let (xfi, _metrics) = cleaner.finish().await.unwrap();
+        let (xfi, _chunk_hashes, _metrics) = cleaner.finish().await.unwrap();
         upload_session.finalize().await.unwrap();
         xfi
     }