Implement Dataset add_column

docs/source/package_reference/main_classes.rst

@@ -14,7 +14,9 @@ Main classes
 The base class :class:`datasets.Dataset` implements a Dataset backed by an Apache Arrow table.
 
 .. autoclass:: datasets.Dataset
-    :members: from_file, from_buffer, from_pandas, from_dict,
+    :members:
+        add_column,
+        from_file, from_buffer, from_pandas, from_dict,
         data, cache_files, num_columns, num_rows, column_names, shape,
         unique,
         flatten_, cast_, remove_columns_, rename_column_,

src/datasets/arrow_dataset.py

@@ -56,7 +56,7 @@
 from .info import DATASET_INFO_FILENAME, DatasetInfo
 from .search import IndexableMixin
 from .splits import NamedSplit
-from .table import InMemoryTable, MemoryMappedTable, Table, concat_tables, list_table_cache_files
+from .table import ConcatenationTable, InMemoryTable, MemoryMappedTable, Table, concat_tables, list_table_cache_files
 from .utils import map_nested
 from .utils.deprecation_utils import deprecated
 from .utils.logging import WARNING, get_logger, get_verbosity, set_verbosity_warning
@@ -2654,6 +2654,26 @@ def to_pandas(
                 for offset in range(0, len(self), batch_size)
             )
 
+    def add_column(self, name: str, column: Union[list, np.array]):
+        """Add column to Dataset.
+
+        .. versionadded:: 1.6
+
+        Args:
+            name (str): Column name.
+            column (list or np.array): Column data to be added.
+
+        Returns:
+            :class:`Dataset`
+        """
+        column_table = InMemoryTable.from_pydict({name: column})
+        # Concatenate tables horizontally
+        table = ConcatenationTable.from_tables([self._data, column_table], axis=1)
+        # Update features
+        info = self.info
+        info.features = Features.from_arrow_schema(table.schema)
+        return Dataset(table, info=info)
+
     def add_faiss_index(
         self,
         column: str,

tests/test_arrow_dataset.py

@@ -1948,6 +1948,27 @@ def test_concatenate_datasets_duplicate_columns(dataset):
     assert "duplicated" in str(excinfo.value)
 
 
+@pytest.mark.parametrize("in_memory", [False, True])
+def test_dataset_add_column(in_memory, dataset_dict, arrow_path):
+    column_name = "col_4"
+    column = ["a", "b", "c", "d"]
+    dataset = (
+        Dataset(InMemoryTable.from_pydict(dataset_dict))
+        if in_memory
+        else Dataset(MemoryMappedTable.from_file(arrow_path))
+    )
+    dataset = dataset.add_column(column_name, column)
+    assert dataset.data.shape == (4, 4)
+    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64", "col_4": "string"}
+    assert dataset.data.column_names == list(expected_features.keys())
+    for feature, expected_dtype in expected_features.items():
+        assert dataset.features[feature].dtype == expected_dtype
+    assert len(dataset.data.blocks) == 1 if in_memory else 2  # multiple InMemoryTables are consolidated as one
+    column_name = "col_5"
+    dataset = dataset.add_column(column_name, column)
+    assert len(dataset.data.blocks) == 1 if in_memory else 2  # multiple InMemoryTables are consolidated as one
+
+
 @pytest.mark.parametrize("keep_in_memory", [False, True])
 @pytest.mark.parametrize(
     "features",