maxframe 2.0.0b2__cp37-cp37m-win32.whl → 2.3.0rc1__cp37-cp37m-win32.whl

maxframe/dataframe/extensions/extract_kv.py

@@ -0,0 +1,177 @@
+# 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 EntityData, OutputType
+from ...serialization.serializables import AnyField, KeyField, StringField
+from ...utils import make_dtype, no_default
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import make_column_list
+
+
+class DataFrameExtractKv(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.EXTRACT_KV
+
+    columns = AnyField("columns", default=None)
+    kv_delim = StringField("kv_delim", default="=")
+    item_delim = StringField("item_delim", default=",")
+    dtype = AnyField("dtype", default=None)
+    fill_value = AnyField("fill_value", default=None)
+    errors = StringField("errors", default="raise")
+    # intermediate agg data
+    agg_results = KeyField("agg_results", default=None)
+
+    def __init__(self, kv_delim="=", item_delim=",", **kw):
+        super().__init__(kv_delim=kv_delim, item_delim=item_delim, **kw)
+        self.output_types = [OutputType.dataframe]
+
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameExtractKv", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        if op.agg_results is not None:
+            op.agg_results = inputs[-1]
+
+    def __call__(self, df):
+        shape = (df.shape[0], np.nan)
+        errors_arg = self.errors
+
+        def get_keys(row, cols, kv_delim, item_delim):
+            for col in cols:
+                if row[col] is not None:
+                    pairs = row[col].split(item_delim)
+                else:
+                    pairs = []
+                for pair in pairs:
+                    result = pair.split(kv_delim, 1)
+                    if len(result) == 2:
+                        yield f"{col}_{result[0]}"
+                    elif errors_arg == "raise":
+                        raise ValueError(f"Malformed data {pair} in column '{col}'.")
+
+        all_keys = df.mf.flatmap(
+            get_keys,
+            dtypes=pd.Series([str], index=["keys_cols"]),
+            cols=self.columns,
+            kv_delim=self.kv_delim,
+            item_delim=self.item_delim,
+        )
+        self.agg_results = all_keys.drop_duplicates().sort_values(by="keys_cols")
+        inputs = [df]
+        inputs.append(self.agg_results)
+        return self.new_dataframe(
+            inputs,
+            shape=shape,
+            dtypes=None,
+            index_value=df.index_value,
+            columns_value=None,
+        )
+
+
+def extract_kv(
+    data,
+    columns=None,
+    kv_delim="=",
+    item_delim=",",
+    dtype="float",
+    fill_value=None,
+    errors="raise",
+):
+    """
+    Extract values in key-value represented columns into standalone columns.
+    New column names will be the name of the key-value column followed by
+    an underscore and the key.
+
+    Parameters
+    ----------
+    columns : list, default None
+        The key-value columns to be extracted.
+    kv_delim : str, default '='
+        Delimiter between key and value.
+    item_delim : str, default ','
+        Delimiter between key-value pairs.
+    dtype : str
+        Type of value columns to generate.
+    fill_value : object, default None
+        Default value for missing key-value pairs.
+    errors : {'ignore', 'raise'}, default 'raise'
+        * If 'raise', then invalid parsing will raise an exception.
+        * If 'ignore', then invalid parsing will return the input.
+
+    Returns
+    -------
+    DataFrame
+        extracted data frame
+
+    See Also
+    --------
+    DataFrame.mf.collect_kv
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import maxframe.dataframe as md
+
+    >>> df = md.DataFrame({"name": ["name1", "name2", "name3", "name4", "name5"],
+    ...                    "kv": ["k1=1.0,k2=3.0,k5=10.0",
+    ...                           "k2=3.0,k3=5.1",
+    ...                           "k1=7.1,k7=8.2",
+    ...                           "k2=1.2,k3=1.5",
+    ...                           "k2=1.0,k9=1.1"]})
+    >>> df.execute()
+       name   kv
+    0  name1  k1=1.0,k2=3.0,k5=10.0
+    1  name2  k2=3.0,k3=5.1
+    2  name3  k1=7.1,k7=8.2
+    3  name4  k2=1.2,k3=1.5
+    4  name5  k2=1.0,k9=1.1
+
+    The field names to be expanded are specified by columns
+    kv_delim is to delimit the key and value and '=' is default
+    item_delim is to delimit the Key-Value pairs, ',' is default
+    The output field name is the original field name connect with the key by "_"
+    fill_value is used to fill missing values, None is default
+
+    >>> df.mf.extract_kv(columns=['kv'], kv_delim='=', item_delim=',').execute()
+       name   kv_k1   kv_k2   kv_k3   kv_k5   kv_k7   kv_k9
+    0  name1  1.0     3.0     NaN     10.0    NaN     NaN
+    1  name2  NaN     3.0     5.1     NaN     NaN     NaN
+    2  name3  7.1     NaN     NaN     NaN     8.2     NaN
+    3  name4  NaN     1.2     1.5     NaN     NaN     NaN
+    4  name5  NaN     1.0     NaN     NaN     NaN     1.1
+    """
+    if columns is None:
+        columns = data.dtypes.index.tolist()
+    columns_list = make_column_list(columns, data.dtypes)
+    non_exist_key = next(
+        (c for c in columns_list if c not in data.dtypes.index), no_default
+    )
+    if non_exist_key is not no_default:
+        raise ValueError(f"Column {non_exist_key} specified is not a valid column.")
+    for col in columns_list:
+        if str(data.dtypes[col]) not in ("object", "string"):
+            raise ValueError(f"Column '{col}' must be of string type.")
+    op = DataFrameExtractKv(
+        columns=columns,
+        kv_delim=kv_delim,
+        item_delim=item_delim,
+        dtype=make_dtype(dtype),
+        fill_value=fill_value,
+        errors=errors,
+    )
+    return op(data)

maxframe/dataframe/extensions/flatjson.py

@@ -39,12 +39,13 @@ class SeriesFlatJSONOperator(DataFrameOperator, DataFrameOperatorMixin):
                 name=name,
                 dtype=make_dtype(dtype),
             )
