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

lance_csr/
config.rs
