From e65d207bc2c1798d63118f035b7965542662d764 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Tue, 13 Oct 2026 19:19:39 +0800
Subject: [PATCH 1/4] [feat] kv_update: add missing keys/fields and defer
 readiness with finalize

kv_update is now a superset of kv_put: a missing key is created, and a field
the key does not hold yet is stored as given without calling merge_fn. Merge
results are validated against the stored old value, so the update request no
longer carries a controller field_schema.

finalize=False stores the update but keeps the fields not ready, so samplers,
the dataloader and kv_batch_get do not see intermediate values; the update
with finalize=True (the default) marks them ready. Selective dump/restore
keeps each field's readiness as it was when dumped.

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 README.md                                     | 18 ++++--
 tests/e2e/test_data_dump_e2e.py               | 33 +++++++++-
 tests/e2e/test_kv_interface_e2e.py            | 63 +++++++++++++-----
 tests/test_controller_data_partitions.py      | 12 ++++
 tests/test_data_dump.py                       | 18 +++++-
 tests/test_kv_update.py                       | 41 ++++++------
 transfer_queue/client.py                      | 19 ++++--
 transfer_queue/controller.py                  | 44 ++++++++-----
 transfer_queue/data_dump.py                   |  4 +-
 transfer_queue/interface.py                   | 53 +++++++++------
 transfer_queue/storage/managers/base.py       | 11 +++-
 .../managers/simple_storage_manager.py        | 13 ++--
 transfer_queue/storage/simple_storage.py      | 64 +++++++++----------
 tutorial/02_kv_interface.py                   |  2 +-
 14 files changed, 266 insertions(+), 129 deletions(-)

diff --git a/README.md b/README.md
index 4df55109..23367a11 100644
--- a/README.md
+++ b/README.md
@@ -116,15 +116,21 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 - **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag.
 - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches.
 - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields).
-- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into already-produced fields with `merge_fn(old, new)`. SimpleStorage only.
+- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`; a superset of put that adds missing keys and fields. SimpleStorage only.
 - **(async_)kv_list**: List keys and tags (metadata) in a partition.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
 
-`kv_update` preserves each field's tensor/non-tensor type and tensor dtype. A
-timeout has an unknown outcome because a storage unit may commit after the
-caller stops waiting, so non-idempotent merge operations must not be retried
-blindly. Updates are atomic within one storage-unit request; a batch spanning
-multiple units is not a distributed transaction.
+`kv_update` stores a field the key does not hold yet as given, without calling
+`merge_fn`; a merge must preserve the stored value's tensor/non-tensor type and
+tensor dtype. A timeout has an unknown outcome because a storage unit may commit
+after the caller stops waiting, so non-idempotent merge operations must not be
+retried blindly. Updates are atomic within one storage-unit request; a batch
+spanning multiple units is not a distributed transaction.
+
+To build a field over several updates without exposing intermediate values, pass
+`finalize=False` to every `kv_update` except the last one, which uses the default
+`finalize=True`. Until then the field is stored but not ready, so samplers, the
+dataloader and `kv_batch_get` do not see it.
 
 **Key Features**
 
diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py
index a4fd5b9c..07fd56f8 100644
--- a/tests/e2e/test_data_dump_e2e.py
+++ b/tests/e2e/test_data_dump_e2e.py
@@ -243,6 +243,35 @@ def test_heterogeneous_field_sets_are_grouped(self, tq_system, dump_dir, control
         retrieved = tq.kv_batch_get(keys=["h1"], partition_id=partition_id, select_fields=["routed_experts"])
         _assert_rows_equal(retrieved["routed_experts"], [torch.tensor([7, 8])])
 
+    def test_restore_keeps_unfinalized_fields_hidden(self, tq_system, dump_dir, controller):
+        partition_id = "d_unfinalized"
+        _put_rows(partition_id, ["u0", "u1"])
+
+        def concat(old, new):
+            return torch.cat([old, new])
+
+        tq.kv_update(
+            key="u1",
+            partition_id=partition_id,
+            fields={"response": torch.tensor([1, 2])},
+            merge_fn=concat,
+            finalize=False,
+        )
+
+        tq.dump_data_by_key(dump_dir, ["u0", "u1"], partition_id)
+        assert _read_row_index(dump_dir)["rows"]["u1"]["unfinalized"] == ["response"]
+        ray.get(controller.clear_partition.remote(partition_id))
+        tq.load_data_by_key(dump_dir)
+
+        snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id))
+        row = snapshot.keys_mapping["u1"]
+        assert snapshot.production_status[row, snapshot.field_name_mapping["input_ids"]] == 1
+        assert snapshot.production_status[row, snapshot.field_name_mapping["response"]] == 0
+
+        tq.kv_update(key="u1", partition_id=partition_id, fields={"response": torch.tensor([3])}, merge_fn=concat)
+        retrieved = tq.kv_batch_get(keys="u1", partition_id=partition_id, select_fields="response")
+        _assert_rows_equal(retrieved["response"], [torch.tensor([1, 2, 3])])
+
     def test_other_partitions_are_untouched(self, tq_system, dump_dir, controller):
         """Restoring merges by key, unlike a checkpoint load which replaces everything."""
         # Define test data
@@ -681,10 +710,10 @@ async def publish_all():
         started, ready = [], asyncio.Event()
 
         async def load(shards, target_storage_unit):
-            updates = [{"global_indexes": [i], "field_schema": {}} for i in range(num_updates)]
+            updates = [{"global_indexes": [i], "field_schema": {}, "ready": True} for i in range(num_updates)]
             return {"bytes_read": 0, "updates": updates}
 
-        async def notify(partition_id, global_indexes, field_schema):
+        async def notify(partition_id, global_indexes, field_schema, **_kwargs):
             started.append(global_indexes)
             if len(started) == num_updates:
                 ready.set()
diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index fa5d488b..ed88a995 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1244,25 +1244,58 @@ def test_batch_update_multiple_fields_across_storage_units(self, controller, tq_
         assert "score" not in partition.field_metadata
         assert set(tq_api.kv_list(partition_id=partition_id)[partition_id]) == set(keys)
 
-    def test_rejects_missing_key_unproduced_field_and_incompatible_result(self, tq_api, backend_name):
+    def test_unfinalized_updates_stay_hidden_until_finalize(self, controller, tq_api, backend_name):
         if backend_name != "SimpleStorage":
             pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
 
-        with pytest.raises(ValueError, match="not found"):
-            tq_api.kv_update(
-                key="no_such",
-                partition_id="test_partition",
-                fields={"tokens": torch.tensor([1])},
-                merge_fn=lambda _old, new: new,
-            )
+        partition_id = "test_partition"
+        key = "streamed"
+
+        def concat(old, new):
+            return torch.cat([old, new])
+
+        def tokens_ready():
+            partition = get_controller_partition(controller, partition_id)
+            row, col = partition.keys_mapping[key], partition.field_name_mapping["tokens"]
+            return bool(partition.production_status[row, col])
+
+        tq_api.kv_put(key=key, partition_id=partition_id, fields={"prompt": torch.tensor([7])})
+        for chunk in ([1], [2]):
+            fields = {"tokens": torch.tensor(chunk)}
+            tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False)
+            assert not tokens_ready()
+        assert set(tq_api.kv_batch_get(keys=key, partition_id=partition_id).keys()) == {"prompt"}
+
+        tq_api.kv_update(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([3])}, merge_fn=concat)
+        assert tokens_ready()
+        data = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
+        assert_tensor_equal(data["tokens"][0], torch.tensor([1, 2, 3]))
+
+        # A non-final update hides a field that was already ready.
+        tq_api.kv_update(
+            key=key, partition_id=partition_id, fields={"tokens": torch.tensor([4])}, merge_fn=concat, finalize=False
+        )
+        assert not tokens_ready()
+
+    def test_creates_missing_key_and_field_and_rejects_incompatible_result(self, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        def fail(_old, _new):
+            raise AssertionError("merge_fn must not run for a missing value")
+
+        meta = tq_api.kv_update(
+            key="new_key", partition_id="test_partition", fields={"tokens": torch.tensor([1, 2])}, merge_fn=fail
+        )
+        assert meta.fields == ["tokens"]
+        tq_api.kv_update(
+            key="new_key", partition_id="test_partition", fields={"fresh": torch.tensor([9])}, merge_fn=fail
+        )
+        created = tq_api.kv_batch_get(keys="new_key", partition_id="test_partition")
+        assert_tensor_equal(created["tokens"][0], torch.tensor([1, 2]))
+        assert_tensor_equal(created["fresh"][0], torch.tensor([9]))
+
         tq_api.kv_put(key="sample", partition_id="test_partition", fields={"tokens": torch.tensor([1, 2])})
-        with pytest.raises(ValueError, match="not produced"):
-            tq_api.kv_update(
-                key="sample",
-                partition_id="test_partition",
-                fields={"fresh": torch.tensor([1])},
-                merge_fn=lambda _old, new: new,
-            )
         with pytest.raises(RuntimeError, match="keep dtype"):
             tq_api.kv_update(
                 key="sample",
diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py
index ef8538cc..1a6cced5 100644
--- a/tests/test_controller_data_partitions.py
+++ b/tests/test_controller_data_partitions.py
@@ -1350,6 +1350,18 @@ def test_clear_fields_drops_fieldmeta_when_no_sample_holds_it(self):
         assert "pv" not in partition.field_metadata
         assert partition.global_indexes == {0, 1}
 
+    def test_unfinalized_write_records_metadata_but_stays_hidden(self):
+        partition = self._make_partition()
+        t4 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False}
+        partition.update_production_status([0, 1], ["pv"], {"pv": t4})
+
+        assert partition.update_production_status([1], ["pv"], {"pv": t4}, ready=False)
+        assert partition.field_metadata["pv"].global_indexes == {0, 1}
+        assert partition.scan_data_status(["pv"], task_name="trainer") == [0]
+
+        partition.update_production_status([1], ["pv"], {"pv": t4})
+        assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1]
+
     def test_update_production_status_updates_field_metadata(self):
         """Test that update_production_status correctly updates field_metadata via _update_field_metadata."""
         partition = self._make_partition()
diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py
index b6eecb55..b2de9cfe 100644
--- a/tests/test_data_dump.py
+++ b/tests/test_data_dump.py
@@ -149,7 +149,14 @@ def read(self, size=-1):
     for source, target in [(1, 101), (6, 106)]:
         offset, length = reply.body["row_offsets"][source]
         records.append(
-            {"source_index": source, "target_index": target, "fields": ["x"], "offset": offset, "length": length}
+            {
+                "source_index": source,
+                "target_index": target,
+                "fields": ["x"],
+                "unfinalized": [],
+                "offset": offset,
+                "length": length,
+            }
         )
     loaded = unit._handle_load_rows(
         ZMQMessage.create(
@@ -182,7 +189,14 @@ def read(self, size=-1):
 def test_unit_rejects_invalid_records(unit, tmp_path, problem):
     path = tmp_path / "row.pkl"
     path.write_bytes(pickle.dumps({"global_index": 1, "fields": {"x": "value"}}))
-    record = {"source_index": 1, "target_index": 2, "offset": 0, "length": path.stat().st_size, "fields": ["x"]}
+    record = {
+        "source_index": 1,
+        "target_index": 2,
+        "offset": 0,
+        "length": path.stat().st_size,
+        "fields": ["x"],
+        "unfinalized": [],
+    }
     if problem == "wrong_index":
         record["source_index"] = 3
     elif problem == "truncated":
diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 567722ff..7a97e0cd 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -35,10 +35,6 @@ def _concat(old, new):
     return torch.cat([old, new])
 
 
-def _tensor_schema(*fields, dtype=torch.int64):
-    return {field: {"dtype": dtype, "shape": None, "is_nested": True, "is_non_tensor": False} for field in fields}
-
-
 def test_single_update_batch_wraps_tensor_and_python_values():
     batch = _single_update_batch({"tokens": torch.tensor([4, 5]), "meta": {"step": 1}})
 
@@ -57,12 +53,7 @@ def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
 
-    described = data.apply_update(
-        [0],
-        {"sequence_ids": [torch.tensor([20, 21])]},
-        _concat,
-        _tensor_schema("sequence_ids"),
-    )
+    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat)
 
     assert described["sequence_ids"] == {"dtype": torch.int64, "shapes": [(5,)]}
     assert torch.equal(data.field_data["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21]))
@@ -78,17 +69,28 @@ def fail(_old, _new):
         raise RuntimeError("merge failed")
 
     with pytest.raises(RuntimeError, match="merge failed"):
-        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, _tensor_schema("tokens"))
+        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail)
     assert torch.equal(data.field_data["tokens"][7], original)
 
 
-def test_apply_update_rejects_unproduced_field():
+def test_apply_update_stores_missing_values_without_merging():
     data = StorageUnitData()
     data.put_data({"tokens": [torch.tensor([1])]}, [3])
+    merged = []
+
+    def concat(old, new):
+        merged.append(old)
+        return torch.cat([old, new])
+
+    described = data.apply_update(
+        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat
+    )
 
-    with pytest.raises(ValueError, match="unproduced field"):
-        data.apply_update([3], {"fresh": [torch.tensor([2])]}, lambda _old, new: new, _tensor_schema("fresh"))
-    assert "fresh" not in data.field_data
+    assert len(merged) == 1
+    assert torch.equal(data.field_data["tokens"][3], torch.tensor([1, 2]))
+    assert torch.equal(data.field_data["tokens"][4], torch.tensor([5]))
+    assert data.field_data["fresh"] == {3: {"a": 1}, 4: {"a": 2}}
+    assert described["tokens"]["shapes"] == [(2,), (1,)]
 
 
 @pytest.mark.parametrize(
@@ -104,17 +106,16 @@ def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_f
     data.put_data({"tokens": [original.clone()]}, [0])
 
     with pytest.raises(TypeError, match=error):
-        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, _tensor_schema("tokens"))
+        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn)
     assert torch.equal(data.field_data["tokens"][0], original)
 
 
 def test_apply_update_rejects_tensor_result_for_non_tensor_field():
     data = StorageUnitData()
     data.put_data({"meta": [{"step": 1}]}, [0])
-    schema = {"meta": {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True}}
 
     with pytest.raises(TypeError, match="must remain non-tensor"):
-        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]), schema)
+        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]))
     assert data.field_data["meta"][0] == {"step": 1}
 
 
@@ -126,7 +127,7 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data.put_data({"tokens": [prompt]}, [0])
     assert data.ssd_active_values == 1
 
-    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, _tensor_schema("tokens"))
+    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat)
 
     assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])]))
     assert data.ssd_active_values == 1
@@ -152,7 +153,7 @@ def fail_second_field(encoded_values):
         monkeypatch.setattr(data._ssd_store, "write_values", fail_second_field)
         new_data = {"tokens": [torch.tensor([99])], "mask": [torch.tensor([1])]}
         with pytest.raises(OSError, match="No space left"):
-            data.apply_update([0], new_data, _concat, _tensor_schema("tokens", "mask"))
+            data.apply_update([0], new_data, _concat)
 
         stored = data.get_data(["tokens", "mask"], [0])
         assert torch.equal(stored["tokens"][0], prompt)
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
index 41ff29d4..b7ee8d27 100644
--- a/transfer_queue/client.py
+++ b/transfer_queue/client.py
@@ -497,13 +497,15 @@ async def async_update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> BatchMeta:
-        """Merge new values into produced fields on SimpleStorage units.
+        """Merge new values into fields on SimpleStorage units, storing them where absent.
 
         Args:
-            metadata: Samples to update. The key must already exist.
+            metadata: Samples to update, already registered with the controller.
             values: New values aligned with ``metadata``.
-            merge_fn: Called per sample per field as ``merge_fn(old, new)``.
+            merge_fn: Called per sample per field as ``merge_fn(old, new)`` when a value exists.
+            finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
             The same metadata with field_schema and production status updated.
@@ -519,8 +521,10 @@ async def async_update(
         if not metadata or metadata.size == 0:
             raise ValueError("metadata cannot be none or empty")
 
-        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn)
-        return metadata.apply_field_schema(field_schema)
+        field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, finalize=finalize)
+        if not finalize:
+            metadata.production_status[:] = 0
+        return metadata.apply_field_schema(field_schema, set_all_ready=finalize)
 
     async def async_get_data(self, metadata: BatchMeta) -> TensorDict:
         """Asynchronously fetch data from storage units and organize into TensorDict.
@@ -1638,12 +1642,13 @@ def update(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> BatchMeta:
-        """Synchronously merge new values into produced fields on SimpleStorage units.
+        """Synchronously merge or add values on SimpleStorage units.
 
         See ``async_update``.
         """
-        return self._update(metadata, values, merge_fn)
+        return self._update(metadata, values, merge_fn, finalize=finalize)
 
     def get_data(self, metadata: BatchMeta) -> TensorDict:
         """Synchronously fetch data from storage units and organize into TensorDict.
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index c2ba668f..ba8199ff 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -500,6 +500,7 @@ def update_production_status(
         field_names: list[str],
         field_schema: dict[str, dict[str, Any]],
         custom_backend_meta: dict[int, dict[str, Any]] | None = None,
+        ready: bool = True,
     ) -> bool:
         """
         Update production status for specific samples and fields.
@@ -515,6 +516,8 @@ def update_production_status(
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}
             custom_backend_meta: Optional per-sample per-field
                 custom metadata provided by storage backend
+            ready: Whether consumers may see these cells. False records the stored
+                value (field metadata) but hides it until a finalizing kv_update.
 
         Returns:
             True if update was successful, False on error
@@ -543,7 +546,7 @@ def update_production_status(
             # Update production status
             if self.production_status is not None and global_indices and field_names:
                 field_indices = [self.field_name_mapping.get(f) for f in field_names]
-                self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = 1
+                self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = int(ready)
 
             # Update field metadata
             self._update_field_metadata(global_indices, field_schema, custom_backend_meta)
@@ -1130,6 +1133,7 @@ def update_production_status(
         global_indexes: list[int],
         field_schema: dict[str, dict[str, Any]],
         custom_backend_meta: dict[int, dict[str, Any]] | None = None,
+        ready: bool = True,
     ) -> bool:
         """
         Update production status for specific samples and fields in a partition.
@@ -1140,6 +1144,7 @@ def update_production_status(
             global_indexes: List of sample indices to update
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}
             custom_backend_meta: Optional custom backend metadata
+            ready: False keeps the written cells hidden from consumers (unfinalized kv_update)
 
         Returns:
             True if update was successful, False otherwise
@@ -1150,11 +1155,13 @@ def update_production_status(
             logger.error(f"Partition {partition_id} not found")
             return False
 
-        success = partition.update_production_status(global_indexes, field_names, field_schema, custom_backend_meta)
+        success = partition.update_production_status(
+            global_indexes, field_names, field_schema, custom_backend_meta, ready=ready
+        )
         if success:
             logger.debug(
                 f"[{self.controller_id}]: Updated production status for partition {partition_id}: "
-                f"samples={global_indexes}, fields={field_names}"
+                f"samples={global_indexes}, fields={field_names}, ready={ready}"
             )
         return success
 
@@ -1783,9 +1790,10 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]:
             partition_id: Partition that owns ``keys``.
 
         Returns:
-            ``{"rows": {key: {"global_index", "fields", "tag"}}, "field_schema": {field:
-            {"is_nested", "is_non_tensor"}}}``. ``fields`` is empty for a row with no
-            produced field yet; ``field_schema`` covers the fields of the described rows.
+            ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema":
+            {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` is empty for a row with no
+            produced field yet; ``unfinalized`` lists the stored fields not yet ready (written
+            by ``kv_update(finalize=False)``); ``field_schema`` covers the described rows' fields.
 
         Raises:
             KeyError: The partition or any key does not exist.
