maxframe 1.3.1__cp311-cp311-macosx_10_9_universal2.whl → 2.0.0b1__cp311-cp311-macosx_10_9_universal2.whl

maxframe/dataframe/missing/replace.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 from ... import opcodes
+from ...core import EntityData
 from ...serialization.serializables import AnyField, Int32Field
 from ...utils import no_default
 from ..operators import SERIES_TYPE, DataFrameOperator, DataFrameOperatorMixin
@@ -28,14 +31,15 @@ class DataFrameReplace(DataFrameOperator, DataFrameOperatorMixin):
     regex = AnyField("regex", default=None)
     method = AnyField("method", default=no_default)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameReplace", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
         input_iter = iter(inputs)
         next(input_iter)
-        if isinstance(self.to_replace, SERIES_TYPE):
-            self.to_replace = next(input_iter)
-        if isinstance(self.value, SERIES_TYPE):
-            self.value = next(input_iter)
+        if isinstance(op.to_replace, SERIES_TYPE):
+            op.to_replace = next(input_iter)
+        if isinstance(op.value, SERIES_TYPE):
+            op.value = next(input_iter)
 
     def __call__(self, df_or_series):
         inputs = [df_or_series]

maxframe/dataframe/missing/tests/test_missing.py

@@ -64,7 +64,7 @@ def test_drop_na():
     df_raw = pd.DataFrame(np.nan, index=range(0, 20), columns=list("ABCDEFGHIJ"))
     for _ in range(30):
         df_raw.iloc[random.randint(0, 19), random.randint(0, 9)] = random.randint(0, 99)
-    for rowid in range(random.randint(1, 5)):
+    for _ in range(random.randint(1, 5)):
         row = random.randint(0, 19)
         for idx in range(0, 10):
             df_raw.iloc[row, idx] = random.randint(0, 99)
@@ -79,7 +79,7 @@ def test_replace():
     df_raw = pd.DataFrame(-1, index=range(0, 20), columns=list("ABCDEFGHIJ"))
     for _ in range(30):
         df_raw.iloc[random.randint(0, 19), random.randint(0, 9)] = random.randint(0, 99)
-    for rowid in range(random.randint(1, 5)):
+    for _ in range(random.randint(1, 5)):
         row = random.randint(0, 19)
         for idx in range(0, 10):
             df_raw.iloc[row, idx] = random.randint(0, 99)

maxframe/dataframe/operators.py

@@ -15,7 +15,7 @@
 import numpy as np
 import pandas as pd
 
-from ..core import ENTITY_TYPE, OutputType
+from ..core import OutputType
 from ..core.operator import Operator, ShuffleProxy, TileableOperatorMixin
 from ..tensor.core import TENSOR_TYPE
 from ..tensor.datasource import tensor as astensor
@@ -208,32 +208,6 @@ class DataFrameOperatorMixin(TileableOperatorMixin):
             inputs, shape=shape, dtype=dtype, categories_value=categories_value, **kw
         )[0]
 
-    @classmethod
-    def _process_groupby_params(cls, groupby_params):
-        new_groupby_params = groupby_params.copy()
-        if isinstance(groupby_params["by"], list):
-            by = []
-            for v in groupby_params["by"]:
-                if isinstance(v, ENTITY_TYPE):
-                    by.append(cls.concat_tileable_chunks(v).chunks[0])
-                else:
-                    by.append(v)
-            new_groupby_params["by"] = by
-        return new_groupby_params
-
-    @classmethod
-    def _get_groupby_inputs(cls, groupby, groupby_params):
-        inputs = [groupby]
-        chunk_inputs = list(groupby.chunks)
-        if isinstance(groupby_params["by"], list):
-            for chunk_v, v in zip(
-                groupby_params["by"], groupby.op.groupby_params["by"]
-            ):
-                if isinstance(v, ENTITY_TYPE):
-                    inputs.append(v)
-                    chunk_inputs.append(chunk_v)
-        return inputs, chunk_inputs
-
     @staticmethod
     def _process_input(x):
         from .initializer import DataFrame, Series

maxframe/dataframe/reduction/aggregation.py

@@ -14,6 +14,7 @@
 
 import copy
 import functools
+import inspect
 import itertools
 from collections import OrderedDict
 from collections.abc import Iterable
@@ -21,10 +22,13 @@ from typing import List
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 
 from ... import opcodes
 from ... import tensor as maxframe_tensor
 from ...core import ENTITY_TYPE, OutputType, enter_mode
+from ...io.odpsio.schema import pandas_dtype_to_arrow_type
+from ...lib.dtypes_extension import ArrowDtype
 from ...serialization.serializables import AnyField, BoolField, DictField, ListField
 from ...typing_ import TileableType
 from ...utils import lazy_import, pd_release_version
@@ -37,6 +41,7 @@ from .core import (
     ReductionPostStep,
     ReductionPreStep,
 )
+from .unique import _unique
 
 cp = lazy_import("cupy", rename="cp")
 cudf = lazy_import("cudf")
@@ -71,6 +76,7 @@ _agg_functions = {
     "kurt": lambda x, skipna=True, bias=False: x.kurt(skipna=skipna, bias=bias),
     "kurtosis": lambda x, skipna=True, bias=False: x.kurtosis(skipna=skipna, bias=bias),
     "nunique": lambda x: x.nunique(),
+    "unique": lambda x: _unique(x, output_list_scalar=True),
     "median": lambda x, skipna=True: x.median(skipna=skipna),
 }
 
@@ -99,6 +105,46 @@ class DataFrameAggregate(DataFrameOperator, DataFrameOperatorMixin):
             [np.number, np.bool_] if op.numeric_only else [np.bool_]
         ).dtypes
 
+    def _fill_df_dtypes(self, in_df, dtypes):
+        if all(dt != np.dtype("O") for dt in dtypes):
+            return dtypes
+
+        if isinstance(self.func, dict):
+            col_func_it = self.func.items()
+        else:
+            assert in_df.ndim == 2
+            col_func_it = itertools.product(in_df.dtypes.index, self.func)
+
+        col_to_dt = dict(in_df.dtypes.items())
+
+        new_dt = OrderedDict()
+        for (col_name, func), (out_col_name, dt) in zip(col_func_it, dtypes.items()):
+            if dt != np.dtype("O"):
+                new_dt[out_col_name] = dt
+            elif func == "unique":
+                in_dt = col_to_dt[col_name]
+                if in_dt == np.dtype("O"):
+                    in_dt = pd.StringDtype()
+                arrow_dt = pandas_dtype_to_arrow_type(in_dt)
+                new_dt[out_col_name] = ArrowDtype(pa.list_(arrow_dt))
+            else:
+                # do nothing as the result might be string
+                new_dt[out_col_name] = dt
+        return pd.Series(list(new_dt.values()), index=new_dt.keys())
+
+    def _fill_series_dtype(self, in_data, dtype):
+        if len(self.func) != 1 or dtype != np.dtype("O") or in_data.ndim > 1:
+            return dtype
+
+        if self.func[0] == "unique":
+            in_dt = in_data.dtype
+            if in_dt == np.dtype("O"):
+                in_dt = pd.StringDtype()
+            arrow_dt = pandas_dtype_to_arrow_type(in_dt)
+            return ArrowDtype(pa.list_(arrow_dt))
+        else:
+            return dtype
+
     def _calc_result_shape(self, df):
         if df.ndim == 2:
             if self.numeric_only:
@@ -114,16 +160,23 @@ class DataFrameAggregate(DataFrameOperator, DataFrameOperatorMixin):
             )
 
         result_df = test_obj.agg(self.raw_func, axis=self.axis, **self.raw_func_kw)
+        if isinstance(result_df, pd.DataFrame):
+            out_dtypes = self._fill_df_dtypes(df, result_df.dtypes)
+        elif isinstance(result_df, pd.Series):
+            dtype = self._fill_series_dtype(df, result_df.dtype)
+            out_dtypes = pd.Series([dtype], index=[result_df.name])
+        else:
+            out_dtypes = pd.Series([np.array(result_df).dtype], index=[None])
 
         if isinstance(result_df, pd.DataFrame):
             self.output_types = [OutputType.dataframe]
-            return result_df.dtypes, result_df.index
+            return out_dtypes, result_df.index
         elif isinstance(result_df, pd.Series):
             self.output_types = [OutputType.series]
-            return pd.Series([result_df.dtype], index=[result_df.name]), result_df.index
+            return out_dtypes, result_df.index
         else:
             self.output_types = [OutputType.scalar]
-            return np.array(result_df).dtype, None
+            return out_dtypes.iloc[0], None
 
     def __call__(self, df, output_type=None, dtypes=None, index=None):
         self._output_types = df.op.output_types
