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

Code before:
        with contextlib.ExitStack() as stack:
            try:
                arrow_formatted_shard = shard.with_format("arrow")

                # Loop over single examples or batches and write to buffer/file if examples are to be updated
                if not batched:
                    shard_iterable = enumerate(arrow_formatted_shard)
                else:
                    num_rows = len(shard) if not drop_last_batch else len(shard) // batch_size * batch_size
                    shard_iterable = zip(
                        range(0, num_rows, batch_size),
                        arrow_formatted_shard.iter(batch_size, drop_last_batch=drop_last_batch),
                    )
                if not batched:
                    _time = time.time()
                    for i, example in shard_iterable:
                        example = apply_function_on_filtered_inputs(example, i, offset=offset)
                        if update_data:
                            if i == 0:
                                buf_writer, writer, tmp_file = init_buffer_and_writer()
                                stack.enter_context(writer)
                            if isinstance(example, pa.Table):
                                writer.write_row(example)
                            elif isinstance(example, pd.DataFrame):
                                writer.write_row(pa.Table.from_pandas(example))
                            elif config.POLARS_AVAILABLE and isinstance(example, pl.DataFrame):
                                writer.write_table(example.to_arrow())
                            else:
                                writer.write(example)
                        num_examples_progress_update += 1
                        if time.time() > _time + config.PBAR_REFRESH_TIME_INTERVAL:
                            _time = time.time()
                            yield rank, False, num_examples_progress_update
                            num_examples_progress_update = 0
                else:
                    _time = time.time()
                    for i, batch in shard_iterable:
                        num_examples_in_batch = len(batch)
                        indices = list(
                            range(*(slice(i, i + batch_size).indices(shard.num_rows)))
                        )  # Something simpler?
                        try:
                            batch = apply_function_on_filtered_inputs(
                                batch,
                                indices,
                                check_same_num_examples=len(shard.list_indexes()) > 0,
                                offset=offset,
                            )
                        except NumExamplesMismatchError:
                            raise DatasetTransformationNotAllowedError(
                                "Using `.map` in batched mode on a dataset with attached indexes is allowed only if it doesn't create or remove existing examples. You can first run `.drop_index() to remove your index and then re-add it."

Review comment (suggestion):
```suggestion
                            elif config.POLARS_AVAILABLE and "polars" in sys.modules and isinstance(example, pl.DataFrame):
```

Code after:
        with contextlib.ExitStack() as stack:
            try:
                arrow_formatted_shard = shard.with_format("arrow")

                # Loop over single examples or batches and write to buffer/file if examples are to be updated
                if not batched:
                    shard_iterable = enumerate(arrow_formatted_shard)
                else:
                    num_rows = len(shard) if not drop_last_batch else len(shard) // batch_size * batch_size
                    shard_iterable = zip(
                        range(0, num_rows, batch_size),
                        arrow_formatted_shard.iter(batch_size, drop_last_batch=drop_last_batch),
                    )
                if not batched:
                    _time = time.time()
                    for i, example in shard_iterable:
                        example = apply_function_on_filtered_inputs(example, i, offset=offset)
                        if update_data:
                            if i == 0:
                                buf_writer, writer, tmp_file = init_buffer_and_writer()
                                stack.enter_context(writer)
                            if isinstance(example, pa.Table):
                                writer.write_row(example)
                            elif isinstance(example, pd.DataFrame):
                                writer.write_row(pa.Table.from_pandas(example))
                            elif (
                                config.POLARS_AVAILABLE
                                and "polars" in sys.modules
                                and isinstance(example, pl.DataFrame)
                            ):
                                writer.write_row(example.to_arrow())
                            else:
                                writer.write(example)
                        num_examples_progress_update += 1
                        if time.time() > _time + config.PBAR_REFRESH_TIME_INTERVAL:
                            _time = time.time()
                            yield rank, False, num_examples_progress_update
                            num_examples_progress_update = 0
                else:
                    _time = time.time()
                    for i, batch in shard_iterable:
                        num_examples_in_batch = len(batch)
                        indices = list(
                            range(*(slice(i, i + batch_size).indices(shard.num_rows)))
                        )  # Something simpler?
                        try:
                            batch = apply_function_on_filtered_inputs(
                                batch,
                                indices,
                                check_same_num_examples=len(shard.list_indexes()) > 0,
                                offset=offset,