From 6d2ea715a7cd60444903fcf1adaa843cff1e0e30 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Fri, 13 Dec 2024 11:43:59 +0100
Subject: [PATCH] Automatically name index

---
 python/ribasim/ribasim/geometry/edge.py |   7 +-
 python/ribasim/ribasim/geometry/node.py |   7 +-
 python/ribasim/ribasim/input_base.py    |   2 -
 python/ribasim/ribasim/schemas.py       | 169 +++++++++++++++++++++++-
 python/ribasim/tests/test_io.py         |   5 +
 python/ribasim/tests/test_model.py      |   6 +
 utils/templates/schemas.py.jinja        |   9 +-
 7 files changed, 191 insertions(+), 14 deletions(-)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index f6285e1c5..4f2d61946 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -49,9 +49,10 @@ class EdgeSchema(_GeoBaseSchema):
     edge_type: Series[str] = pa.Field(default="flow")
     geometry: GeoSeries[LineString] = pa.Field(default=None, nullable=True)
 
-    @classmethod
-    def _index_name(self) -> str:
-        return "edge_id"
+    @pa.dataframe_parser
+    def _name_index(cls, df):
+        df.index.name = "edge_id"
+        return df
 
 
 class EdgeTable(SpatialTableModel[EdgeSchema]):
diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index 691dc4d19..b92354b63 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -27,9 +27,10 @@ class NodeSchema(_GeoBaseSchema):
     )
     geometry: GeoSeries[Point] = pa.Field(default=None, nullable=True)
 
-    @classmethod
-    def _index_name(self) -> str:
-        return "node_id"
+    @pa.dataframe_parser
+    def _name_index(cls, df):
+        df.index.name = "node_id"
+        return df
 
 
 class NodeTable(SpatialTableModel[NodeSchema]):
diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index c6613c908..48e0e03fd 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -217,7 +217,6 @@ def _check_dataframe(cls, value: Any) -> Any:
 
         # Enable initialization with a DataFrame.
         if isinstance(value, pd.DataFrame | gpd.GeoDataFrame):
-            value.index.rename("fid", inplace=True)
             value = {"df": value}
 
         return value
@@ -386,7 +385,6 @@ def _from_db(cls, path: Path, table: str):
                     # tell pyarrow to map to pd.ArrowDtype rather than NumPy
                     arrow_to_pandas_kwargs={"types_mapper": pd.ArrowDtype},
                 )
-                df.index.rename(cls.tableschema()._index_name(), inplace=True)
             else:
                 df = None
 
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index 2bcd22222..801168bf1 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -16,9 +16,12 @@ class Config:
         add_missing_columns = True
         coerce = True
 
-    @classmethod
-    def _index_name(self) -> str:
-        return "fid"
+    @pa.dataframe_parser
+    def _name_index(cls, df):
+        # Node and Edge have different index names, avoid running both parsers
+        if cls.__name__ not in ("NodeSchema", "EdgeSchema"):
+            df.index.name = "fid"
+        return df
 
     @classmethod
     def migrate(cls, df: Any, schema_version: int) -> Any:
@@ -30,15 +33,19 @@ def migrate(cls, df: Any, schema_version: int) -> Any:
 
 class BasinConcentrationExternalSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     substance: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     concentration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -46,12 +53,15 @@ class BasinConcentrationExternalSchema(_BaseSchema):
 
 class BasinConcentrationStateSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     substance: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     concentration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -59,18 +69,23 @@ class BasinConcentrationStateSchema(_BaseSchema):
 
 class BasinConcentrationSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     substance: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     drainage: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     precipitation: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -78,10 +93,13 @@ class BasinConcentrationSchema(_BaseSchema):
 
 class BasinProfileSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     area: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(nullable=False)
+
     level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -89,9 +107,11 @@ class BasinProfileSchema(_BaseSchema):
 
 class BasinStateSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -99,18 +119,23 @@ class BasinStateSchema(_BaseSchema):
 
 class BasinStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     drainage: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     potential_evaporation: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = (
         pa.Field(nullable=True)
     )
