-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathnode.py
More file actions
50 lines (41 loc) · 1.78 KB
/
node.py
File metadata and controls
50 lines (41 loc) · 1.78 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
from abc import ABC, abstractmethod
from typing import List, Optional
from feast.infra.compute_engines.dag.context import ExecutionContext
from feast.infra.compute_engines.dag.value import DAGValue
class DAGNode(ABC):
name: str
inputs: List["DAGNode"]
outputs: List["DAGNode"]
def __init__(self, name: str, inputs: Optional[List["DAGNode"]] = None):
self.name = name
self.inputs: List["DAGNode"] = []
self.outputs: List["DAGNode"] = []
for node in inputs or []:
self.add_input(node)
def add_input(self, node: "DAGNode"):
if node in self.inputs:
raise ValueError(f"Input node {node.name} already added to {self.name}")
self.inputs.append(node)
node.outputs.append(self)
def get_input_values(self, context: ExecutionContext) -> List[DAGValue]:
input_values = []
for input_node in self.inputs:
if input_node.name not in context.node_outputs:
raise KeyError(
f"Missing output for input node '{input_node.name}' in context."
)
input_values.append(context.node_outputs[input_node.name])
return input_values
def get_single_input_value(self, context: ExecutionContext) -> DAGValue:
if len(self.inputs) != 1:
raise RuntimeError(
f"DAGNode '{self.name}' expected exactly 1 input, but got {len(self.inputs)}."
)
input_node = self.inputs[0]
if input_node.name not in context.node_outputs:
raise KeyError(
f"Missing output for input node '{input_node.name}' in context."
)
return context.node_outputs[input_node.name]
@abstractmethod
def execute(self, context: ExecutionContext) -> DAGValue: ...