Skip to main content

uni_store/storage/
edge.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4use anyhow::Result;
5use arrow_schema::{Field, Schema as ArrowSchema};
6use std::sync::Arc;
7use uni_common::core::schema::Schema;
8
9pub struct EdgeDataset {
10    edge_type: String,
11    /// Lance branch for branched reads. `None` = primary.
12    branch: Option<String>,
13}
14
15impl EdgeDataset {
16    /// `_base_uri` is ignored: physical path resolution belongs to the storage
17    /// backend, and this type now holds only the logical identity of the table.
18    pub fn new(_base_uri: &str, edge_type: &str, _src_label: &str, _dst_label: &str) -> Self {
19        Self {
20            edge_type: edge_type.to_string(),
21            branch: None,
22        }
23    }
24
25    /// Construct an edge dataset that reads from a Lance branch.
26    pub fn new_branched(
27        base_uri: &str,
28        edge_type: &str,
29        src_label: &str,
30        dst_label: &str,
31        branch: impl Into<String>,
32    ) -> Self {
33        let mut ds = Self::new(base_uri, edge_type, src_label, dst_label);
34        ds.branch = Some(branch.into());
35        ds
36    }
37
38    pub fn get_arrow_schema(&self, schema: &Schema) -> Result<Arc<ArrowSchema>> {
39        let mut fields = vec![
40            Field::new("eid", arrow_schema::DataType::UInt64, false),
41            Field::new("src_vid", arrow_schema::DataType::UInt64, false),
42            Field::new("dst_vid", arrow_schema::DataType::UInt64, false),
43            Field::new("_deleted", arrow_schema::DataType::Boolean, false),
44            Field::new("_version", arrow_schema::DataType::UInt64, false),
45        ];
46
47        if let Some(type_props) = schema.properties.get(&self.edge_type) {
48            let mut sorted_props: Vec<_> = type_props.iter().collect();
49            sorted_props.sort_by_key(|(name, _)| *name);
50
51            for (name, meta) in sorted_props {
52                fields.push(Field::new(name, meta.r#type.to_arrow(), meta.nullable));
53            }
54        }
55
56        Ok(Arc::new(ArrowSchema::new(fields)))
57    }
58}