diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index 59451a640e6..a4ccb36f7a8 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -48,6 +48,7 @@ BinaryIO, Callable, Optional, + Sequence, Union, overload, ) @@ -1290,7 +1291,7 @@ def from_list( @staticmethod def from_csv( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1428,9 +1429,10 @@ def from_generator( **kwargs, ).read() + @staticmethod @staticmethod def from_json( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1487,9 +1489,32 @@ def from_json( **kwargs, ).read() + @staticmethod + def from_jsonl( + path_or_paths: Union[PathLike, Sequence[PathLike]], + split: Optional[NamedSplit] = None, + features: Optional[Features] = None, + cache_dir: str = None, + keep_in_memory: bool = False, + field: Optional[str] = None, + num_proc: Optional[int] = None, + **kwargs, + ) -> "Dataset": + """Alias of `Dataset.from_json`.""" + return Dataset.from_json( + path_or_paths, + split=split, + features=features, + cache_dir=cache_dir, + keep_in_memory=keep_in_memory, + field=field, + num_proc=num_proc, + **kwargs, + ) + @staticmethod def from_parquet( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1586,7 +1611,7 @@ def from_parquet( @staticmethod def from_text( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py index 4abea0a381a..4bb07f5bf40 100644 --- a/src/datasets/dataset_dict.py +++ b/src/datasets/dataset_dict.py @@ -1444,7 +1444,7 @@ def load_from_disk( @staticmethod def from_csv( - path_or_paths: dict[str, PathLike], + path_or_paths: dict[str, Union[PathLike, Sequence[PathLike]]], features: Optional[Features] = None, cache_dir: str = None, keep_in_memory: bool = False, @@ -1485,9 +1485,10 @@ def from_csv( **kwargs, ).read() + @staticmethod @staticmethod def from_json( - path_or_paths: dict[str, PathLike], + path_or_paths: dict[str, Union[PathLike, Sequence[PathLike]]], features: Optional[Features] = None, cache_dir: str = None, keep_in_memory: bool = False, @@ -1528,9 +1529,26 @@ def from_json( **kwargs, ).read() + @staticmethod + def from_jsonl( + path_or_paths: dict[str, Union[PathLike, Sequence[PathLike]]], + features: Optional[Features] = None, + cache_dir: str = None, + keep_in_memory: bool = False, + **kwargs, + ) -> "DatasetDict": + """Alias of `DatasetDict.from_json`.""" + return DatasetDict.from_json( + path_or_paths, + features=features, + cache_dir=cache_dir, + keep_in_memory=keep_in_memory, + **kwargs, + ) + @staticmethod def from_parquet( - path_or_paths: dict[str, PathLike], + path_or_paths: dict[str, Union[PathLike, Sequence[PathLike]]], features: Optional[Features] = None, cache_dir: str = None, keep_in_memory: bool = False, @@ -1579,7 +1597,7 @@ def from_parquet( @staticmethod def from_text( - path_or_paths: dict[str, PathLike], + path_or_paths: dict[str, Union[PathLike, Sequence[PathLike]]], features: Optional[Features] = None, cache_dir: str = None, keep_in_memory: bool = False, diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 17b9a2020fc..785ed364c6c 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -15,7 +15,7 @@ from functools import partial from itertools import cycle, islice from pathlib import Path -from typing import TYPE_CHECKING, Any, BinaryIO, Callable, Optional, Union +from typing import TYPE_CHECKING, Any, BinaryIO, Callable, Optional, Sequence, Union import fsspec.asyn import multiprocess as mp @@ -3161,7 +3161,7 @@ def from_list( @staticmethod def from_csv( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, @@ -3202,9 +3202,10 @@ def from_csv( **kwargs, ).read() + @staticmethod @staticmethod def from_json( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, @@ -3249,9 +3250,28 @@ def from_json( **kwargs, ).read() + @staticmethod + def from_jsonl( + path_or_paths: Union[PathLike, Sequence[PathLike]], + split: Optional[NamedSplit] = None, + features: Optional[Features] = None, + keep_in_memory: bool = False, + field: Optional[str] = None, + **kwargs, + ) -> "IterableDataset": + """Alias of `IterableDataset.from_json`.""" + return IterableDataset.from_json( + path_or_paths, + split=split, + features=features, + keep_in_memory=keep_in_memory, + field=field, + **kwargs, + ) + @staticmethod def from_parquet( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, @@ -3336,7 +3356,7 @@ def from_parquet( @staticmethod def from_text( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 18a8a6038fe..b382af98883 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -4028,6 +4028,15 @@ def _check_json_dataset(dataset, expected_features): assert dataset.features[feature].dtype == expected_dtype +def test_dataset_from_jsonl(jsonl_path, tmp_path): + cache_dir = tmp_path / "cache" + expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"} + dataset_json = Dataset.from_json(jsonl_path, cache_dir=cache_dir) + dataset_jsonl = Dataset.from_jsonl(jsonl_path, cache_dir=cache_dir) + assert dataset_json.data == dataset_jsonl.data + _check_json_dataset(dataset_jsonl, expected_features) + + @pytest.mark.parametrize("keep_in_memory", [False, True]) def test_dataset_from_json_keep_in_memory(keep_in_memory, jsonl_path, tmp_path): cache_dir = tmp_path / "cache" diff --git a/tests/test_dataset_dict.py b/tests/test_dataset_dict.py index 50a14ccfdb6..cb13a95a114 100644 --- a/tests/test_dataset_dict.py +++ b/tests/test_dataset_dict.py @@ -759,6 +759,14 @@ def _check_json_datasetdict(dataset_dict, expected_features, splits=("train",)): assert dataset.features[feature].dtype == expected_dtype +def test_datasetdict_from_jsonl(jsonl_path, tmp_path): + cache_dir = tmp_path / "cache" + expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"} + dataset_json = DatasetDict.from_json({"train": jsonl_path}, cache_dir=cache_dir) + dataset_jsonl = DatasetDict.from_jsonl({"train": jsonl_path}, cache_dir=cache_dir) + assert dataset_json["train"].data == dataset_jsonl["train"].data + _check_json_datasetdict(dataset_jsonl, expected_features) + @pytest.mark.parametrize("keep_in_memory", [False, True]) def test_datasetdict_from_json_keep_in_memory(keep_in_memory, jsonl_path, tmp_path): cache_dir = tmp_path / "cache"