@@ -1799,18 +1807,23 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]:
         if missing_keys:
             raise KeyError(f"keys not found in partition {partition_id!r}: {missing_keys}")
 
-        rows: dict[str, dict[str, Any]] = {
-            key: {
+        rows: dict[str, dict[str, Any]] = {}
+        for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True):
+            fields = sorted(
+                field_name
+                for field_name, field_meta in partition.field_metadata.items()
+                if global_index in field_meta.global_indexes
+            )
+            rows[key] = {
                 "global_index": global_index,
-                "fields": sorted(
-                    field_name
-                    for field_name, field_meta in partition.field_metadata.items()
-                    if global_index in field_meta.global_indexes
-                ),
+                "fields": fields,
+                "unfinalized": [
+                    name
+                    for name in fields
+                    if not partition.production_status[global_index, partition.field_name_mapping[name]]
+                ],
                 "tag": partition.custom_meta.get(global_index, {}),
             }
-            for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True)
-        }
         # Only the declared types: row dtypes and shapes live with the storage units.
         field_schema = {
             name: {"is_nested": bool(meta.is_nested), "is_non_tensor": bool(meta.is_non_tensor)}
@@ -2097,6 +2110,7 @@ def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMess
             global_indexes=global_indexes,
             field_schema=params.get("field_schema", {}),
             custom_backend_meta=params.get("custom_backend_meta", {}),
+            ready=params.get("ready", True),
         )
         if success:
             if self._metrics is not None:
diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py
index 80fa53ca..9d30d5b7 100644
--- a/transfer_queue/data_dump.py
+++ b/transfer_queue/data_dump.py
@@ -23,7 +23,7 @@
 
     <dump_dir>/
         dump_info.json                # version, partition, counts
-        row_index.pt                  # key -> {global_index, fields, tag}
+        row_index.pt                  # key -> {global_index, fields, unfinalized, tag}
         shards/
             shard_info.json           # file, unit, row count, source index -> [offset, length]
             shard_<N>_<su_id>.pkl      # independent {global_index, fields} records
@@ -215,6 +215,8 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
                     "key": key,
                     "source_index": source_index,
                     "fields": rows[key]["fields"],
+                    # Older dumps lack this key; every field they hold was ready.
+                    "unfinalized": rows[key].get("unfinalized", []),
                     "offset": offset,
                     "length": length,
                 }
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index b3e0ede4..74f257b9 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -629,22 +629,29 @@ def kv_update(
     partition_id: str,
     fields: dict[str, Any],
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Merge new values into produced fields of one existing key.
+    """Merge new values into the fields of one key, adding the key or fields if missing.
+
+    A superset of ``kv_put``: a missing key is created and a field the key does not
+    hold yet is stored as given, without calling ``merge_fn``.
 
     Args:
-        key: Existing user-specified key. The key is not created if it is missing.
+        key: User-specified key. Created if missing.
         partition_id: Partition that holds the key.
-        fields: New values keyed by an already-produced field name.
-        merge_fn: ``merge_fn(old, new) -> stored``. SimpleStorage runs it once
-            per field on the unit that owns the sample.
+        fields: New values keyed by field name.
+        merge_fn: ``merge_fn(old, new) -> stored`` for fields the key already holds.
+            SimpleStorage runs it once per field on the unit that owns the sample.
+        finalize: If False, the updated fields become (or stay) hidden from consumers,
+            so readers never see an intermediate value. Pass True (the default) on the
+            last update to make them ready.
 
     Returns:
         Metadata for the updated sample.
 
     Raises:
-        ValueError: If the key or a field is missing.
-        TypeError: If ``fields`` is not a non-empty dict or ``merge_fn`` is not callable.
+        TypeError: If ``fields`` is not a non-empty dict, ``merge_fn`` is not callable,
+            or a merge result changes the stored value's type or dtype.
         NotImplementedError: If the storage backend is not SimpleStorage.
 
     Note:
@@ -658,6 +665,7 @@ def kv_update(
             partition_id=partition_id,
             fields=fields,
             merge_fn=merge_fn,
+            finalize=finalize,
         )
     )
 
@@ -667,15 +675,16 @@ def kv_batch_update(
     partition_id: str,
     fields: TensorDict,
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Merge batched values into produced fields of existing keys.
+    """Merge batched values into the fields of keys, adding keys or fields if missing.
 
-    A timeout has an unknown outcome because a unit may commit after the caller
-    stops waiting. Do not blindly retry a merge.
+    See ``kv_update``. A timeout has an unknown outcome because a unit may commit
+    after the caller stops waiting. Do not blindly retry a merge.
     """
     tq_client = _maybe_create_tq_client()
     return tq_client._run_coroutine(
-        async_kv_batch_update(keys=keys, partition_id=partition_id, fields=fields, merge_fn=merge_fn)
+        async_kv_batch_update(keys=keys, partition_id=partition_id, fields=fields, merge_fn=merge_fn, finalize=finalize)
     )
 
 
@@ -1091,13 +1100,15 @@ async def async_kv_update(
     partition_id: str,
     fields: dict[str, Any],
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Asynchronously update fields of an existing key. See ``kv_update``."""
+    """Asynchronously merge or add fields of one key. See ``kv_update``."""
     return await async_kv_batch_update(
         keys=[key],
         partition_id=partition_id,
         fields=_single_update_batch(fields),
         merge_fn=merge_fn,
+        finalize=finalize,
     )
 
 
@@ -1106,8 +1117,9 @@ async def async_kv_batch_update(
     partition_id: str,
     fields: TensorDict,
     merge_fn: Callable[[Any, Any], Any],
+    finalize: bool = True,
 ) -> KVBatchMeta:
-    """Asynchronously merge batched values into produced fields of existing keys."""
+    """Asynchronously merge or add batched fields of keys. See ``kv_update``."""
     if not isinstance(keys, list) or not keys or not all(isinstance(key, str) for key in keys):
         raise TypeError("keys must be a non-empty list of strings")
     if not isinstance(fields, TensorDict) or not fields.keys():
@@ -1117,15 +1129,18 @@ async def async_kv_batch_update(
     if not callable(merge_fn):
         raise TypeError("merge_fn must be callable")
     tq_client = _maybe_create_tq_client()
-    batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False)
+    # Checked up front: missing keys are registered before storage is asked to merge.
+    if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager):
+        raise NotImplementedError(
+            f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; "
+            "it requires the SimpleStorage backend"
+        )
+    batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=True)
 
     if batch_meta.size != len(keys):
