|
|
from typing import Callable, Iterator, Optional, TypeVar |
|
|
|
|
|
from torch.utils.data.datapipes._decorator import functional_datapipe |
|
|
from torch.utils.data.datapipes.datapipe import IterDataPipe |
|
|
from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper |
|
|
from torch.utils.data.datapipes.utils.common import ( |
|
|
_check_unpickable_fn, |
|
|
_deprecation_warning, |
|
|
StreamWrapper, |
|
|
validate_input_col |
|
|
) |
|
|
|
|
|
|
|
|
__all__ = ["FilterIterDataPipe", ] |
|
|
|
|
|
T_co = TypeVar('T_co', covariant=True) |
|
|
|
|
|
|
|
|
@functional_datapipe('filter') |
|
|
class FilterIterDataPipe(IterDataPipe[T_co]): |
|
|
r""" |
|
|
Filters out elements from the source datapipe according to input ``filter_fn`` (functional name: ``filter``). |
|
|
|
|
|
Args: |
|
|
datapipe: Iterable DataPipe being filtered |
|
|
filter_fn: Customized function mapping an element to a boolean. |
|
|
drop_empty_batches (Deprecated): By default, drops a batch if it is empty after filtering instead of keeping an empty list |
|
|
input_col: Index or indices of data which ``filter_fn`` is applied, such as: |
|
|
|
|
|
- ``None`` as default to apply ``filter_fn`` to the data directly. |
|
|
- Integer(s) is used for list/tuple. |
|
|
- Key(s) is used for dict. |
|
|
|
|
|
Example: |
|
|
>>> # xdoctest: +SKIP |
|
|
>>> from torchdata.datapipes.iter import IterableWrapper |
|
|
>>> def is_even(n): |
|
|
... return n % 2 == 0 |
|
|
>>> dp = IterableWrapper(range(5)) |
|
|
>>> filter_dp = dp.filter(filter_fn=is_even) |
|
|
>>> list(filter_dp) |
|
|
[0, 2, 4] |
|
|
""" |
|
|
datapipe: IterDataPipe |
|
|
filter_fn: Callable |
|
|
drop_empty_batches: bool |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
datapipe: IterDataPipe, |
|
|
filter_fn: Callable, |
|
|
drop_empty_batches: Optional[bool] = None, |
|
|
input_col=None, |
|
|
) -> None: |
|
|
super().__init__() |
|
|
self.datapipe = datapipe |
|
|
|
|
|
_check_unpickable_fn(filter_fn) |
|
|
self.filter_fn = filter_fn |
|
|
|
|
|
if drop_empty_batches is None: |
|
|
drop_empty_batches = True |
|
|
else: |
|
|
_deprecation_warning( |
|
|
type(self).__name__, |
|
|
deprecation_version="1.12", |
|
|
removal_version="1.14", |
|
|
old_argument_name="drop_empty_batches", |
|
|
) |
|
|
self.drop_empty_batches = drop_empty_batches |
|
|
|
|
|
self.input_col = input_col |
|
|
validate_input_col(filter_fn, input_col) |
|
|
|
|
|
def _apply_filter_fn(self, data) -> bool: |
|
|
if self.input_col is None: |
|
|
return self.filter_fn(data) |
|
|
elif isinstance(self.input_col, (list, tuple)): |
|
|
args = tuple(data[col] for col in self.input_col) |
|
|
return self.filter_fn(*args) |
|
|
else: |
|
|
return self.filter_fn(data[self.input_col]) |
|
|
|
|
|
def __iter__(self) -> Iterator[T_co]: |
|
|
for data in self.datapipe: |
|
|
filtered = self._returnIfTrue(data) |
|
|
if self._isNonEmpty(filtered): |
|
|
yield filtered |
|
|
else: |
|
|
StreamWrapper.close_streams(data) |
|
|
|
|
|
def _returnIfTrue(self, data): |
|
|
condition = self._apply_filter_fn(data) |
|
|
|
|
|
if df_wrapper.is_column(condition): |
|
|
|
|
|
result = [] |
|
|
for idx, mask in enumerate(df_wrapper.iterate(condition)): |
|
|
if mask: |
|
|
result.append(df_wrapper.get_item(data, idx)) |
|
|
if len(result): |
|
|
return df_wrapper.concat(result) |
|
|
else: |
|
|
return None |
|
|
|
|
|
if not isinstance(condition, bool): |
|
|
raise ValueError("Boolean output is required for `filter_fn` of FilterIterDataPipe, got", type(condition)) |
|
|
if condition: |
|
|
return data |
|
|
|
|
|
def _isNonEmpty(self, data): |
|
|
if df_wrapper.is_dataframe(data): |
|
|
return True |
|
|
r = data is not None and \ |
|
|
not (isinstance(data, list) and len(data) == 0 and self.drop_empty_batches) |
|
|
return r |
|
|
|
