Fix Dataset.map writer initialization when early examples return None

[veeceey]

Summary

Fixes #7990

This PR fixes a bug in Dataset.map() where the writer initialization was incorrectly tied to the index being 0, causing crashes when the map function returns None for the first few examples and later returns a dict.

Changes

• Non-batched mode (line 3676): Changed from if i == 0: to if writer is None:
• Batched mode (line 3701): Changed from if i and i[0] == 0: to if writer is None:

Why This Fix Works

The original code assumed that update_data would always be determined by the time the first example (i=0) was processed. However, update_data is set lazily after processing each example - it becomes True when the function first returns a non-None value.

If a function returns None for early examples and a dict for later ones:

1. At i=0, the function returns None, so update_data remains None
2. Writer is NOT initialized (because we're not updating data)
3. At i=2, the function returns a dict, so update_data becomes True
4. Old code: Tries to use writer (still None) because i != 0 → crash
5. New code: Checks if writer is None and initializes it → works correctly

Test Plan

The fix can be verified with this minimal test case from the issue:

from datasets import Dataset

ds = Dataset.from_dict({"x": [1, 2, 3]})

def fn(example, idx):
    if idx < 2:
        return None
    return {"x": [example["x"] * 10]}

# Should work without errors
result = list(ds.map(fn, with_indices=True))
print(result)  # [{'x': 1}, {'x': 2}, {'x': [30]}]

Before this fix: Crashes with AttributeError: 'NoneType' object has no attribute 'write'
After this fix: Works correctly

Related

This fix ensures the writer is initialized the first time a non-None value is returned, regardless of which example index that occurs at. This makes the code more robust to different map function behaviors.

[veeceey]

Friendly ping - any chance someone could take a look at this when they get a chance? Happy to make any changes if needed.