+        dtypes = make_dtypes(dtypes)
         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),
+            dtypes=dtypes,
         )
 
 

maxframe/dataframe/extensions/map_reduce.py

@@ -0,0 +1,263 @@
+# 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.
+
+import inspect
+from typing import Any, Callable, List, Optional, Union
+
+import pandas as pd
+
+
+def _has_end_arg(func) -> bool:
+    f_args = inspect.getfullargspec(func)
+    return "end" in f_args.args or "end" in f_args.kwonlyargs
+
+
+def _gen_combined_mapper(
+    mapper: Callable,
+    combiner: Callable,
+    group_cols: List[Any],
+    order_cols: List[Any],
+    ascending: Union[bool, List[bool]] = True,
+):
+    class CombinedMapper:
+        def __init__(self):
+            if isinstance(mapper, type):
+                self.f = mapper()
+            else:
+                self.f = mapper
+
+            if isinstance(combiner, type):
+                self.combiner = combiner()
+            else:
+                self.combiner = combiner
+
+        def _combine_mapper_result(self, mapper_result, end=False):
+            if mapper_result is None:
+                return None
+            res = mapper_result
+            if order_cols:
+                res = mapper_result.sort_values(order_cols, ascending=ascending)
+
+            kw = {"end": end} if _has_end_arg(self.combiner) else {}
+            gcols = group_cols or list(res.columns)
+            return res.groupby(gcols, group_keys=False)[list(res.columns)].apply(
+                self.combiner, **kw
+            )
+
+        def __call__(self, batch, end=False):
+            kw = {"end": end} if _has_end_arg(self.f) else {}
+            f_ret = self.f(batch, **kw)
+            return self._combine_mapper_result(f_ret, end=end)
+
+        def close(self) -> None:
+            if hasattr(self.f, "close"):
+                self.f.close()
+            if hasattr(self.combiner, "close"):
+                self.combiner.close()
+
+    return CombinedMapper
+
+
+def map_reduce(
+    df,
+    mapper: Optional[Callable] = None,
+    reducer: Optional[Callable] = None,
+    group_cols: Optional[List[Any]] = None,
+    *,
+    order_cols: List[Any] = None,
+    ascending: Union[bool, List[bool]] = True,
+    combiner: Callable = None,
+    batch_rows: Optional[int] = 1024,
+    mapper_dtypes: pd.Series = None,
+    mapper_index: pd.Index = None,
+    mapper_batch_rows: Optional[int] = None,
+    reducer_dtypes: pd.Series = None,
+    reducer_index: pd.Index = None,
+    reducer_batch_rows: Optional[int] = None,
+    ignore_index: bool = False,
+):
+    """
+    Map-reduce API over certain DataFrames. This function is roughly
+    a shortcut for
+
+    .. code-block:: python
+
+        df.mf.apply_chunk(mapper).groupby(group_keys).mf.apply_chunk(reducer)
+
+    Parameters
+    ----------
+    mapper : function or type
+        Mapper function or class.
+    reducer : function or type
+        Reducer function or class.
+    group_cols : str or list[str]
+        The keys to group after mapper. If absent, all columns in the mapped
+        DataFrame will be used.
+    order_cols : str or list[str]
+        The columns to sort after groupby.
+    ascending : bool or list[bool] or None
+        Whether columns should be in ascending order or not, only effective when
+        `order_cols` are specified. If a list of booleans are passed, orders of
+        every column in `order_cols` are specified.
+    combiner : function or class
+        Combiner function or class. Should accept and returns the same schema
+        of mapper outputs.
+    batch_rows : int or None
+        Rows in batches for mappers and reducers. Ignored if `mapper_batch_rows`
+        specified for mappers or `reducer_batch_rows` specified for reducers.
+        1024 by default.
+    mapper_dtypes : pd.Series or dict or None
+        Output dtypes of mapper stage.
+    mapper_index : pd.Index or None
+        Index of DataFrame returned by mappers.
+    mapper_batch_rows : int or None
+        Rows in batches for mappers. If specified, `batch_rows` will be ignored
+        for mappers.
+    reducer_dtypes : pd.Series or dict or None
+        Output dtypes of reducer stage.
+    reducer_index : pd.Index or None
+        Index of DataFrame returned by reducers.
+    reducer_batch_rows : int or None
+        Rows in batches for mappers. If specified, `batch_rows` will be ignored
+        for reducers.
+    ignore_index : bool
+        If true, indexes generated at mapper or reducer functions will be ignored.
+
+    Returns
+    -------
+    output: DataFrame
+        Result DataFrame after map and reduce.
+
+    Examples
+    --------
+
+    We first define a DataFrame with a column of several words.
+
+    >>> from collections import defaultdict
+    >>> import maxframe.dataframe as md
+    >>> from maxframe.udf import with_running_options
+    >>> df = pd.DataFrame(
+    >>>     {
+    >>>         "name": ["name key", "name", "key", "name", "key name"],
+    >>>         "id": [4, 2, 4, 3, 3],
+    >>>         "fid": [5.3, 3.5, 4.2, 2.2, 4.1],
+    >>>     }
+    >>> )
+
+    Then we write a mapper function which accepts batches in the DataFrame
+    and returns counts of words in every row.
+
+    >>> def mapper(batch):
+    >>>     word_to_count = defaultdict(lambda: 0)
+    >>>     for words in batch["name"]:
+    >>>         for w in words.split():
+    >>>             word_to_count[w] += 1
+    >>>     return pd.DataFrame(
+    >>>         [list(tp) for tp in word_to_count.items()], columns=["word", "count"]
+    >>>     )
+
+    After that we write a reducer function which aggregates records with
+    the same word. Running options such as CPU specifications can be supplied
+    as well.
+
+    >>> @with_running_options(cpu=2)
+    >>> class TestReducer:
+    >>>     def __init__(self):
+    >>>         self._word_to_count = defaultdict(lambda: 0)
+    >>>
+    >>>     def __call__(self, batch, end=False):
+    >>>         word = None
+    >>>         for _, row in batch.iterrows():
+    >>>             word = row.iloc[0]
+    >>>             self._word_to_count[row.iloc[0]] += row.iloc[1]
+    >>>         if end:
+    >>>             return pd.DataFrame(
+    >>>                 [[word, self._word_to_count[word]]], columns=["word", "count"]
+    >>>             )
+    >>>
+    >>>     def close(self):
+    >>>         # you can do several cleanups here
+    >>>         print("close")
+
+    Finally we can call `map_reduce` with mappers and reducers specified above.
+
+    >>> res = df.mf.map_reduce(
+    >>>     mapper,
+    >>>     TestReducer,
+    >>>     group_cols=["word"],
+    >>>     mapper_dtypes={"word": "str", "count": "int"},
+    >>>     mapper_index=pd.Index([0]),
+    >>>     reducer_dtypes={"word": "str", "count": "int"},
+    >>>     reducer_index=pd.Index([0]),
+    >>>     ignore_index=True,
+    >>> )
+    >>> res.execute().fetch()
+       word  count
+    0   key      3
+    1  name      4
+
+    See Also
+    --------
+    DataFrame.mf.apply_chunk, DataFrame.groupby.mf.apply_chunk
+    """
+    mapper_batch_rows = mapper_batch_rows or batch_rows
+    reducer_batch_rows = reducer_batch_rows or batch_rows
+
+    def check_arg(arg_type, locals_):
+        if locals_.get(arg_type) is not None:
+            return
+        for suffix in ("dtypes", "index"):
+            arg_name = f"{arg_type}_{suffix}"
+            if locals_.get(arg_name) is not None:
+                raise ValueError(f"Cannot specify {arg_name} when {arg_type} is None")
+
+    if mapper is None:
+        check_arg("mapper", locals())
+        mapped = df
+        group_cols = group_cols or df.dtypes.index
+        if combiner is not None:
+            raise ValueError("Combiner cannot be set when mapper is None")
+    else:
+        if combiner is not None:
+            mapper = _gen_combined_mapper(
+                mapper, combiner, group_cols, order_cols, ascending=ascending
+            )
+        mapped = df.mf.apply_chunk(
+            mapper,
+            batch_rows=mapper_batch_rows,
+            dtypes=mapper_dtypes,
+            output_type="dataframe",
+            index=mapper_index,
+        )
+        group_cols = group_cols or list(df.dtypes.index)
+
+    if reducer is None:
+        check_arg("reducer", locals())
+        res = mapped
+    else:
+        res = mapped.groupby(group_cols, group_keys=False)[
+            list(mapped.dtypes.index)
+        ].mf.apply_chunk(
+            reducer,
+            batch_rows=reducer_batch_rows,
+            dtypes=reducer_dtypes,
+            output_type="dataframe",
+            index=reducer_index,
+            order_cols=order_cols,
+            ascending=ascending,
+        )
+
+    if ignore_index:
+        return res.reset_index(drop=True)
+    return res

