maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/sperunner.py

@@ -0,0 +1,165 @@
+# 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 logging
+import time
+from typing import Any, Dict
+
+import pandas as pd
+from odps import ODPS
+
+from .codegen import CodeGenResult
+from .codegen.spe.core import SPECodeContext
+from .config import option_context
+from .core import TileableGraph
+from .lib.compat import patch_pandas
+from .protocol import ResultInfo
+from .typing_ import PandasObjectTypes, TileableType
+from .utils import build_temp_table_name
+
+logger = logging.getLogger(__name__)
+
+
+class DAGCancelledError(Exception):
+    pass
+
+
+class SPEDagRunner:
+    def __init__(
+        self,
+        session_id: str,
+        subdag_id: str,
+        subdag: TileableGraph,
+        generated: CodeGenResult,
+        settings: Dict[str, Any],
+    ):
+        self._session_id = session_id
+        self._subdag_id = subdag_id
+        self._subdag = subdag
+        self._key_to_tileable = {t.key: t for t in subdag}
+        self._settings = settings
+
+        with option_context(self._settings) as session_options:
+            self._sql_hints = session_options.sql.settings
+
+        self._code = generated.code
+        self._constants = generated.constants
+        self._out_key_to_vars = generated.output_key_to_variables
+        self._out_key_to_infos = generated.output_key_to_result_infos
+        self._input_key_to_vars = generated.input_key_to_variables
+        self._mark_cancelled = False
+        self._odps = ODPS.from_environments()
+
+    def _get_sql_hints(self) -> Dict[str, str]:
+        hints = self._sql_hints.copy()
+        hints["odps.sql.type.system.odps2"] = True
+        return hints
+
+    @staticmethod
+    def _pre_process_pandas_data(data: PandasObjectTypes) -> PandasObjectTypes:
+        """
+        Sort the pandas dataset first to make sure all the tensor can process the
+        records in the same order.
+        """
+        if isinstance(data, (pd.DataFrame, pd.Series)):
+            data.sort_index(
+                axis=0,
+                ascending=True,
+                kind="quicksort",
+                inplace=True,
+                na_position="last",
+            )
+        elif isinstance(data, pd.Index):
+            data.sort_values(return_indexer=False, ascending=True, na_position="last")
+        return data
+
+    def fetch_data_by_tileable(self, t: TileableType) -> Any:
+        raise NotImplementedError
+
+    def store_data(self, key: str, value: Any) -> ResultInfo:
+        raise NotImplementedError
+
+    def run(self) -> Dict[str, ResultInfo]:
+        # make forward compatibility of new pandas methods
+        patch_pandas()
+
+        local_vars = self._constants
+        local_vars[SPECodeContext.logger_var] = logger
+        start_timestamp = time.time()
+        try:
+            # Fetch data
+            start_time = time.time()
+            for key, var_name in self._input_key_to_vars.items():
+                local_vars[var_name] = self.fetch_data_by_tileable(
+                    self._key_to_tileable[key]
+                )
+            logger.info(
+                "[%s][%s] fetch data costs: %.3f",
+                self._session_id,
+                self._subdag_id,
+                time.time() - start_time,
+            )
+
+            # Execute Python codes
+            start_time = time.time()
+            logger.info("Generated codes:\n--------\n%s\n--------", self._code)
+            exec(self._code, globals(), local_vars)
+            logger.info(
+                "[%s][%s] execute costs: %.3f",
+                self._session_id,
+                self._subdag_id,
+                time.time() - start_time,
+            )
+
+            # Store data
+            start_time = time.time()
+            result_dict = dict()
+            for key, var_name in self._out_key_to_vars.items():
+                result_dict[key] = self.store_data(key, local_vars[var_name])
+            logger.info(
+                "[%s][%s] store data costs: %.3f",
+                self._session_id,
+                self._subdag_id,
+                time.time() - start_time,
+            )
+            return result_dict
+        except Exception as ex:
+            local_vars.clear()
+            if (
+                not isinstance(ex, RuntimeError)
+                or not ex.args
+                or ex.args[0] != "CANCELLED"
+            ):
+                raise
+
+            drop_statements = [
+                f"DROP TABLE IF EXISTS {build_temp_table_name(self._session_id, key)};"
+                for key in self._out_key_to_vars.keys()
+            ]
+            if drop_statements:  # pragma: no branch
+                self._odps.run_sql(
+                    "\n".join(drop_statements), hints=self._get_sql_hints()
+                )
+            raise DAGCancelledError from None
+        finally:
+            logger.info(
+                "[%s][%s] run costs: %.3f",
+                self._session_id,
+                self._subdag_id,
+                time.time() - start_timestamp,
+            )
+
+    def mark_cancel(self):
+        self._constants["running"] = False
+        self._mark_cancelled = True

maxframe/tensor/__init__.py

