Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions transfer_queue/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

check if all the backends can handle delete operations that contain fields that do not exist for one sample

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And we need UT to cover this

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if field_name is None, KVStorageManager will check it and return in advance.

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)

Expand Down
5 changes: 2 additions & 3 deletions transfer_queue/storage/managers/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading