maxframe 1.0.0rc3__cp310-cp310-win_amd64.whl → 1.0.0rc4__cp310-cp310-win_amd64.whl

maxframe/dataframe/extensions/flatmap.py

@@ -0,0 +1,326 @@
+# 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 Callable
+
+import numpy as np
+import pandas as pd
+
+from maxframe import opcodes
+from maxframe.core import OutputType
+from maxframe.dataframe.core import DataFrame, IndexValue
+from maxframe.dataframe.operators import DataFrameOperator, DataFrameOperatorMixin
+from maxframe.dataframe.utils import make_dtypes, parse_index
+from maxframe.serialization.serializables import (
+    BoolField,
+    DictField,
+    FunctionField,
+    TupleField,
+)
+
+
+class DataFrameFlatMapOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.FLATMAP
+
+    func = FunctionField("func")
+    raw = BoolField("raw", default=False)
+    args = TupleField("args", default=())
+    kwargs = DictField("kwargs", default={})
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    @staticmethod
+    def _gen_flattening_index_value(index_value, *args) -> IndexValue:
+        pd_index = index_value.to_pandas()
+        if not isinstance(pd_index, pd.MultiIndex):
+            # for func return multi rows, will copy indexes
+            return parse_index(pd.Index([], dtype=pd_index.dtype), *args)
+        # multi index will keep the same level and types
+        return parse_index(
+            pd.MultiIndex.from_arrays([c[:0] for c in pd_index.levels]), *args
+        )
+
+    def _call_dataframe(self, df: DataFrame, dtypes: pd.Series):
+        dtypes = make_dtypes(dtypes)
+        index_value = self._gen_flattening_index_value(
+            df.index_value,
+            (df.key, df.index_value.key, self.func),
+        )
+        return self.new_dataframe(
+            [df],
+            shape=(np.nan, len(dtypes)),
+            index_value=index_value,
+            columns_value=parse_index(dtypes.index, store_data=True),
+            dtypes=dtypes,
+        )
+
+    def _call_series_or_index(self, series, dtypes=None):
+        index_value = self._gen_flattening_index_value(
+            series.index_value,
+            (series.key, series.index_value.key, self.func),
+        )
+
+        if self.output_types[0] == OutputType.series:
+            name, dtype = dtypes
+            return self.new_series(
+                [series],
+                dtype=dtype,
+                shape=(np.nan,),
+                index_value=index_value,
+                name=name,
+            )
+
+        dtypes = make_dtypes(dtypes)
+        columns_value = parse_index(dtypes.index, store_data=True)
+        return self.new_dataframe(
+            [series],
+            shape=(np.nan, len(dtypes)),
+            index_value=index_value,
+            columns_value=columns_value,
+            dtypes=dtypes,
+        )
+
+    def __call__(
+        self,
+        df_or_series,
+        dtypes=None,
+        output_type=None,
+    ):
+        if df_or_series.op.output_types[0] == OutputType.dataframe:
+            return self._call_dataframe(df_or_series, dtypes=dtypes)
+        else:
+            return self._call_series_or_index(df_or_series, dtypes=dtypes)
+
+
+def df_flatmap(dataframe, func: Callable, dtypes=None, raw=False, args=(), **kwargs):
+    """
+    Apply the given function to each row and then flatten results. Use this method if your transformation returns
+    multiple rows for each input row.
+
+    This function applies a transformation to each row of the DataFrame, where the transformation can return zero
+    or multiple values, effectively flattening Python generators, list-like collections, and DataFrames.
+
+    Parameters
+    ----------
+    dataframe : DataFrame
+        The DataFrame to which the function will be applied.
+
+    func : Callable
+        Function to apply to each row of the DataFrame. It should accept a Series (or an array if `raw=True`)
+        representing a row and return a list or iterable of values.
+
+    dtypes : Series, dict or list
+        Specify dtypes of returned DataFrame.
+
+    raw : bool, default False
+        Determines if the row is passed as a Series or as a numpy array:
+
+        * ``False`` : passes each row as a Series to the function.
+        * ``True`` : the passed function will receive numpy array objects instead.
+
+    args : tuple
+        Positional arguments to pass to `func`.
+
+    **kwargs
+        Additional keyword arguments to pass as keywords arguments to `func`.
+
+    Returns
+    -------
+    DataFrame
+        Return DataFrame with specified `dtypes`.
+
+    Notes
+    -----
+    The `func` must return an iterable of values for each input row. The index of the resulting DataFrame will be
+     repeated based on the number of output rows generated by `func`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+    >>> df.execute()
+       A  B
+    0  1  4
+    1  2  5
+    2  3  6
+
+    Define a function that takes a number and returns a list of two numbers:
+
+    >>> def generate_values_array(row):
+    ...     return [row['A'] * 2, row['B'] * 3]
+
+    Define a function that takes a row and return two rows and two columns:
+
+    >>> def generate_values_in_generator(row):
+    ...     yield [row[0] * 2, row[1] * 4]
+    ...     yield [row[0] * 3, row[1] * 5]
+
+    Which equals to the following function return a dataframe:
+
+    >>> def generate_values_in_dataframe(row):
+    ...     return pd.DataFrame([[row[0] * 2, row[1] * 4], [row[0] * 3, row[1] * 5]])
+
+    Specify `dtypes` with a function which returns a DataFrame:
+
+    >>> df.mf.flatmap(generate_values_array, dtypes=pd.Series({'A': 'int'})).execute()
+            A
+        0   2
+        0  12
+        1   4
+        1  15
+        2   6
+        2  18
+
+    Specify raw=True to pass input row as array:
+
+    >>> df.mf.flatmap(generate_values_in_generator, dtypes={"A": "int", "B": "int"}, raw=True).execute()
+           A   B
+        0  2  16
+        0  3  20
+        1  4  20
+        1  6  25
+        2  6  24
+        2  9  30
+    """
+    if dtypes is None or len(dtypes) == 0:
+        raise TypeError(
+            "Cannot determine {dtypes} by calculating with enumerate data, "
+            "please specify it as arguments"
+        )
+
+    if not isinstance(func, Callable):
+        raise TypeError("function must be a callable object")
+
+    output_types = [OutputType.dataframe]
+    op = DataFrameFlatMapOperator(
+        func=func, raw=raw, output_types=output_types, args=args, kwargs=kwargs
+    )
+    return op(
+        dataframe,
+        dtypes=dtypes,
+    )
+
+
+def series_flatmap(
+    series, func: Callable, dtypes=None, dtype=None, name=None, args=(), **kwargs
+):
+    """
+    Apply the given function to each row and then flatten results. Use this method if your transformation returns
+    multiple rows for each input row.
+
+    This function applies a transformation to each element of the Series, where the transformation can return zero
+     or multiple values, effectively flattening Python generator, list-liked collections and DataFrame.
+
+    Parameters
+    ----------
+    series : Series
+        The series to which the function will be applied.
+
+    func : Callable
+        Function to apply to each element of the Series. It should accept a scalar value
+        (or an array if `raw=True`) and return a list or iterable of values.
+
+    dtypes : Series, default None
+        Specify dtypes of returned DataFrame. Can't work with dtype.
+
+    dtype : numpy.dtype, default None
+        Specify dtype of returned Series. Can't work with dtypes.
+
+    name : str, default None
+        Specify name of the returned Series.
+
+    args : tuple
+        Positional arguments to pass to `func`.
+
+    **kwargs
+        Additional keyword arguments to pass as keywords arguments to `func`.
+
+    Returns
+    -------
+    DataFrame or Series
+        Result of DataFrame when dtypes specified, else Series.
+
+    Notes
+    -----
+    The `func` must return an iterable of values for each input element. If `dtypes` is specified,
+    `flatmap` will return a DataFrame, if `dtype` and `name` is specified, a Series will be returned. The index of
+     the resulting DataFrame/Series will be repeated based on the number of output rows generated by `func`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+    >>> df.execute()
+       A  B
+    0  1  4
+    1  2  5
+    2  3  6
+
+    Define a function that takes a number and returns a list of two numbers:
+
+    >>> def generate_values_array(x):
+    ...     return [x * 2, x * 3]
+
+    >>> def generate_values_in_generator(x):
+    ...     yield pd.Series([x * 2, x * 4])
+    ...     yield pd.Series([x * 3, x * 5])
+
+    Specify `dtype` with a function which returns list to return more than one elements as a Series:
+
+    >>> df['A'].mf.flatmap(generate_values_array, dtype="int", name="C").execute()
+        0    2
+        0    3
+        1    4
+        1    6
+        2    6
+        2    9
+        Name: C, dtype: int64
+
+    Specify `dtypes` to return multi columns as a DataFrame:
+
+    >>> df['A'].mf.flatmap(generate_values_in_generator, dtypes={"A": "int", "B": "int"}).execute()
+           A   B
+        0  2   4
+        0  3   5
+        1  4   8
+        1  6  10
+        2  6  12
+        2  9  15
+    """
+
+    if dtypes and dtype:
+        raise ValueError("Both dtypes and dtype cannot be specified at the same time.")
+
+    dtypes = (name, dtype) if dtype is not None else dtypes
+    if dtypes is None:
+        raise TypeError(
+            "Cannot determine {dtypes} or {dtype} by calculating with enumerate data, "
+            "please specify it as arguments"
+        )
+
+    if not isinstance(func, Callable):
+        raise TypeError("function must be a callable object")
+
+    output_type = OutputType.series if dtype is not None else OutputType.dataframe
+
+    op = DataFrameFlatMapOperator(
+        func=func, raw=False, output_types=[output_type], args=args, kwargs=kwargs
+    )
+    return op(
+        series,
+        dtypes=dtypes,
+    )

maxframe/dataframe/extensions/tests/test_extensions.py

@@ -11,11 +11,12 @@
 # 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
 import pytest
 
 from .... import dataframe as md
+from ... import DataFrame
 from ...core import IndexValue
 from ..reshuffle import DataFrameReshuffle
 
@@ -36,3 +37,63 @@ def test_reshuffle():
     r = mdf.mf.reshuffle(ignore_index=True)
     assert isinstance(r.op, DataFrameReshuffle)
     assert isinstance(r.index_value.value, IndexValue.RangeIndex)
+
+
+@pytest.fixture
+def df1():
+    return DataFrame({"a": [1, 2, 3], "b": [1, 2, 3], "c": [1, 2, 3]})
+
+
+@pytest.fixture
+def df2():
+    return DataFrame([[1, 2, 3], [1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+
+
+@pytest.fixture
+def df3():
+    return DataFrame(
+        [[1, 2, 3], [1, 2, 3], [1, 2, 3]],
+        columns=["a", "b", "c"],
+        index=pd.MultiIndex.from_arrays([[1, 2, 3], [1, 2, 3]], names=["A", "B"]),
+    )
+
+
+def test_flatmap(df1, df2, df3):
+    def f(x, keys):
+        if x["a"] in keys:
+            yield [1, 0]
+            yield [0, 1]
+
+    apply_df = df1[["a"]].mf.flatmap(
+        f,
+        dtypes={"a": "int64", "b": "int64"},
+    )
+    assert apply_df.shape == (np.nan, 2)
+    assert df1.index_value.key != apply_df.index_value.key
+    assert isinstance(df1.index_value.to_pandas(), pd.RangeIndex)
+    assert not isinstance(apply_df.index_value.to_pandas(), pd.RangeIndex)
+    apply_df = df2[["a"]].mf.flatmap(
+        f,
+        dtypes=pd.Series(["int64", "int64"]),
+    )
+    assert apply_df.shape == (np.nan, 2)
+    assert df2.index_value.key != apply_df.index_value.key
+    with pytest.raises(TypeError):
+        apply_s = df3["a"].mf.flatmap(
+            f,
+        )
+    apply_s = df3["a"].mf.flatmap(
+        f,
+        dtype="int64",
+    )
+    assert apply_s.shape == (np.nan,)
+    assert df3.index_value.key != apply_s.index_value.key
+    assert df3.key != apply_s.index_value.key
+    apply_s = df3["a"].mf.flatmap(
+        f,
+        output_type="dataframe",
+        dtypes=["int64", "int64"],
+    )
+    assert apply_s.shape == (np.nan, 2)
+    assert df3.index_value.key != apply_s.index_value.key
+    assert df3.key != apply_s.index_value.key

maxframe/dataframe/indexing/add_prefix_suffix.py

@@ -51,7 +51,7 @@ def _get_prefix_suffix_docs(is_prefix: bool):
     Examples
     --------
     >>> import maxframe.dataframe as md
-        >>> s = md.Series([1, 2, 3, 4])
+    >>> s = md.Series([1, 2, 3, 4])
     >>> s.execute()
     0    1
     1    2

maxframe/dataframe/indexing/rename.py

@@ -248,6 +248,7 @@ def df_rename(
     )
 
 
+# fixme https://github.com/aliyun/alibabacloud-odps-maxframe-client/issues/58
 def series_rename(
     series,
     index=None,
@@ -382,6 +383,7 @@ def index_rename(index, name, inplace=False):
         return ret
 
 
+# fixme https://github.com/aliyun/alibabacloud-odps-maxframe-client/issues/59
 def index_set_names(index, names, level=None, inplace=False):
     """
     Set Index or MultiIndex name.
@@ -407,6 +409,15 @@ def index_set_names(index, names, level=None, inplace=False):
     See Also
     --------
     Index.rename : Able to set new names without level.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> idx = md.Index([1, 2, 3, 4])
+    >>> idx.execute()
+    Int64Index([1, 2, 3, 4], dtype='int64')
+    >>> idx.set_names('quarter').execute()
+    Int64Index([1, 2, 3, 4], dtype='int64', name='quarter')
     """
     op = DataFrameRename(
         index_mapper=names, level=level, output_types=get_output_types(index)

maxframe/dataframe/initializer.py

@@ -15,6 +15,7 @@
 from typing import Union
 
 import pandas as pd
+from pandas.api.types import is_list_like
 from pandas.core.dtypes.common import pandas_dtype
 
 from ..core import ENTITY_TYPE
@@ -61,6 +62,8 @@ class DataFrame(_Frame, metaclass=InitializerMeta):
         num_partitions=None,
     ):
         need_repart = False
+        if columns is not None and not is_list_like(columns):
+            raise ValueError("columns must be a list-like object")
         if isinstance(data, TENSOR_TYPE):
             if chunk_size is not None:
                 data = data.rechunk(chunk_size)
@@ -69,7 +72,10 @@ class DataFrame(_Frame, metaclass=InitializerMeta):
             )
             need_repart = num_partitions is not None
         elif isinstance(data, SERIES_TYPE):
-            df = data.to_frame()
+            if columns is not None and len(columns) != 1:
+                raise ValueError("columns' length must be 1 when data is Series")
+            col_name = columns[0] if columns else None
+            df = data.to_frame(name=col_name)
             need_repart = num_partitions is not None
         elif isinstance(data, DATAFRAME_TYPE):
             if not hasattr(data, "data"):
@@ -77,6 +83,10 @@ class DataFrame(_Frame, metaclass=InitializerMeta):
                 df = _Frame(data)
             else:
                 df = data
+            if columns is not None:
+                if len(df.columns) != len(columns):
+                    raise ValueError("columns' length must be equal to the data's")
+                df.columns = columns
             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

maxframe/dataframe/misc/drop_duplicates.py

@@ -104,7 +104,6 @@ def df_drop_duplicates(
 def series_drop_duplicates(
     series, keep="first", inplace=False, ignore_index=False, method="auto"
 ):
-    # FIXME: https://github.com/aliyun/alibabacloud-odps-maxframe-client/issues/12
     """
     Return Series with duplicate values removed.
 
@@ -148,6 +147,24 @@ def series_drop_duplicates(
     5     hippo
     Name: animal, dtype: object
 
+    With the 'keep' parameter, the selection behaviour of duplicated values
+    can be changed. The value 'first' keeps the first occurrence for each
+    set of duplicated entries. The default value of keep is 'first'.
+    >>> s.drop_duplicates().execute()
+    0      lame
+    1       cow
+    3    beetle
+    5     hippo
+    Name: animal, dtype: object
+    The value 'last' for parameter 'keep' keeps the last occurrence for
+    each set of duplicated entries.
+    >>> s.drop_duplicates(keep='last').execute()
+    1       cow
+    3    beetle
+    4      lame
+    5     hippo
+    Name: animal, dtype: object
+
     The value ``False`` for parameter 'keep' discards all sets of
     duplicated entries. Setting the value of 'inplace' to ``True`` performs
     the operation inplace and returns ``None``.

maxframe/dataframe/tests/test_initializer.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 
 import pandas as pd
+import pytest
 
 from ..core import DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
-from ..initializer import read_pandas
+from ..initializer import DataFrame, Series, read_pandas
 
 
-def test_from_pandas():
+def test_read_pandas():
     df_data = pd.DataFrame([["a", 1], ["b", 2]], columns=["a", "b"])
     assert isinstance(read_pandas(df_data), DATAFRAME_TYPE)
 
@@ -27,3 +28,33 @@ def test_from_pandas():
 
     idx_data = pd.Index(["a", "b"])
     assert isinstance(read_pandas(idx_data), INDEX_TYPE)
+
+
+def test_init_dataframe_from_maxframe_series():
+    s = Series([1, 2, 3, 4], index=[1, 2, 3, 4])
+
+    df = DataFrame(s, index=s.index, columns=["col1"])
+
+    assert isinstance(df, DATAFRAME_TYPE)
+    assert df.dtypes.index == ["col1"]
+
+    with pytest.raises(ValueError):
+        DataFrame(s, index=s.index, columns=[])
+
+    with pytest.raises(ValueError):
+        DataFrame(s, index=s.index, columns="col1")
+
+    with pytest.raises(ValueError):
+        DataFrame(s, index=s.index, columns="col2")
+
+
+def test_init_dataframe_from_maxframe_dataframe():
+    df1 = DataFrame({"A": [1, 2, 3, 4], "B": [1, 2, 3, 4]}, index=[1, 2, 3, 4])
+
+    df2 = DataFrame(df1, index=df1.index, columns=["col1", "col2"])
+
+    assert isinstance(df2, DATAFRAME_TYPE)
+    assert list(df2.dtypes.index) == ["col1", "col2"]
+
+    with pytest.raises(ValueError):
+        DataFrame(df1, index=df1.index, columns=["col1", "col2", "col3"])

maxframe/io/odpsio/schema.py

@@ -54,7 +54,9 @@ _odps_type_to_arrow = {
     odps_types.double: pa.float64(),
     odps_types.date: pa.date32(),
     odps_types.datetime: pa.timestamp("ms"),
+    odps_types.json: pa.string(),
     odps_types.timestamp: pa.timestamp("ns"),
+    odps_types.timestamp_ntz: pa.timestamp("ns"),
 }
 
 
@@ -166,7 +168,7 @@ def odps_schema_to_pandas_dtypes(
     return arrow_schema.empty_table().to_pandas().dtypes
 
 
-def _is_scalar_object(df_obj: Any) -> bool:
+def is_scalar_object(df_obj: Any) -> bool:
     return (
         isinstance(df_obj, TENSOR_TYPE) and df_obj.shape == ()
     ) or pd_types.is_scalar(df_obj)
@@ -187,7 +189,7 @@ def pandas_to_odps_schema(
     from ... import dataframe as md
     from .arrow import pandas_to_arrow
 
-    if _is_scalar_object(df_obj):
+    if is_scalar_object(df_obj):
         empty_index = None
     elif hasattr(df_obj, "index_value"):
         empty_index = df_obj.index_value.to_pandas()[:0]
@@ -289,7 +291,7 @@ def build_dataframe_table_meta(
         obj_type = OutputType.series
     elif isinstance(df_obj, (md.Index, pd.Index)):
         obj_type = OutputType.index
-    elif _is_scalar_object(df_obj):
+    elif is_scalar_object(df_obj):
         obj_type = OutputType.scalar
     else:  # pragma: no cover
         raise TypeError(f"Cannot accept type {type(df_obj)}")