feat: preserve field order when replacing existing fields in with_field

src/awkward/operations/ak_with_field.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 from functools import reduce
 
 import awkward as ak
@@ -96,10 +95,6 @@ def _impl(base, what, where, highlevel, behavior, attrs):
                 ),
             )
 
-        keys = copy.copy(base.fields)
-        if where in base.fields:
-            keys.remove(where)
-
         def purelist_is_record(layout):
             result = False
 
@@ -141,20 +136,38 @@ def action(inputs, **kwargs):
                             backend.nplike.shape_item_as_index(base.length),
                         )
                     )
+
+                base_fields = base.fields
+                # Check if we're replacing an existing field
+                field_exists = where is not None and where in base_fields
                 if base.is_tuple:
                     # Preserve tuple-ness
                     if where is None:
                         fields = None
+                        contents = [base[k] for k in base_fields] + [what]
                     # Otherwise the tuple becomes a record
                     else:
-                        fields = [*keys, where]
+                        if field_exists:
+                            fields = list(base_fields)
+                            contents = [what if k == where else base[k] for k in fields]
+                        else:
+                            # Add new field at the end
+                            fields = [*base_fields, where]
+                            contents = [base[k] for k in base.fields] + [what]
                 # Records with `where=None` will create a tuple-like key
                 elif where is None:
-                    fields = [*keys, str(len(keys))]
+                    fields = [*base_fields, str(len(base_fields))]
+                    contents = [base[k] for k in base_fields] + [what]
                 else:
-                    fields = [*keys, where]
+                    if field_exists:
+                        fields = list(base_fields)
+                        contents = [what if k == where else base[k] for k in fields]
+                    else:
+                        fields = [*base_fields, where]
+                        contents = [base[k] for k in base_fields] + [what]
+
                 out = ak.contents.RecordArray(
-                    [base[k] for k in keys] + [what],
+                    contents,
                     fields,
                     parameters=base.parameters,
                 )

tests/test_3681_preserve_field_order_with_field.py