maxframe/dataframe/extensions/rebalance.py

@@ -0,0 +1,62 @@
+# 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 ... import opcodes
+from ...serialization.serializables import Float64Field, Int64Field
+from ...tensor.extensions.rebalance import RebalanceMixin
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import validate_axis
+
+
+class DataFrameRebalance(RebalanceMixin, DataFrameOperatorMixin, DataFrameOperator):
+    _op_type_ = opcodes.REBALANCE
+
+    axis = Int64Field("axis")
+    factor = Float64Field("factor", default=None)
+    num_partitions = Int64Field("num_partitions")
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+
+def rebalance(df_or_series, axis=0, factor=None, num_partitions=None):
+    """
+    Make data more balanced across entire cluster.
+
+    Parameters
+    ----------
+    axis : int
+        The axis to rebalance.
+    factor : float
+        Specified so that number of chunks after balance is
+        total number of input chunks * factor.
+    num_partitions : int
+        Specified so the number of chunks are at most
+        num_partitions.
+
+    Returns
+    -------
+    Series or DataFrame
+        Result of DataFrame or Series after rebalanced.
+    """
+    axis = validate_axis(axis, df_or_series)
+    if num_partitions is None and factor is None:
+        raise ValueError("Need to specify num_partitions or factor")
+    if num_partitions is not None and factor is not None:
+        raise ValueError(
+            "num_partitions and factor cannot be specified at the same time"
+        )
+
+    op = DataFrameRebalance(axis=axis, factor=factor, num_partitions=num_partitions)
+    return op(df_or_series)

