fix class_encode_column issue

huggingface · Apr 27, 2021 · 88676c9 · 88676c9 · github-actions · Apr 27, 2021
1 parent d93bc76
commit 88676c9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
@@ -775,6 +775,7 @@ def class_encode_column(self, column: str) -> "Dataset":
         class_names = sorted(dset.unique(column))
         dst_feat = ClassLabel(names=class_names)
         dset = dset.map(lambda batch: {column: dst_feat.str2int(batch)}, input_columns=column, batched=True)
+        dset = concatenate_datasets([self.remove_columns([column]), dset], axis=1)
 
         new_features = copy.deepcopy(dset.features)
         new_features[column] = dst_feat

diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
@@ -1975,7 +1975,7 @@ def test_dataset_add_item(item, in_memory, dataset_dict, arrow_path, transform):
     dataset = dataset_to_test.add_item(item)
     assert dataset.data.shape == (5, 3)
     expected_features = dataset_to_test.features
-    assert dataset.data.column_names == list(expected_features.keys())
+    assert sorted(dataset.data.column_names) == sorted(expected_features.keys())
     for feature, expected_dtype in expected_features.items():
         assert dataset.features[feature] == expected_dtype
     assert len(dataset.data.blocks) == 1 if in_memory else 2  # multiple InMemoryTables are consolidated as one