maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/dataframe/indexing/where.py

@@ -0,0 +1,308 @@
+# 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 numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import ENTITY_TYPE
+from ...serialization.serializables import AnyField, BoolField, Int32Field, StringField
+from ...tensor.utils import filter_inputs
+from ..core import DATAFRAME_TYPE, SERIES_TYPE
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import build_df, build_series, validate_axis
+
+
+class DataFrameWhere(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.WHERE
+
+    _input = AnyField("input")
+    cond = AnyField("cond", default=None)
+    other = AnyField("other", default=None)
+    axis = Int32Field("axis", default=None)
+    level = AnyField("level", default=None)
+    errors = StringField("errors", default=None)
+    try_cast = BoolField("try_cast", default=None)
+    replace_true = BoolField("replace_true", default=None)
+
+    def __init__(self, input=None, **kw):
+        super().__init__(_input=input, **kw)
+
+    @property
+    def input(self):
+        return self._input
+
+    def __call__(self, df_or_series):
+        def _check_input_index(obj, axis=None):
+            axis = axis if axis is not None else self.axis
+            if isinstance(obj, DATAFRAME_TYPE) and (
+                df_or_series.columns_value.key != obj.columns_value.key
+                or df_or_series.index_value.key != obj.index_value.key
+            ):
+                raise NotImplementedError("Aligning different indices not supported")
+            elif (
+                isinstance(obj, SERIES_TYPE)
+                and df_or_series.axes[axis].index_value.key != obj.index_value.key
+            ):
+                raise NotImplementedError("Aligning different indices not supported")
+
+        _check_input_index(self.cond, axis=0)
+        _check_input_index(self.other)
+
+        if isinstance(df_or_series, DATAFRAME_TYPE):
+            mock_obj = build_df(df_or_series)
+        else:
+            mock_obj = build_series(df_or_series)
+
+        if isinstance(self.other, (pd.DataFrame, DATAFRAME_TYPE)):
+            mock_other = build_df(self.other)
+        elif isinstance(self.other, (pd.Series, SERIES_TYPE)):
+            mock_other = build_series(self.other)
+        else:
+            mock_other = self.other
+
+        result_df = mock_obj.where(
+            np.zeros(mock_obj.shape).astype(bool),
+            other=mock_other,
+            axis=self.axis,
+            level=self.level,
+            errors=self.errors,
+            try_cast=self.try_cast,
+        )
+
+        inputs = filter_inputs([df_or_series, self.cond, self.other])
+        if isinstance(df_or_series, DATAFRAME_TYPE):
+            return self.new_dataframe(
+                inputs,
+                shape=df_or_series.shape,
+                dtypes=result_df.dtypes,
+                index_value=df_or_series.index_value,
+                columns_value=df_or_series.columns_value,
+            )
+        else:
+            return self.new_series(
+                inputs,
+                shape=df_or_series.shape,
+                name=df_or_series.name,
+                dtype=result_df.dtype,
+                index_value=df_or_series.index_value,
+            )
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        inputs_iter = iter(self._inputs)
+        self._input = next(inputs_iter)
+        if isinstance(self.cond, ENTITY_TYPE):
+            self.cond = next(inputs_iter)
+        if isinstance(self.other, ENTITY_TYPE):
+            self.other = next(inputs_iter)
+
+
+_doc_template = """
+Replace values where the condition is {replace_true}.
+
+Parameters
+----------
+cond : bool Series/DataFrame, array-like, or callable
+    Where `cond` is False, keep the original value. Where
+    True, replace with corresponding value from `other`.
+    If `cond` is callable, it is computed on the Series/DataFrame and
+    should return boolean Series/DataFrame or array. The callable must
+    not change input Series/DataFrame (though pandas doesn't check it).
+other : scalar, Series/DataFrame, or callable
+    Entries where `cond` is True are replaced with
+    corresponding value from `other`.
+    If other is callable, it is computed on the Series/DataFrame and
+    should return scalar or Series/DataFrame. The callable must not
+    change input Series/DataFrame (though pandas doesn't check it).
+inplace : bool, default False
+    Whether to perform the operation in place on the data.
+axis : int, default None
+    Alignment axis if needed.
+level : int, default None
+    Alignment level if needed.
+errors : str, {{'raise', 'ignore'}}, default 'raise'
+    Note that currently this parameter won't affect
+    the results and will always coerce to a suitable dtype.
+
+    - 'raise' : allow exceptions to be raised.
+    - 'ignore' : suppress exceptions. On error return original object.
+
+try_cast : bool, default False
+    Try to cast the result back to the input type (if possible).
+
+Returns
+-------
+Same type as caller
+
+See Also
+--------
+:func:`DataFrame.{opposite}` : Return an object of same shape as
+    self.
+
+Notes
+-----
+The mask method is an application of the if-then idiom. For each
+element in the calling DataFrame, if ``cond`` is ``False`` the
+element is used; otherwise the corresponding element from the DataFrame
+``other`` is used.
+
+The signature for :func:`DataFrame.where` differs from
+:func:`numpy.where`. Roughly ``df1.where(m, df2)`` is equivalent to
+``np.where(m, df1, df2)``.
+
+For further details and examples see the ``mask`` documentation in
+:ref:`indexing <indexing.where_mask>`.
+
+Examples
+--------
+>>> import maxframe.tensor as mt
+>>> import maxframe.dataframe as md
+>>> s = md.Series(range(5))
+>>> s.where(s > 0).execute()
+0    NaN
+1    1.0
+2    2.0
+3    3.0
+4    4.0
+dtype: float64
+
+>>> s.mask(s > 0).execute()
+0    0.0
+1    NaN
+2    NaN
+3    NaN
+4    NaN
+dtype: float64
+
+>>> s.where(s > 1, 10).execute()
+0    10
+1    10
+2    2
+3    3
+4    4
+dtype: int64
+
+>>> df = md.DataFrame(mt.arange(10).reshape(-1, 2), columns=['A', 'B'])
+>>> df.execute()
+   A  B
+0  0  1
+1  2  3
+2  4  5
+3  6  7
+4  8  9
+>>> m = df % 3 == 0
+>>> df.where(m, -df).execute()
+   A  B
+0  0 -1
+1 -2  3
+2 -4 -5
+3  6 -7
+4 -8  9
+>>> df.where(m, -df) == mt.where(m, df, -df).execute()
+      A     B
+0  True  True
+1  True  True
+2  True  True
+3  True  True
+4  True  True
+>>> df.where(m, -df) == df.mask(~m, -df).execute()
+      A     B
+0  True  True
+1  True  True
+2  True  True
+3  True  True
+4  True  True
+"""
+
+
+def _where(
+    df_or_series,
+    cond,
+    other=np.nan,
+    inplace=False,
+    axis=None,
+    level=None,
+    errors="raise",
+    try_cast=False,
+    replace_true=False,
+):
+    if df_or_series.ndim == 2 and getattr(other, "ndim", 2) == 1 and axis is None:
+        raise ValueError("Must specify axis=0 or 1")
+
+    axis = validate_axis(axis or 0, df_or_series)
+    op = DataFrameWhere(
+        cond=cond,
+        other=other,
+        axis=axis,
+        level=level,
+        errors=errors,
+        try_cast=try_cast,
+        replace_true=replace_true,
+    )
+    result = op(df_or_series)
+    if inplace:
+        df_or_series.data = result.data
+    else:
+        return result
+
+
+def where(
+    df_or_series,
+    cond,
+    other=np.nan,
+    inplace=False,
+    axis=None,
+    level=None,
+    errors="raise",
+    try_cast=False,
+):
+    return _where(
+        df_or_series,
+        cond,
+        other=other,
+        inplace=inplace,
+        axis=axis,
+        level=level,
+        errors=errors,
+        try_cast=try_cast,
+        replace_true=False,
+    )
+
+
+def mask(
+    df_or_series,
+    cond,
+    other=np.nan,
+    inplace=False,
+    axis=None,
+    level=None,
+    errors="raise",
+    try_cast=False,
+):
+    return _where(
+        df_or_series,
+        cond,
+        other=other,
+        inplace=inplace,
+        axis=axis,
+        level=level,
+        errors=errors,
+        try_cast=try_cast,
+        replace_true=True,
+    )
+
+
+mask.__doc__ = _doc_template.format(replace_true=True, opposite="where")
+where.__doc__ = _doc_template.format(replace_true=False, opposite="mask")

maxframe/dataframe/initializer.py

@@ -0,0 +1,288 @@
+# 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 typing import Union
+
+import pandas as pd
+from pandas.core.dtypes.common import pandas_dtype
+
+from ..core import ENTITY_TYPE
+from ..serialization.serializables import SerializableMeta
+from ..tensor import stack
+from ..tensor import tensor as astensor
+from ..tensor.array_utils import is_cupy
+from ..tensor.core import TENSOR_TYPE
+from ..utils import ceildiv, lazy_import
+from .core import DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
+from .core import DataFrame as _Frame
+from .core import Index as _Index
+from .core import Series as _Series
+from .datasource.dataframe import from_pandas as from_pandas_df
+from .datasource.from_tensor import (
+    dataframe_from_1d_tileables,
+    dataframe_from_tensor,
+    series_from_tensor,
+)
+from .datasource.index import from_pandas as from_pandas_index
+from .datasource.index import from_tileable as from_tileable_index
+from .datasource.series import from_pandas as from_pandas_series
+from .utils import is_cudf, is_index
+
+cudf = lazy_import("cudf")
+
+
+class InitializerMeta(SerializableMeta):
+    def __instancecheck__(cls, instance):
+        return isinstance(instance, (cls.__base__,) + getattr(cls, "_allow_data_type_"))
+
+
+class DataFrame(_Frame, metaclass=InitializerMeta):
+    def __init__(
+        self,
+        data=None,
+        index=None,
+        columns=None,
+        dtype=None,
+        copy=False,
+        chunk_size=None,
+        gpu=None,
+        sparse=None,
+        num_partitions=None,
+    ):
+        need_repart = False
+        if isinstance(data, TENSOR_TYPE):
+            if chunk_size is not None:
+                data = data.rechunk(chunk_size)
+            df = dataframe_from_tensor(
+                data, index=index, columns=columns, gpu=gpu, sparse=sparse
+            )
+            need_repart = num_partitions is not None
+        elif isinstance(data, SERIES_TYPE):
+            df = data.to_frame()
+            need_repart = num_partitions is not None
+        elif isinstance(data, DATAFRAME_TYPE):
+            if not hasattr(data, "data"):
+                # DataFrameData
+                df = _Frame(data)
+            else:
+                df = data
+            need_repart = num_partitions is not None
+        elif isinstance(data, dict) and self._can_process_by_1d_tileables(data):
+            # data is a dict and some value is tensor
+            df = dataframe_from_1d_tileables(
+                data, index=index, columns=columns, gpu=gpu, sparse=sparse
+            )
+            need_repart = num_partitions is not None
+        elif isinstance(data, list) and any(isinstance(v, ENTITY_TYPE) for v in data):
+            # stack data together
+            data = stack(data)
+            df = dataframe_from_tensor(
+                data, index=index, columns=columns, gpu=gpu, sparse=sparse
+            )
+            need_repart = num_partitions is not None
+        elif isinstance(index, (INDEX_TYPE, SERIES_TYPE)):
+            if isinstance(data, dict):
+                data = {k: astensor(v, chunk_size=chunk_size) for k, v in data.items()}
+                df = dataframe_from_1d_tileables(
+                    data, index=index, columns=columns, gpu=gpu, sparse=sparse
+                )
+            else:
+                if data is not None:
+                    data = astensor(data, chunk_size=chunk_size)
+                df = dataframe_from_tensor(
+                    data, index=index, columns=columns, gpu=gpu, sparse=sparse
+                )
+            need_repart = num_partitions is not None
+        else:
+            if is_cudf(data) or is_cupy(data):  # pragma: no cover
+                pdf = cudf.DataFrame(data, index=index, columns=columns, dtype=dtype)
+                if copy:
+                    pdf = pdf.copy()
+            else:
+                pdf = pd.DataFrame(
+                    data, index=index, columns=columns, dtype=dtype, copy=copy
+                )
+            if num_partitions is not None:
+                chunk_size = ceildiv(len(pdf), num_partitions)
+            df = from_pandas_df(pdf, chunk_size=chunk_size, gpu=gpu, sparse=sparse)
+
+        if need_repart:
+            df = df.rebalance(num_partitions=num_partitions)
+        super().__init__(df.data)
+
+    @classmethod
+    def _can_process_by_1d_tileables(cls, data: dict):
+        for value in data.values():
+            if isinstance(value, ENTITY_TYPE):
+                return True
+            elif isinstance(value, (list, tuple)) and any(
+                isinstance(v, ENTITY_TYPE) for v in value
+            ):
+                return True
+        return False
+
+
+class Series(_Series, metaclass=InitializerMeta):
+    def __init__(
+        self,
+        data=None,
+        index=None,
+        dtype=None,
+        name=None,
+        copy=False,
+        chunk_size=None,
+        gpu=None,
+        sparse=None,
+        num_partitions=None,
+    ):
+        if dtype is not None:
+            dtype = pandas_dtype(dtype)
+        need_repart = False
+        if isinstance(data, (TENSOR_TYPE, INDEX_TYPE)):
+            if chunk_size is not None:
+                data = data.rechunk(chunk_size)
+            name = name or getattr(data, "name", None)
+            series = series_from_tensor(
+                data, index=index, name=name, gpu=gpu, sparse=sparse
+            )
+            need_repart = num_partitions is not None
+        elif isinstance(index, INDEX_TYPE):
+            if data is not None:
+                data = astensor(data, chunk_size=chunk_size)
+            series = series_from_tensor(
+                data, index=index, name=name, dtype=dtype, gpu=gpu, sparse=sparse
+            )
+            need_repart = num_partitions is not None
+        elif isinstance(data, SERIES_TYPE):
+            if not hasattr(data, "data"):
+                # SeriesData
+                series = _Series(data)
+            else:
+                series = data
+            need_repart = num_partitions is not None
+        else:
+            if is_cudf(data) or is_cupy(data):  # pragma: no cover
+                pd_series = cudf.Series(data, index=index, dtype=dtype, name=name)
+                if copy:
+                    pd_series = pd_series.copy()
+            else:
+                pd_series = pd.Series(
+                    data, index=index, dtype=dtype, name=name, copy=copy
+                )
+            if num_partitions is not None:
+                chunk_size = ceildiv(len(pd_series), num_partitions)
+            series = from_pandas_series(
+                pd_series, chunk_size=chunk_size, gpu=gpu, sparse=sparse
+            )
+
+        if need_repart:
+            series = series.rebalance(num_partitions=num_partitions)
+        super().__init__(series.data)
+
+
+class Index(_Index, metaclass=InitializerMeta):
+    def __new__(cls, data, **_):
+        # just return cls always until we support other Index's initializers
+        return object.__new__(cls)
+
+    def __init__(
+        self,
+        data=None,
+        dtype=None,
+        copy=False,
+        name=None,
+        tupleize_cols=True,
+        chunk_size=None,
+        gpu=None,
+        sparse=None,
+        names=None,
+        num_partitions=None,
+        store_data=False,
+    ):
+        need_repart = False
+        if isinstance(data, INDEX_TYPE):
+            if not hasattr(data, "data"):
+                # IndexData
+                index = _Index(data)
+            else:
+                index = data
+            need_repart = num_partitions is not None
+        else:
+            if isinstance(data, ENTITY_TYPE):
+                name = name if name is not None else getattr(data, "name", None)
+                index = from_tileable_index(data, dtype=dtype, name=name, names=names)
+                need_repart = num_partitions is not None
+            else:
+                if not is_index(data):
+                    name = name if name is not None else getattr(data, "name", None)
+                    xdf = cudf if is_cudf(data) or is_cupy(data) else pd
+                    try:
+                        pd_index = xdf.Index(
+                            data=data,
+                            dtype=dtype,
+                            copy=copy,
+                            name=name,
+                            tupleize_cols=tupleize_cols,
+                        )
+                    except TypeError:  # pragma: no cover
+                        pd_index = xdf.Index(
+                            data=data, dtype=dtype, copy=copy, name=name
+                        )
+                else:
+                    pd_index = data
+
+                if num_partitions is not None:
+                    chunk_size = ceildiv(len(pd_index), num_partitions)
+                index = from_pandas_index(
+                    pd_index,
+                    chunk_size=chunk_size,
+                    gpu=gpu,
+                    sparse=sparse,
+                    store_data=store_data,
+                )
+
+        if need_repart:
+            index = index.rebalance(num_partitions=num_partitions)
+        super().__init__(index.data)
+
+
+_pd_type_mapping = {
+    pd.DataFrame: DataFrame,
+    pd.Series: Series,
+    pd.Index: Index,
+}
+
+
+def read_pandas(
+    data: Union[pd.DataFrame, pd.Series, pd.Index], **kwargs
+) -> Union[DataFrame, Series, Index]:
+    """
+    Create MaxFrame objects from pandas.
+
+    Parameters
+    ----------
+    data: Union[pd.DataFrame, pd.Series, pd.Index]
+        pandas data
+    kwargs: dict
+        arguments to be passed to initializers.
+
+    Returns
+    -------
+    result: Union[DataFrame, Series, Index]
+        result MaxFrame object
+    """
+    for pd_cls, cls in _pd_type_mapping.items():
+        if isinstance(data, pd_cls):
+            return cls(data, **kwargs)
+    raise ValueError(f"Type {type(data)} not supported")

maxframe/dataframe/merge/__init__.py

@@ -0,0 +1,32 @@
+# 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 .append import DataFrameAppend, append
+from .concat import DataFrameConcat, concat
+from .merge import DataFrameMerge, DataFrameMergeAlign, join, merge
+
+
+def _install():
+    from ..core import DATAFRAME_TYPE, SERIES_TYPE
+
+    for cls in DATAFRAME_TYPE:
+        setattr(cls, "join", join)
+        setattr(cls, "merge", merge)
+
+    for cls in DATAFRAME_TYPE + SERIES_TYPE:
+        setattr(cls, "append", append)
+
+
+_install()
+del _install