maxframe/dataframe/extensions/tests/test_apply_chunk.py

@@ -91,15 +91,22 @@ def test_apply_chunk_infer_dtypes_and_index(df1, df2, df3):
     assert result.index_value is df2.index_value
     assert result.dtypes.equals(df2.dtypes)
 
+    def process(data, param, k) -> pd.DataFrame[df2.dtypes]:
+        return data * param * k
+
+    result = df2.mf.apply_chunk(process, batch_rows=3, 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):
+    def process(data, k) -> pd.DataFrame[df1.dtypes]:
         return data
 
-    result = df1.mf.apply_chunk(process, batch_rows=3, output_type="dataframe", k=1)
+    result = df1.mf.apply_chunk(process, batch_rows=3, k=1)
     assert result.index_value is df1.index_value
     assert result.dtypes.equals(df1.dtypes)
     assert isinstance(result.op.func, MarkedFunction)

maxframe/dataframe/extensions/tests/test_extensions.py

@@ -60,6 +60,40 @@ def df3():
     )
 
 
+@pytest.fixture
+def df4():
+    return DataFrame(
+        {
+            "name1": ["a", "b", "c", "d"],
+            "name2": ["a", "b", "c", "d"],
+            "num": [1, 2, 3, 4],
+            "kv": [
+                "k1=1.1,k2=3.1,k3=1.0",
+                "k1=7.1,k4=8.2",
+                "k5=1.2,k7=1.5",
+                "k3=1.1,k9=1",
+            ],
+            "vk": ["v1=1.1,v2=1.2", "v3=1.1,v4=1.2", "v5=1.1,v6=1.2", "v7=1.1,v8=1.2"],
+        }
+    )
+
+
+@pytest.fixture
+def df5():
+    return DataFrame(
+        {
+            "name1": ["name1", "name2", "name3", "name4", "name5"],
+            "name2": ["name1", "name2", "name3", "name4", "name5"],
+            "k1": [1.0, None, 7.1, None, None],
+            "k2": [3.0, 3.0, None, 1.2, 1.0],
+            "k3": [None, 5.1, None, 1.5, None],
+            "k5": [10.0, None, None, None, None],
+            "k7": [None, None, 8.2, None, None],
+            "k9": [None, None, None, None, 1.1],
+        }
+    )
+
+
 def test_flatmap(df1, df2, df3):
     def f(x, keys):
         if x["a"] in keys:
@@ -142,3 +176,23 @@ def test_flatjson():
         )
     with pytest.raises(ValueError):
         s1.mf.flatjson(["$.a"])
+
+
+def test_extract_kv(df4):
+    extract_kv_df = df4.mf.extract_kv(
+        columns=["kv", "vk"], kv_delim=",", item_delim="="
+    )
+    assert extract_kv_df.shape == (4, np.nan)
+    assert extract_kv_df.index_value.key == df4.index_value.key
+    with pytest.raises(ValueError):
+        df4.mf.extract_kv(columns=["name"])
+    with pytest.raises(ValueError):
+        df4.mf.extract_kv(columns=["num"])
+
+
+def test_collect_kv(df5):
+    collect_kv_df = df5.mf.collect_kv(columns=["k1", "k2", "k3", "k5", "k7", "k9"])
+    assert collect_kv_df.shape == (5, 3)
+    assert collect_kv_df.index_value.key == df5.index_value.key
+    with pytest.raises(ValueError):
+        df5.mf.collect_kv(columns=["num"])