Lazy column (#7614)

* lazy column

* docs

* fix tests

* fix tests

* fix tests

* again

* again

Author: lhoestq

docs/source/access.mdx

@@ -54,7 +54,7 @@ You can combine row and column name indexing to return a specific value at a pos
 'the rock is destined to be the 21st century\'s new " conan " and that he\'s going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'
 ```
 
-But it is important to remember that indexing order matters, especially when working with large audio and image datasets. Indexing by the column name returns all the values in the column first, then loads the value at that position. For large datasets, it may be slower to index by the column name first.
+Indexing order doesn't matter. Indexing by the column name first returns a [`Column`] object that you can index as usual with row indices as usual:
 
 ```py
 >>> import time
@@ -69,7 +69,7 @@ Elapsed time: 0.0031 seconds
 >>> text = dataset["text"][0]
 >>> end_time = time.time()
 >>> print(f"Elapsed time: {end_time - start_time:.4f} seconds")
-Elapsed time: 0.0094 seconds
+Elapsed time: 0.0042 seconds
 ```
 
 ### Slicing

src/datasets/__init__.py

@@ -14,7 +14,7 @@
 
 __version__ = "3.6.0.dev0"
 
-from .arrow_dataset import Dataset
+from .arrow_dataset import Column, Dataset
 from .arrow_reader import ReadInstruction
 from .builder import ArrowBasedBuilder, BuilderConfig, DatasetBuilder, GeneratorBasedBuilder
 from .combine import concatenate_datasets, interleave_datasets
@@ -30,7 +30,7 @@
     get_dataset_infos,
     get_dataset_split_names,
 )
-from .iterable_dataset import IterableDataset
+from .iterable_dataset import IterableColumn, IterableDataset
 from .load import load_dataset, load_dataset_builder, load_from_disk
 from .splits import (
     NamedSplit,

src/datasets/arrow_dataset.py

@@ -627,6 +627,48 @@ class NonExistentDatasetError(Exception):
     pass
 
 
+class Column(Sequence_):
+    """An iterable for a specific column of an [`Dataset`]."""
+
+    def __init__(self, source: Union["Dataset", "Column"], column_name: str):
+        self.source = source
+        self.column_name = column_name
+        if not isinstance(source.features, dict) or column_name not in source.features:
+            raise ValueError(f"Column '{column_name}' doesn't exist.")
+        self.features = source.features[column_name]
+
+    def __iter__(self) -> Iterator[Any]:
+        if isinstance(self.source, Dataset):
+            source = self.source._fast_select_column(self.column_name)
+        for example in source:
+            yield example[self.column_name]
+
+    def __getitem__(self, key: Union[int, str, list[int]]) -> Any:
+        if isinstance(key, str):
+            return Column(self, key)
+        elif isinstance(self.source, Dataset):
+            return self.source._fast_select_column(self.column_name)[key][self.column_name]
+        elif isinstance(key, int):
+            return self.source[key][self.column_name]
+        else:
+            return [item[self.column_name] for item in self.source[key]]
+
+    def __len__(self) -> int:
+        return len(self.source)
+
+    def __repr__(self):
+        return "Column(" + repr(list(self[:5])) + ")"
+
+    def __str__(self):
+        return "Column(" + str(list(self[:5])) + ")"
+
+    def __eq__(self, value):
+        if isinstance(value, Column):
+            return list(self) == list(value)
+        else:
+            return value == list(self)
+
+
 class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
     """A Dataset backed by an Arrow table."""
 
@@ -2354,6 +2396,13 @@ def select_columns(self, column_names: Union[str, list[str]], new_fingerprint: O
         dataset._fingerprint = new_fingerprint
         return dataset
 
+    @transmit_format
+    def _fast_select_column(self, column_name: str) -> "Dataset":
+        dataset = copy.copy(self)
+        dataset._data = dataset._data.select([column_name])
+        dataset._info = DatasetInfo(features=Features({column_name: self._info.features[column_name]}))
+        return dataset
+
     def __len__(self):
         """Number of rows in the dataset.
 
@@ -2776,6 +2825,9 @@ def __getitem__(self, key: str) -> list:  # noqa: F811
 
     def __getitem__(self, key):  # noqa: F811
         """Can be used to index columns (by string names) or rows (by integer index or iterable of indices or bools)."""
+        if isinstance(key, str):
+            if self._format_type is None or self._format_type not in ("arrow", "pandas", "polars"):
+                return Column(self, key)
         return self._getitem(key)
 
     def __getitems__(self, keys: list) -> list:

tests/features/test_array_xd.py

@@ -173,7 +173,7 @@ def get_dict_examples(self, shape_1, shape_2):
         }
 
     def _check_getitem_output_type(self, dataset, shape_1, shape_2, first_matrix):
-        matrix_column = dataset["matrix"]
+        matrix_column = dataset["matrix"][:]
         self.assertIsInstance(matrix_column, list)
         self.assertIsInstance(matrix_column[0], list)
         self.assertIsInstance(matrix_column[0][0], list)
@@ -192,7 +192,7 @@ def _check_getitem_output_type(self, dataset, shape_1, shape_2, first_matrix):
         self.assertTupleEqual(np.array(matrix_field_of_first_two_examples).shape, (2, *shape_2))
 
         with dataset.formatted_as("numpy"):
-            self.assertTupleEqual(dataset["matrix"].shape, (2, *shape_2))
+            self.assertTupleEqual(dataset["matrix"][:].shape, (2, *shape_2))
             self.assertEqual(dataset[0]["matrix"].shape, shape_2)
             self.assertTupleEqual(dataset[:2]["matrix"].shape, (2, *shape_2))
 

tests/features/test_audio.py

@@ -4,7 +4,7 @@
 import pyarrow as pa
 import pytest
 
-from datasets import Dataset, concatenate_datasets, load_dataset
+from datasets import Column, Dataset, concatenate_datasets, load_dataset
 from datasets.features import Audio, Features, Sequence, Value
 
 from ..utils import (
@@ -292,7 +292,7 @@ def test_dataset_with_audio_feature_with_none():
     assert isinstance(batch["audio"], list) and all(item is None for item in batch["audio"])
     column = dset["audio"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(item is None for item in column)
+    assert isinstance(column, Column) and all(item is None for item in column)
 
     # nested tests
 

tests/features/test_features.py

@@ -9,7 +9,7 @@
 import pytest
 
 from datasets import Array2D
-from datasets.arrow_dataset import Dataset
+from datasets.arrow_dataset import Column, Dataset
 from datasets.features import Audio, ClassLabel, Features, Image, LargeList, Sequence, Value
 from datasets.features.features import (
     _align_features,
@@ -492,7 +492,7 @@ def test_dataset_feature_with_none(feature):
     assert isinstance(batch["col"], list) and all(item is None for item in batch["col"])
     column = dset["col"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(item is None for item in column)
+    assert isinstance(column, Column) and all(item is None for item in column)
 
     # nested tests
 

tests/features/test_image.py

@@ -9,7 +9,7 @@
 import pyarrow as pa
 import pytest
 
-from datasets import Dataset, Features, Image, Sequence, Value, concatenate_datasets, load_dataset
+from datasets import Column, Dataset, Features, Image, Sequence, Value, concatenate_datasets, load_dataset
 from datasets.features.image import encode_np_array, image_to_bytes
 
 from ..utils import require_pil
@@ -149,7 +149,7 @@ def test_dataset_with_image_feature(shared_datadir):
     assert batch["image"][0].mode == "RGB"
     column = dset["image"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(isinstance(item, PIL.Image.Image) for item in column)
+    assert isinstance(column, Column) and all(isinstance(item, PIL.Image.Image) for item in column)
     assert os.path.samefile(column[0].filename, image_path)
     assert column[0].format == "JPEG"
     assert column[0].size == (640, 480)
@@ -182,7 +182,7 @@ def test_dataset_with_image_feature_from_pil_image(infer_feature, shared_datadir
     assert batch["image"][0].mode == "RGB"
     column = dset["image"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(isinstance(item, PIL.Image.Image) for item in column)
+    assert isinstance(column, Column) and all(isinstance(item, PIL.Image.Image) for item in column)
     assert os.path.samefile(column[0].filename, image_path)
     assert column[0].format == "JPEG"
     assert column[0].size == (640, 480)
@@ -215,7 +215,7 @@ def test_dataset_with_image_feature_from_np_array():
     assert batch["image"][0].size == (640, 480)
     column = dset["image"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(isinstance(item, PIL.Image.Image) for item in column)
+    assert isinstance(column, Column) and all(isinstance(item, PIL.Image.Image) for item in column)
     np.testing.assert_array_equal(np.array(column[0]), image_array)
     assert column[0].filename == ""
     assert column[0].format in ["PNG", "TIFF"]
@@ -250,7 +250,7 @@ def test_dataset_with_image_feature_tar_jpg(tar_jpg_path):
     assert batch["image"][0].mode == "RGB"
     column = dset["image"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(isinstance(item, PIL.Image.Image) for item in column)
+    assert isinstance(column, Column) and all(isinstance(item, PIL.Image.Image) for item in column)
     assert column[0].filename == ""
     assert column[0].format == "JPEG"
     assert column[0].size == (640, 480)
@@ -271,7 +271,7 @@ def test_dataset_with_image_feature_with_none():
     assert isinstance(batch["image"], list) and all(item is None for item in batch["image"])
     column = dset["image"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(item is None for item in column)
+    assert isinstance(column, Column) and all(item is None for item in column)
 
     # nested tests
 
@@ -527,8 +527,8 @@ def test_formatted_dataset_with_image_feature(shared_datadir):
         assert batch["image"].shape == (1, 480, 640, 3)
         column = dset["image"]
         assert len(column) == 2
-        assert isinstance(column, np.ndarray)
-        assert column.shape == (2, 480, 640, 3)
+        assert isinstance(column[:], np.ndarray)
+        assert column[:].shape == (2, 480, 640, 3)
 
     with dset.formatted_as("pandas"):
         item = dset[0]

tests/features/test_video.py

@@ -1,6 +1,6 @@
 import pytest
 
-from datasets import Dataset, Features, Video
+from datasets import Column, Dataset, Features, Video
 
 from ..utils import require_torchvision
 
@@ -53,7 +53,7 @@ def test_dataset_with_video_feature(shared_datadir):
     assert isinstance(next(batch["video"][0])["data"], torch.Tensor)
     column = dset["video"]
     assert len(column) == 1
-    assert isinstance(column, list) and all(isinstance(item, VideoReader) for item in column)
+    assert isinstance(column, Column) and all(isinstance(item, VideoReader) for item in column)
     assert next(column[0])["data"].shape == (3, 50, 66)
     assert isinstance(next(column[0])["data"], torch.Tensor)