maxframe 1.0.0rc3__cp311-cp311-macosx_10_9_universal2.whl → 1.1.0__cp311-cp311-macosx_10_9_universal2.whl

maxframe/dataframe/extensions/flatjson.py

@@ -0,0 +1,131 @@
+# 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 List
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import ListField
+from ...serialization.serializables.field_type import FieldTypes
+from ..core import DataFrame
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import make_dtypes, parse_index
+
+
+class SeriesFlatJSONOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.FLATJSON
+
+    query_paths = ListField("query_paths", field_type=FieldTypes.string, default=None)
+
+    def __call__(self, series, dtypes):
+        if self._output_types[0] == OutputType.series:
+            name, dtype = dtypes
+            return self.new_series(
+                [series],
+                shape=series.shape,
+                index_value=series.index_value,
+                name=name,
+                dtype=dtype,
+            )
+        return self.new_dataframe(
+            [series],
+            shape=(series.shape[0], len(dtypes)),
+            index_value=series.index_value,
+            columns_value=parse_index(dtypes.index, store_data=True),
+            dtypes=make_dtypes(dtypes),
+        )
+
+
+def series_flatjson(
+    series,
+    query_paths: List[str],
+    dtypes=None,
+    dtype=None,
+    name: str = None,
+) -> DataFrame:
+    """
+    Flat JSON object in the series to a dataframe according to JSON query.
+
+    Parameters
+    ----------
+    series : Series
+        The series of json strings.
+
+    query_paths: List[str] or str
+        The JSON query paths for each generated column. The path format should follow
+        [RFC9535](https://datatracker.ietf.org/doc/rfc9535/).
+
+    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.
+
+    Returns
+    -------
+    DataFrame or Series
+        Result of DataFrame when dtypes specified, else Series.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> import pandas as pd
+    >>> s = md.Series(
+    ...     [
+    ...         '{"age": 24, "gender": "male", "graduated": false}',
+    ...         '{"age": 25, "gender": "female", "graduated": true}',
+    ...     ]
+    ... )
+    >>> s.execute()
+    0    {"age": 24, "gender": "male", "graduated": false}
+    1    {"age": 25, "gender": "female", "graduated": true}
+    dtype: object
+
+    >>> df = s.mf.flatjson(
+    ...    ["$.age", "$.gender", "$.graduated"],
+    ...    dtypes=pd.Series(["int32", "object", "bool"], index=["age", "gender", "graduated"]),
+    ... )
+    >>> df.execute()
+        age  gender  graduated
+    0   24    male       True
+    1   25  female       True
+
+    >>> s2 = s.mf.flatjson("$.age", name="age", dtype="int32")
+    >>> s2.execute()
+    0    24
+    1    25
+    Name: age, dtype: int32
+    """
+    if isinstance(query_paths, str):
+        query_paths = [query_paths]
+    if dtypes is not None and dtype is not None:
+        raise ValueError("Both dtypes and dtype cannot be specified at the same time.")
+    if dtype is not None:
+        if len(query_paths) != 1:
+            raise ValueError("query_paths should have only one path if dtype is set")
+        output_type = OutputType.series
+    elif dtypes is not None:
+        if len(dtypes) != len(query_paths):
+            raise ValueError("query_paths and dtypes should have same length")
+        output_type = OutputType.dataframe
+    else:
+        raise ValueError("dtypes or dtype should be specified")
+
+    dtypes = (name, dtype) if dtype is not None else dtypes
+    return SeriesFlatJSONOperator(query_paths=query_paths, _output_types=[output_type])(
+        series, dtypes
+    )

maxframe/dataframe/extensions/flatmap.py

@@ -0,0 +1,314 @@
+# 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 ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import (
+    BoolField,
+    DictField,
+    FunctionField,
+    TupleField,
+)
+from ..core import DataFrame
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import gen_unknown_index_value, make_dtypes, parse_index
+
+
+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)
+
+    def _call_dataframe(self, df: DataFrame, dtypes: pd.Series):
+        dtypes = make_dtypes(dtypes)
+        index_value = gen_unknown_index_value(
+            df.index_value,
+            (df.key, df.index_value.key, self.func),
+            normalize_range_index=True,
+        )
+        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 = gen_unknown_index_value(
+            series.index_value,
+            (series.key, series.index_value.key, self.func),
+            normalize_range_index=True,
+        )
+
+        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
+    ----------
+    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
+    ----------
+    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]
+
+    Specify ``dtype`` with a function which returns list to return more 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:
+
+
+    >>> def generate_values_in_generator(x):
+    ...     yield pd.Series([x * 2, x * 4])
+    ...     yield pd.Series([x * 3, x * 5])
+
+    >>> 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 is not None and dtype is not None:
+        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/reshuffle.py

@@ -38,7 +38,7 @@ class DataFrameReshuffle(DataFrameOperator, DataFrameOperatorMixin):
         else:
             idx_value = df.index_value
             if isinstance(idx_value.value, IndexValue.RangeIndex):
-                idx_value = parse_index(pd.Int64Index([0]))
+                idx_value = parse_index(pd.RangeIndex(1))
         params = df.params
         params["index_value"] = idx_value
         self._output_types = get_output_types(df)

