The `add_column()` method does not work if used on dataset sliced with `select()`
State: closed
Labels: description: Something isn't working
name: bug

Hello, I posted this as a question on the forums ([here](https://discuss.huggingface.co/t/add-column-does-not-work-if-used-on-dataset-sliced-with-select/13893)): 

I have a dataset with 2000 entries

> dataset = Dataset.from_dict({'colA': list(range(2000))})

and from which I want to extract the first one thousand rows, create a new dataset with these and also add a new column to it:

> dataset2 = dataset.select(list(range(1000)))
> final_dataset = dataset2.add_column('colB', list(range(1000)))

This gives an error
>ArrowInvalid: Added column's length must match table's length. Expected length 2000 but got length 1000

So it looks like even though it is a dataset with 1000 rows, it "remembers" the shape of the one it was sliced from.

## Actual results
```
ArrowInvalid                              Traceback (most recent call last)
<ipython-input-138-e806860f3ce3> in <module>
----> 1 final_dataset = dataset2.add_column('colB', list(range(1000)))

~/.local/lib/python3.8/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    468         }
    469         # apply actual function
--> 470         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    471         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    472         # re-apply format to the output

~/.local/lib/python3.8/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    404             # Call actual function
    405 
--> 406             out = func(self, *args, **kwargs)
    407 
    408             # Update fingerprint of in-place transforms + update in-place history of transforms

~/.local/lib/python3.8/site-packages/datasets/arrow_dataset.py in add_column(self, name, column, new_fingerprint)
   3343         column_table = InMemoryTable.from_pydict({name: column})
   3344         # Concatenate tables horizontally
-> 3345         table = ConcatenationTable.from_tables([self._data, column_table], axis=1)
   3346         # Update features
   3347         info = self.info.copy()

~/.local/lib/python3.8/site-packages/datasets/table.py in from_tables(cls, tables, axis)
    729             table_blocks = to_blocks(table)
    730             blocks = _extend_blocks(blocks, table_blocks, axis=axis)
--> 731         return cls.from_blocks(blocks)
    732 
    733     @property

~/.local/lib/python3.8/site-packages/datasets/table.py in from_blocks(cls, blocks)
    668     @classmethod
    669     def from_blocks(cls, blocks: TableBlockContainer) -> "ConcatenationTable":
--> 670         blocks = cls._consolidate_blocks(blocks)
    671         if isinstance(blocks, TableBlock):
    672             table = blocks

~/.local/lib/python3.8/site-packages/datasets/table.py in _consolidate_blocks(cls, blocks)
    664             return cls._merge_blocks(blocks, axis=0)
    665         else:
--> 666             return cls._merge_blocks(blocks)
    667 
    668     @classmethod

~/.local/lib/python3.8/site-packages/datasets/table.py in _merge_blocks(cls, blocks, axis)
    650                 merged_blocks += list(block_group)
    651         else:  # both
--> 652             merged_blocks = [cls._merge_blocks(row_block, axis=1) for row_block in blocks]
    653             if all(len(row_block) == 1 for row_block in merged_blocks):
    654                 merged_blocks = cls._merge_blocks(

~/.local/lib/python3.8/site-packages/datasets/table.py in <listcomp>(.0)
    650                 merged_blocks += list(block_group)
    651         else:  # both
--> 652             merged_blocks = [cls._merge_blocks(row_block, axis=1) for row_block in blocks]
    653             if all(len(row_block) == 1 for row_block in merged_blocks):
    654                 merged_blocks = cls._merge_blocks(

~/.local/lib/python3.8/site-packages/datasets/table.py in _merge_blocks(cls, blocks, axis)
    647             for is_in_memory, block_group in groupby(blocks, key=lambda x: isinstance(x, InMemoryTable)):
    648                 if is_in_memory:
--> 649                     block_group = [InMemoryTable(cls._concat_blocks(list(block_group), axis=axis))]
    650                 merged_blocks += list(block_group)
    651         else:  # both

~/.local/lib/python3.8/site-packages/datasets/table.py in _concat_blocks(blocks, axis)
    626                 else:
    627                     for name, col in zip(table.column_names, table.columns):
--> 628                         pa_table = pa_table.append_column(name, col)
    629             return pa_table
    630         else:

~/.local/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.append_column()

~/.local/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.add_column()

~/.local/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

~/.local/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid: Added column's length must match table's length. Expected length 2000 but got length 1000
```

A solution provided by @mariosasko is to use `dataset2.flatten_indices()` after the `select()` and before attempting to add the new column:

> dataset = Dataset.from_dict({'colA': list(range(2000))})
> dataset2 = dataset.select(list(range(1000)))
> dataset2 = dataset2.flatten_indices()
> final_dataset = dataset2.add_column('colB', list(range(1000)))

which works.

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 1.13.2 (note: also checked with version 1.17.0, still the same error)
- Platform: Ubuntu 20.04.3
- Python version: 3.8.10
- PyArrow version: 6.0.0


Comments:
similar #3611