concatenate_datasets(axis=1) with Arrow format drops the first source's columns

[ebarkhordar]

Describe the bug

The Arrow fast path of horizontal concatenation (concatenate_datasets([...], axis=1) consumed through .with_format("arrow"), or any Arrow/tensor consumer) drops the columns of the first source. The plain-Python path is correct; only the Arrow path diverges.

HorizontallyConcatenatedMultiSourcesExamplesIterable._iter_arrow builds each output batch like this:

for j, table in enumerate(pa_tables):
    if j == 0:
        new_pa_table = table
    else:
        for name, col in zip(table.column_names, table.columns):
            new_pa_table = pa_table.append_column(name, col)

The inner loop appends onto pa_table, which is the loop variable that leaked from the fetch loop above (key, pa_table = next(pa_table_iterator)), so it holds the last source's table rather than the accumulator new_pa_table. Each iteration also reassigns new_pa_table from scratch. For sources with columns [a, b] and [c], the batch ends up as the second source's table with c appended to itself (columns [c, c]); columns a and b are lost. When the dataset has resolved features, the resulting schema mismatch surfaces as a CastError.

Steps to reproduce the bug

from datasets import Dataset, concatenate_datasets

ds1 = Dataset.from_dict({"a": [1, 2], "b": [3, 4]}).to_iterable_dataset()
ds2 = Dataset.from_dict({"c": [5, 6]}).to_iterable_dataset()

concatenated = concatenate_datasets([ds1, ds2], axis=1)

# plain-Python path is correct
print(next(iter(concatenated)))            # {'a': 1, 'b': 3, 'c': 5}

# Arrow fast path drops columns a and b
print(next(iter(concatenated.with_format("arrow"))))

The Arrow iteration raises:

datasets.table.CastError: Couldn't cast
c: int64
c: int64
a: null
b: null
...
to
{'a': Value('int64'), 'b': Value('int64'), 'c': Value('int64')}
because column names don't match

Expected behavior

The Arrow fast path should produce the same columns as the plain-Python path: all source columns, in order (['a', 'b', 'c']), with matching values.

Environment info

• datasets 5.0.1.dev0 (main, commit 41adfd0)
• pyarrow 25.0.0
• Python 3.11.15, Linux

Reproduced in a clean container against current main. The fault line was introduced in #8068.

[ebarkhordar]

#self-assign