diff --git a/qgis/widgets/dataset_widget.py b/qgis/widgets/dataset_widget.py
index 9d3f9e9ae..10078e534 100644
--- a/qgis/widgets/dataset_widget.py
+++ b/qgis/widgets/dataset_widget.py
@@ -24,11 +24,11 @@
     QVBoxLayout,
     QWidget,
 )
+from ribasim_qgis.core.nodes import Edge, Node, load_nodes_from_geopackage
+from ribasim_qgis.core.topology import derive_connectivity, explode_lines
 
 from qgis.core import QgsMapLayer, QgsProject
 from qgis.core.additions.edit import edit
-from ribasim_qgis.core.nodes import Edge, Node, load_nodes_from_geopackage
-from ribasim_qgis.core.topology import derive_connectivity, explode_lines
 
 
 class DatasetTreeWidget(QTreeWidget):
diff --git a/qgis/widgets/nodes_widget.py b/qgis/widgets/nodes_widget.py
index ae185af0e..a44ebb12f 100644
--- a/qgis/widgets/nodes_widget.py
+++ b/qgis/widgets/nodes_widget.py
@@ -1,7 +1,6 @@
 from functools import partial
 
 from PyQt5.QtWidgets import QGridLayout, QPushButton, QVBoxLayout, QWidget
-
 from ribasim_qgis.core.nodes import NODES
 
 
diff --git a/qgis/widgets/output_widget.py b/qgis/widgets/output_widget.py
index c26ad6c56..c85b8df6c 100644
--- a/qgis/widgets/output_widget.py
+++ b/qgis/widgets/output_widget.py
@@ -5,21 +5,32 @@ class OutputWidget(QWidget):
     def __init__(self, parent):
         super().__init__(parent)
         self.parent = parent
-        self.output_button = QPushButton("Associate Output")
-
-        self.output_button.clicked.connect(self.set_output)
-
+        self.node_output_button = QPushButton("Associate Node Output")
+        self.edge_output_button = QPushButton("Associate Edge Output")
+        self.node_output_button.clicked.connect(self.set_node_output)
+        self.edge_output_button.clicked.connect(self.set_edge_output)
         layout = QVBoxLayout()
-        layout.addWidget(self.output_button)
+        layout.addWidget(self.node_output_button)
+        layout.addWidget(self.edge_output_button)
         layout.addStretch()
         self.setLayout(layout)
 
-    def set_output(self):
+    def _set_output(self, layer, column: str):
         path, _ = QFileDialog.getOpenFileName(self, "Select file", "", "*.arrow")
         if path == "":
             return
+        if layer is not None:
+            layer.setCustomProperty("arrow_type", "timeseries")
+            layer.setCustomProperty("arrow_path", path)
+            layer.setCustomProperty("arrow_fid_column", column)
+        return
+
+    def set_node_output(self):
         node_layer = self.parent.dataset_widget.node_layer
-        if node_layer is not None:
-            node_layer.setCustomProperty("arrow_type", "timeseries")
-            node_layer.setCustomProperty("arrow_path", path)
+        self._set_output(node_layer, "node_id")
+        return
+
+    def set_edge_output(self):
+        edge_layer = self.parent.dataset_widget.edge_layer
+        self._set_output(edge_layer, "edge_id")
         return
diff --git a/qgis/widgets/ribasim_widget.py b/qgis/widgets/ribasim_widget.py
index 01284f439..2a2cc6bde 100644
--- a/qgis/widgets/ribasim_widget.py
+++ b/qgis/widgets/ribasim_widget.py
@@ -9,12 +9,12 @@
 from typing import Any
 
 from PyQt5.QtWidgets import QTabWidget, QVBoxLayout, QWidget
-
-from qgis.core import QgsMapLayer, QgsProject
 from ribasim_qgis.widgets.dataset_widget import DatasetWidget
 from ribasim_qgis.widgets.nodes_widget import NodesWidget
 from ribasim_qgis.widgets.output_widget import OutputWidget
 
+from qgis.core import QgsMapLayer, QgsProject
+
 PYQT_DELETED_ERROR = "wrapped C/C++ object of type QgsLayerTreeGroup has been deleted"