huggingface#3111 Set features correctly when concatenating.

src/datasets/arrow_dataset.py

@@ -3643,6 +3643,15 @@ def concatenate_datasets(
         format = {}
         logger.info("Some of the datasets have disparate format. Resetting the format of the concatenated dataset.")
 
+    # Find column types.
+    if axis == 1:
+        features_d = {}
+        for dset in dsets:
+            features_d.update(dset.features)
+        features = Features(features_d)
+    else:
+        features = dsets[0].features
+
     # Concatenate tables
     tables_to_concat = [dset._data for dset in dsets if len(dset._data) > 0]
     # There might be no table with data left hence return first empty table
@@ -3702,6 +3711,7 @@ def apply_offset_to_indices_table(table, offset):
         fingerprint=fingerprint,
     )
     concatenated_dataset.set_format(**format)
+    concatenated_dataset = concatenated_dataset.cast(features)
     return concatenated_dataset
 
 

tests/test_arrow_dataset.py

@@ -2242,6 +2242,19 @@ def test_concatenate_datasets_duplicate_columns(dataset):
     assert "duplicated" in str(excinfo.value)
 
 
+def test_concatenate_datasets_column_typing(dataset):
+    data = {"label": [0, 1, 1, 0], "col_1": ["a", "b", "c", "d"]}
+    data_2 = {"col2": [0, 1, 1, 0]}
+
+    features = Features({"label": ClassLabel(2, names=["POS", "NEG"]), "col_1": Value("string")})
+    features_2 = Features({"col2": ClassLabel(2, names=["POS", "NEG"])})
+    with Dataset.from_dict(data, features=features, info=DatasetInfo(features=features)) as dset:
+        with Dataset.from_dict(data_2, features=features_2, info=DatasetInfo(features=features_2)) as dset2:
+            concatenated = concatenate_datasets([dset, dset2], axis=1)
+            assert isinstance(concatenated.features["label"], ClassLabel)
+            assert isinstance(concatenated.features["col2"], ClassLabel)
+
+
 def test_interleave_datasets():
     d1 = Dataset.from_dict({"a": [0, 1, 2]})
     d2 = Dataset.from_dict({"a": [10, 11, 12, 13]})