-        raise ValueError("Some keys or the partition were not found")
-    missing_fields = sorted(set(fields.keys()) - set(batch_meta.field_names))
-    if missing_fields:
-        raise ValueError(f"Fields are not produced for every key: {missing_fields}")
+        raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys")
 
-    batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn)
+    batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn, finalize=finalize)
 
     return KVBatchMeta(
         keys=keys,
diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py
index 1015172e..02fc1575 100644
--- a/transfer_queue/storage/managers/base.py
+++ b/transfer_queue/storage/managers/base.py
@@ -226,15 +226,17 @@ async def notify_data_update(
         global_indexes: list[int],
         field_schema: dict[str, dict[str, Any]],
         custom_backend_meta: dict[int, dict[str, Any]] | None = None,
+        ready: bool = True,
     ) -> bool:
         """
-        Notify controller that new data is ready.
+        Notify controller that new data is stored.
 
         Args:
             partition_id: Current data partition id.
             global_indexes: Data update related global_indexes.
             field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}.
             custom_backend_meta: Per-field custom_meta for each sample, in {global_index: {field: custom_meta}} format.
+            ready: Whether consumers may see the data now; False for a kv_update with finalize=False.
         """
 
         if not self.controller_info:
@@ -264,6 +266,7 @@ async def notify_data_update(
                 "global_indexes": global_indexes,
                 "field_schema": normalized_field_schema,
                 "custom_backend_meta": custom_backend_meta,
+                "ready": ready,
             },
         ).serialize()
 
@@ -362,13 +365,15 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
-        """Merge new values on the unit that holds each sample.
+        """Merge new values on the unit that holds each sample, or store them if absent.
 
         Args:
             metadata: Samples to update.
             values: New values as a TensorDict aligned with ``metadata``.
-            merge_fn: Called per sample per field as ``merge_fn(old, new)``.
+            merge_fn: Called per sample per field as ``merge_fn(old, new)`` when a value exists.
+            finalize: If False, keep the updated fields hidden from consumers.
 
         Returns:
             field_schema of the stored values, keyed by field name.
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index b5bf88b2..1f90fbfe 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -544,11 +544,12 @@ async def update_data(
         metadata: BatchMeta,
         values: TensorDict,
         merge_fn: Callable[[Any, Any], Any],
+        finalize: bool = True,
     ) -> dict[str, dict[str, Any]]:
-        """Merge values on each owner unit and notify the controller.
+        """Merge or add values on each owner unit and notify the controller.
 
         Each unit reports the values it stored; this assembles one batch-level
