Repository: huggingface/datasets
Pull request: Preserve encoding/decoding with features in `Iterable.map` call
As described in https://github.com/huggingface/datasets/issues/3505#issuecomment-1004755657, this PR uses a generator expression to encode/decode examples with `features` (which are set to None in `map`) before applying a map transform.

Fix #3505 

Problem statement:
cast_column function not working with map function in streaming mode for Audio features
## Describe the bug
I am trying to use Audio class for loading audio features using custom dataset. I am able to cast 'audio' feature into 'Audio' format with cast_column function. On using map function, I am not getting 'Audio' casted feature but getting path of audio file only.
I am getting features of 'audio' of string type with load_dataset call. After using cast_column 'audio' feature is converted into 'Audio' type. But in map function I am not able to get Audio type for audio feature & getting string type data containing path of file only. So I am not able to use processor in encode function.

## Steps to reproduce the bug
```python
# Sample code to reproduce the bug
from datasets import load_dataset, Audio
from transformers import Wav2Vec2Processor

def encode(batch, processor):
    print("Audio: ",batch['audio'])
    batch["input_values"] = processor(batch["audio"]['array'], sampling_rate=16000).input_values
    return batch

def print_ds(ds):
    iterator = iter(ds)
    for d in iterator:
        print("Data: ",d)
        break

processor = Wav2Vec2Processor.from_pretrained(pretrained_model_path)

dataset = load_dataset("custom_dataset.py","train",data_files={'train':'train_path.txt'},
                                data_dir="data", streaming=True, split="train")

print("Features: ",dataset.features)

print_ds(dataset)

dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))

print("Features: ",dataset.features)

print_ds(dataset)

dataset = dataset.map(lambda x: encode(x,processor))

print("Features: ",dataset.features)

print_ds(dataset)



```

## Expected results

map function not printing Audio type features be used with processor function and getting error in processor call due to this.

## Actual results

# after load_dataset call
Features:  {'sentence': Value(dtype='string', id=None), 'audio': Value(dtype='string', id=None)}
Data:  {'sentence': 'और अपने पेट को माँ की स्वादिष्ट गरमगरम जलेबियाँ हड़पते\n', 'audio': 'data/0116_003.wav'}

# after cast_column call
Features:  {'sentence': Value(dtype='string', id=None), 'audio': Audio(sampling_rate=16000, mono=True, _storage_dtype='string', id=None)}
Data:  {'sentence': 'और अपने पेट को माँ की स्वादिष्ट गरमगरम जलेबियाँ हड़पते\n', 'audio': {'path': 'data/0116_003.wav', 'array': array([ 1.2662281e-06,  1.0264218e-06, -1.3615092e-06, ...,
        1.3017889e-02,  1.0085563e-02,  4.8155054e-03], dtype=float32), 'sampling_rate': 16000}}

# after map call
Features:  None
Audio:  data/0116_003.wav

Traceback (most recent call last):
  File "demo2.py", line 36, in <module>
    print_ds(dataset)
  File "demo2.py", line 11, in print_ds
    for d in iterator:
  File "/opt/conda/lib/python3.7/site-packages/datasets/iterable_dataset.py", line 341, in __iter__
    for key, example in self._iter():
  File "/opt/conda/lib/python3.7/site-packages/datasets/iterable_dataset.py", line 338, in _iter
    yield from ex_iterable
  File "/opt/conda/lib/python3.7/site-packages/datasets/iterable_dataset.py", line 192, in __iter__
    yield key, self.function(example)
  File "demo2.py", line 32, in <lambda>
    dataset = dataset.map(lambda x: batch_encode(x,processor))
  File "demo2.py", line 6, in batch_encode
    batch["input_values"] = processor(batch["audio"]['array'], sampling_rate=16000).input_values
TypeError: string indices must be integers
## Environment info

- `datasets` version: 1.17.0
- Platform: Linux-4.14.243 with-debian-bullseye-sid
- Python version: 3.7.9
- PyArrow version: 6.0.1

Hints from issue discussion:
Hi! This is probably due to the fact that `IterableDataset.map` sets `features` to `None` before mapping examples. We can fix the issue by passing the old `features` dict to the map generator and performing encoding/decoding there (before calling the map transform function).

Commit message: Add test

Patch to review:
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
index 25144b515eb..a8375bbc84d 100644
--- a/src/datasets/iterable_dataset.py
+++ b/src/datasets/iterable_dataset.py
@@ -1,7 +1,7 @@
 import copy
 from dataclasses import dataclass
 from itertools import cycle, islice, repeat
-from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Union
+from typing import Any, Callable, Dict, Generator, Iterable, Iterator, List, Optional, Union
 
 import numpy as np
 import pyarrow as pa
@@ -38,6 +38,24 @@ def _batch_to_examples(batch: Dict[str, list]) -> List[Dict[str, Any]]:
         yield {col: array[i] for col, array in batch.items()}
 
 
+def _features_formatter(features: Optional[Features]):
+    if features is None:
+
+        def _formatter(example):
+            return example
+
+    else:
+
+        def _formatter(example):
+            # we encode the example for ClassLabel feature types for example
+            encoded_example = features.encode_example(example)
+            # Decode example for Audio feature, e.g.
+            decoded_example = features.decode_example(encoded_example)
+            return decoded_example
+
+    return _formatter
+
+
 class _BaseExamplesIterable:
     """Base class for the examples iterable used by an IterableDataset"""
 