+
     infiltration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     precipitation: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -118,15 +143,19 @@ class BasinStaticSchema(_BaseSchema):
 
 class BasinSubgridSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     subgrid_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     basin_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     subgrid_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -134,21 +163,27 @@ class BasinSubgridSchema(_BaseSchema):
 
 class BasinTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     drainage: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     potential_evaporation: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = (
         pa.Field(nullable=True)
     )
+
     infiltration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     precipitation: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -156,15 +191,19 @@ class BasinTimeSchema(_BaseSchema):
 
 class ContinuousControlFunctionSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     input: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     output: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     controlled_variable: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
@@ -172,18 +211,23 @@ class ContinuousControlFunctionSchema(_BaseSchema):
 
 class ContinuousControlVariableSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     listen_node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     variable: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     weight: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     look_ahead: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -191,12 +235,15 @@ class ContinuousControlVariableSchema(_BaseSchema):
 
 class DiscreteControlConditionSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     compound_variable_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     greater_than: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -204,12 +251,15 @@ class DiscreteControlConditionSchema(_BaseSchema):
 
 class DiscreteControlLogicSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     truth_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
@@ -217,21 +267,27 @@ class DiscreteControlLogicSchema(_BaseSchema):
 
 class DiscreteControlVariableSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     compound_variable_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     listen_node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     variable: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     weight: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     look_ahead: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
@@ -239,15 +295,19 @@ class DiscreteControlVariableSchema(_BaseSchema):
 
 class FlowBoundaryConcentrationSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     substance: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     concentration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -255,10 +315,13 @@ class FlowBoundaryConcentrationSchema(_BaseSchema):
 
 class FlowBoundaryStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -266,12 +329,15 @@ class FlowBoundaryStaticSchema(_BaseSchema):
 
 class FlowBoundaryTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -279,12 +345,15 @@ class FlowBoundaryTimeSchema(_BaseSchema):
 
 class FlowDemandStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     demand: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     priority: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=True
     )
@@ -292,15 +361,19 @@ class FlowDemandStaticSchema(_BaseSchema):
 
 class FlowDemandTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     demand: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     priority: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=True
     )
@@ -308,15 +381,19 @@ class FlowDemandTimeSchema(_BaseSchema):
 
 class LevelBoundaryConcentrationSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     substance: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     concentration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -324,10 +401,13 @@ class LevelBoundaryConcentrationSchema(_BaseSchema):
 
 class LevelBoundaryStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -335,12 +415,15 @@ class LevelBoundaryStaticSchema(_BaseSchema):
 
 class LevelBoundaryTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -348,15 +431,19 @@ class LevelBoundaryTimeSchema(_BaseSchema):
 
 class LevelDemandStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     min_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     max_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     priority: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=True
     )
@@ -364,18 +451,23 @@ class LevelDemandStaticSchema(_BaseSchema):
 
 class LevelDemandTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     min_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     max_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     priority: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=True
     )
