anndata_hdf5/
lib.rs

1use anndata::{
2    backend::*,
3    data::{
4        DynArray, DynCowArray, DynScalar, SelectInfoBounds, SelectInfoElem, SelectInfoElemBounds,
5        Shape,
6    },
7};
8
9use anyhow::{Ok, Result, bail};
10use hdf5::{
11    File, Group, H5Type, Location, Selection,
12    dataset::Dataset,
13    types::IntSize::*,
14    types::{FloatSize, TypeDescriptor, VarLenUnicode},
15};
16use itertools::{EitherOrBoth, Itertools};
17use ndarray::{Array, ArrayD, ArrayView, CowArray, Dimension, IxDyn, SliceInfo, SliceInfoElem};
18use std::ops::{Deref, Index};
19use std::path::{Path, PathBuf};
20
21//-----------------------------------------------------------------------------
22// Type definitions
23//-----------------------------------------------------------------------------
24
25pub struct H5;
26
27pub struct H5File(File);
28
29impl Deref for H5File {
30    type Target = File;
31
32    fn deref(&self) -> &Self::Target {
33        &self.0
34    }
35}
36
37pub struct H5Group(Group);
38
39impl Deref for H5Group {
40    type Target = Group;
41
42    fn deref(&self) -> &Self::Target {
43        &self.0
44    }
45}
46
47pub struct H5Dataset(Dataset);
48
49impl Deref for H5Dataset {
50    type Target = Dataset;
51
52    fn deref(&self) -> &Self::Target {
53        &self.0
54    }
55}
56
57//-----------------------------------------------------------------------------
58// Backend implementation
59//-----------------------------------------------------------------------------
60
61impl Backend for H5 {
62    const NAME: &'static str = "hdf5";
63
64    type Store = H5File;
65    type Group = H5Group;
66    type Dataset = H5Dataset;
67
68    fn new<P: AsRef<Path>>(path: P) -> Result<Self::Store> {
69        Ok(H5File(File::create(path)?))
70    }
71
72    /// Opens a file as read-only, file must exist.
73    fn open<P: AsRef<Path>>(path: P) -> Result<Self::Store> {
74        Ok(File::open(path).map(H5File)?)
75    }
76
77    /// Opens a file as read/write, file must exist.
78    fn open_rw<P: AsRef<Path>>(path: P) -> Result<Self::Store> {
79        Ok(File::open_rw(path).map(H5File)?)
80    }
81}
82
83impl StoreOp<H5> for H5File {
84    fn filename(&self) -> PathBuf {
85        hdf5::Location::filename(self).into()
86    }
87
88    fn close(self) -> Result<()> {
89        Ok(self.0.close()?)
90    }
91}
92
93// Generic GroupOp implementations
94
95fn list(group: &Group) -> Result<Vec<String>> {
96    Ok(group.member_names()?)
97}
98
99fn create_group(group: &Group, name: &str) -> Result<H5Group> {
100    Ok(H5Group(group.create_group(name)?))
101}
102
103fn open_group(group: &Group, name: &str) -> Result<H5Group> {
104    Ok(H5Group(group.group(name)?))
105}
106
107fn new_dataset<T: BackendData>(
108    group: &Group,
109    name: &str,
110    shape: &Shape,
111    config: WriteConfig,
112) -> Result<H5Dataset> {
113    let dtype = T::DTYPE;
114    let mut builder = match dtype {
115        ScalarType::U8 => group.new_dataset::<u8>(),
116        ScalarType::U16 => group.new_dataset::<u16>(),
117        ScalarType::U32 => group.new_dataset::<u32>(),
118        ScalarType::U64 => group.new_dataset::<u64>(),
119        ScalarType::I8 => group.new_dataset::<i8>(),
120        ScalarType::I16 => group.new_dataset::<i16>(),
121        ScalarType::I32 => group.new_dataset::<i32>(),
122        ScalarType::I64 => group.new_dataset::<i64>(),
123        ScalarType::F32 => group.new_dataset::<f32>(),
124        ScalarType::F64 => group.new_dataset::<f64>(),
125        ScalarType::Bool => group.new_dataset::<bool>(),
126        ScalarType::String => group.new_dataset::<VarLenUnicode>(),
127    };
128
129    builder = if let Some(compression) = config.compression {
130        match compression {
131            Compression::Gzip(lvl) => builder.deflate(lvl),
132            Compression::Zst(lvl) => match dtype {
133                ScalarType::String => builder.deflate(3),
134                _ => builder.blosc_zstd(lvl, hdf5::filters::BloscShuffle::Byte),
135            },
136        }
137    } else {
138        builder
139    };
140
141    builder = if let Some(s) = config.block_size {
142        if s.as_ref().iter().all(|&x| x > 0) {
143            builder.chunk(s.as_ref())
144        } else {
145            builder
146        }
147    } else {
148        builder
149    };
150
151    let s: hdf5::Extents = hdf5::SimpleExtents::resizable(shape.as_ref()).into();
152    let dataset = builder.shape(s).create(name)?;
153    Ok(H5Dataset(dataset))
154}
155
156fn open_dataset(group: &Group, name: &str) -> Result<H5Dataset> {
157    Ok(H5Dataset(group.dataset(name)?))
158}
159
160fn delete(group: &Group, name: &str) -> Result<()> {
161    Ok(group.unlink(name)?)
162}
163
164fn exists(group: &Group, name: &str) -> Result<bool> {
165    Ok(group.link_exists(name))
166}
167
168fn create_scalar_data<D: BackendData>(group: &Group, name: &str, data: &D) -> Result<H5Dataset> {
169    match data.as_dyn() {
170        DynScalar::U8(x) => {
171            let dataset = group.new_dataset::<u8>().create(name)?;
172            dataset.write_scalar(&x)?;
173            Ok(dataset)
174        }
175        DynScalar::U16(x) => {
176            let dataset = group.new_dataset::<u16>().create(name)?;
177            dataset.write_scalar(&x)?;
178            Ok(dataset)
179        }
180        DynScalar::U32(x) => {
181            let dataset = group.new_dataset::<u32>().create(name)?;
182            dataset.write_scalar(&x)?;
183            Ok(dataset)
184        }
185        DynScalar::U64(x) => {
186            let dataset = group.new_dataset::<u64>().create(name)?;
187            dataset.write_scalar(&x)?;
188            Ok(dataset)
189        }
190        DynScalar::I8(x) => {
191            let dataset = group.new_dataset::<i8>().create(name)?;
192            dataset.write_scalar(&x)?;
193            Ok(dataset)
194        }
195        DynScalar::I16(x) => {
196            let dataset = group.new_dataset::<i16>().create(name)?;
197            dataset.write_scalar(&x)?;
198            Ok(dataset)
199        }
200        DynScalar::I32(x) => {
201            let dataset = group.new_dataset::<i32>().create(name)?;
202            dataset.write_scalar(&x)?;
203            Ok(dataset)
204        }
205        DynScalar::I64(x) => {
206            let dataset = group.new_dataset::<i64>().create(name)?;
207            dataset.write_scalar(&x)?;
208            Ok(dataset)
209        }
210        DynScalar::F32(x) => {
211            let dataset = group.new_dataset::<f32>().create(name)?;
212            dataset.write_scalar(&x)?;
213            Ok(dataset)
214        }
215        DynScalar::F64(x) => {
216            let dataset = group.new_dataset::<f64>().create(name)?;
217            dataset.write_scalar(&x)?;
218            Ok(dataset)
219        }
220        DynScalar::Bool(x) => {
221            let dataset = group.new_dataset::<bool>().create(name)?;
222            dataset.write_scalar(&x)?;
223            Ok(dataset)
224        }
225        DynScalar::String(x) => {
226            let dataset = group.new_dataset::<VarLenUnicode>().create(name)?;
227            dataset.write_scalar(&x.parse::<VarLenUnicode>().unwrap())?;
228            Ok(dataset)
229        }
230    }
231    .map(H5Dataset)
232}
233
234impl DatasetOp<H5> for H5Dataset {
235    fn dtype(&self) -> Result<ScalarType> {
236        let ty = match hdf5::Container::dtype(self)?.to_descriptor()? {
237            TypeDescriptor::Unsigned(U1) => ScalarType::U8,
238            TypeDescriptor::Unsigned(U2) => ScalarType::U16,
239            TypeDescriptor::Unsigned(U4) => ScalarType::U32,
240            TypeDescriptor::Unsigned(U8) => ScalarType::U64,
241            TypeDescriptor::Integer(U1) => ScalarType::I8,
242            TypeDescriptor::Integer(U2) => ScalarType::I16,
243            TypeDescriptor::Integer(U4) => ScalarType::I32,
244            TypeDescriptor::Integer(U8) => ScalarType::I64,
245            TypeDescriptor::Float(FloatSize::U4) => ScalarType::F32,
246            TypeDescriptor::Float(FloatSize::U8) => ScalarType::F64,
247            TypeDescriptor::Boolean => ScalarType::Bool,
248            TypeDescriptor::VarLenAscii => ScalarType::String,
249            TypeDescriptor::VarLenUnicode => ScalarType::String,
250            ty => bail!("Unsupported type: {ty:?}"),
251        };
252        Ok(ty)
253    }
254
255    fn shape(&self) -> Shape {
256        hdf5::Container::shape(self).into()
257    }
258
259    fn reshape(&mut self, shape: &Shape) -> Result<()> {
260        Ok(Dataset::resize(self, shape.as_ref())?)
261    }
262
263    fn read_scalar<T: BackendData>(&self) -> Result<T> {
264        let val = match T::DTYPE {
265            ScalarType::Bool => self.deref().read_scalar::<bool>()?.as_dyn(),
266            ScalarType::U8 => self.deref().read_scalar::<u8>()?.as_dyn(),
267            ScalarType::U16 => self.deref().read_scalar::<u16>()?.as_dyn(),
268            ScalarType::U32 => self.deref().read_scalar::<u32>()?.as_dyn(),
269            ScalarType::U64 => self.deref().read_scalar::<u64>()?.as_dyn(),
270            ScalarType::I8 => self.deref().read_scalar::<i8>()?.as_dyn(),
271            ScalarType::I16 => self.deref().read_scalar::<i16>()?.as_dyn(),
272            ScalarType::I32 => self.deref().read_scalar::<i32>()?.as_dyn(),
273            ScalarType::I64 => self.deref().read_scalar::<i64>()?.as_dyn(),
274            ScalarType::F32 => self.deref().read_scalar::<f32>()?.as_dyn(),
275            ScalarType::F64 => self.deref().read_scalar::<f64>()?.as_dyn(),
276            ScalarType::String => {
277                let s = self.deref().read_scalar::<VarLenUnicode>()?;
278                s.to_string().as_dyn()
279            }
280        };
281        BackendData::from_dyn(val)
282    }
283
284    fn read_array_slice<T, S, D>(&self, selection: &[S]) -> Result<Array<T, D>>
285    where
286        T: BackendData,
287        S: AsRef<SelectInfoElem>,
288        D: Dimension,
289    {
290        fn select<S, T, D>(arr_: &Array<T, D>, info: &[S]) -> Array<T, D>
291        where
292            S: AsRef<SelectInfoElem>,
293            T: Clone,
294            D: Dimension,
295        {
296            let arr = arr_.view().into_dyn();
297            let slices = info
298                .as_ref()
299                .iter()
300                .map(|x| match x.as_ref() {
301                    SelectInfoElem::Slice(slice) => Some(SliceInfoElem::from(*slice)),
302                    _ => None,
303                })
304                .collect::<Option<Vec<_>>>();
305            if let Some(slices) = slices {
306                arr.slice(slices.as_slice()).into_owned()
307            } else {
308                let shape = arr_.shape();
309                let select: Vec<_> = info
310                    .as_ref()
311                    .iter()
312                    .zip_longest(shape)
313                    .map(|ty| match ty {
314                        EitherOrBoth::Both(x, n) => SelectInfoElemBounds::new(x.as_ref(), *n),
315                        EitherOrBoth::Right(n) => SelectInfoElemBounds::new(
316                            &SelectInfoElem::Slice(anndata::data::slice::SLICE_FULL),
317                            *n,
318                        ),
319                        _ => panic!("inconsistent selection length"),
320                    })
321                    .collect();
322                let new_shape = select.iter().map(|x| x.len()).collect::<Vec<_>>();
323                ArrayD::from_shape_fn(new_shape, |idx| {
324                    let new_idx: Vec<_> =
325                        (0..idx.ndim()).map(|i| select[i].index(idx[i])).collect();
326                    arr.index(new_idx.as_slice()).clone()
327                })
328            }
329            .into_dimensionality::<D>()
330            .unwrap()
331        }
332
333        fn read_arr<T, S, D>(dataset: &H5Dataset, selection: &[S]) -> Result<Array<T, D>>
334        where
335            T: H5Type + BackendData,
336            S: AsRef<SelectInfoElem>,
337            D: Dimension,
338        {
339            if selection.iter().any(|x| x.as_ref().is_index()) {
340                // fancy indexing is too slow, just read all
341                let arr = dataset.deref().read::<T, D>()?;
342                Ok(select(&arr, selection))
343            } else {
344                let (select, shape) = into_selection(selection, dataset.shape());
345                if matches!(select, Selection::Points(_)) {
346                    let slice_1d = hdf5::Container::read_slice_1d::<T, _>(dataset, select)?;
347                    Ok(slice_1d
348                        .into_shape_with_order(shape.as_ref())?
349                        .into_dimensionality::<D>()?)
350                } else {
351                    Ok(hdf5::Container::read_slice::<T, _, D>(dataset, select)?)
352                }
353            }
354        }
355
356        let array: DynArray = match T::DTYPE {
357            ScalarType::I8 => read_arr::<i8, _, D>(self, selection)?.into(),
358            ScalarType::I16 => read_arr::<i16, _, D>(self, selection)?.into(),
359            ScalarType::I32 => read_arr::<i32, _, D>(self, selection)?.into(),
360            ScalarType::I64 => read_arr::<i64, _, D>(self, selection)?.into(),
361            ScalarType::U8 => read_arr::<u8, _, D>(self, selection)?.into(),
362            ScalarType::U16 => read_arr::<u16, _, D>(self, selection)?.into(),
363            ScalarType::U32 => read_arr::<u32, _, D>(self, selection)?.into(),
364            ScalarType::U64 => read_arr::<u64, _, D>(self, selection)?.into(),
365            ScalarType::F32 => read_arr::<f32, _, D>(self, selection)?.into(),
366            ScalarType::F64 => read_arr::<f64, _, D>(self, selection)?.into(),
367            ScalarType::Bool => read_arr::<bool, _, D>(self, selection)?.into(),
368            ScalarType::String => {
369                if selection.as_ref().iter().any(|x| x.as_ref().is_index()) {
370                    // fancy indexing is too slow, just read all
371                    let arr = self.deref().read::<VarLenUnicode, D>()?;
372                    let arr_ = arr.map(|s| s.to_string());
373                    let r: Result<_> = Ok(select(&arr_, selection));
374                    r
375                } else {
376                    let (select, shape) = into_selection(selection, self.shape());
377                    let arr: Result<_> = if matches!(select, Selection::Points(_)) {
378                        let slice_1d = self.deref().read_slice_1d::<VarLenUnicode, _>(select)?;
379                        Ok(slice_1d
380                            .into_shape_with_order(shape.as_ref())?
381                            .into_dimensionality::<D>()?)
382                    } else {
383                        Ok(self.deref().read_slice::<VarLenUnicode, _, D>(select)?)
384                    };
385                    Ok(arr?.map(|s| s.to_string()))
386                }?
387                .into()
388                /*
389                let arr = read_arr::<VarLenUnicode, _, _, D>(dataset, selection)?;
390                let arr = arr.map(|s| s.to_string());
391                arr.into()
392                */
393            }
394        };
395        Ok(BackendData::from_dyn_arr(array)?.into_dimensionality::<D>()?)
396    }
397
398    fn write_array_slice<S, T, D>(&self, data: CowArray<'_, T, D>, selection: &[S]) -> Result<()>
399    where
400        T: BackendData,
401        S: AsRef<SelectInfoElem>,
402        D: Dimension,
403    {
404        fn write_array_impl<T, S>(
405            container: &H5Dataset,
406            arr: CowArray<'_, T, IxDyn>,
407            selection: &[S],
408        ) -> Result<()>
409        where
410            T: H5Type + Clone,
411            S: AsRef<SelectInfoElem>,
412        {
413            let (select, _) = into_selection(selection, container.shape());
414            container
415                .deref()
416                .write_slice(&arr.as_standard_layout(), select)?;
417            Ok(())
418        }
419
420        match BackendData::into_dyn_arr(data.into_dyn()) {
421            DynCowArray::U8(x) => write_array_impl(self, x, selection),
422            DynCowArray::U16(x) => write_array_impl(self, x, selection),
423            DynCowArray::U32(x) => write_array_impl(self, x, selection),
424            DynCowArray::U64(x) => write_array_impl(self, x, selection),
425            DynCowArray::I8(x) => write_array_impl(self, x, selection),
426            DynCowArray::I16(x) => write_array_impl(self, x, selection),
427            DynCowArray::I32(x) => write_array_impl(self, x, selection),
428            DynCowArray::I64(x) => write_array_impl(self, x, selection),
429            DynCowArray::F32(x) => write_array_impl(self, x, selection),
430            DynCowArray::F64(x) => write_array_impl(self, x, selection),
431            DynCowArray::Bool(x) => write_array_impl(self, x, selection),
432            DynCowArray::String(x) => {
433                let data: Array<VarLenUnicode, _> = x.map(|x| x.parse().unwrap());
434                write_array_impl(self, data.into(), selection)
435            }
436        }
437    }
438}
439
440// Generic `LocationOp` functions
441
442fn file(loc: &Location) -> Result<H5File> {
443    Ok(H5File(hdf5::Location::file(loc)?))
444}
445
446fn path(loc: &Location) -> PathBuf {
447    hdf5::Location::name(loc).into()
448}
449
450//-----------------------------------------------------------------------------
451// Derived implementations
452//-----------------------------------------------------------------------------
453
454impl GroupOp<H5> for H5File {
455    fn list(&self) -> Result<Vec<String>> {
456        list(self)
457    }
458
459    fn new_group(&self, name: &str) -> Result<<H5 as Backend>::Group> {
460        create_group(self, name)
461    }
462
463    fn open_group(&self, name: &str) -> Result<<H5 as Backend>::Group> {
464        open_group(self, name)
465    }
466
467    fn new_empty_dataset<T: BackendData>(
468        &self,
469        name: &str,
470        shape: &Shape,
471        config: WriteConfig,
472    ) -> Result<<H5 as Backend>::Dataset> {
473        new_dataset::<T>(self, name, shape, config)
474    }
475
476    fn open_dataset(&self, name: &str) -> Result<<H5 as Backend>::Dataset> {
477        open_dataset(self, name)
478    }
479
480    fn delete(&self, name: &str) -> Result<()> {
481        delete(self, name)
482    }
483
484    fn exists(&self, name: &str) -> Result<bool> {
485        exists(self, name)
486    }
487
488    fn new_scalar_dataset<D: BackendData>(
489        &self,
490        name: &str,
491        data: &D,
492    ) -> Result<<H5 as Backend>::Dataset> {
493        create_scalar_data(self, name, data)
494    }
495}
496
497impl GroupOp<H5> for H5Group {
498    fn list(&self) -> Result<Vec<String>> {
499        list(self)
500    }
501
502    fn new_group(&self, name: &str) -> Result<<H5 as Backend>::Group> {
503        create_group(self, name)
504    }
505
506    fn open_group(&self, name: &str) -> Result<<H5 as Backend>::Group> {
507        open_group(self, name)
508    }
509
510    fn new_empty_dataset<T: BackendData>(
511        &self,
512        name: &str,
513        shape: &Shape,
514        config: WriteConfig,
515    ) -> Result<<H5 as Backend>::Dataset> {
516        new_dataset::<T>(self, name, shape, config)
517    }
518
519    fn open_dataset(&self, name: &str) -> Result<<H5 as Backend>::Dataset> {
520        open_dataset(self, name)
521    }
522
523    fn delete(&self, name: &str) -> Result<()> {
524        delete(self, name)
525    }
526
527    fn exists(&self, name: &str) -> Result<bool> {
528        exists(self, name)
529    }
530
531    fn new_scalar_dataset<D: BackendData>(
532        &self,
533        name: &str,
534        data: &D,
535    ) -> Result<<H5 as Backend>::Dataset> {
536        create_scalar_data(self, name, data)
537    }
538}
539
540impl AttributeOp<H5> for H5Group {
541    fn store(&self) -> Result<<H5 as Backend>::Store> {
542        file(self)
543    }
544
545    fn path(&self) -> PathBuf {
546        path(self)
547    }
548
549    fn new_json_attr(&mut self, name: &str, value: &Value) -> Result<()> {
550        match value {
551            Value::Null => Ok(()),
552            Value::Bool(b) => write_scalar_attr(self, name, *b),
553            Value::Number(n) => n
554                .as_u64()
555                .map(|i| write_scalar_attr(self, name, i))
556                .or_else(|| n.as_i64().map(|i| write_scalar_attr(self, name, i)))
557                .or_else(|| n.as_f64().map(|i| write_scalar_attr(self, name, i)))
558                .expect("number cannot be converted to u64, i64 or f64"),
559            Value::String(s) => write_scalar_attr(self, name, s.clone()),
560            Value::Array(_) => json_to_ndarray(value, |x| x.as_i64())?
561                .map(|x| write_array_attr(self, name, &x))
562                .or_else(|| {
563                    json_to_ndarray(value, |x| x.as_f64())
564                        .unwrap()
565                        .map(|x| write_array_attr(self, name, &x))
566                })
567                .or_else(|| {
568                    json_to_ndarray(value, |x| x.as_str().map(|s| s.to_string()))
569                        .unwrap()
570                        .map(|x| write_array_attr(self, name, &x))
571                })
572                .expect("array cannot be converted to i64, f64 or string"),
573            Value::Object(_) => bail!("attributes of object type are not supported"),
574        }
575    }
576
577    fn get_json_attr(&self, name: &str) -> Result<Value> {
578        if self.attr(name)?.is_scalar() {
579            read_scalar_attr(self, name)
580        } else {
581            read_array_attr(self, name)
582        }
583    }
584}
585
586impl AttributeOp<H5> for H5Dataset {
587    fn store(&self) -> Result<<H5 as Backend>::Store> {
588        file(self)
589    }
590
591    fn path(&self) -> PathBuf {
592        path(self)
593    }
594
595    fn new_json_attr(&mut self, name: &str, value: &Value) -> Result<()> {
596        match value {
597            Value::Null => Ok(()),
598            Value::Bool(b) => write_scalar_attr(self, name, *b),
599            Value::Number(n) => n
600                .as_u64()
601                .map(|i| write_scalar_attr(self, name, i))
602                .or_else(|| n.as_i64().map(|i| write_scalar_attr(self, name, i)))
603                .or_else(|| n.as_f64().map(|i| write_scalar_attr(self, name, i)))
604                .expect("number cannot be converted to u64, i64 or f64"),
605            Value::String(s) => write_scalar_attr(self, name, s.clone()),
606            Value::Array(_) => json_to_ndarray(value, |x| x.as_i64())?
607                .map(|x| write_array_attr(self, name, &x))
608                .or_else(|| {
609                    json_to_ndarray(value, |x| x.as_f64())
610                        .unwrap()
611                        .map(|x| write_array_attr(self, name, &x))
612                })
613                .or_else(|| {
614                    json_to_ndarray(value, |x| x.as_str().map(|s| s.to_string()))
615                        .unwrap()
616                        .map(|x| write_array_attr(self, name, &x))
617                })
618                .expect("array cannot be converted to i64, f64 or string"),
619            Value::Object(_) => bail!("attributes of object type are not supported"),
620        }
621    }
622
623    fn get_json_attr(&self, name: &str) -> Result<Value> {
624        if self.attr(name)?.is_scalar() {
625            read_scalar_attr(self, name)
626        } else {
627            read_array_attr(self, name)
628        }
629    }
630}
631
632///////////////////////////////////////////////////////////////////////////////
633/// Auxiliary functions
634///////////////////////////////////////////////////////////////////////////////
635fn read_scalar_attr(loc: &Location, name: &str) -> Result<Value> {
636    let attr = loc.attr(name)?;
637    let result = match attr.dtype()?.to_descriptor()? {
638        TypeDescriptor::VarLenUnicode => attr.read_scalar::<VarLenUnicode>()?.to_string().into(),
639        TypeDescriptor::VarLenAscii => attr.read_scalar::<VarLenUnicode>()?.to_string().into(),
640        TypeDescriptor::Boolean => attr.read_scalar::<bool>()?.into(),
641        TypeDescriptor::Unsigned(_) => attr.read_scalar::<u64>()?.into(),
642        TypeDescriptor::Integer(_) => attr.read_scalar::<i64>()?.into(),
643        TypeDescriptor::Float(_) => attr.read_scalar::<f64>()?.into(),
644        v => bail!("Unsupported type {v}"),
645    };
646    Ok(result)
647}
648
649fn read_array_attr(loc: &Location, name: &str) -> Result<Value> {
650    let attr = loc.attr(name)?;
651    let result = match attr.dtype()?.to_descriptor()? {
652        TypeDescriptor::VarLenUnicode => {
653            ndarray_to_json(&attr.read::<VarLenUnicode, IxDyn>()?.mapv(|x| x.to_string()))
654        }
655        TypeDescriptor::VarLenAscii => {
656            ndarray_to_json(&attr.read::<VarLenUnicode, IxDyn>()?.mapv(|x| x.to_string()))
657        }
658        TypeDescriptor::Boolean => ndarray_to_json(&attr.read::<bool, IxDyn>()?),
659        TypeDescriptor::Unsigned(_) => ndarray_to_json(&attr.read::<u64, IxDyn>()?),
660        TypeDescriptor::Integer(_) => ndarray_to_json(&attr.read::<i64, IxDyn>()?),
661        TypeDescriptor::Float(_) => ndarray_to_json(&attr.read::<f64, IxDyn>()?),
662        v => bail!("Unsupported type {v}"),
663    };
664    Ok(result)
665}
666
667fn write_array_attr<'a, A, D, Dim>(loc: &Location, name: &str, value: A) -> Result<()>
668where
669    A: Into<ArrayView<'a, D, Dim>>,
670    D: BackendData,
671    Dim: Dimension,
672{
673    del_attr(loc, name);
674    let value = value.into().into_dyn().into();
675    match BackendData::into_dyn_arr(value) {
676        DynCowArray::U8(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
677        DynCowArray::U16(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
678        DynCowArray::U32(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
679        DynCowArray::U64(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
680        DynCowArray::I8(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
681        DynCowArray::I16(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
682        DynCowArray::I32(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
683        DynCowArray::I64(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
684        DynCowArray::F32(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
685        DynCowArray::F64(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
686        DynCowArray::Bool(x) => loc.new_attr_builder().with_data(x.view()).create(name)?,
687        DynCowArray::String(x) => {
688            let data: Array<VarLenUnicode, Dim> =
689                x.map(|x| x.parse().unwrap()).into_dimensionality()?;
690            loc.new_attr_builder().with_data(data.view()).create(name)?
691        }
692    };
693    Ok(())
694}
695
696fn write_scalar_attr<D: BackendData>(loc: &Location, name: &str, value: D) -> Result<()> {
697    del_attr(loc, name);
698    match value.as_dyn() {
699        DynScalar::U8(x) => loc.new_attr::<u8>().create(name)?.write_scalar(&x)?,
700        DynScalar::U16(x) => loc.new_attr::<u16>().create(name)?.write_scalar(&x)?,
701        DynScalar::U32(x) => loc.new_attr::<u32>().create(name)?.write_scalar(&x)?,
702        DynScalar::U64(x) => loc.new_attr::<u64>().create(name)?.write_scalar(&x)?,
703        DynScalar::I8(x) => loc.new_attr::<i8>().create(name)?.write_scalar(&x)?,
704        DynScalar::I16(x) => loc.new_attr::<i16>().create(name)?.write_scalar(&x)?,
705        DynScalar::I32(x) => loc.new_attr::<i32>().create(name)?.write_scalar(&x)?,
706        DynScalar::I64(x) => loc.new_attr::<i64>().create(name)?.write_scalar(&x)?,
707        DynScalar::F32(x) => loc.new_attr::<f32>().create(name)?.write_scalar(&x)?,
708        DynScalar::F64(x) => loc.new_attr::<f64>().create(name)?.write_scalar(&x)?,
709        DynScalar::Bool(x) => loc.new_attr::<bool>().create(name)?.write_scalar(&x)?,
710        DynScalar::String(x) => {
711            let value_: VarLenUnicode = x.parse().unwrap();
712            loc.new_attr::<VarLenUnicode>()
713                .create(name)?
714                .write_scalar(&value_)?
715        }
716    };
717    Ok(())
718}
719
720fn into_selection<S, E>(selection: S, shape: Shape) -> (Selection, Shape)
721where
722    S: AsRef<[E]>,
723    E: AsRef<SelectInfoElem>,
724{
725    if selection.as_ref().iter().all(|x| x.as_ref().is_full()) {
726        (Selection::All, shape)
727    } else {
728        let bounded_selection = SelectInfoBounds::new(&selection, &shape);
729        let out_shape = bounded_selection.out_shape();
730        if let Some(idx) = bounded_selection.try_into_indices() {
731            (Selection::from(idx), out_shape)
732        } else {
733            let slice: SliceInfo<_, _, _> = bounded_selection.try_into().unwrap();
734            (Selection::try_from(slice).unwrap(), out_shape)
735        }
736    }
737}
738
739fn del_attr(loc: &Location, name: &str) {
740    unsafe {
741        let c_name = std::ffi::CString::new(name).unwrap().into_raw();
742        if hdf5_sys::h5a::H5Aexists(loc.id(), c_name) != 0 {
743            hdf5_sys::h5a::H5Adelete(loc.id(), c_name);
744        }
745    }
746}
747
748fn json_to_ndarray<F, T>(json: &Value, f: F) -> Result<Option<ArrayD<T>>>
749where
750    F: Fn(&Value) -> Option<T>,
751{
752    // Recursively determine the shape of the JSON array
753    fn get_shape(value: &Value) -> Vec<usize> {
754        let mut shape = Vec::new();
755        let mut current = value;
756        while let Value::Array(arr) = current {
757            shape.push(arr.len());
758            if arr.is_empty() {
759                break;
760            }
761            current = &arr[0];
762        }
763        shape
764    }
765
766    // Recursively flatten the JSON array to a Vec<i32>
767    fn flatten_array<F, T>(value: &Value, f: &F) -> Option<Vec<T>>
768    where
769        F: Fn(&Value) -> Option<T>,
770    {
771        match value {
772            Value::Array(arr) => {
773                let mut flattened = Vec::new();
774                for item in arr {
775                    flattened.extend(flatten_array(item, f)?);
776                }
777                Some(flattened)
778            }
779            v => Some(vec![f(v)?]),
780        }
781    }
782
783    // Get the shape and flatten the JSON
784    let shape = get_shape(json);
785    if let Some(flattened_data) = flatten_array(json, &f) {
786        Ok(Some(ArrayD::from_shape_vec(IxDyn(&shape), flattened_data)?))
787    } else {
788        Ok(None)
789    }
790}
791
792fn ndarray_to_json<T: Into<Value> + Clone>(array: &ArrayD<T>) -> Value {
793    // Helper function to recursively convert ndarray to nested Vecs
794    fn recursive_convert<T: Into<Value> + Clone>(array: &ArrayD<T>) -> Value {
795        if array.ndim() == 1 {
796            // Base case: 1D array, convert to Vec and serialize
797            let vec = array.iter().cloned().collect::<Vec<T>>();
798            vec.into()
799        } else {
800            // Recursive case: split along the first axis and apply recursively
801            let nested_vec = array
802                .outer_iter()
803                .map(|sub_array| recursive_convert(&sub_array.to_owned().into_dyn()))
804                .collect::<Vec<Value>>();
805            Value::Array(nested_vec)
806        }
807    }
808
809    recursive_convert(array)
810}
811
812/// test module
813#[cfg(test)]
814mod tests {
815    use super::*;
816    use anndata::s;
817    use ndarray::{Array1, Axis, Ix1, concatenate};
818    use ndarray_rand::RandomExt;
819    use ndarray_rand::rand_distr::Uniform;
820    use std::path::PathBuf;
821    use tempfile::tempdir;
822
823    pub fn with_tmp_dir<T, F: FnMut(PathBuf) -> T>(mut func: F) -> T {
824        let dir = tempdir().unwrap();
825        let path = dir.path().to_path_buf();
826        func(path)
827    }
828
829    fn with_tmp_path<T, F: Fn(PathBuf) -> T>(func: F) -> T {
830        with_tmp_dir(|dir| func(dir.join("temp.h5")))
831    }
832
833    #[test]
834    fn test_basic() -> Result<()> {
835        with_tmp_path(|path| {
836            let file = H5::new(path.clone())?;
837            let group = file.new_group("group")?;
838            let subgroup = group.new_group("subgroup")?;
839
840            assert_eq!(subgroup.path(), PathBuf::from("/group/subgroup"));
841            Ok(())
842        })
843    }
844
845    #[test]
846    fn test_write_empty() -> Result<()> {
847        with_tmp_path(|path| {
848            let file = H5::new(&path)?;
849            let group = file.new_group("group")?;
850            let config = WriteConfig {
851                ..Default::default()
852            };
853
854            let empty = Array1::<u8>::from_vec(Vec::new());
855            let dataset = group.new_array_dataset("test", empty.view().into(), config)?;
856            assert_eq!(empty, dataset.read_array::<u8, Ix1>()?);
857            Ok(())
858        })
859    }
860
861    #[test]
862    fn test_write_slice() -> Result<()> {
863        with_tmp_path(|path| -> Result<()> {
864            let file = H5::new(&path)?;
865            let config = WriteConfig {
866                ..Default::default()
867            };
868
869            let mut dataset =
870                file.new_empty_dataset::<i32>("test", &[20, 50].as_slice().into(), config)?;
871            let arr = Array::random((20, 50), Uniform::new(0, 100).unwrap());
872
873            // Repeatitive writes
874            dataset.write_array_slice(arr.view().into(), s![.., ..].as_ref())?;
875            dataset.write_array_slice(arr.view().into(), s![.., ..].as_ref())?;
876
877            // Out-of-bounds writes should fail
878            assert!(
879                dataset
880                    .write_array_slice(arr.view().into(), s![20..40, ..].as_ref())
881                    .is_err()
882            );
883
884            // Reshape and write
885            dataset.reshape(&[40, 50].as_slice().into())?;
886            dataset.write_array_slice(arr.view().into(), s![20..40, ..].as_ref())?;
887
888            // Read back is OK
889            let merged = concatenate(Axis(0), &[arr.view(), arr.view()])?;
890            assert_eq!(merged, dataset.read_array::<i32, _>()?);
891
892            // Shrinking is OK
893            dataset.reshape(&[20, 50].as_slice().into())?;
894            assert_eq!(arr, dataset.read_array::<i32, _>()?);
895
896            Ok(())
897        })
898    }
899}
anndata_hdf5/lib.rs

anndata_hdf5/
lib.rs
