anndata_memory/ad/
helpers.rs

1use std::{
2    collections::HashMap,
3    fmt,
4    ops::{Deref, DerefMut},
5};
6
7use anndata::{
8    backend::DataType,
9    container::{Axis, Dim},
10    data::{DataFrameIndex, DynArray, DynCscMatrix, DynCsrMatrix, Element, SelectInfoElem, Shape},
11    ArrayData, Data, HasShape, Selectable,
12};
13use anyhow::bail;
14
15use ndarray::Array2;
16use polars::{
17    frame::DataFrame,
18    prelude::{Column, IdxCa, NamedFrom},
19    series::Series,
20};
21
22use crate::base::RwSlot;
23use crate::{base::DeepClone, utils::subset_dyn_csc_matrix, utils::subset_dyn_csr_matrix};
24
25impl DeepClone for ArrayData {
26    fn deep_clone(&self) -> Self {
27        self.clone()
28    }
29}
30
31pub struct IMArrayElement(pub RwSlot<ArrayData>);
32
33impl IMArrayElement {
34    pub fn new(data: ArrayData) -> Self {
35        IMArrayElement(RwSlot::new(data))
36    }
37
38    pub fn get_type(&self) -> anyhow::Result<DataType> {
39        Ok(self.0.read_inner().data_type())
40    }
41
42    pub fn get_shape(&self) -> anyhow::Result<Shape> {
43        Ok(self.0.read_inner().shape())
44    }
45
46    pub fn get_data(&self) -> anyhow::Result<ArrayData> {
47        Ok(self.0.read_inner().clone())
48    }
49
50    pub fn set_data(&self, data: ArrayData) -> anyhow::Result<()> {
51        let mut write_guard = self.0.lock_write();
52        let d = write_guard.deref_mut();
53        *d = Some(data);
54        Ok(())
55    }
56
57    pub fn subset_inplace(&self, s: &[&SelectInfoElem]) -> anyhow::Result<()> {
58        let mut write_guard = self.0.write_inner();
59        let d = write_guard.deref_mut();
60
61        // Replace the data with a placeholder to take ownership
62        let placeholder = ArrayData::Array(DynArray::from(Array2::<f64>::zeros((0, 0))));
63        let data = std::mem::replace(d, placeholder);
64
65        // Process the data by consuming it
66        let result = match data {
67            ArrayData::Array(arr) => Ok(ArrayData::Array(arr).select(s)),
68            ArrayData::DataFrame(df) => Ok(ArrayData::DataFrame(df).select(s)),
69            ArrayData::CsrNonCanonical(csr) => Ok(ArrayData::CsrNonCanonical(csr).select(s)),
70            ArrayData::CsrMatrix(dyn_csr) => {
71                subset_dyn_csr_matrix(dyn_csr, s).map(ArrayData::CsrMatrix)
72            }
73            ArrayData::CscMatrix(dyn_csc) => {
74                subset_dyn_csc_matrix(dyn_csc, s).map(ArrayData::CscMatrix)
75            }
76        };
77
78        match result {
79            Ok(processed) => {
80                *d = processed;
81                Ok(())
82            }
83            Err(e) => {
84                Err(e)
85            }
86        }
87    }
88
89    pub fn convert_matrix_format(&self) -> anyhow::Result<()> {
90        let mut write_guard = self.0.write_inner();
91        let d = write_guard.deref_mut();
92
93        // Create a placeholder that we can swap with - use an empty dense array as it's likely the smallest
94        let ddata: Array2<f64> = Array2::zeros((0, 0));
95        let placeholder = ArrayData::Array(DynArray::from(ddata));
96
97        // Take ownership using replace
98        let matrix_data = std::mem::replace(d, placeholder);
99
100        let converted = match matrix_data {
101            ArrayData::CsrMatrix(dyn_csr_matrix) => {
102                let csc = match dyn_csr_matrix {
103                    DynCsrMatrix::F64(m) => DynCscMatrix::F64(m.transpose_as_csc()),
104                    DynCsrMatrix::F32(m) => DynCscMatrix::F32(m.transpose_as_csc()),
105                    DynCsrMatrix::I64(m) => DynCscMatrix::I64(m.transpose_as_csc()),
106                    DynCsrMatrix::I32(m) => DynCscMatrix::I32(m.transpose_as_csc()),
107                    DynCsrMatrix::I16(m) => DynCscMatrix::I16(m.transpose_as_csc()),
108                    DynCsrMatrix::I8(m) => DynCscMatrix::I8(m.transpose_as_csc()),
109                    DynCsrMatrix::U64(m) => DynCscMatrix::U64(m.transpose_as_csc()),
110                    DynCsrMatrix::U32(m) => DynCscMatrix::U32(m.transpose_as_csc()),
111                    DynCsrMatrix::U16(m) => DynCscMatrix::U16(m.transpose_as_csc()),
112                    DynCsrMatrix::U8(m) => DynCscMatrix::U8(m.transpose_as_csc()),
113                    DynCsrMatrix::Bool(m) => DynCscMatrix::Bool(m.transpose_as_csc()),
114                    DynCsrMatrix::String(m) => DynCscMatrix::String(m.transpose_as_csc()),
115                };
116                ArrayData::CscMatrix(csc)
117            }
118            ArrayData::CscMatrix(dyn_csc_matrix) => {
119                let csr = match dyn_csc_matrix {
120                    DynCscMatrix::F64(m) => DynCsrMatrix::F64(m.transpose_as_csr()),
121                    DynCscMatrix::F32(m) => DynCsrMatrix::F32(m.transpose_as_csr()),
122                    DynCscMatrix::I64(m) => DynCsrMatrix::I64(m.transpose_as_csr()),
123                    DynCscMatrix::I32(m) => DynCsrMatrix::I32(m.transpose_as_csr()),
124                    DynCscMatrix::I16(m) => DynCsrMatrix::I16(m.transpose_as_csr()),
125                    DynCscMatrix::I8(m) => DynCsrMatrix::I8(m.transpose_as_csr()),
126                    DynCscMatrix::U64(m) => DynCsrMatrix::U64(m.transpose_as_csr()),
127                    DynCscMatrix::U32(m) => DynCsrMatrix::U32(m.transpose_as_csr()),
128                    DynCscMatrix::U16(m) => DynCsrMatrix::U16(m.transpose_as_csr()),
129                    DynCscMatrix::U8(m) => DynCsrMatrix::U8(m.transpose_as_csr()),
130                    DynCscMatrix::Bool(m) => DynCsrMatrix::Bool(m.transpose_as_csr()),
131                    DynCscMatrix::String(m) => DynCsrMatrix::String(m.transpose_as_csr()),
132                };
133                ArrayData::CsrMatrix(csr)
134            }
135            _ => {
136                // Put back the original value since we're erroring
137                *d = matrix_data;
138                bail!("This datatype is not supported, only CSC and CSR matrices are supported.")
139            }
140        };
141
142        *d = converted;
143        Ok(())
144    }
145
146    pub fn subset(&self, s: &[&SelectInfoElem]) -> anyhow::Result<Self> {
147        let read_guard = self.0.read_inner();
148        let d = read_guard.deref();
149
150        // Return a new ArrayData by selecting from d
151        Ok(IMArrayElement::new(d.select(s)))
152    }
153
154    pub fn deep_clone_content(&self) -> anyhow::Result<ArrayData> {
155        Ok(self.0.read_inner().clone())
156    }
157}
158
159impl DeepClone for IMArrayElement {
160    fn deep_clone(&self) -> Self {
161        IMArrayElement(self.0.deep_clone())
162    }
163}
164
165impl Clone for IMArrayElement {
166    fn clone(&self) -> Self {
167        IMArrayElement(self.0.clone())
168    }
169}
170
171pub struct IMDataFrameElement(RwSlot<InnerIMDataFrame>);
172
173pub struct InnerIMDataFrame {
174    df: DataFrame,
175    pub index: DataFrameIndex,
176}
177
178impl DeepClone for InnerIMDataFrame {
179    fn deep_clone(&self) -> Self {
180        self.clone()
181    }
182}
183
184impl Clone for InnerIMDataFrame {
185    fn clone(&self) -> Self {
186        InnerIMDataFrame {
187            df: self.df.clone(),
188            index: self.index.clone(),
189        }
190    }
191}
192
193impl Clone for IMDataFrameElement {
194    /// Shallow clone of the IMDataFrameElement
195    fn clone(&self) -> Self {
196        IMDataFrameElement(self.0.clone())
197    }
198}
199
200impl IMDataFrameElement {
201    pub fn new(df: DataFrame, index: DataFrameIndex) -> Self {
202        if df.height() == 0 {
203            let tmp_df =
204                DataFrame::new(vec![Column::new("index".into(), &index.clone().into_vec())])
205                    .unwrap();
206            return IMDataFrameElement(RwSlot::new(InnerIMDataFrame { df: tmp_df, index }));
207        }
208        if df.height() != index.len() {
209            panic!("Length of index does not match length of DataFrame");
210        }
211        IMDataFrameElement(RwSlot::new(InnerIMDataFrame { df, index }))
212    }
213
214    pub fn get_data(&self) -> DataFrame {
215        self.0.read_inner().df.clone()
216    }
217
218    pub fn get_index(&self) -> DataFrameIndex {
219        self.0.read_inner().index.clone()
220    }
221
222    pub fn set_both(&self, df: DataFrame, index: DataFrameIndex) -> anyhow::Result<()> {
223        let mut write_guard = self.0.lock_write();
224        let d = write_guard.as_mut();
225        match d {
226            Some(data) => {
227                if index.len() != df.height() {
228                    return Err(anyhow::anyhow!(
229                        "Length of index does not match length of DataFrame"
230                    ));
231                }
232
233                data.df = df;
234                data.index = index;
235                Ok(())
236            }
237            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
238        }
239    }
240
241    pub fn set_data(&self, df: DataFrame) -> anyhow::Result<()> {
242        let mut write_guard = self.0.lock_write();
243        let d = write_guard.as_mut();
244        match d {
245            Some(data) => {
246                if data.index.len() != df.height() || data.df.height() != df.height() {
247                    return Err(anyhow::anyhow!(
248                        "Length of index does not match length of DataFrame"
249                    ));
250                }
251                data.df = df;
252                Ok(())
253            }
254            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
255        }
256    }
257
258    pub fn set_index(&self, index: DataFrameIndex) -> anyhow::Result<()> {
259        let mut write_guard = self.0.lock_write();
260        let d = write_guard.as_mut();
261        match d {
262            Some(data) => {
263                if data.df.height() != index.len() || data.index.len() != index.len() {
264                    return Err(anyhow::anyhow!(
265                        "Length of index does not match length of DataFrame"
266                    ));
267                }
268
269                data.index = index;
270                Ok(())
271            }
272            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
273        }
274    }
275
276    pub fn attach_column_to_df(&self, column: Series) -> anyhow::Result<()> {
277        let mut write_guard = self.0.lock_write();
278        let d = write_guard.as_mut();
279        match d {
280            Some(data) => {
281                if data.df.height() != column.len() {
282                    return Err(anyhow::anyhow!(
283                        "Length of column does not match length of DataFrame"
284                    ));
285                }
286                data.df.with_column(column)?;
287                Ok(())
288            }
289            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
290        }
291    }
292
293    pub fn remove_column_from_df(&self, column_name: &str) -> anyhow::Result<()> {
294        let mut write_guard = self.0.lock_write();
295        let d = write_guard.as_mut();
296        match d {
297            Some(data) => {
298                let _ = data.df.drop_in_place(column_name)?;
299                Ok(())
300            }
301            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
302        }
303    }
304
305    pub fn get_column_from_df(&self, column_name: &str) -> anyhow::Result<Column> {
306        let read_guard = self.0.lock_read();
307        let d = read_guard.as_ref();
308        match d {
309            Some(data) => match data.df.column(column_name) {
310                Ok(column) => Ok(column.clone()),
311                Err(e) => Err(anyhow::anyhow!("Column not found: {}", e)),
312            },
313            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
314        }
315    }
316
317    pub fn set_column_in_df(&self, column_name: &str, column: Series) -> anyhow::Result<()> {
318        let mut write_guard = self.0.lock_write();
319        let d = write_guard.as_mut();
320        match d {
321            Some(data) => {
322                data.df.replace(column_name, column)?;
323                Ok(())
324            }
325            None => Err(anyhow::anyhow!("DataFrame is not initialized")),
326        }
327    }
328
329    pub fn subset_inplace(&self, s: &SelectInfoElem) -> anyhow::Result<()> {
330        let read_guard = self.0.lock_read();
331        let d = read_guard.as_ref().unwrap();
332        let indices = crate::utils::select_info_elem_to_indices(s, d.index.len())?;
333        let indices_u32: Vec<u32> = indices.iter().map(|&i| i as u32).collect();
334        let idx = IdxCa::new("idx".into(), &indices_u32);
335        let ind = d.index.clone().into_vec();
336        let ind_subset: Vec<String> = indices.iter().map(|&i| ind[i].clone()).collect();
337        let df_subset = d.df.take(&idx)?;
338        drop(read_guard);
339        self.set_both(df_subset, DataFrameIndex::from(ind_subset))
340    }
341
342    pub fn subset(&self, s: &SelectInfoElem) -> anyhow::Result<Self> {
343        let read_guard = self.0.lock_read();
344        let d = read_guard.as_ref().unwrap();
345        let indices = crate::utils::select_info_elem_to_indices(s, d.index.len())?;
346        let indices_u32: Vec<u32> = indices.iter().map(|&i| i as u32).collect();
347        let idx = IdxCa::new("idx".into(), &indices_u32);
348        let ind = d.index.clone().into_vec();
349        let ind_subset: Vec<String> = indices.iter().map(|&i| ind[i].clone()).collect();
350        let df_subset = d.df.take(&idx)?;
351        Ok(Self::new(df_subset, DataFrameIndex::from(ind_subset)))
352    }
353}
354
355impl DeepClone for IMDataFrameElement {
356    fn deep_clone(&self) -> Self {
357        IMDataFrameElement(self.0.deep_clone())
358    }
359}
360
361pub struct IMAxisArrays(pub RwSlot<InnerIMAxisArray>);
362
363impl Clone for IMAxisArrays {
364    fn clone(&self) -> Self {
365        IMAxisArrays(self.0.clone())
366    }
367}
368
369impl DeepClone for IMAxisArrays {
370    fn deep_clone(&self) -> Self {
371        IMAxisArrays(self.0.deep_clone())
372    }
373}
374
375pub struct InnerIMAxisArray {
376    pub axis: Axis,
377    pub(crate) dim1: Dim,
378    pub(crate) dim2: Option<Dim>,
379    data: HashMap<String, IMArrayElement>,
380}
381
382impl DeepClone for InnerIMAxisArray {
383    fn deep_clone(&self) -> Self {
384        InnerIMAxisArray {
385            axis: self.axis,
386            dim1: self.dim1.clone(),
387            dim2: self.dim2.clone(),
388            data: self
389                .data
390                .iter()
391                .map(|(k, v)| (k.clone(), v.deep_clone()))
392                .collect(),
393        }
394    }
395}
396
397impl Clone for InnerIMAxisArray {
398    fn clone(&self) -> Self {
399        InnerIMAxisArray {
400            axis: self.axis,
401            dim1: self.dim1.clone(),
402            dim2: self.dim2.clone(),
403            data: self.data.clone(),
404        }
405    }
406}
407
408impl fmt::Display for IMAxisArrays {
409    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
410        let read_guard = self.0.read_inner();
411
412        writeln!(f, "IMAxisArrays {{")?;
413        writeln!(f, "    Axis: {:?}", read_guard.axis)?;
414        writeln!(f, "    Dim1: {}", read_guard.dim1)?;
415        if let Some(dim2) = &read_guard.dim2 {
416            writeln!(f, "    Dim2: {}", dim2)?;
417        }
418        writeln!(f, "    Arrays: {{")?;
419        for (key, value) in &read_guard.data {
420            let shape = value.get_shape().map_err(|_| fmt::Error)?;
421            writeln!(f, "        {}: {:?}", key, shape)?;
422        }
423        writeln!(f, "    }}")?;
424        write!(f, "}}")
425    }
426}
427
428impl IMAxisArrays {
429    // Create a new IMAxisArrays
430    pub fn new(axis: Axis, dim1: Dim, dim2: Option<Dim>) -> Self {
431        let inner = InnerIMAxisArray {
432            axis,
433            dim1,
434            dim2,
435            data: HashMap::new(),
436        };
437        IMAxisArrays(RwSlot::new(inner))
438    }
439
440    pub fn new_from(
441        axis: Axis,
442        dim1: Dim,
443        dim2: Option<Dim>,
444        data: HashMap<String, IMArrayElement>,
445    ) -> Self {
446        let inner = InnerIMAxisArray {
447            axis,
448            dim1,
449            dim2,
450            data,
451        };
452        IMAxisArrays(RwSlot::new(inner))
453    }
454
455    pub fn add_array(&self, key: String, element: IMArrayElement) -> anyhow::Result<()> {
456        let mut write_guard = self.0.write_inner();
457        let imarray = write_guard.deref_mut();
458        // Check if the key already exists
459        if imarray.data.contains_key(&key) {
460            return Err(anyhow::anyhow!("Key already exists"));
461        }
462
463        // Get the shape of the input element
464        let shape = element.get_shape()?;
465        let dim1 = imarray.dim1.get();
466        let dim2 = imarray.dim2.clone().unwrap_or(Dim::new(0)).get();
467
468        // Perform dimensionality checks based on the axis type
469        match imarray.axis {
470            Axis::Row => {
471                if shape[0] != dim1 {
472                    return Err(anyhow::anyhow!(
473                        "Data shape {:?} does not match expected row dimension {}",
474                        shape,
475                        dim1
476                    ));
477                }
478            }
479            Axis::RowColumn => {
480                if shape[0] != dim1 || shape[1] != dim2 {
481                    return Err(anyhow::anyhow!(
482                        "Data shape {:?} does not match expected dimensions ({}, {})",
483                        shape,
484                        dim1,
485                        dim2
486                    ));
487                }
488            }
489            Axis::Pairwise => {
490                if shape[0] != dim1 || shape[1] != dim1 {
491                    return Err(anyhow::anyhow!(
492                        "Data shape {:?} does not match expected pairwise dimensions ({}, {})",
493                        shape,
494                        dim1,
495                        dim1
496                    ));
497                }
498            }
499        }
500
501        // If all checks pass, insert the element
502        imarray.data.insert(key, element);
503        Ok(())
504    }
505
506    // Get an array element (returns a deep clone to avoid holding the read lock)
507    pub fn get_array(&self, key: &str) -> anyhow::Result<IMArrayElement> {
508        let read_guard = self.0.read_inner();
509        read_guard
510            .data
511            .get(key)
512            .map(|element| element.deep_clone())
513            .ok_or_else(|| anyhow::anyhow!("Key not found"))
514    }
515
516    // New method: Get an array element (returns a shallow clone)
517    pub fn get_array_shallow(&self, key: &str) -> anyhow::Result<IMArrayElement> {
518        let read_guard = self.0.read_inner();
519        read_guard
520            .data
521            .get(key)
522            .cloned() // This performs a shallow clone
523            .ok_or_else(|| anyhow::anyhow!("Key not found"))
524    }
525
526    // Remove an array element
527    pub fn remove_array(&self, key: &str) -> anyhow::Result<IMArrayElement> {
528        let mut write_guard = self.0.write_inner();
529        write_guard
530            .data
531            .remove(key)
532            .ok_or_else(|| anyhow::anyhow!("Key not found"))
533    }
534
535    // Get the number of arrays
536    pub fn len(&self) -> usize {
537        let read_guard = self.0.read_inner();
538        read_guard.data.len()
539    }
540
541    // Check if there are any arrays
542    pub fn is_empty(&self) -> bool {
543        let read_guard = self.0.read_inner();
544        read_guard.data.is_empty()
545    }
546
547    // Get all keys
548    pub fn keys(&self) -> Vec<String> {
549        let read_guard = self.0.read_inner();
550        read_guard.data.keys().cloned().collect()
551    }
552
553    // Get the axis
554    pub fn axis(&self) -> Axis {
555        let read_guard = self.0.read_inner();
556        read_guard.axis
557    }
558
559    // Get dimensions
560    pub fn dimensions(&self) -> (Dim, Option<Dim>) {
561        let read_guard = self.0.read_inner();
562        (read_guard.dim1.clone(), read_guard.dim2.clone())
563    }
564
565    // Update an existing array element
566    pub fn update_array(&self, key: &str, new_element: IMArrayElement) -> anyhow::Result<()> {
567        let mut write_guard = self.0.write_inner();
568        if let Some(element) = write_guard.data.get_mut(key) {
569            *element = new_element;
570            Ok(())
571        } else {
572            Err(anyhow::anyhow!("Key not found"))
573        }
574    }
575
576    pub fn subset_inplace(&self, s: &[&SelectInfoElem]) -> anyhow::Result<()> {
577        let mut write_guard = self.0.write_inner();
578        let imarray = write_guard.deref_mut();
579        let dim1_indices = crate::utils::select_info_elem_to_indices(s[0], imarray.dim1.get())?;
580        imarray.dim1 = Dim::new(dim1_indices.len());
581
582        if let Some(dim2) = &mut imarray.dim2 {
583            if s.len() < 2 {
584                return Err(anyhow::anyhow!(
585                    "Subset operation requires two selection elements"
586                ));
587            }
588            let dim2_indices = crate::utils::select_info_elem_to_indices(s[1], dim2.get())?;
589            *dim2 = Dim::new(dim2_indices.len());
590        }
591
592        for element in imarray.data.values_mut() {
593            element.subset_inplace(s)?;
594        }
595
596        Ok(())
597    }
598
599    pub fn subset(&self, s: &[&SelectInfoElem]) -> anyhow::Result<Self> {
600        let read_guard = self.0.read_inner();
601        let imarray = read_guard.deref();
602        let dim1_indices = crate::utils::select_info_elem_to_indices(s[0], imarray.dim1.get())?;
603        let new_dim1 = Dim::new(dim1_indices.len());
604        let mut new_dim2: Option<Dim> = None;
605        if imarray.dim2.is_some() {
606            if s.len() < 2 {
607                return Err(anyhow::anyhow!(
608                    "Subset operation requires two selection elements"
609                ));
610            }
611            let dim2_indices = crate::utils::select_info_elem_to_indices(
612                s[1],
613                imarray.dim2.clone().unwrap().get(),
614            )?;
615            new_dim2 = Some(Dim::new(dim2_indices.len()));
616        }
617        let mut new_data = HashMap::new();
618        for (key, element) in &imarray.data {
619            new_data.insert(key.clone(), element.subset(s)?);
620        }
621
622        Ok(IMAxisArrays::new_from(
623            imarray.axis,
624            new_dim1,
625            new_dim2,
626            new_data,
627        ))
628    }
629
630    // Perform an operation on all arrays
631    pub fn map<F>(&self, f: F) -> anyhow::Result<()>
632    where
633        F: Fn(&mut IMArrayElement) -> anyhow::Result<()>,
634    {
635        let mut write_guard = self.0.write_inner();
636        for element in write_guard.data.values_mut() {
637            f(element)?;
638        }
639        Ok(())
640    }
641}
642
643pub struct IMElement(pub RwSlot<Data>);
644
645impl DeepClone for Data {
646    fn deep_clone(&self) -> Self {
647        self.clone()
648    }
649}
650
651impl DeepClone for IMElement {
652    fn deep_clone(&self) -> Self {
653        IMElement(self.0.deep_clone())
654    }
655}
656
657impl Clone for IMElement {
658    fn clone(&self) -> Self {
659        IMElement(self.0.clone())
660    }
661}
662
663impl IMElement {
664    pub fn new(data: Data) -> Self {
665        IMElement(RwSlot::new(data))
666    }
667
668    pub fn get_data(&self) -> anyhow::Result<Data> {
669        Ok(self.0.read_inner().clone())
670    }
671
672    pub fn set_data(&self, data: Data) -> anyhow::Result<()> {
673        let mut write_guard = self.0.lock_write();
674        let d = write_guard.deref_mut();
675        *d = Some(data);
676        Ok(())
677    }
678}
679
680pub struct IMElementCollection(pub RwSlot<HashMap<String, IMElement>>);
681
682impl DeepClone for IMElementCollection {
683    fn deep_clone(&self) -> Self {
684        let temp_data = self.0.read_inner();
685        let data = temp_data.deref();
686        let mut new_data = HashMap::new();
687        for (key, value) in data.iter() {
688            new_data.insert(key.clone(), value.deep_clone());
689        }
690        IMElementCollection(RwSlot::new(new_data))
691    }
692}
693
694impl Clone for IMElementCollection {
695    fn clone(&self) -> Self {
696        IMElementCollection(self.0.clone())
697    }
698}
699
700impl IMElementCollection {
701    pub fn new_empty() -> Self {
702        IMElementCollection(RwSlot::new(HashMap::new()))
703    }
704
705    pub fn add_data(&self, key: String, element: IMElement) -> anyhow::Result<()> {
706        let mut write_guard = self.0.write_inner();
707        let collection = write_guard.deref_mut();
708        if collection.contains_key(&key) {
709            return Err(anyhow::anyhow!("Key already exists"));
710        }
711        collection.insert(key, element);
712        Ok(())
713    }
714
715    pub fn remove_data(&self, key: &str) -> anyhow::Result<IMElement> {
716        let mut write_guard = self.0.write_inner();
717        write_guard
718            .remove(key)
719            .ok_or_else(|| anyhow::anyhow!("Key not found"))
720    }
721
722    pub fn get_data(&self, key: &str) -> anyhow::Result<IMElement> {
723        let read_guard = self.0.read_inner();
724        read_guard
725            .get(key)
726            .cloned()
727            .ok_or_else(|| anyhow::anyhow!("Key not found"))
728    }
729
730    pub fn get_data_deep(&self, key: &str) -> anyhow::Result<IMElement> {
731        let read_guard = self.0.read_inner();
732        read_guard
733            .get(key)
734            .map(|element| element.deep_clone())
735            .ok_or_else(|| anyhow::anyhow!("Key not found"))
736    }
737
738    pub fn keys(&self) -> anyhow::Result<Vec<String>> {
739        let read_guard = self.0.read_inner();
740        let data = read_guard.deref();
741        let keys = data.keys().map(|k| k.clone()).collect();
742        Ok(keys)
743    }
744}
anndata_memory/ad/helpers.rs

anndata_memory/ad/
helpers.rs
