maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/dataframe/indexing/where.py

@@ -0,0 +1,300 @@
+# 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 typing import List
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import ENTITY_TYPE, EntityData
+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,
+        )
+
+        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,
+            )
+
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameWhere", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        inputs_iter = iter(op._inputs)
+        op._input = next(inputs_iter)
+        if isinstance(op.cond, ENTITY_TYPE):
+            op.cond = next(inputs_iter)
+        if isinstance(op.other, ENTITY_TYPE):
+            op.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.
+
+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/indexing/xs.py

@@ -0,0 +1,148 @@
+# 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 pandas.api.types import is_list_like
+
+from ..utils import validate_axis
+
+
+def xs(df_or_series, key, axis=0, level=None, drop_level=True):
+    """
+    Return cross-section from the Series/DataFrame.
+
+    This method takes a `key` argument to select data at a particular
+    level of a MultiIndex.
+
+    Parameters
+    ----------
+    key : label or tuple of label
+        Label contained in the index, or partially in a MultiIndex.
+    axis : {0 or 'index', 1 or 'columns'}, default 0
+        Axis to retrieve cross-section on.
+    level : object, defaults to first n levels (n=1 or len(key))
+        In case of a key partially contained in a MultiIndex, indicate
+        which levels are used. Levels can be referred by label or position.
+    drop_level : bool, default True
+        If False, returns object with same levels as self.
+
+    Returns
+    -------
+    Series or DataFrame
+        Cross-section from the original Series or DataFrame
+        corresponding to the selected index levels.
+
+    See Also
+    --------
+    DataFrame.loc : Access a group of rows and columns
+        by label(s) or a boolean array.
+    DataFrame.iloc : Purely integer-location based indexing
+        for selection by position.
+
+    Notes
+    -----
+    `xs` can not be used to set values.
+
+    MultiIndex Slicers is a generic way to get/set values on
+    any level or levels.
+    It is a superset of `xs` functionality, see
+    :ref:`MultiIndex Slicers <advanced.mi_slicers>`.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> d = {'num_legs': [4, 4, 2, 2],
+    ...      'num_wings': [0, 0, 2, 2],
+    ...      'class': ['mammal', 'mammal', 'mammal', 'bird'],
+    ...      'animal': ['cat', 'dog', 'bat', 'penguin'],
+    ...      'locomotion': ['walks', 'walks', 'flies', 'walks']}
+    >>> df = md.DataFrame(data=d)
+    >>> df = df.set_index(['class', 'animal', 'locomotion'])
+    >>> df.execute()
+                                num_legs  num_wings
+    class  animal  locomotion
+    mammal cat     walks              4          0
+            dog     walks              4          0
+            bat     flies              2          2
+    bird   penguin walks              2          2
+
+    Get values at specified index
+
+    >>> df.xs('mammal').execute()
+                        num_legs  num_wings
+    animal locomotion
+    cat    walks              4          0
+    dog    walks              4          0
+    bat    flies              2          2
+
+    Get values at several indexes
+
+    >>> df.xs(('mammal', 'dog')).execute()
+                num_legs  num_wings
+    locomotion
+    walks              4          0
+
+    Get values at specified index and level
+
+    >>> df.xs('cat', level=1).execute()
+                        num_legs  num_wings
+    class  locomotion
+    mammal walks              4          0
+
+    Get values at several indexes and levels
+
+    >>> df.xs(('bird', 'walks'),
+    ...       level=[0, 'locomotion']).execute()
+                num_legs  num_wings
+    animal
+    penguin         2          2
+
+    Get values at specified column and axis
+
+    >>> df.xs('num_wings', axis=1).execute()
+    class   animal   locomotion
+    mammal  cat      walks         0
+            dog      walks         0
+            bat      flies         2
+    bird    penguin  walks         2
+    Name: num_wings, dtype: int64
+    """
+    axis = validate_axis(axis, df_or_series)
+    if level is None:
+        level = range(df_or_series.axes[axis].nlevels)
+    elif not is_list_like(level):
+        level = [level]
+
+    slc = [slice(None)] * df_or_series.axes[axis].nlevels
+    if not is_list_like(key):
+        key = (key,)
+
+    level_set = set()
+    for k, level_ in zip(key, level):
+        slc[level_] = k
+        level_set.add(level_)
+    left_levels = set(range(df_or_series.axes[axis].nlevels)) - level_set
+
+    if len(slc) > 1:
+        slc = tuple(slc)
+
+    res = df_or_series.loc(axis=axis)[slc]
+    if drop_level:
+        if len(left_levels) == 0:
+            if res.ndim > 1:
+                res = res.iloc[0, :] if axis == 0 else res.iloc[:, 0]
+            else:
+                res = res.iloc[0]
+        else:
+            res = res.droplevel(list(level_set), axis=axis)
+    return res