maxframe 1.3.1__cp310-cp310-win32.whl → 2.0.0b1__cp310-cp310-win32.whl

maxframe/codegen/spe/dataframe/reduction.py

@@ -0,0 +1,160 @@
+# 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 Any, Dict, List, Union
+
+from ....dataframe.reduction import (
+    DataFrameAggregate,
+    DataFrameAll,
+    DataFrameAny,
+    DataFrameCount,
+    DataFrameKurtosis,
+    DataFrameMax,
+    DataFrameMean,
+    DataFrameMedian,
+    DataFrameMin,
+    DataFrameNunique,
+    DataFrameProd,
+    DataFrameSem,
+    DataFrameSkew,
+    DataFrameSum,
+    DataFrameUnique,
+    DataFrameVar,
+)
+from ....dataframe.reduction.core import (
+    DataFrameCumReductionOperator,
+    DataFrameReductionOperator,
+)
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+
+@register_op_adapter(DataFrameCumReductionOperator)
+class DataFrameCumsumAdapter(SPEOperatorAdapter):
+    """
+    TODO: Refine this in window functions
+    """
+
+    def generate_code(
+        self, op: DataFrameCumReductionOperator, context: SPECodeContext
+    ) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        args = []
+        if op.axis is not None:
+            args.append(f"axis={op.axis!r}")
+        if op.skipna is not None:
+            args.append(f"skipna={op.skipna!r}")
+        args_str = ", ".join(args)
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        return [f"{res_var_name} = {input_var_name}.{op._func_name}({args_str})"]
+
+
+@register_op_adapter(DataFrameAggregate)
+class DataFrameAggregateAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: DataFrameAggregate, context: SPECodeContext):
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        agg_func_str = self.translate_var(context, op.raw_func)
+        axis_str = self.translate_var(context, op.axis)
+
+        return [
+            f"{res_var_name} = {input_var_name}.agg({agg_func_str}, axis={axis_str})"
+        ]
+
+
+# TODO: DataFrameStrConcat, DataFrameReductionSize and DataFrameCustomReduction
+@register_op_adapter(
+    [
+        DataFrameAll,
+        DataFrameAny,
+        DataFrameCount,
+        DataFrameMax,
+        DataFrameMean,
+        DataFrameMedian,
+        DataFrameMin,
+        DataFrameProd,
+        DataFrameSum,
+    ]
+)
+class DataFrameReductionAdapter(SPEOperatorAdapter):
+    _common_args = ["axis", "skipna", "numeric_only", "bool_only", "level", "min_count"]
+
+    def extra_args(self, op: DataFrameReductionOperator) -> Dict[str, Any]:
+        """
+        Get the extra arguments of the API call.
+
+        Parameters
+        ----------
+        op : DataFrameReductionOperator
+            The DataFrameReductionOperator instance.
+
+        Returns
+        -------
+        Dict[str, Any]:
+            Extra arguments key and values.
+        """
+        return dict()
+
+    def generate_code(
+        self, op: DataFrameReductionOperator, context: SPECodeContext
+    ) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        kwargs = dict()
+        for k in self._common_args:
+            v = getattr(op, k)
+            if v is not None:
+                kwargs[k] = v
+        kwargs.update(self.extra_args(op))
+        args_str = ", ".join(self._translate_call_args(context, **kwargs))
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        return [f"{res_var_name} = {input_var_name}.{op._func_name}({args_str})"]
+
+
+@register_op_adapter(DataFrameKurtosis)
+class DataFrameKurtosisAdapter(DataFrameReductionAdapter):
+    def extra_args(self, op: DataFrameKurtosis) -> Dict[str, Any]:
+        return {"bias": op.bias, "fisher": op.fisher}
+
+
+@register_op_adapter(DataFrameNunique)
+class DataFrameNuniqueAdapter(DataFrameReductionAdapter):
+    _common_args = ["level", "min_count"]
+
+    def extra_args(self, op: DataFrameNunique) -> Dict[str, Any]:
+        if op.inputs[0].ndim == 2:
+            return {"axis": op.axis, "dropna": op.dropna}
+        return {"dropna": op.dropna}
+
+
+@register_op_adapter(
+    [
+        DataFrameSem,
+        DataFrameSkew,
+        DataFrameVar,
+    ]
+)
+class DataFrameVarAdapter(DataFrameReductionAdapter):
+    def extra_args(
+        self, op: Union[DataFrameSem, DataFrameSkew, DataFrameVar]
+    ) -> Dict[str, Any]:
+        return {"ddof": op.ddof}
+
+
+@register_op_adapter(DataFrameUnique)
+class DataFrameUniqueAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: DataFrameUnique, context: SPECodeContext) -> List[str]:
+        context.register_import("pandas", "pd")
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        return [f"{res_var_name} = pd.unique({input_var_name})"]

