maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/dataframe/datasource/read_odps_table.py

@@ -0,0 +1,295 @@
+# Copyright 1999-2025 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 logging
+from typing import List, MutableMapping, Optional, Union
+
+import numpy as np
+import pandas as pd
+from odps import ODPS
+from odps.models import Table
+from odps.utils import to_timestamp
+
+from ... import opcodes
+from ...config import option_context, options
+from ...core import OutputType
+from ...io.odpsio import odps_schema_to_pandas_dtypes
+from ...serialization.serializables import (
+    AnyField,
+    BoolField,
+    FieldTypes,
+    Int64Field,
+    ListField,
+    SeriesField,
+    StringField,
+)
+from ...utils import estimate_table_size, is_empty
+from ..core import DataFrame  # noqa: F401
+from ..utils import parse_index, validate_dtype_backend
+from .core import (
+    ColumnPruneSupportedDataSourceMixin,
+    DtypeBackendCompatibleMixin,
+    IncrementalIndexDatasource,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class DataFrameReadODPSTable(
+    IncrementalIndexDatasource,
+    ColumnPruneSupportedDataSourceMixin,
+    DtypeBackendCompatibleMixin,
+):
+    __slots__ = ("_odps_entry",)
+    _op_type_ = opcodes.READ_ODPS_TABLE
+
+    table_name = StringField("table_name")
+    partitions = ListField("partitions", field_type=FieldTypes.string, default=None)
+    dtypes = SeriesField("dtypes", default=None)
+    columns = AnyField("columns", default=None)
+    nrows = Int64Field("nrows", default=None)
+    dtype_backend = StringField("dtype_backend", default=None)
+    string_as_binary = BoolField("string_as_binary", default=None)
+    append_partitions = BoolField("append_partitions", default=None)
+    last_modified_time = Int64Field("last_modified_time", default=None)
+    index_columns = ListField("index_columns", FieldTypes.string, default=None)
+    index_dtypes = SeriesField("index_dtypes", default=None)
+
+    def __init__(self, memory_scale=None, dtype_backend=None, **kw):
+        output_type = kw.pop("output_type", OutputType.dataframe)
+        self._odps_entry = kw.pop("odps_entry", None)
+        dtype_backend = validate_dtype_backend(dtype_backend)
+        super(DataFrameReadODPSTable, self).__init__(
+            memory_scale=memory_scale,
+            dtype_backend=dtype_backend,
+            _output_types=[output_type],
+            **kw,
+        )
+
+    @property
+    def partition(self):
+        return getattr(self, "partition_spec", None)
+
+    def get_columns(self):
+        return self.columns or list(self.dtypes.index)
+
+    def set_pruned_columns(self, columns, *, keep_order=None):  # pragma: no cover
+        self.columns = columns
+
+    def __call__(self, shape, chunk_bytes=None, chunk_size=None):
+        if is_empty(self.index_columns):
+            if np.isnan(shape[0]):
+                index_value = parse_index(pd.RangeIndex(0))
+            else:
+                index_value = parse_index(pd.RangeIndex(shape[0]))
+        elif len(self.index_columns) == 1:
+            index_value = parse_index(
+                pd.Index([], name=self.index_columns[0]).astype(
+                    self.index_dtypes.iloc[0]
+                )
+            )
+        else:
+            idx = pd.MultiIndex.from_frame(
+                pd.DataFrame([], columns=self.index_columns).astype(self.index_dtypes)
+            )
+            index_value = parse_index(idx)
+
+        if self.output_types[0] == OutputType.dataframe:
+            columns_value = parse_index(self.dtypes.index, store_data=True)
+            return self.new_tileable(
+                [],
+                None,
+                shape=shape,
+                dtypes=self.dtypes,
+                index_value=index_value,
+                columns_value=columns_value,
+                chunk_bytes=chunk_bytes,
+                chunk_size=chunk_size,
+            )
+        elif self.output_types[0] == OutputType.series:
+            return self.new_tileable(
+                [],
+                None,
+                shape=shape[:1],
+                name=self.dtypes.index[0],
+                dtype=self.dtypes.iloc[0],
+                index_value=index_value,
+                chunk_bytes=chunk_bytes,
+                chunk_size=chunk_size,
+            )
+        elif self.output_types[0] == OutputType.index:
+            return self.new_tileable(
+                [],
+                None,
+                shape=shape[:1],
+                name=getattr(index_value, "name", None),
+                names=getattr(index_value, "names", None),
+                dtype=self.index_dtypes.iloc[0],
+                index_value=index_value,
+                chunk_bytes=chunk_bytes,
+                chunk_size=chunk_size,
+            )
+
+    @classmethod
+    def estimate_size(
+        cls, ctx: MutableMapping[str, Union[int, float]], op: "DataFrameReadODPSTable"
+    ) -> None:
+        odps_entry = op._odps_entry or ODPS.from_global() or ODPS.from_environments()
+        if not odps_entry:  # pragma: no cover
+            ctx[op.outputs[0].key] = float("inf")
+            return
+        ctx[op.outputs[0].key] = estimate_table_size(
+            odps_entry, op.table_name, op.partitions
+        )
+
+
+def read_odps_table(
+    table_name: Union[str, Table],
+    partitions: Union[None, str, List[str]] = None,
+    columns: Optional[List[str]] = None,
+    index_col: Union[None, str, List[str]] = None,
+    odps_entry: ODPS = None,
+    string_as_binary: bool = None,
+    append_partitions: bool = False,
+    dtype_backend: str = None,
+    **kw,
+):
+    """
+    Read data from a MaxCompute (ODPS) table into DataFrame.
+
+    Supports specifying some columns as indexes. If not specified, RangeIndex
+    will be generated.
+
+    Parameters
+    ----------
+    table_name: Union[str, Table]
+        Name of the table to read from.
+    partitions: Union[None, str, List[str]]
+        Table partition or list of partitions to read from.
+    columns: Optional[List[str]]
+        Table columns to read from. You may also specify partition columns here.
+        If not specified, all table columns (or include partition columns if
+        `append_partitions` is True) will be included.
+    index_col: Union[None, str, List[str]]
+        Columns to be specified as indexes.
+    append_partitions: bool
+        If True, will add all partition columns as selected columns when
+        `columns` is not specified,
+    dtype_backend: {'numpy', 'pyarrow'}, default 'numpy'
+        Back-end data type applied to the resultant DataFrame (still experimental).
+
+    Returns
+    -------
+    result: DataFrame
+        DataFrame read from MaxCompute (ODPS) table
+    """
+    odps_entry = odps_entry or ODPS.from_global() or ODPS.from_environments()
+    schema = options.session.default_schema or odps_entry.schema
+    if odps_entry is None:
+        raise ValueError("Missing odps_entry parameter")
+    if isinstance(table_name, Table):
+        table = table_name
+    else:
+        table = odps_entry.get_table(table_name, schema=schema)
+
+    if not table.table_schema.partitions and (
+        partitions is not None or append_partitions
+    ):
+        raise ValueError("Cannot handle partitions on non-partition tables")
+
+    cols = (
+        table.table_schema.columns
+        if append_partitions or columns is not None
+        else table.table_schema.simple_columns
+    )
+    table_columns = [c.name.lower() for c in cols]
+
+    if "use_arrow_dtype" in kw:
+        dtype_backend = dtype_backend or validate_dtype_backend(
+            kw.pop("use_arrow_dtype")
+        )
+    dtype_backend = validate_dtype_backend(
+        dtype_backend or options.dataframe.dtype_backend
+    )
+
+    with option_context():
+        options.dataframe.dtype_backend = dtype_backend
+        table_dtypes = odps_schema_to_pandas_dtypes(
+            table.table_schema, with_partitions=True
+        )
+    df_types = [table_dtypes[c] for c in table_columns]
+
+    if isinstance(index_col, str):
+        index_col = [index_col]
+    if index_col and columns is None:
+        index_col_set = set([c.lower() for c in index_col])
+        col_diff = sorted(index_col_set - set(table_columns))
+        if col_diff:
+            raise ValueError(
+                f"Cannot found column {', '.join(col_diff)} specified "
+                f"in index_columns argument in source table"
+            )
+        columns = [c for c in table_columns if c not in index_col_set]
+
+    if not index_col:
+        index_dtypes = None
+    else:
+        table_index_types = [df_types[table_columns.index(col)] for col in index_col]
+        index_dtypes = pd.Series(table_index_types, index=index_col)
+
+    if columns is not None:
+        new_columns = [c.lower() for c in columns]
+        table_col_set = set(new_columns)
+        col_diff = sorted(table_col_set - set(table_columns))
+        if col_diff:
+            raise ValueError(
+                f"Cannot found column {', '.join(col_diff)} specified "
+                f"in columns argument in source table"
+            )
+        if any(col in table_col_set for col in index_col or ()):
+            raise ValueError("Index columns and columns shall not overlap.")
+
+        # reorder columns
+        df_types = [df_types[table_columns.index(col)] for col in new_columns]
+        table_columns = new_columns
+        columns = new_columns
+
+    shape = kw.pop("shape", None) or (np.nan, len(df_types))
+    dtypes = pd.Series(df_types, index=table_columns)
+    chunk_bytes = kw.pop("chunk_bytes", None)
+    chunk_size = kw.pop("chunk_size", None)
+
+    partitions = partitions or kw.get("partition")
+    if isinstance(partitions, str):
+        partitions = [partitions]
+
+    append_partitions = append_partitions or any(
+        pt.name in (columns if not is_empty(columns) else ())
+        for pt in (table.table_schema.partitions or ())
+    )
+    op = DataFrameReadODPSTable(
+        table_name=table.full_table_name,
+        partitions=partitions,
+        dtypes=dtypes,
+        columns=columns,
+        dtype_backend=dtype_backend,
+        string_as_binary=string_as_binary,
+        append_partitions=append_partitions,
+        last_modified_time=to_timestamp(table.last_data_modified_time),
+        index_columns=index_col,
+        index_dtypes=index_dtypes,
+        odps_entry=odps_entry,
+        **kw,
+    )
+    return op(shape, chunk_bytes=chunk_bytes, chunk_size=chunk_size)

maxframe/dataframe/datasource/read_parquet.py

@@ -0,0 +1,278 @@
+# Copyright 1999-2025 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 glob
+import json
+from collections import OrderedDict
+from typing import MutableMapping, Union
+from urllib.parse import urlparse
+
+import numpy as np
+import pandas as pd
+
+from ...utils import make_dtypes
+
+try:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+except ImportError:
+    pa = pq = None
+
+from ... import opcodes
+from ...config import options
+from ...serialization.serializables import (
+    BoolField,
+    DictField,
+    Int32Field,
+    Int64Field,
+    ListField,
+    StringField,
+)
+from ...utils import no_default
+from ..operators import OutputType
+from ..utils import parse_index, to_arrow_dtypes
+from .core import (
+    ColumnPruneSupportedDataSourceMixin,
+    DtypeBackendCompatibleMixin,
+    LakeDataSource,
+)
+
+PARQUET_MEMORY_SCALE = 15
+STRING_FIELD_OVERHEAD = 50
+
+
+def check_engine(engine):
+    if engine == "auto":
+        return "pyarrow"
+    elif engine in ("pyarrow",):
+        return engine
+    else:  # pragma: no cover
+        raise RuntimeError("Unsupported engine {} to read parquet.".format(engine))
+
+
+class DataFrameReadParquet(
+    LakeDataSource,
+    ColumnPruneSupportedDataSourceMixin,
+    DtypeBackendCompatibleMixin,
+):
+    _op_type_ = opcodes.READ_PARQUET
+
+    engine = StringField("engine")
+    columns = ListField("columns")
+    groups_as_chunks = BoolField("groups_as_chunks", default=None)
+    group_index = Int32Field("group_index", default=None)
+    read_kwargs = DictField("read_kwargs", default=None)
+    # 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", default=None)
+    first_chunk_raw_bytes = Int64Field("first_chunk_raw_bytes", default=None)
+
+    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):
+        if self.read_stage is not None:
+            # output for planning or meta fetching
+            self._output_types = [OutputType.scalar]
+            return self.new_tileable(None, shape=(), dtype=np.dtype("O"))
+        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,
+        )
+
+    @classmethod
+    def estimate_size(
+        cls, ctx: MutableMapping[str, Union[int, float]], op: "DataFrameReadParquet"
+    ):  # pragma: no cover
+        # todo implement this to facilitate local computation
+        ctx[op.outputs[0].key] = float("inf")
+
+
+def _resolve_dict_dtypes(dtypes_dict) -> pd.Series:
+    names = list(dtypes_dict.keys())
+    vals = [pd.api.types.pandas_dtype(dt) for dt in dtypes_dict.values()]
+    return pd.Series(vals, index=names)
+
+
+def read_parquet(
+    path,
+    engine: str = "auto",
+    columns: list = None,
+    groups_as_chunks: bool = False,
+    dtype_backend: str = no_default,
+    incremental_index: bool = False,
+    storage_options: dict = None,
+    use_nullable_dtypes: bool = no_default,
+    *,
+    dtypes: pd.Series = None,
+    index_dtypes: pd.Series = None,
+    memory_scale: int = None,
+    merge_small_files: bool = True,
+    gpu: bool = None,
+    session=None,
+    run_kwargs: dict = 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'}, default 'auto'
+        Parquet library to use. The default behavior is to try 'pyarrow',
+    storage_options: dict, optional
+        Options for storage connection.
+    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.
+    dtype_backend: {'numpy', 'pyarrow'}, default 'numpy'
+        Back-end data type applied to the resultant DataFrame (still experimental).
+    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.
+    **kwargs
+        Any additional kwargs are passed to the engine.
+
+    Returns
+    -------
+    MaxFrame DataFrame
+    """
+    from .dataframe import from_pandas
+
+    engine_type = check_engine(engine)
+
+    single_path = path[0] if isinstance(path, list) else path
+    parsed_path = urlparse(single_path)
+    local_test_mode = kwargs.pop("_local_test_mode", False)
+    if not local_test_mode and (
+        not parsed_path.scheme or parsed_path.scheme.lower() == "file"
+    ):
+        # todo chunk with multiple files and / or row groups?
+        # just read locally when path is not remote
+        local_dfs = []
+        paths = path if isinstance(path, list) else [path]
+        real_paths = []
+        for path in paths:
+            if "*" in path or "?" in path:
+                real_paths.extend(glob.glob(path))
+            else:
+                real_paths.append(path)
+        for path in real_paths:
+            kw = {}
+            if use_nullable_dtypes is not no_default:
+                kw = {"use_nullable_dtypes": use_nullable_dtypes}
+            if dtype_backend is not no_default:
+                kw = {"dtype_backend": dtype_backend}
+            local_dfs.append(
+                pd.read_parquet(path, engine=engine_type, columns=columns, **kw)
+            )
+        df = pd.concat(local_dfs) if len(local_dfs) > 1 else local_dfs[0]
+        return from_pandas(df)
+
+    is_partitioned = None
+    if dtypes is not None:
+        dtypes = make_dtypes(dtypes)
+        if index_dtypes is not None:
+            index_dtypes = make_dtypes(index_dtypes)
+    else:
+        # need to submit a job to get dtypes
+        dt_op = DataFrameReadParquet(
+            path=path,
+            engine=engine_type,
+            columns=columns,
+            groups_as_chunks=groups_as_chunks,
+            dtype_backend=dtype_backend,
+            read_kwargs=kwargs,
+            incremental_index=incremental_index,
+            storage_options=storage_options,
+            memory_scale=memory_scale,
+            merge_small_files=merge_small_files,
+            read_stage="get_dtypes",
+        )
+        run_kwargs = run_kwargs or {}
+        dt_result = json.loads(
+            dt_op().execute(session=session, **run_kwargs).fetch(session=session),
+            object_pairs_hook=OrderedDict,
+        )
+        is_partitioned = dt_result["is_partitioned"]
+        dtypes = _resolve_dict_dtypes(dt_result["dtypes"])
+        if dt_result.get("index_dtypes"):
+            index_dtypes = _resolve_dict_dtypes(dt_result["index_dtypes"])
+
+    if columns:
+        dtypes = dtypes[columns]
+
+    if dtype_backend is None:
+        dtype_backend = options.dataframe.dtype_backend
+    if dtype_backend == "pyarrow":
+        dtypes = to_arrow_dtypes(dtypes)
+
+    if index_dtypes is None:
+        index_value = parse_index(pd.RangeIndex(-1))
+    else:
+        idx_df = pd.DataFrame([], columns=index_dtypes.index).astype(index_dtypes)
+        pd_idx = pd.MultiIndex.from_frame(idx_df)
+        if len(index_dtypes) == 1:
+            pd_idx = pd_idx.get_level_values(0)
+        index_value = parse_index(pd_idx, store_data=False)
+        incremental_index = False
+
+    columns_value = parse_index(dtypes.index, store_data=True)
+    op = DataFrameReadParquet(
+        path=path,
+        engine=engine_type,
+        columns=columns,
+        groups_as_chunks=groups_as_chunks,
+        dtype_backend=dtype_backend,
+        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,
+        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-2025 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-2025 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.