lance_csr/
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///
12/// Node labels and relationship types are case-insensitive: the mappings are keyed by their
13/// lowercase form. Build a config with [`GraphConfigBuilder`], which normalizes keys, and read
14/// it through the lookup and iterator accessors. Deserialization goes through the builder too,
15/// so a deserialized config is validated like a built one.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17#[serde(try_from = "GraphConfigBuilder")]
18pub struct GraphConfig {
19    /// Node labels (lowercase) to their field configurations
20    node_mappings: HashMap<String, NodeMapping>,
21
22    /// Relationship types (lowercase) to their field configurations
23    relationship_mappings: HashMap<String, RelationshipMapping>,
24
25    /// Default node ID field if not specified in mappings
26    pub default_node_id_field: String,
27
28    /// Default relationship type field if not specified in mappings
29    pub default_relationship_type_field: String,
30}
31
32/// Configuration for mapping node labels to dataset fields
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct NodeMapping {
35    /// The node label (e.g., "Person", "Product")
36    pub label: String,
37    /// Field name that serves as the node identifier
38    pub id_field: String,
39    /// Optional fields that define node properties
40    pub property_fields: Vec<String>,
41    /// Optional filter conditions for this node type
42    pub filter_conditions: Option<String>,
43}
44
45/// Configuration for mapping relationship types to dataset fields
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct RelationshipMapping {
48    /// The relationship type (e.g., "KNOWS", "PURCHASED")
49    pub relationship_type: String,
50    /// Field containing the source node ID
51    pub source_id_field: String,
52    /// Field containing the target node ID
53    pub target_id_field: String,
54    /// Optional field containing the relationship type
55    pub type_field: Option<String>,
56    /// Optional fields that define relationship properties
57    pub property_fields: Vec<String>,
58    /// Optional filter conditions for this relationship type
59    pub filter_conditions: Option<String>,
60    /// Node labels at the source and target end, like `FROM Post TO Person` in a Kuzu
61    /// REL table. Only needed to plan relationship-type alternation (`[:A|B]`) or
62    /// multi-label nodes, where they decide which label goes with which type.
63    #[serde(default)]
64    pub source_label: Option<String>,
65    #[serde(default)]
66    pub target_label: Option<String>,
67}
68
69impl Default for GraphConfig {
70    fn default() -> Self {
71        Self {
72            node_mappings: HashMap::new(),
73            relationship_mappings: HashMap::new(),
74            default_node_id_field: "id".to_string(),
75            default_relationship_type_field: "type".to_string(),
76        }
77    }
78}
79
80impl GraphConfig {
81    /// Create a new builder for GraphConfig
82    pub fn builder() -> GraphConfigBuilder {
83        GraphConfigBuilder::new()
84    }
85
86    /// Get node mapping for a given label (case-insensitive)
87    ///
88    /// Looks up the node mapping using case-insensitive comparison.
89    /// For example, "Person", "PERSON", and "person" all refer to the same label.
90    pub fn get_node_mapping(&self, label: &str) -> Option<&NodeMapping> {
91        self.node_mappings.get(&label.to_lowercase())
92    }
93
94    /// Get relationship mapping for a given type (case-insensitive)
95    ///
96    /// Looks up the relationship mapping using case-insensitive comparison.
97    /// For example, "FOLLOWS", "follows", and "Follows" all refer to the same type.
98    pub fn get_relationship_mapping(&self, rel_type: &str) -> Option<&RelationshipMapping> {
99        self.relationship_mappings.get(&rel_type.to_lowercase())
100    }
101
102    /// Iterate over normalized node labels and their mappings without allowing mutation.
103    pub fn iter_node_mappings(&self) -> impl ExactSizeIterator<Item = (&String, &NodeMapping)> {
104        self.node_mappings.iter()
105    }
106
107    /// Iterate over normalized relationship types and their mappings without allowing mutation.
108    pub fn iter_relationship_mappings(
109        &self,
110    ) -> impl ExactSizeIterator<Item = (&String, &RelationshipMapping)> {
111        self.relationship_mappings.iter()
112    }
113
114    /// Validate the configuration
115    ///
116    /// Checks for:
117    /// - Empty ID fields
118    /// - Non-normalized keys (must be lowercase)
119    /// - Case-insensitive duplicates
120    pub fn validate(&self) -> Result<()> {
121        // Validate node mappings
122        for (label, mapping) in &self.node_mappings {
123            // Check that keys are normalized (lowercase)
124            if label != &label.to_lowercase() {
125                return Err(GraphError::ConfigError {
126                    message: format!(
127                        "Node mapping key '{}' is not normalized. \
128                         Keys must be lowercase. Use GraphConfigBuilder to ensure proper normalization.",
129                        label
130                    ),
131                    location: snafu::Location::new(file!(), line!(), column!()),
132                });
133            }
134
135            if mapping.id_field.is_empty() {
136                return Err(GraphError::ConfigError {
137                    message: format!("Node mapping for '{}' has empty id_field", label),
138                    location: snafu::Location::new(file!(), line!(), column!()),
139                });
140            }
141        }
142
143        // Validate relationship mappings
144        for (rel_type, mapping) in &self.relationship_mappings {
145            // Check that keys are normalized (lowercase)
146            if rel_type != &rel_type.to_lowercase() {
147                return Err(GraphError::ConfigError {
148                    message: format!(
149                        "Relationship mapping key '{}' is not normalized. \
150                         Keys must be lowercase. Use GraphConfigBuilder to ensure proper normalization.",
151                        rel_type
152                    ),
153                    location: snafu::Location::new(file!(), line!(), column!()),
154                });
155            }
156
157            if mapping.source_id_field.is_empty() || mapping.target_id_field.is_empty() {
158                return Err(GraphError::ConfigError {
159                    message: format!(
160                        "Relationship mapping for '{}' has empty source or target id field",
161                        rel_type
162                    ),
163                    location: snafu::Location::new(file!(), line!(), column!()),
164                });
165            }
166
167            for label in [&mapping.source_label, &mapping.target_label]
168                .into_iter()
169                .flatten()
170            {
171                if self.get_node_mapping(label).is_none() {
172                    return Err(GraphError::ConfigError {
173                        message: format!(
174                            "Relationship '{}' names endpoint label '{}', which has no node mapping",
175                            rel_type, label
176                        ),
177                        location: snafu::Location::new(file!(), line!(), column!()),
178                    });
179                }
180            }
181        }
182
183        Ok(())
184    }
185}
186
187/// Builder for GraphConfig
188#[derive(Debug, Default, Clone, Deserialize)]
189#[serde(default)]
190pub struct GraphConfigBuilder {
191    node_mappings: HashMap<String, NodeMapping>,
192    relationship_mappings: HashMap<String, RelationshipMapping>,
193    default_node_id_field: Option<String>,
194    default_relationship_type_field: Option<String>,
195}
196
197impl GraphConfigBuilder {
198    /// Create a new builder
199    pub fn new() -> Self {
200        Self::default()
201    }
202
203    /// A node table: the nodes labelled `label`, one per row of the dataset of the same
204    /// name, identified by its `primary_key` column.
205    ///
206    /// Labels are case-insensitive: adding "Person" and then "person" keeps the second.
207    pub fn node_table<S: Into<String>>(mut self, label: S, primary_key: S) -> Self {
208        let label_str = label.into();
209        let normalized_key = label_str.to_lowercase();
210        self.node_mappings.insert(
211            normalized_key,
212            NodeMapping {
213                label: label_str, // Keep original case for display
214                id_field: primary_key.into(),
215                property_fields: Vec::new(),
216                filter_conditions: None,
217            },
218        );
219        self
220    }
221
222    /// A relationship table: edges of type `rel_type` from `from_table` nodes to
223    /// `to_table` nodes, one per row of the dataset of the same name, with the source
224    /// node's key in column `src` and the target node's key in column `dst`. Both node
225    /// tables must be defined by the time [`Self::build`] runs.
226    pub fn rel_table<S: Into<String>>(
227        mut self,
228        rel_type: S,
229        from_table: S,
230        to_table: S,
231        src: S,
232        dst: S,
233    ) -> Self {
234        let type_str = rel_type.into();
235        let normalized_key = type_str.to_lowercase();
236        self.relationship_mappings.insert(
237            normalized_key,
238            RelationshipMapping {
239                relationship_type: type_str, // Keep original case for display
240                source_id_field: src.into(),
241                target_id_field: dst.into(),
242                type_field: None,
243                property_fields: Vec::new(),
244                filter_conditions: None,
245                source_label: Some(from_table.into()),
246                target_label: Some(to_table.into()),
247            },
248        );
249        self
250    }
251
252    /// Add a node label mapping.
253    #[deprecated(since = "0.2.0", note = "use `node_table`")]
254    pub fn with_node_label<S: Into<String>>(self, label: S, id_field: S) -> Self {
255        self.node_table(label, id_field)
256    }
257
258    /// Add a node mapping with additional configuration
259    pub fn with_node_mapping(mut self, mapping: NodeMapping) -> Self {
260        let normalized_key = mapping.label.to_lowercase();
261        self.node_mappings.insert(normalized_key, mapping);
262        self
263    }
264
265    /// Add a relationship type mapping, without endpoint node tables.
266    #[deprecated(
267        since = "0.2.0",
268        note = "use `rel_table`, which also takes the node tables at both ends"
269    )]
270    pub fn with_relationship<S: Into<String>>(
271        mut self,
272        rel_type: S,
273        source_field: S,
274        target_field: S,
275    ) -> Self {
276        let type_str = rel_type.into();
277        let normalized_key = type_str.to_lowercase();
278        self.relationship_mappings.insert(
279            normalized_key,
280            RelationshipMapping {
281                relationship_type: type_str, // Keep original case for display
282                source_id_field: source_field.into(),
283                target_id_field: target_field.into(),
284                type_field: None,
285                property_fields: Vec::new(),
286                filter_conditions: None,
287                source_label: None,
288                target_label: None,
289            },
290        );
291        self
292    }
293
294    /// Record the node labels a configured relationship type connects (see
295    /// [`RelationshipMapping::source_label`]). No-op for an unknown type.
296    #[deprecated(
297        since = "0.2.0",
298        note = "use `rel_table`, which takes the endpoint node tables"
299    )]
300    pub fn with_relationship_endpoints<S: Into<String>>(
301        mut self,
302        rel_type: &str,
303        source_label: S,
304        target_label: S,
305    ) -> Self {
306        if let Some(mapping) = self.relationship_mappings.get_mut(&rel_type.to_lowercase()) {
307            mapping.source_label = Some(source_label.into());
308            mapping.target_label = Some(target_label.into());
309        }
310        self
311    }
312
313    /// Add a relationship mapping with additional configuration
314    pub fn with_relationship_mapping(mut self, mapping: RelationshipMapping) -> Self {
315        let normalized_key = mapping.relationship_type.to_lowercase();
316        self.relationship_mappings.insert(normalized_key, mapping);
317        self
318    }
319
320    /// Set the default node ID field
321    pub fn with_default_node_id_field<S: Into<String>>(mut self, field: S) -> Self {
322        self.default_node_id_field = Some(field.into());
323        self
324    }
325
326    /// Set the default relationship type field
327    pub fn with_default_relationship_type_field<S: Into<String>>(mut self, field: S) -> Self {
328        self.default_relationship_type_field = Some(field.into());
329        self
330    }
331
332    /// Build the GraphConfig
333    pub fn build(self) -> Result<GraphConfig> {
334        let config = GraphConfig {
335            node_mappings: self.node_mappings,
336            relationship_mappings: self.relationship_mappings,
337            default_node_id_field: self
338                .default_node_id_field
339                .unwrap_or_else(|| "id".to_string()),
340            default_relationship_type_field: self
341                .default_relationship_type_field
342                .unwrap_or_else(|| "type".to_string()),
343        };
344
345        config.validate()?;
346        Ok(config)
347    }
348}
349
350impl TryFrom<GraphConfigBuilder> for GraphConfig {
351    type Error = GraphError;
352
353    fn try_from(builder: GraphConfigBuilder) -> Result<Self> {
354        builder.build()
355    }
356}
357
358impl NodeMapping {
359    /// Create a new node mapping
360    pub fn new<S: Into<String>>(label: S, id_field: S) -> Self {
361        Self {
362            label: label.into(),
363            id_field: id_field.into(),
364            property_fields: Vec::new(),
365            filter_conditions: None,
366        }
367    }
368
369    /// Add property fields to the mapping
370    pub fn with_properties(mut self, fields: Vec<String>) -> Self {
371        self.property_fields = fields;
372        self
373    }
374
375    /// Add filter conditions for this node type
376    pub fn with_filter<S: Into<String>>(mut self, filter: S) -> Self {
377        self.filter_conditions = Some(filter.into());
378        self
379    }
380}
381
382impl RelationshipMapping {
383    /// Create a new relationship mapping
384    pub fn new<S: Into<String>>(rel_type: S, source_field: S, target_field: S) -> Self {
385        Self {
386            relationship_type: rel_type.into(),
387            source_id_field: source_field.into(),
388            target_id_field: target_field.into(),
389            type_field: None,
390            property_fields: Vec::new(),
391            filter_conditions: None,
392            source_label: None,
393            target_label: None,
394        }
395    }
396
397    /// Set the node labels this relationship type connects
398    pub fn with_endpoints<S: Into<String>>(mut self, source_label: S, target_label: S) -> Self {
399        self.source_label = Some(source_label.into());
400        self.target_label = Some(target_label.into());
401        self
402    }
403
404    /// Set the type field for this relationship
405    pub fn with_type_field<S: Into<String>>(mut self, type_field: S) -> Self {
406        self.type_field = Some(type_field.into());
407        self
408    }
409
410    /// Add property fields to the mapping
411    pub fn with_properties(mut self, fields: Vec<String>) -> Self {
412        self.property_fields = fields;
413        self
414    }
415
416    /// Add filter conditions for this relationship type
417    pub fn with_filter<S: Into<String>>(mut self, filter: S) -> Self {
418        self.filter_conditions = Some(filter.into());
419        self
420    }
421}
422
423#[cfg(test)]
424mod tests {
425    use super::*;
426
427    #[test]
428    fn test_mapping_iterators_keep_normalized_keys_and_original_labels() {
429        let config = GraphConfig::builder()
430            .node_table("Person", "id")
431            .rel_table("KNOWS", "Person", "Person", "src", "dst")
432            .build()
433            .unwrap();
434        let (key, node) = config.iter_node_mappings().next().unwrap();
435        assert_eq!(key, "person");
436        assert_eq!(node.label, "Person");
437        assert!(std::ptr::eq(
438            node,
439            config.get_node_mapping("PERSON").unwrap()
440        ));
441        let (key, relationship) = config.iter_relationship_mappings().next().unwrap();
442        assert_eq!(key, "knows");
443        assert_eq!(relationship.relationship_type, "KNOWS");
444        assert!(std::ptr::eq(
445            relationship,
446            config.get_relationship_mapping("KnOwS").unwrap()
447        ));
448    }
449
450    #[test]
451    fn test_graph_config_builder() {
452        let config = GraphConfig::builder()
453            .node_table("Person", "person_id")
454            .node_table("Company", "company_id")
455            .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id")
456            .build()
457            .unwrap();
458
459        assert_eq!(config.iter_node_mappings().len(), 2);
460        assert_eq!(config.iter_relationship_mappings().len(), 1);
461
462        let person_mapping = config.get_node_mapping("Person").unwrap();
463        assert_eq!(person_mapping.id_field, "person_id");
464
465        let works_for_mapping = config.get_relationship_mapping("WORKS_FOR").unwrap();
466        assert_eq!(works_for_mapping.source_id_field, "person_id");
467        assert_eq!(works_for_mapping.target_id_field, "company_id");
468        assert_eq!(works_for_mapping.source_label.as_deref(), Some("Person"));
469        assert_eq!(works_for_mapping.target_label.as_deref(), Some("Company"));
470    }
471
472    #[test]
473    fn rel_table_endpoints_must_be_node_tables() {
474        let result = GraphConfig::builder()
475            .node_table("Person", "id")
476            .rel_table("WORKS_FOR", "Person", "Company", "src", "dst")
477            .build();
478        let message = result.unwrap_err().to_string();
479        assert!(message.contains("Company"), "{message}");
480    }
481
482    #[test]
483    #[allow(deprecated)]
484    fn deprecated_builder_methods_still_configure_the_graph() {
485        let config = GraphConfig::builder()
486            .with_node_label("Person", "id")
487            .with_relationship("KNOWS", "src", "dst")
488            .with_relationship_endpoints("KNOWS", "Person", "Person")
489            .build()
490            .unwrap();
491        let knows = config.get_relationship_mapping("knows").unwrap();
492        assert_eq!(config.get_node_mapping("person").unwrap().id_field, "id");
493        assert_eq!(knows.source_label.as_deref(), Some("Person"));
494    }
495
496    #[test]
497    fn test_validation_empty_id_field() {
498        let mut config = GraphConfig::default();
499        config.node_mappings.insert(
500            "Person".to_string(),
501            NodeMapping {
502                label: "Person".to_string(),
503                id_field: "".to_string(),
504                property_fields: Vec::new(),
505                filter_conditions: None,
506            },
507        );
508
509        assert!(config.validate().is_err());
510    }
511
512    #[test]
513    fn deserialization_validates_through_the_builder() {
514        let config = GraphConfig::builder()
515            .node_table("Person", "id")
516            .rel_table("KNOWS", "Person", "Person", "src", "dst")
517            .build()
518            .unwrap();
519        let json = serde_json::to_string(&config).unwrap();
520        let restored: GraphConfig = serde_json::from_str(&json).unwrap();
521        assert_eq!(restored.get_node_mapping("PERSON").unwrap().label, "Person");
522        assert!(restored.get_relationship_mapping("knows").is_some());
523
524        let unnormalized = json.replace("\"person\":", "\"Person\":");
525        let message = serde_json::from_str::<GraphConfig>(&unnormalized)
526            .unwrap_err()
527            .to_string();
528        assert!(message.contains("not normalized"), "{message}");
529    }
530
531    #[test]
532    fn test_node_mapping_with_properties() {
533        let mapping = NodeMapping::new("Person", "id")
534            .with_properties(vec!["name".to_string(), "age".to_string()])
535            .with_filter("age > 18".to_string());
536
537        assert_eq!(mapping.property_fields.len(), 2);
538        assert!(mapping.filter_conditions.is_some());
539    }
540
541    #[test]
542    fn test_case_insensitive_node_label_lookup() {
543        // Test that node label lookups are case-insensitive
544        let config = GraphConfig::builder()
545            .node_table("Person", "person_id")
546            .node_table("Company", "company_id")
547            .build()
548            .unwrap();
549
550        // All case variations should work
551        assert!(config.get_node_mapping("Person").is_some());
552        assert!(config.get_node_mapping("person").is_some());
553        assert!(config.get_node_mapping("PERSON").is_some());
554        assert!(config.get_node_mapping("PeRsOn").is_some());
555
556        assert!(config.get_node_mapping("Company").is_some());
557        assert!(config.get_node_mapping("company").is_some());
558        assert!(config.get_node_mapping("COMPANY").is_some());
559
560        // Non-existent labels should return None
561        assert!(config.get_node_mapping("Unknown").is_none());
562        assert!(config.get_node_mapping("unknown").is_none());
563
564        // Verify we get the same mapping regardless of case
565        let mapping1 = config.get_node_mapping("Person").unwrap();
566        let mapping2 = config.get_node_mapping("person").unwrap();
567        let mapping3 = config.get_node_mapping("PERSON").unwrap();
568
569        assert_eq!(mapping1.id_field, mapping2.id_field);
570        assert_eq!(mapping2.id_field, mapping3.id_field);
571        assert_eq!(mapping1.id_field, "person_id");
572    }
573
574    #[test]
575    fn test_case_insensitive_relationship_type_lookup() {
576        // Test that relationship type lookups are case-insensitive
577        let config = GraphConfig::builder()
578            .node_table("Person", "person_id")
579            .node_table("Company", "company_id")
580            .rel_table("FOLLOWS", "Person", "Person", "src_id", "dst_id")
581            .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id")
582            .build()
583            .unwrap();
584
585        // All case variations should work
586        assert!(config.get_relationship_mapping("FOLLOWS").is_some());
587        assert!(config.get_relationship_mapping("follows").is_some());
588        assert!(config.get_relationship_mapping("Follows").is_some());
589
590        assert!(config.get_relationship_mapping("WORKS_FOR").is_some());
591        assert!(config.get_relationship_mapping("works_for").is_some());
592        assert!(config.get_relationship_mapping("Works_For").is_some());
593
594        // Non-existent types should return None
595        assert!(config.get_relationship_mapping("UNKNOWN").is_none());
596        assert!(config.get_relationship_mapping("unknown").is_none());
597
598        // Verify we get the same mapping regardless of case
599        let mapping1 = config.get_relationship_mapping("FOLLOWS").unwrap();
600        let mapping2 = config.get_relationship_mapping("follows").unwrap();
601        let mapping3 = config.get_relationship_mapping("Follows").unwrap();
602
603        assert_eq!(mapping1.source_id_field, mapping2.source_id_field);
604        assert_eq!(mapping2.source_id_field, mapping3.source_id_field);
605        assert_eq!(mapping1.source_id_field, "src_id");
606    }
607
608    #[test]
609    fn test_duplicate_label_different_case_should_overwrite() {
610        let builder = GraphConfig::builder()
611            .node_table("Person", "id")
612            .node_table("person", "id2"); // Should overwrite first entry
613
614        // Should have only 1 entry (second overwrites first due to lowercase key normalization)
615        assert_eq!(builder.node_mappings.len(), 1);
616
617        // The second one should have won (id2)
618        let mapping = builder.node_mappings.get("person").unwrap();
619        assert_eq!(mapping.id_field, "id2");
620    }
621}
lance_csr/config.rs

lance_csr/
config.rs
