# 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 abc import base64 import dataclasses import logging from collections import defaultdict from enum import Enum from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union from odps.types import OdpsSchema from odps.utils import camel_to_underline from .core import OperatorType, Tileable, TileableGraph 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 from .lib import wrapped_pickle as pickle from .protocol import DataFrameTableMeta, ResultInfo from .serialization import PickleContainer from .serialization.serializables import Serializable, StringField from .typing_ import PandasObjectTypes from .udf import MarkedFunction, PythonPackOptions if TYPE_CHECKING: from odpsctx import ODPSSessionContext logger = logging.getLogger(__name__) @dataclasses.dataclass class CodeGenResult: code: str input_key_to_variables: Dict[str, str] output_key_to_variables: Dict[str, str] output_key_to_result_infos: Dict[str, ResultInfo] constants: Dict[str, Any] class AbstractUDF(Serializable): _session_id: str = StringField("session_id") def __init__(self, session_id: Optional[str] = None, **kw): super().__init__(_session_id=session_id, **kw) @property def name(self) -> str: return camel_to_underline(type(self).__name__) @property def session_id(self): return getattr(self, "_session_id", None) @session_id.setter def session_id(self, value: str): self._session_id = value @abc.abstractmethod def register(self, odps: "ODPSSessionContext", overwrite: bool = False): raise NotImplementedError @abc.abstractmethod def unregister(self, odps: "ODPSSessionContext"): raise NotImplementedError @abc.abstractmethod def collect_pythonpack(self) -> List[PythonPackOptions]: raise NotImplementedError @abc.abstractmethod def load_pythonpack_resources(self, odps_ctx: "ODPSSessionContext") -> None: raise NotImplementedError class UserCodeMixin: __slots__ = () @classmethod def obj_to_python_expr(cls, obj: Any = None) -> str: """ Parameters ---------- obj The object to convert to python expr. Returns ------- str : The str type content equals to the object when use in the python code directly. """ if obj is None: return "None" if isinstance(obj, (int, float)): return repr(obj) if isinstance(obj, bool): return "True" if obj else "False" if isinstance(obj, bytes): base64_bytes = base64.b64encode(obj) return f"base64.b64decode({base64_bytes})" if isinstance(obj, str): return repr(obj) if isinstance(obj, list): return ( f"[{', '.join([cls.obj_to_python_expr(element) for element in obj])}]" ) if isinstance(obj, dict): items = ( f"{repr(key)}: {cls.obj_to_python_expr(value)}" for key, value in obj.items() ) return f"{{{', '.join(items)}}}" if isinstance(obj, tuple): return f"({', '.join([cls.obj_to_python_expr(sub_obj) for sub_obj in obj])}{',' if len(obj) == 1 else ''})" if isinstance(obj, set): return ( f"{{{', '.join([cls.obj_to_python_expr(sub_obj) for sub_obj in obj])}}}" if obj else "set()" ) if isinstance(obj, PickleContainer): return UserCodeMixin.generate_pickled_codes(obj, None) raise ValueError(f"not support arg type {type(obj)}") @classmethod def generate_pickled_codes( cls, code_to_pickle: Any, unpicked_data_var_name: Union[str, None] = "pickled_data", ) -> str: """ Generate pickled codes. The final pickled variable is called 'pickled_data'. Parameters ---------- code_to_pickle: Any The code to be pickled. unpicked_data_var_name: str The variables in code used to hold the loads object from the cloudpickle Returns ------- str : The code snippets of pickling, the final variable is called 'pickled_data' by default. """ pickled, buffers = cls.dump_pickled_data(code_to_pickle) pickle_loads_expr = f"cloudpickle.loads({cls.obj_to_python_expr(pickled)}, buffers={cls.obj_to_python_expr(buffers)})" if unpicked_data_var_name: return f"{unpicked_data_var_name} = {pickle_loads_expr}" return pickle_loads_expr @staticmethod def dump_pickled_data( code_to_pickle: Any, ) -> Tuple[List[bytes], List[bytes]]: if isinstance(code_to_pickle, MarkedFunction): code_to_pickle = code_to_pickle.func if isinstance(code_to_pickle, PickleContainer): buffers = code_to_pickle.get_buffers() pickled = buffers[0] buffers = buffers[1:] else: pickled = pickle.dumps(code_to_pickle, protocol=pickle.DEFAULT_PROTOCOL) buffers = [] return pickled, buffers class BigDagCodeContext(metaclass=abc.ABCMeta): def __init__(self, session_id: str = None, subdag_id: str = None): self._session_id = session_id self._subdag_id = subdag_id self._tileable_key_to_variables = dict() self.constants = dict() self._data_table_meta_cache = dict() self._odps_schema_cache = dict() self._udfs = dict() self._tileable_key_to_result_infos = dict() self._next_var_id = 0 self._next_const_id = 0 @property def session_id(self) -> str: return self._session_id def register_udf(self, udf: AbstractUDF): udf.session_id = self._session_id self._udfs[udf.name] = udf def get_udfs(self) -> List[AbstractUDF]: return list(self._udfs.values()) 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: var_name = self.next_var_name() self._tileable_key_to_variables[tileable.key] = var_name return var_name def next_var_name(self) -> str: var_name = f"var_{self._next_var_id}" self._next_var_id += 1 return var_name def get_odps_schema( self, data: PandasObjectTypes, unknown_as_string: bool = False ) -> OdpsSchema: """ Get the corresponding ODPS schema of the input df_obj. Parameters ---------- data : The pandas data object. unknown_as_string : Whether mapping the unknown data type to a temp string value. Returns ------- OdpsSchema : The OdpsSchema of df_obj. """ if data.key not in self._odps_schema_cache: odps_schema, table_meta = pandas_to_odps_schema(data, unknown_as_string) self._data_table_meta_cache[data.key] = table_meta self._odps_schema_cache[data.key] = odps_schema return self._odps_schema_cache[data.key] def get_pandas_data_table_meta(self, data: PandasObjectTypes) -> DataFrameTableMeta: if data.key not in self._data_table_meta_cache: self._data_table_meta_cache[data.key] = build_dataframe_table_meta(data) return self._data_table_meta_cache[data.key] def register_operator_constants(self, const_val, var_name: str = None) -> str: if var_name is None: if ( isinstance(const_val, (int, str, bytes, bool, float)) or const_val is None ): return repr(const_val) var_name = f"const_{self._next_const_id}" self._next_const_id += 1 self.constants[var_name] = const_val return var_name def put_tileable_result_info( self, tileable: Tileable, result_info: ResultInfo ) -> None: self._tileable_key_to_result_infos[tileable.key] = result_info def get_tileable_result_infos(self) -> Dict[str, ResultInfo]: return self._tileable_key_to_result_infos class EngineAcceptance(Enum): """ DENY: The operator is not accepted by the current engine. ACCEPT: The operator is accepted by the current engine, and doesn't break from here. BREAK: The operator is accepted by the current engine, but should break from here. PREDECESSOR: The acceptance of the operator is decided by engines of its predecessors. If acceptance of all predecessors are SUCCESSOR, the acceptance of current operator is SUCCESSOR. Otherwise the engine selected in predecessors with highest priority is used. SUCCESSOR: The acceptance of the operator is decided by engines of its successors. If the operator has no successors, the acceptance will be treated as ACCEPT. Otherwise the engine selected in successors with highest priority is used. """ DENY = 0 ACCEPT = 1 BREAK = 2 PREDECESSOR = 3 SUCCESSOR = 4 @classmethod def _missing_(cls, pred: bool) -> "EngineAcceptance": """ A convenience method to get ACCEPT or DENY result via the input predicate. Parameters ---------- pred : bool The predicate variable. Returns ------- EngineAcceptance : Returns ACCEPT if the predicate is true, otherwise returns DENY. """ return cls.ACCEPT if pred else cls.DENY class BigDagOperatorAdapter(metaclass=abc.ABCMeta): # todo handle refcount issue when generated code is being executed def accepts(self, op: Operator) -> EngineAcceptance: return EngineAcceptance.ACCEPT @abc.abstractmethod def generate_code(self, op: OperatorType, context: BigDagCodeContext) -> List[str]: raise NotImplementedError def generate_comment( self, op: OperatorType, context: BigDagCodeContext ) -> List[str]: """ Generate the comment codes before actual ones. Parameters ---------- op : Operator The operator instance. context : BigDagCodeContext The BigDagCodeContext instance. Returns ------- result: List[str] The comment codes, one per line. """ 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() def register_engine_codegen(type_: Type["BigDagCodeGenerator"]): _engine_to_codegen[type_.engine_type] = type_ return type_ BUILTIN_ENGINE_SPE = "SPE" BUILTIN_ENGINE_MCSQL = "MCSQL" class BigDagCodeGenerator(metaclass=abc.ABCMeta): _context: BigDagCodeContext 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): if cls._extension_loaded: return for name, ep in iter_extensions(): _engine_to_codegen[name.upper()] = ep.get_codegen() cls._extension_loaded = True @classmethod def get_engine_types(cls) -> List[str]: cls._load_engine_extensions() engines = sorted( _engine_to_codegen.values(), key=lambda x: x.engine_priority, reverse=True ) return [e.engine_type for e in engines] @classmethod def get_by_engine_type(cls, engine_type: str) -> Type["BigDagCodeGenerator"]: cls._load_engine_extensions() return _engine_to_codegen[engine_type] @abc.abstractmethod def get_op_adapter( self, op_type: Type[OperatorType] ) -> Type[BigDagOperatorAdapter]: raise NotImplementedError @abc.abstractmethod def _init_context(self, session_id: str, subdag_id: str) -> BigDagCodeContext: raise NotImplementedError def _generate_delete_code(self, var_name: str) -> List[str]: return [] def generate_code(self, dag: TileableGraph) -> List[str]: """ Generate the code of the input dag. Parameters ---------- dag : TileableGraph The input DAG instance. Returns ------- List[str] : The code lines. """ code_lines = [] visited_op_key = set() result_key_set = set(t.key for t in dag.result_tileables) out_refcounts = dict() for tileable in dag.topological_iter(): op: OperatorType = tileable.op if op.key in visited_op_key or isinstance(op, Fetch): continue visited_op_key.add(op.key) adapter = self.get_op_adapter(type(op))() 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 for out_t in op.outputs: if out_t.key in result_key_set: continue if dag.count_successors(out_t) == 0: delete_code = self._generate_delete_code( self._context.get_input_tileable_variable(out_t) ) code_lines.extend(delete_code) else: out_refcounts[out_t.key] = dag.count_successors(out_t) # check if refs of inputs are no longer needed for inp_t in op.inputs: if inp_t.key not in out_refcounts: continue out_refcounts[inp_t.key] -= 1 if out_refcounts[inp_t.key] == 0: delete_code = self._generate_delete_code( self._context.get_input_tileable_variable(inp_t) ) code_lines.extend(delete_code) out_refcounts.pop(inp_t.key) return code_lines def generate(self, dag: TileableGraph) -> CodeGenResult: code_lines = self.generate_code(dag) input_key_to_vars = dict() for tileable in dag.topological_iter(): op: OperatorType = tileable.op if isinstance(op, Fetch): 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_input_tileable_variable(t) for t in dag.results } return CodeGenResult( code="\n".join(code_lines), input_key_to_variables=input_key_to_vars, output_key_to_variables=result_variables, constants=self._context.constants, output_key_to_result_infos=self._context.get_tileable_result_infos(), ) def run_pythonpacks( self, odps_ctx: "ODPSSessionContext", python_tag: str, is_production: bool = False, schedule_id: Optional[str] = None, hints: Optional[dict] = None, priority: Optional[int] = None, ) -> Dict[str, PythonPackOptions]: key_to_packs = defaultdict(list) for udf in self._context.get_udfs(): for pack in udf.collect_pythonpack(): key_to_packs[pack.key].append(pack) distinct_packs = [] for packs in key_to_packs.values(): distinct_packs.append(packs[0]) inst_id_to_req = {} for pack in distinct_packs: inst = odps_ctx.run_pythonpack( requirements=pack.requirements, prefer_binary=pack.prefer_binary, pre_release=pack.pre_release, force_rebuild=pack.force_rebuild, no_audit_wheel=pack.no_audit_wheel, python_tag=python_tag, is_production=is_production, schedule_id=schedule_id, hints=hints, priority=priority, ) # fulfill instance id of pythonpacks with same keys for same_pack in key_to_packs[pack.key]: same_pack.pack_instance_id = inst.id inst_id_to_req[inst.id] = pack return inst_id_to_req def register_udfs(self, odps_ctx: "ODPSSessionContext"): for udf in self._context.get_udfs(): logger.info("[Session=%s] Registering UDF %s", self._session_id, udf.name) udf.register(odps_ctx, True) def unregister_udfs(self, odps_ctx: "ODPSSessionContext"): for udf in self._context.get_udfs(): logger.info("[Session=%s] Unregistering UDF %s", self._session_id, udf.name) udf.unregister(odps_ctx) def get_udfs(self) -> List[AbstractUDF]: return self._context.get_udfs()