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

Code before:
                    _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."
                            ) from None
                        if update_data:
                            if i == 0:
                                buf_writer, writer, tmp_file = init_buffer_and_writer()
                                stack.enter_context(writer)
                            if isinstance(batch, pa.Table):
                                writer.write_table(batch)
                            elif isinstance(batch, pd.DataFrame):
                                writer.write_table(pa.Table.from_pandas(batch))
                            elif config.POLARS_AVAILABLE and isinstance(batch, pl.DataFrame):
                                writer.write_table(batch.to_arrow())
                            else:
                                writer.write_batch(batch)
                        num_examples_progress_update += num_examples_in_batch
                        if time.time() > _time + config.PBAR_REFRESH_TIME_INTERVAL:
                            _time = time.time()
                            yield rank, False, num_examples_progress_update
                            num_examples_progress_update = 0
                if update_data and writer is not None:
                    writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file
            except (Exception, KeyboardInterrupt):
                yield rank, False, num_examples_progress_update
                if update_data:
                    if writer is not None:
                        writer.finalize()
                    if tmp_file is not None:
                        tmp_file.close()
                        if os.path.exists(tmp_file.name):
                            os.remove(tmp_file.name)
                raise

        yield rank, False, num_examples_progress_update
        if update_data and tmp_file is not None:
            tmp_file.close()
            shutil.move(tmp_file.name, cache_file_name)

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

Code after:
                            _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."
                            ) from None
                        if update_data:
                            if i == 0:
                                buf_writer, writer, tmp_file = init_buffer_and_writer()
                                stack.enter_context(writer)
                            if isinstance(batch, pa.Table):
                                writer.write_table(batch)
                            elif isinstance(batch, pd.DataFrame):
                                writer.write_table(pa.Table.from_pandas(batch))
                            elif (
                                config.POLARS_AVAILABLE and "polars" in sys.modules and isinstance(batch, pl.DataFrame)
                            ):
                                writer.write_table(batch.to_arrow())
                            else:
                                writer.write_batch(batch)
                        num_examples_progress_update += num_examples_in_batch
                        if time.time() > _time + config.PBAR_REFRESH_TIME_INTERVAL:
                            _time = time.time()
                            yield rank, False, num_examples_progress_update
                            num_examples_progress_update = 0
                if update_data and writer is not None:
                    writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file
            except (Exception, KeyboardInterrupt):
                yield rank, False, num_examples_progress_update
                if update_data:
                    if writer is not None:
                        writer.finalize()
                    if tmp_file is not None:
                        tmp_file.close()
                        if os.path.exists(tmp_file.name):
                            os.remove(tmp_file.name)