raphtory-storage 0.17.0

Raphtory storage interface
Documentation
use crate::disk::{
    graph_impl::DiskEdge, storage_interface::edges_ref::DiskEdgesRef, DiskGraphStorage,
};
use itertools::Itertools;
use raphtory_api::{
    core::entities::{edges::edge_ref::EdgeRef, LayerIds, LayerVariants, EID},
    iter::IntoDynBoxed,
};
use raphtory_core::utils::iter::GenLockedIter;
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use std::{iter, sync::Arc};

#[derive(Clone, Debug)]
pub struct DiskEdges {
    graph: Arc<DiskGraphStorage>,
}

impl DiskEdges {
    pub(crate) fn new(graph: &DiskGraphStorage) -> Self {
        Self {
            graph: Arc::new(graph.clone()),
        }
    }

    pub fn as_ref(&self) -> DiskEdgesRef<'_> {
        DiskEdgesRef {
            graph: &self.graph.inner,
        }
    }

    pub fn into_iter_refs(self, layer_ids: LayerIds) -> impl Iterator<Item = EdgeRef> {
        match layer_ids {
            LayerIds::None => LayerVariants::None(iter::empty()),
            LayerIds::All => LayerVariants::All(GenLockedIter::from(self.graph, |graph| {
                graph
                    .inner
                    .all_edge_ids()
                    .map(|(eid, src, dst)| EdgeRef::new_outgoing(eid, src, dst))
                    .into_dyn_boxed()
            })),
            LayerIds::One(layer_id) => {
                LayerVariants::One(GenLockedIter::from(self.graph, move |graph| {
                    graph
                        .inner
                        .layer_edge_ids(layer_id)
                        .map(|(eid, src, dst)| EdgeRef::new_outgoing(eid, src, dst))
                        .into_dyn_boxed()
                }))
            }
            LayerIds::Multiple(ids) => LayerVariants::Multiple(
                ids.into_iter()
                    .map(move |layer_id| {
                        GenLockedIter::from(self.graph.clone(), move |graph| {
                            graph.inner.layer_edge_ids(layer_id).into_dyn_boxed()
                        })
                    })
                    .kmerge_by(|(eid1, _, _), (eid2, _, _)| eid1 < eid2)
                    .dedup()
                    .map(move |(eid, src, dst)| EdgeRef::new_outgoing(eid, src, dst)),
            ),
        }
    }

    pub fn into_par_iter_refs(self, layer_ids: LayerIds) -> impl ParallelIterator<Item = EID> {
        match layer_ids {
            LayerIds::None => LayerVariants::None(rayon::iter::empty()),
            LayerIds::One(layer_id) => {
                LayerVariants::One(self.graph.inner.all_edge_ids_par(layer_id))
            }
            LayerIds::All => {
                LayerVariants::All((0..self.graph.inner.num_edges()).into_par_iter().map(EID))
            }
            LayerIds::Multiple(ids) => LayerVariants::Multiple(
                (0..self.graph.inner.num_edges())
                    .into_par_iter()
                    .map(EID)
                    .filter(move |e| {
                        ids.into_iter()
                            .any(|layer_id| self.graph.inner.edge(*e).has_layer_inner(layer_id))
                    }),
            ),
        }
    }

    pub fn get(&self, eid: EID) -> DiskEdge<'_> {
        self.graph.inner.edge(eid)
    }
}