Repository navigation
Expand file tree
/
Copy pathtest_to_sql.py
More file actions
174 lines (146 loc) · 5.3 KB
/
Copy pathtest_to_sql.py
File metadata and controls
174 lines (146 loc) · 5.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors
"""Tests for the to_sql API that converts Cypher queries to SQL."""
import pyarrow as pa
import pytest
from lance_graph import CypherQuery, GraphConfig
@pytest.fixture
def knowledge_graph_env():
"""Create a complex knowledge graph with multiple entity types and relationships."""
# Authors and their publications
authors_table = pa.table(
{
"author_id": [1, 2, 3, 4, 5],
"name": [
"Alice Chen",
"Bob Smith",
"Carol Wang",
"David Lee",
"Eve Martinez",
],
"institution": ["MIT", "Stanford", "CMU", "Berkeley", "MIT"],
"h_index": [45, 38, 52, 41, 29],
"country": ["USA", "USA", "USA", "USA", "Spain"],
}
)
papers_table = pa.table(
{
"paper_id": [101, 102, 103, 104, 105, 106],
"title": [
"Deep Learning Advances",
"Graph Neural Networks",
"Transformer Architecture",
"Reinforcement Learning",
"Computer Vision Methods",
"Natural Language Processing",
],
"year": [2020, 2021, 2019, 2022, 2021, 2020],
"citations": [450, 320, 890, 210, 380, 520],
"venue": ["NeurIPS", "ICML", "NeurIPS", "ICLR", "CVPR", "ACL"],
}
)
authorship_table = pa.table(
{
"author_id": [1, 1, 2, 2, 3, 3, 4, 5, 5],
"paper_id": [101, 102, 102, 103, 103, 104, 105, 105, 106],
"author_position": [1, 1, 2, 1, 2, 1, 1, 2, 1],
}
)
citations_table = pa.table(
{
"citing_paper_id": [102, 103, 104, 104, 105, 106],
"cited_paper_id": [101, 101, 102, 103, 103, 101],
}
)
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")
.build()
)
datasets = {
"Author": authors_table,
"Paper": papers_table,
"AUTHORED": authorship_table,
"CITES": citations_table,
}
return config, datasets
def test_multi_hop_relationship_with_aggregation(knowledge_graph_env):
"""Test complex multi-hop query with aggregation and filtering.
Find authors who have written highly cited papers (>400 citations) and count
how many such papers they have, filtering for prolific authors.
"""
config, datasets = knowledge_graph_env
query = CypherQuery(
"""
MATCH (a:Author)-[:AUTHORED]->(p:Paper)
WHERE p.citations > 400
RETURN a.name, a.institution, COUNT(*) AS high_impact_papers
ORDER BY high_impact_papers DESC
"""
).with_config(config)
sql = query.to_sql(datasets)
assert isinstance(sql, str)
sql_upper = sql.upper()
assert "SELECT" in sql_upper
assert "JOIN" in sql_upper
assert "WHERE" in sql_upper
assert "COUNT" in sql_upper
assert "GROUP BY" in sql_upper
assert "ORDER BY" in sql_upper
def test_citation_network_analysis(knowledge_graph_env):
"""Test citation network traversal with multiple joins.
Find papers that cite other papers, along with author information,
filtered by venue and year range.
"""
config, datasets = knowledge_graph_env
query = CypherQuery(
"""
MATCH (citing:Paper)-[:CITES]->(cited:Paper)
WHERE citing.year >= 2020 AND citing.venue = 'NeurIPS'
RETURN citing.title, cited.title, citing.year, cited.citations
ORDER BY cited.citations DESC
LIMIT 10
"""
).with_config(config)
sql = query.to_sql(datasets)
assert isinstance(sql, str)
sql_upper = sql.upper()
assert "SELECT" in sql_upper
assert "JOIN" in sql_upper
assert "WHERE" in sql_upper
assert "ORDER BY" in sql_upper
assert "LIMIT" in sql_upper
def test_collaborative_network_query(knowledge_graph_env):
"""Test finding collaboration patterns through shared papers.
Find pairs of authors who have co-authored papers, with filtering
on institution and h-index.
"""
config, datasets = knowledge_graph_env
query = CypherQuery(
"""
MATCH (a1:Author)-[:AUTHORED]->(p:Paper)<-[:AUTHORED]-(a2:Author)
WHERE a1.author_id < a2.author_id
AND a1.institution = 'MIT'
AND a2.h_index > 30
RETURN DISTINCT a1.name, a2.name, p.title, p.year
ORDER BY p.year DESC
"""
).with_config(config)
sql = query.to_sql(datasets)
assert isinstance(sql, str)
sql_upper = sql.upper()
assert "SELECT" in sql_upper
# DISTINCT may be converted to GROUP BY by the SQL unparser
assert "DISTINCT" in sql_upper or "GROUP BY" in sql_upper
assert "JOIN" in sql_upper
assert "WHERE" in sql_upper
assert "ORDER BY" in sql_upper
def test_to_sql_without_config_raises_error(knowledge_graph_env):
"""Test that to_sql fails gracefully without config."""
_, datasets = knowledge_graph_env
query = CypherQuery("MATCH (a:Author) RETURN a.name")
with pytest.raises(Exception):
query.to_sql(datasets)