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

Code before:
                        num_examples_to_skip -= 1
                        continue
                    yield new_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 iterator:
                # If not batched, we can apply the transform and yield the example directly
                # first copy the example, since we might drop some keys
                example = dict(example)
                example = format_dict(example) if format_dict else example
                # then apply the transform
                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)
                processed_inputs = self.function(*function_args, **self.fn_kwargs)
                inputs_to_merge = dict(example)
                # this logic mimics the one in Dataset.map
                if self.remove_columns:
                    for c in self.remove_columns:
                        if c in inputs_to_merge:
                            del inputs_to_merge[c]
                        if processed_inputs == inputs and c in processed_inputs:
                            del processed_inputs[c]
                transformed_example = {**inputs_to_merge, **processed_inputs}
                current_idx += 1
                if self._state_dict:
                    self._state_dict["previous_state_example_idx"] += 1
                yield key, transformed_example

    def _iter_arrow(self, max_chunksize: Optional[int] = None) -> Iterator[Tuple[Key, pa.Table]]:
        if self.ex_iterable.iter_arrow:
            iterator = self.ex_iterable.iter_arrow()
        else:
            iterator = _convert_to_arrow(
                self.ex_iterable,
                batch_size=self.batch_size if self.batched else 1,
                drop_last_batch=self.drop_last_batch,
            )
        if self._state_dict and self._state_dict["previous_state"]:
            self.ex_iterable.load_state_dict(self._state_dict["previous_state"])
            num_examples_to_skip = self._state_dict["num_examples_since_previous_state"]
        else:
            num_examples_to_skip = 0
        if self._state_dict and max_chunksize is not None:
            self._state_dict["previous_state"] = self.ex_iterable.state_dict()
            self._state_dict["num_examples_since_previous_state"] = 0
        current_idx = self._state_dict["previous_state_example_idx"] if self._state_dict else 0

Review comment (suggestion):
```suggestion
                        if processed_inputs is inputs and c in processed_inputs:
```

Code after:
                        num_examples_to_skip -= 1
                        continue
                    yield new_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 iterator:
                # If not batched, we can apply the transform and yield the example directly
                # first copy the example, since we might drop some keys
                example = dict(example)
                example = format_dict(example) if format_dict else example
                # then apply the transform
                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)
                processed_inputs = self.function(*function_args, **self.fn_kwargs)
                inputs_to_merge = dict(example)
                # this logic mimics the one in Dataset.map
                if self.remove_columns:
                    for c in self.remove_columns:
                        if c in inputs_to_merge:
                            del inputs_to_merge[c]
                        if processed_inputs is inputs and c in processed_inputs:
                            del processed_inputs[c]
                transformed_example = {**inputs_to_merge, **processed_inputs}
                current_idx += 1
                if self._state_dict:
                    self._state_dict["previous_state_example_idx"] += 1
                yield key, transformed_example

    def _iter_arrow(self, max_chunksize: Optional[int] = None) -> Iterator[Tuple[Key, pa.Table]]:
        if self.ex_iterable.iter_arrow:
            iterator = self.ex_iterable.iter_arrow()
        else:
            iterator = _convert_to_arrow(
                self.ex_iterable,
                batch_size=self.batch_size if self.batched else 1,
                drop_last_batch=self.drop_last_batch,
            )
        if self._state_dict and self._state_dict["previous_state"]:
            self.ex_iterable.load_state_dict(self._state_dict["previous_state"])
            num_examples_to_skip = self._state_dict["num_examples_since_previous_state"]
        else:
            num_examples_to_skip = 0
        if self._state_dict and max_chunksize is not None:
            self._state_dict["previous_state"] = self.ex_iterable.state_dict()
            self._state_dict["num_examples_since_previous_state"] = 0
        current_idx = self._state_dict["previous_state_example_idx"] if self._state_dict else 0