-        field_schema so production status stays ready.
+        field_schema. ``finalize`` decides whether the fields become ready.
         """
         logger.debug(f"[{self.storage_manager_id}]: receive update_data request, updating {metadata.size} samples.")
 
@@ -574,7 +575,6 @@ async def update_data(
                         group.global_indexes,
                         storage_data,
                         merge_fn,
-                        {name: metadata.field_schema[name] for name in field_names},
                         target_storage_unit=su_id,
                     ),
                     # A merge reads stored state, so a replay after a lost answer
@@ -603,6 +603,7 @@ async def update_data(
             metadata.partition_ids[0],
             metadata.global_indexes,
             field_schema,
+            ready=finalize,
         )
         if not published:
             raise RuntimeError(
@@ -674,7 +675,6 @@ async def _update_to_single_storage_unit(
         global_indexes: list[int],
         storage_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        field_schema: dict[str, dict[str, Any]],
         target_storage_unit: str,
         socket: zmq.Socket = None,
     ) -> dict[str, dict[str, Any]]:
@@ -683,7 +683,6 @@ async def _update_to_single_storage_unit(
             "global_indexes": global_indexes,
             "data": storage_data,
             "merge_fn": merge_fn,
-            "field_schema": field_schema,
         }
 
         request_msg = ZMQMessage.create(
@@ -1139,7 +1138,9 @@ async def load_rows_by_index(self, partition_id: str, shards: list[dict[str, Any
         # dominate a wide cross-topology load. The notify pool bounds the concurrency.
         await asyncio.gather(
             *(
-                self.notify_data_update(partition_id, update["global_indexes"], update["field_schema"])
+                self.notify_data_update(
+                    partition_id, update["global_indexes"], update["field_schema"], ready=update["ready"]
+                )
                 for result in loaded
                 for update in result["updates"]
             )
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index ca8ad853..a8d56bb5 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -235,11 +235,12 @@ def apply_update(
         global_indexes: list[int],
         new_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
-        field_schema: dict[str, dict[str, Any]],
     ) -> dict[str, dict[str, Any]]:
-        """Merge produced values, validate their types, then write once.
+        """Merge into stored values, store new ones as is, validate, then write once.
 
-        All merge calls and type checks finish before one ``put_data`` call, so
+        A row that has no value for a field takes ``new`` without calling ``merge_fn``,
+        as a put would. A merge result must keep the old value's tensor/non-tensor kind
+        and tensor dtype. All merges and checks finish before one ``put_data`` call, so
         a merge error or incompatible result leaves storage unchanged.
 
         Returns:
@@ -248,23 +249,21 @@ def apply_update(
         computed: dict[str, list] = {}
         for field, new_values in new_data.items():
             stored = self.field_data.get(field, {})
-            missing = [idx for idx in global_indexes if idx not in stored]
-            if missing:
-                raise ValueError(f"Cannot update unproduced field {field!r} for indexes {missing[:20]}")
-            old_values = self.get_data([field], global_indexes)[field]
-            computed[field] = [merge_fn(old, new_values[i]) for i, old in enumerate(old_values)]
-
-            expected = field_schema[field]
-            if expected.get("is_non_tensor"):
-                if any(isinstance(value, torch.Tensor) for value in computed[field]):
-                    raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor")
-            else:
-                for value in computed[field]:
-                    if not isinstance(value, torch.Tensor) or value.dtype != expected["dtype"]:
-                        actual = value.dtype if isinstance(value, torch.Tensor) else type(value).__name__
+            present = [i for i, idx in enumerate(global_indexes) if idx in stored]
+            old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else []
+            values = list(new_values)
+            for i, old in zip(present, old_values, strict=True):
+                merged = merge_fn(old, new_values[i])
+                if isinstance(old, torch.Tensor):
+                    if not isinstance(merged, torch.Tensor) or merged.dtype != old.dtype:
+                        actual = merged.dtype if isinstance(merged, torch.Tensor) else type(merged).__name__
                         raise TypeError(
-                            f"Merge result for tensor field {field!r} must keep dtype {expected['dtype']}, got {actual}"
+                            f"Merge result for tensor field {field!r} must keep dtype {old.dtype}, got {actual}"
                         )
+                elif isinstance(merged, torch.Tensor):
+                    raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor")
+                values[i] = merged
+            computed[field] = values
         self.put_data(computed, global_indexes)
         return {field: _describe_stored_values(values) for field, values in computed.items()}
 
@@ -1243,12 +1242,11 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage:
             )
 
     def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
-        """Merge existing fields, write once, and describe what was stored."""
+        """Merge or add fields, write once, and describe what was stored."""
         try:
             global_indexes = data_parts.body["global_indexes"]
             merge_fn = data_parts.body.get("merge_fn")
             new_data = data_parts.body.get("data")
-            field_schema = data_parts.body.get("field_schema")
 
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update"
@@ -1257,9 +1255,7 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
                     raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}")
                 if not isinstance(new_data, dict) or not new_data:
                     raise TypeError("update data must be a non-empty dict")
-                if not isinstance(field_schema, dict) or set(field_schema) != set(new_data):
-                    raise ValueError("field_schema must describe every update field")
-                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, field_schema)
+                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn)
 
             return ZMQMessage.create(
                 request_type=ZMQRequestType.UPDATE_DATA_RESPONSE,  # type: ignore[arg-type]
@@ -1511,19 +1507,23 @@ def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage:
                                     f, row["offset"], row["length"], row["source_index"], row["fields"]
                                 )
                                 validate_dump_values(fields, shard["field_schema"], row["source_index"])
-                                groups[tuple(row["fields"])].append((row, fields))
+                                groups[(tuple(row["fields"]), tuple(row["unfinalized"]))].append((row, fields))
                                 bytes_read += row["length"]
-                            for signature, rows in groups.items():
+                            for (signature, unfinalized), rows in groups.items():
                                 indexes = [row["target_index"] for row, _ in rows]
                                 values = {name: [fields[name] for _, fields in rows] for name in signature}
-                                schema = select_dump_schema(
-                                    shard["field_schema"],
-                                    [row["source_index"] for row, _ in rows],
-                                    indexes,
-                                    signature,
-                                )
                                 self.storage_data.put_data(values, indexes)
-                                updates.append({"global_indexes": indexes, "field_schema": schema})
+                                source_indexes = [row["source_index"] for row, _ in rows]
+                                finalized = tuple(name for name in signature if name not in unfinalized)
+                                # Publish each field with the readiness it had when dumped.
+                                for names, ready in ((finalized, True), (unfinalized, False)):
+                                    if names:
+                                        schema = select_dump_schema(
+                                            shard["field_schema"], source_indexes, indexes, names
+                                        )
+                                        updates.append(
+                                            {"global_indexes": indexes, "field_schema": schema, "ready": ready}
+                                        )
             logger.info(
                 "[%s]: loaded %s rows (%s bytes)",
                 self.storage_unit_id,
diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py
index 999a8193..ef9e62b4 100644
--- a/tutorial/02_kv_interface.py
+++ b/tutorial/02_kv_interface.py
@@ -215,7 +215,7 @@ def main():
         Key Methods:
         1. (async_)kv_put          - Insert/Update a multi-column sample by key, with optional metadata tag
         2. (async_)kv_batch_put    - Put multiple key-value pairs efficiently in batch
-        3. (async_)kv_update       - Merge new values into produced fields with merge_fn(old, new)
+        3. (async_)kv_update       - Merge into fields with merge_fn(old, new), adding missing keys/fields
         4. (async_)kv_batch_get    - Retrieve samples (by keys), supporting column selection (by fields)
         5. (async_)kv_list         - List keys and tags (metadata) in a partition
         6. (async_)kv_clear        - Remove key-value pairs, or only selected fields, from storage

From 3891aa47d18761d1fa9295da32063f5b915ef2cd Mon Sep 17 00:00:00 2001
From: Azency <azency_1128@outlook.com>
Date: Tue, 13 Oct 2026 20:35:19 +0800
Subject: [PATCH 2/4] [fix] kv_update: type-check new rows, surface unfinalized
 fields, bump dump format

- Storage units check rows kv_update adds against the controller's column type,
  so a dtype conflict fails before anything is stored or marked ready.
- kv_batch_get raises when a selected field is not ready in every key instead of
  returning an empty result.
- Add kv_describe / async_kv_describe to list each key's unfinalized fields.
- Dump format version 4 records unfinalized fields; version 3 dumps still load.

Signed-off-by: Azency <azency_1128@outlook.com>
---
 README.md                                     |  1 +
 docs/data_dump.md                             |  5 +-
 tests/e2e/test_data_dump_e2e.py               | 19 +++++++
 tests/e2e/test_kv_interface_e2e.py            | 33 ++++++++++-
 tests/test_kv_update.py                       | 32 ++++++++---
 transfer_queue/__init__.py                    |  4 ++
 transfer_queue/client.py                      |  7 ++-
 transfer_queue/controller.py                  |  7 +++
 transfer_queue/data_dump.py                   | 10 ++--
 transfer_queue/interface.py                   | 57 ++++++++++++++++++-
 .../managers/simple_storage_manager.py        |  9 +++
 transfer_queue/storage/simple_storage.py      | 39 ++++++++-----
 12 files changed, 194 insertions(+), 29 deletions(-)

diff --git a/README.md b/README.md
index 23367a11..242b577c 100644
--- a/README.md
+++ b/README.md
@@ -118,6 +118,7 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API
 - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields).
 - **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`; a superset of put that adds missing keys and fields. SimpleStorage only.
 - **(async_)kv_list**: List keys and tags (metadata) in a partition.
+- **(async_)kv_describe**: Show each key's stored fields, including those `kv_update(finalize=False)` keeps hidden.
 - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys.
 
 `kv_update` stores a field the key does not hold yet as given, without calling
diff --git a/docs/data_dump.md b/docs/data_dump.md
index f3829b5b..4365ff63 100644
--- a/docs/data_dump.md
+++ b/docs/data_dump.md
@@ -55,7 +55,10 @@ These count application reads, not filesystem read-ahead or physical disk traffi
 
 ## Format and compatibility
 
-Dumps use `format_version: 3`, the only version this build reads:
+Dumps use `format_version: 4`. Version 4 adds each row's `unfinalized` fields, which
+restore keeps hidden. This build also reads version 3 dumps, whose fields restore as
+ready. Builds that only read version 3 reject version 4 dumps instead of exposing
+hidden fields. The layout:
 
 ```text
 dump_info.json
diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py
index 07fd56f8..c60b50bf 100644
--- a/tests/e2e/test_data_dump_e2e.py
+++ b/tests/e2e/test_data_dump_e2e.py
@@ -272,6 +272,25 @@ def concat(old, new):
         retrieved = tq.kv_batch_get(keys="u1", partition_id=partition_id, select_fields="response")
         _assert_rows_equal(retrieved["response"], [torch.tensor([1, 2, 3])])
 
+    def test_restores_a_version_3_dump_as_ready(self, tq_system, dump_dir, controller):
+        partition_id = "d_version3"
+        _put_rows(partition_id, ["w0"])
+        tq.dump_data_by_key(dump_dir, ["w0"], partition_id)
+
+        # Rewrite it as version 3 wrote it: no ``unfinalized`` in the row index.
+        row_index = _read_row_index(dump_dir)
+        del row_index["rows"]["w0"]["unfinalized"]
+        torch.save(row_index, dump_dir / "row_index.pt")
+        info_path = dump_dir / "dump_info.json"
+        info = json.loads(info_path.read_text(encoding="utf-8"))
+        info["format_version"] = 3
+        info_path.write_text(json.dumps(info), encoding="utf-8")
+
+        ray.get(controller.clear_partition.remote(partition_id))
+        tq.load_data_by_key(dump_dir)
+        retrieved = tq.kv_batch_get(keys="w0", partition_id=partition_id, select_fields="input_ids")
+        _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(0)])
+
     def test_other_partitions_are_untouched(self, tq_system, dump_dir, controller):
         """Restoring merges by key, unlike a checkpoint load which replaces everything."""
         # Define test data
diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py
index ed88a995..69c85036 100644
--- a/tests/e2e/test_kv_interface_e2e.py
+++ b/tests/e2e/test_kv_interface_e2e.py
@@ -1262,12 +1262,21 @@ def tokens_ready():
         tq_api.kv_put(key=key, partition_id=partition_id, fields={"prompt": torch.tensor([7])})
         for chunk in ([1], [2]):
             fields = {"tokens": torch.tensor(chunk)}
-            tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False)
+            meta = tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False)
             assert not tokens_ready()
         assert set(tq_api.kv_batch_get(keys=key, partition_id=partition_id).keys()) == {"prompt"}
+        with pytest.raises(ValueError, match="not finalized"):
+            tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
+        with pytest.raises(ValueError, match="not finalized"):
+            tq_api.kv_batch_get_by_meta(meta)
+        row = tq_api.kv_describe(keys=key, partition_id=partition_id)[key]
+        assert (row["fields"], row["unfinalized"]) == (["prompt", "tokens"], ["tokens"])
+        with pytest.raises(RuntimeError, match="not found"):
+            tq_api.kv_describe(keys=[key, "absent"], partition_id=partition_id)
 
         tq_api.kv_update(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([3])}, merge_fn=concat)
         assert tokens_ready()
+        assert tq_api.kv_describe(keys=key, partition_id=partition_id)[key]["unfinalized"] == []
         data = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens")
         assert_tensor_equal(data["tokens"][0], torch.tensor([1, 2, 3]))
 
@@ -1276,6 +1285,7 @@ def tokens_ready():
             key=key, partition_id=partition_id, fields={"tokens": torch.tensor([4])}, merge_fn=concat, finalize=False
         )
         assert not tokens_ready()
+        tq_api.kv_clear(keys=key, partition_id=partition_id)
 
     def test_creates_missing_key_and_field_and_rejects_incompatible_result(self, tq_api, backend_name):
         if backend_name != "SimpleStorage":
@@ -1306,6 +1316,27 @@ def fail(_old, _new):
         unchanged = tq_api.kv_batch_get(keys="sample", partition_id="test_partition", select_fields="tokens")
         assert_tensor_equal(unchanged["tokens"][0], torch.tensor([1, 2]))
 
+    def test_new_row_must_match_column_type_held_by_other_keys(self, tq_api, backend_name):
+        if backend_name != "SimpleStorage":
+            pytest.skip("merge-backed kv_update is implemented only for SimpleStorage")
+
+        def fail(_old, _new):
+            raise AssertionError("merge_fn must not run for a missing value")
+
+        partition_id = "test_partition"
+        # The column is known only through a hidden write on another key.
+        tq_api.kv_update(
+            key="k1", partition_id=partition_id, fields={"reward": torch.tensor([0.5])}, merge_fn=fail, finalize=False
+        )
+        with pytest.raises(RuntimeError, match="keep dtype"):
+            tq_api.kv_update(key="k2", partition_id=partition_id, fields={"reward": torch.tensor([1])}, merge_fn=fail)
+
+        # Nothing was stored for k2, so a compatible value is still a first write.
+        tq_api.kv_update(key="k2", partition_id=partition_id, fields={"reward": torch.tensor([0.25])}, merge_fn=fail)
+        data = tq_api.kv_batch_get(keys="k2", partition_id=partition_id, select_fields="reward")
+        assert_tensor_equal(data["reward"][0], torch.tensor([0.25]))
+        tq_api.kv_clear(keys=["k1", "k2"], partition_id=partition_id)
+
     def test_kv_backend_rejects(self, tq_api, backend_name):
         if backend_name == "SimpleStorage":
             pytest.skip("rejection is the KV-backend contract")
diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 7a97e0cd..09574289 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -53,7 +53,7 @@ def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
 
-    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat)
+    described = data.apply_update([0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat, {})
 
     assert described["sequence_ids"] == {"dtype": torch.int64, "shapes": [(5,)]}
     assert torch.equal(data.field_data["sequence_ids"][0], torch.tensor([10, 11, 12, 20, 21]))
@@ -69,7 +69,7 @@ def fail(_old, _new):
         raise RuntimeError("merge failed")
 
     with pytest.raises(RuntimeError, match="merge failed"):
-        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail)
+        data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, {})
     assert torch.equal(data.field_data["tokens"][7], original)
 
 
@@ -83,7 +83,7 @@ def concat(old, new):
         return torch.cat([old, new])
 
     described = data.apply_update(
-        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat
+        [3, 4], {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, concat, {}
     )
 
     assert len(merged) == 1
@@ -106,7 +106,7 @@ def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_f
     data.put_data({"tokens": [original.clone()]}, [0])
 
     with pytest.raises(TypeError, match=error):
-        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn)
+        data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, {})
     assert torch.equal(data.field_data["tokens"][0], original)
 
 
@@ -115,10 +115,28 @@ def test_apply_update_rejects_tensor_result_for_non_tensor_field():
     data.put_data({"meta": [{"step": 1}]}, [0])
 
     with pytest.raises(TypeError, match="must remain non-tensor"):
-        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]))
+        data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]), {})
     assert data.field_data["meta"][0] == {"step": 1}
 
 
+@pytest.mark.parametrize(
+    "column, value, error",
+    [
+        ({"dtype": torch.float32, "is_non_tensor": False}, torch.tensor([1]), "keep dtype"),
+        ({"dtype": torch.float32, "is_non_tensor": False}, {"not": "a tensor"}, "keep dtype"),
+        ({"dtype": None, "is_non_tensor": True}, torch.tensor([1.0]), "must remain non-tensor"),
+    ],
+)
+def test_apply_update_checks_new_rows_against_column_type_before_writing(column, value, error):
+    data = StorageUnitData()
+    data.put_data({"score": [torch.tensor([0.5])]}, [0])
+
+    with pytest.raises(TypeError, match=error):
+        data.apply_update([0, 1], {"score": [torch.tensor([0.25]), value]}, _concat, {"score": column})
+    assert set(data.field_data["score"]) == {0}
+    assert torch.equal(data.field_data["score"][0], torch.tensor([0.5]))
+
+
 def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data = HybridStorageUnitData(
         storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit"
@@ -127,7 +145,7 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path):
     data.put_data({"tokens": [prompt]}, [0])
     assert data.ssd_active_values == 1
 
-    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat)
+    data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, {})
 
     assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])]))
     assert data.ssd_active_values == 1
@@ -153,7 +171,7 @@ def fail_second_field(encoded_values):
         monkeypatch.setattr(data._ssd_store, "write_values", fail_second_field)
         new_data = {"tokens": [torch.tensor([99])], "mask": [torch.tensor([1])]}
         with pytest.raises(OSError, match="No space left"):
-            data.apply_update([0], new_data, _concat)
+            data.apply_update([0], new_data, _concat, {})
 
         stored = data.get_data(["tokens", "mask"], [0])
         assert torch.equal(stored["tokens"][0], prompt)
diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py
index 6c514b2b..05e63cdf 100644
--- a/transfer_queue/__init__.py
+++ b/transfer_queue/__init__.py
@@ -23,6 +23,7 @@
     async_kv_batch_put,
     async_kv_batch_update,
     async_kv_clear,
+    async_kv_describe,
     async_kv_list,
     async_kv_put,
     async_kv_update,
@@ -36,6 +37,7 @@
     kv_batch_put,
     kv_batch_update,
     kv_clear,
+    kv_describe,
     kv_list,
     kv_put,
     kv_update,
@@ -72,6 +74,7 @@
         "kv_batch_get",
         "kv_batch_get_by_meta",
         "kv_list",
+        "kv_describe",
         "kv_clear",
         "kv_update",
         "async_kv_put",
@@ -80,6 +83,7 @@
         "async_kv_batch_get",
         "async_kv_batch_get_by_meta",
         "async_kv_list",
+        "async_kv_describe",
         "async_kv_clear",
         "async_kv_update",
         "KVBatchMeta",
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
index b7ee8d27..85faa74f 100644
--- a/transfer_queue/client.py
+++ b/transfer_queue/client.py
@@ -972,6 +972,7 @@ async def async_kv_retrieve_meta(
         keys: list[str] | str,
         partition_id: str,
         create: bool = False,
+        include_fields: list[str] | None = None,
         socket: zmq.asyncio.Socket | None = None,
     ) -> BatchMeta:
         """Asynchronously retrieve BatchMeta by user-defined keys.
