From c257792eaaaf0ec226dbe1bcc1d22277b6de61e8 Mon Sep 17 00:00:00 2001 From: Mehant Kammakomati Date: Tue, 9 Jul 2024 17:01:09 +0530 Subject: [PATCH] use generator Co-authored-by: Quentin Lhoest <42851186+lhoestq@users.noreply.github.com> --- src/datasets/packaged_modules/arrow/arrow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/datasets/packaged_modules/arrow/arrow.py b/src/datasets/packaged_modules/arrow/arrow.py index 60c7634adb20..2613fa058bce 100644 --- a/src/datasets/packaged_modules/arrow/arrow.py +++ b/src/datasets/packaged_modules/arrow/arrow.py @@ -67,7 +67,7 @@ def _generate_tables(self, files): batches = pa.ipc.open_stream(f) except pa.lib.ArrowInvalid: reader = pa.ipc.open_file(f) - batches = [reader.get_batch(i) for i in range(reader.num_record_batches)] + batches = (reader.get_batch(i) for i in range(reader.num_record_batches)) for batch_idx, record_batch in enumerate(batches): pa_table = pa.Table.from_batches([record_batch]) # Uncomment for debugging (will print the Arrow table size and elements)