FEAT: Dataset from_export

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):

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

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
 

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,
+)

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 <op={type(self.op).__name__}, key={self.key}>"
+        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 <op={type(self.op).__name__}, key={self.key}>"
+
+    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)

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()

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"]

python/xorbits/datasets/tests/test_iterable_dataset.py

@@ -22,6 +22,7 @@
 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 +246,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)