Merge pull request #1115 from misrasaurabh1/codeflash/optimize-Batch.remove_by_indices-m7dwflwc

⚡️ Speed up method `Batch.remove_by_indices` by 38%

Author: PawelPeczek-Roboflow

inference/core/workflows/execution_engine/entities/base.py

@@ -132,11 +132,7 @@ def init(
 
         return cls(content=content, indices=indices)
 
-    def __init__(
-        self,
-        content: List[B],
-        indices: Optional[List[Tuple[int, ...]]],
-    ):
+    def __init__(self, content: List[B], indices: Optional[List[Tuple[int, ...]]]):
         self._content = content
         self._indices = indices
 
@@ -157,20 +153,19 @@ def __iter__(self) -> Iterator[B]:
         yield from self._content
 
     def remove_by_indices(self, indices_to_remove: Set[tuple]) -> "Batch":
-        content, new_indices = [], []
-        for index, element in self.iter_with_indices():
-            if index in indices_to_remove:
-                continue
-            content.append(element)
-            new_indices.append(index)
-        return Batch(
-            content=content,
-            indices=new_indices,
-        )
+        filtered_content = [
+            element
+            for index, element in zip(self._indices, self._content)
+            if index not in indices_to_remove
+        ]
+        filtered_indices = [
+            index for index in self._indices if index not in indices_to_remove
+        ]
+
+        return Batch(content=filtered_content, indices=filtered_indices)
 
     def iter_with_indices(self) -> Iterator[Tuple[Tuple[int, ...], B]]:
-        for index, element in zip(self._indices, self._content):
-            yield index, element
+        return zip(self._indices, self._content)
 
     def broadcast(self, n: int) -> "Batch":
         if n <= 0: