maxframe 2.4.0rc1__cp312-cp312-win32.whl

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

@@ -0,0 +1,75 @@
+# 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 pandas as pd
+import pyarrow as pa
+import pytest
+
+from ...... import dataframe as md
+from ......utils import ARROW_DTYPE_NOT_SUPPORTED
+from ....core import SPECodeContext
+from ...accessors.struct_ import SeriesStructMethodAdapter
+
+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 series():
+    version_type = pa.struct(
+        [
+            ("major", pa.int64()),
+            ("minor", pa.int64()),
+        ]
+    )
+    return md.Series(
+        [
+            {"version": {"major": 1, "minor": 5}, "project": "pandas"},
+            {"version": {"major": 2, "minor": 1}, "project": "pandas"},
+            {"version": {"major": 1, "minor": 26}, "project": "numpy"},
+        ],
+        dtype=pd.ArrowDtype(
+            pa.struct([("version", version_type), ("project", pa.string())])
+        ),
+    )
+
+
+def test_field(series):
+    s1 = series.struct.field(["version", "minor"])
+    context = SPECodeContext()
+    adapter = SeriesStructMethodAdapter()
+    results = adapter.generate_code(s1.op, context)
+
+    expected_results = [
+        """
+var_1 = var_0.struct.field(['version', 'minor'])
+"""
+    ]
+    assert results == expected_results
+    local_vars = _run_generated_code(results[0], context, series.op.data)
+    expected_series = pd.Series(
+        [5, 1, 26], dtype=pd.ArrowDtype(pa.int64()), name="minor"
+    )
+    pd.testing.assert_series_equal(expected_series, local_vars["var_1"])

maxframe/codegen/spe/dataframe/tests/indexing/__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/indexing/conftest.py

@@ -0,0 +1,58 @@
+# 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 pytest
+
+from ...... import dataframe as md
+
+
+@pytest.fixture
+def multi_index():
+    arrays = [["m1", "m2", "m3"], ["n1", "n2", "n3"]]
+    return pd.MultiIndex.from_arrays(arrays, names=("idx_1", "idx_2"))
+
+
+@pytest.fixture
+def df1():
+    return md.DataFrame(
+        np.random.randn(6, 4),
+        index=md.Index(["1", "2", "3", "4", "5", "6"], name="test_idx"),
+        columns=list("ABCD"),
+    )
+
+
+@pytest.fixture
+def df2():
+    return md.DataFrame(
+        np.random.randn(6, 4),
+        index=md.Index(["6", "5", "4", "3", "2", "1"], name="test_idx"),
+        columns=list("ABCD"),
+    )
+
+
+@pytest.fixture
+def s1():
+    return md.Series([1, 2, 3], index=list("abc"))
+
+
+@pytest.fixture()
+def s2(multi_index):
+    return md.Series([1, 2, 3], index=multi_index)
+
+
+@pytest.fixture
+def idx1(multi_index):
+    return md.Index(multi_index)

maxframe/codegen/spe/dataframe/tests/indexing/test_getitem.py

@@ -0,0 +1,124 @@
+# 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 pytest
+
+from .... import SPECodeGenerator
+from ....core import SPECodeContext
+from ...indexing import DataFrameIndexAdapter, SeriesIndexAdapter
+
+
+@pytest.mark.skip(reason="Mars bug: https://github.com/mars-project/mars/issues/3356")
+def test_series_index_with_selecting_attributes(s1):
+    s = s1.b
+    results = SeriesIndexAdapter().generate_code(s.op, SPECodeContext())
+    expected_results = []
+    assert results == expected_results
+
+
+@pytest.mark.skip(reason="Mars bug: https://github.com/mars-project/mars/issues/3357")
+def test_series_index_with_multi_index_and_selecting_single_row(s2):
+    s = s2[("m1", "n1")]
+    results = SeriesIndexAdapter().generate_code(s.op, SPECodeContext())
+    assert results
+
+
+@pytest.mark.skip(reason="Mars bug: https://github.com/mars-project/mars/issues/3357")
+def test_series_index_with_multi_index_and_selecting_multi_row(s2):
+    s = s2[("m1", "n1")]
+    results = SeriesIndexAdapter().generate_code(s.op, SPECodeContext())
+    assert results
+
+
+def test_series_index_with_selecting_single_row(s1):
+    s = s1["b"]
+    results = SeriesIndexAdapter().generate_code(s.op, SPECodeContext())
+    expected_results = ["var_1 = var_0['b']"]
+    assert results == expected_results
+
+
+def test_series_index_with_selecting_multi_row(s1):
+    s = s1[["a", "b"]]
+    context = SPECodeContext()
+    results = SeriesIndexAdapter().generate_code(s.op, context)
+    expected_results = ["var_1 = var_0[['a', 'b']]"]
+    assert results == expected_results
+
+
+def test_dataframe_index_with_selecting_attributes(df1):
+    df = df1.A
+    results = DataFrameIndexAdapter().generate_code(df.op, SPECodeContext())
+    expected_results = ["var_1 = var_0['A']"]
+    assert results == expected_results
+
+
+def test_dataframe_index_with_selecting_column(df1):
+    df = df1["B"]
+    results = DataFrameIndexAdapter().generate_code(df.op, SPECodeContext())
+    expected_results = ["var_1 = var_0['B']"]
+    assert results == expected_results
+
+
+def test_dataframe_index_with_simple_filter(df1):
+    v1 = df1["A"]
+    v2 = v1 > 0
+    v3 = v1 + 1
+    v4 = v3 > 0
+    v5 = v2 & v4
+    v6 = df1["B"]
+    v7 = v1 + v6
+    v8 = df1["C"]
+    v9 = v7 > v8
+    v10 = v5 | v9
+    # df1[(df1["A"] > 0) & (df1["A"] + 1 > 0) & (df1["A"] + df1["B"] > df1["C"])]
+    df = df1[v10]
+    graph = df.build_graph()
+    generator = SPECodeGenerator("1")
+    context = SPECodeContext()
+    results = []
+    skipped = True
+    for v in graph.topological_iter():
+        if skipped:
+            if v.key == df1.key:
+                skipped = False
+            continue
+        r = generator.get_op_adapter(type(v.op))().generate_code(v.op, context)
+        results.extend(r)
+    out_vars = [
+        context.get_output_tileable_variable(v.op.outputs[0])
+        for v in [df1, v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, df]
+    ]
+    idx1 = results.index(f"{out_vars[1]} = {out_vars[0]}['A']")
+    idx2 = results.index(f"{out_vars[2]} = {out_vars[1]}.gt(0)")
+    idx3 = results.index(f"{out_vars[3]} = {out_vars[1]}.add(1)")
+    idx4 = results.index(f"{out_vars[4]} = {out_vars[3]}.gt(0)")
+    idx5 = results.index(f"{out_vars[5]} = {out_vars[2]} & {out_vars[4]}")
+    idx6 = results.index(f"{out_vars[6]} = {out_vars[0]}['B']")
+    idx7 = results.index(f"{out_vars[7]} = {out_vars[1]}.add({out_vars[6]})")
+    idx8 = results.index(f"{out_vars[8]} = {out_vars[0]}['C']")
+    idx9 = results.index(f"{out_vars[9]} = {out_vars[7]}.gt({out_vars[8]})")
+    idx10 = results.index(f"{out_vars[10]} = {out_vars[5]} | {out_vars[9]}")
+    idx_out = results.index(f"{out_vars[11]} = {out_vars[0]}[{out_vars[10]}]")
+    assert idx1 < idx2
+    assert idx1 < idx3
+    assert idx3 < idx4
+    assert idx2 < idx5
+    assert idx4 < idx5
+    assert idx1 < idx7
+    assert idx6 < idx7
+    assert idx7 < idx9
+    assert idx8 < idx9
+    assert idx5 < idx10
+    assert idx9 < idx10
+    assert idx10 < idx_out

