maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/codegen/spe/dataframe/sort.py

@@ -0,0 +1,104 @@
+# 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 DataFrameRank, DataFrameSortIndex, DataFrameSortValues
+from ....dataframe.sort.core import DataFrameSortOperator
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+from ..utils import build_method_call_adapter
+
+DataFrameRankAdapter = build_method_call_adapter(
+    DataFrameRank,
+    "rank",
+    kw_keys=["axis", "method", "numeric_only", "na_option", "ascending", "pct"],
+    skip_none=True,
+)
+
+
+class _DataFrameSortAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: DataFrameSortOperator, context: SPECodeContext
+    ) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        if op.keep_kind == "head" or not op.nrows:
+            args = []
+            kwargs = {
+                "axis": op.axis,
+                "ascending": op.ascending,
+                "inplace": op.inplace,
+                "na_position": op.na_position,
+                "ignore_index": op.ignore_index,
+            }
+        else:
+            args = [op.nrows]
+            kwargs = {"keep": op.keep_kind}
+        kwargs.update(self.extra_args(op, context))
+        args_str = ", ".join(self._translate_call_args(context, *args, **kwargs))
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        ret = [
+            f"{res_var_name} = {input_var_name}.{self.get_method_name(op)}({args_str})"
+        ]
+        if op.keep_kind == "head" and op.nrows:
+            ret.append(f"{res_var_name} = {res_var_name}.head({op.nrows})")
+        return ret
+
+    @abstractmethod
+    def extra_args(
+        self, op: DataFrameSortOperator, context: SPECodeContext
+    ) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_method_name(self, op: DataFrameSortOperator) -> str:
+        raise NotImplementedError
+
+
+@register_op_adapter(DataFrameSortIndex)
+class DataFrameSortIndexAdapter(_DataFrameSortAdapter):
+    _method_name = "sort_index"
+
+    def get_method_name(self, op: DataFrameSortIndex) -> 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"
+
+    def get_method_name(self, op: DataFrameSortValues) -> str:
+        if op.keep_kind == "head" or not op.nrows:
+            return self._method_name
+        else:
+            return "nsmallest" if op.ascending else "nlargest"
+
+    def extra_args(
+        self, op: DataFrameSortValues, context: SPECodeContext
+    ) -> Dict[str, Any]:
+        if isinstance(op.outputs[0], DATAFRAME_TYPE):
+            if op.keep_kind == "head" or not op.nrows:
+                return {"by": op.by}
+            else:
+                return {"columns": 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,304 @@
+# 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 SeriesDictMethodAdapter
+
+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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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 = SeriesDictMethodAdapter()
+    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"])

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

@@ -0,0 +1,134 @@
+# 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 list_
+from ......utils import ARROW_DTYPE_NOT_SUPPORTED
+from ....core import SPECodeContext
+from ...accessors.list_ import SeriesListMethodAdapter
+
+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(
+                [[1, 2, 3], [4, 5, 6, 7], None],
+                index=[1, 2, 3],
+                dtype=list_(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"].list.get(0)
+    context = SPECodeContext()
+    adapter = SeriesListMethodAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_get(data):
+    try:
+        return data[0]
+    except IndexError:
+        if True:
+            return None
+        else:
+            raise
+
+var_1 = var_0.map(_inner_get, 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(
+        [1, 4, None],
+        index=[1, 2, 3],
+        dtype=pd.ArrowDtype(pa.int32()),
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])
+
+
+def test_getitem_with_index_error(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].list[3]
+    context = SPECodeContext()
+    adapter = SeriesListMethodAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+def _inner_get(data):
+    try:
+        return data[3]
+    except IndexError:
+        if False:
+            return None
+        else:
+            raise
+
+var_1 = var_0.map(_inner_get, na_action="ignore").astype(const_0)
+var_1.name = None
+"""
+    ]
+    assert results == expected_results
+    with pytest.raises(IndexError):
+        _run_generated_code(results[0], context, pd_df_1["A"])
+
+
+def test_length(md_df_1, pd_df_1):
+    s1 = md_df_1["A"].list.len()
+    context = SPECodeContext()
+    adapter = SeriesListMethodAdapter()
+    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(
+        [3, 4, pd.NA], index=[1, 2, 3], name=None, dtype=pd.ArrowDtype(pa.int64())
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])