Speed up default Python row extraction

src/datasets/formatting/formatting.py

@@ -141,7 +141,14 @@ def extract_batch(self, pa_table: pa.Table) -> pa.Table:
 
 class PythonArrowExtractor(BaseArrowExtractor[dict, list, dict]):
     def extract_row(self, pa_table: pa.Table) -> dict:
-        return _unnest(pa_table.to_pydict())
+        # Convert only the single row that is asked for. to_pydict() builds a
+        # one-element list per column and _unnest then throws those lists away.
+        # ArrayXD implements additional shape and null conversion in to_pylist(),
+        # so keep that conversion consistent with the column and batch paths.
+        return {
+            name: col.to_pylist()[0] if isinstance(col.type, _ArrayXDExtensionType) else col[0].as_py()
+            for name, col in zip(pa_table.column_names, pa_table.columns)
+        }
 
     def extract_column(self, pa_table: pa.Table) -> list:
         return pa_table.column(0).to_pylist()
@@ -468,23 +475,26 @@ class PythonFormatter(Formatter[Mapping, list, Mapping]):
     def __init__(self, features=None, lazy=False, token_per_repo_id=None):
         super().__init__(features, token_per_repo_id)
         self.lazy = lazy
+        # PythonArrowExtractor is stateless, so build it once instead of once per
+        # row. format_row() is called for every row of every iteration.
+        self._python_arrow_extractor = self.python_arrow_extractor()
 
     def format_row(self, pa_table: pa.Table) -> Mapping:
         if self.lazy:
             return LazyRow(pa_table, self)
-        row = self.python_arrow_extractor().extract_row(pa_table)
+        row = self._python_arrow_extractor.extract_row(pa_table)
         row = self.python_features_decoder.decode_row(row)
         return row
 
     def format_column(self, pa_table: pa.Table) -> list:
-        column = self.python_arrow_extractor().extract_column(pa_table)
+        column = self._python_arrow_extractor.extract_column(pa_table)
         column = self.python_features_decoder.decode_column(column, pa_table.column_names[0])
         return column
 
     def format_batch(self, pa_table: pa.Table) -> Mapping:
         if self.lazy:
             return LazyBatch(pa_table, self)
-        batch = self.python_arrow_extractor().extract_batch(pa_table)
+        batch = self._python_arrow_extractor.extract_batch(pa_table)
         batch = self.python_features_decoder.decode_batch(batch)
         return batch
 

tests/test_formatting.py

@@ -7,7 +7,7 @@
 import pyarrow as pa
 import pytest
 
-from datasets import Audio, Features, Image, IterableDataset
+from datasets import Array2D, Audio, Dataset, Features, Image, IterableDataset
 from datasets.formatting import NumpyFormatter, PandasFormatter, PythonFormatter, query_table
 from datasets.formatting.formatting import (
     LazyBatch,
@@ -73,6 +73,20 @@ def test_python_extractor(self):
         batch = extractor.extract_batch(pa_table)
         self.assertEqual(batch, {"a": _COL_A, "b": _COL_B, "c": _COL_C, "d": _COL_D})
 
+    def test_python_extractor_array_xd_row_matches_batch(self):
+        dataset = Dataset.from_dict(
+            {"a": [None, [[1, 2, 3]]]},
+            features=Features({"a": Array2D(shape=(None, 3), dtype="int64")}),
+        )
+        pa_table = dataset._data.fast_slice(0, 1)
+        extractor = PythonArrowExtractor()
+
+        row = extractor.extract_row(pa_table)
+        batch = extractor.extract_batch(pa_table)
+
+        self.assertTrue(np.isnan(row["a"]))
+        self.assertTrue(np.isnan(batch["a"][0]))
+
     def test_numpy_extractor(self):
         pa_table = self._create_dummy_table().drop(["c", "d"])
         extractor = NumpyArrowExtractor()