maxframe/dataframe/extensions/tests/test_apply_chunk.py

@@ -0,0 +1,186 @@
+# 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
+import pytest
+
+from ....udf import MarkedFunction
+from ... import DataFrame
+from ...core import DATAFRAME_TYPE, SERIES_TYPE
+
+
+@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_apply_chunk_infer_dtypes_and_index(df1, df2, df3):
+    # dataframe -> dataframe filter
+    result = df3.mf.apply_chunk(
+        lambda data: data.query("A > 1"), batch_rows=2, output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert df3.index_value.key != result.index_value.key
+    assert df3.index_value.to_pandas().names == result.index_value.to_pandas().names
+
+    # dataframe -> dataframe keep same
+    result = df1.mf.apply_chunk(
+        lambda data: data, batch_rows=2, output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert result.index_value is df1.index_value
+
+    # dataframe -> dataframe ufunc with arguments
+    result = df1.mf.apply_chunk(
+        np.add, batch_rows=2, args=(2,), output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert result.index_value is df1.index_value
+    assert result.dtypes.equals(df1.dtypes)
+    assert result.shape == df1.shape
+
+    # dataframe -> series ufunc return series
+    result = df1.mf.apply_chunk(np.sum, batch_rows=2)
+    assert isinstance(result, SERIES_TYPE)
+    assert result.index_value is not df1.index_value
+
+    # series -> series
+    result = df3.a.mf.apply_chunk(lambda data: data, batch_rows=2, output_type="series")
+    assert isinstance(result, SERIES_TYPE)
+    assert df3.a.index_value is result.index_value
+
+    result = df3.a.mf.apply_chunk(
+        np.sum, batch_rows=2, output_type="series", dtype=np.int64, name="sum"
+    )
+    assert isinstance(result, SERIES_TYPE)
+    assert isinstance(result.index_value.to_pandas(), pd.RangeIndex)
+
+    # general functions
+    def process(data, param, k):
+        return data * param * k
+
+    result = df2.mf.apply_chunk(
+        process, batch_rows=3, output_type="dataframe", args=(4,), k=1
+    )
+    assert result.index_value is df2.index_value
+    assert result.dtypes.equals(df2.dtypes)
+
+    # mark functions
+    from ....udf import with_python_requirements, with_resources
+
+    @with_resources("empty.txt")
+    @with_python_requirements("numpy")
+    def process(data, k):
+        return data
+
+    result = df1.mf.apply_chunk(process, batch_rows=3, output_type="dataframe", k=1)
+    assert result.index_value is df1.index_value
+    assert result.dtypes.equals(df1.dtypes)
+    assert isinstance(result.op.func, MarkedFunction)
+    assert result.op.func is not process
+    assert result.op.func.resources is process.resources
+    assert result.op.func.pythonpacks is process.pythonpacks
+
+    def func_series_ret_series(data):
+        return pd.DataFrame([data, data])
+
+    result = df3.a.mf.apply_chunk(
+        func_series_ret_series, batch_rows=2, output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert result.op.func is func_series_ret_series
+
+
+def test_apply_test(df1):
+    def process(x, param):
+        return x * param
+
+    result = df1.a.mf.apply_chunk(
+        process, batch_rows=2, output_type="series", args=(5,)
+    )
+    assert isinstance(result, SERIES_TYPE)
+
+
+def test_apply_chunk(df1):
+    keys = [1, 2]
+
+    def f(x, keys):
+        if x["a"] in keys:
+            return [1, 0]
+        else:
+            return [0, 1]
+
+    result = df1[["a"]].mf.apply_chunk(
+        f,
+        output_type="dataframe",
+        dtypes=pd.Series(["int64", "int64"]),
+        batch_rows=5,
+        keys=keys,
+    )
+    assert result.shape == (np.nan, 2)
+    assert df1.index_value.key != result.index_value.key
+
+    # dataframe return series
+    result = df1.mf.apply_chunk(
+        lambda x: x.a,
+        output_type="series",
+        dtype="int64",
+        batch_rows=5,
+    )
+    assert result.shape == (np.nan,)
+    assert df1.index_value.key == result.index_value.key
+    assert df1.a.index_value.key == result.index_value.key
+
+    # return dataframe with given dtypes
+    result = df1.a.mf.apply_chunk(
+        lambda x: pd.concat([x, x], axis=1),
+        output_type="dataframe",
+        dtypes=pd.Series(["int64", "int64"]),
+        batch_rows=5,
+    )
+    assert result.shape == (np.nan, 2)
+    assert df1.a.index_value.key != result.index_value.key
+
+    # return series but as dataframe
+    result = df1.a.mf.apply_chunk(
+        lambda x: pd.concat([x, x], axis=0),
+        output_type="dataframe",
+        dtypes={"c": np.int_},
+        batch_rows=5,
+    )
+    assert result.shape == (np.nan, 1)
+
+
+def test_apply_chunk_exception(df1):
+    with pytest.raises(ValueError):
+        df1.mf.apply_chunk(lambda data: data, batch_rows=-1, output_type="dataframe")
+
+    with pytest.raises(TypeError):
+        df1.mf.apply_chunk(
+            lambda data: data, batch_rows=object(), output_type="dataframe"
+        )