apply formatting after iter_arrow to speed up format -> map, filter for iterable datasets

src/datasets/formatting/formatting.py

@@ -215,11 +215,14 @@ def extract_batch(self, pa_table: pa.Table) -> pd.DataFrame:
 
 
 class PythonFeaturesDecoder:
-    def __init__(self, features: Optional[Features]):
+    def __init__(
+        self, features: Optional[Features], token_per_repo_id: Optional[Dict[str, Union[str, bool, None]]] = None
+    ):
         self.features = features
+        self.token_per_repo_id = token_per_repo_id
 
     def decode_row(self, row: dict) -> dict:
-        return self.features.decode_example(row) if self.features else row
+        return self.features.decode_example(row, token_per_repo_id=self.token_per_repo_id) if self.features else row
 
     def decode_column(self, column: list, column_name: str) -> list:
         return self.features.decode_column(column, column_name) if self.features else column
@@ -393,9 +396,14 @@ class Formatter(Generic[RowFormat, ColumnFormat, BatchFormat]):
     numpy_arrow_extractor = NumpyArrowExtractor
     pandas_arrow_extractor = PandasArrowExtractor
 
-    def __init__(self, features: Optional[Features] = None):
+    def __init__(
+        self,
+        features: Optional[Features] = None,
+        token_per_repo_id: Optional[Dict[str, Union[str, bool, None]]] = None,
+    ):
         self.features = features
-        self.python_features_decoder = PythonFeaturesDecoder(self.features)
+        self.token_per_repo_id = token_per_repo_id
+        self.python_features_decoder = PythonFeaturesDecoder(self.features, self.token_per_repo_id)
         self.pandas_features_decoder = PandasFeaturesDecoder(self.features)
 
     def __call__(self, pa_table: pa.Table, query_type: str) -> Union[RowFormat, ColumnFormat, BatchFormat]:
@@ -433,8 +441,8 @@ def format_batch(self, pa_table: pa.Table) -> pa.Table:
 
 
 class PythonFormatter(Formatter[Mapping, list, Mapping]):
-    def __init__(self, features=None, lazy=False):
-        super().__init__(features)
+    def __init__(self, features=None, lazy=False, token_per_repo_id=None):
+        super().__init__(features, token_per_repo_id)
         self.lazy = lazy
 
     def format_row(self, pa_table: pa.Table) -> Mapping:
@@ -484,8 +492,8 @@ class CustomFormatter(Formatter[dict, ColumnFormat, dict]):
     to return.
     """
 
-    def __init__(self, transform: Callable[[dict], dict], features=None, **kwargs):
-        super().__init__(features=features)
+    def __init__(self, transform: Callable[[dict], dict], features=None, token_per_repo_id=None, **kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
         self.transform = transform
 
     def format_row(self, pa_table: pa.Table) -> dict:

src/datasets/formatting/jax_formatter.py

@@ -36,8 +36,8 @@
 
 
 class JaxFormatter(TensorFormatter[Mapping, "jax.Array", Mapping]):
-    def __init__(self, features=None, device=None, **jnp_array_kwargs):
-        super().__init__(features=features)
+    def __init__(self, features=None, device=None, token_per_repo_id=None, **jnp_array_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
         import jax
         from jaxlib.xla_client import Device
 

src/datasets/formatting/np_formatter.py

@@ -24,8 +24,8 @@
 
 
 class NumpyFormatter(TensorFormatter[Mapping, np.ndarray, Mapping]):
-    def __init__(self, features=None, **np_array_kwargs):
-        super().__init__(features=features)
+    def __init__(self, features=None, token_per_repo_id=None, **np_array_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
         self.np_array_kwargs = np_array_kwargs
 
     def _consolidate(self, column):

src/datasets/formatting/tf_formatter.py

@@ -30,8 +30,8 @@
 
 
 class TFFormatter(TensorFormatter[Mapping, "tf.Tensor", Mapping]):
-    def __init__(self, features=None, **tf_tensor_kwargs):
-        super().__init__(features=features)
+    def __init__(self, features=None, token_per_repo_id=None, **tf_tensor_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
         self.tf_tensor_kwargs = tf_tensor_kwargs
         import tensorflow as tf  # noqa: F401 - import tf at initialization
 

src/datasets/formatting/torch_formatter.py

@@ -30,8 +30,8 @@
 
 
 class TorchFormatter(TensorFormatter[Mapping, "torch.Tensor", Mapping]):
-    def __init__(self, features=None, **torch_tensor_kwargs):
-        super().__init__(features=features)
+    def __init__(self, features=None, token_per_repo_id=None, **torch_tensor_kwargs):
+        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
         self.torch_tensor_kwargs = torch_tensor_kwargs
         import torch  # noqa import torch at initialization