#!/usr/bin/env python # -*- coding: utf-8 -*- # Copyright 1999-2022 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 itertools import json import os import tarfile import uuid import time from collections import OrderedDict from .... import options, tempobj from ....compat import BytesIO, six from ....utils import TEMP_TABLE_PREFIX from ....errors import ODPSError UDF_CLASS_NAME = 'PyOdpsFunc' expr_to_sql = dict() expr_deps = dict() expr_ref_name = dict() class ODPSContext(object): def __init__(self, odps, indent_size=2): self._odps = odps self._index = itertools.count(1) self._expr_alias = dict() self._compiled_exprs = dict() self._expr_raw_args = dict() self._col_index = itertools.count(1) self._need_alias_columns = dict() self._need_alias_column_indexes = dict() self._select_index = itertools.count(1) self._expr_ref_index = itertools.count(1) self._func_to_udfs = OrderedDict() self._registered_funcs = OrderedDict() self._func_to_functions = OrderedDict() self._func_to_resources = OrderedDict() self._func_to_cus = OrderedDict() self._indent_size = indent_size self._mapjoin_hints = [] self._skewjoin_hints = [] self._default_schema = odps.schema self._path_to_resources = dict() self._to_drops = [] def next_select_id(self): return next(self._select_index) def new_alias_table_name(self): return 't%s' % next(self._index) def new_ref_name(self): return '@c%d' % next(self._expr_ref_index) def register_collection(self, expr): expr_id = id(expr) if expr_id in self._expr_alias: return self._expr_alias[expr_id] self._expr_alias[expr_id] = self.new_alias_table_name() return self._expr_alias[expr_id] def get_collection_alias(self, expr, create=False, silent=False): try: return self._expr_alias[id(expr)], False except KeyError as e: if create: return self.register_collection(expr), True if not silent: raise e def remove_collection_alias(self, expr): if id(expr) in self._expr_alias: del self._expr_alias[id(expr)] def add_expr_compiled(self, expr, compiled): self._compiled_exprs[id(expr)] = compiled def get_expr_compiled(self, expr): return self._compiled_exprs[id(expr)] def _gen_udf_name(self, func): identifier = getattr(func, "_identifier", None) or "udf" return 'pyodps_%s_%s_%s' % ( identifier, int(time.time()), str(uuid.uuid4()).replace('-', '_') ) def _gen_resource_name(self): return 'pyodps_res_%s_%s' % (int(time.time()), str(uuid.uuid4()).replace('-', '_')) def _gen_table_name(self): return '%s_%s_%s' % (TEMP_TABLE_PREFIX, int(time.time()), str(uuid.uuid4()).replace('-', '_')) def register_udfs(self, func_to_udfs, func_to_resources, func_to_cus): self._func_to_udfs = func_to_udfs for func in func_to_udfs.keys(): self._registered_funcs[func] = self._gen_udf_name(func) self._func_to_resources = func_to_resources self._func_to_cus = func_to_cus def get_udf(self, func): return self._registered_funcs[func] def get_udf_count(self): return len(self._registered_funcs) def get_udf_sql_hints(self): cu_dict = OrderedDict() for func, name in self._registered_funcs.items(): if func in self._func_to_cus: prefix = self._odps.project + "." if self._default_schema: prefix += self._default_schema + "." cu_dict[prefix + name] = self._func_to_cus[func] if not cu_dict: return {} return {"odps.sql.udf.cu": json.dumps(cu_dict)} def prepare_resources(self, libraries): from ....models import Resource if libraries is None: return None ret_libs = [] for lib in libraries: if isinstance(lib, Resource): ret_libs.append(lib) continue elif lib in self._path_to_resources: ret_libs.append(self._path_to_resources[lib]) continue tarbinary = BytesIO() tar = tarfile.open(fileobj=tarbinary, mode='w:gz') if os.path.isfile(lib): with open(lib, 'rb') as fo: finfo = tarfile.TarInfo('pyodps_files/' + os.path.basename(lib)) finfo.size = os.path.getsize(lib) tar.addfile(finfo, fo) else: base_dir = os.path.dirname(os.path.abspath(lib)) for root, dirs, files in os.walk(lib): for f in files: fpath = os.path.join(root, f) rpath = os.path.relpath(fpath, base_dir).replace(os.path.sep, '/') with open(fpath, 'rb') as fo: finfo = tarfile.TarInfo(rpath) finfo.size = os.path.getsize(fpath) tar.addfile(finfo, fo) tar.close() res_name = self._gen_resource_name() + '.tar.gz' res = self._odps.create_resource( res_name, 'archive', fileobj=tarbinary.getvalue(), schema=self._default_schema, temp=True ) if options.df.delete_udfs: tempobj.register_temp_resource( self._odps, res_name, schema=self._default_schema ) self._to_drops.append(res) self._path_to_resources[lib] = res ret_libs.append(res) return ret_libs def create_udfs(self, libraries=None): self._func_to_functions.clear() for func, udf in six.iteritems(self._func_to_udfs): udf_name = self._registered_funcs[func] py_resource = self._odps.create_resource( udf_name + '.py', 'py', file_obj=udf, schema=self._default_schema, temp=True ) if options.df.delete_udfs: tempobj.register_temp_resource( self._odps, udf_name + '.py', schema=self._default_schema ) self._to_drops.append(py_resource) resources = [py_resource, ] if func in self._func_to_resources: for _, name, _, create, table_name in self._func_to_resources[func]: if not create: resources.append(name) else: res = self._odps.create_resource( name, 'table', table_name=table_name, schema=self._default_schema, temp=True ) if options.df.delete_udfs: tempobj.register_temp_resource(self._odps, name, schema=self._default_schema) self._to_drops.append(res) resources.append(res) if libraries is not None: resources.extend(libraries) function_name = '{0}.{1}'.format(udf_name, UDF_CLASS_NAME) function = self._odps.create_function( udf_name, class_type=function_name, resources=resources, schema=self._default_schema ) if options.df.delete_udfs: tempobj.register_temp_function(self._odps, udf_name, schema=self._default_schema) self._to_drops.append(function) self._func_to_functions[func] = function def add_need_alias_column(self, column): if id(column) in self._need_alias_column_indexes: return self._need_alias_column_indexes[id(column)] symbol = 'col_%s' % next(self._col_index) self._need_alias_columns[id(column)] = column self._need_alias_column_indexes[id(column)] = symbol return symbol def get_all_need_alias_column_symbols(self): for col_id, column in six.iteritems(self._need_alias_columns): symbol = self._need_alias_column_indexes[col_id] yield symbol, column ##################################### # mem cache relatives ##################################### def register_mem_cache_sql(self, expr, sql): expr_to_sql[expr._id] = sql def is_expr_mem_cached(self, expr): return expr._id in expr_to_sql def get_mem_cached_sql(self, expr): return expr_to_sql[expr._id] def register_mem_cache_dep(self, expr, dep): if expr._id not in expr_deps: expr_deps[expr._id] = set() expr_deps[expr._id].add(dep._id) def get_mem_cache_ref_name(self, expr, create=True): if not create: return expr_ref_name[expr._id] ref_name = expr_ref_name.get(expr._id, self.new_ref_name()) if expr._id not in expr_ref_name: expr_ref_name[expr._id] = ref_name return ref_name def get_mem_cache_dep_sqls(self, *expr_ids): sqls = [] fetched = dict() def h(eid): if eid in expr_deps: for dep_id in expr_deps[eid]: if dep_id not in fetched: h(dep_id) if eid not in fetched: origin_sql = expr_to_sql[eid] sql = '{0} := CACHE ON {1}'.format( expr_ref_name[eid], origin_sql ) sqls.append(sql) fetched[eid] = True for expr_id in expr_ids: h(expr_id) return sqls def _drop_silent(self, obj): try: obj.drop() except ODPSError: pass @property def default_schema(self): return self._default_schema @default_schema.setter def default_schema(self, value): self._default_schema = value def close(self): [self._drop_silent(it) for it in self._to_drops]