@@ -0,0 +1,228 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import awkward as ak
+
+
+def test_simple_field_replacement_preserves_order():
+    """Test that replacing an existing field preserves its position in the field order."""
+    array = ak.Array([{"x": 1, "y": 2, "z": 3}])
+
+    # Replace middle field
+    result = ak.with_field(array, [100], "y")
+    assert ak.fields(result) == ["x", "y", "z"]
+    assert result.to_list() == [{"x": 1, "y": 100, "z": 3}]
+
+    # Replace first field
+    result = ak.with_field(array, [100], "x")
+    assert ak.fields(result) == ["x", "y", "z"]
+    assert result.to_list() == [{"x": 100, "y": 2, "z": 3}]
+
+    # Replace last field
+    result = ak.with_field(array, [100], "z")
+    assert ak.fields(result) == ["x", "y", "z"]
+    assert result.to_list() == [{"x": 1, "y": 2, "z": 100}]
+
+
+def test_simple_field_addition_appends():
+    """Test that adding a new field appends it to the end."""
+    array = ak.Array([{"x": 1, "y": 2, "z": 3}])
+
+    result = ak.with_field(array, [100], "w")
+    assert ak.fields(result) == ["x", "y", "z", "w"]
+    assert result.to_list() == [{"x": 1, "y": 2, "z": 3, "w": 100}]
+
+
+def test_setitem_replacement_preserves_order():
+    """Test that using array['field'] = value preserves field order."""
+    array = ak.Array([{"x": 1, "y": 2, "z": 3}])
+
+    # Replace middle field using setitem
+    array["y"] = [100]
+    assert ak.fields(array) == ["x", "y", "z"]
+    assert array.to_list() == [{"x": 1, "y": 100, "z": 3}]
+
+
+def test_setitem_addition_appends():
+    """Test that using array['new_field'] = value appends the field."""
+    array = ak.Array([{"x": 1, "y": 2, "z": 3}])
+
+    array["w"] = [100]
+    assert ak.fields(array) == ["x", "y", "z", "w"]
+    assert array.to_list() == [{"x": 1, "y": 2, "z": 3, "w": 100}]
+
+
+def test_nested_field_replacement_preserves_order():
+    """Test that replacing a nested field preserves order at both levels."""
+    array = ak.Array([{"a": {"x": 1, "y": 2, "z": 3}, "b": 10, "c": 20}])
+
+    # Replace nested field "y" inside "a"
+    result = ak.with_field(array, [100], ["a", "y"])
+    assert ak.fields(result) == ["a", "b", "c"]
+    assert ak.fields(result["a"]) == ["x", "y", "z"]
+    assert result.to_list() == [{"a": {"x": 1, "y": 100, "z": 3}, "b": 10, "c": 20}]
+
+
+def test_nested_field_replacement_with_setitem():
+    """Test nested field replacement using setitem syntax."""
+    array = ak.Array([{"a": {"x": 1, "y": 2, "z": 3}, "b": 10, "c": 20}])
+
+    # Replace nested field using setitem
+    array["a", "y"] = [100]
+    assert ak.fields(array) == ["a", "b", "c"]
+    assert ak.fields(array["a"]) == ["x", "y", "z"]
+    assert array.to_list() == [{"a": {"x": 1, "y": 100, "z": 3}, "b": 10, "c": 20}]
+
+
+def test_nested_field_addition():
+    """Test that adding a new nested field appends it."""
+    array = ak.Array([{"a": {"x": 1, "y": 2, "z": 3}, "b": 10, "c": 20}])
+
+    # Add new nested field
+    result = ak.with_field(array, [100], ["a", "w"])
+    assert ak.fields(result) == ["a", "b", "c"]
+    assert ak.fields(result["a"]) == ["x", "y", "z", "w"]
+    assert result.to_list() == [
+        {"a": {"x": 1, "y": 2, "z": 3, "w": 100}, "b": 10, "c": 20}
+    ]
+
+
+def test_multiple_replacements_preserve_order():
+    """Test that multiple sequential replacements preserve order."""
+    array = ak.Array([{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}])
+
+    # Replace multiple fields in different orders
+    result = array
+    result = ak.with_field(result, [20], "b")
+    result = ak.with_field(result, [40], "d")
+    result = ak.with_field(result, [10], "a")
+
+    assert ak.fields(result) == ["a", "b", "c", "d", "e"]
+    assert result.to_list() == [{"a": 10, "b": 20, "c": 3, "d": 40, "e": 5}]
+
+
+def test_replacement_with_array_data():
+    """Test field replacement with various array structures."""
+    array = ak.Array(
+        [
+            {"x": 1, "y": [1, 2], "z": 3},
+            {"x": 2, "y": [3, 4, 5], "z": 4},
+        ]
+    )
+
+    # Replace field with different array structure
+    result = ak.with_field(array, [[10, 20], [30, 40, 50]], "y")
+    assert ak.fields(result) == ["x", "y", "z"]
+    assert result.to_list() == [
+        {"x": 1, "y": [10, 20], "z": 3},
+        {"x": 2, "y": [30, 40, 50], "z": 4},
+    ]
+
+
+def test_replacement_with_scalar_broadcast():
+    """Test field replacement with scalar broadcasting."""
+    array = ak.Array(
+        [
+            {"x": 1, "y": 2, "z": 3},
+            {"x": 4, "y": 5, "z": 6},
+            {"x": 7, "y": 8, "z": 9},
+        ]
+    )
+
+    # Replace with scalar (should broadcast)
+    result = ak.with_field(array, 100, "y")
+    assert ak.fields(result) == ["x", "y", "z"]
+    assert result.to_list() == [
+        {"x": 1, "y": 100, "z": 3},
+        {"x": 4, "y": 100, "z": 6},
+        {"x": 7, "y": 100, "z": 9},
+    ]
+
+
+def test_tuple_field_replacement():
+    """Test replacing a field in a tuple (which converts it to a record)."""
+    array = ak.Array([(1, 2, 3)])
+
+    # Replacing field "1" should preserve position
+    result = ak.with_field(array, [100], "1")
+    assert ak.fields(result) == ["0", "1", "2"]
+    assert result.to_list() == [{"0": 1, "1": 100, "2": 3}]
+
+
+def test_single_field_replacement():
+    """Test replacing the only field in a record."""
+    array = ak.Array([{"x": 1}, {"x": 2}, {"x": 3}])
+
+    result = ak.with_field(array, [100, 200, 300], "x")
+    assert ak.fields(result) == ["x"]
+    assert result.to_list() == [{"x": 100}, {"x": 200}, {"x": 300}]
+
+
+def test_deeply_nested_replacement():
+    """Test replacing a deeply nested field."""
+    array = ak.Array([{"a": {"b": {"x": 1, "y": 2, "z": 3}, "c": 10}, "d": 100}])
+
+    # Replace deeply nested field
+    result = ak.with_field(array, [999], ["a", "b", "y"])
+    assert ak.fields(result) == ["a", "d"]
+    assert ak.fields(result["a"]) == ["b", "c"]
+    assert ak.fields(result["a", "b"]) == ["x", "y", "z"]
+    assert result.to_list() == [
+        {"a": {"b": {"x": 1, "y": 999, "z": 3}, "c": 10}, "d": 100}
+    ]
+
+
+def test_where_none_always_appends():
+    """Test that where=None always appends a new field."""
+    array = ak.Array([{"x": 1, "y": 2}])
+
+    # where=None should always append
+    result = ak.with_field(array, [100], where=None)
+    assert ak.fields(result) == ["x", "y", "2"]
+    assert result.to_list() == [{"x": 1, "y": 2, "2": 100}]
+
+    # Multiple where=None should keep appending
+    result = ak.with_field(result, [200], where=None)
+    assert ak.fields(result) == ["x", "y", "2", "3"]
+    assert result.to_list() == [{"x": 1, "y": 2, "2": 100, "3": 200}]
+
+
+def test_mixed_operations():
+    """Test a mix of additions and replacements."""
+    array = ak.Array([{"x": 1, "y": 2, "z": 3}])
+
+    # Add a new field
+    result = ak.with_field(array, [10], "a")
+    assert ak.fields(result) == ["x", "y", "z", "a"]
+
+    # Replace an existing field
+    result = ak.with_field(result, [20], "y")
+    assert ak.fields(result) == ["x", "y", "z", "a"]
+
+    # Add another new field
+    result = ak.with_field(result, [30], "b")
+    assert ak.fields(result) == ["x", "y", "z", "a", "b"]
+
+    # Replace the first field
+    result = ak.with_field(result, [100], "x")
+    assert ak.fields(result) == ["x", "y", "z", "a", "b"]
+
+    assert result.to_list() == [{"x": 100, "y": 20, "z": 3, "a": 10, "b": 30}]
+
+
+def test_replacement_in_list_of_records():
+    """Test field replacement in a list of records."""
+    array = ak.Array(
+        [
+            [{"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}],
+            [{"x": 7, "y": 8, "z": 9}],
+        ]
+    )
+
+    result = ak.with_field(array, [[100, 200], [300]], "y")
+    assert ak.fields(result) == ["x", "y", "z"]
+    assert result.to_list() == [
+        [{"x": 1, "y": 100, "z": 3}, {"x": 4, "y": 200, "z": 6}],
+        [{"x": 7, "y": 300, "z": 9}],
+    ]