maxframe 1.2.0__cp39-cp39-macosx_10_9_universal2.whl → 1.3.0__cp39-cp39-macosx_10_9_universal2.whl

maxframe/codegen.py

@@ -24,7 +24,7 @@ from odps.types import OdpsSchema
 from odps.utils import camel_to_underline
 
 from .core import OperatorType, Tileable, TileableGraph
-from .core.operator import Fetch
+from .core.operator import Fetch, Operator
 from .extension import iter_extensions
 from .io.odpsio import build_dataframe_table_meta
 from .io.odpsio.schema import pandas_to_odps_schema
@@ -211,7 +211,21 @@ class BigDagCodeContext(metaclass=abc.ABCMeta):
     def get_udfs(self) -> List[AbstractUDF]:
         return list(self._udfs.values())
 
-    def get_tileable_variable(self, tileable: Tileable) -> str:
+    def get_input_tileable_variable(self, tileable: Tileable) -> str:
+        """
+        Get or create the variable name for an input tileable. It should be used on the
+        RIGHT side of the assignment.
+        """
+        return self._get_tileable_variable(tileable)
+
+    def get_output_tileable_variable(self, tileable: Tileable) -> str:
+        """
+        Get or create the variable name for an output tileable. It should be used on the
+        LEFT side of the assignment.
+        """
+        return self._get_tileable_variable(tileable)
+
+    def _get_tileable_variable(self, tileable: Tileable) -> str:
         try:
             return self._tileable_key_to_variables[tileable.key]
         except KeyError:
@@ -315,7 +329,7 @@ class EngineAcceptance(Enum):
 
 class BigDagOperatorAdapter(metaclass=abc.ABCMeta):
     # todo handle refcount issue when generated code is being executed
-    def accepts(self, op: OperatorType) -> EngineAcceptance:
+    def accepts(self, op: Operator) -> EngineAcceptance:
         return EngineAcceptance.ACCEPT
 
     @abc.abstractmethod
@@ -330,7 +344,7 @@ class BigDagOperatorAdapter(metaclass=abc.ABCMeta):
 
         Parameters
         ----------
-        op : OperatorType
+        op : Operator
             The operator instance.
         context : BigDagCodeContext
             The BigDagCodeContext instance.