@@ -379,6 +432,15 @@ def aggregate(df, func=None, axis=0, **kw):
     min    1
     """
     axis = validate_axis(axis, df)
+    if func == "unique":
+        # workaround for direct call of unique function which
+        #  returns a tensor directly
+        func = getattr(df, func)
+        if "axis" in inspect.getfullargspec(func).args:
+            kw = kw.copy()
+            kw["axis"] = axis
+        return func(**kw)
+
     if (
         df.ndim == 2
         and isinstance(func, dict)

maxframe/dataframe/reduction/core.py

@@ -226,7 +226,9 @@ class DataFrameReductionMixin(DataFrameOperatorMixin):
                     #  handle pandas Dtypes in the future more carefully.
                     reduced_dtype = np.dtype("O")
                 else:
-                    reduced_dtype = np.find_common_type(dtypes, [])
+                    from pandas.core.dtypes.cast import find_common_type
+
+                    reduced_dtype = find_common_type(dtypes)
 
         if level is not None:
             return self._call_groupby_level(df[reduced_cols], level)

maxframe/dataframe/reduction/median.py

@@ -31,7 +31,7 @@ def median_series(df, axis=None, skipna=True, level=None, method=None):
         axis=axis,
         skipna=skipna,
         level=level,
-        output_types=[OutputType.scalar if level is not None else OutputType.scalar],
+        output_types=[OutputType.series if level is not None else OutputType.scalar],
         method=method,
     )
     return op(df)

maxframe/dataframe/reduction/tests/test_reduction.py

@@ -19,9 +19,11 @@ from typing import NamedTuple
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import pytest
 
 from .... import dataframe as md
+from ....lib.dtypes_extension import ArrowDtype
 from ....tensor import Tensor
 from ....tests.utils import assert_mf_index_dtype
 from ...core import DataFrame, IndexValue, OutputType, Series
@@ -29,6 +31,7 @@ from ...datasource.dataframe import from_pandas as from_pandas_df
 from ...datasource.series import from_pandas as from_pandas_series
 from .. import (
     CustomReduction,
+    DataFrameAggregate,
     DataFrameAll,
     DataFrameAny,
     DataFrameCount,
@@ -196,6 +199,36 @@ def test_nunique():
     assert isinstance(result2.op, DataFrameNunique)
 
 
+def test_unique():
+    pd_df = pd.DataFrame(
+        {
+            "col1": pd.Series(np.random.choice(["a", "b", "c", "d"], 100)),
+            "col2": pd.Series(np.random.choice([0, 1, 2, 3], 100)),
+        }
+    )
+    df = from_pandas_df(pd_df, chunk_size=3)
+    result = df.agg(["unique"])
+
+    assert result.shape == (1, 2)
+    assert result.op.output_types[0] == OutputType.dataframe
+    assert isinstance(result.op, DataFrameAggregate)
+    pd.testing.assert_series_equal(
+        result.dtypes,
+        pd.Series(
+            [ArrowDtype(pa.list_(pa.string())), ArrowDtype(pa.list_(pa.int64()))],
+            index=pd_df.columns,
+        ),
+    )
+
+    pd_s = pd.Series(np.random.choice(["a", "b", "c", "d"], 100))
+    ms = from_pandas_series(pd_s, chunk_size=3)
+    result = ms.agg(["unique"])
+    assert result.shape == (1,)
+    assert result.op.output_types[0] == OutputType.series
+    assert isinstance(result.op, DataFrameAggregate)
+    assert result.dtype == ArrowDtype(pa.list_(pa.string()))
+
+
 def test_dataframe_aggregate():
     data = pd.DataFrame(np.random.rand(20, 19))
     agg_funcs = [

maxframe/dataframe/reduction/unique.py

@@ -12,14 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 
 from ... import opcodes
 from ...core import ENTITY_TYPE, OutputType
+from ...io.odpsio.schema import (
+    pandas_dtype_to_arrow_type,
+    pandas_dtypes_to_arrow_schema,
+)
+from ...lib.dtypes_extension import ArrowDtype
+from ...serialization.serializables import BoolField
 from ...tensor.core import TensorOrder
 from ...utils import lazy_import
+from ..core import DATAFRAME_TYPE
 from ..initializer import Series as asseries
 from .core import CustomReduction, DataFrameReductionMixin, DataFrameReductionOperator
 
@@ -40,6 +47,12 @@ class DataFrameUnique(DataFrameReductionOperator, DataFrameReductionMixin):
     _op_type_ = opcodes.UNIQUE
     _func_name = "unique"
 
+    output_list_scalar = BoolField("output_list_scalar", default=False)
+
+    @property
+    def is_atomic(self):
+        return True
+
     @classmethod
     def get_reduction_callable(cls, op):
         return UniqueReduction(name=cls._func_name, is_gpu=op.is_gpu())
@@ -47,10 +60,43 @@ class DataFrameUnique(DataFrameReductionOperator, DataFrameReductionMixin):
     def __call__(self, a):
         if not isinstance(a, ENTITY_TYPE):
             a = asseries(a)
-        self.output_types = [OutputType.tensor]
-        return self.new_tileables(
-            [a], shape=(np.nan,), dtype=a.dtype, order=TensorOrder.C_ORDER
-        )[0]
+        self.axis = 0
+        if isinstance(a, DATAFRAME_TYPE):
+            assert self.output_list_scalar and self.axis == 0
+            pa_schema = pandas_dtypes_to_arrow_schema(a.dtypes, unknown_as_string=True)
+            if len(set(pa_schema.types)) == 1:
+                out_dtype = ArrowDtype(pa.list_(pa_schema.types[0]))
+            else:
+                out_dtype = np.dtype("O")
+            kw = {
+                "dtype": out_dtype,
+                "index_value": a.columns_value,
+                "shape": (a.shape[1],),
+            }
+            self.output_types = [OutputType.series]
+            return self.new_tileables([a], **kw)[0]
+        else:
+            if self.output_list_scalar:
+                arrow_type = pa.list_(
+                    pandas_dtype_to_arrow_type(a.dtype, unknown_as_string=True)
+                )
+                kw = {
+                    "dtype": ArrowDtype(arrow_type),
+                    "shape": (),
+                }
+                self.output_types = [OutputType.scalar]
+            else:
+                kw = {
+                    "dtype": a.dtype,
+                    "shape": (np.nan,),
+                }
+                self.output_types = [OutputType.tensor]
+            return self.new_tileables([a], order=TensorOrder.C_ORDER, **kw)[0]
+
+
+def _unique(values, method="tree", **kwargs):
+    op = DataFrameUnique(method=method, **kwargs)
+    return op(values)
 
 
 def unique(values, method="tree"):
@@ -62,6 +108,7 @@ def unique(values, method="tree"):
     values : 1d array-like
     method : 'shuffle' or 'tree', 'tree' method provide a better performance, 'shuffle'
     is recommended if the number of unique values is very large.
+
     See Also
     --------
     Index.unique
@@ -86,5 +133,4 @@ def unique(values, method="tree"):
     array([Timestamp('2016-01-01 00:00:00-0500', tz='US/Eastern')],
           dtype=object)
     """