@@ -383,16 +475,21 @@ class LevelDemandTimeSchema(_BaseSchema):
 
 class LinearResistanceStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     resistance: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     max_flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -400,22 +497,29 @@ class LinearResistanceStaticSchema(_BaseSchema):
 
 class ManningResistanceStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     length: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     manning_n: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     profile_width: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     profile_slope: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -423,25 +527,33 @@ class ManningResistanceStaticSchema(_BaseSchema):
 
 class OutletStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     min_flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     max_flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     min_upstream_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     max_downstream_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = (
         pa.Field(nullable=True)
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -449,25 +561,33 @@ class OutletStaticSchema(_BaseSchema):
 
 class PidControlStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     listen_node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     target: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     proportional: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     integral: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     derivative: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -475,27 +595,35 @@ class PidControlStaticSchema(_BaseSchema):
 
 class PidControlTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     listen_node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     target: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     proportional: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     integral: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     derivative: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -503,25 +631,33 @@ class PidControlTimeSchema(_BaseSchema):
 
 class PumpStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     min_flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     max_flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     min_upstream_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     max_downstream_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = (
         pa.Field(nullable=True)
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -529,19 +665,25 @@ class PumpStaticSchema(_BaseSchema):
 
 class TabulatedRatingCurveStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     max_downstream_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = (
         pa.Field(nullable=True)
     )
+
     control_state: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=True
     )
@@ -549,18 +691,23 @@ class TabulatedRatingCurveStaticSchema(_BaseSchema):
 
 class TabulatedRatingCurveTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     flow_rate: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     max_downstream_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = (
         pa.Field(nullable=True)
     )
@@ -568,15 +715,19 @@ class TabulatedRatingCurveTimeSchema(_BaseSchema):
 
 class UserDemandConcentrationSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     substance: Series[Annotated[pd.ArrowDtype, pyarrow.string()]] = pa.Field(
         nullable=False
     )
+
     concentration: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
@@ -584,19 +735,25 @@ class UserDemandConcentrationSchema(_BaseSchema):
 
 class UserDemandStaticSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     active: Series[Annotated[pd.ArrowDtype, pyarrow.bool_()]] = pa.Field(nullable=True)
+
     demand: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=True
     )
+
     return_factor: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     min_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     priority: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=True
     )
@@ -604,21 +761,27 @@ class UserDemandStaticSchema(_BaseSchema):
 
 class UserDemandTimeSchema(_BaseSchema):
     fid: Index[Int32] = pa.Field(default=1, check_name=True, coerce=True)
+
     node_id: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=False, default=0
     )
+
     time: Series[Annotated[pd.ArrowDtype, pyarrow.timestamp("ms")]] = pa.Field(
         nullable=False
     )
+
     demand: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     return_factor: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     min_level: Series[Annotated[pd.ArrowDtype, pyarrow.float64()]] = pa.Field(
         nullable=False
     )
+
     priority: Series[Annotated[pd.ArrowDtype, pyarrow.int32()]] = pa.Field(
         nullable=True
     )
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index 0320dca87..7c2dcf7b7 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -97,6 +97,11 @@ def test_extra_columns():
 def test_index_tables():
     p = pump.Static(flow_rate=[1.2])
     assert p.df.index.name == "fid"
+    # Index name is applied by _name_index
+    df = p.df.reset_index(drop=True)
+    assert df.index.name is None
+    p.df = df
+    assert p.df.index.name == "fid"
 
 
 def test_extra_spatial_columns():
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index f5e341561..9336b5275 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -112,6 +112,12 @@ def test_write_adds_fid_in_tables(basic, tmp_path):
     assert model_orig.edge.df.index.name == "edge_id"
     assert model_orig.edge.df.index.equals(pd.RangeIndex(1, nrow + 1))
 
+    # Index name is applied by _name_index
+    df = model_orig.edge.df.copy()
+    df.index.name = "other"
+    model_orig.edge.df = df
+    assert model_orig.edge.df.index.name == "edge_id"
+
     model_orig.write(tmp_path / "basic/ribasim.toml")
     with connect(tmp_path / "basic/database.gpkg") as connection:
         query = f"select * from {esc_id('Basin / profile')}"
diff --git a/utils/templates/schemas.py.jinja b/utils/templates/schemas.py.jinja
index 19b53c7ce..6e79f6953 100644
--- a/utils/templates/schemas.py.jinja
+++ b/utils/templates/schemas.py.jinja
@@ -15,9 +15,12 @@ class _BaseSchema(pa.DataFrameModel):
         add_missing_columns = True
         coerce = True
 
-    @classmethod
-    def _index_name(self) -> str:
-        return "fid"
+    @pa.dataframe_parser
+    def _name_index(cls, df):
+        # Node and Edge have different index names, avoid running both parsers
+        if cls.__name__ not in ("NodeSchema", "EdgeSchema"):
+            df.index.name = "fid"
+        return df
 
     @classmethod
     def migrate(cls, df: Any, schema_version: int) -> Any: