maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/dataframe/datasource/read_parquet.py

@@ -0,0 +1,421 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 1999-2024 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 typing import Dict
+from urllib.parse import urlparse
+
+import numpy as np
+import pandas as pd
+
+try:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+except ImportError:
+    pa = None
+
+try:
+    import fastparquet
+except ImportError:
+    fastparquet = None
+
+from ... import opcodes
+from ...config import options
+from ...lib.filesystem import FileSystem, get_fs, glob, open_file
+from ...serialization.serializables import (
+    AnyField,
+    BoolField,
+    DictField,
+    Int32Field,
+    Int64Field,
+    ListField,
+    StringField,
+)
+from ...utils import lazy_import
+from ..arrays import ArrowStringDtype
+from ..operators import OutputType
+from ..utils import parse_index, to_arrow_dtypes
+from .core import ColumnPruneSupportedDataSourceMixin, IncrementalIndexDatasource
+
+PARQUET_MEMORY_SCALE = 15
+STRING_FIELD_OVERHEAD = 50
+cudf = lazy_import("cudf")
+
+
+def check_engine(engine):
+    if engine == "auto":
+        if pa is not None:
+            return "pyarrow"
+        elif fastparquet is not None:  # pragma: no cover
+            return "fastparquet"
+        else:  # pragma: no cover
+            raise RuntimeError("Please install either pyarrow or fastparquet.")
+    elif engine == "pyarrow":
+        if pa is None:  # pragma: no cover
+            raise RuntimeError("Please install pyarrow first.")
+        return engine
+    elif engine == "fastparquet":
+        if fastparquet is None:  # pragma: no cover
+            raise RuntimeError("Please install fastparquet first.")
+        return engine
+    else:  # pragma: no cover
+        raise RuntimeError("Unsupported engine {} to read parquet.".format(engine))
+
+
+def get_engine(engine):
+    if engine == "pyarrow":
+        return ArrowEngine()
+    elif engine == "fastparquet":
+        return FastpaquetEngine()
+    else:  # pragma: no cover
+        raise RuntimeError("Unsupported engine {}".format(engine))
+
+
+class ParquetEngine:
+    def get_row_num(self, f):
+        raise NotImplementedError
+
+    def read_dtypes(self, f, **kwargs):
+        raise NotImplementedError
+
+    def read_to_pandas(
+        self, f, columns=None, nrows=None, use_arrow_dtype=None, **kwargs
+    ):
+        raise NotImplementedError
+
+    def read_group_to_pandas(
+        self, f, group_index, columns=None, nrows=None, use_arrow_dtype=None, **kwargs
+    ):
+        raise NotImplementedError
+
+    def read_partitioned_to_pandas(
+        self,
+        f,
+        partitions: Dict,
+        partition_keys: Dict,
+        columns=None,
+        nrows=None,
+        use_arrow_dtype=None,
+        **kwargs,
+    ):
+        raw_df = self.read_to_pandas(
+            f, columns=columns, nrows=nrows, use_arrow_dtype=use_arrow_dtype, **kwargs
+        )
+        for col, value in partition_keys.items():
+            dictionary = partitions[col]
+            raw_df[col] = pd.Series(
+                value,
+                dtype=pd.CategoricalDtype(categories=dictionary.tolist()),
+                index=raw_df.index,
+            )
+        return raw_df
+
+    def read_partitioned_dtypes(self, fs: FileSystem, directory, storage_options):
+        # As ParquetDataset will iterate all files,
+        # here we just find one file to infer dtypes
+        current_path = directory
+        partition_cols = []
+        while fs.isdir(current_path):
+            _, dirs, files = next(fs.walk(current_path))
+            dirs = [d for d in dirs if not d.startswith(".")]
+            files = [f for f in files if not f.startswith(".")]
+            if len(files) == 0:
+                # directory as partition
+                partition_cols.append(dirs[0].split("=", 1)[0])
+                current_path = os.path.join(current_path, dirs[0])
+            elif len(dirs) == 0:
+                # parquet files in deepest directory
+                current_path = os.path.join(current_path, files[0])
+            else:  # pragma: no cover
+                raise ValueError(
+                    "Files and directories are mixed in an intermediate directory"
+                )
+
+        # current path is now a parquet file
+        with open_file(current_path, storage_options=storage_options) as f:
+            dtypes = self.read_dtypes(f)
+        for partition in partition_cols:
+            dtypes[partition] = pd.CategoricalDtype()
+        return dtypes
+
+
+def _parse_prefix(path):
+    path_prefix = ""
+    if isinstance(path, str):
+        parsed_path = urlparse(path)
+        if parsed_path.scheme:
+            path_prefix = f"{parsed_path.scheme}://{parsed_path.netloc}"
+    return path_prefix
+
+
+class ArrowEngine(ParquetEngine):
+    def get_row_num(self, f):
+        file = pq.ParquetFile(f)
+        return file.metadata.num_rows
+
+    def read_dtypes(self, f, **kwargs):
+        file = pq.ParquetFile(f)
+        return file.schema_arrow.empty_table().to_pandas().dtypes
+
+    @classmethod
+    def _table_to_pandas(cls, t, nrows=None, use_arrow_dtype=None):
+        if nrows is not None:
+            t = t.slice(0, nrows)
+        if use_arrow_dtype:
+            df = t.to_pandas(types_mapper={pa.string(): ArrowStringDtype()}.get)
+        else:
+            df = t.to_pandas()
+        return df
+
+    def read_to_pandas(
+        self, f, columns=None, nrows=None, use_arrow_dtype=None, **kwargs
+    ):
+        file = pq.ParquetFile(f)
+        t = file.read(columns=columns, **kwargs)
+        return self._table_to_pandas(t, nrows=nrows, use_arrow_dtype=use_arrow_dtype)
+
+    def read_group_to_pandas(
+        self, f, group_index, columns=None, nrows=None, use_arrow_dtype=None, **kwargs
+    ):
+        file = pq.ParquetFile(f)
+        t = file.read_row_group(group_index, columns=columns, **kwargs)
+        return self._table_to_pandas(t, nrows=nrows, use_arrow_dtype=use_arrow_dtype)
+
+
+class FastpaquetEngine(ParquetEngine):
+    def get_row_num(self, f):
+        file = fastparquet.ParquetFile(f)
+        return file.count()
+
+    def read_dtypes(self, f, **kwargs):
+        file = fastparquet.ParquetFile(f)
+        dtypes_dict = file._dtypes()
+        return pd.Series(dict((c, dtypes_dict[c]) for c in file.columns))
+
+    def read_to_pandas(
+        self, f, columns=None, nrows=None, use_arrow_dtype=None, **kwargs
+    ):
+        file = fastparquet.ParquetFile(f)
+        df = file.to_pandas(columns, **kwargs)
+        if nrows is not None:
+            df = df.head(nrows)
+        if use_arrow_dtype:
+            df = df.astype(to_arrow_dtypes(df.dtypes).to_dict())
+        return df
+
+
+class CudfEngine:
+    @classmethod
+    def read_to_cudf(cls, file, columns: list = None, nrows: int = None, **kwargs):
+        df = cudf.read_parquet(file, columns=columns, **kwargs)
+        if nrows is not None:
+            df = df.head(nrows)
+        return df
+
+    def read_group_to_cudf(
+        self, file, group_index: int, columns: list = None, nrows: int = None, **kwargs
+    ):
+        return self.read_to_cudf(
+            file, columns=columns, nrows=nrows, row_groups=group_index, **kwargs
+        )
+
+    @classmethod
+    def read_partitioned_to_cudf(
+        cls,
+        file,
+        partitions: Dict,
+        partition_keys: Dict,
+        columns=None,
+        nrows=None,
+        **kwargs,
+    ):
+        # cudf will read entire partitions even if only one partition provided,
+        # so we just read with pyarrow and convert to cudf DataFrame
+        file = pq.ParquetFile(file)
+        t = file.read(columns=columns, **kwargs)
+        t = t.slice(0, nrows) if nrows is not None else t
+        t = pa.table(t.columns, names=t.column_names)
+        raw_df = cudf.DataFrame.from_arrow(t)
+        for col, value in partition_keys.items():
+            dictionary = partitions[col].tolist()
+            codes = cudf.core.column.as_column(
+                dictionary.index(value), length=len(raw_df)
+            )
+            raw_df[col] = cudf.core.column.build_categorical_column(
+                categories=dictionary,
+                codes=codes,
+                size=codes.size,
+                offset=codes.offset,
+                ordered=False,
+            )
+        return raw_df
+
+
+class DataFrameReadParquet(
+    IncrementalIndexDatasource,
+    ColumnPruneSupportedDataSourceMixin,
+):
+    _op_type_ = opcodes.READ_PARQUET
+
+    path = AnyField("path")
+    engine = StringField("engine")
+    columns = ListField("columns")
+    use_arrow_dtype = BoolField("use_arrow_dtype")
+    groups_as_chunks = BoolField("groups_as_chunks")
+    group_index = Int32Field("group_index")
+    read_kwargs = DictField("read_kwargs")
+    incremental_index = BoolField("incremental_index")
+    storage_options = DictField("storage_options")
+    is_partitioned = BoolField("is_partitioned")
+    merge_small_files = BoolField("merge_small_files")
+    merge_small_file_options = DictField("merge_small_file_options")
+    # for chunk
+    partitions = DictField("partitions", default=None)
+    partition_keys = DictField("partition_keys", default=None)
+    num_group_rows = Int64Field("num_group_rows", default=None)
+    # as read meta may be too time-consuming when number of files is large,
+    # thus we only read first file to get row number and raw file size
+    first_chunk_row_num = Int64Field("first_chunk_row_num")
+    first_chunk_raw_bytes = Int64Field("first_chunk_raw_bytes")
+
+    def get_columns(self):
+        return self.columns
+
+    def set_pruned_columns(self, columns, *, keep_order=None):
+        self.columns = columns
+
+    def __call__(self, index_value=None, columns_value=None, dtypes=None):
+        self._output_types = [OutputType.dataframe]
+        shape = (np.nan, len(dtypes))
+        return self.new_dataframe(
+            None,
+            shape,
+            dtypes=dtypes,
+            index_value=index_value,
+            columns_value=columns_value,
+        )
+
+
+def read_parquet(
+    path,
+    engine: str = "auto",
+    columns: list = None,
+    groups_as_chunks: bool = False,
+    use_arrow_dtype: bool = None,
+    incremental_index: bool = False,
+    storage_options: dict = None,
+    memory_scale: int = None,
+    merge_small_files: bool = True,
+    merge_small_file_options: dict = None,
+    gpu: bool = None,
+    **kwargs,
+):
+    """
+    Load a parquet object from the file path, returning a DataFrame.
+
+    Parameters
+    ----------
+    path : str, path object or file-like object
+        Any valid string path is acceptable. The string could be a URL.
+        For file URLs, a host is expected. A local file could be:
+        ``file://localhost/path/to/table.parquet``.
+        A file URL can also be a path to a directory that contains multiple
+        partitioned parquet files. Both pyarrow and fastparquet support
+        paths to directories as well as file URLs. A directory path could be:
+        ``file://localhost/path/to/tables``.
+        By file-like object, we refer to objects with a ``read()`` method,
+        such as a file handler (e.g. via builtin ``open`` function)
+        or ``StringIO``.
+    engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto'
+        Parquet library to use. The default behavior is to try 'pyarrow',
+        falling back to 'fastparquet' if 'pyarrow' is unavailable.
+    columns : list, default=None
+        If not None, only these columns will be read from the file.
+    groups_as_chunks : bool, default False
+        if True, each row group correspond to a chunk.
+        if False, each file correspond to a chunk.
+        Only available for 'pyarrow' engine.
+    incremental_index: bool, default False
+        If index_col not specified, ensure range index incremental,
+        gain a slightly better performance if setting False.
+    use_arrow_dtype: bool, default None
+        If True, use arrow dtype to store columns.
+    storage_options: dict, optional
+        Options for storage connection.
+    memory_scale: int, optional
+        Scale that real memory occupation divided with raw file size.
+    merge_small_files: bool, default True
+        Merge small files whose size is small.
+    merge_small_file_options: dict
+        Options for merging small files
+    **kwargs
+        Any additional kwargs are passed to the engine.
+
+    Returns
+    -------
+    MaxFrame DataFrame
+    """
+
+    engine_type = check_engine(engine)
+    engine = get_engine(engine_type)
+
+    single_path = path[0] if isinstance(path, list) else path
+    fs = get_fs(single_path, storage_options)
+    is_partitioned = False
+    if fs.isdir(single_path):
+        paths = fs.ls(path)
+        if all(fs.isdir(p) for p in paths):
+            # If all are directories, it is read as a partitioned dataset.
+            dtypes = engine.read_partitioned_dtypes(fs, path, storage_options)
+            is_partitioned = True
+        else:
+            with fs.open(paths[0], mode="rb") as f:
+                dtypes = engine.read_dtypes(f)
+    else:
+        if not isinstance(path, list):
+            file_path = glob(path, storage_options=storage_options)[0]
+        else:
+            file_path = path[0]
+
+        with open_file(file_path, storage_options=storage_options) as f:
+            dtypes = engine.read_dtypes(f)
+
+    if columns:
+        dtypes = dtypes[columns]
+
+    if use_arrow_dtype is None:
+        use_arrow_dtype = options.dataframe.use_arrow_dtype
+    if use_arrow_dtype:
+        dtypes = to_arrow_dtypes(dtypes)
+
+    index_value = parse_index(pd.RangeIndex(-1))
+    columns_value = parse_index(dtypes.index, store_data=True)
+    op = DataFrameReadParquet(
+        path=path,
+        engine=engine_type,
+        columns=columns,
+        groups_as_chunks=groups_as_chunks,
+        use_arrow_dtype=use_arrow_dtype,
+        read_kwargs=kwargs,
+        incremental_index=incremental_index,
+        storage_options=storage_options,
+        is_partitioned=is_partitioned,
+        memory_scale=memory_scale,
+        merge_small_files=merge_small_files,
+        merge_small_file_options=merge_small_file_options,
+        gpu=gpu,
+    )
+    return op(index_value=index_value, columns_value=columns_value, dtypes=dtypes)

maxframe/dataframe/datasource/series.py

@@ -0,0 +1,55 @@
+# Copyright 1999-2024 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 ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import DataTypeField, SeriesField
+from ..operators import DataFrameOperatorMixin
+from ..utils import is_cudf, parse_index
+from .core import PandasDataSourceOperator
+
+
+class SeriesDataSource(PandasDataSourceOperator, DataFrameOperatorMixin):
+    """
+    Represents data from pandas Series
+    """
+
+    _op_type_ = opcodes.SERIES_DATA_SOURCE
+
+    data = SeriesField("data")
+    dtype = DataTypeField("dtype")
+
+    def __init__(self, data=None, dtype=None, gpu=None, **kw):
+        if dtype is None and data is not None:
+            dtype = data.dtype
+        if gpu is None and is_cudf(data):  # pragma: no cover
+            gpu = True
+        super().__init__(
+            data=data, dtype=dtype, gpu=gpu, _output_types=[OutputType.series], **kw
+        )
+
+    def __call__(self, shape, chunk_size=None):
+        return self.new_series(
+            None,
+            shape=shape,
+            dtype=self.dtype,
+            index_value=parse_index(self.data.index),
+            name=self.data.name,
+            raw_chunk_size=chunk_size,
+        )
+
+
+def from_pandas(data, chunk_size=None, gpu=None, sparse=False):
+    op = SeriesDataSource(data=data, gpu=gpu, sparse=sparse)
+    return op(data.shape, chunk_size=chunk_size)

maxframe/dataframe/datasource/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2024 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.