unify explain api and update tests

Author: ChunxuTang

python/python/tests/test_explain.py

@@ -1,4 +1,4 @@
-"""Tests for explain_datafusion API."""
+"""Tests for explain API."""
 
 import pyarrow as pa
 import pytest
@@ -25,7 +25,7 @@ def test_explain_simple_query(person_data):
     """Test explain output contains all expected sections."""
     config, people = person_data
     query = CypherQuery("MATCH (p:Person) RETURN p.name, p.age").with_config(config)
-    plan = query.explain_datafusion({"Person": people})
+    plan = query.explain({"Person": people})
 
     # Verify the plan is a non-empty string
     assert isinstance(plan, str)
@@ -48,7 +48,7 @@ def test_explain_with_clauses(person_data):
     query = CypherQuery(
         "MATCH (p:Person) WHERE p.age > 30 RETURN p.name ORDER BY p.age LIMIT 2"
     ).with_config(config)
-    plan = query.explain_datafusion({"Person": people})
+    plan = query.explain({"Person": people})
 
     assert isinstance(plan, str)
     assert "WHERE p.age > 30" in plan
@@ -63,11 +63,11 @@ def test_explain_error_handling(person_data):
     # Missing config
     query_no_config = CypherQuery("MATCH (p:Person) RETURN p.name")
     with pytest.raises(ValueError, match="Graph configuration is required"):
-        query_no_config.explain_datafusion({"Person": people})
+        query_no_config.explain({"Person": people})
 
     # Missing datasets
     query_with_config = CypherQuery("MATCH (p:Person) RETURN p.name").with_config(
         config
     )
     with pytest.raises(ValueError, match="No input datasets provided"):
-        query_with_config.explain_datafusion({})
+        query_with_config.explain({})

python/python/tests/test_graph.py

@@ -62,41 +62,38 @@ def graph_env(tmp_path):
     return config, datasets, people_table
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_basic_node_selection(graph_env, execute_method):
+def test_basic_node_selection(graph_env):
     config, datasets, _ = graph_env
     query = CypherQuery("MATCH (p:Person) RETURN p.name, p.age").with_config(config)
-    result = getattr(query, execute_method)({"Person": datasets["Person"]})
+    result = query.execute({"Person": datasets["Person"]})
     data = result.to_pydict()
 
     assert set(data.keys()) == {"p.name", "p.age"}
     assert len(data["p.name"]) == 4
     assert "Alice" in set(data["p.name"])
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_filtered_query(graph_env, execute_method):
+def test_filtered_query(graph_env):
     config, datasets, _ = graph_env
     query = CypherQuery(
         "MATCH (p:Person) WHERE p.age > 30 RETURN p.name, p.age"
     ).with_config(config)
-    result = getattr(query, execute_method)({"Person": datasets["Person"]})
+    result = query.execute({"Person": datasets["Person"]})
     data = result.to_pydict()
 
     assert len(data["p.name"]) == 2
     assert set(data["p.name"]) == {"Bob", "David"}
     assert all(age > 30 for age in data["p.age"])
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_relationship_query(graph_env, execute_method):
+def test_relationship_query(graph_env):
     config, datasets, _ = graph_env
     query = CypherQuery(
         "MATCH (p:Person)-[:WORKS_FOR]->(c:Company) "
         "RETURN p.person_id AS person_id, p.name AS name, c.company_id AS company_id"
     ).with_config(config)
 
