kitedb 0.2.6__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- kitedb/__init__.py +273 -0
- kitedb/_kitedb.cpython-312-aarch64-linux-gnu.so +0 -0
- kitedb/_kitedb.pyi +677 -0
- kitedb/builders.py +901 -0
- kitedb/fluent.py +850 -0
- kitedb/schema.py +327 -0
- kitedb/traversal.py +1523 -0
- kitedb/vector_index.py +472 -0
- kitedb-0.2.6.dist-info/METADATA +216 -0
- kitedb-0.2.6.dist-info/RECORD +12 -0
- kitedb-0.2.6.dist-info/WHEEL +5 -0
- kitedb-0.2.6.dist-info/licenses/LICENSE +21 -0
kitedb/traversal.py
ADDED
|
@@ -0,0 +1,1523 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Traversal Builder for KiteDB
|
|
3
|
+
|
|
4
|
+
Provides a fluent API for graph traversals with lazy property loading.
|
|
5
|
+
|
|
6
|
+
By default, traversals load all properties. Use `.select([...])` or
|
|
7
|
+
`.load_props(...)` to load a subset for performance.
|
|
8
|
+
|
|
9
|
+
Example:
|
|
10
|
+
>>> # Default traversal - properties loaded
|
|
11
|
+
>>> friends = db.from_(alice).out(knows).to_list()
|
|
12
|
+
>>>
|
|
13
|
+
>>> # Load specific properties only
|
|
14
|
+
>>> friends = db.from_(alice).out(knows).select(["name", "age"]).to_list()
|
|
15
|
+
>>>
|
|
16
|
+
>>> # Filter requires properties - auto-loads them
|
|
17
|
+
>>> young_friends = (
|
|
18
|
+
... db.from_(alice)
|
|
19
|
+
... .out(knows)
|
|
20
|
+
... .where_node(lambda n: n.age is not None and n.age < 35)
|
|
21
|
+
... .to_list()
|
|
22
|
+
... )
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
from __future__ import annotations
|
|
26
|
+
|
|
27
|
+
from dataclasses import dataclass, field
|
|
28
|
+
from typing import (
|
|
29
|
+
TYPE_CHECKING,
|
|
30
|
+
Any,
|
|
31
|
+
Callable,
|
|
32
|
+
Dict,
|
|
33
|
+
Generator,
|
|
34
|
+
Generic,
|
|
35
|
+
Iterator,
|
|
36
|
+
List,
|
|
37
|
+
Literal,
|
|
38
|
+
Optional,
|
|
39
|
+
Set,
|
|
40
|
+
Tuple,
|
|
41
|
+
TypeVar,
|
|
42
|
+
Union,
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
from .builders import NodeRef, from_prop_value
|
|
46
|
+
from .schema import EdgeDef, NodeDef
|
|
47
|
+
|
|
48
|
+
if TYPE_CHECKING:
|
|
49
|
+
from kitedb._kitedb import Database
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
N = TypeVar("N", bound=NodeDef)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
# ============================================================================
|
|
56
|
+
# Traversal Step Types
|
|
57
|
+
# ============================================================================
|
|
58
|
+
|
|
59
|
+
@dataclass
|
|
60
|
+
class OutStep:
|
|
61
|
+
"""Traverse outgoing edges."""
|
|
62
|
+
type: Literal["out"] = "out"
|
|
63
|
+
edge_def: Optional[EdgeDef] = None
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
@dataclass
|
|
67
|
+
class InStep:
|
|
68
|
+
"""Traverse incoming edges."""
|
|
69
|
+
type: Literal["in"] = "in"
|
|
70
|
+
edge_def: Optional[EdgeDef] = None
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
@dataclass
|
|
74
|
+
class BothStep:
|
|
75
|
+
"""Traverse both directions."""
|
|
76
|
+
type: Literal["both"] = "both"
|
|
77
|
+
edge_def: Optional[EdgeDef] = None
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@dataclass
|
|
81
|
+
class TraverseOptions:
|
|
82
|
+
"""Options for variable-depth traversal."""
|
|
83
|
+
max_depth: int
|
|
84
|
+
min_depth: int = 1
|
|
85
|
+
direction: Literal["out", "in", "both"] = "out"
|
|
86
|
+
unique: bool = True
|
|
87
|
+
where_edge: Optional[Callable[["EdgeResult"], bool]] = None
|
|
88
|
+
where_node: Optional[Callable[[NodeRef[Any]], bool]] = None
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
@dataclass
|
|
92
|
+
class TraverseStep:
|
|
93
|
+
"""Variable-depth traversal step."""
|
|
94
|
+
type: Literal["traverse"] = "traverse"
|
|
95
|
+
edge_def: Optional[EdgeDef] = None
|
|
96
|
+
options: TraverseOptions = field(default_factory=lambda: TraverseOptions(max_depth=1))
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
TraversalStep = Union[OutStep, InStep, BothStep, TraverseStep]
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
# ============================================================================
|
|
103
|
+
# Property Loading Strategy
|
|
104
|
+
# ============================================================================
|
|
105
|
+
|
|
106
|
+
@dataclass
|
|
107
|
+
class PropLoadStrategy:
|
|
108
|
+
"""Strategy for loading properties."""
|
|
109
|
+
load_all: bool = False
|
|
110
|
+
prop_names: Optional[Set[str]] = None
|
|
111
|
+
|
|
112
|
+
@staticmethod
|
|
113
|
+
def none() -> PropLoadStrategy:
|
|
114
|
+
"""Don't load any properties."""
|
|
115
|
+
return PropLoadStrategy(load_all=False, prop_names=None)
|
|
116
|
+
|
|
117
|
+
@staticmethod
|
|
118
|
+
def all() -> PropLoadStrategy:
|
|
119
|
+
"""Load all properties."""
|
|
120
|
+
return PropLoadStrategy(load_all=True, prop_names=None)
|
|
121
|
+
|
|
122
|
+
@staticmethod
|
|
123
|
+
def only(*names: str) -> PropLoadStrategy:
|
|
124
|
+
"""Load only specified properties."""
|
|
125
|
+
return PropLoadStrategy(load_all=False, prop_names=set(names))
|
|
126
|
+
|
|
127
|
+
def should_load(self, prop_name: str) -> bool:
|
|
128
|
+
"""Check if a property should be loaded."""
|
|
129
|
+
if self.load_all:
|
|
130
|
+
return True
|
|
131
|
+
if self.prop_names is not None:
|
|
132
|
+
return prop_name in self.prop_names
|
|
133
|
+
return False
|
|
134
|
+
|
|
135
|
+
def needs_any_props(self) -> bool:
|
|
136
|
+
"""Check if any properties need to be loaded."""
|
|
137
|
+
return self.load_all or (self.prop_names is not None and len(self.prop_names) > 0)
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
# ============================================================================
|
|
141
|
+
# Edge Results
|
|
142
|
+
# ============================================================================
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
@dataclass
|
|
146
|
+
class EdgeResult:
|
|
147
|
+
"""Edge result with optional properties."""
|
|
148
|
+
src: int
|
|
149
|
+
etype: int
|
|
150
|
+
dst: int
|
|
151
|
+
props: Dict[str, Any] = field(default_factory=dict)
|
|
152
|
+
|
|
153
|
+
def __getattr__(self, name: str) -> Any:
|
|
154
|
+
props = object.__getattribute__(self, "props")
|
|
155
|
+
if name in props:
|
|
156
|
+
return props[name]
|
|
157
|
+
if name == "$src":
|
|
158
|
+
return self.src
|
|
159
|
+
if name == "$dst":
|
|
160
|
+
return self.dst
|
|
161
|
+
if name == "$etype":
|
|
162
|
+
return self.etype
|
|
163
|
+
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
|
|
164
|
+
|
|
165
|
+
def __getitem__(self, key: str) -> Any:
|
|
166
|
+
if key == "$src":
|
|
167
|
+
return self.src
|
|
168
|
+
if key == "$dst":
|
|
169
|
+
return self.dst
|
|
170
|
+
if key == "$etype":
|
|
171
|
+
return self.etype
|
|
172
|
+
props = object.__getattribute__(self, "props")
|
|
173
|
+
if key in props:
|
|
174
|
+
return props[key]
|
|
175
|
+
raise KeyError(key)
|
|
176
|
+
|
|
177
|
+
def to_dict(self) -> Dict[str, Any]:
|
|
178
|
+
data = {"$src": self.src, "$dst": self.dst, "$etype": self.etype}
|
|
179
|
+
data.update(self.props)
|
|
180
|
+
return data
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
@dataclass
|
|
184
|
+
class RawEdge:
|
|
185
|
+
"""Raw edge data without property loading."""
|
|
186
|
+
src: int
|
|
187
|
+
etype: int
|
|
188
|
+
dst: int
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
# ============================================================================
|
|
192
|
+
# Traversal Result
|
|
193
|
+
# ============================================================================
|
|
194
|
+
|
|
195
|
+
class TraversalResult(Generic[N]):
|
|
196
|
+
"""
|
|
197
|
+
Result of a traversal that can be iterated or collected.
|
|
198
|
+
|
|
199
|
+
This is a lazy iterator - it doesn't execute until you call
|
|
200
|
+
to_list(), first(), or iterate over it.
|
|
201
|
+
|
|
202
|
+
By default, all properties are loaded. Use `.select([...])` or
|
|
203
|
+
`.load_props("name", "age")` to load specific properties.
|
|
204
|
+
"""
|
|
205
|
+
|
|
206
|
+
def __init__(
|
|
207
|
+
self,
|
|
208
|
+
db: Database,
|
|
209
|
+
start_nodes: List[NodeRef[Any]],
|
|
210
|
+
steps: List[TraversalStep],
|
|
211
|
+
node_filter: Optional[Callable[[NodeRef[Any]], bool]],
|
|
212
|
+
edge_filter: Optional[Callable[[EdgeResult], bool]],
|
|
213
|
+
limit: Optional[int],
|
|
214
|
+
resolve_etype_id: Callable[[EdgeDef], int],
|
|
215
|
+
resolve_prop_key_id: Callable[[NodeDef, str], int],
|
|
216
|
+
get_node_def: Callable[[int], Optional[NodeDef]],
|
|
217
|
+
prop_strategy: PropLoadStrategy,
|
|
218
|
+
):
|
|
219
|
+
self._db = db
|
|
220
|
+
self._start_nodes = start_nodes
|
|
221
|
+
self._steps = steps
|
|
222
|
+
self._node_filter = node_filter
|
|
223
|
+
self._edge_filter = edge_filter
|
|
224
|
+
self._limit = limit
|
|
225
|
+
self._resolve_etype_id = resolve_etype_id
|
|
226
|
+
self._resolve_prop_key_id = resolve_prop_key_id
|
|
227
|
+
self._get_node_def = get_node_def
|
|
228
|
+
self._prop_strategy = prop_strategy
|
|
229
|
+
|
|
230
|
+
def _load_node_props(
|
|
231
|
+
self,
|
|
232
|
+
node_id: int,
|
|
233
|
+
node_def: NodeDef,
|
|
234
|
+
prop_strategy: Optional[PropLoadStrategy] = None,
|
|
235
|
+
) -> Dict[str, Any]:
|
|
236
|
+
"""Load properties for a node based on strategy using single FFI call."""
|
|
237
|
+
props: Dict[str, Any] = {}
|
|
238
|
+
|
|
239
|
+
strategy = prop_strategy or self._prop_strategy
|
|
240
|
+
|
|
241
|
+
if not strategy.needs_any_props():
|
|
242
|
+
return props
|
|
243
|
+
|
|
244
|
+
# Use get_node_props() for single FFI call instead of per-property calls
|
|
245
|
+
all_props = self._db.get_node_props(node_id)
|
|
246
|
+
if all_props is None:
|
|
247
|
+
return props
|
|
248
|
+
|
|
249
|
+
# Build reverse mapping: prop_key_id -> prop_name
|
|
250
|
+
key_id_to_name = {v: k for k, v in node_def._prop_key_ids.items()}
|
|
251
|
+
|
|
252
|
+
for node_prop in all_props:
|
|
253
|
+
prop_name = key_id_to_name.get(node_prop.key_id)
|
|
254
|
+
if prop_name is not None and strategy.should_load(prop_name):
|
|
255
|
+
props[prop_name] = from_prop_value(node_prop.value)
|
|
256
|
+
|
|
257
|
+
return props
|
|
258
|
+
|
|
259
|
+
def _load_edge_props(self, edge_def: Optional[EdgeDef], src: int, etype: int, dst: int) -> Dict[str, Any]:
|
|
260
|
+
"""Load edge properties using edge definition mapping."""
|
|
261
|
+
if edge_def is None or not edge_def.props:
|
|
262
|
+
return {}
|
|
263
|
+
|
|
264
|
+
props: Dict[str, Any] = {}
|
|
265
|
+
for prop_name in edge_def.props.keys():
|
|
266
|
+
try:
|
|
267
|
+
prop_key_id = self._resolve_prop_key_id(edge_def, prop_name)
|
|
268
|
+
except Exception:
|
|
269
|
+
continue
|
|
270
|
+
prop_value = self._db.get_edge_prop(src, etype, dst, prop_key_id)
|
|
271
|
+
if prop_value is not None:
|
|
272
|
+
props[prop_name] = from_prop_value(prop_value)
|
|
273
|
+
return props
|
|
274
|
+
|
|
275
|
+
def _build_edge_result(
|
|
276
|
+
self,
|
|
277
|
+
edge_def: Optional[EdgeDef],
|
|
278
|
+
src: int,
|
|
279
|
+
etype: int,
|
|
280
|
+
dst: int,
|
|
281
|
+
) -> EdgeResult:
|
|
282
|
+
props = self._load_edge_props(edge_def, src, etype, dst)
|
|
283
|
+
return EdgeResult(src=src, etype=etype, dst=dst, props=props)
|
|
284
|
+
|
|
285
|
+
def _create_node_ref(
|
|
286
|
+
self,
|
|
287
|
+
node_id: int,
|
|
288
|
+
load_props: bool = False,
|
|
289
|
+
prop_strategy: Optional[PropLoadStrategy] = None,
|
|
290
|
+
) -> Optional[NodeRef[Any]]:
|
|
291
|
+
"""Create a NodeRef from a node ID."""
|
|
292
|
+
node_def = self._get_node_def(node_id)
|
|
293
|
+
if node_def is None:
|
|
294
|
+
return None
|
|
295
|
+
|
|
296
|
+
key = self._db.get_node_key(node_id)
|
|
297
|
+
if key is None:
|
|
298
|
+
key = f"node:{node_id}"
|
|
299
|
+
|
|
300
|
+
if load_props:
|
|
301
|
+
props = self._load_node_props(node_id, node_def, prop_strategy=prop_strategy)
|
|
302
|
+
else:
|
|
303
|
+
props = {}
|
|
304
|
+
|
|
305
|
+
return NodeRef(id=node_id, key=key, node_def=node_def, props=props)
|
|
306
|
+
|
|
307
|
+
def _create_node_ref_fast(self, node_id: int, node_def: NodeDef) -> NodeRef[Any]:
|
|
308
|
+
"""Create a minimal NodeRef without loading key or properties."""
|
|
309
|
+
return NodeRef(id=node_id, key="", node_def=node_def, props={})
|
|
310
|
+
|
|
311
|
+
def _has_traverse_step(self) -> bool:
|
|
312
|
+
return any(isinstance(step, TraverseStep) for step in self._steps)
|
|
313
|
+
|
|
314
|
+
def _needs_full_execution(self) -> bool:
|
|
315
|
+
return (
|
|
316
|
+
self._edge_filter is not None
|
|
317
|
+
or self._node_filter is not None
|
|
318
|
+
or self._prop_strategy.needs_any_props()
|
|
319
|
+
or self._limit is not None
|
|
320
|
+
or self._has_traverse_step()
|
|
321
|
+
)
|
|
322
|
+
|
|
323
|
+
def _needs_full_execution_for_scalar(self) -> bool:
|
|
324
|
+
return (
|
|
325
|
+
self._edge_filter is not None
|
|
326
|
+
or self._node_filter is not None
|
|
327
|
+
or self._limit is not None
|
|
328
|
+
or self._has_traverse_step()
|
|
329
|
+
)
|
|
330
|
+
|
|
331
|
+
def _execute_single_hop(
|
|
332
|
+
self,
|
|
333
|
+
node: NodeRef[Any],
|
|
334
|
+
step: Union[OutStep, InStep, BothStep],
|
|
335
|
+
) -> Generator[Tuple[NodeRef[Any], EdgeResult], None, None]:
|
|
336
|
+
"""Execute a single-hop step and yield (node, edge) pairs."""
|
|
337
|
+
directions: List[str]
|
|
338
|
+
if step.type == "both":
|
|
339
|
+
directions = ["out", "in"]
|
|
340
|
+
else:
|
|
341
|
+
directions = [step.type]
|
|
342
|
+
|
|
343
|
+
etype_id: Optional[int] = None
|
|
344
|
+
if step.edge_def is not None:
|
|
345
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
346
|
+
|
|
347
|
+
for direction in directions:
|
|
348
|
+
if direction == "out":
|
|
349
|
+
edges = self._db.get_out_edges(node.id)
|
|
350
|
+
for edge in edges:
|
|
351
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
352
|
+
continue
|
|
353
|
+
neighbor_id = edge.node_id
|
|
354
|
+
neighbor_ref = self._create_node_ref(neighbor_id, load_props=self._prop_strategy.needs_any_props())
|
|
355
|
+
if neighbor_ref is None:
|
|
356
|
+
continue
|
|
357
|
+
edge_result = self._build_edge_result(step.edge_def, node.id, edge.etype, neighbor_id)
|
|
358
|
+
yield neighbor_ref, edge_result
|
|
359
|
+
else:
|
|
360
|
+
edges = self._db.get_in_edges(node.id)
|
|
361
|
+
for edge in edges:
|
|
362
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
363
|
+
continue
|
|
364
|
+
neighbor_id = edge.node_id
|
|
365
|
+
neighbor_ref = self._create_node_ref(neighbor_id, load_props=self._prop_strategy.needs_any_props())
|
|
366
|
+
if neighbor_ref is None:
|
|
367
|
+
continue
|
|
368
|
+
edge_result = self._build_edge_result(step.edge_def, neighbor_id, edge.etype, node.id)
|
|
369
|
+
yield neighbor_ref, edge_result
|
|
370
|
+
|
|
371
|
+
def _iter_traverse_edges(
|
|
372
|
+
self,
|
|
373
|
+
node_id: int,
|
|
374
|
+
direction: Literal["out", "in", "both"],
|
|
375
|
+
edge_def: Optional[EdgeDef],
|
|
376
|
+
etype_id: Optional[int],
|
|
377
|
+
) -> Generator[Tuple[int, EdgeResult], None, None]:
|
|
378
|
+
"""Yield (neighbor_id, edge_result) for a traversal step."""
|
|
379
|
+
directions: List[str]
|
|
380
|
+
if direction == "both":
|
|
381
|
+
directions = ["out", "in"]
|
|
382
|
+
else:
|
|
383
|
+
directions = [direction]
|
|
384
|
+
|
|
385
|
+
for dir_ in directions:
|
|
386
|
+
if dir_ == "out":
|
|
387
|
+
edges = self._db.get_out_edges(node_id)
|
|
388
|
+
for edge in edges:
|
|
389
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
390
|
+
continue
|
|
391
|
+
neighbor_id = edge.node_id
|
|
392
|
+
edge_result = self._build_edge_result(edge_def, node_id, edge.etype, neighbor_id)
|
|
393
|
+
yield neighbor_id, edge_result
|
|
394
|
+
else:
|
|
395
|
+
edges = self._db.get_in_edges(node_id)
|
|
396
|
+
for edge in edges:
|
|
397
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
398
|
+
continue
|
|
399
|
+
neighbor_id = edge.node_id
|
|
400
|
+
edge_result = self._build_edge_result(edge_def, neighbor_id, edge.etype, node_id)
|
|
401
|
+
yield neighbor_id, edge_result
|
|
402
|
+
|
|
403
|
+
def _execute_traverse_filtered(
|
|
404
|
+
self,
|
|
405
|
+
node: NodeRef[Any],
|
|
406
|
+
step: TraverseStep,
|
|
407
|
+
etype_id: Optional[int],
|
|
408
|
+
) -> Generator[Tuple[NodeRef[Any], EdgeResult], None, None]:
|
|
409
|
+
"""Execute variable-depth traversal with filters applied during traversal."""
|
|
410
|
+
from collections import deque
|
|
411
|
+
|
|
412
|
+
options = step.options
|
|
413
|
+
node_filter = options.where_node
|
|
414
|
+
edge_filter = options.where_edge
|
|
415
|
+
|
|
416
|
+
prop_strategy = self._prop_strategy
|
|
417
|
+
if node_filter is not None:
|
|
418
|
+
prop_strategy = PropLoadStrategy.all()
|
|
419
|
+
|
|
420
|
+
load_props = prop_strategy.needs_any_props()
|
|
421
|
+
|
|
422
|
+
visited: Set[int] = set()
|
|
423
|
+
if options.unique:
|
|
424
|
+
visited.add(node.id)
|
|
425
|
+
|
|
426
|
+
queue = deque([(node, 0)])
|
|
427
|
+
|
|
428
|
+
while queue:
|
|
429
|
+
current, depth = queue.popleft()
|
|
430
|
+
if depth >= options.max_depth:
|
|
431
|
+
continue
|
|
432
|
+
|
|
433
|
+
for neighbor_id, edge_result in self._iter_traverse_edges(
|
|
434
|
+
current.id,
|
|
435
|
+
options.direction,
|
|
436
|
+
step.edge_def,
|
|
437
|
+
etype_id,
|
|
438
|
+
):
|
|
439
|
+
if options.unique and neighbor_id in visited:
|
|
440
|
+
continue
|
|
441
|
+
|
|
442
|
+
neighbor_ref = self._create_node_ref(
|
|
443
|
+
neighbor_id,
|
|
444
|
+
load_props=load_props,
|
|
445
|
+
prop_strategy=prop_strategy,
|
|
446
|
+
)
|
|
447
|
+
if neighbor_ref is None:
|
|
448
|
+
continue
|
|
449
|
+
|
|
450
|
+
if edge_filter is not None and not edge_filter(edge_result):
|
|
451
|
+
continue
|
|
452
|
+
if node_filter is not None and not node_filter(neighbor_ref):
|
|
453
|
+
continue
|
|
454
|
+
|
|
455
|
+
if options.unique:
|
|
456
|
+
visited.add(neighbor_id)
|
|
457
|
+
|
|
458
|
+
next_depth = depth + 1
|
|
459
|
+
if next_depth >= options.min_depth:
|
|
460
|
+
yield neighbor_ref, edge_result
|
|
461
|
+
|
|
462
|
+
if next_depth < options.max_depth:
|
|
463
|
+
queue.append((neighbor_ref, next_depth))
|
|
464
|
+
|
|
465
|
+
def _execute_traverse(
|
|
466
|
+
self,
|
|
467
|
+
node: NodeRef[Any],
|
|
468
|
+
step: TraverseStep,
|
|
469
|
+
) -> Generator[Tuple[NodeRef[Any], EdgeResult], None, None]:
|
|
470
|
+
"""Execute variable-depth traversal from a node."""
|
|
471
|
+
options = step.options
|
|
472
|
+
etype_id: Optional[int] = None
|
|
473
|
+
if step.edge_def is not None:
|
|
474
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
475
|
+
|
|
476
|
+
if options.where_node is not None or options.where_edge is not None:
|
|
477
|
+
yield from self._execute_traverse_filtered(node, step, etype_id)
|
|
478
|
+
return
|
|
479
|
+
|
|
480
|
+
results = self._db.traverse(
|
|
481
|
+
node.id,
|
|
482
|
+
options.max_depth,
|
|
483
|
+
etype_id,
|
|
484
|
+
options.min_depth,
|
|
485
|
+
options.direction,
|
|
486
|
+
options.unique,
|
|
487
|
+
)
|
|
488
|
+
|
|
489
|
+
for result in results:
|
|
490
|
+
node_ref = self._create_node_ref(
|
|
491
|
+
result.node_id,
|
|
492
|
+
load_props=self._prop_strategy.needs_any_props(),
|
|
493
|
+
)
|
|
494
|
+
if node_ref is None:
|
|
495
|
+
continue
|
|
496
|
+
|
|
497
|
+
if result.edge_src is None or result.edge_dst is None or result.edge_type is None:
|
|
498
|
+
continue
|
|
499
|
+
|
|
500
|
+
edge_result = self._build_edge_result(
|
|
501
|
+
step.edge_def,
|
|
502
|
+
int(result.edge_src),
|
|
503
|
+
int(result.edge_type),
|
|
504
|
+
int(result.edge_dst),
|
|
505
|
+
)
|
|
506
|
+
yield node_ref, edge_result
|
|
507
|
+
|
|
508
|
+
def _iter_results(self) -> Generator[Tuple[NodeRef[Any], Optional[EdgeResult]], None, None]:
|
|
509
|
+
"""Full execution path that yields (node, edge) results."""
|
|
510
|
+
current_results: List[Tuple[NodeRef[Any], Optional[EdgeResult]]] = [
|
|
511
|
+
(node, None) for node in self._start_nodes
|
|
512
|
+
]
|
|
513
|
+
|
|
514
|
+
for step in self._steps:
|
|
515
|
+
next_results: List[Tuple[NodeRef[Any], EdgeResult]] = []
|
|
516
|
+
for node, _ in current_results:
|
|
517
|
+
if isinstance(step, TraverseStep):
|
|
518
|
+
for result in self._execute_traverse(node, step):
|
|
519
|
+
next_results.append(result)
|
|
520
|
+
else:
|
|
521
|
+
for result in self._execute_single_hop(node, step):
|
|
522
|
+
next_results.append(result)
|
|
523
|
+
current_results = [(n, e) for n, e in next_results]
|
|
524
|
+
|
|
525
|
+
count = 0
|
|
526
|
+
for node, edge in current_results:
|
|
527
|
+
if edge is not None and self._edge_filter is not None:
|
|
528
|
+
if not self._edge_filter(edge):
|
|
529
|
+
continue
|
|
530
|
+
if self._node_filter is not None and not self._node_filter(node):
|
|
531
|
+
continue
|
|
532
|
+
if self._limit is not None and count >= self._limit:
|
|
533
|
+
break
|
|
534
|
+
yield node, edge
|
|
535
|
+
count += 1
|
|
536
|
+
|
|
537
|
+
def _execute_edges(self) -> Generator[EdgeResult, None, None]:
|
|
538
|
+
"""Execute traversal and yield edge results."""
|
|
539
|
+
for _, edge in self._iter_results():
|
|
540
|
+
if edge is not None:
|
|
541
|
+
yield edge
|
|
542
|
+
|
|
543
|
+
def _build_steps_for_rust(self) -> List[Tuple[str, Optional[int]]]:
|
|
544
|
+
"""Build step tuples for Rust traverse_multi call."""
|
|
545
|
+
rust_steps = []
|
|
546
|
+
for step in self._steps:
|
|
547
|
+
etype_id = None
|
|
548
|
+
if step.edge_def is not None:
|
|
549
|
+
etype_id = step.edge_def._etype_id
|
|
550
|
+
if etype_id is None:
|
|
551
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
552
|
+
rust_steps.append((step.type, etype_id))
|
|
553
|
+
return rust_steps
|
|
554
|
+
|
|
555
|
+
def _execute_fast(self) -> Generator[int, None, None]:
|
|
556
|
+
"""Execute traversal and yield only node IDs (fastest path)."""
|
|
557
|
+
if self._has_traverse_step():
|
|
558
|
+
for node, _ in self._iter_results():
|
|
559
|
+
yield node.id
|
|
560
|
+
return
|
|
561
|
+
if not self._steps:
|
|
562
|
+
for node in self._start_nodes:
|
|
563
|
+
yield node.id
|
|
564
|
+
return
|
|
565
|
+
|
|
566
|
+
# For single step, use direct call (lower overhead)
|
|
567
|
+
# For multi-step, use Rust batch traversal
|
|
568
|
+
if len(self._steps) == 1:
|
|
569
|
+
step = self._steps[0]
|
|
570
|
+
etype_id = None
|
|
571
|
+
if step.edge_def is not None:
|
|
572
|
+
etype_id = step.edge_def._etype_id
|
|
573
|
+
if etype_id is None:
|
|
574
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
575
|
+
|
|
576
|
+
visited: Set[int] = set()
|
|
577
|
+
for node in self._start_nodes:
|
|
578
|
+
if step.type == "out":
|
|
579
|
+
neighbor_ids = self._db.traverse_out(node.id, etype_id)
|
|
580
|
+
elif step.type == "in":
|
|
581
|
+
neighbor_ids = self._db.traverse_in(node.id, etype_id)
|
|
582
|
+
else: # both
|
|
583
|
+
out_ids = self._db.traverse_out(node.id, etype_id)
|
|
584
|
+
in_ids = self._db.traverse_in(node.id, etype_id)
|
|
585
|
+
neighbor_ids = list(set(out_ids) | set(in_ids))
|
|
586
|
+
|
|
587
|
+
for neighbor_id in neighbor_ids:
|
|
588
|
+
if neighbor_id not in visited:
|
|
589
|
+
visited.add(neighbor_id)
|
|
590
|
+
yield neighbor_id
|
|
591
|
+
else:
|
|
592
|
+
# Multi-step: use Rust batch traversal
|
|
593
|
+
start_ids = [node.id for node in self._start_nodes]
|
|
594
|
+
rust_steps = self._build_steps_for_rust()
|
|
595
|
+
results = self._db.traverse_multi(start_ids, rust_steps)
|
|
596
|
+
for node_id, _ in results:
|
|
597
|
+
yield node_id
|
|
598
|
+
|
|
599
|
+
def _execute_fast_with_keys(self) -> Generator[Tuple[int, str], None, None]:
|
|
600
|
+
"""Execute traversal and yield (node_id, key) pairs."""
|
|
601
|
+
if self._has_traverse_step():
|
|
602
|
+
for node, _ in self._iter_results():
|
|
603
|
+
yield (node.id, node.key)
|
|
604
|
+
return
|
|
605
|
+
# No steps - just yield start nodes
|
|
606
|
+
if not self._steps:
|
|
607
|
+
for node in self._start_nodes:
|
|
608
|
+
yield (node.id, node.key)
|
|
609
|
+
return
|
|
610
|
+
|
|
611
|
+
# For single step, use direct batch call (lower overhead)
|
|
612
|
+
if len(self._steps) == 1:
|
|
613
|
+
step = self._steps[0]
|
|
614
|
+
etype_id = None
|
|
615
|
+
if step.edge_def is not None:
|
|
616
|
+
etype_id = step.edge_def._etype_id
|
|
617
|
+
if etype_id is None:
|
|
618
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
619
|
+
|
|
620
|
+
visited: Set[int] = set()
|
|
621
|
+
for node in self._start_nodes:
|
|
622
|
+
if step.type == "out":
|
|
623
|
+
pairs = self._db.traverse_out_with_keys(node.id, etype_id)
|
|
624
|
+
elif step.type == "in":
|
|
625
|
+
pairs = self._db.traverse_in_with_keys(node.id, etype_id)
|
|
626
|
+
else: # both
|
|
627
|
+
out_pairs = self._db.traverse_out_with_keys(node.id, etype_id)
|
|
628
|
+
in_pairs = self._db.traverse_in_with_keys(node.id, etype_id)
|
|
629
|
+
seen = set()
|
|
630
|
+
pairs = []
|
|
631
|
+
for nid, key in out_pairs + in_pairs:
|
|
632
|
+
if nid not in seen:
|
|
633
|
+
seen.add(nid)
|
|
634
|
+
pairs.append((nid, key))
|
|
635
|
+
|
|
636
|
+
for neighbor_id, key in pairs:
|
|
637
|
+
if neighbor_id not in visited:
|
|
638
|
+
visited.add(neighbor_id)
|
|
639
|
+
yield (neighbor_id, key or f"node:{neighbor_id}")
|
|
640
|
+
else:
|
|
641
|
+
# Multi-step: use Rust batch traversal
|
|
642
|
+
start_ids = [node.id for node in self._start_nodes]
|
|
643
|
+
rust_steps = self._build_steps_for_rust()
|
|
644
|
+
results = self._db.traverse_multi(start_ids, rust_steps)
|
|
645
|
+
for node_id, key in results:
|
|
646
|
+
yield (node_id, key or f"node:{node_id}")
|
|
647
|
+
|
|
648
|
+
def _execute_fast_count(self) -> int:
|
|
649
|
+
"""Execute traversal and return just the count."""
|
|
650
|
+
if self._has_traverse_step():
|
|
651
|
+
return sum(1 for _ in self._iter_results())
|
|
652
|
+
if not self._steps:
|
|
653
|
+
return len(self._start_nodes)
|
|
654
|
+
|
|
655
|
+
# For single step, count directly
|
|
656
|
+
if len(self._steps) == 1:
|
|
657
|
+
step = self._steps[0]
|
|
658
|
+
etype_id = None
|
|
659
|
+
if step.edge_def is not None:
|
|
660
|
+
etype_id = step.edge_def._etype_id
|
|
661
|
+
if etype_id is None:
|
|
662
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
663
|
+
|
|
664
|
+
visited: Set[int] = set()
|
|
665
|
+
for node in self._start_nodes:
|
|
666
|
+
if step.type == "out":
|
|
667
|
+
neighbor_ids = self._db.traverse_out(node.id, etype_id)
|
|
668
|
+
elif step.type == "in":
|
|
669
|
+
neighbor_ids = self._db.traverse_in(node.id, etype_id)
|
|
670
|
+
else: # both
|
|
671
|
+
out_ids = self._db.traverse_out(node.id, etype_id)
|
|
672
|
+
in_ids = self._db.traverse_in(node.id, etype_id)
|
|
673
|
+
neighbor_ids = list(set(out_ids) | set(in_ids))
|
|
674
|
+
|
|
675
|
+
for neighbor_id in neighbor_ids:
|
|
676
|
+
visited.add(neighbor_id)
|
|
677
|
+
|
|
678
|
+
return len(visited)
|
|
679
|
+
else:
|
|
680
|
+
# Multi-step: use Rust batch traversal count
|
|
681
|
+
start_ids = [node.id for node in self._start_nodes]
|
|
682
|
+
rust_steps = self._build_steps_for_rust()
|
|
683
|
+
return self._db.traverse_multi_count(start_ids, rust_steps)
|
|
684
|
+
|
|
685
|
+
def _execute(self) -> Generator[NodeRef[Any], None, None]:
|
|
686
|
+
"""Execute the traversal and yield results."""
|
|
687
|
+
if not self._needs_full_execution():
|
|
688
|
+
# Ultra-fast path: use batch operation to get IDs + keys in one call
|
|
689
|
+
for node_id, key in self._execute_fast_with_keys():
|
|
690
|
+
node_def = self._get_node_def(node_id)
|
|
691
|
+
if node_def is not None:
|
|
692
|
+
yield NodeRef(id=node_id, key=key, node_def=node_def, props={})
|
|
693
|
+
return
|
|
694
|
+
|
|
695
|
+
for node_ref, _ in self._iter_results():
|
|
696
|
+
yield node_ref
|
|
697
|
+
|
|
698
|
+
def __iter__(self) -> Iterator[NodeRef[Any]]:
|
|
699
|
+
"""Iterate over the traversal results."""
|
|
700
|
+
return iter(self._execute())
|
|
701
|
+
|
|
702
|
+
def to_list(self) -> List[NodeRef[N]]:
|
|
703
|
+
"""
|
|
704
|
+
Execute the traversal and collect results into a list.
|
|
705
|
+
|
|
706
|
+
Returns:
|
|
707
|
+
List of NodeRef objects
|
|
708
|
+
"""
|
|
709
|
+
if not self._needs_full_execution():
|
|
710
|
+
results: List[NodeRef[N]] = []
|
|
711
|
+
for node_id, key in self._execute_fast_with_keys():
|
|
712
|
+
node_def = self._get_node_def(node_id)
|
|
713
|
+
if node_def is not None:
|
|
714
|
+
results.append(NodeRef(id=node_id, key=key, node_def=node_def, props={})) # type: ignore
|
|
715
|
+
return results
|
|
716
|
+
|
|
717
|
+
return list(self._execute()) # type: ignore
|
|
718
|
+
|
|
719
|
+
def first(self) -> Optional[NodeRef[N]]:
|
|
720
|
+
"""
|
|
721
|
+
Execute the traversal and return the first result.
|
|
722
|
+
|
|
723
|
+
Returns:
|
|
724
|
+
First NodeRef or None if no results
|
|
725
|
+
"""
|
|
726
|
+
for node in self._execute():
|
|
727
|
+
return node # type: ignore
|
|
728
|
+
return None
|
|
729
|
+
|
|
730
|
+
def count(self) -> int:
|
|
731
|
+
"""
|
|
732
|
+
Execute the traversal and count results.
|
|
733
|
+
|
|
734
|
+
This is optimized to not load properties when counting.
|
|
735
|
+
|
|
736
|
+
Returns:
|
|
737
|
+
Number of matching nodes
|
|
738
|
+
"""
|
|
739
|
+
if not self._needs_full_execution_for_scalar() and self._node_filter is None:
|
|
740
|
+
return self._execute_fast_count()
|
|
741
|
+
return sum(1 for _ in self._execute())
|
|
742
|
+
|
|
743
|
+
def ids(self) -> List[int]:
|
|
744
|
+
"""
|
|
745
|
+
Get just the node IDs (fastest possible).
|
|
746
|
+
|
|
747
|
+
Returns:
|
|
748
|
+
List of node IDs
|
|
749
|
+
"""
|
|
750
|
+
if not self._needs_full_execution_for_scalar():
|
|
751
|
+
return list(self._execute_fast())
|
|
752
|
+
return [node.id for node, _ in self._iter_results()]
|
|
753
|
+
|
|
754
|
+
def keys(self) -> List[str]:
|
|
755
|
+
"""
|
|
756
|
+
Get just the node keys.
|
|
757
|
+
|
|
758
|
+
Returns:
|
|
759
|
+
List of node keys
|
|
760
|
+
"""
|
|
761
|
+
if not self._needs_full_execution_for_scalar():
|
|
762
|
+
result: List[str] = []
|
|
763
|
+
for node_id in self._execute_fast():
|
|
764
|
+
key = self._db.get_node_key(node_id)
|
|
765
|
+
if key:
|
|
766
|
+
result.append(key)
|
|
767
|
+
return result
|
|
768
|
+
|
|
769
|
+
return [node.key for node, _ in self._iter_results() if node.key]
|
|
770
|
+
|
|
771
|
+
|
|
772
|
+
class EdgeTraversalResult:
|
|
773
|
+
"""Traversal result that yields edges."""
|
|
774
|
+
|
|
775
|
+
def __init__(self, traversal: TraversalResult[Any]):
|
|
776
|
+
self._traversal = traversal
|
|
777
|
+
|
|
778
|
+
def __iter__(self) -> Iterator[EdgeResult]:
|
|
779
|
+
return iter(self._traversal._execute_edges())
|
|
780
|
+
|
|
781
|
+
def to_list(self) -> List[EdgeResult]:
|
|
782
|
+
return list(self._traversal._execute_edges())
|
|
783
|
+
|
|
784
|
+
def first(self) -> Optional[EdgeResult]:
|
|
785
|
+
for edge in self._traversal._execute_edges():
|
|
786
|
+
return edge
|
|
787
|
+
return None
|
|
788
|
+
|
|
789
|
+
def count(self) -> int:
|
|
790
|
+
return sum(1 for _ in self._traversal._execute_edges())
|
|
791
|
+
|
|
792
|
+
|
|
793
|
+
# ============================================================================
|
|
794
|
+
# Traversal Builder
|
|
795
|
+
# ============================================================================
|
|
796
|
+
|
|
797
|
+
class TraversalBuilder(Generic[N]):
|
|
798
|
+
"""
|
|
799
|
+
Builder for graph traversals.
|
|
800
|
+
|
|
801
|
+
By default, traversals load all properties. Use `.select([...])` or
|
|
802
|
+
`.load_props(...)` to load a subset.
|
|
803
|
+
|
|
804
|
+
Example:
|
|
805
|
+
>>> # Default traversal
|
|
806
|
+
>>> friend_refs = db.from_(alice).out(knows).to_list()
|
|
807
|
+
>>>
|
|
808
|
+
>>> # Load specific properties only
|
|
809
|
+
>>> friends = db.from_(alice).out(knows).load_props("name").to_list()
|
|
810
|
+
>>>
|
|
811
|
+
>>> # Filter automatically loads properties
|
|
812
|
+
>>> young = db.from_(alice).out(knows).where_node(lambda n: n.age < 35).to_list()
|
|
813
|
+
"""
|
|
814
|
+
|
|
815
|
+
def __init__(
|
|
816
|
+
self,
|
|
817
|
+
db: Database,
|
|
818
|
+
start_nodes: List[NodeRef[Any]],
|
|
819
|
+
resolve_etype_id: Callable[[EdgeDef], int],
|
|
820
|
+
resolve_prop_key_id: Callable[[NodeDef, str], int],
|
|
821
|
+
get_node_def: Callable[[int], Optional[NodeDef]],
|
|
822
|
+
):
|
|
823
|
+
self._db = db
|
|
824
|
+
self._start_nodes = start_nodes
|
|
825
|
+
self._resolve_etype_id = resolve_etype_id
|
|
826
|
+
self._resolve_prop_key_id = resolve_prop_key_id
|
|
827
|
+
self._get_node_def = get_node_def
|
|
828
|
+
self._steps: List[TraversalStep] = []
|
|
829
|
+
self._node_filter: Optional[Callable[[NodeRef[Any]], bool]] = None
|
|
830
|
+
self._edge_filter: Optional[Callable[[EdgeResult], bool]] = None
|
|
831
|
+
self._limit: Optional[int] = None
|
|
832
|
+
self._prop_strategy: PropLoadStrategy = PropLoadStrategy.all()
|
|
833
|
+
|
|
834
|
+
def out(self, edge: Optional[EdgeDef] = None) -> TraversalBuilder[N]:
|
|
835
|
+
"""
|
|
836
|
+
Traverse outgoing edges.
|
|
837
|
+
|
|
838
|
+
Args:
|
|
839
|
+
edge: Optional edge definition to filter by type
|
|
840
|
+
|
|
841
|
+
Returns:
|
|
842
|
+
Self for chaining
|
|
843
|
+
"""
|
|
844
|
+
self._steps.append(OutStep(edge_def=edge))
|
|
845
|
+
return self
|
|
846
|
+
|
|
847
|
+
def in_(self, edge: Optional[EdgeDef] = None) -> TraversalBuilder[N]:
|
|
848
|
+
"""
|
|
849
|
+
Traverse incoming edges.
|
|
850
|
+
|
|
851
|
+
Note: Named `in_` because `in` is a Python reserved word.
|
|
852
|
+
|
|
853
|
+
Args:
|
|
854
|
+
edge: Optional edge definition to filter by type
|
|
855
|
+
|
|
856
|
+
Returns:
|
|
857
|
+
Self for chaining
|
|
858
|
+
"""
|
|
859
|
+
self._steps.append(InStep(edge_def=edge))
|
|
860
|
+
return self
|
|
861
|
+
|
|
862
|
+
def both(self, edge: Optional[EdgeDef] = None) -> TraversalBuilder[N]:
|
|
863
|
+
"""
|
|
864
|
+
Traverse both incoming and outgoing edges.
|
|
865
|
+
|
|
866
|
+
Args:
|
|
867
|
+
edge: Optional edge definition to filter by type
|
|
868
|
+
|
|
869
|
+
Returns:
|
|
870
|
+
Self for chaining
|
|
871
|
+
"""
|
|
872
|
+
self._steps.append(BothStep(edge_def=edge))
|
|
873
|
+
return self
|
|
874
|
+
|
|
875
|
+
def traverse(self, edge: EdgeDef, options: TraverseOptions) -> TraversalBuilder[N]:
|
|
876
|
+
"""
|
|
877
|
+
Variable-depth traversal.
|
|
878
|
+
|
|
879
|
+
Args:
|
|
880
|
+
edge: Edge definition to traverse
|
|
881
|
+
options: TraverseOptions (max_depth required)
|
|
882
|
+
"""
|
|
883
|
+
self._steps.append(TraverseStep(edge_def=edge, options=options))
|
|
884
|
+
return self
|
|
885
|
+
|
|
886
|
+
def with_props(self) -> TraversalBuilder[N]:
|
|
887
|
+
"""
|
|
888
|
+
Load all properties for traversed nodes.
|
|
889
|
+
|
|
890
|
+
This is the default behavior; use load_props/select to limit properties.
|
|
891
|
+
|
|
892
|
+
Returns:
|
|
893
|
+
Self for chaining
|
|
894
|
+
|
|
895
|
+
Example:
|
|
896
|
+
>>> friends = db.from_(alice).out(knows).with_props().to_list()
|
|
897
|
+
>>> for f in friends:
|
|
898
|
+
... print(f.name, f.email)
|
|
899
|
+
"""
|
|
900
|
+
self._prop_strategy = PropLoadStrategy.all()
|
|
901
|
+
return self
|
|
902
|
+
|
|
903
|
+
def load_props(self, *prop_names: str) -> TraversalBuilder[N]:
|
|
904
|
+
"""
|
|
905
|
+
Load only specific properties for traversed nodes.
|
|
906
|
+
|
|
907
|
+
This is faster than with_props() when you only need a few properties.
|
|
908
|
+
|
|
909
|
+
Args:
|
|
910
|
+
*prop_names: Names of properties to load
|
|
911
|
+
|
|
912
|
+
Returns:
|
|
913
|
+
Self for chaining
|
|
914
|
+
|
|
915
|
+
Example:
|
|
916
|
+
>>> friends = db.from_(alice).out(knows).load_props("name").to_list()
|
|
917
|
+
>>> for f in friends:
|
|
918
|
+
... print(f.name) # Available
|
|
919
|
+
... print(f.email) # Will be None
|
|
920
|
+
"""
|
|
921
|
+
self._prop_strategy = PropLoadStrategy.only(*prop_names)
|
|
922
|
+
return self
|
|
923
|
+
|
|
924
|
+
def select(self, props: List[str]) -> TraversalBuilder[N]:
|
|
925
|
+
"""
|
|
926
|
+
Select specific properties to load.
|
|
927
|
+
|
|
928
|
+
This mirrors the TypeScript `select([...])` behavior.
|
|
929
|
+
"""
|
|
930
|
+
self._prop_strategy = PropLoadStrategy.only(*props)
|
|
931
|
+
return self
|
|
932
|
+
|
|
933
|
+
def where_edge(self, predicate: Callable[[EdgeResult], bool]) -> TraversalBuilder[N]:
|
|
934
|
+
"""
|
|
935
|
+
Filter results by edge properties.
|
|
936
|
+
|
|
937
|
+
Args:
|
|
938
|
+
predicate: Function that returns True for edges to include
|
|
939
|
+
"""
|
|
940
|
+
self._edge_filter = predicate
|
|
941
|
+
return self
|
|
942
|
+
|
|
943
|
+
def take(self, limit: int) -> TraversalBuilder[N]:
|
|
944
|
+
"""Limit the number of results."""
|
|
945
|
+
self._limit = limit
|
|
946
|
+
return self
|
|
947
|
+
|
|
948
|
+
def where_node(self, predicate: Callable[[NodeRef[Any]], bool]) -> TraversalBuilder[N]:
|
|
949
|
+
"""
|
|
950
|
+
Filter nodes by a predicate.
|
|
951
|
+
|
|
952
|
+
Note: Using a filter will automatically load all properties
|
|
953
|
+
since the predicate may access any property.
|
|
954
|
+
|
|
955
|
+
Args:
|
|
956
|
+
predicate: Function that returns True for nodes to include
|
|
957
|
+
|
|
958
|
+
Returns:
|
|
959
|
+
Self for chaining
|
|
960
|
+
|
|
961
|
+
Example:
|
|
962
|
+
>>> young_friends = (
|
|
963
|
+
... db.from_(alice)
|
|
964
|
+
... .out(knows)
|
|
965
|
+
... .where_node(lambda n: n.age is not None and n.age < 35)
|
|
966
|
+
... .to_list()
|
|
967
|
+
... )
|
|
968
|
+
"""
|
|
969
|
+
self._node_filter = predicate
|
|
970
|
+
# Filter needs properties to work, so enable loading all
|
|
971
|
+
self._prop_strategy = PropLoadStrategy.all()
|
|
972
|
+
return self
|
|
973
|
+
|
|
974
|
+
def _build_result(self) -> TraversalResult[N]:
|
|
975
|
+
"""Build the traversal result."""
|
|
976
|
+
return TraversalResult(
|
|
977
|
+
db=self._db,
|
|
978
|
+
start_nodes=self._start_nodes,
|
|
979
|
+
steps=self._steps,
|
|
980
|
+
node_filter=self._node_filter,
|
|
981
|
+
edge_filter=self._edge_filter,
|
|
982
|
+
limit=self._limit,
|
|
983
|
+
resolve_etype_id=self._resolve_etype_id,
|
|
984
|
+
resolve_prop_key_id=self._resolve_prop_key_id,
|
|
985
|
+
get_node_def=self._get_node_def,
|
|
986
|
+
prop_strategy=self._prop_strategy,
|
|
987
|
+
)
|
|
988
|
+
|
|
989
|
+
def nodes(self) -> TraversalResult[N]:
|
|
990
|
+
"""
|
|
991
|
+
Return node results.
|
|
992
|
+
|
|
993
|
+
Returns:
|
|
994
|
+
TraversalResult that can be iterated or collected
|
|
995
|
+
"""
|
|
996
|
+
return self._build_result()
|
|
997
|
+
|
|
998
|
+
def edges(self) -> "EdgeTraversalResult":
|
|
999
|
+
"""Return edge results from the traversal."""
|
|
1000
|
+
return EdgeTraversalResult(self._build_result())
|
|
1001
|
+
|
|
1002
|
+
def raw_edges(self) -> Generator[RawEdge, None, None]:
|
|
1003
|
+
"""Return raw edge data without property loading."""
|
|
1004
|
+
if any(isinstance(step, TraverseStep) for step in self._steps):
|
|
1005
|
+
raise ValueError("raw_edges() does not support variable-depth traverse()")
|
|
1006
|
+
|
|
1007
|
+
current_ids = [node.id for node in self._start_nodes]
|
|
1008
|
+
|
|
1009
|
+
for step in self._steps:
|
|
1010
|
+
if isinstance(step, TraverseStep):
|
|
1011
|
+
raise ValueError("raw_edges() does not support variable-depth traverse()")
|
|
1012
|
+
|
|
1013
|
+
directions: List[str]
|
|
1014
|
+
if step.type == "both":
|
|
1015
|
+
directions = ["out", "in"]
|
|
1016
|
+
else:
|
|
1017
|
+
directions = [step.type]
|
|
1018
|
+
|
|
1019
|
+
etype_id: Optional[int] = None
|
|
1020
|
+
if step.edge_def is not None:
|
|
1021
|
+
etype_id = self._resolve_etype_id(step.edge_def)
|
|
1022
|
+
|
|
1023
|
+
next_ids: List[int] = []
|
|
1024
|
+
|
|
1025
|
+
for node_id in current_ids:
|
|
1026
|
+
for direction in directions:
|
|
1027
|
+
if direction == "out":
|
|
1028
|
+
edges = self._db.get_out_edges(node_id)
|
|
1029
|
+
for edge in edges:
|
|
1030
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
1031
|
+
continue
|
|
1032
|
+
yield RawEdge(src=node_id, etype=edge.etype, dst=edge.node_id)
|
|
1033
|
+
next_ids.append(edge.node_id)
|
|
1034
|
+
else:
|
|
1035
|
+
edges = self._db.get_in_edges(node_id)
|
|
1036
|
+
for edge in edges:
|
|
1037
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
1038
|
+
continue
|
|
1039
|
+
yield RawEdge(src=edge.node_id, etype=edge.etype, dst=node_id)
|
|
1040
|
+
next_ids.append(edge.node_id)
|
|
1041
|
+
|
|
1042
|
+
current_ids = next_ids
|
|
1043
|
+
|
|
1044
|
+
def to_list(self) -> List[NodeRef[N]]:
|
|
1045
|
+
"""
|
|
1046
|
+
Shortcut for .nodes().to_list()
|
|
1047
|
+
|
|
1048
|
+
Returns:
|
|
1049
|
+
List of NodeRef objects
|
|
1050
|
+
"""
|
|
1051
|
+
return self._build_result().to_list()
|
|
1052
|
+
|
|
1053
|
+
def first(self) -> Optional[NodeRef[N]]:
|
|
1054
|
+
"""
|
|
1055
|
+
Shortcut for .nodes().first()
|
|
1056
|
+
|
|
1057
|
+
Returns:
|
|
1058
|
+
First NodeRef or None
|
|
1059
|
+
"""
|
|
1060
|
+
return self._build_result().first()
|
|
1061
|
+
|
|
1062
|
+
def count(self) -> int:
|
|
1063
|
+
"""
|
|
1064
|
+
Shortcut for .nodes().count()
|
|
1065
|
+
|
|
1066
|
+
This is optimized to not load properties when counting
|
|
1067
|
+
(unless a filter is set).
|
|
1068
|
+
|
|
1069
|
+
Returns:
|
|
1070
|
+
Number of matching nodes
|
|
1071
|
+
"""
|
|
1072
|
+
return self._build_result().count()
|
|
1073
|
+
|
|
1074
|
+
def ids(self) -> List[int]:
|
|
1075
|
+
"""
|
|
1076
|
+
Get just the node IDs (fastest possible).
|
|
1077
|
+
|
|
1078
|
+
Returns:
|
|
1079
|
+
List of node IDs
|
|
1080
|
+
"""
|
|
1081
|
+
return self._build_result().ids()
|
|
1082
|
+
|
|
1083
|
+
def keys(self) -> List[str]:
|
|
1084
|
+
"""
|
|
1085
|
+
Get just the node keys.
|
|
1086
|
+
|
|
1087
|
+
Returns:
|
|
1088
|
+
List of node keys
|
|
1089
|
+
"""
|
|
1090
|
+
return self._build_result().keys()
|
|
1091
|
+
|
|
1092
|
+
|
|
1093
|
+
# ============================================================================
|
|
1094
|
+
# Pathfinding Builder (simplified version)
|
|
1095
|
+
# ============================================================================
|
|
1096
|
+
|
|
1097
|
+
@dataclass
|
|
1098
|
+
class PathResult(Generic[N]):
|
|
1099
|
+
"""
|
|
1100
|
+
Result of a pathfinding query.
|
|
1101
|
+
|
|
1102
|
+
Attributes:
|
|
1103
|
+
nodes: List of node references in the path
|
|
1104
|
+
edges: List of edges in the path
|
|
1105
|
+
found: Whether a path was found
|
|
1106
|
+
total_weight: Total path weight (for weighted paths)
|
|
1107
|
+
"""
|
|
1108
|
+
nodes: List[NodeRef[N]]
|
|
1109
|
+
found: bool
|
|
1110
|
+
total_weight: float = 0.0
|
|
1111
|
+
edges: List[EdgeResult] = field(default_factory=list)
|
|
1112
|
+
|
|
1113
|
+
@property
|
|
1114
|
+
def path(self) -> List[NodeRef[N]]:
|
|
1115
|
+
return self.nodes
|
|
1116
|
+
|
|
1117
|
+
@property
|
|
1118
|
+
def totalWeight(self) -> float:
|
|
1119
|
+
return self.total_weight
|
|
1120
|
+
|
|
1121
|
+
def __bool__(self) -> bool:
|
|
1122
|
+
return self.found
|
|
1123
|
+
|
|
1124
|
+
def __len__(self) -> int:
|
|
1125
|
+
return len(self.nodes)
|
|
1126
|
+
|
|
1127
|
+
|
|
1128
|
+
WeightSpec = Union[str, Callable[[EdgeResult], float]]
|
|
1129
|
+
|
|
1130
|
+
|
|
1131
|
+
class PathFindingBuilder(Generic[N]):
|
|
1132
|
+
"""
|
|
1133
|
+
Builder for pathfinding queries.
|
|
1134
|
+
|
|
1135
|
+
Example:
|
|
1136
|
+
>>> path = db.shortest_path(alice).to(bob).find()
|
|
1137
|
+
>>> if path:
|
|
1138
|
+
... for node in path.nodes:
|
|
1139
|
+
... print(node.key)
|
|
1140
|
+
"""
|
|
1141
|
+
|
|
1142
|
+
def __init__(
|
|
1143
|
+
self,
|
|
1144
|
+
db: Database,
|
|
1145
|
+
source: NodeRef[N],
|
|
1146
|
+
resolve_etype_id: Callable[[EdgeDef], int],
|
|
1147
|
+
resolve_prop_key_id: Callable[[NodeDef, str], int],
|
|
1148
|
+
get_node_def: Callable[[int], Optional[NodeDef]],
|
|
1149
|
+
):
|
|
1150
|
+
self._db = db
|
|
1151
|
+
self._source = source
|
|
1152
|
+
self._resolve_etype_id = resolve_etype_id
|
|
1153
|
+
self._resolve_prop_key_id = resolve_prop_key_id
|
|
1154
|
+
self._get_node_def = get_node_def
|
|
1155
|
+
self._targets: Optional[List[NodeRef[Any]]] = None
|
|
1156
|
+
self._edge_type: Optional[EdgeDef] = None
|
|
1157
|
+
self._max_depth: Optional[int] = None
|
|
1158
|
+
self._direction: str = "out"
|
|
1159
|
+
self._load_props: bool = True
|
|
1160
|
+
self._weight_spec: Optional[WeightSpec] = None
|
|
1161
|
+
|
|
1162
|
+
def to(self, target: NodeRef[Any]) -> PathFindingBuilder[N]:
|
|
1163
|
+
"""Set the target node."""
|
|
1164
|
+
self._targets = [target]
|
|
1165
|
+
return self
|
|
1166
|
+
|
|
1167
|
+
def to_any(self, targets: List[NodeRef[Any]]) -> PathFindingBuilder[N]:
|
|
1168
|
+
"""Set multiple target nodes (find path to any)."""
|
|
1169
|
+
if not targets:
|
|
1170
|
+
raise ValueError("to_any requires at least one target")
|
|
1171
|
+
self._targets = targets
|
|
1172
|
+
return self
|
|
1173
|
+
|
|
1174
|
+
def via(self, edge: EdgeDef) -> PathFindingBuilder[N]:
|
|
1175
|
+
"""Filter by edge type."""
|
|
1176
|
+
self._edge_type = edge
|
|
1177
|
+
return self
|
|
1178
|
+
|
|
1179
|
+
def max_depth(self, depth: int) -> PathFindingBuilder[N]:
|
|
1180
|
+
"""Set maximum path length."""
|
|
1181
|
+
self._max_depth = depth
|
|
1182
|
+
return self
|
|
1183
|
+
|
|
1184
|
+
def direction(self, dir: Literal["out", "in", "both"]) -> PathFindingBuilder[N]:
|
|
1185
|
+
"""Set traversal direction."""
|
|
1186
|
+
self._direction = dir
|
|
1187
|
+
return self
|
|
1188
|
+
|
|
1189
|
+
def with_props(self) -> PathFindingBuilder[N]:
|
|
1190
|
+
"""Load properties for nodes in the path (default behavior)."""
|
|
1191
|
+
self._load_props = True
|
|
1192
|
+
return self
|
|
1193
|
+
|
|
1194
|
+
def weight(self, spec: WeightSpec) -> PathFindingBuilder[N]:
|
|
1195
|
+
"""Set weight specification (property name or function)."""
|
|
1196
|
+
self._weight_spec = spec
|
|
1197
|
+
return self
|
|
1198
|
+
|
|
1199
|
+
def _create_node_ref(self, node_id: int) -> Optional[NodeRef[Any]]:
|
|
1200
|
+
"""Create a NodeRef from a node ID."""
|
|
1201
|
+
node_def = self._get_node_def(node_id)
|
|
1202
|
+
if node_def is None:
|
|
1203
|
+
return None
|
|
1204
|
+
|
|
1205
|
+
key = self._db.get_node_key(node_id)
|
|
1206
|
+
if key is None:
|
|
1207
|
+
key = f"node:{node_id}"
|
|
1208
|
+
|
|
1209
|
+
props: Dict[str, Any] = {}
|
|
1210
|
+
if self._load_props:
|
|
1211
|
+
for prop_name, prop_def in node_def.props.items():
|
|
1212
|
+
prop_key_id = self._resolve_prop_key_id(node_def, prop_name)
|
|
1213
|
+
prop_value = self._db.get_node_prop(node_id, prop_key_id)
|
|
1214
|
+
if prop_value is not None:
|
|
1215
|
+
props[prop_name] = from_prop_value(prop_value)
|
|
1216
|
+
|
|
1217
|
+
return NodeRef(id=node_id, key=key, node_def=node_def, props=props)
|
|
1218
|
+
|
|
1219
|
+
def _get_targets(self) -> List[NodeRef[Any]]:
|
|
1220
|
+
if not self._targets:
|
|
1221
|
+
raise ValueError("Target node required. Use .to(target) or .to_any(targets) first.")
|
|
1222
|
+
if self._edge_type is None:
|
|
1223
|
+
raise ValueError("Must specify at least one edge type with via()")
|
|
1224
|
+
return self._targets
|
|
1225
|
+
|
|
1226
|
+
def _max_depth_value(self) -> int:
|
|
1227
|
+
return self._max_depth if self._max_depth is not None else 100
|
|
1228
|
+
|
|
1229
|
+
def _build_edge_result(self, edge_def: Optional[EdgeDef], src: int, etype: int, dst: int) -> EdgeResult:
|
|
1230
|
+
props: Dict[str, Any] = {}
|
|
1231
|
+
if edge_def is not None and edge_def.props:
|
|
1232
|
+
for prop_name in edge_def.props.keys():
|
|
1233
|
+
try:
|
|
1234
|
+
prop_key_id = self._resolve_prop_key_id(edge_def, prop_name)
|
|
1235
|
+
except Exception:
|
|
1236
|
+
continue
|
|
1237
|
+
prop_value = self._db.get_edge_prop(src, etype, dst, prop_key_id)
|
|
1238
|
+
if prop_value is not None:
|
|
1239
|
+
props[prop_name] = from_prop_value(prop_value)
|
|
1240
|
+
return EdgeResult(src=src, etype=etype, dst=dst, props=props)
|
|
1241
|
+
|
|
1242
|
+
def _coerce_weight(self, value: Any) -> float:
|
|
1243
|
+
try:
|
|
1244
|
+
weight = float(value)
|
|
1245
|
+
except Exception:
|
|
1246
|
+
return 1.0
|
|
1247
|
+
if not weight or weight <= 0:
|
|
1248
|
+
return 1.0
|
|
1249
|
+
return weight
|
|
1250
|
+
|
|
1251
|
+
def _edge_weight(self, edge: EdgeResult) -> float:
|
|
1252
|
+
if self._weight_spec is None:
|
|
1253
|
+
return 1.0
|
|
1254
|
+
if isinstance(self._weight_spec, str):
|
|
1255
|
+
prop_name = self._weight_spec
|
|
1256
|
+
if prop_name in edge.props:
|
|
1257
|
+
return self._coerce_weight(edge.props[prop_name])
|
|
1258
|
+
return 1.0
|
|
1259
|
+
return self._coerce_weight(self._weight_spec(edge))
|
|
1260
|
+
|
|
1261
|
+
def _iter_neighbors(self, node_id: int) -> Generator[Tuple[int, EdgeResult], None, None]:
|
|
1262
|
+
directions: List[str]
|
|
1263
|
+
if self._direction == "both":
|
|
1264
|
+
directions = ["out", "in"]
|
|
1265
|
+
else:
|
|
1266
|
+
directions = [self._direction]
|
|
1267
|
+
|
|
1268
|
+
etype_id: Optional[int] = None
|
|
1269
|
+
if self._edge_type is not None:
|
|
1270
|
+
etype_id = self._resolve_etype_id(self._edge_type)
|
|
1271
|
+
|
|
1272
|
+
for direction in directions:
|
|
1273
|
+
if direction == "out":
|
|
1274
|
+
edges = self._db.get_out_edges(node_id)
|
|
1275
|
+
for edge in edges:
|
|
1276
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
1277
|
+
continue
|
|
1278
|
+
neighbor_id = edge.node_id
|
|
1279
|
+
edge_result = self._build_edge_result(self._edge_type, node_id, edge.etype, neighbor_id)
|
|
1280
|
+
yield neighbor_id, edge_result
|
|
1281
|
+
else:
|
|
1282
|
+
edges = self._db.get_in_edges(node_id)
|
|
1283
|
+
for edge in edges:
|
|
1284
|
+
if etype_id is not None and edge.etype != etype_id:
|
|
1285
|
+
continue
|
|
1286
|
+
neighbor_id = edge.node_id
|
|
1287
|
+
edge_result = self._build_edge_result(self._edge_type, neighbor_id, edge.etype, node_id)
|
|
1288
|
+
yield neighbor_id, edge_result
|
|
1289
|
+
|
|
1290
|
+
def _reconstruct_path(
|
|
1291
|
+
self,
|
|
1292
|
+
parents: Dict[int, Tuple[Optional[int], Optional[EdgeResult]]],
|
|
1293
|
+
target_id: int,
|
|
1294
|
+
) -> PathResult[N]:
|
|
1295
|
+
path_nodes: List[NodeRef[N]] = []
|
|
1296
|
+
path_edges: List[EdgeResult] = []
|
|
1297
|
+
|
|
1298
|
+
current: Optional[int] = target_id
|
|
1299
|
+
while current is not None:
|
|
1300
|
+
parent, edge = parents.get(current, (None, None))
|
|
1301
|
+
node_ref = self._create_node_ref(current)
|
|
1302
|
+
if node_ref is not None:
|
|
1303
|
+
path_nodes.append(node_ref) # type: ignore
|
|
1304
|
+
if edge is not None:
|
|
1305
|
+
path_edges.append(edge)
|
|
1306
|
+
current = parent
|
|
1307
|
+
|
|
1308
|
+
path_nodes.reverse()
|
|
1309
|
+
path_edges.reverse()
|
|
1310
|
+
|
|
1311
|
+
return PathResult(nodes=path_nodes, found=True, total_weight=0.0, edges=path_edges)
|
|
1312
|
+
|
|
1313
|
+
def find(self) -> PathResult[N]:
|
|
1314
|
+
"""
|
|
1315
|
+
Find the shortest path using BFS.
|
|
1316
|
+
|
|
1317
|
+
Returns:
|
|
1318
|
+
PathResult containing the path if found
|
|
1319
|
+
"""
|
|
1320
|
+
return self.bfs()
|
|
1321
|
+
|
|
1322
|
+
def find_weighted(self) -> PathResult[N]:
|
|
1323
|
+
"""
|
|
1324
|
+
Find the shortest weighted path using Dijkstra.
|
|
1325
|
+
|
|
1326
|
+
Returns:
|
|
1327
|
+
PathResult containing the path if found
|
|
1328
|
+
"""
|
|
1329
|
+
return self.dijkstra()
|
|
1330
|
+
|
|
1331
|
+
def exists(self) -> bool:
|
|
1332
|
+
"""
|
|
1333
|
+
Check if a path exists between source and target.
|
|
1334
|
+
|
|
1335
|
+
Returns:
|
|
1336
|
+
True if a path exists
|
|
1337
|
+
"""
|
|
1338
|
+
return self.bfs().found
|
|
1339
|
+
|
|
1340
|
+
def bfs(self) -> PathResult[N]:
|
|
1341
|
+
"""Execute BFS (unweighted shortest path)."""
|
|
1342
|
+
targets = self._get_targets()
|
|
1343
|
+
target_ids = {t.id for t in targets}
|
|
1344
|
+
max_depth = self._max_depth_value()
|
|
1345
|
+
|
|
1346
|
+
if self._source.id in target_ids:
|
|
1347
|
+
node_ref = self._create_node_ref(self._source.id)
|
|
1348
|
+
if node_ref is None:
|
|
1349
|
+
return PathResult(nodes=[], found=False)
|
|
1350
|
+
return PathResult(nodes=[node_ref], found=True, total_weight=0.0)
|
|
1351
|
+
|
|
1352
|
+
from collections import deque
|
|
1353
|
+
|
|
1354
|
+
queue = deque([(self._source.id, 0)])
|
|
1355
|
+
visited = {self._source.id}
|
|
1356
|
+
parents: Dict[int, Tuple[Optional[int], Optional[EdgeResult]]] = {
|
|
1357
|
+
self._source.id: (None, None)
|
|
1358
|
+
}
|
|
1359
|
+
|
|
1360
|
+
while queue:
|
|
1361
|
+
node_id, depth = queue.popleft()
|
|
1362
|
+
if depth >= max_depth:
|
|
1363
|
+
continue
|
|
1364
|
+
|
|
1365
|
+
for neighbor_id, edge in self._iter_neighbors(node_id):
|
|
1366
|
+
if neighbor_id in visited:
|
|
1367
|
+
continue
|
|
1368
|
+
visited.add(neighbor_id)
|
|
1369
|
+
parents[neighbor_id] = (node_id, edge)
|
|
1370
|
+
|
|
1371
|
+
if neighbor_id in target_ids:
|
|
1372
|
+
result = self._reconstruct_path(parents, neighbor_id)
|
|
1373
|
+
result.total_weight = float(len(result.edges))
|
|
1374
|
+
return result
|
|
1375
|
+
|
|
1376
|
+
queue.append((neighbor_id, depth + 1))
|
|
1377
|
+
|
|
1378
|
+
return PathResult(nodes=[], found=False)
|
|
1379
|
+
|
|
1380
|
+
def dijkstra(self) -> PathResult[N]:
|
|
1381
|
+
"""Execute Dijkstra's algorithm."""
|
|
1382
|
+
targets = self._get_targets()
|
|
1383
|
+
target_ids = {t.id for t in targets}
|
|
1384
|
+
max_depth = self._max_depth_value()
|
|
1385
|
+
|
|
1386
|
+
if isinstance(self._weight_spec, str) and self._edge_type is None:
|
|
1387
|
+
raise ValueError("weight by property requires via(edge)")
|
|
1388
|
+
|
|
1389
|
+
if self._source.id in target_ids:
|
|
1390
|
+
node_ref = self._create_node_ref(self._source.id)
|
|
1391
|
+
if node_ref is None:
|
|
1392
|
+
return PathResult(nodes=[], found=False)
|
|
1393
|
+
return PathResult(nodes=[node_ref], found=True, total_weight=0.0)
|
|
1394
|
+
|
|
1395
|
+
import heapq
|
|
1396
|
+
|
|
1397
|
+
dist: Dict[int, float] = {self._source.id: 0.0}
|
|
1398
|
+
depth_map: Dict[int, int] = {self._source.id: 0}
|
|
1399
|
+
parents: Dict[int, Tuple[Optional[int], Optional[EdgeResult]]] = {
|
|
1400
|
+
self._source.id: (None, None)
|
|
1401
|
+
}
|
|
1402
|
+
heap: List[Tuple[float, int, int]] = [(0.0, 0, self._source.id)]
|
|
1403
|
+
visited: Set[int] = set()
|
|
1404
|
+
|
|
1405
|
+
while heap:
|
|
1406
|
+
cost, depth, node_id = heapq.heappop(heap)
|
|
1407
|
+
if node_id in visited:
|
|
1408
|
+
continue
|
|
1409
|
+
visited.add(node_id)
|
|
1410
|
+
|
|
1411
|
+
if node_id in target_ids:
|
|
1412
|
+
result = self._reconstruct_path(parents, node_id)
|
|
1413
|
+
result.total_weight = cost
|
|
1414
|
+
return result
|
|
1415
|
+
|
|
1416
|
+
if depth >= max_depth:
|
|
1417
|
+
continue
|
|
1418
|
+
|
|
1419
|
+
for neighbor_id, edge in self._iter_neighbors(node_id):
|
|
1420
|
+
next_depth = depth + 1
|
|
1421
|
+
if next_depth > max_depth:
|
|
1422
|
+
continue
|
|
1423
|
+
|
|
1424
|
+
new_cost = cost + self._edge_weight(edge)
|
|
1425
|
+
if new_cost < dist.get(neighbor_id, float("inf")):
|
|
1426
|
+
dist[neighbor_id] = new_cost
|
|
1427
|
+
depth_map[neighbor_id] = next_depth
|
|
1428
|
+
parents[neighbor_id] = (node_id, edge)
|
|
1429
|
+
heapq.heappush(heap, (new_cost, next_depth, neighbor_id))
|
|
1430
|
+
|
|
1431
|
+
return PathResult(nodes=[], found=False)
|
|
1432
|
+
|
|
1433
|
+
def a_star(self, heuristic: Callable[[NodeRef[N], NodeRef[N]], float]) -> PathResult[N]:
|
|
1434
|
+
"""Execute A* algorithm with a heuristic."""
|
|
1435
|
+
targets = self._get_targets()
|
|
1436
|
+
target_ids = {t.id for t in targets}
|
|
1437
|
+
target_ref = targets[0]
|
|
1438
|
+
max_depth = self._max_depth_value()
|
|
1439
|
+
|
|
1440
|
+
if isinstance(self._weight_spec, str) and self._edge_type is None:
|
|
1441
|
+
raise ValueError("weight by property requires via(edge)")
|
|
1442
|
+
|
|
1443
|
+
if self._source.id in target_ids:
|
|
1444
|
+
node_ref = self._create_node_ref(self._source.id)
|
|
1445
|
+
if node_ref is None:
|
|
1446
|
+
return PathResult(nodes=[], found=False)
|
|
1447
|
+
return PathResult(nodes=[node_ref], found=True, total_weight=0.0)
|
|
1448
|
+
|
|
1449
|
+
import heapq
|
|
1450
|
+
|
|
1451
|
+
def safe_heuristic(current: NodeRef[N]) -> float:
|
|
1452
|
+
try:
|
|
1453
|
+
return float(heuristic(current, target_ref))
|
|
1454
|
+
except Exception:
|
|
1455
|
+
return 0.0
|
|
1456
|
+
|
|
1457
|
+
g_score: Dict[int, float] = {self._source.id: 0.0}
|
|
1458
|
+
parents: Dict[int, Tuple[Optional[int], Optional[EdgeResult]]] = {
|
|
1459
|
+
self._source.id: (None, None)
|
|
1460
|
+
}
|
|
1461
|
+
heap: List[Tuple[float, float, int, int]] = []
|
|
1462
|
+
|
|
1463
|
+
source_ref = self._create_node_ref(self._source.id)
|
|
1464
|
+
if source_ref is None:
|
|
1465
|
+
return PathResult(nodes=[], found=False)
|
|
1466
|
+
heapq.heappush(heap, (safe_heuristic(source_ref), 0.0, self._source.id, 0))
|
|
1467
|
+
|
|
1468
|
+
visited: Set[int] = set()
|
|
1469
|
+
|
|
1470
|
+
while heap:
|
|
1471
|
+
f_score, g_score_val, node_id, depth = heapq.heappop(heap)
|
|
1472
|
+
if node_id in visited:
|
|
1473
|
+
continue
|
|
1474
|
+
visited.add(node_id)
|
|
1475
|
+
|
|
1476
|
+
if node_id in target_ids:
|
|
1477
|
+
result = self._reconstruct_path(parents, node_id)
|
|
1478
|
+
result.total_weight = g_score_val
|
|
1479
|
+
return result
|
|
1480
|
+
|
|
1481
|
+
if depth >= max_depth:
|
|
1482
|
+
continue
|
|
1483
|
+
|
|
1484
|
+
for neighbor_id, edge in self._iter_neighbors(node_id):
|
|
1485
|
+
next_depth = depth + 1
|
|
1486
|
+
if next_depth > max_depth:
|
|
1487
|
+
continue
|
|
1488
|
+
|
|
1489
|
+
tentative_g = g_score_val + self._edge_weight(edge)
|
|
1490
|
+
if tentative_g < g_score.get(neighbor_id, float("inf")):
|
|
1491
|
+
neighbor_ref = self._create_node_ref(neighbor_id)
|
|
1492
|
+
if neighbor_ref is None:
|
|
1493
|
+
continue
|
|
1494
|
+
g_score[neighbor_id] = tentative_g
|
|
1495
|
+
parents[neighbor_id] = (node_id, edge)
|
|
1496
|
+
h = safe_heuristic(neighbor_ref)
|
|
1497
|
+
heapq.heappush(heap, (tentative_g + h, tentative_g, neighbor_id, next_depth))
|
|
1498
|
+
|
|
1499
|
+
return PathResult(nodes=[], found=False)
|
|
1500
|
+
|
|
1501
|
+
def all_paths(self, max_paths: Optional[int] = None) -> Iterator[PathResult[N]]:
|
|
1502
|
+
"""Yield shortest paths (currently returns at most one)."""
|
|
1503
|
+
result = self.dijkstra() if self._weight_spec is not None else self.bfs()
|
|
1504
|
+
if result.found:
|
|
1505
|
+
yield result
|
|
1506
|
+
|
|
1507
|
+
|
|
1508
|
+
__all__ = [
|
|
1509
|
+
"TraversalBuilder",
|
|
1510
|
+
"TraversalResult",
|
|
1511
|
+
"EdgeTraversalResult",
|
|
1512
|
+
"PathFindingBuilder",
|
|
1513
|
+
"PathResult",
|
|
1514
|
+
"EdgeResult",
|
|
1515
|
+
"RawEdge",
|
|
1516
|
+
"TraverseOptions",
|
|
1517
|
+
"PropLoadStrategy",
|
|
1518
|
+
"OutStep",
|
|
1519
|
+
"InStep",
|
|
1520
|
+
"BothStep",
|
|
1521
|
+
"TraverseStep",
|
|
1522
|
+
"TraversalStep",
|
|
1523
|
+
]
|