lance_graph/
config.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Graph configuration for mapping Lance datasets to property graphs
5
6use crate::error::{GraphError, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Configuration for mapping Lance datasets to property graphs
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct GraphConfig {
13    /// Mapping of node labels to their field configurations
14    pub node_mappings: HashMap<String, NodeMapping>,
15    /// Mapping of relationship types to their field configurations  
16    pub relationship_mappings: HashMap<String, RelationshipMapping>,
17    /// Default node ID field if not specified in mappings
18    pub default_node_id_field: String,
19    /// Default relationship type field if not specified in mappings
20    pub default_relationship_type_field: String,
21}
22
23/// Configuration for mapping node labels to dataset fields
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct NodeMapping {
26    /// The node label (e.g., "Person", "Product")
27    pub label: String,
28    /// Field name that serves as the node identifier
29    pub id_field: String,
30    /// Optional fields that define node properties
31    pub property_fields: Vec<String>,
32    /// Optional filter conditions for this node type
33    pub filter_conditions: Option<String>,
34}
35
36/// Configuration for mapping relationship types to dataset fields
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct RelationshipMapping {
39    /// The relationship type (e.g., "KNOWS", "PURCHASED")
40    pub relationship_type: String,
41    /// Field containing the source node ID
42    pub source_id_field: String,
43    /// Field containing the target node ID
44    pub target_id_field: String,
45    /// Optional field containing the relationship type
46    pub type_field: Option<String>,
47    /// Optional fields that define relationship properties
48    pub property_fields: Vec<String>,
49    /// Optional filter conditions for this relationship type
50    pub filter_conditions: Option<String>,
51}
52
53impl Default for GraphConfig {
54    fn default() -> Self {
55        Self {
56            node_mappings: HashMap::new(),
57            relationship_mappings: HashMap::new(),
58            default_node_id_field: "id".to_string(),
59            default_relationship_type_field: "type".to_string(),
60        }
61    }
62}
63
64impl GraphConfig {
65    /// Create a new builder for GraphConfig
66    pub fn builder() -> GraphConfigBuilder {
67        GraphConfigBuilder::new()
68    }
69
70    /// Get node mapping for a given label
71    pub fn get_node_mapping(&self, label: &str) -> Option<&NodeMapping> {
72        self.node_mappings.get(label)
73    }
74
75    /// Get relationship mapping for a given type
76    pub fn get_relationship_mapping(&self, rel_type: &str) -> Option<&RelationshipMapping> {
77        self.relationship_mappings.get(rel_type)
78    }
79
80    /// Validate the configuration
81    pub fn validate(&self) -> Result<()> {
82        // Check for conflicting field names
83        for (label, mapping) in &self.node_mappings {
84            if mapping.id_field.is_empty() {
85                return Err(GraphError::ConfigError {
86                    message: format!("Node mapping for '{}' has empty id_field", label),
87                    location: snafu::Location::new(file!(), line!(), column!()),
88                });
89            }
90        }
91
92        for (rel_type, mapping) in &self.relationship_mappings {
93            if mapping.source_id_field.is_empty() || mapping.target_id_field.is_empty() {
94                return Err(GraphError::ConfigError {
95                    message: format!(
96                        "Relationship mapping for '{}' has empty source or target id field",
97                        rel_type
98                    ),
99                    location: snafu::Location::new(file!(), line!(), column!()),
100                });
101            }
102        }
103
104        Ok(())
105    }
106}
107
108/// Builder for GraphConfig
109#[derive(Debug, Default, Clone)]
110pub struct GraphConfigBuilder {
111    node_mappings: HashMap<String, NodeMapping>,
112    relationship_mappings: HashMap<String, RelationshipMapping>,
113    default_node_id_field: Option<String>,
114    default_relationship_type_field: Option<String>,
115}
116
117impl GraphConfigBuilder {
118    /// Create a new builder
119    pub fn new() -> Self {
120        Self::default()
121    }
122
123    /// Add a node label mapping
124    pub fn with_node_label<S: Into<String>>(mut self, label: S, id_field: S) -> Self {
125        let label_str = label.into();
126        self.node_mappings.insert(
127            label_str.clone(),
128            NodeMapping {
129                label: label_str,
130                id_field: id_field.into(),
131                property_fields: Vec::new(),
132                filter_conditions: None,
133            },
134        );
135        self
136    }
137
138    /// Add a node mapping with additional configuration
139    pub fn with_node_mapping(mut self, mapping: NodeMapping) -> Self {
140        self.node_mappings.insert(mapping.label.clone(), mapping);
141        self
142    }
143
144    /// Add a relationship type mapping
145    pub fn with_relationship<S: Into<String>>(
146        mut self,
147        rel_type: S,
148        source_field: S,
149        target_field: S,
150    ) -> Self {
151        let type_str = rel_type.into();
152        self.relationship_mappings.insert(
153            type_str.clone(),
154            RelationshipMapping {
155                relationship_type: type_str,
156                source_id_field: source_field.into(),
157                target_id_field: target_field.into(),
158                type_field: None,
159                property_fields: Vec::new(),
160                filter_conditions: None,
161            },
162        );
163        self
164    }
165
166    /// Add a relationship mapping with additional configuration
167    pub fn with_relationship_mapping(mut self, mapping: RelationshipMapping) -> Self {
168        self.relationship_mappings
169            .insert(mapping.relationship_type.clone(), mapping);
170        self
171    }
172
173    /// Set the default node ID field
174    pub fn with_default_node_id_field<S: Into<String>>(mut self, field: S) -> Self {
175        self.default_node_id_field = Some(field.into());
176        self
177    }
178
179    /// Set the default relationship type field
180    pub fn with_default_relationship_type_field<S: Into<String>>(mut self, field: S) -> Self {
181        self.default_relationship_type_field = Some(field.into());
182        self
183    }
184
185    /// Build the GraphConfig
186    pub fn build(self) -> Result<GraphConfig> {
187        let config = GraphConfig {
188            node_mappings: self.node_mappings,
189            relationship_mappings: self.relationship_mappings,
190            default_node_id_field: self
191                .default_node_id_field
192                .unwrap_or_else(|| "id".to_string()),
193            default_relationship_type_field: self
194                .default_relationship_type_field
195                .unwrap_or_else(|| "type".to_string()),
196        };
197
198        config.validate()?;
199        Ok(config)
200    }
201}
202
203impl NodeMapping {
204    /// Create a new node mapping
205    pub fn new<S: Into<String>>(label: S, id_field: S) -> Self {
206        Self {
207            label: label.into(),
208            id_field: id_field.into(),
209            property_fields: Vec::new(),
210            filter_conditions: None,
211        }
212    }
213
214    /// Add property fields to the mapping
215    pub fn with_properties(mut self, fields: Vec<String>) -> Self {
216        self.property_fields = fields;
217        self
218    }
219
220    /// Add filter conditions for this node type
221    pub fn with_filter<S: Into<String>>(mut self, filter: S) -> Self {
222        self.filter_conditions = Some(filter.into());
223        self
224    }
225}
226
227impl RelationshipMapping {
228    /// Create a new relationship mapping
229    pub fn new<S: Into<String>>(rel_type: S, source_field: S, target_field: S) -> Self {
230        Self {
231            relationship_type: rel_type.into(),
232            source_id_field: source_field.into(),
233            target_id_field: target_field.into(),
234            type_field: None,
235            property_fields: Vec::new(),
236            filter_conditions: None,
237        }
238    }
239
240    /// Set the type field for this relationship
241    pub fn with_type_field<S: Into<String>>(mut self, type_field: S) -> Self {
242        self.type_field = Some(type_field.into());
243        self
244    }
245
246    /// Add property fields to the mapping
247    pub fn with_properties(mut self, fields: Vec<String>) -> Self {
248        self.property_fields = fields;
249        self
250    }
251
252    /// Add filter conditions for this relationship type
253    pub fn with_filter<S: Into<String>>(mut self, filter: S) -> Self {
254        self.filter_conditions = Some(filter.into());
255        self
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[test]
264    fn test_graph_config_builder() {
265        let config = GraphConfig::builder()
266            .with_node_label("Person", "person_id")
267            .with_node_label("Company", "company_id")
268            .with_relationship("WORKS_FOR", "person_id", "company_id")
269            .build()
270            .unwrap();
271
272        assert_eq!(config.node_mappings.len(), 2);
273        assert_eq!(config.relationship_mappings.len(), 1);
274
275        let person_mapping = config.get_node_mapping("Person").unwrap();
276        assert_eq!(person_mapping.id_field, "person_id");
277
278        let works_for_mapping = config.get_relationship_mapping("WORKS_FOR").unwrap();
279        assert_eq!(works_for_mapping.source_id_field, "person_id");
280        assert_eq!(works_for_mapping.target_id_field, "company_id");
281    }
282
283    #[test]
284    fn test_validation_empty_id_field() {
285        let mut config = GraphConfig::default();
286        config.node_mappings.insert(
287            "Person".to_string(),
288            NodeMapping {
289                label: "Person".to_string(),
290                id_field: "".to_string(),
291                property_fields: Vec::new(),
292                filter_conditions: None,
293            },
294        );
295
296        assert!(config.validate().is_err());
297    }
298
299    #[test]
300    fn test_node_mapping_with_properties() {
301        let mapping = NodeMapping::new("Person", "id")
302            .with_properties(vec!["name".to_string(), "age".to_string()])
303            .with_filter("age > 18".to_string());
304
305        assert_eq!(mapping.property_fields.len(), 2);
306        assert!(mapping.filter_conditions.is_some());
307    }
308}
lance_graph/config.rs

lance_graph/
config.rs
