xlin 0.1.24__py2.py3-none-any.whl → 0.1.26__py2.py3-none-any.whl

xlin/multiprocess_mapping.py

@@ -15,21 +15,25 @@ from xlin.util import ls
 
 
 def element_mapping(
-    iterator: List[Any],
+    iterator: list[Any],
     mapping_func: Callable[[Any], Tuple[bool, Any]],
     use_multiprocessing=True,
     thread_pool_size=int(os.getenv("THREAD_POOL_SIZE", 5)),
 ):
     rows = []
+    # 转换为列表以获取长度，用于进度条显示
+    items = list(iterator)
+    total = len(items)
+
     if use_multiprocessing:
         pool = ThreadPool(thread_pool_size)
-        results = pool.map(mapping_func, iterator)
-        pool.close()
-        for ok, row in results:
+        # 使用imap替代map，结合tqdm显示进度
+        for ok, row in tqdm(pool.imap(mapping_func, items), total=total, desc="Processing"):
             if ok:
                 rows.append(row)
+        pool.close()
     else:
-        for row in tqdm(iterator):
+        for row in tqdm(items, desc="Processing"):
             ok, row = mapping_func(row)
             if ok:
                 rows.append(row)
@@ -37,8 +41,8 @@ def element_mapping(
 
 
 def batch_mapping(
-    iterator: List[Any],
-    mapping_func: Callable[[List[Any]], Tuple[bool, List[Any]]],
+    iterator: list[Any],
+    mapping_func: Callable[[list[Any]], Tuple[bool, list[Any]]],
     use_multiprocessing=True,
     thread_pool_size=int(os.getenv("THREAD_POOL_SIZE", 5)),
     batch_size=4,
@@ -68,59 +72,124 @@ def dataframe_with_row_mapping(
     return df
 
 
-def multiprocessing_mapping_jsonlist(
-    jsonlist: List[Any],
-    output_path: Optional[Union[str, Path]],
-    partial_func,
+def xmap(
+    jsonlist: list[Any],
+    work_func: Union[Callable[[Any], dict], Callable[[list[Any]], list[dict]]],
+    output_path: Optional[Union[str, Path]]=None,  # 输出路径，None表示不缓存
     batch_size=multiprocessing.cpu_count(),
     cache_batch_num=1,
-    thread_pool_size=int(os.getenv("THREAD_POOL_SIZE", 5)),
+    thread_pool_size=int(os.getenv("THREAD_POOL_SIZE", 8)),
+    use_process_pool=True,  # CPU密集型任务时设为True
+    preserve_order=True,  # 是否保持结果顺序
+    chunksize=None,  # 自动计算最佳分块大小
+    retry_count=0,  # 失败重试次数
+    force_overwrite=False,  # 是否强制覆盖输出文件
+    is_batch_work_func=False,  # 是否批量处理函数
 ):
-    """mapping a column to another column
+    """高效处理JSON列表，支持多进程/多线程
 
     Args:
-        df (DataFrame): [description]
-        output_path (Path): 数据量大的时候需要缓存
-        partial_func (function): (Dict[str, str]) -> Dict[str, str]
+        jsonlist (list[Any]): 要处理的JSON对象列表
+        output_path (Optional[Union[str, Path]]): 输出路径，None表示不缓存
+        work_func (Callable): 处理函数，接收dict返回dict
+        batch_size (int): 批处理大小
+        cache_batch_num (int): 缓存批次数量
+        thread_pool_size (int): 线程/进程池大小
+        use_process_pool (bool): 是否使用进程池(CPU密集型任务)
+        preserve_order (bool): 是否保持结果顺序
+        chunksize (Optional[int]): 单个任务分块大小，None为自动计算
+        retry_count (int): 任务失败重试次数
     """
     need_caching = output_path is not None
-    tmp_list, output_list = list(), list()
+    output_list = []
     start_idx = 0
+
+    # 自动计算最佳chunksize
+    if chunksize is None:
+        chunksize = max(1, min(batch_size // thread_pool_size, 100))
+
+    # 处理缓存
     if need_caching:
         output_path = Path(output_path)
         if output_path.exists():
-            output_list = load_json_list(output_path)
-            start_idx = len(output_list)
-            logger.warning(f"Cache found {output_path} has {start_idx} rows. This process will continue at row index {start_idx}.")
-            logger.warning(f"缓存 {output_path} 存在 {start_idx} 行. 本次处理将从第 {start_idx} 行开始.")
+            if force_overwrite:
+                logger.warning(f"强制覆盖输出文件: {output_path}")
+                output_path.unlink()
+            else:
+                output_list = load_json_list(output_path)
+                start_idx = len(output_list)
+                logger.info(f"继续处理: 已有{start_idx}条记录，共{len(jsonlist)}条")
         else:
             output_path.parent.mkdir(parents=True, exist_ok=True)
-    pool = ThreadPool(thread_pool_size)
-    logger.debug(f"pool size: {thread_pool_size}, cpu count: {multiprocessing.cpu_count()}")
-    start_time = time.time()
-    last_save_time = start_time
-    for i, line in tqdm(list(enumerate(jsonlist))):
-        if i < start_idx:
-            continue
-        tmp_list.append(line)
-        if len(tmp_list) == batch_size:
-            results = pool.map(partial_func, tmp_list)
-            output_list.extend([x for x in results])
-            tmp_list = list()
-        if need_caching and (i // batch_size) % cache_batch_num == 0:
-            current_time = time.time()
-            if current_time - last_save_time < 3:
-                # 如果多进程处理太快，为了不让 IO 成为瓶颈拉慢进度，不足 3 秒的批次都忽略，也不缓存中间结果
-                last_save_time = current_time
-                continue
-            save_json_list(output_list, output_path)
-            last_save_time = time.time()
-    if len(tmp_list) > 0:
-        results = pool.map(partial_func, tmp_list)
-        output_list.extend([x for x in results])
-    pool.close()
+
+    # 选择线程池或进程池
+    if use_process_pool:
+        pool_cls = multiprocessing.Pool
+        logger.info(f"使用进程池(ProcessPool)，适用于CPU密集型任务")
+    else:
+        pool_cls = ThreadPool
+        logger.info(f"使用线程池(ThreadPool)，适用于IO密集型任务")
+
+    with pool_cls(thread_pool_size) as pool:
+        logger.info(f"池大小: {thread_pool_size}, 批处理大小: {batch_size}, 分块大小: {chunksize}")
+
+        # 准备要处理的数据
+        remaining_items = jsonlist[start_idx:]
+        total_items = len(remaining_items)
+
+        # 批量处理逻辑
+        def process_batch(items_batch, retry_remaining=retry_count):
+            try:
+                if is_batch_work_func:
+                    # 批量处理函数
+                    return work_func(items_batch)
+                else:
+                    # 选择合适的映射方法
+                    map_func = pool.imap_unordered if not preserve_order else pool.imap
+                    return list(map_func(work_func, items_batch, chunksize))
+            except Exception as e:
+                if retry_remaining > 0:
+                    logger.warning(f"批处理失败，重试中 ({retry_count-retry_remaining+1}/{retry_count}): {e}")
+                    return process_batch(items_batch, retry_remaining - 1)
+                else:
+                    logger.error(f"批处理失败: {e}")
+                    raise
+
+        # 处理数据
+        with tqdm(total=total_items, desc="处理数据", unit="项") as pbar:
+            # 跳过已处理的项目
+            pbar.update(start_idx)
+
+            # 分批处理
+            for i in range(0, total_items, batch_size):
+                batch = remaining_items[i : i + batch_size]
+
+                # 处理当前批次
+                batch_start_time = time.time()
+                results = process_batch(batch)
+                batch_time = time.time() - batch_start_time
+
+                # 更新结果
+                output_list.extend(results)
+                pbar.update(len(batch))
+
+                # 性能统计
+                items_per_second = len(batch) / batch_time if batch_time > 0 else 0
+                pbar.set_postfix_str(f"速率: {items_per_second:.1f}项/秒")
+
+                # 缓存逻辑
+                if need_caching and (i // batch_size) % cache_batch_num == 0:
+                    # 仅当处理速度足够慢时才保存缓存，避免IO成为瓶颈
+                    if batch_time > 3 or i + batch_size >= total_items:
+                        save_json_list(output_list, output_path)
+                        logger.debug(f"已保存{len(output_list)}条记录到{output_path}")
+
+    # 最终保存
     if need_caching:
         save_json_list(output_list, output_path)
+    drop_count = len(jsonlist) - len(output_list)
+    logger.info(f"处理完成，共处理{len(jsonlist)}条记录" + ", 丢弃{len(jsonlist) - len(output_list)}条记录" if drop_count > 0 else "")
+
     return output_list
 
 
@@ -186,35 +255,6 @@ def multiprocessing_mapping(
     return output_df, output_list
 
 
-def continue_run(
-    jsonfiles: List[str],
-    save_dir: str,
-    mapping_func,
-    load_func=load_json,
-    save_func=save_json,
-    batch_size=1024,
-    cache_size=8,
-):
-    save_dir: Path = Path(save_dir)
-    save_dir.mkdir(parents=True, exist_ok=True)
-    new_jsonfiles = []
-    for jsonfile in ls(jsonfiles):
-        jsonlist = load_func(jsonfile)
-        output_filepath = save_dir / jsonfile.name
-        for row in jsonlist:
-            row["来源"] = jsonfile.name
-        new_jsonlist = multiprocessing_mapping_jsonlist(
-            jsonlist,
-            output_filepath,
-            mapping_func,
-            batch_size,
-            cache_size,
-        )
-        save_func(new_jsonlist, output_filepath)
-        new_jsonfiles.append(output_filepath)
-    return new_jsonfiles
-
-
 def dataframe_mapping(
     df: pd.DataFrame,
     row_func: Callable[[dict], dict],

{xlin-0.1.24.dist-info → xlin-0.1.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xlin
-Version: 0.1.24
+Version: 0.1.26
 Summary: toolbox for LinXueyuan
 License: MIT
 Author: LinXueyuanStdio

{xlin-0.1.24.dist-info → xlin-0.1.26.dist-info}/RECORD

@@ -2,14 +2,14 @@ xlin/__init__.py,sha256=MWWCNPgJFS_oV2US52ULa4yg4Ku61qjn40NVKqcp9-c,248
 xlin/ischinese.py,sha256=Ia9IMQ6q-UHkdLwqS70L1fTnfSPbluFrv_I1UqsKquo,293
 xlin/jsonl.py,sha256=IDRydHh2x-8iAGCxt9ScK2wfNLNA40PxNxR5hhr4v6k,7903
 xlin/metric.py,sha256=N7wJ35y-C-IaBr1I1CJ_37lTG7gA69zmn9Xg6xSwKoI,1690
-xlin/multiprocess_mapping.py,sha256=dRXQoLaG1dK_qZ8B3bJblV0RKM2gqIeSW1EaOZbIdD0,14251
+xlin/multiprocess_mapping.py,sha256=M_d8G-apwZFcWuQxYexOpN8a4SX3OMUuqdJ4JHHgQHw,16480
 xlin/read_as_dataframe.py,sha256=MqY57L7Wp9UoWTRlZLSBKQNaZa-dKw51-ufrKvHKf8s,9041
 xlin/statistic.py,sha256=2DCUgzf7xkMFH4Pk9v82bFDNeSxCTjwPh9Y4IPJBHCE,9300
 xlin/timing.py,sha256=XMT8dMcMolOMohDvAZOIM_BAiPMREhGQKnO1kc5s6PU,998
 xlin/util.py,sha256=TTWJaqF5D_r-gAZ_fj0kyHomvCagjwHXQZ2OPSgwd54,10976
 xlin/xls2xlsx.py,sha256=uSmXcDvIhi5Sq0LGidMXy0wErNBXdjaoa6EftYVjTXs,947
 xlin/yaml.py,sha256=kICi7G3Td5q2MaSXXt85qNTWoHMgjzt7pvn7r3C4dME,183
-xlin-0.1.24.dist-info/LICENSE,sha256=60ys6rRtc1dZOP8UjSUr9fAqhZudT3WpKe5WbMCralM,1066
-xlin-0.1.24.dist-info/METADATA,sha256=Xm8lU2owbTVPZt1G9baKLiWMIqUZB3THuoCgGGJrQ58,1098
-xlin-0.1.24.dist-info/WHEEL,sha256=IrRNNNJ-uuL1ggO5qMvT1GGhQVdQU54d6ZpYqEZfEWo,92
-xlin-0.1.24.dist-info/RECORD,,
+xlin-0.1.26.dist-info/LICENSE,sha256=60ys6rRtc1dZOP8UjSUr9fAqhZudT3WpKe5WbMCralM,1066
+xlin-0.1.26.dist-info/METADATA,sha256=suP1OEbJd7a0-mhqkh3thQVNBTGUi2amRQVaSGbYafk,1098
+xlin-0.1.26.dist-info/WHEEL,sha256=IrRNNNJ-uuL1ggO5qMvT1GGhQVdQU54d6ZpYqEZfEWo,92
+xlin-0.1.26.dist-info/RECORD,,