From b03d11f176cfd354665f1424ae453a63159ea2e7 Mon Sep 17 00:00:00 2001 From: JoshuaTang <1240604020@qq.com> Date: Sun, 2 Nov 2025 17:15:56 -0800 Subject: [PATCH 1/2] feat: add datafusion pipeline python apis --- python/src/graph.rs | 68 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/python/src/graph.rs b/python/src/graph.rs index 35f2a9036..7914faee0 100644 --- a/python/src/graph.rs +++ b/python/src/graph.rs @@ -292,6 +292,74 @@ 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 { + // 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 + /// + /// 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 + /// ------- + /// str + /// Query graph logical plan, DataFusion logical plan, DataFusion physical plan as string + /// + /// Raises + /// ------ + /// ValueError + /// If the query is invalid or datasets are missing + /// RuntimeError + /// If query explain fails + fn explain_datafusion(&self, py: Python, datasets: &Bound<'_, PyDict>) -> PyResult { + // 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 plan = RT + .block_on(Some(py), inner_query.explain_datafusion(arrow_datasets))? + .map_err(graph_error_to_pyerr)?; + + Ok(plan) + } + /// Get variables used in the query fn variables(&self) -> Vec { self.inner.variables() From 38b86d930704fa7df5803a832c9f3fbed69df174 Mon Sep 17 00:00:00 2001 From: JoshuaTang <1240604020@qq.com> Date: Sun, 2 Nov 2025 17:16:18 -0800 Subject: [PATCH 2/2] test: add tests for datafusion-related python apis --- python/python/tests/test_explain.py | 73 +++++++++++++++++++++++ python/python/tests/test_graph.py | 89 ++++++++++++++++++++++------- 2 files changed, 141 insertions(+), 21 deletions(-) create mode 100644 python/python/tests/test_explain.py diff --git a/python/python/tests/test_explain.py b/python/python/tests/test_explain.py new file mode 100644 index 000000000..7cbd73baf --- /dev/null +++ b/python/python/tests/test_explain.py @@ -0,0 +1,73 @@ +"""Tests for explain_datafusion API.""" + +import pyarrow as pa +import pytest +from lance_graph import CypherQuery, GraphConfig + + +@pytest.fixture +def person_data(): + """Create simple Person dataset for testing.""" + people_table = pa.table( + { + "person_id": [1, 2, 3, 4], + "name": ["Alice", "Bob", "Carol", "David"], + "age": [28, 34, 29, 42], + } + ) + + config = GraphConfig.builder().with_node_label("Person", "person_id").build() + + return config, people_table + + +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}) + + # Verify the plan is a non-empty string + assert isinstance(plan, str) + assert len(plan) > 0 + + # Verify it contains expected sections + assert "Cypher Query:" in plan + assert "MATCH (p:Person) RETURN p.name, p.age" in plan + assert "graph_logical_plan" in plan + assert "logical_plan" in plan + assert "physical_plan" in plan + + # Verify table format + assert "+" in plan and "|" in plan + + +def test_explain_with_clauses(person_data): + """Test explain output includes query clauses (WHERE, ORDER BY, LIMIT).""" + config, people = 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}) + + assert isinstance(plan, str) + assert "WHERE p.age > 30" in plan + assert "ORDER BY" in plan + assert "LIMIT" in plan + + +def test_explain_error_handling(person_data): + """Test explain error handling for missing config and datasets.""" + config, people = 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}) + + # 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({}) diff --git a/python/python/tests/test_graph.py b/python/python/tests/test_graph.py index fda44b10f..a045abc61 100644 --- a/python/python/tests/test_graph.py +++ b/python/python/tests/test_graph.py @@ -62,37 +62,55 @@ def graph_env(tmp_path): return config, datasets, people_table -def test_basic_node_selection(graph_env): +@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"]) +def test_basic_node_selection(graph_env, execute_method): config, datasets, _ = graph_env query = CypherQuery("MATCH (p:Person) RETURN p.name, p.age").with_config(config) - result = query.execute({"Person": datasets["Person"]}) + result = getattr(query, execute_method)({"Person": datasets["Person"]}) data = result.to_pydict() - assert len(data["name"]) == 4 - assert set(data.keys()) == {"name", "age"} - assert "Alice" in set(data["name"]) - -def test_filtered_query(graph_env): + # TODO: remove this if/else statements when the execute() also returns + # Cypher dot notation + if execute_method == "execute": + # execute() returns unqualified names for simple queries + assert set(data.keys()) == {"name", "age"} + assert len(data["name"]) == 4 + assert "Alice" in set(data["name"]) + else: + # execute_datafusion() returns Cypher dot notation + 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): config, datasets, _ = graph_env query = CypherQuery( "MATCH (p:Person) WHERE p.age > 30 RETURN p.name, p.age" ).with_config(config) - result = query.execute({"Person": datasets["Person"]}) + result = getattr(query, execute_method)({"Person": datasets["Person"]}) data = result.to_pydict() - assert len(data["name"]) == 2 - assert set(data["name"]) == {"Bob", "David"} - assert all(age > 30 for age in data["age"]) + + if execute_method == "execute": + assert len(data["name"]) == 2 + assert set(data["name"]) == {"Bob", "David"} + assert all(age > 30 for age in data["age"]) + else: + assert len(data["p.name"]) == 2 + assert set(data["p.name"]) == {"Bob", "David"} + assert all(age > 30 for age in data["p.age"]) -def test_relationship_query(graph_env): +@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"]) +def test_relationship_query(graph_env, execute_method): config, datasets, _ = graph_env - # Alias outputs to stable column names regardless of internal qualification 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 = query.execute( + result = getattr(query, execute_method)( { "Person": datasets["Person"], "Company": datasets["Company"], @@ -105,7 +123,8 @@ def test_relationship_query(graph_env): assert data["company_id"] == [101, 101, 102, 103] -def test_friendship_direct_and_network(graph_env): +@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"]) +def test_friendship_direct_and_network(graph_env, execute_method): config, datasets, _ = graph_env # Direct friends of Alice (person_id = 1) query_direct = CypherQuery( @@ -114,7 +133,7 @@ def test_friendship_direct_and_network(graph_env): "RETURN b.person_id AS friend_id" ).with_config(config) - result_direct = query_direct.execute( + result_direct = getattr(query_direct, execute_method)( { "Person": datasets["Person"], "FRIEND_OF": datasets["FRIEND_OF"], @@ -129,7 +148,7 @@ def test_friendship_direct_and_network(graph_env): "RETURN f.person_id AS person1_id, t.person_id AS person2_id" ).with_config(config) - result_edges = query_edges.execute( + result_edges = getattr(query_edges, execute_method)( { "Person": datasets["Person"], "FRIEND_OF": datasets["FRIEND_OF"], @@ -140,7 +159,8 @@ def test_friendship_direct_and_network(graph_env): assert got == {(1, 2), (1, 3), (2, 4), (3, 4)} -def test_two_hop_friends_of_friends(graph_env): +@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"]) +def test_two_hop_friends_of_friends(graph_env, execute_method): config, datasets, _ = graph_env query = CypherQuery( "MATCH (a:Person)-[:FRIEND_OF]->(b:Person)-[:FRIEND_OF]->(c:Person) " @@ -148,7 +168,7 @@ def test_two_hop_friends_of_friends(graph_env): "RETURN a.person_id AS a_id, b.person_id AS b_id, c.person_id AS c_id" ).with_config(config) - result = query.execute( + result = getattr(query, execute_method)( { "Person": datasets["Person"], "FRIEND_OF": datasets["FRIEND_OF"], @@ -158,15 +178,42 @@ def test_two_hop_friends_of_friends(graph_env): assert set(data["c_id"]) == {4} -def test_variable_length_path(graph_env): +@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"]) +def test_variable_length_path(graph_env, execute_method): 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) - _ = query.execute( + _ = getattr(query, execute_method)( { "Person": datasets["Person"], "FRIEND_OF": datasets["FRIEND_OF"], } ) + + +@pytest.mark.parametrize("execute_method", ["execute", "execute_datafusion"]) +def test_distinct_clause(graph_env, execute_method): + 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)( + { + "Person": datasets["Person"], + "Company": datasets["Company"], + "WORKS_FOR": datasets["WORKS_FOR"], + } + ) + data = result.to_pydict() + + if execute_method == "execute": + # execute() returns qualified column names for relationship queries + assert len(data["c__company_name"]) == 3 + assert set(data["c__company_name"]) == {"TechCorp", "DataInc", "CloudSoft"} + else: + # execute_datafusion() returns Cypher dot notation + assert len(data["c.company_name"]) == 3 + assert set(data["c.company_name"]) == {"TechCorp", "DataInc", "CloudSoft"}