maxframe 1.3.0__cp38-cp38-win_amd64.whl → 2.0.0b1__cp38-cp38-win_amd64.whl

maxframe/codegen/spe/dataframe/datastore.py

@@ -0,0 +1,204 @@
+# 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 Dict, List, Optional
+
+import pyarrow as pa
+from odps import ODPS
+from odps.types import Column, OdpsSchema, PartitionSpec
+
+from ....config import options
+from ....dataframe.datastore.to_odps import DataFrameToODPSTable
+from ....io.odpsio import ODPSTableIO, pandas_to_arrow, pandas_to_odps_schema
+from ....protocol import DataFrameTableMeta, ResultInfo
+from ....typing_ import PandasObjectTypes
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+
+@register_op_adapter(DataFrameToODPSTable)
+class DataFrameToODPSTableAdapter(SPEOperatorAdapter):
+    @classmethod
+    def _get_index_mapping(
+        cls,
+        index: bool,
+        index_label: Optional[List[str]],
+        input_meta: DataFrameTableMeta,
+    ) -> Dict[str, str]:
+        if not index:
+            return {}
+        labels = DataFrameToODPSTable.get_index_mapping(
+            index_label, input_meta.pd_index_level_names
+        )
+        return {
+            src_col: idx_col
+            for src_col, idx_col in zip(input_meta.table_index_column_names, labels)
+        }
+
+    @classmethod
+    def _get_result_schema(
+        cls,
+        partition_spec: Optional[str],
+        partition_columns: Optional[List[str]],
+        src_schema: OdpsSchema,
+        exclude_cols: List[str],
+        index_mappings: Dict[str, str],
+    ) -> OdpsSchema:
+        part_spec = (
+            PartitionSpec(partition_spec) if partition_spec is not None else None
+        )
+        exclude_col_set = set(exclude_cols or ())
+
+        parts = []
+        if part_spec is not None:
+            parts.extend(Column(k, "string") for k in part_spec.keys())
+        if partition_columns is not None:
+            part_column_set = set(partition_columns)
+            cols = [
+                col
+                for col in src_schema.columns
+                if col.name not in part_column_set and col.name not in exclude_col_set
+            ]
+            parts.extend(
+                col for col in src_schema.columns if col.name in part_column_set
+            )
+        else:
+            cols = [
+                col for col in src_schema.columns if col.name not in exclude_col_set
+            ]
+        odps_schema = OdpsSchema(cols, parts)
+
+        for col in odps_schema.columns:
+            col.name = index_mappings.get(col.name, col.name)
+        return odps_schema
+
+    @classmethod
+    def _check_table_schema(cls, table_schema: OdpsSchema, data_schema: OdpsSchema):
+        table_cols = [col.name for col in table_schema.simple_columns]
+        data_cols = [col.name for col in data_schema.simple_columns]
+        if set(table_cols) != set(data_cols):
+            raise ValueError(
+                f"Source and dest columns does not match: "
+                f"{sorted(data_cols)} does not equal to {sorted(table_cols)}"
+            )
+        table_part_cols = [col.name for col in table_schema.partitions]
+        data_part_cols = [col.name for col in data_schema.partitions]
+        if set(table_part_cols) != set(data_part_cols):
+            raise ValueError(
+                f"Source and dest partitions does not match: "
+                f"{sorted(data_part_cols)} does not equal to {sorted(table_part_cols)}"
+            )
+
+    @classmethod
+    def _write_as_table(
+        cls,
+        df_obj: PandasObjectTypes,
+        table_name: str,
+        partition_spec: Optional[str] = None,
+        partition_columns: Optional[List[str]] = None,
+        index: bool = True,
+        index_label: Optional[List[str]] = None,
+        overwrite: bool = True,
+        lifecycle: Optional[int] = None,
+        table_properties: Optional[Dict] = None,
+    ):
+        partition_spec = (
+            PartitionSpec(partition_spec) if partition_spec is not None else None
+        )
+        odps_schema, odps_meta = pandas_to_odps_schema(df_obj, unknown_as_string=True)
+        index_mapping = cls._get_index_mapping(index, index_label, odps_meta)
+        exclude_cols = [] if index else odps_meta.table_index_column_names
+        result_schema = cls._get_result_schema(
+            partition_spec, partition_columns, odps_schema, exclude_cols, index_mapping
+        )
+
+        odps_entry = ODPS.from_environments() or ODPS.from_global()
+        if not odps_entry.exist_table(table_name):
+            odps_entry.create_table(
+                table_name,
+                result_schema,
+                lifecycle=lifecycle,
+                table_properties=table_properties,
+            )
+            part_cols = [col.name for col in result_schema.partitions]
+        else:
+            table = odps_entry.get_table(table_name)
+            cls._check_table_schema(table.table_schema, result_schema)
+            result_schema = table.table_schema
+            part_cols = [col.name for col in table.table_schema.partitions]
+
+        part_and_group = []
+        if partition_columns is not None:
+            groupby_obj = df_obj.groupby(partition_columns)
+            for group in groupby_obj.groups:
+                group_data = groupby_obj.get_group(group)
+                group_data = group_data.loc[
+                    :, ~group_data.columns.isin(partition_columns)
+                ]
+
+                group = (group,) if len(partition_columns) == 1 else group
+                part_dict = dict(zip(partition_columns, group))
+
+                part_dict.update(dict(partition_spec.items() if partition_spec else {}))
+                part_spec = ",".join(f"{col}={part_dict[col]}" for col in part_cols)
+                part_and_group.append((part_spec, group_data))
+        else:
+            part_and_group.append((partition_spec, df_obj))
+
+        for part, data in part_and_group:
+            if part is not None:
+                odps_entry.get_table(table_name).create_partition(
+                    part, if_not_exists=True
+                )
+
+            with ODPSTableIO(odps_entry).open_writer(
+                table_name, part, overwrite=overwrite
+            ) as writer:
+                arrow_val, _ = pandas_to_arrow(data)
+                data_array_dict = dict(zip(arrow_val.column_names, arrow_val.columns))
+                cols, arrays = [], []
+                col_name_to_src_name = {v: k for k, v in index_mapping.items()}
+                for col in result_schema.simple_columns:
+                    cols.append(col.name)
+                    arrays.append(
+                        data_array_dict[col_name_to_src_name.get(col.name, col.name)]
+                    )
+                writer.write(pa.Table.from_arrays(arrays, cols))
+
+    def generate_code(
+        self, op: DataFrameToODPSTable, context: SPECodeContext
+    ) -> List[str]:
+        # to_odps_table should return None
+        context.put_tileable_result_info(op.outputs[0], ResultInfo())
+
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        context.register_import(__name__, from_item=type(self).__name__)
+
+        args = self._translate_call_args(
+            context,
+            op.table_name,
+            op.partition_spec,
+            op.partition_columns,
+            op.index,
+            op.index_label,
+            overwrite=op.overwrite,
+            lifecycle=op.lifecycle or options.session.table_lifecycle,
+            table_properties=op.table_properties,
+        )
+        return [
+            f"{type(self).__name__}._write_as_table("
+            f"{input_var_name}, {', '.join(args)})",
+            f"{res_var_name} = None",
+        ]