-    result = getattr(query, execute_method)(
+    result = query.execute(
         {
             "Person": datasets["Person"],
             "Company": datasets["Company"],
@@ -109,8 +106,7 @@ def test_relationship_query(graph_env, execute_method):
     assert data["company_id"] == [101, 101, 102, 103]
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_friendship_direct_and_network(graph_env, execute_method):
+def test_friendship_direct_and_network(graph_env):
     config, datasets, _ = graph_env
     # Direct friends of Alice (person_id = 1)
     query_direct = CypherQuery(
@@ -119,7 +115,7 @@ def test_friendship_direct_and_network(graph_env, execute_method):
         "RETURN b.person_id AS friend_id"
     ).with_config(config)
 
-    result_direct = getattr(query_direct, execute_method)(
+    result_direct = query_direct.execute(
         {
             "Person": datasets["Person"],
             "FRIEND_OF": datasets["FRIEND_OF"],
@@ -134,7 +130,7 @@ def test_friendship_direct_and_network(graph_env, execute_method):
         "RETURN f.person_id AS person1_id, t.person_id AS person2_id"
     ).with_config(config)
 
-    result_edges = getattr(query_edges, execute_method)(
+    result_edges = query_edges.execute(
         {
             "Person": datasets["Person"],
             "FRIEND_OF": datasets["FRIEND_OF"],
@@ -145,16 +141,15 @@ def test_friendship_direct_and_network(graph_env, execute_method):
     assert got == {(1, 2), (1, 3), (2, 4), (3, 4)}
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_two_hop_friends_of_friends(graph_env, execute_method):
+def test_two_hop_friends_of_friends(graph_env):
     config, datasets, _ = graph_env
     query = CypherQuery(
         "MATCH (a:Person)-[:FRIEND_OF]->(b:Person)-[:FRIEND_OF]->(c:Person) "
         "WHERE a.person_id = 1 "
         "RETURN a.person_id AS a_id, b.person_id AS b_id, c.person_id AS c_id"
     ).with_config(config)
 
-    result = getattr(query, execute_method)(
+    result = query.execute(
         {
             "Person": datasets["Person"],
             "FRIEND_OF": datasets["FRIEND_OF"],
@@ -164,29 +159,31 @@ def test_two_hop_friends_of_friends(graph_env, execute_method):
     assert set(data["c_id"]) == {4}
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_variable_length_path(graph_env, execute_method):
+def test_variable_length_path(graph_env):
     config, datasets, _ = graph_env
     query = CypherQuery(
         "MATCH (p1:Person)-[:FRIEND_OF*1..2]-(p2:Person) "
         "RETURN p1.person_id AS p1, p2.person_id AS p2"
     ).with_config(config)
-    _ = getattr(query, execute_method)(
+
+    result = query.execute(
         {
             "Person": datasets["Person"],
             "FRIEND_OF": datasets["FRIEND_OF"],
         }
     )
+    data = result.to_pydict()
+    got = set(zip(data["p1"], data["p2"]))
+    assert got == {(1, 2), (1, 3), (2, 4), (3, 4), (1, 4)}
 
 
-@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"])
-def test_distinct_clause(graph_env, execute_method):
+def test_distinct_clause(graph_env):
     config, datasets, _ = graph_env
     query = CypherQuery(
         "MATCH (p:Person)-[:WORKS_FOR]->(c:Company) RETURN DISTINCT c.company_name"
     ).with_config(config)
 
-    result = getattr(query, execute_method)(
+    result = query.execute(
         {
             "Person": datasets["Person"],
             "Company": datasets["Company"],

python/src/graph.rs

@@ -335,41 +335,7 @@ impl CypherQuery {
         record_batch_to_python_table(py, &result_batch)
     }
 
-    /// Execute query using the DataFusion planner with in-memory datasets
-    ///
-    /// Parameters
-    /// ----------
-    /// datasets : dict
-    ///     Dictionary mapping table names to in-memory tables (pyarrow.Table, LanceDataset, etc.)
-    ///     Keys should match node labels and relationship types in the graph config.
-    ///
-    /// Returns
-    /// -------
-    /// pyarrow.Table
-    ///     Query results as Arrow table
-    ///
-    /// Raises
-    /// ------
-    /// ValueError
-    ///     If the query is invalid or datasets are missing
-    /// RuntimeError
-    ///     If query execution fails
-    fn execute_datafusion(&self, py: Python, datasets: &Bound<'_, PyDict>) -> PyResult<PyObject> {
-        // Convert datasets to Arrow RecordBatch map
-        let arrow_datasets = python_datasets_to_batches(datasets)?;
-
-        // Clone for async move
-        let inner_query = self.inner.clone();
-
-        // Execute via runtime
-        let result_batch = RT
-            .block_on(Some(py), inner_query.execute_datafusion(arrow_datasets))?
-            .map_err(graph_error_to_pyerr)?;
-
-        record_batch_to_python_table(py, &result_batch)
-    }
-
-    /// Explain query uusing the DataFusion planner with in-memory datasets
+    /// Explain query using the DataFusion planner with in-memory datasets
     ///
     /// Parameters
     /// ----------
@@ -388,7 +354,7 @@ impl CypherQuery {
     ///     If the query is invalid or datasets are missing
     /// RuntimeError
     ///     If query explain fails
-    fn explain_datafusion(&self, py: Python, datasets: &Bound<'_, PyDict>) -> PyResult<String> {
+    fn explain(&self, py: Python, datasets: &Bound<'_, PyDict>) -> PyResult<String> {
         // Convert datasets to Arrow RecordBatch map
         let arrow_datasets = python_datasets_to_batches(datasets)?;
 
@@ -397,7 +363,7 @@ impl CypherQuery {
 
         // Execute via runtime
         let plan = RT
-            .block_on(Some(py), inner_query.explain_datafusion(arrow_datasets))?
+            .block_on(Some(py), inner_query.explain(arrow_datasets))?
             .map_err(graph_error_to_pyerr)?;
 
         Ok(plan)