@@ -164,7 +182,11 @@ def shuffle_data_sources(self, seed: Optional[int]) -> "RandomlyCyclingMultiSour
 
 class MappedExamplesIterable(_BaseExamplesIterable):
     def __init__(
-        self, ex_iterable: _BaseExamplesIterable, function: Callable, batched: bool = False, batch_size: int = 1000
+        self,
+        ex_iterable: Union[_BaseExamplesIterable, Generator[Dict[str, Any], None, None]],
+        function: Callable,
+        batched: bool = False,
+        batch_size: int = 1000,
     ):
         self.ex_iterable = ex_iterable
         self.function = function
@@ -338,15 +360,9 @@ def _iter(self):
         yield from ex_iterable
 
     def __iter__(self):
+        formatter = _features_formatter(self.features)
         for key, example in self._iter():
-            if self.features:
-                # we encode the example for ClassLabel feature types for example
-                encoded_example = self.features.encode_example(example)
-                # Decode example for Audio feature, e.g.
-                decoded_example = self.features.decode_example(encoded_example)
-                yield decoded_example
-            else:
-                yield example
+            yield formatter(example)
 
     def with_format(
         self,
@@ -396,8 +412,12 @@ def map(self, function: Callable, batched: bool = False, batch_size: int = 1000)
         """
         info = copy.deepcopy(self._info)
         info.features = None
+        formatter = _features_formatter(self.features)
         ex_iterable = MappedExamplesIterable(
-            self._ex_iterable, function=function, batched=batched, batch_size=batch_size
+            ((key, formatter(example)) for key, example in self._ex_iterable) if self.features else self._ex_iterable,
+            function=function,
+            batched=batched,
+            batch_size=batch_size,
         )
         return iterable_dataset(
             ex_iterable=ex_iterable,
diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py
index ca2da0dd9b6..699a1e3c4a7 100644
--- a/tests/test_iterable_dataset.py
+++ b/tests/test_iterable_dataset.py
@@ -290,6 +290,25 @@ def test_iterable_dataset_map_batched(dataset: IterableDataset, generate_example
     assert next(iter(dataset)) == _func_unbatched(next(iter(generate_examples_fn()))[1])
 
 
+def test_iterable_dataset_map_complex_features(dataset: IterableDataset, generate_examples_fn):
+    # https://github.com/huggingface/datasets/issues/3505
+    ex_iterable = ExamplesIterable(generate_examples_fn, {"label": "positive"})
+    features = Features(
+        {
+            "id": Value("int64"),
+            "label": Value("string"),
+        }
+    )
+    dataset = IterableDataset(ex_iterable, info=DatasetInfo(features=features))
+    dataset = dataset.cast_column("label", ClassLabel(names=["negative", "positive"]))
+    dataset = dataset.map(lambda x: {"id+1": x["id"] + 1, **x})
+    assert isinstance(dataset._ex_iterable, MappedExamplesIterable)
+    features["label"] = ClassLabel(names=["negative", "positive"])
+    assert [{k: v for k, v in ex.items() if k != "id+1"} for ex in dataset] == [
+        features.encode_example(ex) for _, ex in ex_iterable
+    ]
+
+
 @pytest.mark.parametrize("seed", [42, 1337, 101010, 123456])
 @pytest.mark.parametrize("epoch", [None, 0, 1])
 def test_iterable_dataset_shuffle(dataset: IterableDataset, generate_examples_fn, seed, epoch):


Code review comments:
diff_hunk: @@ -164,7 +182,11 @@ def shuffle_data_sources(self, seed: Optional[int]) -> "RandomlyCyclingMultiSour
 
 class MappedExamplesIterable(_BaseExamplesIterable):
     def __init__(
-        self, ex_iterable: _BaseExamplesIterable, function: Callable, batched: bool = False, batch_size: int = 1000
+        self,
+        ex_iterable: Union[_BaseExamplesIterable, Generator[Dict[str, Any], None, None]],
path: src/datasets/iterable_dataset.py
text: @user1:
Generators are iterators, and so have the same issue as in the previous comment.

diff_hunk: @@ -396,8 +412,12 @@ def map(self, function: Callable, batched: bool = False, batch_size: int = 1000)
         """
         info = copy.deepcopy(self._info)
         info.features = None
+        formatter = _features_formatter(self.features)
         ex_iterable = MappedExamplesIterable(
-            self._ex_iterable, function=function, batched=batched, batch_size=batch_size
+            ((key, formatter(example)) for key, example in self._ex_iterable) if self.features else self._ex_iterable,
path: src/datasets/iterable_dataset.py
text: @user1:
If I'm not mistaken here you are passing an iterator (more specifcically a generator) instead of an iterable.
It means that if you iterate over the dataset once, then the second time your iterate over it, it will be empty.
Maybe we can have a TypedExamplesIterable class that does the feature formatting ?

@author:
Oh, I see. Yes, a new class that does the formatting makes sense.