Streaming dataset looses `.feature` method after `.add_column`

[sanchit-gandhi]

Describe the bug

After appending a new column to a streaming dataset using .add_column, we can no longer access the list of dataset features using the .feature method.

Steps to reproduce the bug

from datasets import load_dataset

original_dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)
print(original_dataset.features.keys())

# now add a new column to our streaming dataset
modified_dataset = original_dataset.add_column("new_column", ["some random text" for _ in range(50)])
print(modified_dataset.features.keys())

Print Output:

dict_keys(['file', 'audio', 'text', 'speaker_id', 'chapter_id', 'id'])
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 8
      6 # now add a new column to our streaming dataset
      7 modified_dataset = original_dataset.add_column("new_column", ["some random text" for _ in range(50)])
----> 8 print(modified_dataset.features.keys())

AttributeError: 'NoneType' object has no attribute 'keys'

We see that we get the features for the original dataset, but not the modified one with the added column.

Expected behavior

Features should be persevered after adding a new column, i.e. calling:

print(modified_dataset.features.keys())

Should return:

dict_keys(['file', 'audio', 'text', 'speaker_id', 'chapter_id', 'id', 'new_column'])

Environment info

• datasets version: 2.10.2.dev0
• Platform: Linux-4.19.0-23-cloud-amd64-x86_64-with-glibc2.28
• Python version: 3.9.16
• Huggingface_hub version: 0.13.3
• PyArrow version: 10.0.1
• Pandas version: 1.5.2

[sanchit-gandhi]

I believe the issue resides in this line:

datasets/src/datasets/iterable_dataset.py

Line 1415 in 7c3a9b0

return self.map(add_column_fn, with_indices=True)

If we pass the new features of the dataset to the .map method we can return the features after adding a column, e.g.:

from datasets import load_dataset, Value

original_dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)
print(original_dataset.features.keys())

# now add a new column to our streaming dataset using our hack
name = "new_column"
column = ["some random text" for _ in range(50)]

new_features = original_dataset.features.copy()
new_features[name] = Value(dtype="string", id=None)  # I know the correct column type for this feature

def add_column_fn(example, idx):
    if name in example:
        raise ValueError(f"Error when adding {name}: column {name} is already in the dataset.")
    return {name: column[idx]}

modified_dataset = original_dataset.map(add_column_fn, with_indices=True, features=new_features)

print(modified_dataset.features.keys())

Print Output:

dict_keys(['file', 'audio', 'text', 'speaker_id', 'chapter_id', 'id'])
dict_keys(['file', 'audio', 'text', 'speaker_id', 'chapter_id', 'id', 'new_column'])

[voidful]

It seems that map will also cause this issue

Steps to reproduce the bug

from datasets import load_dataset

original_dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)
print(original_dataset.features.keys())

def test(data):
    return data

modified_dataset = original_dataset.map(test)
print(modified_dataset.features.keys())

Output

dict_keys(['file', 'audio', 'text', 'speaker_id', 'chapter_id', 'id'])
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[5], line 10
      7     return data
      9 modified_dataset = original_dataset.map(test)
---> 10 print(modified_dataset.features.keys())

AttributeError: 'NoneType' object has no attribute 'keys'