@@ -0,0 +1,325 @@
+# 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 extensions, fetch, fft, linalg, random, spatial, special, statistics
+from .arithmetic import absolute
+from .arithmetic import absolute as abs
+from .arithmetic import (
+    add,
+    angle,
+    arccos,
+    arccosh,
+    arcsin,
+    arcsinh,
+    arctan,
+    arctan2,
+    arctanh,
+    around,
+)
+from .arithmetic import bitand as bitwise_and
+from .arithmetic import bitor as bitwise_or
+from .arithmetic import bitxor as bitwise_xor
+from .arithmetic import (
+    cbrt,
+    ceil,
+    clip,
+    conj,
+    conjugate,
+    copysign,
+    cos,
+    cosh,
+    deg2rad,
+    degrees,
+    divide,
+    equal,
+    exp,
+    exp2,
+    expm1,
+    fabs,
+    fix,
+    float_power,
+    floor,
+)
+from .arithmetic import floordiv as floor_divide
+from .arithmetic import fmax, fmin, fmod, frexp, greater, greater_equal, hypot, i0, imag
+from .arithmetic import invert
+from .arithmetic import invert as bitwise_not
+from .arithmetic import (
+    isclose,
+    iscomplex,
+    iscomplexobj,
+    isfinite,
+    isinf,
+    isnan,
+    isreal,
+    ldexp,
+    less,
+    less_equal,
+    log,
+    log1p,
+    log2,
+    log10,
+    logaddexp,
+    logaddexp2,
+    logical_and,
+    logical_not,
+    logical_or,
+    logical_xor,
+)
+from .arithmetic import lshift as left_shift
+from .arithmetic import (
+    maximum,
+    minimum,
+    mod,
+    modf,
+    multiply,
+    nan_to_num,
+    negative,
+    nextafter,
+    not_equal,
+    positive,
+    power,
+    rad2deg,
+    radians,
+    real,
+    reciprocal,
+    remainder,
+    rint,
+)
+from .arithmetic import round_
+from .arithmetic import round_ as round
+from .arithmetic import rshift as right_shift
+from .arithmetic import (
+    sign,
+    signbit,
+    sin,
+    sinc,
+    sinh,
+    spacing,
+    sqrt,
+    square,
+    subtract,
+    tan,
+    tanh,
+)
+from .arithmetic import truediv as true_divide
+from .arithmetic import trunc
+from .core import Tensor
+from .datasource import (
+    arange,
+    array,
+    asarray,
+    ascontiguousarray,
+    asfortranarray,
+    diag,
+    diagflat,
+    empty,
+    empty_like,
+    eye,
+    from_dataframe,
+    full,
+    full_like,
+    identity,
+    indices,
+    linspace,
+    meshgrid,
+    ones,
+    ones_like,
+    scalar,
+    tensor,
+    tril,
+    triu,
+    zeros,
+    zeros_like,
+)
+from .indexing import (
+    choose,
+    compress,
+    extract,
+    fill_diagonal,
+    flatnonzero,
+    nonzero,
+    take,
+    unravel_index,
+)
+from .linalg import dot, einsum, inner, innerproduct, matmul, tensordot, vdot
+from .merge import append, column_stack, concatenate, dstack, hstack, stack, vstack
+from .misc import (
+    argwhere,
+    array_split,
+    astype,
+    atleast_1d,
+    atleast_2d,
+    atleast_3d,
+    broadcast_arrays,
+    broadcast_to,
+    copy,
+    copyto,
+    delete,
+    diff,
+    dsplit,
+    ediff1d,
+    expand_dims,
+    flip,
+    fliplr,
+    flipud,
+    hsplit,
+    in1d,
+    insert,
+    isin,
+    moveaxis,
+    ndim,
+    ravel,
+    repeat,
+    result_type,
+    roll,
+    rollaxis,
+    searchsorted,
+    setdiff1d,
+    shape,
+    split,
+    squeeze,
+    swapaxes,
+    tile,
+    transpose,
+    trapezoid,
+    unique,
+    vsplit,
+    where,
+)
+from .rechunk import rechunk
+from .reduction import (
+    all,
+    allclose,
+    any,
+    argmax,
+    argmin,
+    array_equal,
+    count_nonzero,
+    cumprod,
+    cumsum,
+)
+from .reduction import max
+from .reduction import max as amax
+from .reduction import mean
+from .reduction import min
+from .reduction import min as amin
+from .reduction import (
+    nanargmax,
+    nanargmin,
+    nancumprod,
+    nancumsum,
+    nanmax,
+    nanmean,
+    nanmin,
+    nanprod,
+    nanstd,
+    nansum,
+    nanvar,
+)
+from .reduction import prod
+from .reduction import prod as product
+from .reduction import std, sum, var
+from .reshape import reshape
+from .sort import argpartition, argsort, partition, sort
+from .statistics import (
+    average,
+    bincount,
+    corrcoef,
+    cov,
+    digitize,
+    histogram,
+    histogram_bin_edges,
+    median,
+    percentile,
+    ptp,
+    quantile,
+)
+from .ufunc import ufunc
+
+# isort: off
+# noinspection PyUnresolvedReferences
+from .lib.index_tricks import mgrid, ogrid, ndindex, r_, c_
+
+# noinspection PyUnresolvedReferences
+from numpy import (
+    e,
+    errstate,
+    geterr,
+    inf,
+    nan,
+    newaxis,
+    pi,
+    seterr,
+)
+
+try:
+    from numpy.exceptions import AxisError
+except ImportError:
+    from numpy import AxisError
+
+NAN = nan
+NINF = -inf
+Inf = inf
+NaN = nan
+
+# import numpy types
+# noinspection PyUnresolvedReferences
+from numpy import (
+    bool_ as bool,
+    bytes_,
+    character,
+    complex64,
+    complex128,
+    complexfloating,
+    datetime64,
+    double,
+    dtype,
+    flexible,
+    float16,
+    float32,
+    float64,
+    floating,
+    generic,
+    inexact,
+    int8,
+    int16,
+    int32,
+    int64,
+    intc,
+    intp,
+    number,
+    integer,
+    object_ as object,
+    signedinteger,
+    timedelta64,
+    uint,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    unsignedinteger,
+    void,
+)
+
+try:
+    from numpy import cfloat
+except ImportError:
+    from numpy import cdouble as cfloat
+try:
+    from numpy import str_ as unicode_
+except ImportError:
+    from numpy import unicode_
+
+del fetch, ufunc