Skip to content

Commit 95010b4

Browse files
alexander-beediealexander-beedie
authored
Support use of QueryResult as a context manager (#3009)
1 parent b3ffc83 commit 95010b4

3 files changed

Lines changed: 135 additions & 46 deletions

File tree

‎src_py/query_result.py‎

Lines changed: 86 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,12 @@ def __init__(self, connection, query_result):
2323
self._query_result = query_result
2424
self.is_closed = False
2525

26+
def __enter__(self):
27+
return self
28+
29+
def __exit__(self, exc_type, exc_value, traceback):
30+
self.close()
31+
2632
def __del__(self):
2733
self.close()
2834

@@ -73,18 +79,22 @@ def close(self):
7379
Close the query result.
7480
"""
7581

76-
if self.is_closed:
77-
return
78-
self._query_result.close()
79-
# Allows the connection to be garbage collected if the query result
80-
# is closed manually by the user.
81-
self.connection = None
82-
self.is_closed = True
82+
if not self.is_closed:
83+
# Allows the connection to be garbage collected if the query result
84+
# is closed manually by the user.
85+
self._query_result.close()
86+
self.connection = None
87+
self.is_closed = True
8388

8489
def get_as_df(self):
8590
"""
8691
Get the query result as a Pandas DataFrame.
8792
93+
See Also
94+
--------
95+
get_as_pl : Get the query result as a Polars DataFrame.
96+
get_as_arrow : Get the query result as a PyArrow Table.
97+
8898
Returns
8999
-------
90100
pandas.DataFrame
@@ -102,14 +112,26 @@ def get_as_pl(self):
102112
"""
103113
Get the query result as a Polars DataFrame.
104114
115+
See Also
116+
--------
117+
get_as_df : Get the query result as a Pandas DataFrame.
118+
get_as_arrow : Get the query result as a PyArrow Table.
119+
105120
Returns
106121
-------
107122
polars.DataFrame
108123
Query result as a Polars DataFrame.
109124
"""
110125

111126
import polars as pl
112-
return pl.from_arrow(data=self.get_as_arrow(10_000))
127+
128+
target_n_elems = (
129+
10_000_000 # adaptive chunk_size; target 10m elements per chunk
130+
)
131+
target_chunk_size = max(target_n_elems // len(self.get_column_names()), 10)
132+
return pl.from_arrow(
133+
data=self.get_as_arrow(chunk_size=target_chunk_size),
134+
)
113135

114136
def get_as_arrow(self, chunk_size):
115137
"""
@@ -120,6 +142,11 @@ def get_as_arrow(self, chunk_size):
120142
chunk_size : int
121143
Number of rows to include in each chunk.
122144
145+
See Also
146+
--------
147+
get_as_pl : Get the query result as a Polars DataFrame.
148+
get_as_df : Get the query result as a Pandas DataFrame.
149+
123150
Returns
124151
-------
125152
pyarrow.Table
@@ -159,6 +186,26 @@ def get_column_names(self):
159186
self.check_for_query_result_close()
160187
return self._query_result.getColumnNames()
161188

189+
def get_schema(self):
190+
"""
191+
Get the column schema of the query result.
192+
193+
Returns
194+
-------
195+
dict
196+
Schema of the query result.
197+
198+
"""
199+
200+
self.check_for_query_result_close()
201+
return {
202+
name: dtype
203+
for name, dtype in zip(
204+
self._query_result.getColumnNames(),
205+
self._query_result.getColumnDataTypes(),
206+
)
207+
}
208+
162209
def reset_iterator(self):
163210
"""
164211
Reset the iterator of the query result.
@@ -203,7 +250,9 @@ def get_as_networkx(self, directed=True):
203250
table_primary_key_dict = {}
204251

205252
def encode_node_id(node, table_primary_key_dict):
206-
return node['_label'] + "_" + str(node[table_primary_key_dict[node['_label']]])
253+
return (
254+
node["_label"] + "_" + str(node[table_primary_key_dict[node["_label"]]])
255+
)
207256

208257
# De-duplicate nodes and rels
209258
while self.has_next():
@@ -218,36 +267,42 @@ def encode_node_id(node, table_primary_key_dict):
218267
elif column_type == Type.REL.value:
219268
_src = row[i]["_src"]
220269
_dst = row[i]["_dst"]
221-
rels[(_src["table"], _src["offset"], _dst["table"],
222-
_dst["offset"])] = row[i]
270+
rels[
271+
(_src["table"], _src["offset"], _dst["table"], _dst["offset"])
272+
] = row[i]
223273

224274
elif column_type == Type.RECURSIVE_REL.value:
225-
for node in row[i]['_nodes']:
275+
for node in row[i]["_nodes"]:
226276
_id = node["_id"]
227277
nodes[(_id["table"], _id["offset"])] = node
228278
table_to_label_dict[_id["table"]] = node["_label"]
229-
for rel in row[i]['_rels']:
279+
for rel in row[i]["_rels"]:
230280
for key in rel:
231281
if rel[key] is None:
232282
del rel[key]
233283
_src = rel["_src"]
234284
_dst = rel["_dst"]
235-
rels[(_src["table"], _src["offset"], _dst["table"],
236-
_dst["offset"])] = rel
285+
rels[
286+
(
287+
_src["table"],
288+
_src["offset"],
289+
_dst["table"],
290+
_dst["offset"],
291+
)
292+
] = rel
237293

238294
# Add nodes
239295
for node in nodes.values():
240296
_id = node["_id"]
241-
node_id = node['_label'] + "_" + str(_id["offset"])
242-
if node['_label'] not in table_primary_key_dict:
243-
props = self.connection._get_node_property_names(
244-
node['_label'])
297+
node_id = node["_label"] + "_" + str(_id["offset"])
298+
if node["_label"] not in table_primary_key_dict:
299+
props = self.connection._get_node_property_names(node["_label"])
245300
for prop_name in props:
246-
if props[prop_name]['is_primary_key']:
247-
table_primary_key_dict[node['_label']] = prop_name
301+
if props[prop_name]["is_primary_key"]:
302+
table_primary_key_dict[node["_label"]] = prop_name
248303
break
249304
node_id = encode_node_id(node, table_primary_key_dict)
250-
node[node['_label']] = True
305+
node[node["_label"]] = True
251306
nx_graph.add_node(node_id, **node)
252307

253308
# Add rels
@@ -270,7 +325,11 @@ def _get_properties_to_extract(self):
270325
for i in range(len(column_names)):
271326
column_name = column_names[i]
272327
column_type = column_types[i]
273-
if column_type in [Type.NODE.value, Type.REL.value, Type.RECURSIVE_REL.value]:
328+
if column_type in [
329+
Type.NODE.value,
330+
Type.REL.value,
331+
Type.RECURSIVE_REL.value,
332+
]:
274333
properties_to_extract[i] = (column_type, column_name)
275334
return properties_to_extract
276335

@@ -280,7 +339,7 @@ def get_as_torch_geometric(self):
280339
torch_geometric.data.Data or torch_geometric.data.HeteroData.
281340
282341
For node conversion, numerical and boolean properties are directly converted into tensor and
283-
stored in Data/HeteroData. For properties cannot be converted into tensor automatically
342+
stored in Data/HeteroData. For properties cannot be converted into tensor automatically
284343
(please refer to the notes below for more detail), they are returned as unconverted_properties.
285344
286345
For rel conversion, rel is converted into edge_index tensor director. Edge properties are returned
@@ -290,9 +349,9 @@ def get_as_torch_geometric(self):
290349
- If the type of a node property is not one of INT64, DOUBLE, or BOOL, it cannot be converted
291350
automatically.
292351
- If a node property contains a null value, it cannot be converted automatically.
293-
- If a node property contains a nested list of variable length (e.g. [[1,2],[3]]), it cannot be
352+
- If a node property contains a nested list of variable length (e.g. [[1,2],[3]]), it cannot be
294353
converted automatically.
295-
- If a node property is a list or nested list, but the shape is inconsistent (e.g. the list length
354+
- If a node property is a list or nested list, but the shape is inconsistent (e.g. the list length
296355
is 6 for one node but 5 for another node), it cannot be converted automatically.
297356
298357
Additional conversion rules:
@@ -363,7 +422,7 @@ def get_num_tuples(self):
363422
-------
364423
int
365424
Number of tuples.
366-
425+
367426
"""
368427
self.check_for_query_result_close()
369428
return self._query_result.getNumTuples()

‎test/test_get_header.py‎

Lines changed: 39 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,46 @@
11
def test_get_column_names(establish_connection):
22
conn, db = establish_connection
3-
result = conn.execute("MATCH (a:person)-[e:knows]->(b:person) RETURN a.fName, e.date, b.ID;")
4-
column_names = result.get_column_names()
5-
assert column_names[0] == 'a.fName'
6-
assert column_names[1] == 'e.date'
7-
assert column_names[2] == 'b.ID'
8-
result.close()
3+
with conn.execute(
4+
"MATCH (a:person)-[e:knows]->(b:person) RETURN a.fName, e.date, b.ID;"
5+
) as result:
6+
column_names = result.get_column_names()
7+
assert column_names[0] == 'a.fName'
8+
assert column_names[1] == 'e.date'
9+
assert column_names[2] == 'b.ID'
910

1011

1112
def test_get_column_data_types(establish_connection):
1213
conn, db = establish_connection
13-
result = conn.execute(
14+
with conn.execute(
1415
"MATCH (p:person) RETURN p.ID, p.fName, p.isStudent, p.eyeSight, p.birthdate, p.registerTime, "
15-
"p.lastJobDuration, p.workedHours, p.courseScoresPerTerm;")
16-
column_data_types = result.get_column_data_types()
17-
assert column_data_types[0] == 'INT64'
18-
assert column_data_types[1] == 'STRING'
19-
assert column_data_types[2] == 'BOOL'
20-
assert column_data_types[3] == 'DOUBLE'
21-
assert column_data_types[4] == 'DATE'
22-
assert column_data_types[5] == 'TIMESTAMP'
23-
assert column_data_types[6] == 'INTERVAL'
24-
assert column_data_types[7] == 'INT64[]'
25-
assert column_data_types[8] == 'INT64[][]'
26-
result.close()
16+
"p.lastJobDuration, p.workedHours, p.courseScoresPerTerm;"
17+
) as result:
18+
column_data_types = result.get_column_data_types()
19+
assert column_data_types[0] == 'INT64'
20+
assert column_data_types[1] == 'STRING'
21+
assert column_data_types[2] == 'BOOL'
22+
assert column_data_types[3] == 'DOUBLE'
23+
assert column_data_types[4] == 'DATE'
24+
assert column_data_types[5] == 'TIMESTAMP'
25+
assert column_data_types[6] == 'INTERVAL'
26+
assert column_data_types[7] == 'INT64[]'
27+
assert column_data_types[8] == 'INT64[][]'
28+
29+
30+
def test_get_schema(establish_connection):
31+
conn, db = establish_connection
32+
with conn.execute(
33+
"MATCH (p:person) RETURN p.ID, p.fName, p.isStudent, p.eyeSight, p.birthdate, p.registerTime, "
34+
"p.lastJobDuration, p.workedHours, p.courseScoresPerTerm;"
35+
) as result:
36+
assert result.get_schema() == {
37+
'p.ID': 'INT64',
38+
'p.fName': 'STRING',
39+
'p.isStudent': 'BOOL',
40+
'p.eyeSight': 'DOUBLE',
41+
'p.birthdate': 'DATE',
42+
'p.registerTime': 'TIMESTAMP',
43+
'p.lastJobDuration': 'INTERVAL',
44+
'p.workedHours': 'INT64[]',
45+
'p.courseScoresPerTerm': 'INT64[][]'
46+
}

‎test/test_query_result.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,13 @@ def test_explain(establish_connection):
2424
result = conn.execute("EXPLAIN MATCH (a:person) WHERE a.ID = 0 RETURN a")
2525
assert result.get_num_tuples() == 1
2626
result.close()
27+
28+
def test_context_manager(establish_connection):
29+
conn, db = establish_connection
30+
with conn.execute("MATCH (a:person) WHERE a.ID = 0 RETURN a") as result:
31+
assert result.get_num_tuples() == 1
32+
assert result.get_compiling_time() > 0
33+
34+
# context exit guarantees immediately 'close' of the underlying QueryResult
35+
# (don't have to wait for __del__, which may not ever actually get called)
36+
assert result.is_closed

0 commit comments

Comments
 (0)