chronaxis.results.merge_runs
源码基线:8904c2f744951857659abd40aaa3535d57faf53c
1"""在固定内存和文件句柄数量内归并多份已排序结果,减少中间文件重复写入。23输入为同一ResultSpool持有的1至64份run及一个全新目标路径;输出保留原JSONL4的全部JSONL字节。每路只保留当前一行,按原排序键及输入路序稳定归并。5读取与写入沿用owner的严格校验和累计磁盘配额,RAM先预约到owner._merge_cost。6本模块不删除输入、不修改结果模型,也不持有全量run清单或独立预算。ExitStack7逐个关闭实际文件,未关闭句柄和未退费用仍保留在owner._merge_run_resources,8调用方可通过同一owner.close重试;关闭失败不能被当成归并成功或资源已经释放。9"""1011from contextlib import ExitStack12from heapq import heappop, heappush13from pathlib import Path1415from .budget import _reserve_owned_ram, _release_owned_ram16from .errors import AnalysisError, ResultFormatError171819_MAX_ROUTES = 64 # 仅限制本次打开的输入数,不按全部run数量建立路径列表。20_OUTPUT_WORK_BYTES = 65536 # 两个8KiB输出缓冲、输出FD和有限临时状态的余量。21_ROUTE_WORK_BYTES = 1024 # 每路Path/heap节点/ExitStack登记的保守有限描述费用。222324class _MergeResources:25 """保存本次所有实际文件与已付归并费,持续关闭失败也不丢真实责任。2627 handles每槽只持一个实际FD或None;pending先安装新开FD再扩展登记列表。28 ExitStack完成首轮关闭,失败后按仍存活的槽重试;只有全部FD与费用均清理才closed。29 """3031 def __init__(self, owner):32 """建立已预付的有限责任对象。参数: owner为真实暂存。返回: None;初始化实例状态。异常: 分配错误。"""33 self.owner = owner34 self.handles, self.pending, self.stack = [], None, ExitStack()35 self.drained = self.closed = False3637 def open(self, path, mode, buffering):38 """把新FD先安装到真实槽,再登记ExitStack回调,覆盖登记时取消。3940 参数: path/mode/buffering为原文件参数。返回: 本对象实际持有的句柄。41 异常: 原打开/登记/取消错误;已打开的FD仍由pending或handles持有。42 """43 self.pending = path.open(mode, buffering=buffering)44 self.handles.append(self.pending)45 position = len(self.handles) - 146 handle, self.pending = self.pending, None47 self.stack.callback(self._close_slot, position)48 return handle4950 def _close_slot(self, position):51 """实际关闭一个已登记FD,关闭后取消仍记录真实关闭事实。5253 参数: position为有限句柄槽。返回: None。异常: 原close错误;未关对象保原槽。54 """55 handle = self.handles[position]56 if handle is None:57 return58 try:59 handle.close()60 finally:61 if handle.closed:62 self.handles[position] = None63 handle = None6465 def close(self):66 """先关闭所有实际FD再退自有费用;失败后同一责任对象可重复关闭。6768 参数: 无。返回: None;closed后幂等。异常: 首个清理错误为主,其他错误附note;69 存活FD或未抵销费用时保owner._merge_run_resources,绝不释放借用余额。70 """71 if self.closed:72 return73 failure = handle = None74 try:75 # 新FD尚未完成列表登记时仍由pending持有,也必须先关闭才可归还费用。76 if self.pending is not None:77 try:78 self.pending.close()79 except BaseException as exc:80 failure = exc81 finally:82 if self.pending.closed:83 self.pending = None84 if not self.drained:85 try:86 self.stack.close()87 except BaseException as exc:88 if failure is None:89 failure = exc90 else:91 failure.add_note(f'归并ExitStack关闭附加失败: {type(exc).__name__}')92 finally:93 # ExitStack虽会耗尽回调,失败FD仍在handles真实槽,不从回调是否存在猜关闭。94 self.drained = True95 # 包括callback登记之前取消的FD;首轮有关闭故障时留待显式下一次close重试。96 if failure is None:97 for position, handle in enumerate(self.handles):98 if handle is not None:99 try:100 self._close_slot(position)101 except BaseException as exc:102 if failure is None:103 failure = exc104 else:105 failure.add_note(f'归并FD关闭附加失败: {type(exc).__name__}')106 handle = None107 unclosed = self.pending is not None or any(handle is not None for handle in self.handles)108 if not unclosed:109 try:110 _release_owned_ram(self.owner, self.owner._merge_cost, '_merge_cost')111 except BaseException as exc:112 if failure is None:113 failure = exc114 else:115 failure.add_note(f'归并费用清理附加失败: {type(exc).__name__}')116 finally:117 # 释放回调可能已更新账本再取消;按实际槽值确认,不重复抵销其他人的钱。118 if self.owner._merge_cost == 0:119 self.closed = True120 self.owner._merge_run_resources = None121 self.handles.clear()122 self.stack = None123 elif failure is None:124 failure = AnalysisError('归并文件尚未确认关闭,保留实际责任与费用')125 if failure is not None:126 if not self.closed:127 failure.result_merge_cleanup_owner = self128 raise failure129 finally:130 handle = None131132133def merge_runs(owner, paths_tuple, destination, *, table):134 """将同一暂存的有限有序输入稳定合并到新文件,失败仍保原错误与实际费用。135136 参数: owner为当前ResultSpool;paths_tuple为1至64个Path且无重复;destination137 为不存在的Path;table为固定表名或使用全局硬限的None。输入须按原键升序排列。138 返回: None;目标JSONL经真实关闭后才成功返回。139 异常: 配置/预算/行损坏、排序回退、文件错误和原取消;清理副错只附note,140 未退RAM仍留在同一owner._merge_cost,可由owner.close继续处理。141 边界: 相等键按输入路序稳定取出;每次写后先丢旧行再读该路下一行。142 """143 if (type(paths_tuple) is not tuple or not 1 <= len(paths_tuple) <= _MAX_ROUTES144 or any(not isinstance(path, Path) for path in paths_tuple)145 or not isinstance(destination, Path)):146 raise ValueError('多路归并需要1至64个Path组成的tuple及新目标Path')147 if owner._merge_cost or getattr(owner, '_merge_run_resources', None) is not None:148 raise AnalysisError('当前暂存仍持有前次归并费用,必须先完成清理')149 # row_limit只取成功编码事实,不能按损坏文件大小自行放宽读取范围。150 maximum = owner.max_row_bytes if table is None else owner.row_limit(table)151 reader_cost = owner.reader_cost if table is None else owner.reader_cost_for(table)152 # 单路复制及双路保原费用,维持已验证的最低可用预算;多路按统一公式预约。153 cost = (2 * reader_cost if len(paths_tuple) <= 2 else len(paths_tuple) * reader_cost154 + _OUTPUT_WORK_BYTES + _ROUTE_WORK_BYTES * len(paths_tuple))155 failure = resources = heap = inputs = entry = key = previous_key = None156 output = handle = None157 try:158 # 先登记实际费用,再建立路径集合、FD、heap和行;成功后取消仍由原owner负责。159 _reserve_owned_ram(owner, cost, '_merge_cost')160 if len(set(paths_tuple)) != len(paths_tuple) or destination in paths_tuple:161 raise ValueError('多路归并输入不得重复,目标不得复用输入路径')162 resources = _MergeResources(owner)163 owner._merge_run_resources = resources164 heap, inputs = [], []165 for path in paths_tuple:166 handle = resources.open(path, 'rb', buffering=8192)167 inputs.append(handle)168 output = resources.open(destination, 'xb', buffering=0)169 # 每路先取得一条;空run在当前文件EOF处结束。170 for route, handle in enumerate(inputs):171 entry = owner._read_entry(handle, maximum)172 if entry is not None:173 heappush(heap, (entry[0], route, entry))174 entry = None175 while heap:176 key, route, entry = heappop(heap)177 owner._write_bytes(output, entry[1])178 # 清旧payload后再读下一行;只留下比较顺序需要的有限键。179 previous_key, entry, key = key, None, None180 handle = inputs[route]181 entry = owner._read_entry(handle, maximum)182 if entry is not None:183 if entry[0] < previous_key:184 raise ResultFormatError('多路归并输入run的排序键回退')185 heappush(heap, (entry[0], route, entry))186 entry = previous_key = None187 except BaseException as exc:188 failure = exc189 raise190 finally:191 # 清掉大行与局部引用,实际FD仍在resources槽,持续失败不会因ExitStack耗尽回调而丢失。192 heap = inputs = entry = key = previous_key = None193 output = handle = None194 try:195 if resources is not None:196 resources.close()197 else:198 # 预约后、资源对象建立前的失败尚无FD,费用仍须按真实槽值抵销。199 _release_owned_ram(owner, owner._merge_cost, '_merge_cost')200 except BaseException as cleanup:201 if failure is None:202 raise203 failure.add_note(f'多路归并资源清理附加失败: {type(cleanup).__name__}')204 if resources is not None and not resources.closed:205 failure.result_merge_cleanup_owner = resources206 if owner._merge_cost:207 failure.result_spool_cleanup_owner = owner208