From d5a3a8a5a9d20143322c1e3925709064fe7af445 Mon Sep 17 00:00:00 2001
From: Dylan Pulver <dylanpulver@users.noreply.github.com>
Date: Tue, 25 Aug 2026 16:33:26 +0300
Subject: [PATCH] fix: preserve column order in PointsModel.parse()

---
 src/spatialdata/models/models.py | 14 ++++++++----
 tests/models/test_models.py      | 39 ++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/spatialdata/models/models.py b/src/spatialdata/models/models.py
index a818bdadc..5b44fe520 100644
--- a/src/spatialdata/models/models.py
+++ b/src/spatialdata/models/models.py
@@ -790,8 +790,10 @@ def _(
                 df_dict[instance_key] = annotation[instance_key]
             if Z not in axes and Z in annotation.columns:
                 logger.info(f"Column `{Z}` in `annotation` will be ignored since the data is 2D.")
-            for c in set(annotation.columns) - {feature_key, instance_key, X, Y, Z}:
-                df_dict[c] = annotation[c]
+            handled_columns = {feature_key, instance_key, X, Y, Z}
+            for c in annotation.columns:
+                if c not in handled_columns:
+                    df_dict[c] = annotation[c]
 
         table: DaskDataFrame = dd.from_pandas(pd.DataFrame(**df_kwargs), **kwargs)
         return cls._add_metadata_and_validate(
@@ -881,15 +883,17 @@ def _(
                 )
         if Z not in axes and Z in data.columns:
             logger.info(f"Column `{Z}` in `data` will be ignored since the data is 2D.")
-        for c in set(data.columns) - {
+        handled_columns = {
             feature_key,
             instance_key,
             *coordinates.values(),
             X,
             Y,
             Z,
-        }:
-            table[c] = data[c]
+        }
+        for c in data.columns:
+            if c not in handled_columns:
+                table[c] = data[c]
 
         validated = cls._add_metadata_and_validate(
             table,
diff --git a/tests/models/test_models.py b/tests/models/test_models.py
index 33a98a434..4ee3210b7 100644
--- a/tests/models/test_models.py
+++ b/tests/models/test_models.py
@@ -814,6 +814,45 @@ def test_force2d():
     assert_elements_are_identical(multipolygons_3d, expected_multipolygons_2d)
 
 
+def test_points_model_preserves_column_order():
+    # the extra columns used to be added iterating over a set, so their order in the parsed element depended on
+    # PYTHONHASHSEED. Renaming the coordinate columns changes the set of column names, which used to bypass the
+    # reordering safeguard at the end of PointsModel.parse().
+    extra_columns = ["qv", "intensity", "radius", "z_score", "nucleus_distance", "codeword_index"]
+    n = 10
+    data = pd.DataFrame(
+        {
+            "my_x": np.arange(n, dtype=float),
+            "my_y": np.arange(n, dtype=float),
+            "target": pd.Categorical(["a", "b"] * (n // 2)),
+            "cell_id": np.arange(n),
+            **{c: np.arange(n, dtype=float) for c in extra_columns},
+        }
+    )
+    expected = ["x", "y", "target", "cell_id", *extra_columns]
+
+    from_pandas = PointsModel.parse(
+        data, coordinates={"x": "my_x", "y": "my_y"}, feature_key="target", instance_key="cell_id"
+    )
+    assert list(from_pandas.columns) == expected
+
+    from_dask = PointsModel.parse(
+        dd.from_pandas(data, npartitions=2),
+        coordinates={"x": "my_x", "y": "my_y"},
+        feature_key="target",
+        instance_key="cell_id",
+    )
+    assert list(from_dask.columns) == expected
+
+    from_numpy = PointsModel.parse(
+        data[["my_x", "my_y"]].to_numpy(),
+        annotation=data.drop(columns=["my_x", "my_y"]),
+        feature_key="target",
+        instance_key="cell_id",
+    )
+    assert list(from_numpy.columns) == expected
+
+
 def test_dask_points_unsorted_index_with_warning(points):
     chunksize = 300
     element = points["points_0"]