-
Notifications
You must be signed in to change notification settings - Fork 35
Expand file tree
/
Copy pathtest_sql.py
More file actions
175 lines (139 loc) · 5.55 KB
/
Copy pathtest_sql.py
File metadata and controls
175 lines (139 loc) · 5.55 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
175
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors
import pyarrow as pa
import pytest
from lance_graph import SqlEngine, SqlQuery
@pytest.fixture
def person_table():
return pa.table(
{
"id": [1, 2, 3, 4],
"name": ["Alice", "Bob", "Carol", "David"],
"age": [28, 34, 29, 42],
"city": ["New York", "San Francisco", "New York", "Chicago"],
}
)
@pytest.fixture
def knows_table():
return pa.table(
{
"src_id": [1, 1, 2, 3],
"dst_id": [2, 3, 4, 4],
"since_year": [2015, 2018, 2020, 2021],
}
)
@pytest.fixture
def datasets(person_table, knows_table):
return {"person": person_table, "knows": knows_table}
# ==========================================================================
# SqlQuery tests
# ==========================================================================
class TestSqlQuery:
def test_basic_select(self, datasets):
query = SqlQuery("SELECT name, age FROM person WHERE age > 30 ORDER BY age")
result = query.execute(datasets)
data = result.to_pydict()
assert data["name"] == ["Bob", "David"]
assert data["age"] == [34, 42]
def test_select_star(self, datasets):
query = SqlQuery("SELECT * FROM person")
result = query.execute(datasets)
assert result.num_rows == 4
assert result.num_columns == 4
def test_limit(self, datasets):
query = SqlQuery("SELECT name FROM person ORDER BY name LIMIT 2")
result = query.execute(datasets)
assert result.num_rows == 2
def test_join(self, datasets):
query = SqlQuery(
"SELECT p.name, k.dst_id "
"FROM person p "
"JOIN knows k ON p.id = k.src_id "
"ORDER BY p.name, k.dst_id"
)
result = query.execute(datasets)
data = result.to_pydict()
assert data["name"] == ["Alice", "Alice", "Bob", "Carol"]
def test_self_join(self, datasets):
query = SqlQuery(
"SELECT p1.name AS person, p2.name AS friend "
"FROM person p1 "
"JOIN knows k ON p1.id = k.src_id "
"JOIN person p2 ON p2.id = k.dst_id "
"ORDER BY p1.name, p2.name"
)
result = query.execute(datasets)
data = result.to_pydict()
assert data["person"] == ["Alice", "Alice", "Bob", "Carol"]
assert data["friend"] == ["Bob", "Carol", "David", "David"]
def test_count(self, datasets):
query = SqlQuery("SELECT COUNT(*) AS cnt FROM person")
result = query.execute(datasets)
assert result.to_pydict()["cnt"] == [4]
def test_sum(self, datasets):
query = SqlQuery("SELECT SUM(age) AS total FROM person")
result = query.execute(datasets)
assert result.to_pydict()["total"] == [28 + 34 + 29 + 42]
def test_avg(self, datasets):
query = SqlQuery("SELECT AVG(age) AS avg_age FROM person")
result = query.execute(datasets)
avg = result.to_pydict()["avg_age"][0]
assert abs(avg - 33.25) < 0.01
def test_group_by(self, datasets):
query = SqlQuery(
"SELECT city, COUNT(*) AS cnt "
"FROM person GROUP BY city ORDER BY cnt DESC, city"
)
result = query.execute(datasets)
data = result.to_pydict()
assert data["city"][0] == "New York"
assert data["cnt"][0] == 2
def test_explain(self, datasets):
query = SqlQuery("SELECT name FROM person WHERE age > 30")
plan = query.explain(datasets)
assert "Logical Plan" in plan
assert "Physical Plan" in plan
def test_sql_accessor(self):
query = SqlQuery("SELECT 1")
assert query.sql() == "SELECT 1"
def test_repr(self):
query = SqlQuery("SELECT 1")
assert "SqlQuery" in repr(query)
def test_invalid_sql(self, datasets):
query = SqlQuery("INVALID SQL")
with pytest.raises((RuntimeError, ValueError)):
query.execute(datasets)
def test_case_insensitive_table_names(self, person_table):
"""Table name 'Person' should be lowercased to 'person'."""
query = SqlQuery("SELECT name FROM person LIMIT 1")
result = query.execute({"Person": person_table})
assert result.num_rows == 1
# ==========================================================================
# SqlEngine tests
# ==========================================================================
class TestSqlEngine:
def test_basic_query(self, datasets):
engine = SqlEngine(datasets)
result = engine.execute(
"SELECT name, age FROM person WHERE age > 30 ORDER BY age"
)
data = result.to_pydict()
assert data["name"] == ["Bob", "David"]
def test_multiple_queries(self, datasets):
engine = SqlEngine(datasets)
r1 = engine.execute("SELECT COUNT(*) AS cnt FROM person")
r2 = engine.execute("SELECT name FROM person WHERE age > 30 ORDER BY name")
r3 = engine.execute(
"SELECT p.name, k.dst_id "
"FROM person p JOIN knows k ON p.id = k.src_id "
"ORDER BY p.name LIMIT 2"
)
assert r1.to_pydict()["cnt"] == [4]
assert r2.to_pydict()["name"] == ["Bob", "David"]
assert r3.num_rows == 2
def test_repr(self, datasets):
engine = SqlEngine(datasets)
assert "SqlEngine" in repr(engine)
def test_empty_datasets_raises(self):
with pytest.raises(ValueError, match="No input datasets"):
SqlEngine({})