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

Code before:
        # Batched
        def func_return_single_row_pd_dataframe_batched(x):
            batch_size = len(x[next(iter(x))])
            return pd.DataFrame({"id": [0] * batch_size, "text": ["a"] * batch_size})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                with dset.map(func_return_single_row_pd_dataframe_batched, batched=True) as dset_test:
                    self.assertEqual(len(dset_test), 30)
                    self.assertDictEqual(
                        dset_test.features,
                        Features({"id": Value("int64"), "text": Value("string")}),
                    )
                    self.assertEqual(dset_test[0]["id"], 0)
                    self.assertEqual(dset_test[0]["text"], "a")

        # Error when returning a table with more than one row in the non-batched mode
        def func_return_multi_row_pd_dataframe(x):
            return pd.DataFrame({"id": [0, 1], "text": ["a", "b"]})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                self.assertRaises(ValueError, dset.map, func_return_multi_row_pd_dataframe)

    def test_map_return_pl_dataframe(self, in_memory):
        import polars as pl

        def func_return_single_row_pl_dataframe(x):
            return pl.DataFrame({"id": [0], "text": ["a"]})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                with dset.map(func_return_single_row_pl_dataframe) as dset_test:
                    self.assertEqual(len(dset_test), 30)
                    self.assertDictEqual(
                        dset_test.features,
                        Features({"id": Value("int64"), "text": Value("large_string")}),
                    )
                    self.assertEqual(dset_test[0]["id"], 0)
                    self.assertEqual(dset_test[0]["text"], "a")

        # Batched
        def func_return_single_row_pl_dataframe_batched(x):
            batch_size = len(x[next(iter(x))])
            return pl.DataFrame({"id": [0] * batch_size, "text": ["a"] * batch_size})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                with dset.map(func_return_single_row_pl_dataframe_batched, batched=True) as dset_test:
                    self.assertEqual(len(dset_test), 30)
                    self.assertDictEqual(

Review comment (suggestion):
```suggestion
    @require_polars
    def test_map_return_pl_dataframe(self, in_memory):
```

Code after:
        # Batched
        def func_return_single_row_pd_dataframe_batched(x):
            batch_size = len(x[next(iter(x))])
            return pd.DataFrame({"id": [0] * batch_size, "text": ["a"] * batch_size})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                with dset.map(func_return_single_row_pd_dataframe_batched, batched=True) as dset_test:
                    self.assertEqual(len(dset_test), 30)
                    self.assertDictEqual(
                        dset_test.features,
                        Features({"id": Value("int64"), "text": Value("string")}),
                    )
                    self.assertEqual(dset_test[0]["id"], 0)
                    self.assertEqual(dset_test[0]["text"], "a")

        # Error when returning a table with more than one row in the non-batched mode
        def func_return_multi_row_pd_dataframe(x):
            return pd.DataFrame({"id": [0, 1], "text": ["a", "b"]})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                self.assertRaises(ValueError, dset.map, func_return_multi_row_pd_dataframe)

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

        def func_return_single_row_pl_dataframe(x):
            return pl.DataFrame({"id": [0], "text": ["a"]})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                with dset.map(func_return_single_row_pl_dataframe) as dset_test:
                    self.assertEqual(len(dset_test), 30)
                    self.assertDictEqual(
                        dset_test.features,
                        Features({"id": Value("int64"), "text": Value("large_string")}),
                    )
                    self.assertEqual(dset_test[0]["id"], 0)
                    self.assertEqual(dset_test[0]["text"], "a")

        # Batched
        def func_return_single_row_pl_dataframe_batched(x):
            batch_size = len(x[next(iter(x))])
            return pl.DataFrame({"id": [0] * batch_size, "text": ["a"] * batch_size})

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                with dset.map(func_return_single_row_pl_dataframe_batched, batched=True) as dset_test:
                    self.assertEqual(len(dset_test), 30)