Repository: huggingface/datasets
Pull request: handle data alteration when trying type
Fix #649 

The bug came from the type inference that didn't handle a weird case in Pyarrow.
Indeed this code runs without error but alters the data in arrow:
```python
import pyarrow as pa

type = pa.struct({"a": pa.struct({"b": pa.string()})})
array_with_altered_data = pa.array([{"a": {"b": "foo", "c": "bar"}}] * 10, type=type)
print(array_with_altered_data[0].as_py())
# {'a': {'b': 'foo'}} -> the sub-field "c" is missing
```
(I don't know if this is intended in pyarrow tbh)

We didn't take this case into account during type inference. By default it was keeping old features and maybe alter data.
To fix that I added a line that checks that the first element of the array is not altered.

Problem statement:
Inconsistent behavior in map
I'm observing inconsistent behavior when applying .map(). This happens specifically when I'm incrementally adding onto a feature that is a nested dictionary. Here's a simple example that reproduces the problem.

```python
import datasets

# Dataset with a single feature called 'field' consisting of two examples
dataset = datasets.Dataset.from_dict({'field': ['a', 'b']})
print(dataset[0])
# outputs
{'field': 'a'}

# Map this dataset to create another feature called 'otherfield', which is a dictionary containing a key called 'capital'
dataset = dataset.map(lambda example: {'otherfield': {'capital': example['field'].capitalize()}})
print(dataset[0])
# output is okay
{'field': 'a', 'otherfield': {'capital': 'A'}}

# Now I want to map again to modify 'otherfield', by adding another key called 'append_x' to the dictionary under 'otherfield'
print(dataset.map(lambda example: {'otherfield': {'append_x': example['field'] + 'x'}})[0])
# printing out the first example after applying the map shows that the new key 'append_x' doesn't get added
# it also messes up the value stored at 'capital'
{'field': 'a', 'otherfield': {'capital': None}}

# Instead, I try to do the same thing by using a different mapped fn
print(dataset.map(lambda example:  {'otherfield': {'append_x': example['field'] + 'x', 'capital': example['otherfield']['capital']}})[0])
# this preserves the value under capital, but still no 'append_x'
{'field': 'a', 'otherfield': {'capital': 'A'}}

# Instead, I try to pass 'otherfield' to remove_columns
print(dataset.map(lambda example:  {'otherfield': {'append_x': example['field'] + 'x', 'capital': example['otherfield']['capital']}}, remove_columns=['otherfield'])[0])
# this still doesn't fix the problem
{'field': 'a', 'otherfield': {'capital': 'A'}}

# Alternately, here's what happens if I just directly map both 'capital' and 'append_x' on a fresh dataset.

# Recreate the dataset
dataset = datasets.Dataset.from_dict({'field': ['a', 'b']})
# Now map the entire 'otherfield' dict directly, instead of incrementally as before
print(dataset.map(lambda example:  {'otherfield': {'append_x': example['field'] + 'x', 'capital': example['field'].capitalize()}})[0])
# This looks good!
{'field': 'a', 'otherfield': {'append_x': 'ax', 'capital': 'A'}}
```

This might be a new issue, because I didn't see this behavior in the `nlp` library. 

Any help is appreciated!

Hints from issue discussion:
Thanks for reporting !

This issue must have appeared when we refactored type inference in `nlp`
By default the library tries to keep the same feature types when applying `map` but apparently it has troubles with nested structures. I'll try to fix that next week

Commit message: handle data alteration when trying type

Patch to review:
diff --git a/src/datasets/arrow_writer.py b/src/datasets/arrow_writer.py
index 17ac8ada6b6..a9fdfab2c81 100644
--- a/src/datasets/arrow_writer.py
+++ b/src/datasets/arrow_writer.py
@@ -94,9 +94,12 @@ def __arrow_array__(self, type=None):
             type = self.type
         try:
             if isinstance(type, _ArrayXDExtensionType):
-                return pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
+                out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
             else:
-                return pa.array(self.data, type=type)
+                out = pa.array(self.data, type=type)
+            if trying_type and out[0].as_py() != self.data[0]:
+                raise TypeError("Specified try_type alters data")
+            return out
         except (TypeError, pa.lib.ArrowInvalid) as e:  # handle type errors and overflows
             if trying_type:
                 try:
diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
index 7057bc8c08c..cc29055c16b 100644
--- a/tests/test_arrow_dataset.py
+++ b/tests/test_arrow_dataset.py
@@ -583,6 +583,14 @@ def map_batched_with_indices(example, idx):
                 Features({"filename": Value("string"), "filename_new": Value("string")}),
             )
 
+    def test_map_nested(self, in_memory):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            dset = Dataset.from_dict({"field": ["a", "b"]})
+            dset = self._to(in_memory, tmp_dir, dset)
+            dset = dset.map(lambda example: {"otherfield": {"capital": example["field"].capitalize()}})
+            dset = dset.map(lambda example: {"otherfield": {"append_x": example["field"] + "x"}})
+            self.assertEqual(dset[0], {"field": "a", "otherfield": {"append_x": "ax"}})
+
     @require_torch
     def test_map_torch(self, in_memory):
         import torch


Code review comments:
diff_hunk: @@ -94,9 +94,12 @@ def __arrow_array__(self, type=None):
             type = self.type
         try:
             if isinstance(type, _ArrayXDExtensionType):
-                return pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
+                out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
             else:
-                return pa.array(self.data, type=type)
+                out = pa.array(self.data, type=type)
+            if trying_type and out[0].as_py() != self.data[0]:
+                raise TypeError("Specified try_type alters data")
path: src/datasets/arrow_writer.py
text: @user1:
```suggestion
                raise TypeError("Specified try_type alters data. Please check that the type/feature that you provided match the type/features of the data.")
```