Repository: huggingface/datasets
File: tests/test_arrow_dataset.py (Python)

Code before:
                for col_name in dset.column_names:
                    self.assertEqual(len(dset_to_pandas[col_name]), len(dset))

                # With index mapping
                with dset.select([1, 0, 3]) as dset:
                    dset_to_pandas = dset.to_pandas()
                    self.assertIsInstance(dset_to_pandas, pd.DataFrame)
                    self.assertEqual(len(dset_to_pandas), 3)
                    self.assertListEqual(sorted(dset_to_pandas.columns), sorted(dset.column_names))

                    for col_name in dset.column_names:
                        self.assertEqual(len(dset_to_pandas[col_name]), dset.num_rows)

    @require_polars
    def test_to_polars(self, in_memory):
        import polars as pl

        with tempfile.TemporaryDirectory() as tmp_dir:
            # Batched
            with self._create_dummy_dataset(in_memory, tmp_dir, multiple_columns=True) as dset:
                batch_size = dset.num_rows - 1
                to_polars_generator = dset.to_polars(batched=True, batch_size=batch_size)

                for batch in to_polars_generator:
                    self.assertIsInstance(batch, pl.DataFrame)
                    self.assertListEqual(sorted(batch.columns), sorted(dset.column_names))
                    for col_name in dset.column_names:
                        self.assertLessEqual(len(batch[col_name]), batch_size)

                # Full
                dset_to_polars = dset.to_polars()
                self.assertIsInstance(dset_to_polars, pl.DataFrame)
                self.assertListEqual(sorted(dset_to_polars.columns), sorted(dset.column_names))
                for col_name in dset.column_names:
                    self.assertEqual(len(dset_to_polars[col_name]), len(dset))

                # With index mapping
                with dset.select([1, 0, 3]) as dset:
                    dset_to_polars = dset.to_polars()
                    self.assertIsInstance(dset_to_polars, pl.DataFrame)
                    self.assertEqual(len(dset_to_polars), 3)
                    self.assertListEqual(sorted(dset_to_polars.columns), sorted(dset.column_names))

                    for col_name in dset.column_names:
                        self.assertEqual(len(dset_to_polars[col_name]), dset.num_rows)

    def test_to_parquet(self, in_memory):
        with tempfile.TemporaryDirectory() as tmp_dir:
            # File path argument
            with self._create_dummy_dataset(in_memory, tmp_dir, multiple_columns=True) as dset:
                file_path = os.path.join(tmp_dir, "test_path.parquet")

Review comment (performance):
```suggestion
                        self.assertLessEqual(len(batch[col_name]), batch_size)
                del to_polars_generator
```

Code after:
                for col_name in dset.column_names:
                    self.assertEqual(len(dset_to_pandas[col_name]), len(dset))

                # With index mapping
                with dset.select([1, 0, 3]) as dset:
                    dset_to_pandas = dset.to_pandas()
                    self.assertIsInstance(dset_to_pandas, pd.DataFrame)
                    self.assertEqual(len(dset_to_pandas), 3)
                    self.assertListEqual(sorted(dset_to_pandas.columns), sorted(dset.column_names))

                    for col_name in dset.column_names:
                        self.assertEqual(len(dset_to_pandas[col_name]), dset.num_rows)

    @require_polars
    def test_to_polars(self, in_memory):
        with tempfile.TemporaryDirectory() as tmp_dir:
            # Batched
            with self._create_dummy_dataset(in_memory, tmp_dir, multiple_columns=True) as dset:
                batch_size = dset.num_rows - 1
                to_polars_generator = dset.to_polars(batched=True, batch_size=batch_size)

                for batch in to_polars_generator:
                    self.assertIsInstance(batch, sys.modules["polars"].DataFrame)
                    self.assertListEqual(sorted(batch.columns), sorted(dset.column_names))
                    for col_name in dset.column_names:
                        self.assertLessEqual(len(batch[col_name]), batch_size)
                    del batch

                # Full
                dset_to_polars = dset.to_polars()
                self.assertIsInstance(dset_to_polars, sys.modules["polars"].DataFrame)
                self.assertListEqual(sorted(dset_to_polars.columns), sorted(dset.column_names))
                for col_name in dset.column_names:
                    self.assertEqual(len(dset_to_polars[col_name]), len(dset))

                # With index mapping
                with dset.select([1, 0, 3]) as dset:
                    dset_to_polars = dset.to_polars()
                    self.assertIsInstance(dset_to_polars, sys.modules["polars"].DataFrame)
                    self.assertEqual(len(dset_to_polars), 3)
                    self.assertListEqual(sorted(dset_to_polars.columns), sorted(dset.column_names))

                    for col_name in dset.column_names:
                        self.assertEqual(len(dset_to_polars[col_name]), dset.num_rows)

    def test_to_parquet(self, in_memory):
        with tempfile.TemporaryDirectory() as tmp_dir:
            # File path argument
            with self._create_dummy_dataset(in_memory, tmp_dir, multiple_columns=True) as dset:
                file_path = os.path.join(tmp_dir, "test_path.parquet")
                dset.to_parquet(path_or_buf=file_path)