hdf5_pure_rust/hl/mutable_file/
write_chunk.rs

1use std::borrow::Cow;
2use std::io::{Seek, SeekFrom, Write};
3
4use crate::error::{Error, Result};
5use crate::format::btree_v2::BTreeV2Header;
6use crate::format::messages::data_layout::{ChunkIndexType, LayoutClass};
7use crate::format::messages::filter_pipeline::{FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_SHUFFLE};
8
9use super::MutableFile;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12enum WritableChunkIndexKind {
13    BTreeV1,
14    FixedArray,
15    ExtensibleArray,
16    BTreeV2,
17}
18
19impl MutableFile {
20    /// Write a full uncompressed chunk and update the dataset's chunk index.
21    ///
22    /// This supports chunked datasets written by this crate where the chunk
23    /// index variant is one of the currently mutable writer-side subsets.
24    pub fn write_chunk(
25        &mut self,
26        dataset_path: &str,
27        chunk_coords: &[u64],
28        data: &[u8],
29    ) -> Result<()> {
30        let ds = self.dataset(dataset_path)?;
31        let info = ds.info()?;
32        let index_kind = Self::writable_chunk_index_kind(&info);
33        if self.superblock.base_addr != 0 && index_kind != WritableChunkIndexKind::BTreeV1 {
34            return Err(Error::Unsupported(
35                "write_chunk supports userblock files only for v1 B-tree chunk indexes".into(),
36            ));
37        }
38
39        if info.layout.layout_class != LayoutClass::Chunked {
40            return Err(Error::InvalidFormat(
41                "write_chunk only supports chunked datasets".into(),
42            ));
43        }
44        if info.layout.version > 3
45            && !matches!(
46                info.layout.chunk_index_type,
47                Some(ChunkIndexType::BTreeV1)
48                    | Some(ChunkIndexType::FixedArray)
49                    | Some(ChunkIndexType::ExtensibleArray)
50                    | Some(ChunkIndexType::BTreeV2)
51            )
52        {
53            return Err(Error::Unsupported(
54                "write_chunk currently supports only v1 B-tree, fixed-array, simple extensible-array, and simple v2 B-tree chunk indexes".into(),
55            ));
56        }
57
58        let chunk_data_dims = Self::chunk_data_dims(&info)?;
59        Self::validate_chunk_coords(chunk_coords, &info.dataspace.dims, chunk_data_dims)?;
60
61        let element_size = Self::u64_to_usize(u64::from(info.datatype.size), "datatype size")?;
62        let expected_len = Self::expected_chunk_len(chunk_data_dims, element_size)?;
63        Self::validate_chunk_write_len(data.len(), expected_len)?;
64        let filtered = Self::encode_chunk_write_data(&info, data, element_size)?;
65
66        let index_addr = info
67            .layout
68            .chunk_index_addr
69            .ok_or_else(|| Error::InvalidFormat("chunked dataset missing B-tree address".into()))?;
70
71        self.validate_chunk_index_rewrite(
72            index_kind,
73            index_addr,
74            &info,
75            chunk_coords,
76            chunk_data_dims,
77            expected_len,
78        )?;
79
80        let chunk_physical_addr = self.write_handle.seek(SeekFrom::End(0))?;
81        let chunk_addr =
82            self.logical_addr_from_physical(chunk_physical_addr, "appended chunk address")?;
83        self.write_handle.write_all(&filtered)?;
84        self.rewrite_chunk_index(
85            index_kind,
86            index_addr,
87            ds.addr(),
88            &info,
89            chunk_coords,
90            chunk_data_dims,
91            Self::usize_to_u64(filtered.len(), "filtered chunk size")?,
92            chunk_addr,
93            expected_len,
94            element_size,
95        )?;
96        let physical_eof = self.write_handle.seek(SeekFrom::End(0))?;
97        let logical_eof = physical_eof
98            .checked_sub(self.superblock.base_addr)
99            .ok_or_else(|| Error::InvalidFormat("file EOF is before HDF5 base address".into()))?;
100        self.rewrite_superblock_eof(logical_eof)?;
101        self.write_handle.flush()?;
102        self.reopen_reader()?;
103
104        Ok(())
105    }
106
107    pub(super) fn chunk_data_dims(info: &crate::hl::dataset::DatasetInfo) -> Result<&[u64]> {
108        let chunk_dims = info
109            .layout
110            .chunk_dims
111            .as_ref()
112            .ok_or_else(|| Error::InvalidFormat("chunked layout missing chunk dims".into()))?;
113        if chunk_dims.len() == info.dataspace.dims.len() + 1 {
114            Ok(&chunk_dims[..info.dataspace.dims.len()])
115        } else if chunk_dims.len() == info.dataspace.dims.len() {
116            Ok(chunk_dims)
117        } else {
118            Err(Error::InvalidFormat(format!(
119                "chunk dimension rank {} does not match dataset rank {}",
120                chunk_dims.len(),
121                info.dataspace.dims.len()
122            )))
123        }
124    }
125
126    fn validate_chunk_coords(
127        chunk_coords: &[u64],
128        data_dims: &[u64],
129        chunk_data_dims: &[u64],
130    ) -> Result<()> {
131        if chunk_coords.len() != chunk_data_dims.len() {
132            return Err(Error::InvalidFormat(format!(
133                "chunk coordinate rank {} does not match dataset rank {}",
134                chunk_coords.len(),
135                chunk_data_dims.len()
136            )));
137        }
138        if data_dims.len() != chunk_data_dims.len() {
139            return Err(Error::InvalidFormat(format!(
140                "chunk dimension rank {} does not match dataset rank {}",
141                chunk_data_dims.len(),
142                data_dims.len()
143            )));
144        }
145        for (((idx, &coord), &dim), &chunk) in chunk_coords
146            .iter()
147            .enumerate()
148            .zip(data_dims)
149            .zip(chunk_data_dims)
150        {
151            if chunk == 0 || coord % chunk != 0 {
152                return Err(Error::InvalidFormat(format!(
153                    "chunk coordinate {idx}={coord} is not aligned to chunk size {chunk}"
154                )));
155            }
156            if coord >= dim {
157                return Err(Error::Unsupported(format!(
158                    "write_chunk coordinate {idx}={coord} is outside dataset extent {dim}"
159                )));
160            }
161        }
162        Ok(())
163    }
164
165    fn expected_chunk_len(chunk_data_dims: &[u64], element_size: usize) -> Result<usize> {
166        let chunk_elements = chunk_data_dims.iter().try_fold(1usize, |acc, &dim| {
167            let dim = Self::u64_to_usize(dim, "chunk dimension")?;
168            acc.checked_mul(dim)
169                .ok_or_else(|| Error::InvalidFormat("chunk element count overflow".into()))
170        })?;
171        chunk_elements
172            .checked_mul(element_size)
173            .ok_or_else(|| Error::InvalidFormat("chunk byte size overflow".into()))
174    }
175
176    fn validate_chunk_write_len(actual_len: usize, expected_len: usize) -> Result<()> {
177        if actual_len != expected_len {
178            return Err(Error::InvalidFormat(format!(
179                "chunk data has {actual_len} bytes, expected {expected_len}",
180            )));
181        }
182        Ok(())
183    }
184
185    fn encode_chunk_write_data<'a>(
186        info: &crate::hl::dataset::DatasetInfo,
187        data: &'a [u8],
188        element_size: usize,
189    ) -> Result<Cow<'a, [u8]>> {
190        let mut filtered = Cow::Borrowed(data);
191        let mut filter_buf = Vec::new();
192        if let Some(ref pipeline) = info.filter_pipeline {
193            for filter in &pipeline.filters {
194                match filter.id {
195                    FILTER_SHUFFLE => {
196                        filter_buf.clear();
197                        filter_buf.resize(filtered.len(), 0);
198                        crate::filters::shuffle::shuffle_into(
199                            filtered.as_ref(),
200                            element_size,
201                            &mut filter_buf,
202                        )?;
203                        Self::replace_encoded_filter_buffer(&mut filtered, &mut filter_buf);
204                    }
205                    FILTER_DEFLATE => {
206                        let level = filter.client_data.first().copied().unwrap_or(6);
207                        filter_buf.clear();
208                        crate::filters::deflate::compress_into(
209                            filtered.as_ref(),
210                            level,
211                            &mut filter_buf,
212                        )?;
213                        Self::replace_encoded_filter_buffer(&mut filtered, &mut filter_buf);
214                    }
215                    FILTER_FLETCHER32 => {
216                        crate::filters::fletcher32::append_checksum_in_place(filtered.to_mut())?;
217                    }
218                    other => {
219                        return Err(Error::Unsupported(format!(
220                            "write_chunk cannot encode filter {other}"
221                        )));
222                    }
223                }
224            }
225        }
226        Ok(filtered)
227    }
228
229    fn replace_encoded_filter_buffer<'a>(filtered: &mut Cow<'a, [u8]>, filter_buf: &mut Vec<u8>) {
230        let previous = std::mem::replace(filtered, Cow::Owned(std::mem::take(filter_buf)));
231        if let Cow::Owned(mut previous) = previous {
232            previous.clear();
233            *filter_buf = previous;
234        }
235    }
236
237    fn validate_chunk_index_rewrite(
238        &mut self,
239        index_kind: WritableChunkIndexKind,
240        index_addr: u64,
241        info: &crate::hl::dataset::DatasetInfo,
242        chunk_coords: &[u64],
243        chunk_data_dims: &[u64],
244        unfiltered_chunk_bytes: usize,
245    ) -> Result<()> {
246        match index_kind {
247            WritableChunkIndexKind::FixedArray | WritableChunkIndexKind::ExtensibleArray => {
248                Self::linear_chunk_index(chunk_coords, &info.dataspace.dims, chunk_data_dims)?;
249                Ok(())
250            }
251            WritableChunkIndexKind::BTreeV1 => self.validate_btree_v1_chunk_index(index_addr),
252            WritableChunkIndexKind::BTreeV2 => self.validate_btree_v2_chunk_index(
253                index_addr,
254                info,
255                chunk_coords.len(),
256                unfiltered_chunk_bytes,
257            ),
258        }
259    }
260
261    fn validate_btree_v1_chunk_index(&mut self, btree_addr: u64) -> Result<()> {
262        let mut guard = self.inner.lock();
263        guard.reader.seek(btree_addr)?;
264        let mut magic = [0u8; 4];
265        guard.reader.read_bytes_into(&mut magic)?;
266        if magic != [b'T', b'R', b'E', b'E'] {
267            return Err(Error::InvalidFormat("invalid chunk B-tree magic".into()));
268        }
269        let node_type = guard.reader.read_u8()?;
270        if node_type != 1 {
271            return Err(Error::InvalidFormat(format!(
272                "expected raw-data chunk B-tree, got type {node_type}"
273            )));
274        }
275        Ok(())
276    }
277
278    fn validate_btree_v2_chunk_index(
279        &mut self,
280        index_addr: u64,
281        info: &crate::hl::dataset::DatasetInfo,
282        ndims: usize,
283        unfiltered_chunk_bytes: usize,
284    ) -> Result<()> {
285        let filtered = info
286            .filter_pipeline
287            .as_ref()
288            .map(|pipeline| !pipeline.filters.is_empty())
289            .unwrap_or(false);
290        let chunk_size_len = if filtered {
291            Self::filtered_chunk_size_len(
292                info.layout.version,
293                unfiltered_chunk_bytes,
294                self.superblock.sizeof_size,
295            )
296        } else {
297            0
298        };
299        let expected_record_size = Self::btree_v2_preflight_record_size(
300            usize::from(self.superblock.sizeof_addr),
301            filtered,
302            chunk_size_len,
303            ndims,
304        )?;
305
306        let mut guard = self.inner.lock();
307        let header = BTreeV2Header::read_at(&mut guard.reader, index_addr)?;
308        if usize::from(header.record_size) != expected_record_size {
309            return Err(Error::InvalidFormat(format!(
310                "v2 B-tree chunk record size {} does not match expected {expected_record_size}",
311                header.record_size
312            )));
313        }
314        if header.tree_type != 10 && header.tree_type != 11 {
315            return Err(Error::Unsupported(format!(
316                "write_chunk cannot update v2 B-tree type {} chunk indexes",
317                header.tree_type
318            )));
319        }
320        Ok(())
321    }
322
323    fn btree_v2_preflight_record_size(
324        sizeof_addr: usize,
325        filtered: bool,
326        chunk_size_len: usize,
327        ndims: usize,
328    ) -> Result<usize> {
329        let filter_bytes = if filtered {
330            chunk_size_len
331                .checked_add(4)
332                .ok_or_else(|| Error::InvalidFormat("v2 B-tree record size overflow".into()))?
333        } else {
334            0
335        };
336        let coord_bytes = ndims
337            .checked_mul(8)
338            .ok_or_else(|| Error::InvalidFormat("v2 B-tree record size overflow".into()))?;
339        sizeof_addr
340            .checked_add(filter_bytes)
341            .and_then(|value| value.checked_add(coord_bytes))
342            .ok_or_else(|| Error::InvalidFormat("v2 B-tree record size overflow".into()))
343    }
344
345    #[allow(clippy::too_many_arguments)]
346    fn rewrite_chunk_index(
347        &mut self,
348        index_kind: WritableChunkIndexKind,
349        index_addr: u64,
350        dataset_addr: u64,
351        info: &crate::hl::dataset::DatasetInfo,
352        chunk_coords: &[u64],
353        chunk_data_dims: &[u64],
354        filtered_len: u64,
355        chunk_addr: u64,
356        expected_len: usize,
357        element_size: usize,
358    ) -> Result<()> {
359        match index_kind {
360            WritableChunkIndexKind::FixedArray => self.rewrite_fixed_array_chunk(
361                index_addr,
362                dataset_addr,
363                info,
364                chunk_coords,
365                chunk_data_dims,
366                filtered_len,
367                chunk_addr,
368                expected_len,
369            ),
370            WritableChunkIndexKind::ExtensibleArray => self.rewrite_extensible_array_chunk(
371                index_addr,
372                info,
373                chunk_coords,
374                chunk_data_dims,
375                filtered_len,
376                chunk_addr,
377                expected_len,
378            ),
379            WritableChunkIndexKind::BTreeV2 => self.rewrite_btree_v2_chunk(
380                index_addr,
381                info,
382                chunk_coords,
383                chunk_data_dims,
384                filtered_len,
385                chunk_addr,
386                expected_len,
387            ),
388            WritableChunkIndexKind::BTreeV1 => self.rewrite_leaf_chunk_btree(
389                index_addr,
390                chunk_coords,
391                chunk_data_dims,
392                Self::u64_to_u32(filtered_len, "filtered chunk size")?,
393                chunk_addr,
394                Self::usize_to_u32(element_size, "datatype size")?,
395            ),
396        }
397    }
398
399    fn writable_chunk_index_kind(info: &crate::hl::dataset::DatasetInfo) -> WritableChunkIndexKind {
400        match info.layout.chunk_index_type {
401            Some(ChunkIndexType::FixedArray) => WritableChunkIndexKind::FixedArray,
402            Some(ChunkIndexType::ExtensibleArray) => WritableChunkIndexKind::ExtensibleArray,
403            Some(ChunkIndexType::BTreeV2) => WritableChunkIndexKind::BTreeV2,
404            _ => WritableChunkIndexKind::BTreeV1,
405        }
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412
413    #[test]
414    fn expected_chunk_len_rejects_element_count_overflow() {
415        let err = MutableFile::expected_chunk_len(&[u64::MAX, 2], 1).unwrap_err();
416        assert!(
417            err.to_string().contains("chunk element count"),
418            "unexpected error: {err}"
419        );
420    }
421
422    #[test]
423    fn validate_chunk_coords_rejects_out_of_extent_chunk() {
424        let err = MutableFile::validate_chunk_coords(&[10], &[10], &[5]).unwrap_err();
425        assert!(
426            err.to_string().contains("outside dataset extent"),
427            "unexpected error: {err}"
428        );
429    }
430}
hdf5_pure_rust/hl/mutable_file/write_chunk.rs

hdf5_pure_rust/hl/mutable_file/
write_chunk.rs
