energy-data-trans/etl/wind_power/min_sec/ReadAndSaveTmp.py

import datetime
import multiprocessing
import os
import traceback

import pandas as pd

from conf.constants import ParallelProcessing
from etl.common.PathsAndTable import PathsAndTable
from etl.wind_power.min_sec import TransParam
from service.trans_conf_service import update_trans_transfer_progress
from utils.file.trans_methods import read_excel_files, split_array, del_blank, \
    create_file_path, read_file_to_df, valid_eval
from utils.log.trans_log import info, debug, error
from utils.systeminfo.sysinfo import use_files_get_max_cpu_count, get_dir_size


class ReadAndSaveTmp(object):
    """读取并保存临时文件类"""

    def __init__(self, pathsAndTable: PathsAndTable, trans_param: TransParam):
        """
        初始化读取并保存临时文件类
        
        Args:
            pathsAndTable: 路径和表对象
            trans_param: 转换参数对象
        """
        self.pathsAndTable = pathsAndTable
        self.trans_param = trans_param
        self.exist_wind_names = multiprocessing.Manager().list()
        self.lock = multiprocessing.Manager().Lock()
        self.file_lock = multiprocessing.Manager().dict()

    def _save_to_tmp_csv_by_name(self, df: pd.DataFrame, name: str):
        """
        根据风机名称保存到临时CSV文件
        
        Args:
            df: 数据帧
            name: 风机名称
        """
        save_name = str(name) + '.csv'
        save_path = os.path.join(self.pathsAndTable.get_read_tmp_path(), save_name)
        create_file_path(save_path, is_file_path=True)

        with self.lock:
            if name in self.exist_wind_names:
                contains_name = True
            else:
                contains_name = False
                self.exist_wind_names.append(name)

            if contains_name:
                df.to_csv(save_path, index=False, encoding='utf8', mode='a',
                          header=False)
            else:
                df.to_csv(save_path, index=False, encoding='utf8')

    def save_merge_data(self, file_path: str):
        """
        保存合并数据
        
        Args:
            file_path: 文件路径
        """
        df = self.read_excel_to_df(file_path)
        if self.trans_param.wind_name_exec:
            if valid_eval(self.trans_param.wind_name_exec):
                exec_str = f"df['wind_turbine_number'].apply(lambda wind_name: {self.trans_param.wind_name_exec} )"
                df['wind_turbine_number'] = eval(exec_str)

        df = self.trans_df_cols(df)

        wind_names = set(df['wind_turbine_number'].values)
        cols = list(df.columns)
        cols.sort()
        for wind_name in wind_names:
            for col in df.columns:
                if col not in ['wind_turbine_number', 'time_stamp']:
                    csv_name = str(col) + ".csv"
                    exist_name = wind_name + '-' + csv_name
                    merge_path = self.pathsAndTable.get_merge_tmp_path(wind_name)
                    create_file_path(merge_path)

                    with self.lock:
                        if exist_name in self.exist_wind_names:
                            contains_name = True
                        else:
                            contains_name = False
                            self.exist_wind_names.append(exist_name)
                        save_path = os.path.join(merge_path, csv_name)
                        now_df = df[df['wind_turbine_number'] == wind_name][['time_stamp', col]]
                        if contains_name:
                            now_df.to_csv(save_path, index=False, encoding='utf-8', mode='a',
                                          header=False)
                        else:
                            now_df.to_csv(save_path, index=False, encoding='utf-8')

    def trans_df_cols(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        转换数据帧列名
        
        Args:
            df: 数据帧
        
        Returns:
            转换后的数据帧
        """
        if self.trans_param.is_vertical_table:
            pass
        else:
            # 转换字段
            same_col = {}
            if self.trans_param.cols_tran:
                cols_tran = self.trans_param.cols_tran
                real_cols_trans = dict()
                for k, v in cols_tran.items():
                    if v and not v.startswith("$"):
                        if v not in real_cols_trans.keys():
                            real_cols_trans[v] = k
                        else:
                            value = real_cols_trans[v]
                            if value in same_col.keys():
                                same_col[value].append(k)
                            else:
                                same_col[value] = [k]
                    if v and v.find('|') > -1:
                        vs = v.split('|')
                        for s in vs:
                            if s not in real_cols_trans.keys():
                                real_cols_trans[s] = k
                            else:
                                value = real_cols_trans[s]
                                if value in same_col.keys():
                                    same_col[value].append(k)
                                else:
                                    same_col[value] = [k]

                df.rename(columns=real_cols_trans, inplace=True)

                # 添加使用同一个excel字段的值
                for key in same_col.keys():
                    for col in same_col[key]:
                        df[col] = df[key]

                del_keys = set(df.columns) - set(cols_tran.keys())

                for key in del_keys:
                    df.drop(key, axis=1, inplace=True)

        return df

    def df_save_to_tmp_file(self, df: pd.DataFrame = pd.DataFrame()):
        """
        保存数据帧到临时文件
        
        Args:
            df: 数据帧
        """
        df = self.trans_df_cols(df)

        df = del_blank(df, ['wind_turbine_number'])
        df = df[df['time_stamp'].isna() == False]
        if self.trans_param.wind_name_exec and not self.trans_param.merge_columns:
            if valid_eval(self.trans_param.wind_name_exec):
                exec_str = f"df['wind_turbine_number'].apply(lambda wind_name: {self.trans_param.wind_name_exec} )"
                df['wind_turbine_number'] = eval(exec_str)

        self.save_to_tmp_csv(df)

    def save_to_tmp_csv(self, df: pd.DataFrame):
        """
        保存到临时CSV文件
        
        Args:
            df: 数据帧
        """
        names = set(df['wind_turbine_number'].values)
        if names:
            debug("开始保存", str(names), "到临时文件", df.shape)

            for name in names:
                self._save_to_tmp_csv_by_name(df[df['wind_turbine_number'] == name], name)
            del df
            debug("保存", str(names), "到临时文件成功, 风机数量", len(names))

    def merge_df(self, dir_path: str) -> pd.DataFrame:
        """
        合并数据帧
        
        Args:
            dir_path: 目录路径
        
        Returns:
            合并后的数据帧
        """
        all_files = read_excel_files(dir_path)
        wind_turbine_number = os.path.basename(dir_path)
        df = pd.DataFrame()
        for file in all_files:
            now_df = read_file_to_df(file)
            now_df['wind_turbine_number'] = wind_turbine_number
            now_df.dropna(subset=['time_stamp'], inplace=True)
            now_df.drop_duplicates(subset=['time_stamp'], inplace=True)
            now_df.set_index(keys=['time_stamp', 'wind_turbine_number'], inplace=True)
            df = pd.concat([df, now_df], axis=1)
        df.reset_index(inplace=True)

        self.df_save_to_tmp_file(df)

        return df

    def read_file_and_save_tmp(self):
        """
        读取文件并保存到临时文件
        """
        all_files = read_excel_files(self.pathsAndTable.get_excel_tmp_path())
        split_count = use_files_get_max_cpu_count(all_files)
        # 限制最大进程数
        split_count = min(split_count, ParallelProcessing.MAX_PROCESSES)
        all_arrays = split_array(all_files, split_count)

        if self.trans_param.merge_columns:
            for index, arr in enumerate(all_arrays):
                try:
                    with multiprocessing.Pool(split_count) as pool:
                        pool.starmap(self.save_merge_data, [(ar,) for ar in arr])

                except Exception as e:
                    error(traceback.format_exc())
                    message = "整理临时文件,系统返回错误:" + str(e)
                    raise ValueError(message)

                update_trans_transfer_progress(self.pathsAndTable.id,
                                               round(20 + 20 * (index + 1) / len(all_arrays), 2),
                                               self.pathsAndTable.save_db)

            dirs = [os.path.join(self.pathsAndTable.get_merge_tmp_path(), dir_name) for dir_name in
                    os.listdir(self.pathsAndTable.get_merge_tmp_path())]
            if dirs:
                dir_total_size = get_dir_size(dirs[0])
                # 限制最大进程数
                split_count = min(dir_total_size, ParallelProcessing.MAX_PROCESSES)
                all_arrays = split_array(dirs, split_count)
                for index, arr in enumerate(all_arrays):
                    try:
                        with multiprocessing.Pool(split_count) as pool:
                            pool.starmap(self.merge_df, [(ar,) for ar in arr])

                    except Exception as e:
                        error(traceback.format_exc())
                        message = "整理临时文件,系统返回错误:" + str(e)
                        raise ValueError(message)

                    update_trans_transfer_progress(self.pathsAndTable.id,
                                                   round(20 + 30 * (index + 1) / len(all_arrays), 2),
                                                   self.pathsAndTable.save_db)

        else:
            for index, arr in enumerate(all_arrays):
                try:
                    with multiprocessing.Pool(split_count) as pool:
                        dfs = pool.starmap(self.read_excel_to_df, [(ar,) for ar in arr])
                    for df in dfs:
                        self.df_save_to_tmp_file(df)
                except Exception as e:
                    error(traceback.format_exc())
                    message = "整理临时文件,系统返回错误:" + str(e)
                    raise ValueError(message)

                update_trans_transfer_progress(self.pathsAndTable.id,
                                               round(20 + 30 * (index + 1) / len(all_arrays), 2),
                                               self.pathsAndTable.save_db)

    def read_excel_to_df(self, file_path: str) -> pd.DataFrame:
        """
        读取Excel文件到数据帧
        
        Args:
            file_path: 文件路径
        
        Returns:
            数据帧
        """
        read_cols = [v.split(",")[0] for k, v in self.trans_param.cols_tran.items() if v and not v.startswith("$")]

        trans_dict = {}
        for k, v in self.trans_param.cols_tran.items():
            if v and not str(v).startswith("$"):
                trans_dict[v] = k

        if self.trans_param.is_vertical_table:
            vertical_cols = self.trans_param.vertical_cols
            df = read_file_to_df(file_path, vertical_cols, trans_cols=self.trans_param.vertical_cols)
            df = df[df[self.trans_param.vertical_key].isin(read_cols)]
            df.rename(columns={self.trans_param.cols_tran['wind_turbine_number']: 'wind_turbine_number',
                               self.trans_param.cols_tran['time_stamp']: 'time_stamp'}, inplace=True)
            df[self.trans_param.vertical_key] = df[self.trans_param.vertical_key].map(trans_dict).fillna(
                df[self.trans_param.vertical_key])

            return df

        else:
            trans_dict = dict()
            trans_cols = []
            for k, v in self.trans_param.cols_tran.items():
                if v and v.startswith("$") or v.find(",") > 0:
                    trans_dict[v] = k

                if v.find("|") > -1:
                    vs = v.split("|")
                    trans_cols.extend(vs)
                else:
                    trans_cols.append(v)
            trans_cols = list(set(trans_cols))
            if self.trans_param.merge_columns:
                df = read_file_to_df(file_path, trans_cols=trans_cols, not_find_header='ignore',
                                     resolve_col_prefix=self.trans_param.resolve_col_prefix)
            else:
                if self.trans_param.need_valid_cols:
                    if self.trans_param.resolve_col_prefix:
                        df = read_file_to_df(file_path, trans_cols=trans_cols,
                                             resolve_col_prefix=self.trans_param.resolve_col_prefix)
                    else:
                        df = read_file_to_df(file_path, read_cols, trans_cols=trans_cols,
                                             resolve_col_prefix=self.trans_param.resolve_col_prefix)
                else:
                    df = read_file_to_df(file_path, trans_cols=trans_cols,
                                         resolve_col_prefix=self.trans_param.resolve_col_prefix)

            # 处理列名前缀问题
            if self.trans_param.resolve_col_prefix:
                columns_dict = dict()
                if valid_eval(self.trans_param.resolve_col_prefix):
                    for column in df.columns:
                        columns_dict[column] = eval(self.trans_param.resolve_col_prefix)
                    df.rename(columns=columns_dict, inplace=True)

            if self.trans_param.merge_columns:
                select_cols = [self.trans_param.cols_tran['wind_turbine_number'],
                               self.trans_param.cols_tran['time_stamp'],
                               'wind_turbine_number', 'time_stamp']
                select_cols.extend(trans_cols)

                rename_dict = dict()
                wind_turbine_number_col = self.trans_param.cols_tran['wind_turbine_number']
                if wind_turbine_number_col.find("|") > -1:
                    cols = wind_turbine_number_col.split("|")
                    for col in cols:
                        rename_dict[col] = 'wind_turbine_number'

                time_stamp_col = self.trans_param.cols_tran['time_stamp']
                if time_stamp_col.find("|") > -1:
                    cols = time_stamp_col.split("|")
                    for col in cols:
                        rename_dict[col] = 'time_stamp'

                df.rename(columns=rename_dict, inplace=True)

                for col in df.columns:
                    if col not in select_cols:
                        del df[col]

            for k, v in trans_dict.items():
                if k.startswith("$file"):
                    file = ".".join(os.path.basename(file_path).split(".")[0:-1])
                    if k == "$file":
                        ks = k.split("|")
                        bool_contains = False
                        for k_data in ks:
                            if k_data in df.columns or v in df.columns:
                                bool_contains = True
                        if not bool_contains:
                            df[v] = str(file)

                    elif k.startswith("$file["):
                        ks = k.split("|")
                        bool_contains = False
                        for k_data in ks:
                            if k_data in df.columns or v in df.columns:
                                bool_contains = True
                        if not bool_contains:
                            datas = str(ks[0].replace("$file", "").replace("[", "").replace("]", "")).split(":")
                            if len(datas) != 2:
                                raise Exception("字段映射出现错误 :" + str(trans_dict))
                            df[v] = str(file[int(datas[0]):int(datas[1])]).strip()
                    elif k.startswith("$file.split"):
                        ks = k.split("|")
                        bool_contains = False
                        for k_data in ks:
                            if k_data in df.columns or v in df.columns:
                                bool_contains = True
                        if not bool_contains:
                            datas = str(ks[0]).replace("$file.split(", "").replace(")", "").split(",")
                            split_str = str(datas[0])
                            split_index = int(datas[1])
                            df[v] = str(file.split(split_str)[split_index])

                elif k.find("$file_date") > 0:
                    datas = str(k.split(",")[1].replace("$file_date", "").replace("[", "").replace("]", "")).split(":")
                    if len(datas) != 2:
                        raise Exception("字段映射出现错误 :" + str(trans_dict))
                    file = ".".join(os.path.basename(file_path).split(".")[0:-1])
                    date_str = str(file[int(datas[0]):int(datas[1])]).strip()
                    df[v] = df[k.split(",")[0]].apply(lambda x: date_str + " " + str(x))

                elif k.startswith("$folder"):
                    folder = file_path
                    ks = k.split("|")
                    bool_contains = False
                    for k_data in ks:
                        if k_data in df.columns or v in df.columns:
                            bool_contains = True
                    if not bool_contains:
                        cengshu = int(str(ks[0].replace("$folder", "").replace("[", "").replace("]", "")))
                        for i in range(cengshu):
                            folder = os.path.dirname(folder)
                        df[v] = str(str(folder).split(os.sep)[-1]).strip()
                elif k.startswith("$sheet_name"):
                    df[v] = df['sheet_name']

            if 'time_stamp' not in df.columns:
                cols_trans = [i for i in self.trans_param.cols_tran['time_stamp'].split('|')]
                cols_dict = dict()
                for col in cols_trans:
                    cols_dict[col] = 'time_stamp'
                df.rename(columns=cols_dict, inplace=True)

            if 'wind_turbine_number' not in df.columns:
                cols_trans = [i for i in self.trans_param.cols_tran['wind_turbine_number'].split('|')]
                cols_dict = dict()
                for col in cols_trans:
                    cols_dict[col] = 'wind_turbine_number'

                df.rename(columns=cols_dict, inplace=True)

            return df

    def run(self):
        """
        """
        info("开始保存数据到临时文件")
        begin = datetime.datetime.now()
        self.read_file_and_save_tmp()
        update_trans_transfer_progress(self.pathsAndTable.id, 50,
                                       self.pathsAndTable.save_db)
        info("保存数据到临时文件结束,耗时:", datetime.datetime.now() - begin)