import numpy as np import time import pickle import os import datetime import pandas as pd import sys sys.path.append(os.path.dirname(os.path.abspath(__file__))+ '/..') from model.partition_tree import PartitionTree from model.join_eval import JoinEvaluator from model.query_example import synetic_queries base_dir= os.path.dirname(os.path.abspath(__file__)) """ The core partitioning algorithm class Attribute: partition_tree is constructed partition tree. """ class PartitionAlgorithm: ''' The partition algorithms, inlcuding NORA, QdTree and kd-tree. ''' def __init__(self, block_size = 10000, benchmark = 'tpch', table_name = 'lineitem'): self.block_size = block_size self.partition_tree = None self.benchmark = benchmark self.table_name = table_name self.machine_num=10 def load_data(self): ''' Load data from the data source. ''' # 读取表的阈值 metadata = pickle.load(open(f'{base_dir}/../dataset/{self.benchmark}/metadata.pkl', 'rb')) self.used_columns = metadata[self.table_name]['numeric_columns'] # 获取表域 table_min_domains,table_max_domains = [],[] for _, col_range in metadata[self.table_name]['ranges'].items(): # table_domains.append([ # datetime.datetime.combine(col_range['min'], datetime.datetime.min.time()).timestamp() if isinstance(col_range['min'], datetime.date) else float(col_range['min']), # datetime.datetime.combine(col_range['max'], datetime.datetime.min.time()).timestamp() if isinstance(col_range['max'], datetime.date) else float(col_range['max']), # ]) # datetime类型转化为'yyyy-mm-dd'的形式，其他类别不变 table_min_domains.append( col_range['min'].strftime('%Y-%m-%d') if isinstance(col_range['min'], datetime.date) else float(col_range['min']) ) table_max_domains.append( col_range['max'].strftime('%Y-%m-%d') if isinstance(col_range['max'], datetime.date) else float(col_range['max']) ) self.table_domains = table_min_domains+table_max_domains self.column_width=metadata[self.table_name]['width'] # 读取csv数据 data = pd.read_csv(f'{base_dir}/../dataset/{self.benchmark}/{self.table_name}.csv', usecols=self.used_columns) self.tabledata = data.values def load_join_query(self,using_example=False,join_indeuced='MTO'): ''' Load join query from the data source. ''' metadata = pickle.load(open(f'{base_dir}/../dataset/{self.benchmark}/metadata.pkl', 'rb')) if using_example: query_dict=synetic_queries else: if join_indeuced=='MTO': query_dict=pickle.load(open(f'{base_dir}/../queryset/{self.benchmark}/mto_queries.pkl','rb')) elif join_indeuced=='PAC': query_dict=pickle.load(open(f'{base_dir}/../queryset/{self.benchmark}/pac_queries.pkl','rb')) else: query_dict=pickle.load(open(f'{base_dir}/../queryset/{self.benchmark}/paw_queries.pkl','rb')) join_freqs={table:{} for table in metadata.keys()} join_query_vectors=[] # for queryid, item in query_dict.items(): # query_vectors={} # for table, col_range in item['ranges'].items(): # min_vector,max_vector=[],[] # for col in col_range: # min_vector.append( # col_range[col]['min'].strftime('%Y-%m-%d') if isinstance(col_range[col]['min'], datetime.date) else col_range[col]['min'] if isinstance(col_range[col]['min'], str) else float(col_range[col]['min']) # ) # max_vector.append( # col_range[col]['max'].strftime('%Y-%m-%d') if isinstance(col_range[col]['max'], datetime.date) else col_range[col]['max'] if isinstance(col_range[col]['max'], str) else float(col_range[col]['max']) # ) # if min_vector: # query_vectors[table]=min_vector+max_vector for queryid, item in query_dict.items(): join_vector={'vectors':[],'join_relations':[]} for join_info in item['join_info']: for join_preds in join_info['join_keys']: left_cond,right_cond=join_preds.split('=') left_table,left_col=left_cond.split('.') left_col=metadata[left_table]['numeric_columns'].index(left_col) right_table,right_col=right_cond.split('.') right_col=metadata[right_table]['numeric_columns'].index(right_col) join_dict={left_table:left_col,right_table:right_col} if left_table not in item['ranges']: item['ranges'][left_table]={} if right_table not in item['ranges']: item['ranges'][right_table]={} join_freqs[left_table][left_col] = join_freqs[left_table].setdefault(left_col, 0) + 1 join_freqs[right_table][right_col] = join_freqs[right_table].setdefault(right_col, 0) + 1 if join_dict not in join_vector['join_relations']: join_vector['join_relations'].append(join_dict) query_vectors={} for table, col_ranges in item['ranges'].items(): min_vector,max_vector=[],[] for col_name,col_domain in metadata[table]['ranges'].items(): if col_name not in col_ranges: min_vector.append(col_domain['min'].strftime('%Y-%m-%d') if isinstance(col_domain['min'], datetime.date) else col_domain['min'] if isinstance(col_domain['min'], str) else float(col_domain['min'])) max_vector.append(col_domain['max'].strftime('%Y-%m-%d') if isinstance(col_domain['max'], datetime.date) else col_domain['max'] if isinstance(col_domain['max'], str) else float(col_domain['max'])) else: min_vector.append( col_ranges[col_name]['min'].strftime('%Y-%m-%d') if isinstance(col_ranges[col_name]['min'], datetime.date) else col_ranges[col_name]['min'] if isinstance(col_ranges[col_name]['min'], str) else float(col_ranges[col_name]['min']) ) max_vector.append( col_ranges[col_name]['max'].strftime('%Y-%m-%d') if isinstance(col_ranges[col_name]['max'], datetime.date) else col_ranges[col_name]['max'] if isinstance(col_ranges[col_name]['max'], str) else float(col_ranges[col_name]['max']) ) if min_vector: query_vectors[table]=min_vector+max_vector join_vector['vectors']=query_vectors join_query_vectors.append(join_vector) self.join_queries = join_query_vectors # join_freqs按照 value的值对key进行降序排列 self.join_freqs= {table:dict(sorted(join_freqs[table].items(), key=lambda item: item[1], reverse=True)) for table in join_freqs.keys()} def load_query(self,using_example=False,join_indeuced='MTO'): ''' Load query from the data source. ''' if using_example: query_dict=synetic_queries else: if join_indeuced=='MTO': query_dict=pickle.load(open(f'{base_dir}/../queryset/{self.benchmark}/mto_queries.pkl','rb')) elif join_indeuced=='PAC': query_dict=pickle.load(open(f'{base_dir}/../queryset/{self.benchmark}/pac_queries.pkl','rb')) else: query_dict=pickle.load(open(f'{base_dir}/../queryset/{self.benchmark}/paw_queries.pkl','rb')) metadata = pickle.load(open(f'{base_dir}/../dataset/{self.benchmark}/metadata.pkl', 'rb')) query_vectors=[] # for queryid, item in query_dict.items(): # min_vector,max_vector=[],[] # for table, col_range in item['ranges'].items(): # if table == self.table_name: # for col in col_range: # min_vector.append( # col_range[col]['min'].strftime('%Y-%m-%d') if isinstance(col_range[col]['min'], datetime.date) else col_range[col]['min'] if isinstance(col_range[col]['min'], str) else float(col_range[col]['min']) # ) # max_vector.append( # col_range[col]['max'].strftime('%Y-%m-%d') if isinstance(col_range[col]['max'], datetime.date) else col_range[col]['max'] if isinstance(col_range[col]['max'], str) else float(col_range[col]['max']) # ) # if min_vector: query_vectors.append(min_vector+max_vector) for queryid, item in query_dict.items(): min_vector,max_vector=[],[] for table, col_ranges in item['ranges'].items(): if table!=self.table_name: continue for col_name,col_domain in metadata[table]['ranges'].items(): if col_name not in col_ranges: min_vector.append(col_domain['min'].strftime('%Y-%m-%d') if isinstance(col_domain['min'], datetime.date) else col_domain['min'] if isinstance(col_domain['min'], str) else float(col_domain['min'])) max_vector.append(col_domain['max'].strftime('%Y-%m-%d') if isinstance(col_domain['max'], datetime.date) else col_domain['max'] if isinstance(col_domain['max'], str) else float(col_domain['max'])) else: min_vector.append( col_ranges[col_name]['min'].strftime('%Y-%m-%d') if isinstance(col_ranges[col_name]['min'], datetime.date) else col_ranges[col_name]['min'] if isinstance(col_ranges[col_name]['min'], str) else float(col_ranges[col_name]['min']) ) max_vector.append( col_ranges[col_name]['max'].strftime('%Y-%m-%d') if isinstance(col_ranges[col_name]['max'], datetime.date) else col_ranges[col_name]['max'] if isinstance(col_ranges[col_name]['max'], str) else float(col_ranges[col_name]['max']) ) if min_vector: query_vectors.append(min_vector+max_vector) # join_vector=[] # for join_info in item['join_info']: # for join_preds in join_info['join_keys']: # left_cond,right_cond=join_preds.split('=') # left_table,left_col=left_cond.split('.') # right_table,right_col=right_cond.split('.') # if left_table==table_name: # join_vector.append(self.used_columns.index(left_col)) # if right_table==table_name: # join_vector.append(self.used_columns.index(right_col)) # if min_vector: # join_query_vectors.append(list(set(join_vector))) self.queries = query_vectors # self.join_queries = join_query_vectors # MTO layout def InitializeWithQDT(self): ''' # should I also store the candidate cut positions in Partition Node ? The dimension of queries should match the dimension of boundary and dataset! ''' start_time = time.time() num_dims=len(self.used_columns) boundary=self.table_domains self.partition_tree = PartitionTree(num_dims, boundary) self.partition_tree.name='QDT' self.partition_tree.used_columns=self.used_columns self.partition_tree.column_width=self.column_width self.partition_tree.pt_root.node_size = len(self.tabledata) self.partition_tree.pt_root.dataset = self.tabledata self.partition_tree.pt_root.queryset = self.queries # assume all queries overlap with the boundary self.__QDT(self.block_size) end_time = time.time() print(f"{self.table_name} Build Time (s):", end_time-start_time) self.partition_tree.build_time = end_time - start_time self.partition_tree.save_tree(f'{base_dir}/../layouts/{self.benchmark}/{self.table_name}-QDT.pkl') def __QDT(self, block_size): ''' the QdTree partition algorithm ''' CanSplit = True print_s=True while CanSplit: CanSplit = False # for leaf in self.partition_tree.get_leaves(): leaves = self.partition_tree.get_leaves() #print("# number of leaf nodes:",len(leaves)) for leaf in leaves: # print("current leaf node id:",leaf.nid, "leaf node dataset size:",len(leaf.dataset)) if leaf.node_size < 2 * block_size: continue candidate_cuts = leaf.get_candidate_cuts() # get best candidate cut position skip, max_skip, max_skip_split_dim, max_skip_split_value = 0, -1, 0, 0 for split_dim, split_value in candidate_cuts: valid,skip,_,_ = leaf.if_split(split_dim, split_value, block_size) if valid and skip > max_skip: max_skip = skip max_skip_split_dim = split_dim max_skip_split_value = split_value if max_skip >= 0: # if the cost become smaller, apply the cut if print_s: print("QDTREE CUT!") print_s=False child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, max_skip_split_dim, max_skip_split_value) # print(" Split on node id:", leaf.nid) CanSplit = True def LogicalJoinTree(self,join_col,join_max_depth=3): ''' # should I also store the candidate cut positions in Partition Node ? The dimension of queries should match the dimension of boundary and dataset! ''' start_time = time.time() tree_path=f'{base_dir}/../layouts/{self.benchmark}/logical/{self.table_name}-{join_col}-tree.pkl' # if not os.path.exists(tree_path): for i in range(1,self.machine_num): if 2**i>=self.machine_num: join_max_depth=i+1 break num_dims=len(self.used_columns) boundary=self.table_domains self.partition_tree = PartitionTree(num_dims, boundary) self.partition_tree.name='LJT' self.partition_tree.join_attr = self.used_columns.index(join_col) self.partition_tree.join_depth = join_max_depth self.partition_tree.used_columns=self.used_columns self.partition_tree.column_width=self.column_width self.partition_tree.pt_root.node_size = len(self.tabledata) self.partition_tree.pt_root.dataset = self.tabledata self.partition_tree.pt_root.queryset = self.queries canSplit = True cur_depth = 0 # 仅在join_cols不为空时执行top-layer tree construction while canSplit: canSplit = False leaves = self.partition_tree.get_leaves() cur_depth += 1 for leaf in leaves: if leaf.node_size < 2 * self.block_size: continue if cur_depth<=join_max_depth: split_dim=self.partition_tree.join_attr split_value = np.median(leaf.dataset[:, split_dim]) valid, skip, _, _ = leaf.if_split(split_dim, split_value, self.block_size) if valid: child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, split_dim, split_value) canSplit = True else: candidate_cuts = leaf.get_candidate_cuts(extended=True) # get best candidate cut position skip, max_skip, max_skip_split_dim, max_skip_split_value = 0, -1, 0, 0 for split_dim, split_value in candidate_cuts: valid, skip, _, _ = leaf.if_split(split_dim, split_value, self.block_size) if valid and skip > max_skip: max_skip = skip max_skip_split_dim = split_dim max_skip_split_value = split_value if max_skip >= 0: # if the cost become smaller, apply the cut child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, max_skip_split_dim, max_skip_split_value) canSplit = True end_time = time.time() self.partition_tree.build_time = end_time - start_time self.partition_tree.save_tree(tree_path) # AdaptDB layout def InitializeWithADP(self,join_depth=3): ''' # should I also store the candidate cut positions in Partition Node ? The dimension of queries should match the dimension of boundary and dataset! ''' join_cols=list(self.join_freqs[self.table_name].keys()) num_dims=len(self.used_columns) boundary=self.table_domains self.partition_tree = PartitionTree(num_dims, boundary) self.partition_tree.name='AdaptDB' self.partition_tree.join_attr = join_cols self.partition_tree.join_depth = join_depth self.partition_tree.used_columns=self.used_columns self.partition_tree.column_width=self.column_width self.partition_tree.pt_root.node_size = len(self.tabledata) self.partition_tree.pt_root.dataset = self.tabledata self.partition_tree.pt_root.queryset = self.queries # assume all queries overlap with the boundary start_time = time.time() self.__ADP(self.block_size, join_cols, join_depth) end_time = time.time() print(f"{self.table_name} Build Time (s):", end_time-start_time) self.partition_tree.build_time = end_time - start_time self.partition_tree.save_tree(f'{base_dir}/../layouts/{self.benchmark}/{self.table_name}-ADP.pkl') def __ADP(self, block_size, join_cols, join_depth): print_s=True # top-layer tree construction canSplit = True cur_depth = 1 # 仅在join_cols不为空时执行top-layer tree construction if join_cols: while canSplit: canSplit = False leaves = self.partition_tree.get_leaves() cur_depth += 1 if cur_depth>join_depth: break for leaf in leaves: if leaf.node_size < 2 * block_size or leaf.depth < cur_depth: continue # 选择剩余可分配数据最大的属性作为下一个切割点 # temp_allocations = self.partition_tree.allocations.copy() split_dim=join_cols[0] split_value = np.median(leaf.dataset[:, split_dim]) valid, skip, _, _ = leaf.if_split(split_dim, split_value, block_size) if valid: if print_s: print("AdaptDB CUT!") print_s=False child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, split_dim, split_value) child_node1.depth = leaf.depth + 1 child_node2.depth = leaf.depth + 1 # self.partition_tree.allocations[split_dim] -= 2.0 / pow(2, cur_depth) canSplit = True # bottom-layer tree construction CanSplit = True while CanSplit: CanSplit = False # for leaf in self.partition_tree.get_leaves(): leaves = self.partition_tree.get_leaves() # print("# number of leaf nodes:",len(leaves)) for leaf in leaves: # print("current leaf node id:",leaf.nid, "leaf node dataset size:",len(leaf.dataset)) if leaf.node_size < 2 * block_size: continue candidate_cuts = leaf.get_candidate_cuts() # get best candidate cut position skip, max_skip, max_skip_split_dim, max_skip_split_value = 0, -1, 0, 0 for split_dim, split_value in candidate_cuts: valid, skip, _, _ = leaf.if_split(split_dim, split_value, block_size) if valid and skip > max_skip: max_skip = skip max_skip_split_dim = split_dim max_skip_split_value = split_value if max_skip >= 0: # if the cost become smaller, apply the cut child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, max_skip_split_dim, max_skip_split_value) CanSplit = True # join tree layout def InitializeWithJT(self,join_depth=3): ''' # should I also store the candidate cut positions in Partition Node ? The dimension of queries should match the dimension of boundary and dataset! ''' join_cols=list(self.join_freqs[self.table_name].keys()) # 1->2 2->4 3->8 4->16 for i in range(1,self.machine_num): if 2**i>=self.machine_num: join_depth=i+1 break num_dims=len(self.used_columns) boundary=self.table_domains self.partition_tree = PartitionTree(num_dims, boundary) self.partition_tree.name='JoinTree' self.partition_tree.join_attr = join_cols self.partition_tree.join_depth = join_depth self.partition_tree.used_columns=self.used_columns self.partition_tree.column_width=self.column_width self.partition_tree.pt_root.node_size = len(self.tabledata) self.partition_tree.pt_root.dataset = self.tabledata self.partition_tree.pt_root.queryset = self.queries # assume all queries overlap with the boundary start_time = time.time() self.__JT(self.block_size, join_cols, join_depth) end_time = time.time() print(f"{self.table_name} Build Time (s):", end_time-start_time) self.partition_tree.build_time = end_time - start_time self.partition_tree.save_tree(f'{base_dir}/../layouts/{self.benchmark}/{self.table_name}-JT.pkl') def __JT(self, block_size, join_cols, join_depth): print_s=True # top-layer tree construction canSplit = True cur_depth = 1 # 仅在join_cols不为空时执行top-layer tree construction if join_cols: while canSplit: canSplit = False leaves = self.partition_tree.get_leaves() cur_depth += 1 if cur_depth>join_depth: break for leaf in leaves: if leaf.node_size < 2 * block_size or leaf.depth < cur_depth: continue # 选择剩余可分配数据最大的属性作为下一个切割点 # temp_allocations = self.partition_tree.allocations.copy() split_dim=join_cols[0] split_value = np.median(leaf.dataset[:, split_dim]) valid, skip, _, _ = leaf.if_split(split_dim, split_value, block_size) if valid: if print_s: print("JoinTree CUT!") print_s=False child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, split_dim, split_value) child_node1.depth = leaf.depth + 1 child_node2.depth = leaf.depth + 1 # self.partition_tree.allocations[split_dim] -= 2.0 / pow(2, cur_depth) canSplit = True # bottom-layer tree construction CanSplit = True while CanSplit: CanSplit = False # for leaf in self.partition_tree.get_leaves(): leaves = self.partition_tree.get_leaves() # print("# number of leaf nodes:",len(leaves)) for leaf in leaves: # print("current leaf node id:",leaf.nid, "leaf node dataset size:",len(leaf.dataset)) if leaf.node_size < 2 * block_size: continue candidate_cuts = leaf.get_candidate_cuts() # get best candidate cut position skip, max_skip, max_skip_split_dim, max_skip_split_value = 0, -1, 0, 0 for split_dim, split_value in candidate_cuts: valid, skip, _, _ = leaf.if_split(split_dim, split_value, block_size) if valid and skip > max_skip: max_skip = skip max_skip_split_dim = split_dim max_skip_split_value = split_value if max_skip >= 0: # if the cost become smaller, apply the cut child_node1, child_node2 = self.partition_tree.apply_split(leaf.nid, max_skip_split_dim, max_skip_split_value) CanSplit = True def evaluate_single_table_access_cost(self): self.partition_tree.evaluate_query_cost(self.queries,len(self.tabledata),self.column_width,self.used_columns) def evaluate_multiple_table_access_cost(self,trees): group_type=0 if trees[list(trees.keys())[0]].name=='AdaptDB': group_type=0 elif trees[trees.keys()[0]].name=='JoinTree': group_type=1 tot_cost={} for join_query in self.join_queries: # 分别处理每条查询 if join_query['join_relations']: for join_op in join_query['join_relations']: # 分别处理查询中的每个join操作 join_queryset,joined_cols,joined_trees=[],[],[] for join_table,join_col in join_op.items(): # overlapped_leaf_ids=trees[join_table].query_single(join_query['vectors'][join_table]) join_queryset.append(join_query['vectors'][join_table]) joined_cols.append(join_col) joined_trees.append(trees[join_table]) join_eval = JoinEvaluator(join_queryset,joined_cols,joined_trees) #jobv2 hyper_read_cost, shuffle_read_cost, hyper_cost_list = join_eval.compute_total_shuffle_hyper_cost(group_type) tot_cost+=hyper_read_cost+shuffle_read_cost else: for tablename in join_query['vectors']: tot_cost[tablename]=tot_cost.setdefault(tablename,0)+trees[tablename].evaluate_query_cost([join_query['vectors'][tablename]]) return tot_cost