maxframe/codegen/spe/dataframe/extensions.py

@@ -0,0 +1,63 @@
+# 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.extensions import DataFrameApplyChunk, DataFrameReshuffle
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+
+@register_op_adapter(DataFrameReshuffle)
+class DataFrameReshuffleAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: DataFrameReshuffle, context: SPECodeContext
+    ) -> List[str]:
+        input_var = context.get_input_tileable_variable(op.inputs[0])
+        if not op.group_by:
+            kwargs = {}
+            if op.ignore_index:
+                kwargs["ignore_index"] = op.ignore_index
+            args = self._translate_call_args(context, frac=1, **kwargs)
+            shuffled = f"{input_var}.sample({', '.join(args)})"
+        else:
+            sort_keys = list(op.group_by) + list(op.sort_by or [])
+            kwargs = {"ascending": op.ascending}
+            if op.ignore_index:
+                kwargs["ignore_index"] = op.ignore_index
+            args = self._translate_call_args(context, sort_keys, **kwargs)
+            shuffled = f"{input_var}.sort_values({', '.join(args)})"
+        output_var = context.get_output_tileable_variable(op.outputs[0])
+        return [f"{output_var} = {shuffled}"]
+
+
+@register_op_adapter(DataFrameApplyChunk)
+class DataFrameApplyChunkAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: DataFrameApplyChunk, context: SPECodeContext
+    ) -> List[str]:
+        context.register_import("pandas", "pd")
+        input_var = context.get_input_tileable_variable(op.inputs[0])
+        output_var = context.get_output_tileable_variable(op.outputs[0])
+        func_var = self.translate_var(context, op.func)
+        if not op.batch_rows:
+            return [f"{output_var} = {func_var}({input_var})"]
+        else:
+            return [
+                f"{input_var}_batches = [{func_var}({input_var}.iloc[i : i + {op.batch_rows}]) "
+                f"for i in range(0, len({input_var}), {op.batch_rows})]",
+                f"if len({input_var}_batches) > 1:",
+                f"    {output_var} = pd.concat({input_var}_batches)",
+                f"else:",
+                f"    {output_var} = {input_var}_batches[0]",
+            ]

