Reject an invalid axis in concatenate_datasets

[LeSingh1]

concatenate_datasets() documents axis as {0, 1}, but does not check it. Any other value silently returns only the first dataset:

from datasets import Dataset, concatenate_datasets

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

concatenate_datasets([a, b], axis=2).to_dict()
# {'a': [1, 2, 3]}          <- b is gone, no error

concatenate_datasets([a, b], axis="1").to_dict()
# {'a': [1, 2, 3]}          <- a string axis does the same

In _concatenate_map_style_datasets every branch is written as if axis == 0: ... else: ..., so an invalid value takes the axis=1 path through the checks (row counts must match, no duplicate column names — both pass here) and is then forwarded to concat_tables(..., axis=axis), which builds nothing for it.

axis="1" is the realistic version of this: it comes straight out of a config file or a CLI argument, and instead of an error you get a dataset that is missing half its columns.

The iterable implementation concatenates horizontally for axis=2, so the map-style and streaming paths disagree on top of that.

This PR raises a ValueError for anything but 0 or 1 at the top of _concatenate_map_style_datasets.

Added tests/test_arrow_dataset.py::test_concatenate_datasets_with_invalid_axis, parametrized over 2, -1 and "1"; all three fail on main.

[shashvat-singham]

The check works for map-style datasets — 2, -1, "1" and None all raise cleanly now:

axis=2      -> ValueError: axis must be 0 (rows) or 1 (columns), but got 2.
axis=-1     -> ValueError: axis must be 0 (rows) or 1 (columns), but got -1.
axis='1'    -> ValueError: axis must be 0 (rows) or 1 (columns), but got 1.
axis=None   -> ValueError: axis must be 0 (rows) or 1 (columns), but got None.

Nice touch that it catches "1" and None, not just out-of-range ints.

The iterable path is untouched, and it accepts everything. The check went into _concatenate_map_style_datasets, but concatenate_datasets dispatches to _concatenate_iterable_datasets for IterableDataset. Same four values, same branch:

axis=2      -> NO ERROR, 3 rows, first={'col_1': 0, 'col_2': 3}
axis=-1     -> NO ERROR, 3 rows, first={'col_1': 0, 'col_2': 3}
axis='1'    -> NO ERROR, 3 rows, first={'col_1': 0, 'col_2': 3}
axis=None   -> NO ERROR, 3 rows, first={'col_1': 0, 'col_2': 3}

Every invalid axis is silently treated as axis=1 — the dispatch is effectively if axis == 0: vertical else: horizontal, so anything that is not 0 falls into the horizontal branch, including None and a string.

So as it stands this PR creates a divergence: after it lands, concatenate_datasets([...], axis=2) raises on Dataset and quietly concatenates columns on IterableDataset. That is the same eager/streaming mismatch your #8451 exists to remove, in a new place.

Putting the check in concatenate_datasets in combine.py instead of in _concatenate_map_style_datasets would cover both dispatch targets from one place, and it is arguably where an argument-validation check belongs anyway — it validates the public signature rather than one of the two private implementations.

If you would rather keep it where it is, the mirror check in _concatenate_iterable_datasets would at least close the gap, and the existing @pytest.mark.parametrize("axis", [2, -1, "1"]) test extends to iterable datasets almost for free.

Verified on Windows 11 / Python 3.11.9, pr-8447.