@@ -983,6 +984,7 @@ async def async_kv_retrieve_meta(
             keys: List of keys to retrieve.
             partition_id: The ID of the logical partition to search for keys.
             create:  If True, automatically create entries for missing keys.
+            include_fields: Known fields to include even if not every key holds them ready.
             socket: ZMQ socket injected by @with_controller_socket.
 
         Returns:
@@ -1008,6 +1010,7 @@ async def async_kv_retrieve_meta(
                     "keys": keys,
                     "partition_id": partition_id,
                     "create": create,
+                    "include_fields": include_fields,
                 },
             )
             return response_msg.body.get("metadata", BatchMeta.empty())
@@ -1137,8 +1140,8 @@ async def async_kv_describe(
         key raises instead of being dropped.
 
         Returns:
-            ``{"rows": {key: {"global_index", "fields", "tag"}}, "field_schema": {field:
-            {"is_nested", "is_non_tensor"}}}``.
+            ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema":
+            {field: {"is_nested", "is_non_tensor"}}}``.
         """
         response = await self._request_controller(
             socket=socket,
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
index ba8199ff..36c34302 100644
--- a/transfer_queue/controller.py
+++ b/transfer_queue/controller.py
@@ -1652,6 +1652,7 @@ def kv_retrieve_meta(
         keys: list[str],
         partition_id: str,
         create: bool = False,
+        include_fields: list[str] | None = None,
     ) -> BatchMeta:
         """
         Retrieve BatchMeta from the controller using a list of keys.