maxframe/codegen/spe/dataframe/sort.py

@@ -0,0 +1,83 @@
+# 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 abc import abstractmethod
+from typing import Any, Dict, List
+
+from ....dataframe.core import DATAFRAME_TYPE
+from ....dataframe.sort import DataFrameSortIndex, DataFrameSortValues
+from ....dataframe.sort.core import DataFrameSortOperator
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+
+class _DataFrameSortAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: DataFrameSortOperator, context: SPECodeContext
+    ) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        kwargs = {
+            "axis": op.axis,
+            "ascending": op.ascending,
+            "inplace": op.inplace,
+            "na_position": op.na_position,
+            "ignore_index": op.ignore_index,
+        }
+        kwargs.update(self.extra_args(op, context))
+        args_str = ", ".join(self._translate_call_args(context, **kwargs))
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        return [f"{res_var_name} = {input_var_name}.{self.method_name}({args_str})"]
+
+    @abstractmethod
+    def extra_args(
+        self, op: DataFrameSortOperator, context: SPECodeContext
+    ) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def method_name(self) -> str:
+        raise NotImplementedError
+
+
+@register_op_adapter(DataFrameSortIndex)
+class DataFrameSortIndexAdapter(_DataFrameSortAdapter):
+    _method_name = "sort_index"
+
+    @property
+    def method_name(self) -> str:
+        return self._method_name
+
+    def extra_args(
+        self, op: DataFrameSortIndex, context: SPECodeContext
+    ) -> Dict[str, Any]:
+        return {
+            "level": op.level,
+            "sort_remaining": op.sort_remaining,
+        }
+
+
+@register_op_adapter(DataFrameSortValues)
+class DataFrameSortValuesAdapter(_DataFrameSortAdapter):
+    _method_name = "sort_values"
+
+    @property
+    def method_name(self) -> str:
+        return self._method_name
+
+    def extra_args(
+        self, op: DataFrameSortValues, context: SPECodeContext
+    ) -> Dict[str, Any]:
+        if isinstance(op.outputs[0], DATAFRAME_TYPE):
+            return {"by": op.by}
+        return {}

maxframe/codegen/spe/dataframe/statistics.py

@@ -0,0 +1,46 @@
+# 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
+
+from ....dataframe.statistics.corr import DataFrameCorr
+from ....dataframe.statistics.quantile import DataFrameQuantile
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+from ..utils import build_method_call_adapter
+
+
+@register_op_adapter(DataFrameCorr)
+class DataFrameCorrAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: DataFrameCorr, context: SPECodeContext) -> List[str]:
+        if op.other is None:
+            func_name = "corr"
+            args = []
+            kw_keys = ["method", "min_periods"]
+        else:
+            func_name = "corrwith"
+            args = ["other"]
+            kw_keys = ["axis", "drop", "method"]
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        args = self.generate_call_args_with_attributes(
+            op, context, *args, kw_keys=kw_keys
+        )
+        return [f"{res_var_name} = {input_var_name}.{func_name}({args})"]
+
+
+DataFrameQuantileAdapter = build_method_call_adapter(
+    DataFrameQuantile,
+    "quantile",
+    kw_keys=["q", "axis", "numeric_only", "interpolation"],
+)

maxframe/codegen/spe/dataframe/tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.

