graphar 0.1.2

Apache GraphAr format reader/writer
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

use std::path::{Path, PathBuf};

use arrow_array::RecordBatch;

use crate::{
    error::{GraphArError, Result},
    io,
    metadata::{PropertyGroup, VertexInfo},
};

/// Reads vertex chunks from disk for a given vertex type.
pub struct VertexReader {
    info: VertexInfo,
    base_dir: PathBuf,
    vertex_count: u64,
}

impl VertexReader {
    pub fn new(info: VertexInfo, base_dir: impl AsRef<Path>, vertex_count: u64) -> Self {
        Self {
            info,
            base_dir: base_dir.as_ref().to_path_buf(),
            vertex_count,
        }
    }

    pub fn vertex_info(&self) -> &VertexInfo {
        &self.info
    }

    pub fn chunk_count(&self) -> u64 {
        self.info.chunk_count(self.vertex_count)
    }

    /// Read all chunks for a property group, returning concatenated batches.
    pub fn read_property_group(&self, group: &PropertyGroup) -> Result<Vec<RecordBatch>> {
        let mut all = Vec::new();
        for chunk_idx in 0..self.chunk_count() {
            let mut batches = self.read_property_group_chunk(group, chunk_idx)?;
            all.append(&mut batches);
        }
        Ok(all)
    }

    /// Read a single chunk for a property group.
    pub fn read_property_group_chunk(
        &self,
        group: &PropertyGroup,
        chunk_index: u64,
    ) -> Result<Vec<RecordBatch>> {
        let rel = self.info.chunk_path(group, chunk_index);
        let full = self.base_dir.join(&rel);
        if !full.exists() {
            return Err(GraphArError::ChunkOutOfRange(
                chunk_index,
                self.chunk_count(),
            ));
        }
        io::read_chunk(&full, &group.file_type)
    }

    /// Read every property group and merge them into one row-aligned batch
    /// holding all columns of this vertex type.
    pub fn read_merged(&self) -> Result<RecordBatch> {
        let groups = self.info.property_groups.clone();
        self.read_merged_groups(&groups)
    }

    /// Read the given property groups and merge them into one batch.
    pub fn read_merged_groups(&self, groups: &[PropertyGroup]) -> Result<RecordBatch> {
        let per_group = match groups
            .iter()
            .map(|g| self.read_property_group(g))
            .collect::<Result<Vec<_>>>()
        {
            Ok(v) => v,
            Err(e) => {
                // Emit the vertex-load read path's health into the nornir matrix.
                crate::functional_status(
                    "graphar/vertex_reader",
                    "read_merged",
                    false,
                    "property-group chunk read failed",
                );
                return Err(e);
            }
        };
        let merged = crate::merge::merge_property_groups(&per_group);
        crate::functional_status(
            "graphar/vertex_reader",
            "read_merged",
            merged.is_ok(),
            &format!("{} groups over {} chunks", groups.len(), self.chunk_count()),
        );
        merged
    }

    /// Read chunks for the property group containing the named property.
    pub fn read_by_property(&self, property_name: &str) -> Result<Vec<RecordBatch>> {
        let group = self
            .info
            .find_property_group(property_name)
            .ok_or_else(|| GraphArError::PropertyNotFound(property_name.to_string()))?
            .clone();
        self.read_property_group(&group)
    }
}

/// Iterator over vertex chunks for a property group.
pub struct VertexChunkIter<'a> {
    reader: &'a VertexReader,
    group: &'a PropertyGroup,
    current: u64,
    total: u64,
}

impl<'a> VertexChunkIter<'a> {
    pub fn new(reader: &'a VertexReader, group: &'a PropertyGroup) -> Self {
        let total = reader.chunk_count();
        Self {
            reader,
            group,
            current: 0,
            total,
        }
    }
}

impl Iterator for VertexChunkIter<'_> {
    type Item = Result<Vec<RecordBatch>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current >= self.total {
            return None;
        }
        let idx = self.current;
        self.current += 1;
        Some(self.reader.read_property_group_chunk(self.group, idx))
    }
}