@@ -1661,6 +1662,8 @@ def kv_retrieve_meta(
             keys: List of keys to retrieve from the controller
             partition_id: Partition id to retrieve from the controller
             create: Whether to register new keys if not found.
+            include_fields: Fields whose column schema to include even if not every key holds
+                them ready, so kv_update can type-check the values it adds to them.
 
         Returns:
             metadata: BatchMeta of the requested keys
@@ -1721,6 +1724,9 @@ def kv_retrieve_meta(
         for field_name, col_idx in partition.field_name_mapping.items():
             if col_idx < len(col_mask) and col_mask[col_idx]:
                 data_fields.append(field_name)
+        for field_name in include_fields or []:
+            if field_name in partition.field_metadata and field_name not in data_fields:
+                data_fields.append(field_name)
 
         return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch")
 
@@ -2241,6 +2247,7 @@ def _handle_kv_retrieve_meta_request(self, request_msg: ZMQMessage) -> ZMQMessag
             keys=params["keys"],
             partition_id=params["partition_id"],
             create=params["create"],
+            include_fields=params.get("include_fields"),
         )
         return self._make_response(
             request_msg,
diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py
index 9d30d5b7..ad3791fb 100644
--- a/transfer_queue/data_dump.py
+++ b/transfer_queue/data_dump.py
@@ -42,7 +42,10 @@
 
 logger = get_logger(__name__)
 
-DUMP_FORMAT_VERSION = 3
+DUMP_FORMAT_VERSION = 4
+# Version 3 predates ``unfinalized``; every field it holds was ready. A version-3 reader
+# rejects version 4 instead of restoring hidden fields as ready.
+_READABLE_FORMAT_VERSIONS = (3, DUMP_FORMAT_VERSION)
 SHARD_SUBDIR = "shards"
 
 _DUMP_INFO_FILE = "dump_info.json"
@@ -171,10 +174,10 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
         raise FileNotFoundError(f"{_DUMP_INFO_FILE} not found in {dump_dir}")
     with open(info_path, encoding="utf-8") as f:
         dump_info = json.load(f)
-    if dump_info["format_version"] != DUMP_FORMAT_VERSION:
+    if dump_info["format_version"] not in _READABLE_FORMAT_VERSIONS:
         raise ValueError(
             f"Unsupported dump format version {dump_info['format_version']} in {dump_dir}; "
-            f"this build reads version {DUMP_FORMAT_VERSION}"
+            f"this build reads versions {list(_READABLE_FORMAT_VERSIONS)}"
         )
 
     row_index = _read_row_index(dump_dir)
@@ -215,7 +218,6 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
                     "key": key,
                     "source_index": source_index,
                     "fields": rows[key]["fields"],
-                    # Older dumps lack this key; every field they hold was ready.
                     "unfinalized": rows[key].get("unfinalized", []),
                     "offset": offset,
                     "length": length,
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 74f257b9..10356651 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -539,6 +539,7 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s
     Raises:
         ValueError: If keys or partition are not found
         ValueError: If empty fields exist in any key (sample)
+        ValueError: If a field in ``select_fields`` is not ready in every found key
 
     Example:
         >>> import transfer_queue as tq
@@ -596,6 +597,21 @@ def kv_list(partition_id: str | None = None) -> dict[str, dict[str, Any]]:
     return tq_client._run_coroutine(async_kv_list(partition_id=partition_id))
 
 
+def kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, Any]]:
+    """Describe which fields each key holds, including unfinalized ones.
+
+    See ``async_kv_describe``.
+
+    Example:
+        >>> import transfer_queue as tq
+        >>> tq.init()
+        >>> rows = tq.kv_describe(keys=["sample_1", "sample_2"], partition_id="train")
+        >>> hidden = {key: row["unfinalized"] for key, row in rows.items() if row["unfinalized"]}
+    """
+    tq_client = _maybe_create_tq_client()
+    return tq_client._run_coroutine(async_kv_describe(keys=keys, partition_id=partition_id))
+
+
 def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None:
     """Clear key-value pairs from TransferQueue.
 
@@ -964,6 +980,7 @@ async def async_kv_batch_get(
     Raises:
         ValueError: If keys or partition are not found
         ValueError: If empty fields exist in any key (sample)
+        ValueError: If a field in ``select_fields`` is not ready in every found key
 
     Example:
         >>> import transfer_queue as tq
@@ -989,6 +1006,13 @@ async def async_kv_batch_get(
             fields_to_fetch = [select_fields]
         else:
             fields_to_fetch = select_fields
+        # The meta keeps only fields ready in every key; select_fields would drop the rest silently.
+        not_ready = [f for f in fields_to_fetch if f not in batch_meta.field_names]
+        if not_ready:
+            raise ValueError(
+                f"Fields {not_ready} are not ready in all the requested keys: missing, or written by "
+                "kv_update(finalize=False) and not finalized yet"
+            )
         batch_meta = batch_meta.select_fields(fields_to_fetch)
 
     if not batch_meta.is_ready:
@@ -1040,6 +1064,35 @@ async def async_kv_list(partition_id: str | None = None) -> dict[str, dict[str,
     return partition_info
 
 
+async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, Any]]:
+    """Asynchronously describe which fields each key holds, including unfinalized ones.
+
+    ``kv_batch_get`` returns only fields that are ready in every key, so this is the way
+    to find fields written by ``kv_update(finalize=False)`` and not yet finalized.
+
+    Args:
+        keys: Single key or list of keys to describe.
+        partition_id: Partition containing the keys.
+
+    Returns:
+        ``{key: {"global_index", "fields", "unfinalized", "tag"}}``. ``fields`` lists every
+        stored field of the key and ``unfinalized`` the subset still hidden from consumers.
+
+    Raises:
+        RuntimeError: If the partition or any key does not exist.
+
+    Example:
+        >>> import transfer_queue as tq
+        >>> tq.init()
+        >>> rows = await tq.async_kv_describe(keys="sample_1", partition_id="train")
+        >>> print(rows["sample_1"]["unfinalized"])
+    """
+    keys = [keys] if isinstance(keys, str) else keys
+    tq_client = _maybe_create_tq_client()
+    description = await tq_client.async_kv_describe(list(dict.fromkeys(keys)), partition_id)
+    return description["rows"]
+
+
 async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None:
     """Asynchronously clear key-value pairs from TransferQueue.
 
@@ -1135,7 +1188,9 @@ async def async_kv_batch_update(
             f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; "
             "it requires the SimpleStorage backend"
         )
-    batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=True)
+    batch_meta = await tq_client.async_kv_retrieve_meta(
+        keys=keys, partition_id=partition_id, create=True, include_fields=list(fields.keys())
+    )
 
     if batch_meta.size != len(keys):
         raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys")
diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py
index 1f90fbfe..681125a5 100644
--- a/transfer_queue/storage/managers/simple_storage_manager.py
+++ b/transfer_queue/storage/managers/simple_storage_manager.py
@@ -560,6 +560,12 @@ async def update_data(
                 f"Batch size of values ({values.batch_size[0]}) does not match metadata size ({metadata.size})"
             )
         field_names = list(values.keys())
+        # Fields new to the partition have no column type yet; the rest must keep theirs.
+        column_types = {
+            name: {key: metadata.field_schema[name][key] for key in ("dtype", "is_non_tensor")}
+            for name in field_names
+            if name in metadata.field_schema
+        }
 
         routing = self._group_by_hash(metadata.global_indexes)
         tasks = []
@@ -575,6 +581,7 @@ async def update_data(
                         group.global_indexes,
                         storage_data,
                         merge_fn,
+                        column_types,
                         target_storage_unit=su_id,
                     ),
                     # A merge reads stored state, so a replay after a lost answer
