diff --git a/transfer_queue/client.py b/transfer_queue/client.py index bd1e57b..3524058 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -475,6 +475,19 @@ async def async_clear_samples(self, metadata: BatchMeta): if not self._controller: raise RuntimeError("No controller registered") + # If field_names is empty, + # query the controller for the partition's fields before clearing meta. + if not metadata.field_names: + merged_field_schema: dict = {} + for partition_id in dict.fromkeys(metadata.partition_ids): + partition_meta = await self._get_partition_meta(partition_id) + if partition_meta: + merged_field_schema.update(partition_meta.field_schema) + if merged_field_schema: + metadata = metadata.copy() + metadata.field_schema = merged_field_schema + metadata._field_names = sorted(merged_field_schema.keys()) + # Clear the controller metadata await self._clear_meta_in_controller(metadata) diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index f4d545d..2753574 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -695,9 +695,8 @@ async def clear_data(self, metadata: BatchMeta) -> None: """Remove stored data associated with the given metadata.""" if not metadata.field_names: - raise RuntimeError( - "Fail to clear_data for key-value based backends due to lack of `field_names` in BatchMeta" - ) + logger.debug("clear_data called with no field_names; nothing to clear for KV-based backend.") + return keys = self._generate_keys(metadata.field_names, metadata.global_indexes) _, _, custom_meta = self._get_shape_type_custom_backend_meta_list(metadata)