-    op = DataFrameUnique(method=method)
-    return op(values)
+    return _unique(values, method=method)

maxframe/dataframe/statistics/corr.py

@@ -12,11 +12,13 @@
 # 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
+from ...core import ENTITY_TYPE, EntityData
 from ...serialization.serializables import AnyField, BoolField, Int32Field, KeyField
 from ...tensor.utils import filter_inputs
 from ..core import DATAFRAME_TYPE, SERIES_TYPE
@@ -33,12 +35,13 @@ class DataFrameCorr(DataFrameOperator, DataFrameOperatorMixin):
     axis = Int32Field("axis", default=None)
     drop = BoolField("drop", default=None)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        inputs_iter = iter(self._inputs)
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameCorr", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        inputs_iter = iter(op._inputs)
         next(inputs_iter)
-        if isinstance(self.other, ENTITY_TYPE):
-            self.other = next(inputs_iter)
+        if isinstance(op.other, ENTITY_TYPE):
+            op.other = next(inputs_iter)
 
     def __call__(self, df_or_series):
         if isinstance(df_or_series, SERIES_TYPE):

maxframe/dataframe/statistics/quantile.py

@@ -12,12 +12,14 @@
 # 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 pandas.core.dtypes.cast import find_common_type
 
 from ... import opcodes
-from ...core import ENTITY_TYPE
+from ...core import ENTITY_TYPE, EntityData
 from ...serialization.serializables import (
     AnyField,
     BoolField,
@@ -50,11 +52,12 @@ class DataFrameQuantile(DataFrameOperator, DataFrameOperatorMixin):
     def __init__(self, output_types=None, **kw):
         super().__init__(_output_types=output_types, **kw)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.input = self._inputs[0]
-        if isinstance(self.q, TENSOR_TYPE):
-            self.q = self._inputs[-1]
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameQuantile", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.input = op._inputs[0]
+        if isinstance(op.q, TENSOR_TYPE):
+            op.q = op._inputs[-1]
 
     def _calc_dtype_on_axis_1(self, a, dtypes):
         quantile_dtypes = []

maxframe/dataframe/tseries/to_datetime.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+from typing import Any, List
 
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_dict_like, is_scalar
 
 from ... import opcodes
+from ...core import EntityData
 from ...serialization.serializables import AnyField, BoolField, KeyField, StringField
 from ...tensor import tensor as astensor
 from ..core import DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
@@ -54,9 +55,10 @@ class DataFrameToDatetime(DataFrameOperator, DataFrameOperatorMixin):
             if k not in self._no_copy_attrs_ and k != "arg" and hasattr(self, k)
         )
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.arg = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameToDatetime", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.arg = op._inputs[0]
 
     def __call__(self, arg):
         if is_scalar(arg):