Repository: huggingface/datasets
File: src/datasets/iterable_dataset.py (Python)

Code before:
                self.ex_iterable, batch_size=self.batch_size, formatter=formatter, batched=self.batched
            )

        if self.batched:
            if self._state_dict:
                self._state_dict["previous_state"] = self.ex_iterable.state_dict()
                self._state_dict["num_examples_since_previous_state"] = 0
                self._state_dict["previous_state_example_idx"] = current_idx
            for combined_key, batch in batched_examples_iterator:
                if batch:
                    batch_len = len(batch[next(iter(batch))])
                else:
                    batch_len = 0
                # then compute the mask for the batch
                inputs = batch
                function_args = [inputs] if self.input_columns is None else [inputs[col] for col in self.input_columns]
                if self.with_indices:
                    function_args.append([current_idx + i for i in range(batch_len)])
                mask = self.function(*function_args, **self.fn_kwargs)
                # yield one example at a time from the batch
                examples = _batch_to_examples(batch)
                # TODO: nicer way to handle keys?
                if not self.formatting:
                    keys = combined_key.split("_")
                else:
                    keys = [combined_key] * len(mask)
                for key, example, to_keep in zip(keys, examples, mask):
                    current_idx += 1
                    if self._state_dict:
                        self._state_dict["num_examples_since_previous_state"] += 1
                    if num_examples_to_skip > 0:
                        num_examples_to_skip -= 1
                        continue
                    if to_keep:
                        yield key, example
                if self._state_dict:
                    self._state_dict["previous_state"] = self.ex_iterable.state_dict()
                    self._state_dict["num_examples_since_previous_state"] = 0
                    self._state_dict["previous_state_example_idx"] = current_idx
        else:
            for key, example in batched_examples_iterator:
                # If not batched, we can apply the filtering function direcly
                example = dict(example)
                inputs = example
                function_args = [inputs] if self.input_columns is None else [inputs[col] for col in self.input_columns]
                if self.with_indices:
                    function_args.append(current_idx)
                to_keep = self.function(*function_args, **self.fn_kwargs)
                current_idx += 1
                if self._state_dict:
                    self._state_dict["previous_state_example_idx"] += 1

Review comment (suggestion):
this should be fine
```suggestion
                keys = [combined_key] * len(mask)
```

Code after:
        return self.ex_iterable.num_shards


class FilteredExamplesIterable(_BaseExamplesIterable):
    def __init__(
        self,
        ex_iterable: _BaseExamplesIterable,
        function: Callable,
        with_indices: bool = False,
        input_columns: Optional[List[str]] = None,
        batched: bool = False,
        batch_size: Optional[int] = 1000,
        fn_kwargs: Optional[dict] = None,
        formatting: Optional["FormattingConfig"] = None,
    ):
        super().__init__()
        self.ex_iterable = ex_iterable
        self.function = function
        self.batched = batched
        self.batch_size = batch_size
        self.with_indices = with_indices
        self.input_columns = input_columns
        self.fn_kwargs = fn_kwargs or {}
        self.formatting = formatting  # required for iter_arrow
        # sanity checks
        if formatting and formatting.format_type == "arrow":
            # batch_size should match for iter_arrow
            if not isinstance(ex_iterable, RebatchedArrowExamplesIterable):
                raise ValueError(
                    "The Arrow-formatted FilteredExamplesIterable has underlying iterable"
                    f"that is a {type(ex_iterable).__name__} instead of a RebatchedArrowExamplesIterable."
                )
            elif ex_iterable.batch_size != (batch_size if batched else 1):
                raise ValueError(
                    f"The Arrow-formatted FilteredExamplesIterable has batch_size={batch_size if batched else 1} which is"
                    f"different from {ex_iterable.batch_size=} from its underlying iterable."
                )

    @property
    def iter_arrow(self):
        if self.formatting and self.formatting.format_type == "arrow":
            return self._iter_arrow

    @property
    def is_typed(self):
        return self.ex_iterable.is_typed

    @property
    def features(self):
        return self.ex_iterable.features