From ce93e1b72fa4ea8f8398a5dd74230fc41d3e5713 Mon Sep 17 00:00:00 2001 From: codingl2k1 Date: Wed, 23 Aug 2023 17:53:45 +0800 Subject: [PATCH 1/2] Impl from_export --- .../xorbits/_mars/core/entity/output_types.py | 1 + python/xorbits/_mars/services/meta/metas.py | 8 ++ python/xorbits/datasets/__init__.py | 1 + .../datasets/backends/arrow/__init__.py | 20 +++ .../xorbits/datasets/backends/arrow/core.py | 94 +++++++++++++ .../datasets/backends/arrow/from_export.py | 127 ++++++++++++++++++ python/xorbits/datasets/iterable_dataset.py | 4 + .../datasets/tests/test_iterable_dataset.py | 18 +++ 8 files changed, 273 insertions(+) create mode 100644 python/xorbits/datasets/backends/arrow/__init__.py create mode 100644 python/xorbits/datasets/backends/arrow/core.py create mode 100644 python/xorbits/datasets/backends/arrow/from_export.py diff --git a/python/xorbits/_mars/core/entity/output_types.py b/python/xorbits/_mars/core/entity/output_types.py index 5c7137c49..ba3a4d5b1 100644 --- a/python/xorbits/_mars/core/entity/output_types.py +++ b/python/xorbits/_mars/core/entity/output_types.py @@ -32,6 +32,7 @@ class OutputType(Enum): series_groupby = 9 df_or_series = 10 huggingface_dataset = 11 + arrow_dataset = 12 @classmethod def serialize_list(cls, output_types): diff --git a/python/xorbits/_mars/services/meta/metas.py b/python/xorbits/_mars/services/meta/metas.py index e6df8b9ba..18b3e49a3 100644 --- a/python/xorbits/_mars/services/meta/metas.py +++ b/python/xorbits/_mars/services/meta/metas.py @@ -18,6 +18,7 @@ import numpy as np +from ....datasets.backends.arrow.core import ArrowDatasetChunk, ArrowDatasetChunkData from ....datasets.backends.huggingface.core import ( HuggingfaceDatasetChunk, HuggingfaceDatasetChunkData, @@ -202,6 +203,13 @@ class DatasetChunkMeta(_ChunkMeta): shape: Tuple[int] = None +@register_meta_type((ArrowDatasetChunk, ArrowDatasetChunkData)) +@dataslots +@dataclass +class ArrowChunkMeta(_ChunkMeta): + shape: Tuple[int] = None + + @register_meta_type(DATAFRAME_OR_SERIES_TYPE) @dataslots @dataclass diff --git a/python/xorbits/datasets/__init__.py b/python/xorbits/datasets/__init__.py index a265e6cc8..bf8470093 100644 --- a/python/xorbits/datasets/__init__.py +++ b/python/xorbits/datasets/__init__.py @@ -13,6 +13,7 @@ # limitations under the License. from .adapter import MARS_DATASET_CALLABLES, MARS_DATASET_TYPE, _install +from .backends.arrow.from_export import from_export from .backends.huggingface.from_huggingface import from_huggingface from .dataset import Dataset diff --git a/python/xorbits/datasets/backends/arrow/__init__.py b/python/xorbits/datasets/backends/arrow/__init__.py new file mode 100644 index 000000000..b163239ec --- /dev/null +++ b/python/xorbits/datasets/backends/arrow/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2022-2023 XProbe Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import ( + ArrowDataset, + ArrowDatasetChunk, + ArrowDatasetChunkData, + ArrowDatasetData, +) diff --git a/python/xorbits/datasets/backends/arrow/core.py b/python/xorbits/datasets/backends/arrow/core.py new file mode 100644 index 000000000..91462c8f8 --- /dev/null +++ b/python/xorbits/datasets/backends/arrow/core.py @@ -0,0 +1,94 @@ +# Copyright 2022-2023 XProbe Inc. +# derived from copyright 1999-2021 Alibaba Group Holding Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict + +from ...._mars.core import is_build_mode +from ...._mars.core.entity import ( + OutputType, + register_fetch_class, + register_output_types, +) +from ...._mars.core.entity.utils import refresh_tileable_shape +from ...._mars.core.operand.objects import ObjectFetch +from ...._mars.serialization.serializables import FieldTypes, ListField +from ...dataset import Dataset, DatasetChunk, DatasetChunkData, DatasetData + + +class ArrowDatasetChunkData(DatasetChunkData): + __slots__ = () + type_name = "ArrowDatasetChunkData" + + @classmethod + def get_params_from_data(cls, data) -> Dict[str, Any]: + """For updating chunk shape from data.""" + return {"shape": data.shape} + + +class ArrowDatasetChunk(DatasetChunk): + __slots__ = () + _allow_data_type_ = (ArrowDatasetChunkData,) + type_name = "ArrowDatasetChunk" + + +class ArrowDatasetData(DatasetData): + __slots__ = () + type_name = "Arrow Dataset" + + _chunks = ListField( + "chunks", + FieldTypes.reference(ArrowDatasetChunk), + on_serialize=lambda x: [it.data for it in x] if x is not None else x, + on_deserialize=lambda x: [ArrowDatasetChunk(it) for it in x] + if x is not None + else x, + ) + + def __repr__(self): + if is_build_mode() or len(self._executed_sessions) == 0: + # in build mode, or not executed, just return representation + return f"Arrow Dataset " + else: + try: + return f"Dataset({{\n features: {self.dtypes.index.values.tolist()},\n num_rows: {self.shape[0]}\n}})" + except: # noqa: E722 # nosec # pylint: disable=bare-except # pragma: no cover + return f"Arrow Dataset " + + def refresh_params(self): + refresh_tileable_shape(self) + # TODO(codingl2k1): update dtypes. + + +class ArrowDataset(Dataset): + __slots__ = () + _allow_data_type_ = (ArrowDatasetData,) + type_name = "Arrow Dataset" + + def to_dataset(self): + return Dataset(self.data) + + +register_output_types( + OutputType.arrow_dataset, + (ArrowDataset, ArrowDatasetData), + (ArrowDatasetChunk, ArrowDatasetChunkData), +) + + +class ArrowDatasetFetch(ObjectFetch): + _output_type_ = OutputType.arrow_dataset + + +register_fetch_class(OutputType.arrow_dataset, ArrowDatasetFetch, None) diff --git a/python/xorbits/datasets/backends/arrow/from_export.py b/python/xorbits/datasets/backends/arrow/from_export.py new file mode 100644 index 000000000..893bdd6e9 --- /dev/null +++ b/python/xorbits/datasets/backends/arrow/from_export.py @@ -0,0 +1,127 @@ +# Copyright 2022-2023 XProbe Inc. +# derived from copyright 1999-2021 Alibaba Group Holding Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Union + +import fsspec +import pyarrow as pa + +from ...._mars.core.entity import OutputType +from ...._mars.serialization.serializables import ( + AnyField, + ListField, + StringField, + TupleField, +) +from ...iterable_dataset import IterableDataset +from ...operand import DataOperand, DataOperandMixin + + +class FromExport(DataOperand, DataOperandMixin): + path = StringField("path") + groups = ListField("groups") + index = TupleField("index") + iterable_dataset = AnyField("iterable_dataset") + + def __call__(self): + ids = self.iterable_dataset + return self.new_tileable( + [], dtypes=ids.schema.empty_table().to_pandas().dtypes, shape=ids.shape + ) + + @classmethod + def tile(cls, op: "FromExport"): + assert len(op.inputs) == 0 + out = op.outputs[0] + ids: IterableDataset = op.iterable_dataset + op.iterable_dataset = None + + chunks = [] + default_group = ids.group_infos()[0] + splits = default_group.index[1] + for cdx, index in enumerate(default_group.index[0]): + chunk_op = op.copy().reset_key() + chunk_op.index = index.as_py() + c = chunk_op.new_chunk( + inputs=[], index=(cdx, 0), shape=(splits[cdx].as_py(), ids.shape[1]) + ) + chunks.append(c) + return op.copy().new_tileable( + op.inputs, + chunks=chunks, + nsplits=(splits.to_pylist(), (ids.shape[1],)), + **out.params, + ) + + @classmethod + def execute(cls, ctx, op: "FromExport"): + arrow_file_paths = [ + os.path.join( + op.path, name, IterableDataset._FILE_NAME_FORMATTER.format(*op.index) + ) + for name in op.groups + ] + + def _load_arrow_table(filepath): + # TODO(codingl2k1): mmap if local. + with fsspec.open(filepath, "rb") as f: + with pa.ipc.RecordBatchStreamReader(f) as reader: + return reader.read_all() + + futures = [] + with ThreadPoolExecutor( + thread_name_prefix=IterableDataset._get_infos.__qualname__ + ) as executor: + for arrow_file in arrow_file_paths: + futures.append(executor.submit(_load_arrow_table, arrow_file)) + + arrow_tables = [fut.result() for fut in futures] + result_table = arrow_tables[0] + # TODO(codingl2k1): Better way to concat table columns. + for table in arrow_tables[1:]: + for idx, col in enumerate(table.itercolumns()): + result_table = result_table.append_column(table.field(idx), col) + ctx[op.outputs[0].key] = result_table + + +def from_export( + path: Union[str, os.PathLike], + storage_options: Optional[dict] = None, + version: Optional[str] = None, +): + """Create a dataset from exported Dataset. + + Parameters + ---------- + path: str + The export path. + storage_options: dict + Key/value pairs to be passed on to the caching file-system backend, if any. + version: str + The dataset version. + + Returns + ------- + Dataset + """ + ids = IterableDataset(path=path, storage_options=storage_options, version=version) + op = FromExport( + output_types=[OutputType.arrow_dataset], + iterable_dataset=ids, + path=ids.path, + groups=ids.groups, + ) + return op().to_dataset() diff --git a/python/xorbits/datasets/iterable_dataset.py b/python/xorbits/datasets/iterable_dataset.py index 34c0ffc45..9ac13e987 100644 --- a/python/xorbits/datasets/iterable_dataset.py +++ b/python/xorbits/datasets/iterable_dataset.py @@ -518,6 +518,10 @@ def _get_formatter(self) -> Formatter: # Fast path without decoding. return Formatter() + @functools.cached_property + def path(self): + return self._path + @functools.cached_property def groups(self) -> List[str]: return self._info["groups"] diff --git a/python/xorbits/datasets/tests/test_iterable_dataset.py b/python/xorbits/datasets/tests/test_iterable_dataset.py index b55ca9d49..5c42b1466 100644 --- a/python/xorbits/datasets/tests/test_iterable_dataset.py +++ b/python/xorbits/datasets/tests/test_iterable_dataset.py @@ -19,9 +19,11 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import datasets import pytest from ..._mars.tests.core import mock +from ..backends.arrow.from_export import from_export from ..backends.huggingface.from_huggingface import from_huggingface from ..iterable_dataset import IterableDataset, map_retry @@ -245,3 +247,19 @@ def test_iterable_dataset(): ) finally: shutil.rmtree(export_dir, ignore_errors=True) + + +def test_from_export(setup): + tmp_dir = Path(tempfile.gettempdir()) + export_dir = tmp_dir.joinpath("test_iterable_dataset") + shutil.rmtree(export_dir, ignore_errors=True) + db = from_huggingface("cifar10", split="train") + db.export(export_dir) + try: + ds = from_export(export_dir) + ds.execute() + table = ds.fetch() + assert table.shape == db.shape + # TODO(codingl2k1): Add more tests + finally: + shutil.rmtree(export_dir, ignore_errors=True) From 17a51ffd244d452d5c49646c1e01b94c102ddd36 Mon Sep 17 00:00:00 2001 From: codingl2k1 Date: Wed, 23 Aug 2023 18:34:46 +0800 Subject: [PATCH 2/2] Fix lint --- python/xorbits/datasets/tests/test_iterable_dataset.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/xorbits/datasets/tests/test_iterable_dataset.py b/python/xorbits/datasets/tests/test_iterable_dataset.py index 5c42b1466..5b304c4ce 100644 --- a/python/xorbits/datasets/tests/test_iterable_dataset.py +++ b/python/xorbits/datasets/tests/test_iterable_dataset.py @@ -19,7 +19,6 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path -import datasets import pytest from ..._mars.tests.core import mock