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

Code before:
                    else islice(iterator, self.batch_size - 1)
                )
                key_examples_list = [(key, example)] + list(iterator_batch)
                keys, examples = zip(*key_examples_list)
                if (
                    self.drop_last_batch
                    and self.batch_size is not None
                    and self.batch_size > 0
                    and len(examples) < self.batch_size
                ):  # ignore last batch
                    return
                batch = _examples_to_batch(examples)
                batch = format_dict(batch) if format_dict else batch
                # then apply the transform
                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(len(key_examples_list))])
                inputs_to_merge = dict(batch)
                processed_inputs = self.function(*function_args, **self.fn_kwargs)
                # 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_batch = {**inputs_to_merge, **processed_inputs}
                if transformed_batch:
                    first_col = next(iter(transformed_batch))
                    bad_cols = [
                        col
                        for col in transformed_batch
                        if len(transformed_batch[col]) != len(transformed_batch[first_col])
                    ]
                    if bad_cols:
                        raise ValueError(
                            f"Column lengths mismatch: columns {bad_cols} have length {[len(transformed_batch[col]) for col in bad_cols]} while {first_col} has length {len(transformed_batch[first_col])}."
                        )
                # the new key is the concatenation of the examples keys from the batch
                new_key = "_".join(str(key) for key in keys)
                # yield one example at a time from the transformed batch
                for example in _batch_to_examples(transformed_batch):
                    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
                    yield new_key, example
                if self._state_dict:

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

Code after:
                    else islice(iterator, self.batch_size - 1)
                )
                key_examples_list = [(key, example)] + list(iterator_batch)
                keys, examples = zip(*key_examples_list)
                if (
                    self.drop_last_batch
                    and self.batch_size is not None
                    and self.batch_size > 0
                    and len(examples) < self.batch_size
                ):  # ignore last batch
                    return
                batch = _examples_to_batch(examples)
                batch = format_dict(batch) if format_dict else batch
                # then apply the transform
                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(len(key_examples_list))])
                inputs_to_merge = dict(batch)
                processed_inputs = self.function(*function_args, **self.fn_kwargs)
                # 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_batch = {**inputs_to_merge, **processed_inputs}
                if transformed_batch:
                    first_col = next(iter(transformed_batch))
                    bad_cols = [
                        col
                        for col in transformed_batch
                        if len(transformed_batch[col]) != len(transformed_batch[first_col])
                    ]
                    if bad_cols:
                        raise ValueError(
                            f"Column lengths mismatch: columns {bad_cols} have length {[len(transformed_batch[col]) for col in bad_cols]} while {first_col} has length {len(transformed_batch[first_col])}."
                        )
                # the new key is the concatenation of the examples keys from the batch
                new_key = "_".join(str(key) for key in keys)
                # yield one example at a time from the transformed batch
                for example in _batch_to_examples(transformed_batch):
                    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
                    yield new_key, example
                if self._state_dict: