diff --git a/README.md b/README.md index 75c113a..476a3fc 100644 --- a/README.md +++ b/README.md @@ -119,7 +119,7 @@ people = pa.table({ config = ( GraphConfig.builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() ) diff --git a/ci/check_wheel.py b/ci/check_wheel.py index 240e1c8..3d3e31c 100644 --- a/ci/check_wheel.py +++ b/ci/check_wheel.py @@ -39,8 +39,8 @@ def main(): config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) people = pa.table({"id": [1, 2, 3, 4], "name": ["Alice", "Bob", "Carol", "Dave"]}) diff --git a/crates/lance-graph-benches/benches/graph_execution.rs b/crates/lance-graph-benches/benches/graph_execution.rs index fab1a94..2b6831b 100644 --- a/crates/lance-graph-benches/benches/graph_execution.rs +++ b/crates/lance-graph-benches/benches/graph_execution.rs @@ -130,8 +130,8 @@ fn bench_cypher_execution(c: &mut Criterion) { // Helper function to create graph config let make_config = || { GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_relationship("FRIEND_OF", "person1_id", "person2_id") + .node_table("Person", "person_id") + .rel_table("FRIEND_OF", "Person", "Person", "person1_id", "person2_id") .build() .unwrap() }; diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 569b47f..778d688 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -377,6 +377,12 @@ pub struct GraphConfigBuilder { inner: lance_graph::config::GraphConfigBuilder, } +/// Raise a ``DeprecationWarning`` pointing at the caller's line. +fn deprecated(py: Python, message: &std::ffi::CStr) -> PyResult<()> { + let category = py.get_type::(); + PyErr::warn(py, &category, message, 1) +} + #[pymethods] impl GraphConfigBuilder { #[new] @@ -386,57 +392,89 @@ impl GraphConfigBuilder { } } - /// Add a node label mapping + /// Add a node table: the nodes with ``label``, one per row of the dataset of the + /// same name, identified by its ``primary_key`` column. /// /// Parameters /// ---------- /// label : str - /// The node label (e.g., "Person") - /// id_field : str - /// The field in the dataset that serves as the node ID + /// The node label, and the name of its dataset (e.g., "Person") + /// primary_key : str + /// The column that identifies each node /// /// Returns /// ------- /// GraphConfigBuilder - /// A new builder with the node mapping applied - fn with_node_label(&self, label: &str, id_field: &str) -> Self { + /// A new builder with the node table added + fn node_table(&self, label: &str, primary_key: &str) -> Self { Self { - inner: self.inner.clone().with_node_label(label, id_field), + inner: self.inner.clone().node_table(label, primary_key), } } - /// Add a relationship mapping + /// Add a relationship table: edges of type ``rel_type`` from ``from_table`` nodes to + /// ``to_table`` nodes, one per row of the dataset of the same name. /// /// Parameters /// ---------- /// rel_type : str - /// The relationship type (e.g., "KNOWS") - /// source_field : str - /// The field containing source node IDs - /// target_field : str - /// The field containing target node IDs - /// source_label, target_label : str, optional - /// Node labels at each end (like ``FROM Post TO Person``). Needed only for - /// relationship-type alternation (``[:A|B]``) next to nodes whose label - /// depends on the type, and for multi-label nodes. Give both or neither. + /// The relationship type, and the name of its dataset (e.g., "KNOWS") + /// from_table : str + /// The node table at the source end + /// to_table : str + /// The node table at the target end + /// src : str + /// The column holding the source node's key + /// dst : str + /// The column holding the target node's key /// /// Returns /// ------- /// GraphConfigBuilder - /// A new builder with the relationship mapping applied + /// A new builder with the relationship table added + fn rel_table( + &self, + rel_type: &str, + from_table: &str, + to_table: &str, + src: &str, + dst: &str, + ) -> Self { + Self { + inner: self + .inner + .clone() + .rel_table(rel_type, from_table, to_table, src, dst), + } + } + + /// Deprecated: use ``node_table``. + fn with_node_label(&self, py: Python, label: &str, id_field: &str) -> PyResult { + deprecated(py, c"with_node_label() is deprecated; use node_table()")?; + Ok(self.node_table(label, id_field)) + } + + /// Deprecated: use ``rel_table``, which also takes the node tables at both ends. #[pyo3(signature = (rel_type, source_field, target_field, source_label=None, target_label=None))] fn with_relationship( &self, + py: Python, rel_type: &str, source_field: &str, target_field: &str, source_label: Option<&str>, target_label: Option<&str>, ) -> PyResult { + deprecated( + py, + c"with_relationship() is deprecated; use rel_table(), which also takes the node tables at both ends", + )?; + #[allow(deprecated)] let inner = self .inner .clone() .with_relationship(rel_type, source_field, target_field); + #[allow(deprecated)] let inner = match (source_label, target_label) { (Some(src), Some(dst)) => inner.with_relationship_endpoints(rel_type, src, dst), (None, None) => inner, @@ -761,7 +799,7 @@ impl CypherQuery { /// -------- /// >>> from lance_csr import CypherQuery, VectorSearch, DistanceMetric, GraphConfig /// >>> - /// >>> config = GraphConfig.builder().with_node_label("Document", "id").build() + /// >>> config = GraphConfig.builder().node_table("Document", "id").build() /// >>> query = CypherQuery("MATCH (d:Document) WHERE d.category = 'tech' RETURN d.id, d.name, d.embedding") /// >>> query = query.with_config(config) /// >>> @@ -1851,8 +1889,8 @@ impl RecordBatchReader for BlockingBatchReader { /// >>> /// >>> # Setup /// >>> config = GraphConfig.builder() \\ -/// ... .with_node_label("Person", "id") \\ -/// ... .with_relationship("KNOWS", "src_id", "dst_id") \\ +/// ... .node_table("Person", "id") \\ +/// ... .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") \\ /// ... .build() /// >>> /// >>> datasets = { diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index b2da6f8..5fecb3a 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -23,8 +23,8 @@ use arrow_schema::{DataType, Field, Schema}; use lance_graph::{CypherQuery, ExecutionStrategy, GraphConfig}; let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "person_id") + .rel_table("KNOWS", "Person", "Person", "src_person_id", "dst_person_id") .build()?; let schema = Arc::new(Schema::new(vec![ @@ -63,8 +63,8 @@ Graph mappings are declared with `GraphConfig::builder()`: use lance_graph::{GraphConfig, NodeMapping, RelationshipMapping}; let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "person_id") + .rel_table("KNOWS", "Person", "Person", "src_person_id", "dst_person_id") .build()?; ``` diff --git a/crates/lance-graph/src/config.rs b/crates/lance-graph/src/config.rs index 0ed7756..e35eb75 100644 --- a/crates/lance-graph/src/config.rs +++ b/crates/lance-graph/src/config.rs @@ -214,18 +214,18 @@ impl GraphConfigBuilder { Self::default() } - /// Add a node label mapping + /// A node table: the nodes labelled `label`, one per row of the dataset of the same + /// name, identified by its `primary_key` column. /// - /// Note: Labels are case-insensitive. If you add "Person" and "person", the second - /// will overwrite the first. Keys are stored as lowercase to prevent duplicates. - pub fn with_node_label>(mut self, label: S, id_field: S) -> Self { + /// Labels are case-insensitive: adding "Person" and then "person" keeps the second. + pub fn node_table>(mut self, label: S, primary_key: S) -> Self { let label_str = label.into(); let normalized_key = label_str.to_lowercase(); self.node_mappings.insert( normalized_key, NodeMapping { label: label_str, // Keep original case for display - id_field: id_field.into(), + id_field: primary_key.into(), property_fields: Vec::new(), filter_conditions: None, }, @@ -233,6 +233,42 @@ impl GraphConfigBuilder { self } + /// A relationship table: edges of type `rel_type` from `from_table` nodes to + /// `to_table` nodes, one per row of the dataset of the same name, with the source + /// node's key in column `src` and the target node's key in column `dst`. Both node + /// tables must be defined by the time [`Self::build`] runs. + pub fn rel_table>( + mut self, + rel_type: S, + from_table: S, + to_table: S, + src: S, + dst: S, + ) -> Self { + let type_str = rel_type.into(); + let normalized_key = type_str.to_lowercase(); + self.relationship_mappings.insert( + normalized_key, + RelationshipMapping { + relationship_type: type_str, // Keep original case for display + source_id_field: src.into(), + target_id_field: dst.into(), + type_field: None, + property_fields: Vec::new(), + filter_conditions: None, + source_label: Some(from_table.into()), + target_label: Some(to_table.into()), + }, + ); + self + } + + /// Add a node label mapping. + #[deprecated(since = "0.2.0", note = "use `node_table`")] + pub fn with_node_label>(self, label: S, id_field: S) -> Self { + self.node_table(label, id_field) + } + /// Add a node mapping with additional configuration pub fn with_node_mapping(mut self, mapping: NodeMapping) -> Self { let normalized_key = mapping.label.to_lowercase(); @@ -240,7 +276,11 @@ impl GraphConfigBuilder { self } - /// Add a relationship type mapping + /// Add a relationship type mapping, without endpoint node tables. + #[deprecated( + since = "0.2.0", + note = "use `rel_table`, which also takes the node tables at both ends" + )] pub fn with_relationship>( mut self, rel_type: S, @@ -267,6 +307,10 @@ impl GraphConfigBuilder { /// Record the node labels a configured relationship type connects (see /// [`RelationshipMapping::source_label`]). No-op for an unknown type. + #[deprecated( + since = "0.2.0", + note = "use `rel_table`, which takes the endpoint node tables" + )] pub fn with_relationship_endpoints>( mut self, rel_type: &str, @@ -389,9 +433,9 @@ mod tests { #[test] fn test_graph_config_builder() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -404,6 +448,32 @@ mod tests { let works_for_mapping = config.get_relationship_mapping("WORKS_FOR").unwrap(); assert_eq!(works_for_mapping.source_id_field, "person_id"); assert_eq!(works_for_mapping.target_id_field, "company_id"); + assert_eq!(works_for_mapping.source_label.as_deref(), Some("Person")); + assert_eq!(works_for_mapping.target_label.as_deref(), Some("Company")); + } + + #[test] + fn rel_table_endpoints_must_be_node_tables() { + let result = GraphConfig::builder() + .node_table("Person", "id") + .rel_table("WORKS_FOR", "Person", "Company", "src", "dst") + .build(); + let message = result.unwrap_err().to_string(); + assert!(message.contains("Company"), "{message}"); + } + + #[test] + #[allow(deprecated)] + fn deprecated_builder_methods_still_configure_the_graph() { + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src", "dst") + .with_relationship_endpoints("KNOWS", "Person", "Person") + .build() + .unwrap(); + let knows = config.get_relationship_mapping("knows").unwrap(); + assert_eq!(config.get_node_mapping("person").unwrap().id_field, "id"); + assert_eq!(knows.source_label.as_deref(), Some("Person")); } #[test] @@ -436,8 +506,8 @@ mod tests { fn test_case_insensitive_node_label_lookup() { // Test that node label lookups are case-insensitive let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") .build() .unwrap(); @@ -469,8 +539,10 @@ mod tests { fn test_case_insensitive_relationship_type_lookup() { // Test that relationship type lookups are case-insensitive let config = GraphConfig::builder() - .with_relationship("FOLLOWS", "src_id", "dst_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("FOLLOWS", "Person", "Person", "src_id", "dst_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -500,8 +572,8 @@ mod tests { #[test] fn test_duplicate_label_different_case_should_overwrite() { let builder = GraphConfig::builder() - .with_node_label("Person", "id") - .with_node_label("person", "id2"); // Should overwrite first entry + .node_table("Person", "id") + .node_table("person", "id2"); // Should overwrite first entry // Should have only 1 entry (second overwrites first due to lowercase key normalization) assert_eq!(builder.node_mappings.len(), 1); diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index e573d8e..89abee0 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -289,8 +289,8 @@ mod tests { }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); let _planner = crate::datafusion_planner::DataFusionPlanner::new(cfg); @@ -340,8 +340,8 @@ mod tests { }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); let _planner = crate::datafusion_planner::DataFusionPlanner::new(cfg); @@ -380,8 +380,14 @@ mod tests { }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let _planner = crate::datafusion_planner::DataFusionPlanner::new(cfg); diff --git a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs index 3e1f981..9b4ad83 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs @@ -800,8 +800,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -843,8 +849,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -885,8 +897,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -928,8 +946,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -971,8 +995,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -1025,8 +1055,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 173cfbf..f2b8333 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -195,8 +195,8 @@ mod tests { fn test_collect_variables_includes_relationship_variables() { // Test that collect_variables now captures relationship variables let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); diff --git a/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs b/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs index 107ef87..7fc3ac3 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs @@ -531,8 +531,14 @@ mod tests { // 2. Generates composite keys (src_id + dst_id) for relationship variables // 3. Generates single keys for node variables let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -624,8 +630,8 @@ mod tests { fn test_shared_relationship_variable_detected() { // Test that shared relationship variables are detected let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); diff --git a/crates/lance-graph/src/datafusion_planner/config_helpers.rs b/crates/lance-graph/src/datafusion_planner/config_helpers.rs index fd94433..af29dad 100644 --- a/crates/lance-graph/src/datafusion_planner/config_helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/config_helpers.rs @@ -108,8 +108,8 @@ mod tests { fn planner_with_basic_config() -> DataFusionPlanner { let cfg = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() .unwrap(); DataFusionPlanner::with_catalog(cfg, make_catalog()) @@ -195,7 +195,7 @@ mod tests { #[test] fn test_get_target_node_mapping_missing_config_entry() { let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -222,7 +222,7 @@ mod tests { #[test] fn test_get_catalog_missing() { let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let planner = DataFusionPlanner::new(cfg); diff --git a/crates/lance-graph/src/datafusion_planner/join_ops.rs b/crates/lance-graph/src/datafusion_planner/join_ops.rs index e4e5110..da9ed1c 100644 --- a/crates/lance-graph/src/datafusion_planner/join_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/join_ops.rs @@ -385,8 +385,14 @@ mod tests { }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -430,8 +436,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -474,8 +486,14 @@ mod tests { }], }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -526,8 +544,14 @@ mod tests { predicate: pred, }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); diff --git a/crates/lance-graph/src/datafusion_planner/scan_ops.rs b/crates/lance-graph/src/datafusion_planner/scan_ops.rs index 4d877e5..33434f8 100644 --- a/crates/lance-graph/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/scan_ops.rs @@ -96,7 +96,7 @@ impl DataFusionPlanner { return Err(crate::error::GraphError::ConfigError { message: format!( "Node label '{}' not found in catalog. \ - Ensure the label is registered in your GraphConfig with .with_node_label()", + Ensure the label is registered in your GraphConfig with .node_table()", label ), location: snafu::Location::new(file!(), line!(), column!()), @@ -316,7 +316,7 @@ mod tests { }; let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -335,8 +335,14 @@ mod tests { fn test_temp_variable_with_underscores_in_source() { // Test that temporary variables work correctly when source variable contains underscores let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -375,7 +381,7 @@ mod tests { // Test that when a catalog is attached, scanning a non-existent label fails fast // This catches typos and configuration issues at planning time let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); @@ -413,7 +419,7 @@ mod tests { // Test that when no catalog is attached, scanning creates an empty source fallback // This allows DataFusionPlanner::new() to work without requiring a catalog let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let planner = DataFusionPlanner::new(cfg); // No catalog attached @@ -437,8 +443,8 @@ mod tests { fn test_expand_with_missing_relationship() { // Test that expanding with non-existent relationship type handles gracefully let cfg = crate::config::GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); diff --git a/crates/lance-graph/src/datafusion_planner/test_fixtures.rs b/crates/lance-graph/src/datafusion_planner/test_fixtures.rs index 7aceba6..90ef19e 100644 --- a/crates/lance-graph/src/datafusion_planner/test_fixtures.rs +++ b/crates/lance-graph/src/datafusion_planner/test_fixtures.rs @@ -33,15 +33,21 @@ pub fn make_catalog() -> Arc { pub fn person_config() -> GraphConfig { GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap() } pub fn person_knows_config() -> GraphConfig { GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap() } diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index df63f63..37083e7 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -22,8 +22,8 @@ //! //! # fn example() -> Result<()> { //! let config = GraphConfig::builder() -//! .with_node_label("Person", "person_id") -//! .with_relationship("KNOWS", "src_person_id", "dst_person_id") +//! .node_table("Person", "person_id") +//! .rel_table("KNOWS", "Person", "Person", "src_person_id", "dst_person_id") //! .build()?; //! //! let query = CypherQuery::new("MATCH (p:Person) WHERE p.age > 30 RETURN p.name")? diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index ef6520f..ed89e08 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -2276,7 +2276,7 @@ mod tests { let ast = parse_cypher_query(query_text).unwrap(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let mut planner = LogicalPlanner::new(&config); @@ -2343,11 +2343,11 @@ mod tests { #[test] fn continued_path_starts_from_the_bound_end_with_smaller_fan_out() { let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_node_label("Post", "id") - .with_node_label("Forum", "id") - .with_relationship("postHasCreator", "src", "dst") - .with_relationship("containerOf", "src", "dst") + .node_table("Person", "id") + .node_table("Post", "id") + .node_table("Forum", "id") + .rel_table("postHasCreator", "Post", "Person", "src", "dst") + .rel_table("containerOf", "Forum", "Post", "src", "dst") .build() .unwrap(); let q = "MATCH (friend:Person), (forum:Forum) WITH friend, forum \ @@ -2381,8 +2381,8 @@ mod tests { fn continued_path_bound_only_at_its_far_end_starts_there() { // After `WITH DISTINCT b`, `a` is a new variable: expansion must start from `b`. let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() .unwrap(); let q = "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH DISTINCT b \ diff --git a/crates/lance-graph/src/pattern_alternatives.rs b/crates/lance-graph/src/pattern_alternatives.rs index 4169a75..f2d1356 100644 --- a/crates/lance-graph/src/pattern_alternatives.rs +++ b/crates/lance-graph/src/pattern_alternatives.rs @@ -283,8 +283,8 @@ pub fn expand( for rel in rels.iter().filter(|r| r.left == id || r.right == id) { if let Some(t) = rel.types.iter().find(|t| endpoints(t).is_none()) { return Err(plan_error(format!( - "relationship type '{t}' needs endpoint labels (GraphConfig \ - with_relationship_endpoints) to be planned next to a node that can \ + "relationship type '{t}' needs endpoint labels (GraphConfig rel_table's \ + from_table and to_table) to be planned next to a node that can \ have several labels" ))); } @@ -388,8 +388,8 @@ fn unroll_variable_length_alternation( m.and_then(|m| m.target_label.clone()), ) else { return Err(plan_error(format!( - "relationship type '{t}' needs endpoint labels (GraphConfig \ - with_relationship_endpoints) for alternation in a variable-length pattern" + "relationship type '{t}' needs endpoint labels (GraphConfig rel_table's \ + from_table and to_table) for alternation in a variable-length pattern" ))); }; // Labels in traversal order: the segment is read from its left node. @@ -565,7 +565,7 @@ mod tests { fn ldbc_config() -> GraphConfig { let mut b = GraphConfig::builder(); for label in ["Person", "Post", "Comment"] { - b = b.with_node_label(label, "id"); + b = b.node_table(label, "id"); } for (t, s, d) in [ ("postHasCreator", "Post", "Person"), @@ -573,11 +573,13 @@ mod tests { ("replyOfPost", "Comment", "Post"), ("replyOfComment", "Comment", "Comment"), ] { - b = b - .with_relationship(t, "src", "dst") - .with_relationship_endpoints(t, s, d); + b = b.rel_table(t, s, d, "src", "dst"); + } + // Without endpoint labels, which only the deprecated builder can still configure. + #[allow(deprecated)] + { + b = b.with_relationship("knows", "src", "dst"); } - b = b.with_relationship("knows", "src", "dst"); b.build().unwrap() } diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index 6f8e004..b5333e4 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -475,13 +475,13 @@ mod tests { fn config() -> GraphConfig { GraphConfig::builder() - .with_node_label("Forum", "id") - .with_node_label("Person", "id") - .with_node_label("Tag", "id") - .with_node_label("Tagclass", "id") - .with_relationship("hasMember", "src", "dst") - .with_relationship("forumHasTag", "src", "dst") - .with_relationship("hasType", "src", "dst") + .node_table("Forum", "id") + .node_table("Person", "id") + .node_table("Tag", "id") + .node_table("Tagclass", "id") + .rel_table("hasMember", "Forum", "Person", "src", "dst") + .rel_table("forumHasTag", "Forum", "Tag", "src", "dst") + .rel_table("hasType", "Tag", "Tagclass", "src", "dst") .build() .unwrap() } diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 70313ec..3b43c20 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -447,8 +447,8 @@ impl CypherQuery { /// /// // Step 1: Create GraphConfig /// let config = GraphConfig::builder() - /// .with_node_label("Person", "person_id") - /// .with_relationship("KNOWS", "src_id", "dst_id") + /// .node_table("Person", "person_id") + /// .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") /// .build()?; /// /// // Step 2: Register data sources in DataFusion @@ -1360,7 +1360,7 @@ mod tests { #[test] fn test_query_builder() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -1407,7 +1407,7 @@ mod tests { .unwrap(); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1472,8 +1472,14 @@ mod tests { // Config: Person(id) and KNOWS(src_person_id -> dst_person_id) let cfg = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); @@ -1523,7 +1529,7 @@ mod tests { .unwrap(); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1561,7 +1567,7 @@ mod tests { .unwrap(); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1595,7 +1601,7 @@ mod tests { .unwrap(); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1635,7 +1641,7 @@ mod tests { .unwrap(); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1704,7 +1710,7 @@ mod tests { // Create query let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1761,7 +1767,7 @@ mod tests { // Create query with filter let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -1847,8 +1853,8 @@ mod tests { // Create query let cfg = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); @@ -1919,7 +1925,7 @@ mod tests { // Create query with ORDER BY and LIMIT let cfg = GraphConfig::builder() - .with_node_label("Student", "id") + .node_table("Student", "id") .build() .unwrap(); @@ -1973,7 +1979,7 @@ mod tests { datasets.insert("Person".to_string(), batch); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2057,8 +2063,8 @@ mod tests { .await; let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_relationship("FRIEND_OF", "person1_id", "person2_id") + .node_table("Person", "person_id") + .rel_table("FRIEND_OF", "Person", "Person", "person1_id", "person2_id") .build() .expect("valid graph config"); @@ -2100,7 +2106,7 @@ mod tests { datasets.insert("Person".to_string(), batch); let cfg = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index dbe4311..b8cf649 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -970,10 +970,10 @@ mod tests { fn test_config() -> GraphConfig { GraphConfig::builder() - .with_node_label("Person", "id") - .with_node_label("Employee", "id") - .with_node_label("Company", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .node_table("Employee", "id") + .node_table("Company", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap() } @@ -1246,7 +1246,7 @@ mod tests { .with_node_mapping( NodeMapping::new("Person", "id").with_properties(vec!["name".to_string()]), ) - .with_relationship("KNOWS", "src_id", "dst_id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); @@ -1368,9 +1368,9 @@ mod tests { // Custom config that knows both relationship types to avoid warnings muddying the assertion let custom_config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") - .with_relationship("FRIEND", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") + .rel_table("FRIEND", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); diff --git a/crates/lance-graph/tests/test_case_insensitivity.rs b/crates/lance-graph/tests/test_case_insensitivity.rs index 6cd1c95..8fd59f8 100644 --- a/crates/lance-graph/tests/test_case_insensitivity.rs +++ b/crates/lance-graph/tests/test_case_insensitivity.rs @@ -40,7 +40,7 @@ fn create_test_person_dataset() -> RecordBatch { async fn test_identifiers_case_insensitive() { // Test labels, properties, and table names with various case combinations let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let person_batch = create_test_person_dataset(); @@ -111,7 +111,7 @@ async fn test_identifiers_case_insensitive() { async fn test_clauses_case_insensitive() { // Test WHERE, ORDER BY, and variable references with case variations let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let person_batch = create_test_person_dataset(); @@ -207,8 +207,8 @@ async fn test_clauses_case_insensitive() { async fn test_relationships_case_insensitive() { // Test relationship types and relationship variables with case variations let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_id", "dst_id") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() .unwrap(); @@ -278,7 +278,7 @@ async fn test_relationships_case_insensitive() { async fn test_complex_mixed_case_query() { // Comprehensive end-to-end test with mixed case everywhere let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); let person_batch = create_test_person_dataset(); diff --git a/crates/lance-graph/tests/test_datafusion_pipeline.rs b/crates/lance-graph/tests/test_datafusion_pipeline.rs index 906f27e..2b2041f 100644 --- a/crates/lance-graph/tests/test_datafusion_pipeline.rs +++ b/crates/lance-graph/tests/test_datafusion_pipeline.rs @@ -109,8 +109,14 @@ fn create_knows_dataset() -> RecordBatch { /// Helper function to create graph config fn create_graph_config() -> GraphConfig { GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap() } @@ -2589,7 +2595,7 @@ async fn test_datafusion_varlength_count() { async fn test_count_star_all_nodes() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2621,7 +2627,7 @@ async fn test_count_star_all_nodes() { async fn test_count_distinct_star_rejected() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2653,7 +2659,7 @@ async fn test_count_distinct_star_rejected() { async fn test_count_variable() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2684,7 +2690,7 @@ async fn test_count_variable() { async fn test_count_with_filter() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2716,7 +2722,7 @@ async fn test_count_with_filter() { async fn test_count_property() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2746,7 +2752,7 @@ async fn test_count_property() { async fn test_count_with_grouping() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2804,7 +2810,7 @@ async fn test_count_with_grouping() { async fn test_count_without_alias_has_descriptive_name() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2834,7 +2840,7 @@ async fn test_count_without_alias_has_descriptive_name() { async fn test_count_property_without_alias_has_descriptive_name() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -2863,8 +2869,14 @@ async fn test_count_property_without_alias_has_descriptive_name() { #[tokio::test] async fn test_count_distinct_basic() { let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); @@ -2901,8 +2913,14 @@ async fn test_count_distinct_basic() { #[tokio::test] async fn test_count_vs_count_distinct() { let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); @@ -2965,8 +2983,14 @@ async fn test_count_vs_count_distinct() { #[tokio::test] async fn test_count_distinct_with_grouping() { let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); @@ -3027,7 +3051,7 @@ async fn test_count_distinct_with_grouping() { async fn test_sum_property() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3058,7 +3082,7 @@ async fn test_sum_property() { async fn test_sum_with_filter() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3090,7 +3114,7 @@ async fn test_sum_with_filter() { async fn test_sum_with_grouping() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3145,7 +3169,7 @@ async fn test_sum_with_grouping() { async fn test_sum_without_alias_has_descriptive_name() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3175,7 +3199,7 @@ async fn test_sum_without_alias_has_descriptive_name() { async fn test_avg_property() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3206,7 +3230,7 @@ async fn test_avg_property() { async fn test_avg_with_filter() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3238,7 +3262,7 @@ async fn test_avg_with_filter() { async fn test_avg_with_grouping() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3294,7 +3318,7 @@ async fn test_avg_with_grouping() { async fn test_avg_without_alias_has_descriptive_name() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3324,7 +3348,7 @@ async fn test_avg_without_alias_has_descriptive_name() { async fn test_min_property() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3357,7 +3381,7 @@ async fn test_min_property() { async fn test_max_property() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -3390,7 +3414,7 @@ async fn test_max_property() { async fn test_min_max_with_grouping() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4447,7 +4471,7 @@ async fn test_collect_property() { // Test COLLECT aggregation - collects values into an array let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4474,7 +4498,7 @@ async fn test_collect_with_grouping() { // Test COLLECT with GROUP BY - collect names grouped by city let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4515,7 +4539,7 @@ async fn test_collect_with_null_values() { // David has NULL city, so collecting cities should include the null let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4550,7 +4574,7 @@ async fn test_with_simple_projection() { // Test WITH as a simple projection pass-through let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4591,7 +4615,7 @@ async fn test_with_aggregation() { // Test WITH for aggregation: count people by city let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4622,7 +4646,7 @@ async fn test_with_order_by_limit_and_where() { // Test WITH with ORDER BY, LIMIT, and post-WITH WHERE filter let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); @@ -4673,8 +4697,14 @@ async fn test_with_post_match_chaining() { let person_batch = create_person_dataset(); let knows_batch = create_knows_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap(); @@ -4712,7 +4742,7 @@ async fn test_with_post_match_chaining() { async fn test_unimplemented_scalar_function_errors() { let person_batch = create_person_dataset(); let config = GraphConfig::builder() - .with_node_label("Person", "id") + .node_table("Person", "id") .build() .unwrap(); diff --git a/crates/lance-graph/tests/test_datafusion_scenarios.rs b/crates/lance-graph/tests/test_datafusion_scenarios.rs index 287054f..b77ae44 100644 --- a/crates/lance-graph/tests/test_datafusion_scenarios.rs +++ b/crates/lance-graph/tests/test_datafusion_scenarios.rs @@ -50,7 +50,7 @@ fn animals_graph() -> TestGraph { datasets.insert("Animal".to_string(), batch); let config = GraphConfig::builder() - .with_node_label("Animal", "id") + .node_table("Animal", "id") .build() .unwrap(); @@ -90,7 +90,7 @@ fn planets_graph() -> TestGraph { datasets.insert("Planet".to_string(), batch); let config = GraphConfig::builder() - .with_node_label("Planet", "id") + .node_table("Planet", "id") .build() .unwrap(); @@ -169,9 +169,9 @@ fn books_graph() -> TestGraph { datasets.insert("WROTE".to_string(), wrote_batch); let config = GraphConfig::builder() - .with_node_label("Author", "id") - .with_node_label("Book", "id") - .with_relationship("WROTE", "author_id", "book_id") + .node_table("Author", "id") + .node_table("Book", "id") + .rel_table("WROTE", "Author", "Book", "author_id", "book_id") .build() .unwrap(); @@ -223,8 +223,14 @@ fn airports_graph() -> TestGraph { datasets.insert("ROUTE".to_string(), route_batch); let config = GraphConfig::builder() - .with_node_label("Airport", "code") - .with_relationship("ROUTE", "src_airport_code", "dst_airport_code") + .node_table("Airport", "code") + .rel_table( + "ROUTE", + "Airport", + "Airport", + "src_airport_code", + "dst_airport_code", + ) .build() .unwrap(); @@ -291,9 +297,15 @@ fn courses_graph() -> TestGraph { datasets.insert("TAUGHT_BY".to_string(), taught_batch); let config = GraphConfig::builder() - .with_node_label("Course", "id") - .with_node_label("Professor", "id") - .with_relationship("TAUGHT_BY", "course_id", "professor_id") + .node_table("Course", "id") + .node_table("Professor", "id") + .rel_table( + "TAUGHT_BY", + "Course", + "Professor", + "course_id", + "professor_id", + ) .build() .unwrap(); @@ -380,11 +392,11 @@ fn engineers_graph() -> TestGraph { datasets.insert("USES".to_string(), uses_batch); let config = GraphConfig::builder() - .with_node_label("Engineer", "id") - .with_node_label("Team", "id") - .with_node_label("Tool", "id") - .with_relationship("MEMBER_OF", "engineer_id", "team_id") - .with_relationship("USES", "team_id", "tool_id") + .node_table("Engineer", "id") + .node_table("Team", "id") + .node_table("Tool", "id") + .rel_table("MEMBER_OF", "Engineer", "Team", "engineer_id", "team_id") + .rel_table("USES", "Team", "Tool", "team_id", "tool_id") .build() .unwrap(); @@ -454,9 +466,9 @@ fn suppliers_graph() -> TestGraph { datasets.insert("SUPPLIES".to_string(), supplies_batch); let config = GraphConfig::builder() - .with_node_label("Supplier", "id") - .with_node_label("Part", "id") - .with_relationship("SUPPLIES", "supplier_id", "part_id") + .node_table("Supplier", "id") + .node_table("Part", "id") + .rel_table("SUPPLIES", "Supplier", "Part", "supplier_id", "part_id") .build() .unwrap(); @@ -497,7 +509,7 @@ fn museums_graph() -> TestGraph { datasets.insert("Museum".to_string(), batch); let config = GraphConfig::builder() - .with_node_label("Museum", "id") + .node_table("Museum", "id") .build() .unwrap(); @@ -549,8 +561,14 @@ fn languages_graph() -> TestGraph { datasets.insert("RELATED_TO".to_string(), related_batch); let config = GraphConfig::builder() - .with_node_label("Language", "id") - .with_relationship("RELATED_TO", "src_language_id", "dst_language_id") + .node_table("Language", "id") + .rel_table( + "RELATED_TO", + "Language", + "Language", + "src_language_id", + "dst_language_id", + ) .build() .unwrap(); @@ -622,9 +640,9 @@ fn projects_graph() -> TestGraph { datasets.insert("HAS".to_string(), has_batch); let config = GraphConfig::builder() - .with_node_label("Project", "id") - .with_node_label("Milestone", "id") - .with_relationship("HAS", "project_id", "milestone_id") + .node_table("Project", "id") + .node_table("Milestone", "id") + .rel_table("HAS", "Project", "Milestone", "project_id", "milestone_id") .build() .unwrap(); diff --git a/crates/lance-graph/tests/test_datafusion_varlength_complex.rs b/crates/lance-graph/tests/test_datafusion_varlength_complex.rs index 46c220a..cbe51f3 100644 --- a/crates/lance-graph/tests/test_datafusion_varlength_complex.rs +++ b/crates/lance-graph/tests/test_datafusion_varlength_complex.rs @@ -136,8 +136,14 @@ fn create_complex_knows_dataset() -> RecordBatch { /// Helper function to create graph config fn create_complex_graph_config() -> GraphConfig { GraphConfig::builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .node_table("Person", "id") + .rel_table( + "KNOWS", + "Person", + "Person", + "src_person_id", + "dst_person_id", + ) .build() .unwrap() } diff --git a/crates/lance-graph/tests/test_datafusion_vector_search.rs b/crates/lance-graph/tests/test_datafusion_vector_search.rs index 6975269..b8d18a7 100644 --- a/crates/lance-graph/tests/test_datafusion_vector_search.rs +++ b/crates/lance-graph/tests/test_datafusion_vector_search.rs @@ -60,7 +60,7 @@ fn create_person_graph_with_embeddings() -> (GraphConfig, HashMap (GraphConfig, HashMap RecordBatch { #[tokio::test] async fn test_to_sql_spark_simple_node_scan() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -109,7 +109,7 @@ async fn test_to_sql_spark_simple_node_scan() { #[tokio::test] async fn test_to_sql_spark_with_filter() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -135,9 +135,9 @@ async fn test_to_sql_spark_with_filter() { #[tokio::test] async fn test_to_sql_spark_with_relationship() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -167,9 +167,9 @@ async fn test_to_sql_spark_with_relationship() { #[tokio::test] async fn test_to_sql_spark_complex_query() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -209,7 +209,7 @@ async fn test_to_sql_spark_complex_query() { #[tokio::test] async fn test_to_sql_default_dialect_no_backticks() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -234,7 +234,7 @@ async fn test_to_sql_default_dialect_no_backticks() { #[tokio::test] async fn test_spark_sql_differs_from_default_sql() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -267,7 +267,7 @@ async fn test_spark_sql_differs_from_default_sql() { #[tokio::test] async fn test_to_sql_postgresql_dialect() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); diff --git a/crates/lance-graph/tests/test_to_sql.rs b/crates/lance-graph/tests/test_to_sql.rs index 9f39d4c..90947fe 100644 --- a/crates/lance-graph/tests/test_to_sql.rs +++ b/crates/lance-graph/tests/test_to_sql.rs @@ -86,7 +86,7 @@ fn create_works_for_table() -> RecordBatch { #[tokio::test] async fn test_to_sql_simple_node_scan() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -118,7 +118,7 @@ async fn test_to_sql_simple_node_scan() { #[tokio::test] async fn test_to_sql_with_filter() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -146,7 +146,7 @@ async fn test_to_sql_with_filter() { #[tokio::test] async fn test_to_sql_with_multiple_properties() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -170,9 +170,9 @@ async fn test_to_sql_with_multiple_properties() { #[tokio::test] async fn test_to_sql_with_relationship() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -206,9 +206,9 @@ async fn test_to_sql_with_relationship() { #[tokio::test] async fn test_to_sql_with_relationship_filter() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -235,7 +235,7 @@ async fn test_to_sql_with_relationship_filter() { #[tokio::test] async fn test_to_sql_with_order_by() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -261,7 +261,7 @@ async fn test_to_sql_with_order_by() { #[tokio::test] async fn test_to_sql_with_limit() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -287,7 +287,7 @@ async fn test_to_sql_with_limit() { #[tokio::test] async fn test_to_sql_with_distinct() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -311,7 +311,7 @@ async fn test_to_sql_with_distinct() { #[tokio::test] async fn test_to_sql_with_alias() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); @@ -337,9 +337,9 @@ async fn test_to_sql_with_alias() { #[tokio::test] async fn test_to_sql_complex_query() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") .build() .unwrap(); @@ -403,7 +403,7 @@ async fn test_to_sql_missing_config() { #[tokio::test] async fn test_to_sql_empty_datasets() { let config = GraphConfig::builder() - .with_node_label("Person", "person_id") + .node_table("Person", "person_id") .build() .unwrap(); diff --git a/examples/basic_cypher.py b/examples/basic_cypher.py index 81cdaa2..0ef57b8 100644 --- a/examples/basic_cypher.py +++ b/examples/basic_cypher.py @@ -27,7 +27,7 @@ def make_people_batch() -> pa.RecordBatch: def main() -> None: - config = GraphConfigBuilder().with_node_label("Person", "person_id").build() + config = GraphConfigBuilder().node_table("Person", "person_id").build() query = CypherQuery("MATCH (n:Person) WHERE n.age > 30 RETURN n.name").with_config( config diff --git a/examples/kg_traversal.py b/examples/kg_traversal.py index 8edd982..96a7c10 100644 --- a/examples/kg_traversal.py +++ b/examples/kg_traversal.py @@ -38,8 +38,8 @@ def make_friendship_batch(n: int = 6) -> pa.RecordBatch: def main() -> None: config = ( GraphConfigBuilder() - .with_node_label("Person", "person_id") - .with_relationship("FRIEND_OF", "person1_id", "person2_id") + .node_table("Person", "person_id") + .rel_table("FRIEND_OF", "Person", "Person", "person1_id", "person2_id") .build() ) diff --git a/examples/lance_snapshot_refresh.py b/examples/lance_snapshot_refresh.py index a527aff..1588a0c 100644 --- a/examples/lance_snapshot_refresh.py +++ b/examples/lance_snapshot_refresh.py @@ -36,8 +36,8 @@ def main() -> None: config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) engine = CypherEngine( diff --git a/python/README.md b/python/README.md index 386bbe1..99bd547 100644 --- a/python/README.md +++ b/python/README.md @@ -38,9 +38,9 @@ from lance_csr import CypherQuery, GraphConfig cfg = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_node_label("City", "id") - .with_relationship("lives_in", "src", "dst") + .node_table("Person", "id") + .node_table("City", "id") + .rel_table("lives_in", "Person", "City", "src", "dst") .build() ) @@ -71,9 +71,9 @@ from lance_csr import CypherEngine, GraphConfig cfg = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_node_label("City", "id") - .with_relationship("lives_in", "src", "dst") + .node_table("Person", "id") + .node_table("City", "id") + .rel_table("lives_in", "Person", "City", "src", "dst") .build() ) @@ -100,10 +100,15 @@ result4 = engine.execute( ) ``` -Giving each relationship its endpoint labels, as in -`.with_relationship("lives_in", "src", "dst", "Person", "City")`, lets queries leave -nodes unlabeled where the relationship decides the label (`(p:Person)-[:lives_in]->(c)`), -and is required for relationship-type alternation such as `[:likes_post|likes_comment]`. +A node table maps a label to the dataset of the same name and its primary key. A +relationship table maps a type to its dataset, the node tables at its two ends +(`from_table` and `to_table`, like `FROM Person TO City` in Kuzu), and the columns that +hold each end's key (`src` and `dst`). Knowing the ends lets queries leave nodes +unlabeled where the relationship decides the label (`(p:Person)-[:lives_in]->(c)`), and +lets relationship-type alternation (`[:likes_post|likes_comment]`) pick the right table. + +`with_node_label` and `with_relationship`, the names in lance-csr 0.1, still work but are +deprecated and will be removed in a later release. ### 3. Lance datasets: snapshots, refresh and streaming @@ -115,7 +120,7 @@ scan). Each dataset is pinned to the version you pass in; later writes are visib import lance from lance_csr import CypherEngine, GraphConfig -config = GraphConfig.builder().with_node_label("Person", "id").build() +config = GraphConfig.builder().node_table("Person", "id").build() engine = CypherEngine(config, {"Person": lance.dataset("people.lance")}) engine.execute("MATCH (p:Person) RETURN count(*)") # sees the pinned version @@ -146,8 +151,8 @@ people, knows = db.open_table("Person"), db.open_table("KNOWS") config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) engine = CypherEngine(config, {"Person": people.to_lance(), "KNOWS": knows.to_lance()}) diff --git a/python/benchmarks/bounded_memory.py b/python/benchmarks/bounded_memory.py index 0d4b314..a2b5768 100644 --- a/python/benchmarks/bounded_memory.py +++ b/python/benchmarks/bounded_memory.py @@ -69,8 +69,8 @@ def measure(data: Path, query: str) -> dict: config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) start = time.perf_counter() diff --git a/python/benchmarks/csr_delta.py b/python/benchmarks/csr_delta.py index 85e1b87..e74ccbe 100644 --- a/python/benchmarks/csr_delta.py +++ b/python/benchmarks/csr_delta.py @@ -29,8 +29,8 @@ CONFIG = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) diff --git a/python/benchmarks/csr_delta_accumulated.py b/python/benchmarks/csr_delta_accumulated.py index eec3093..e2dcf73 100644 --- a/python/benchmarks/csr_delta_accumulated.py +++ b/python/benchmarks/csr_delta_accumulated.py @@ -24,8 +24,8 @@ CONFIG = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) diff --git a/python/benchmarks/csr_expand.py b/python/benchmarks/csr_expand.py index 73da1db..e7d591c 100644 --- a/python/benchmarks/csr_expand.py +++ b/python/benchmarks/csr_expand.py @@ -75,8 +75,8 @@ def child(root: Path, mode: str, repeats: int) -> dict: config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) nodes = lance.dataset(root / "person.lance").count_rows() diff --git a/python/benchmarks/engine_cache.py b/python/benchmarks/engine_cache.py index 78ab091..05a1ff2 100644 --- a/python/benchmarks/engine_cache.py +++ b/python/benchmarks/engine_cache.py @@ -76,9 +76,9 @@ def main() -> None: builder = mod.GraphConfig.builder() for i in range(labels): - builder = builder.with_node_label(f"L{i}", "id") + builder = builder.node_table(f"L{i}", "id") for i in range(labels - 1): - builder = builder.with_relationship(f"R{i}", "src", "dst") + builder = builder.rel_table(f"R{i}", f"L{i}", f"L{i + 1}", "src", "dst") config = builder.build() def open_datasets(): diff --git a/python/python/tests/test_config_builder.py b/python/python/tests/test_config_builder.py new file mode 100644 index 0000000..3ba984c --- /dev/null +++ b/python/python/tests/test_config_builder.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""The graph config builder: node and relationship tables, and the deprecated names.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + +TABLES = { + "Person": pa.table({"id": [1, 2], "name": ["a", "b"]}), + "Company": pa.table({"id": [10], "name": ["c"]}), + "WORKS_AT": pa.table({"src": [1, 2], "dst": [10, 10]}), +} +QUERY = "MATCH (p:Person)-[:WORKS_AT]->(c) RETURN p.name, c.name ORDER BY p.name" + + +def answer(config): + return CypherEngine(config, TABLES).execute(QUERY).to_pylist() + + +def test_tables_configure_the_graph(): + config = ( + GraphConfig.builder() + .node_table("Person", primary_key="id") + .node_table("Company", primary_key="id") + .rel_table( + "WORKS_AT", from_table="Person", to_table="Company", src="src", dst="dst" + ) + .build() + ) + # `c` needs no label: WORKS_AT's to_table makes it a Company. + assert answer(config) == [ + {"p.name": "a", "c.name": "c"}, + {"p.name": "b", "c.name": "c"}, + ] + + +def test_rel_table_endpoints_must_be_node_tables(): + builder = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("WORKS_AT", "Person", "Company", "src", "dst") + ) + with pytest.raises(Exception, match="Company"): + builder.build() + + +def test_deprecated_names_warn_and_still_work(): + builder = GraphConfig.builder() + with pytest.deprecated_call(match="node_table"): + builder = builder.with_node_label("Person", "id").with_node_label( + "Company", "id" + ) + with pytest.deprecated_call(match="rel_table"): + builder = builder.with_relationship( + "WORKS_AT", "src", "dst", "Person", "Company" + ) + assert answer(builder.build()) == [ + {"p.name": "a", "c.name": "c"}, + {"p.name": "b", "c.name": "c"}, + ] diff --git a/python/python/tests/test_cypher_engine.py b/python/python/tests/test_cypher_engine.py index cb57029..40c0d29 100644 --- a/python/python/tests/test_cypher_engine.py +++ b/python/python/tests/test_cypher_engine.py @@ -46,10 +46,10 @@ def graph_env(): config = ( GraphConfig.builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") - .with_relationship("FRIEND_OF", "person1_id", "person2_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") + .rel_table("FRIEND_OF", "Person", "Person", "person1_id", "person2_id") .build() ) diff --git a/python/python/tests/test_engine_alternation.py b/python/python/tests/test_engine_alternation.py index 55c2837..2d8d9fa 100644 --- a/python/python/tests/test_engine_alternation.py +++ b/python/python/tests/test_engine_alternation.py @@ -33,10 +33,14 @@ def config(endpoints=True): builder = GraphConfig.builder() for label in ("Person", "Post", "Comment"): - builder = builder.with_node_label(label, "id") + builder = builder.node_table(label, "id") for name, (src, dst, _, _) in RELS.items(): - labels = (src, dst) if endpoints else (None, None) - builder = builder.with_relationship(name, "src", "dst", *labels) + if endpoints: + builder = builder.rel_table(name, src, dst, "src", "dst") + else: + # Only the deprecated method can still leave the endpoint tables out. + with pytest.deprecated_call(): + builder = builder.with_relationship(name, "src", "dst") return builder.build() @@ -108,10 +112,9 @@ def test_label_choice_needs_endpoint_labels(datasets): def test_endpoint_labels_must_be_given_together(): - with pytest.raises(ValueError, match="together"): - GraphConfig.builder().with_node_label("Post", "id").with_relationship( - "r", "src", "dst", "Post" - ) + builder = GraphConfig.builder().node_table("Post", "id") + with pytest.raises(ValueError, match="together"), pytest.deprecated_call(): + builder.with_relationship("r", "src", "dst", "Post") @pytest.mark.parametrize("use_csr", [False, True]) diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index c7e96f2..7741a55 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -14,8 +14,8 @@ CONFIG = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) @@ -213,9 +213,9 @@ def test_target_lookup_replaces_target_join(datasets): def test_cross_label_target_lookup(tmp_path): config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_node_label("Company", "cid") - .with_relationship("WORKS_AT", "pid", "cid") + .node_table("Person", "id") + .node_table("Company", "cid") + .rel_table("WORKS_AT", "Person", "Company", "pid", "cid") .build() ) company = pa.table({"cid": pa.array([10, 20], pa.int32()), "cname": ["x", "y"]}) diff --git a/python/python/tests/test_engine_csr_delta.py b/python/python/tests/test_engine_csr_delta.py index 68a4e91..ad168c8 100644 --- a/python/python/tests/test_engine_csr_delta.py +++ b/python/python/tests/test_engine_csr_delta.py @@ -14,8 +14,8 @@ CONFIG = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) N = 40 diff --git a/python/python/tests/test_engine_exists.py b/python/python/tests/test_engine_exists.py index 00fd962..f96ce97 100644 --- a/python/python/tests/test_engine_exists.py +++ b/python/python/tests/test_engine_exists.py @@ -81,7 +81,7 @@ def test_date_part_and_epoch_millis(): ), } ) - config = GraphConfig.builder().with_node_label("Person", "id").build() + config = GraphConfig.builder().node_table("Person", "id").build() engine = CypherEngine(config, {"Person": people}) got = engine.execute( "MATCH (p:Person) RETURN p.id, date_part('month', p.birthday) AS m, " diff --git a/python/python/tests/test_engine_lance.py b/python/python/tests/test_engine_lance.py index fd2605a..1b815d3 100644 --- a/python/python/tests/test_engine_lance.py +++ b/python/python/tests/test_engine_lance.py @@ -22,8 +22,8 @@ CONFIG = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_relationship("KNOWS", "src", "dst") + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") .build() ) diff --git a/python/python/tests/test_engine_optional.py b/python/python/tests/test_engine_optional.py index 2002047..c92689d 100644 --- a/python/python/tests/test_engine_optional.py +++ b/python/python/tests/test_engine_optional.py @@ -85,11 +85,11 @@ def test_continued_path_expands_from_the_cheaper_bound_end(): # `friend`, should expand from `forum` and check the creator at the end. config = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_node_label("Post", "id") - .with_node_label("Forum", "id") - .with_relationship("postHasCreator", "src", "dst") - .with_relationship("containerOf", "src", "dst") + .node_table("Person", "id") + .node_table("Post", "id") + .node_table("Forum", "id") + .rel_table("postHasCreator", "Post", "Person", "src", "dst") + .rel_table("containerOf", "Forum", "Post", "src", "dst") .build() ) i64 = pa.int64() diff --git a/python/python/tests/test_engine_values.py b/python/python/tests/test_engine_values.py index fa35a5c..b66bb3d 100644 --- a/python/python/tests/test_engine_values.py +++ b/python/python/tests/test_engine_values.py @@ -45,12 +45,12 @@ def utc(*args): } CONFIG = ( GraphConfig.builder() - .with_node_label("Person", "id") - .with_node_label("City", "id") - .with_node_label("Post", "id") - .with_relationship("KNOWS", "src", "dst", "Person", "Person") - .with_relationship("LIVES_IN", "src", "dst", "Person", "City") - .with_relationship("HAS_CREATOR", "src", "dst", "Post", "Person") + .node_table("Person", "id") + .node_table("City", "id") + .node_table("Post", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .rel_table("LIVES_IN", "Person", "City", "src", "dst") + .rel_table("HAS_CREATOR", "Post", "Person", "src", "dst") .build() ) diff --git a/python/python/tests/test_explain.py b/python/python/tests/test_explain.py index 20032c1..d176f22 100644 --- a/python/python/tests/test_explain.py +++ b/python/python/tests/test_explain.py @@ -16,7 +16,7 @@ def person_data(): } ) - config = GraphConfig.builder().with_node_label("Person", "person_id").build() + config = GraphConfig.builder().node_table("Person", "person_id").build() return config, people_table diff --git a/python/python/tests/test_graph.py b/python/python/tests/test_graph.py index deab745..a7b00db 100644 --- a/python/python/tests/test_graph.py +++ b/python/python/tests/test_graph.py @@ -45,10 +45,10 @@ def graph_env(tmp_path): config = ( GraphConfig.builder() - .with_node_label("Person", "person_id") - .with_node_label("Company", "company_id") - .with_relationship("WORKS_FOR", "person_id", "company_id") - .with_relationship("FRIEND_OF", "person1_id", "person2_id") + .node_table("Person", "person_id") + .node_table("Company", "company_id") + .rel_table("WORKS_FOR", "Person", "Company", "person_id", "company_id") + .rel_table("FRIEND_OF", "Person", "Person", "person1_id", "person2_id") .build() ) diff --git a/python/python/tests/test_scalar_functions.py b/python/python/tests/test_scalar_functions.py index e3e878a..83eb1b5 100644 --- a/python/python/tests/test_scalar_functions.py +++ b/python/python/tests/test_scalar_functions.py @@ -7,7 +7,7 @@ def test_unimplemented_scalar_function_errors() -> None: - cfg = GraphConfig.builder().with_node_label("Person", "name").build() + cfg = GraphConfig.builder().node_table("Person", "name").build() datasets = {"Person": pa.table({"name": ["Alice", "BOB", "CaSeY"]})} query = CypherQuery( diff --git a/python/python/tests/test_to_sql.py b/python/python/tests/test_to_sql.py index 1fad4d2..8806d8f 100644 --- a/python/python/tests/test_to_sql.py +++ b/python/python/tests/test_to_sql.py @@ -62,10 +62,10 @@ def knowledge_graph_env(): config = ( GraphConfig.builder() - .with_node_label("Author", "author_id") - .with_node_label("Paper", "paper_id") - .with_relationship("AUTHORED", "author_id", "paper_id") - .with_relationship("CITES", "citing_paper_id", "cited_paper_id") + .node_table("Author", "author_id") + .node_table("Paper", "paper_id") + .rel_table("AUTHORED", "Author", "Paper", "author_id", "paper_id") + .rel_table("CITES", "Paper", "Paper", "citing_paper_id", "cited_paper_id") .build() ) diff --git a/python/python/tests/test_vector_search.py b/python/python/tests/test_vector_search.py index ed4bde7..07c6601 100644 --- a/python/python/tests/test_vector_search.py +++ b/python/python/tests/test_vector_search.py @@ -40,7 +40,7 @@ def vector_env(): } ) - config = GraphConfig.builder().with_node_label("Document", "id").build() + config = GraphConfig.builder().node_table("Document", "id").build() datasets = {"Document": documents_table}