maxframe/codegen/spe/dataframe/tests/indexing/test_iloc.py

@@ -0,0 +1,95 @@
+# 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 ....core import SPECodeContext
+from ...indexing import DataFrameIlocGetItemAdapter, DataFrameIlocSetItemAdapter
+
+
+def test_series_iloc_single_row(s1):
+    df = s1.iloc[2]
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = ["var_1 = var_0.iloc[2]"]
+    assert results == expected_results
+
+
+def test_series_iloc_set_single_row(s1):
+    s1.iloc[2] = 2
+    context = SPECodeContext()
+    results = DataFrameIlocSetItemAdapter().generate_code(s1.op, context)
+    expected_results = ["var_1 = var_0.copy()", "var_1.iloc[2] = 2"]
+    assert results == expected_results
+
+
+def test_series_iloc_multi_rows(s1):
+    df = s1.iloc[:2]
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = ["var_1 = var_0.iloc[slice(None, 2, None)]"]
+    assert results == expected_results
+
+
+def test_series_head(s1):
+    df = s1.head(-3)
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = ["var_1 = var_0.iloc[slice(0, -3, None)]"]
+    assert results == expected_results
+
+
+def test_dataframe_iloc_only_columns(df1):
+    df = df1.iloc[:, [0, 1]]
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = ["var_1 = var_0.iloc[slice(None, None, None), [0, 1]]"]
+    assert results == expected_results
+
+
+def test_dataframe_iloc_set_only_column(df1):
+    df1.iloc[:, 0] = 5
+    context = SPECodeContext()
+    results = DataFrameIlocSetItemAdapter().generate_code(df1.op, context)
+    expected_results = [
+        "var_1 = var_0.copy()",
+        "var_1.iloc[slice(None, None, None), 0] = 5",
+    ]
+    assert results == expected_results
+
+
+def test_dataframe_iloc_only_rows(df1):
+    df = df1.iloc[1]
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = ["var_1 = var_0.iloc[1, slice(None, None, None)]"]
+    assert results == expected_results
+
+
+def test_dataframe_iloc_rows_and_columns(df1):
+    df = df1.iloc[1:4:2, [True, False, False, True]]
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = [
+        "var_1 = var_0.iloc[slice(1, 4, 2), [True, False, False, True]]"
+    ]
+    assert results == expected_results
+
+
+def test_dataframe_head(df1):
+    df = df1.head(3)
+    context = SPECodeContext()
+    results = DataFrameIlocGetItemAdapter().generate_code(df.op, context)
+    expected_results = [
+        "var_1 = var_0.iloc[slice(0, 3, None), slice(None, None, None)]"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/dataframe/tests/indexing/test_indexing.py

@@ -0,0 +1,39 @@
+# 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 ....core import SPECodeContext
+from ...indexing import DataFrameAlignAdapter, DataFrameInsertAdapter
+
+
+def test_dataframe_align(df1, df2):
+    context = SPECodeContext()
+    adapter = DataFrameAlignAdapter()
+
+    aligned1, aligned2 = df1.align(df2, join="outer")
+    results = adapter.generate_code(aligned1.op, context)
+    expected_results = ["var_1, var_2 = var_0.align(var_3, join='outer', fill_axis=0)"]
+    assert results == expected_results
+
+
+def test_dataframe_insert(df1, df2):
+    context = SPECodeContext()
+    adapter = DataFrameInsertAdapter()
+
+    df1.insert(1, "newcol", [4, 9, 2, 3, 5, 7])
+    results = adapter.generate_code(df1.op, context)
+    expected_results = [
+        "var_1 = var_0.insert(loc=1, column='newcol', "
+        "value=[4, 9, 2, 3, 5, 7], allow_duplicates=False)"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/dataframe/tests/indexing/test_loc.py

@@ -0,0 +1,35 @@
+# 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 ....core import SPECodeContext
+from ...indexing import DataFrameLocGetItemAdapter, DataFrameLocSetItemAdapter
+
+
+def test_dataframe_iloc_only_columns(df1):
+    df = df1.iloc[:, [0, 1]]
+    context = SPECodeContext()
+    results = DataFrameLocGetItemAdapter().generate_code(df.op, context)
+    expected_results = ["var_1 = var_0.loc[slice(None, None, None), [0, 1]]"]
+    assert results == expected_results
+
+
+def test_dataframe_iloc_set_only_column(df1):
+    df1.iloc[:, 0] = 5
+    context = SPECodeContext()
+    results = DataFrameLocSetItemAdapter().generate_code(df1.op, context)
+    expected_results = [
+        "var_1 = var_0.copy()",
+        "var_1.loc[slice(None, None, None), 0] = 5",
+    ]
+    assert results == expected_results

maxframe/codegen/spe/dataframe/tests/indexing/test_rename.py

@@ -0,0 +1,51 @@
+# 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 ....core import SPECodeContext
+from ...indexing import DataFrameRenameAdapter
+
+
+def test_dataframe_rename(df1):
+    context = SPECodeContext()
+    adapter = DataFrameRenameAdapter()
+
+    df1 = df1.rename(columns={"A": "x"})
+    results = adapter.generate_code(df1.op, context)
+    expected_results = ["var_1 = var_0.rename(columns={'A': 'x'}, errors='ignore')"]
+    assert results == expected_results
+
+    df1 = df1.rename(index=lambda x: x + "_")
+    results = adapter.generate_code(df1.op, context)
+    expected_results = [
+        f"var_2 = var_1.rename(index={context.get_udfs()[0].name}, " "errors='ignore')"
+    ]
+    assert results == expected_results
+
+
+def test_series_rename(s1):
+    s1.name = "series_name"
+    context = SPECodeContext()
+    adapter = DataFrameRenameAdapter()
+    results = adapter.generate_code(s1.op, context)
+    expected_results = ["var_1 = var_0.rename(index='series_name')"]
+    assert results == expected_results
+
+
+def test_index_rename(idx1):
+    idx1.set_names("renamed", level=0, inplace=True)
+    context = SPECodeContext()
+    adapter = DataFrameRenameAdapter()
+    results = adapter.generate_code(idx1.op, context)
+    expected_results = ["var_1 = var_0.set_names('renamed', level=0)"]
+    assert results == expected_results

maxframe/codegen/spe/dataframe/tests/indexing/test_reset_index.py

@@ -0,0 +1,88 @@
+# 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 ....core import SPECodeContext
+from ...indexing import DataFrameResetIndexAdapter
+
+
+def test_dataframe_reset_and_drop(df1):
+    # index
+    df1 = df1.reset_index(drop=True)
+    context = SPECodeContext()
+    adapter = DataFrameResetIndexAdapter()
+    results = adapter.generate_code(df1.op, context)
+    expected_results = [
+        "var_1 = var_0.reset_index(level=None, drop=True, inplace=False, col_level=0, col_fill='')"
+    ]
+    assert results == expected_results
+
+
+def test_series_reset_and_drop(s1):
+    # index
+    s1 = s1.reset_index(drop=True)
+    context = SPECodeContext()
+    adapter = DataFrameResetIndexAdapter()
+    results = adapter.generate_code(s1.op, context)
+    expected_results = [
+        "var_1 = var_0.reset_index(level=None, drop=True, inplace=False, name=0)"
+    ]
+    assert results == expected_results
+
+
+def test_series_reset_and_name(s1):
+    # index
+    s1 = s1.reset_index(name="gen_by_reset")
+    context = SPECodeContext()
+    adapter = DataFrameResetIndexAdapter()
+    results = adapter.generate_code(s1.op, context)
+    expected_results = [
+        "var_1 = var_0.reset_index(level=None, drop=False, inplace=False, name='gen_by_reset')"
+    ]
+    assert results == expected_results
+
+
+def test_series_reset_and_name_and_drop(s1):
+    # index
+    s1 = s1.reset_index(name="gen_by_reset", drop=True)
+    context = SPECodeContext()
+    adapter = DataFrameResetIndexAdapter()
+    results = adapter.generate_code(s1.op, context)
+    # should ignore `name`
+    expected_results = [
+        "var_1 = var_0.reset_index(level=None, drop=True, inplace=False, name='gen_by_reset')"
+    ]
+    assert results == expected_results
+
+
+def test_dataframe_reset(df1):
+    df1 = df1.reset_index()
+    context = SPECodeContext()
+    adapter = DataFrameResetIndexAdapter()
+    results = adapter.generate_code(df1.op, context)
+    expected_results = [
+        "var_1 = var_0.reset_index(level=None, drop=False, inplace=False, col_level=0, col_fill='')"
+    ]
+    assert results == expected_results
+
+
+def test_dataframe_reset_index_multi_level_column(df2):
+    col_fill = "fill_auto_gen"
+    reset_df2 = df2.reset_index(level=["test_idx"], col_fill=col_fill, col_level=1)
+    context = SPECodeContext()
+    adapter = DataFrameResetIndexAdapter()
+    results = adapter.generate_code(reset_df2.op, context)
+    expected_results = [
+        "var_1 = var_0.reset_index(level=['test_idx'], drop=False, inplace=False, col_level=1, col_fill='fill_auto_gen')"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/dataframe/tests/indexing/test_sample.py

@@ -0,0 +1,45 @@
+# 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
+
+from ....core import SPECodeContext
+from ...indexing import DataFrameSampleAdapter
+
+
+def test_dataframe_sample_size(s1):
+    df = s1.sample(n=2)
+    adapter = DataFrameSampleAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(df.op, context)
+    assert results[0] == "var_1_rs = np.random.RandomState()"
+    assert results[1].startswith("var_1_rs.set_state(")
+    assert results[2] == (
+        "var_1 = var_0.sample(n=2, frac=None, replace=False, weights=None,"
+        " axis=0, random_state=var_1_rs)"
+    )
+
+
+def test_dataframe_sample_frac(s1):
+    s1._shape = (np.NAN, np.NAN)
+    df = s1.sample(frac=1.1, weights=[0.1, 0.2, 0.3, 0.4], replace=True, random_state=1)
+    adapter = DataFrameSampleAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(df.op, context)
+    assert results[0] == "var_1_rs = np.random.RandomState()"
+    assert results[1].startswith("var_1_rs.set_state(")
+    assert results[2] == (
+        "var_1 = var_0.sample(n=None, frac=1.1, replace=True,"
+        " weights=[0.1, 0.2, 0.3, 0.4], axis=0, random_state=var_1_rs)"
+    )

maxframe/codegen/spe/dataframe/tests/indexing/test_set_axis.py

@@ -0,0 +1,45 @@
+# 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
+
+from ....core import SPECodeContext
+from ...indexing import DataFrameSetAxisAdapter
+
+
+def test_dataframe_set_axis_with_np_index(df1):
+    df1.index = np.arange(1, 7)
+    context = SPECodeContext()
+    adapter = DataFrameSetAxisAdapter()
+    results = adapter.generate_code(df1.op, context)
+    expected_results = ["var_0 = var_0.set_axis(const_0, axis=0)"]
+    assert results == expected_results
+
+
+def test_dataframe_set_axis_with_array(df1):
+    df1.columns = list("WXYZ")
+    context = SPECodeContext()
+    adapter = DataFrameSetAxisAdapter()
+    results = adapter.generate_code(df1.op, context)
+    expected_results = ["var_0 = var_0.set_axis(const_0, axis=1)"]
+    assert results == expected_results
+
+
+def test_series_set_axis_with_np_index(s1):
+    s1.index = np.arange(1, 4)
+    context = SPECodeContext()
+    adapter = DataFrameSetAxisAdapter()
+    results = adapter.generate_code(s1.op, context)
+    expected_results = ["var_0 = var_0.set_axis(const_0, axis=0)"]
+    assert results == expected_results