maxframe/codegen/spe/dataframe/fetch.py

@@ -0,0 +1,26 @@
+# 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.fetch import DataFrameFetch
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+
+@register_op_adapter(DataFrameFetch)
+class DataFrameFetchAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: DataFrameFetch, context: SPECodeContext) -> List[str]:
+        # Simply register Fetch as SPE-executable.
+        # Actual codegen is done inside code generator itself.
+        return []

maxframe/codegen/spe/dataframe/groupby.py

@@ -0,0 +1,224 @@
+# 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
+
+import pandas as pd
+
+from ....dataframe.groupby.aggregation import DataFrameGroupByAgg
+from ....dataframe.groupby.apply import GroupByApply
+from ....dataframe.groupby.apply_chunk import GroupByApplyChunk
+from ....dataframe.groupby.core import DataFrameGroupByOp
+from ....dataframe.groupby.cum import GroupByCumcount, GroupByCumReductionOperator
+from ....dataframe.groupby.fill import GroupByFill
+from ....dataframe.groupby.getitem import GroupByIndex
+from ....dataframe.groupby.head import GroupByHead
+from ....dataframe.groupby.sample import GroupBySample
+from ....dataframe.groupby.transform import GroupByTransform
+from ....dataframe.utils import make_column_list
+from ....lib.version import parse as parse_version
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+from ..utils import build_method_call_adapter
+
+_need_enforce_group_keys = parse_version(pd.__version__) < parse_version("1.5.0")
+
+
+class SPEGroupByOperatorAdapter(SPEOperatorAdapter):
+    def build_groupby_call(
+        self, context: SPECodeContext, groupby_params: Dict[str, Any]
+    ) -> str:
+        groupby_params = {k: v for k, v in groupby_params.items() if v is not None}
+        selection = groupby_params.pop("selection", None)
+        by_args_list = self._translate_call_args(context, **groupby_params)
+        groupby_args = ", ".join(by_args_list)
+        groupby_call = f"groupby({groupby_args})"
+        if selection is not None:
+            sel_str = ", ".join(repr(s) for s in selection)
+            groupby_call += f"[[{sel_str}]]"
+        return groupby_call
+
+
+@register_op_adapter(DataFrameGroupByAgg)
+class DataFrameGroupByAggAdapter(SPEGroupByOperatorAdapter):
+    def generate_code(
+        self, op: DataFrameGroupByAgg, context: SPECodeContext
+    ) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        groupby_call = self.build_groupby_call(context, op.groupby_params)
+        agg_func_str = self.translate_var(context, op.raw_func)
+
+        return [f"{res_var_name} = {input_var_name}.{groupby_call}.agg({agg_func_str})"]
+
+
+DataFrameGroupByOpAdapter = build_method_call_adapter(
+    DataFrameGroupByOp,
+    "groupby",
+    kw_keys=["by", "level", "as_index", "sort", "group_keys"],
+)
+
+
+@register_op_adapter(GroupByApply)
+class GroupByApplyAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: GroupByApply, context: SPECodeContext) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        args_list = self._translate_call_args(context, op.func, *op.args, **op.kwds)
+        args = ", ".join(args_list)
+        return [f"{res_var_name} = {input_var_name}.apply({args})"]
+
+
+@register_op_adapter(GroupByApplyChunk)
+class GroupByApplyChunkAdapter(SPEGroupByOperatorAdapter):
+    def generate_code(
+        self, op: GroupByApplyChunk, context: SPECodeContext
+    ) -> List[str]:
+        input_var = context.get_input_tileable_variable(op.inputs[0])
+        output_var = context.get_output_tileable_variable(op.outputs[0])
+        func_var = self.translate_var(context, op.func)
+
+        by_cols = make_column_list(op.groupby_params.get("by", []), op.inputs[0].dtypes)
+
+        need_group_keys = False
+        if isinstance(op.inputs[0].op, (GroupByIndex, DataFrameGroupByOp)):
+            input_obj_code = input_var
+        else:
+            assert op.groupby_params
+            groupby_params = (op.groupby_params or {}).copy()
+            if op.groupby_params.get("selection") is None:
+                by_col_set = set(by_cols)
+                sel = [
+                    col for col in op.inputs[0].dtypes.index if col not in by_col_set
+                ]
+
+                if _need_enforce_group_keys and groupby_params.get("group_keys"):
+                    need_group_keys = True
+                    sel.extend(by_cols)
+
+                groupby_params["selection"] = sel
+            groupby_call = self.build_groupby_call(context, groupby_params)
+            input_obj_code = f"{input_var}.{groupby_call}"
+
+        args_list = self._translate_call_args(context, *op.args, **op.kwargs)
+        args = ", ".join(args_list)
+        if args:
+            args = f", {args}"
+
+        if not op.batch_rows:
+            lines = [f"{output_var} = {input_obj_code}.apply({func_var}{args})"]
+        else:
+            lines = [
+                f"def batch_apply_{input_var}(frame, *args, **kw):",
+                f"    batches = [{func_var}(frame.iloc[i : i + {op.batch_rows}], *args, **kw) "
+                f"    for i in range(0, len(frame), {op.batch_rows})]",
+                f"    if len(batches) > 1:",
+                f"        return pd.concat(batches)",
+                f"    else:",
+                f"        return batches[0]",
+                f"{output_var} = {input_obj_code}.apply(batch_apply_{input_var}{args})",
+            ]
+        if need_group_keys:
+            by_cols_str = ", ".join(repr(col) for col in by_cols)
+            lines.append(
+                f"{output_var} = {output_var}.set_index("
+                f"[{by_cols_str}, {output_var}.index])",
+            )
+        return lines
+
+
+@register_op_adapter(GroupByCumReductionOperator)
+class GroupByCumReductionAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: GroupByCumReductionOperator, context: SPECodeContext
+    ) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        if isinstance(op, GroupByCumcount):
+            kw_keys = ["ascending"]
+        else:
+            kw_keys = ["axis"]
+
+        args = self.generate_call_args_with_attributes(
+            op, context, kw_keys=kw_keys, skip_none=True
+        )
+        return [f"{res_var_name} = {input_var_name}.{op._func_name}({args})"]
+
+
+@register_op_adapter(GroupByFill)
+class GroupByFillOperatorAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: GroupByFill, context: SPECodeContext) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        kw_keys = ["method", "axis", "limit", "downcast"]
+        if op._func_name == "fillna":
+            kw_args = ["value"]
+        else:
+            kw_args = []
+        args = self.generate_call_args_with_attributes(
+            op, context, *kw_args, kw_keys=kw_keys, skip_none=True
+        )
+        return [f"{res_var_name} = {input_var_name}.{op._func_name}({args})"]
+
+
+@register_op_adapter(GroupByIndex)
+class GroupByIndexAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: GroupByIndex, context: SPECodeContext) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        selection = self.translate_var(context, op.selection)
+        return [f"{res_var_name} = {input_var_name}[{selection}]"]
+
+
+@register_op_adapter(GroupByHead)
+class GroupByHeadAdapter(SPEGroupByOperatorAdapter):
+    def generate_code(self, op: GroupByHead, context: SPECodeContext) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        groupby_call = self.build_groupby_call(context, op.groupby_params)
+        args = self.generate_call_args_with_attributes(op, context, "row_count")
+        return [f"{res_var_name} = {input_var_name}.{groupby_call}.head({args})"]
+
+
+@register_op_adapter(GroupBySample)
+class GroupBySampleAdapter(SPEGroupByOperatorAdapter):
+    def generate_code(self, op: GroupBySample, context: SPECodeContext) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        groupby_call = self.build_groupby_call(context, op.groupby_params)
+        args = self.generate_call_args_with_attributes(
+            op,
+            context,
+            size="n",
+            kw_keys=["frac", "replace", "weights"],
+            skip_none=True,
+        )
+        return [f"{res_var_name} = {input_var_name}.{groupby_call}.agg({args})"]
+
+
+@register_op_adapter(GroupByTransform)
+class GroupByTransformAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: GroupByTransform, context: SPECodeContext) -> List[str]:
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+
+        args_list = self._translate_call_args(context, op.func, *op.args, **op.kwds)
+        args = ", ".join(args_list)
+        return [f"{res_var_name} = {input_var_name}.transform({args})"]