@@ -675,6 +682,7 @@ async def _update_to_single_storage_unit(
         global_indexes: list[int],
         storage_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
+        column_types: dict[str, dict[str, Any]],
         target_storage_unit: str,
         socket: zmq.Socket = None,
     ) -> dict[str, dict[str, Any]]:
@@ -683,6 +691,7 @@ async def _update_to_single_storage_unit(
             "global_indexes": global_indexes,
             "data": storage_data,
             "merge_fn": merge_fn,
+            "column_types": column_types,
         }
 
         request_msg = ZMQMessage.create(
diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index a8d56bb5..abc3b3a9 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -112,6 +112,16 @@ class _SSDValueRef:
 TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0))
 
 
+def _check_value_type(field: str, value: Any, is_tensor: bool, dtype: Any) -> None:
+    """Reject a value that would change its field's tensor/non-tensor kind or tensor dtype."""
+    if not is_tensor:
+        if isinstance(value, torch.Tensor):
+            raise TypeError(f"Value for non-tensor field {field!r} must remain non-tensor")
+    elif not isinstance(value, torch.Tensor) or (dtype is not None and value.dtype != dtype):
+        actual = value.dtype if isinstance(value, torch.Tensor) else type(value).__name__
+        raise TypeError(f"Value for tensor field {field!r} must keep dtype {dtype}, got {actual}")
+
+
 def _describe_stored_values(values: list) -> dict[str, Any]:
     """Describe per-sample stored values for the manager.
 
@@ -235,13 +245,15 @@ def apply_update(
         global_indexes: list[int],
         new_data: dict[str, Any],
         merge_fn: Callable[[Any, Any], Any],
+        column_types: dict[str, dict[str, Any]],
     ) -> dict[str, dict[str, Any]]:
         """Merge into stored values, store new ones as is, validate, then write once.
 
         A row that has no value for a field takes ``new`` without calling ``merge_fn``,
         as a put would. A merge result must keep the old value's tensor/non-tensor kind
-        and tensor dtype. All merges and checks finish before one ``put_data`` call, so
-        a merge error or incompatible result leaves storage unchanged.
+        and tensor dtype; a new row must match ``column_types``, the controller's column
+        type for fields it already knows. All merges and checks finish before one
+        ``put_data`` call, so a merge error or incompatible value leaves storage unchanged.
 
         Returns:
             Per-sample description of the stored values, keyed by field name.
@@ -253,16 +265,14 @@ def apply_update(
             old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else []
             values = list(new_values)
             for i, old in zip(present, old_values, strict=True):
-                merged = merge_fn(old, new_values[i])
-                if isinstance(old, torch.Tensor):
-                    if not isinstance(merged, torch.Tensor) or merged.dtype != old.dtype:
-                        actual = merged.dtype if isinstance(merged, torch.Tensor) else type(merged).__name__
-                        raise TypeError(
-                            f"Merge result for tensor field {field!r} must keep dtype {old.dtype}, got {actual}"
-                        )
-                elif isinstance(merged, torch.Tensor):
-                    raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor")
-                values[i] = merged
+                values[i] = merge_fn(old, new_values[i])
+                _check_value_type(field, values[i], isinstance(old, torch.Tensor), getattr(old, "dtype", None))
+            column = column_types.get(field)
+            if column is not None:
+                merged = set(present)
+                for i, value in enumerate(values):
+                    if i not in merged:
+                        _check_value_type(field, value, not column["is_non_tensor"], column["dtype"])
             computed[field] = values
         self.put_data(computed, global_indexes)
         return {field: _describe_stored_values(values) for field, values in computed.items()}
@@ -1247,6 +1257,7 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
             global_indexes = data_parts.body["global_indexes"]
             merge_fn = data_parts.body.get("merge_fn")
             new_data = data_parts.body.get("data")
+            column_types = data_parts.body.get("column_types")
 
             with limit_pytorch_auto_parallel_threads(
                 target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update"
@@ -1255,7 +1266,9 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage:
                     raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}")
                 if not isinstance(new_data, dict) or not new_data:
                     raise TypeError("update data must be a non-empty dict")
-                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn)
+                if not isinstance(column_types, dict):
+                    raise TypeError("update request must carry column_types")
+                stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, column_types)
 
             return ZMQMessage.create(
                 request_type=ZMQRequestType.UPDATE_DATA_RESPONSE,  # type: ignore[arg-type]

From 63a751d204c640de17228586e6b2d438b2c1a1be Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Tue, 13 Oct 2026 21:16:28 +0800
Subject: [PATCH 3/4] [fix] kv_batch_update: reject duplicate keys before
 registering them

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 tests/test_kv_update.py     | 11 ++++++++++-
 transfer_queue/interface.py | 10 ++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py
index 09574289..94c6edff 100644
--- a/tests/test_kv_update.py
+++ b/tests/test_kv_update.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import errno
 from unittest.mock import MagicMock, patch
 
@@ -21,7 +22,7 @@
 import torch
 from tensordict import TensorDict
 
-from transfer_queue.interface import _single_update_batch
+from transfer_queue.interface import _single_update_batch, async_kv_batch_update
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage.managers.base import KVStorageManager
 from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager
@@ -49,6 +50,14 @@ def test_single_update_batch_rejects_invalid_fields(fields):
         _single_update_batch(fields)
 
 
+@patch("transfer_queue.interface._maybe_create_tq_client", side_effect=AssertionError("must fail before any RPC"))
+def test_batch_update_rejects_duplicate_keys_before_registering(_client):
+    fields = TensorDict({"tokens": torch.tensor([[1], [2], [3]])}, batch_size=3)
+
+    with pytest.raises(ValueError, match=r"duplicated: \['k'\]"):
+        asyncio.run(async_kv_batch_update(["k", "other", "k"], "p", fields, _concat))
+
+
 def test_apply_update_concatenates_and_keeps_field_name():
     data = StorageUnitData()
     data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0])
diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py
index 10356651..c3e40431 100644
--- a/transfer_queue/interface.py
+++ b/transfer_queue/interface.py
@@ -695,8 +695,9 @@ def kv_batch_update(
 ) -> KVBatchMeta:
     """Merge batched values into the fields of keys, adding keys or fields if missing.
 
-    See ``kv_update``. A timeout has an unknown outcome because a unit may commit
-    after the caller stops waiting. Do not blindly retry a merge.
+    See ``kv_update``. ``keys`` must be unique; a repeated key raises ``ValueError``
+    before anything is registered or written. A timeout has an unknown outcome because
+    a unit may commit after the caller stops waiting. Do not blindly retry a merge.
     """
     tq_client = _maybe_create_tq_client()
     return tq_client._run_coroutine(
@@ -1175,6 +1176,11 @@ async def async_kv_batch_update(
     """Asynchronously merge or add batched fields of keys. See ``kv_update``."""
     if not isinstance(keys, list) or not keys or not all(isinstance(key, str) for key in keys):
         raise TypeError("keys must be a non-empty list of strings")
+    # A repeated missing key would get one row per occurrence, orphaning all but the last;
+    # a repeated existing key would merge twice from the same old value, losing one.
+    if len(set(keys)) != len(keys):
+        duplicates = sorted({key for key in keys if keys.count(key) > 1})
+        raise ValueError(f"keys must be unique in one kv_batch_update; duplicated: {duplicates}")
     if not isinstance(fields, TensorDict) or not fields.keys():
         raise TypeError("fields must be a non-empty TensorDict")
     if fields.batch_size != torch.Size([len(keys)]):

From 045bdb0a6d17955b26a23bf02dc69e9b920225b7 Mon Sep 17 00:00:00 2001
From: OutstanderWang <wangweiyanster@gmail.com>
Date: Tue, 13 Oct 2026 21:16:28 +0800
Subject: [PATCH 4/4] style: annotate load-rows updates so mypy accepts mixed
 value types

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
---
 transfer_queue/storage/simple_storage.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py
index abc3b3a9..76fd6821 100644
--- a/transfer_queue/storage/simple_storage.py
+++ b/transfer_queue/storage/simple_storage.py
@@ -1505,7 +1505,7 @@ def _handle_dump_rows(self, request: ZMQMessage) -> ZMQMessage:
 
     def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage:
         """Seek to assigned records and merge them into local storage at current indexes."""
-        updates = []
+        updates: list[dict[str, Any]] = []
         bytes_read = 0
         try:
             with limit_pytorch_auto_parallel_threads(TQ_NUM_THREADS):