@@ -342,6 +356,48 @@ class BigDagOperatorAdapter(metaclass=abc.ABCMeta):
         """
         return list()
 
+    def generate_pre_op_code(
+        self, op: Operator, context: BigDagCodeContext
+    ) -> List[str]:
+        """
+        Generate the codes before actually handling the operator.
+        This method is usually implemented in the base class of each engine.
+
+        Parameters
+        ----------
+        op : Operator
+            The operator instance.
+        context : BigDagCodeContext
+            The BigDagCodeContext instance.
+
+        Returns
+        -------
+        result: List[str]
+            The codes generated before one operator actually handled, one per line.
+        """
+        return list()
+
+    def generate_post_op_code(
+        self, op: Operator, context: BigDagCodeContext
+    ) -> List[str]:
+        """
+        Generate the codes after actually handling the operator.
+        This method is usually implemented in the base class of each engine.
+
+        Parameters
+        ----------
+        op : Operator
+            The operator instance.
+        context : BigDagCodeContext
+            The BigDagCodeContext instance.
+
+        Returns
+        -------
+        result: List[str]
+            The codes generated after one operator actually handled, one per line.
+        """
+        return list()
+
 
 _engine_to_codegen: Dict[str, Type["BigDagCodeGenerator"]] = dict()
 
@@ -354,9 +410,6 @@ def register_engine_codegen(type_: Type["BigDagCodeGenerator"]):
 BUILTIN_ENGINE_SPE = "SPE"
 BUILTIN_ENGINE_MCSQL = "MCSQL"
 
-FAST_RANGE_INDEX_ENABLED = "codegen.fast_range_index_enabled"
-ROW_NUMBER_WINDOW_INDEX_ENABLED = "codegen.row_number_window_index_enabled"
-
 
 class BigDagCodeGenerator(metaclass=abc.ABCMeta):
     _context: BigDagCodeContext
@@ -364,11 +417,13 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
     engine_type: Optional[str] = None
     engine_priority: int = 0
     _extension_loaded = False
+    _generate_comments_enabled: bool = True
 
     def __init__(self, session_id: str, subdag_id: str = None):
         self._session_id = session_id
         self._subdag_id = subdag_id
         self._context = self._init_context(session_id, subdag_id)
+        self._generate_comments_enabled = True
 
     @classmethod
     def _load_engine_extensions(cls):
@@ -401,14 +456,6 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
     def _init_context(self, session_id: str, subdag_id: str) -> BigDagCodeContext:
         raise NotImplementedError
 
-    def _generate_comments(
-        self, op: OperatorType, adapter: BigDagOperatorAdapter
-    ) -> List[str]:
-        return adapter.generate_comment(op, self._context)
-
-    def _generate_pre_op_code(self, op: OperatorType) -> List[str]:
-        return []
-
     def _generate_delete_code(self, var_name: str) -> List[str]:
         return []
 
@@ -438,9 +485,11 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
             visited_op_key.add(op.key)
 
             adapter = self.get_op_adapter(type(op))()
-            code_lines.extend(self._generate_pre_op_code(op))
-            code_lines.extend(self._generate_comments(op, adapter))
+            code_lines.extend(adapter.generate_pre_op_code(op, self._context))
+            if self._generate_comments_enabled:
+                code_lines.extend(adapter.generate_comment(op, self._context))
             code_lines.extend(adapter.generate_code(op, self._context))
+            code_lines.extend(adapter.generate_post_op_code(op, self._context))
             code_lines.append("")  # Append an empty line to separate operators
 
             # record refcounts
@@ -449,7 +498,7 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
                     continue
                 if dag.count_successors(out_t) == 0:
                     delete_code = self._generate_delete_code(
-                        self._context.get_tileable_variable(out_t)
+                        self._context.get_input_tileable_variable(out_t)
                     )
                     code_lines.extend(delete_code)
                 else:
@@ -462,7 +511,7 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
                 out_refcounts[inp_t.key] -= 1
                 if out_refcounts[inp_t.key] == 0:
                     delete_code = self._generate_delete_code(
-                        self._context.get_tileable_variable(inp_t)
+                        self._context.get_input_tileable_variable(inp_t)
                     )
                     code_lines.extend(delete_code)
                     out_refcounts.pop(inp_t.key)
@@ -475,11 +524,11 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
         for tileable in dag.topological_iter():
             op: OperatorType = tileable.op
             if isinstance(op, Fetch):
-                fetch_tileable = self._context.get_tileable_variable(tileable)
+                fetch_tileable = self._context.get_input_tileable_variable(tileable)
                 input_key_to_vars[op.outputs[0].key] = fetch_tileable
 
         result_variables = {
-            t.key: self._context.get_tileable_variable(t) for t in dag.results
+            t.key: self._context.get_input_tileable_variable(t) for t in dag.results
         }
 
         return CodeGenResult(

maxframe/config/config.py

@@ -407,6 +407,12 @@ default_options.register_option(
     validator=is_integer,
     remote=True,
 )
+default_options.register_option(
+    "session.temp_table_properties",
+    None,
+    validator=is_null | is_dict,
+    remote=True,
+)
 default_options.register_option(
     "session.auto_purge_temp_tables",
     False,

maxframe/core/accessor.py

@@ -11,6 +11,7 @@
 # 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 functools
 
 

maxframe/dataframe/accessors/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import datetime_, dict_, plotting, string_
+from . import datetime_, dict_, list_, plotting, string_

maxframe/dataframe/accessors/dict_/accessor.py

@@ -11,6 +11,7 @@
 # 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 TYPE_CHECKING
 
 import pandas as pd

maxframe/dataframe/accessors/dict_/length.py

@@ -11,6 +11,7 @@
 # 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
 

maxframe/dataframe/accessors/dict_/setitem.py

@@ -11,6 +11,7 @@
 # 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 ....core.entity.output_types import OutputType
 from ....serialization.serializables.field import AnyField

maxframe/dataframe/accessors/dict_/tests/test_dict_accessor.py

@@ -11,6 +11,7 @@
 # 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
@@ -25,6 +26,10 @@ from ..length import SeriesDictLengthOperator
 from ..remove import SeriesDictRemoveOperator
 from ..setitem import SeriesDictSetItemOperator
 
+pytestmark = pytest.mark.skipif(
+    ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported"
+)
+
 
 @pytest.fixture
 def df():
@@ -40,13 +45,11 @@ def df():
     )
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_invalid_dtype(df):
     with pytest.raises(AttributeError):
         df["C"].dict.contains("k1")
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_getitem(df):
     s1 = df["A"].dict["k1"]
     assert isinstance(s1, md.Series)
@@ -61,7 +64,6 @@ def test_getitem(df):
     assert op.ignore_key_error is False
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_getitem_with_default_value(df):
     s1 = df["B"].dict.get("k1", 1)
     assert isinstance(s1, md.Series)
@@ -76,7 +78,6 @@ def test_getitem_with_default_value(df):
     assert op.ignore_key_error is True
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_setitem(df):
     s1 = df["A"]
     s1.dict["k1"] = "v3"
@@ -91,7 +92,6 @@ def test_setitem(df):
     assert op.value == "v3"
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_length(df):
     s1 = df["A"].dict.len()
     assert isinstance(s1, md.Series)
@@ -103,7 +103,6 @@ def test_length(df):
     assert isinstance(op, SeriesDictLengthOperator)
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_remove(df):
     s1 = df["A"].dict.remove("k1", ignore_key_error=True)
     assert isinstance(s1, md.Series)
@@ -117,7 +116,6 @@ def test_remove(df):
     assert op.ignore_key_error is True
 
 
-@pytest.mark.skipif(ARROW_DTYPE_NOT_SUPPORTED, reason="Arrow Dtype is not supported")
 def test_contains(df):
     s1 = df["A"].dict.contains("k1")
     assert isinstance(s1, md.Series)

maxframe/dataframe/accessors/list_/__init__.py

@@ -0,0 +1,37 @@
+# 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.
+
+
+def _install():
+    from ....core import CachedAccessor
+    from ...core import SERIES_TYPE
+    from .accessor import ListAccessor
+    from .getitem import series_list_getitem, series_list_getitem_with_index_error
+    from .length import series_list_length
+
+    dict_method_to_handlers = {
+        "__getitem__": series_list_getitem_with_index_error,
+        "get": series_list_getitem,
+        "len": series_list_length,
+    }
+
+    for name, handler in dict_method_to_handlers.items():
+        ListAccessor._register(name, handler)
+
+    for series in SERIES_TYPE:
+        series.list = CachedAccessor("list", ListAccessor)
+
+
+_install()
+del _install

maxframe/dataframe/accessors/list_/accessor.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 typing import TYPE_CHECKING
+
+import pandas as pd
+import pyarrow as pa
+
+from ....core import BaseMaxFrameAccessor
+from ....utils import ARROW_DTYPE_NOT_SUPPORTED
+
+if TYPE_CHECKING:
+    from ...core import Series
+
+
+class ListAccessor(BaseMaxFrameAccessor):
+    obj: "Series"
+    _api_count: int = 0
+
+    def __init__(self, series):
+        super().__init__(series)
+        if ARROW_DTYPE_NOT_SUPPORTED:
+            raise ImportError("pd.ArrowDtype is not supported in current environment")
+
+        if not isinstance(series.dtype, pd.ArrowDtype) or not isinstance(
+            series.dtype.pyarrow_dtype, pa.ListType
+        ):
+            raise AttributeError("Can only use .list accessor with list values")

maxframe/dataframe/accessors/list_/getitem.py

@@ -0,0 +1,135 @@
+# 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
+
+from .... import opcodes
+from ....core.entity.output_types import OutputType
+from ....serialization.serializables.field import AnyField, BoolField
+from ...operators import DataFrameOperator, DataFrameOperatorMixin
+
+
+class SeriesListGetItemOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.SERIES_LIST_GETITEM
+    query_index = AnyField("query_index", default=None)
+    ignore_index_error = BoolField("ignore_index_error", default=False)
+
+    def __init__(self, **kw):
+        super().__init__(_output_types=[OutputType.series], **kw)
+
+    def __call__(self, series):
+        arrow_list_type = series.dtype.pyarrow_dtype
+        return self.new_series(
+            [series],
+            shape=series.shape,
+            dtype=pd.ArrowDtype(arrow_list_type.value_type),
+            index_value=series.index_value,
+        )
+
+
+def series_list_getitem(series, query_index):
+    """
+    Get the value by the index of each list in the Series.
+
+    Parameters
+    ----------
+    query_index : Any
+        The key to check, must be index.
+
+    Returns
+    -------
+    Series :
+        A Series with the list value's data type. The value will be
+        ``None`` if the list is None.
+
+    Examples
+    --------
+    Create a series with list type data.
+
+    >>> import maxframe.dataframe as md
+    >>> import pyarrow as pa
+    >>> from maxframe.lib.dtypes_extension import list_
+    >>> s = md.Series(
+    ...     data=[[1, 2, 3], [4, 5, 6], None],
+    ...     index=[1, 2, 3],
+    ...     dtype=list_(pa.int64()),
+    ... )
+    >>> s.execute()
+    1    [1, 2, 3]
+    2    [4, 5, 6]
+    3         <NA>
+    dtype: list<int64>[pyarrow]
+
+    >>> s.list.get(0).execute()
+    1       1
+    2       4
+    3    <NA>
+    dtype: int64[pyarrow]
+    """
+    return SeriesListGetItemOperator(query_index=query_index, ignore_index_error=True)(
+        series
+    )
+
+
+def series_list_getitem_with_index_error(series, query_index):
+    """
+    Get the value by the index of each list in the Series. If the index
+    is not in the list, raise IndexError.
+
+    Parameters
+    ----------
+    query_index : Any
+        The index to check, must be integer.
+
+    Returns
+    -------
+    Series :
+        A Series with the list value's data type. Return ``None`` if the list is None.
+
+    Raises
+    ------
+    KeyError
+        If the index is not in one list.
+
+    See Also
+    --------
+    Series.list.get: Get the value by the index of each list in the Series.
+
+    Examples
+    --------
+    Create a series with list type data.
+
+    >>> import maxframe.dataframe as md
+    >>> import pyarrow as pa
+    >>> from maxframe.lib.dtypes_extension import list_
+    >>> s = md.Series(
+    ...     data=[[1, 2, 3], [4, 5, 6], None],
+    ...     index=[1, 2, 3],
+    ...     dtype=list_(pa.int64()),
+    ... )
+    >>> s.execute()
+    1    [1, 2, 3]
+    2    [4, 5, 6]
+    3         <NA>
+    dtype: list<int64>[pyarrow]
+
+    >>> s.list.get(0).execute()
+    1       1
+    2       4
+    3    <NA>
+    dtype: int64[pyarrow]
+    """
+    return SeriesListGetItemOperator(query_index=query_index, ignore_index_error=False)(
+        series
+    )

maxframe/dataframe/accessors/list_/length.py

@@ -0,0 +1,73 @@
+# 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
+
+from .... import opcodes
+from ....core.entity.output_types import OutputType
+from ...operators import DataFrameOperator, DataFrameOperatorMixin
+
+
+class SeriesListLengthOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.SERIES_LIST_LENGTH
+
+    def __init__(self, **kw):
+        super().__init__(_output_types=[OutputType.series], **kw)
+
+    def __call__(self, series):
+        return self.new_series(
+            [series],
+            shape=series.shape,
+            index_value=series.index_value,
+            dtype=pd.ArrowDtype(pa.int64()),
+            name=None,
+        )
+
+
+def series_list_length(series):
+    """
+    Get the length of each list of the Series.
+
+    Returns
+    -------
+    Series :
+        A Series with data type ``pandas.ArrowDtype(pyarrow.int64)``. Each element
+        represents the length of the list, or ``None`` if the list is ``None``.
+
+    Examples
+    --------
+    Create a series with list type data.
+
+    >>> import maxframe.dataframe as md
+    >>> import pyarrow as pa
+    >>> from maxframe.lib.dtypes_extension import list_
+    >>> s = md.Series(
+    ...     data=[[1, 2, 3], [4, 5, 6], None],
+    ...     index=[1, 2, 3],
+    ...     dtype=list_(pa.int64()),
+    ... )
+    >>> s.execute()
+    1    [1, 2, 3]
+    2    [4, 5, 6]
+    3         <NA>
+    dtype: list<int64>[pyarrow]
+
+    >>> s.list.len().execute()
+    1       2
+    2       1
+    3    <NA>
+    dtype: int64[pyarrow]
+    """
+    return SeriesListLengthOperator()(series)

maxframe/dataframe/accessors/list_/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.