maxframe/codegen/spe/dataframe/tests/accessors/__init__.py

@@ -0,0 +1,13 @@
+# 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.

maxframe/codegen/spe/dataframe/tests/accessors/test_base.py

@@ -0,0 +1,33 @@
+# 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 dataframe as md
+from ....core import SPECodeContext
+from ...accessors.base import SeriesDatetimeMethodAdapter, SeriesStringMethodAdapter
+
+
+def test_datetime_methods():
+    v0 = md.Series([md.Timestamp("2023-10-11 22:11:32.123")])
+    v1 = v0.dt.year
+    context = SPECodeContext()
+    results = SeriesDatetimeMethodAdapter().generate_code(v1.op, context)
+    assert results[0] == "var_1 = var_0.dt.year"
+
+
+def test_string_methods():
+    v0 = md.Series(["hello, world"])
+    v1 = v0.str.upper()
+    context = SPECodeContext()
+    results = SeriesStringMethodAdapter().generate_code(v1.op, context)
+    assert results[0] == "var_1 = var_0.str.upper()"

maxframe/codegen/spe/dataframe/tests/accessors/test_dict.py

@@ -0,0 +1,310 @@
+# 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 numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from ...... import dataframe as md
+from ......lib.dtypes_extension import dict_
+from ......utils import ARROW_DTYPE_NOT_SUPPORTED
+from ....core import SPECodeContext
+from ...accessors.dict_ import (
+    SeriesDictContainsOperatorAdapter,
+    SeriesDictGetItemOperatorAdapter,
+    SeriesDictLengthOperatorAdapter,
+    SeriesDictRemoveOperatorAdapter,
+    SeriesDictSetItemOperatorAdapter,
+)
+
+pytestmark = pytest.mark.skipif(
+    ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported"
+)
+
+
+def _run_generated_code(
+    code: str, ctx: SPECodeContext, input_val: pd.DataFrame
+) -> dict:
+    local_vars = ctx.constants.copy()
+    local_vars["var_0"] = input_val
+    import_code = "import pandas as pd\nimport numpy as np\n"
+    exec(import_code + code, local_vars, local_vars)
+    return local_vars
+
+
+@pytest.fixture
+def pd_df_1():
+    return pd.DataFrame(
+        {
+            "A": pd.Series(
+                [[("k1", 1), ("k2", 2)], [("k1", 11)], None],
+                index=[1, 2, 3],
+                dtype=dict_(pa.string(), pa.int32()),
+            ),
+            "B": pd.Series([1, 2, 3], index=[1, 2, 3], dtype=np.dtype("int64")),
+        },
+    )
+
+
+@pytest.fixture
+def md_df_1(pd_df_1):
+    return md.DataFrame(pd_df_1)
+
+
+def test_getitem(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict["k1"]
+    context = SPECodeContext()
+    adapter = SeriesDictGetItemOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_get(data):
+    found = False
+    for tup in data:
+        if tup[0] == 'k1':
+            found = True
+            return tup[1]
+    if not found:
+        if False:
+            return None
+        else:
+            raise KeyError('k1')
+
+var_1 = var_0.map(_inner_get, na_action="ignore").astype(const_0)
+var_1.name = 'k1'
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, pd_df_1["A"])
+    expected_series = pd.Series(
+        [1, 11, None],
+        index=[1, 2, 3],
+        name="k1",
+        dtype=pd.ArrowDtype(pa.int32()),
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])
+
+
+@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
+def test_getitem_with_default_value(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict.get("k2", 9)
+    context = SPECodeContext()
+    adapter = SeriesDictGetItemOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_get(data):
+    found = False
+    for tup in data:
+        if tup[0] == 'k2':
+            found = True
+            return tup[1]
+    if not found:
+        if True:
+            return 9
+        else:
+            raise KeyError('k2')
+
+var_1 = var_0.map(_inner_get, na_action="ignore").astype(const_0)
+var_1.name = 'k2'
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, pd_df_1["A"])
+    expected_series = pd.Series(
+        [2, 9, None],
+        index=[1, 2, 3],
+        name="k2",
+        dtype=pd.ArrowDtype(pa.int32()),
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])
+
+
+def test_getitem_with_key_error(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict["k2"]
+    context = SPECodeContext()
+    adapter = SeriesDictGetItemOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_get(data):
+    found = False
+    for tup in data:
+        if tup[0] == 'k2':
+            found = True
+            return tup[1]
+    if not found:
+        if False:
+            return None
+        else:
+            raise KeyError('k2')
+
+var_1 = var_0.map(_inner_get, na_action="ignore").astype(const_0)
+var_1.name = 'k2'
+"""
+    ]
+    assert results == expected_results
+    with pytest.raises(KeyError):
+        _run_generated_code(results[0], context, pd_df_1["A"])
+
+
+def test_setitem(md_df_1, pd_df_1):
+    s1 = md_df_1["A"]
+    s1.dict["k2"] = 9
+    context = SPECodeContext()
+    adapter = SeriesDictSetItemOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_set(row):
+    found = False
+    value = list()
+    for tup in row:
+        if tup[0] == 'k2':
+            value.append((tup[0], 9))
+            found = True
+        else:
+            value.append(tup)
+    if not found:
+        value.append(('k2', 9))
+    return value
+
+var_1 = var_0.map(_inner_set, na_action="ignore").astype(const_0)
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, pd_df_1["A"])
+    expected_series = pd.Series(
+        [
+            [("k1", 1), ("k2", 9)],
+            [("k1", 11), ("k2", 9)],
+            None,
+        ],
+        index=[1, 2, 3],
+        name="A",
+        dtype=pd.ArrowDtype(pa.map_(pa.string(), pa.int32())),
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])
+
+
+@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
+def test_length(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict.len()
+    context = SPECodeContext()
+    adapter = SeriesDictLengthOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+var_1 = var_0.map(len, na_action="ignore").astype(const_0)
+var_1.name = None
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, pd_df_1["A"])
+    expected_series = pd.Series(
+        [2, 1, pd.NA], index=[1, 2, 3], name=None, dtype=pd.ArrowDtype(pa.int64())
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])
+
+
+def test_remove_with_ignore_key_error(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict.remove("k2", ignore_key_error=True)
+    context = SPECodeContext()
+    adapter = SeriesDictRemoveOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_remove(value):
+    row = list()
+    found = False
+    for tup in value:
+        if tup[0] == 'k2':
+            found = True
+        else:
+            row.append(tup)
+    if not found and not True:
+        raise KeyError('k2')
+    return row
+
+var_1 = var_0.map(_inner_remove, na_action="ignore").astype(const_0)
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, pd_df_1["A"])
+    expected_series = pd.Series(
+        [[("k1", 1)], [("k1", 11)], None],
+        index=[1, 2, 3],
+        name="A",
+        dtype=pd.ArrowDtype(pa.map_(pa.string(), pa.int32())),
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])
+
+
+@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
+def test_remove_with_key_error(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict.remove("k2", ignore_key_error=False)
+    context = SPECodeContext()
+    adapter = SeriesDictRemoveOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_remove(value):
+    row = list()
+    found = False
+    for tup in value:
+        if tup[0] == 'k2':
+            found = True
+        else:
+            row.append(tup)
+    if not found and not False:
+        raise KeyError('k2')
+    return row
+
+var_1 = var_0.map(_inner_remove, na_action="ignore").astype(const_0)
+"""
+    ]
+    assert results == expected_results
+    with pytest.raises(KeyError):
+        _run_generated_code(results[0], context, pd_df_1["A"])
+
+
+def test_contains(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].dict.contains("k2")
+    context = SPECodeContext()
+    adapter = SeriesDictContainsOperatorAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+var_1 = var_0.map(lambda x: any('k2' in tup[0] for tup in x), na_action="ignore").astype(const_0)
+var_1.name = None
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, pd_df_1["A"])
+    expected_series = pd.Series(
+        [True, False, None],
+        index=[1, 2, 3],
+        name=None,
+        dtype=pd.ArrowDtype(pa.bool_()),
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])