re-common 10.0.39__py3-none-any.whl → 10.0.41__py3-none-any.whl

re_common/v2/baselibrary/tools/hdfs_base_processor.py

@@ -1,204 +1,204 @@
-import abc
-import asyncio
-import sys
-from concurrent.futures import ProcessPoolExecutor
-import gzip
-
-import multiprocessing
-from pathlib import Path
-
-from io import BytesIO
-import time
-from typing import Awaitable, Callable, Any, Generator, List, Literal, Union
-
-from hdfs import InsecureClient
-
-
-from re_common.v2.baselibrary.tools.resume_tracker import ResumeTracker
-
-
-_pool = None
-_loop = None
-
-
-class HDFSBaseProcessor(abc.ABC):
-    def __init__(
-        self,
-        hdfs_dir: str,
-        hdfs_url: str = "http://VIP-DC-MASTER-2:9870",
-        hdfs_user: str = "root",
-        db_path: Union[str, Path] = "processed_files.db",
-        concurrency: int = 50,
-        batch_size: int = 50,
-        encoding: str = "utf-8",
-        read_mode: Literal["all", "stream"] = "all",
-        retries: int = 3,
-        pool_factory: Callable[[], Awaitable[Any]] = None,
-        max_processes: int = None,  # 添加多进程支持
-        result_dir: str = None,
-    ):
-        self.hdfs_dir = hdfs_dir
-        self.hdfs_url = hdfs_url
-        self.hdfs_user = hdfs_user
-        self.tracker = ResumeTracker(db_path)
-        self.concurrency = concurrency
-        self.batch_size = batch_size
-        self.encoding = encoding
-        self.read_mode = read_mode
-        self.retries = retries
-        self.result_dir = result_dir
-        self.pool_factory = pool_factory
-        self.max_processes = max_processes or max(multiprocessing.cpu_count() - 1, 1)  # 默认使用CPU核心数-1
-        self._client = None
-
-        self.tracker.init_db()
-
-    @property
-    def client(self):
-        if self._client is None:
-            self._client = InsecureClient(self.hdfs_url, user=self.hdfs_user)
-        return self._client
-
-    async def _get_pool(self):
-        if self.pool_factory is None:
-            return None
-        global _pool
-        if _pool is None:
-            _pool = await self.pool_factory()
-        return _pool
-
-    def _list_gz_files(self) -> List[str]:
-        """列出 HDFS 目录中的所有 gzip 文件"""
-        return [
-            f"{self.hdfs_dir}/{file[0]}"
-            for file in self.client.list(self.hdfs_dir, status=True)
-            if file[0].endswith(".gz")
-        ]
-
-    def _count_total_lines(self, gz_file_path: str) -> int:
-        with self.client.read(gz_file_path) as hdfs_file:
-            with gzip.GzipFile(fileobj=hdfs_file) as gz:
-                return sum(1 for _ in gz)
-
-    def _batch_read_gz_stream(self, gz_file_path: str) -> Generator[List[str], Any, None]:
-        """流式读取gz文件，分批yield返回"""
-        with self.client.read(gz_file_path) as hdfs_file:
-            with gzip.GzipFile(fileobj=hdfs_file) as gz:
-                while True:
-                    lines = []
-                    for _ in range(self.batch_size):
-                        try:
-                            line = next(gz)
-                            if line.strip():  # 移除空行
-                                lines.append(line.decode(self.encoding))  # 解码
-                        except StopIteration:  # 文件已读完
-                            break
-                    if not lines:
-                        break
-                    yield lines
-
-    def _batch_read_gz_all(self, gz_file_path: str) -> List[List[str]]:
-        """一次读取gz文件全部内容，二维数组批量返回"""
-        with self.client.read(gz_file_path) as reader:  # 以二进制模式读取
-            compressed_data = reader.read()  # 读取压缩数据
-            with gzip.GzipFile(fileobj=BytesIO(compressed_data)) as gz_file:  # 解压缩
-                content = gz_file.read().decode(self.encoding)  # 解码为字符串
-                print(f"文件读取成功: {gz_file_path}")
-                lines = [i for i in content.split("\n") if i.strip()]
-                batch_lines = [lines[i : i + self.batch_size] for i in range(0, len(lines), self.batch_size)]
-                return batch_lines
-
-    def _batch_read_gz(self, gz_file_path: str):
-        # 这里根据不同的配置选用不同的读取文件的方法
-        if self.read_mode == "stream":
-            return self._batch_read_gz_stream(gz_file_path)
-        else:
-            return self._batch_read_gz_all(gz_file_path)
-
-    def _generate_write_data(self, results):
-        for res in results:
-            yield str(res) + "\n"
-
-    def _print_progress(self, file_path, processed_lines, total_lines, start_time, processing_start_time):
-        elapsed_time = time.perf_counter() - start_time  # 已用时间
-        processing_time = time.perf_counter() - processing_start_time  # 本次处理时间
-        avg_processing_time = (
-            (elapsed_time * 1000) / processed_lines if processed_lines > 0 else float("inf")
-        )  # 平均每条数据的处理时间（毫秒）
-        # 估算剩余时间
-        remaining_time = (
-            ((avg_processing_time / 1000) * (total_lines - processed_lines)) if processed_lines > 0 else float("inf")
-        )
-        # 显示进度信息
-        print(
-            f"文件: {file_path} 总进度: {processed_lines}/{total_lines} 行 | "
-            f"已用时间: {elapsed_time:.2f}秒 | 本次处理时间: {processing_time:.2f}秒 | "
-            f"预估剩余时间: {remaining_time:.2f}秒 | 平均每条处理时间: {avg_processing_time:.2f}毫秒"
-        )
-
-    def _print_final_progress(self, file_path, processed_lines, total_lines, start_time):
-        final_elapsed_time = time.perf_counter() - start_time  # 最终已用时间
-        print(
-            f"文件: {file_path} 处理完成 | 总进度: {processed_lines}/{total_lines} 行 | "
-            f"总已用时间: {final_elapsed_time:.2f}秒 | "
-            f"平均每条处理时间: {(final_elapsed_time * 1000) / processed_lines:.2f}毫秒"
-            if processed_lines > 0
-            else "处理无数据"
-        )
-
-    @abc.abstractmethod
-    async def _process_file(self, hdfs_file_path, process_func):
-        pass
-
-    async def _retry_process_file(self, hdfs_file_path, process_func):
-        """带重试机制的文件处理"""
-        retry_count = 0
-        while retry_count < self.retries:
-            try:
-                if self.tracker.is_processed(hdfs_file_path):
-                    print(f"跳过已处理文件: {hdfs_file_path}")
-                    return True
-                await self._process_file(hdfs_file_path, process_func)
-                self.tracker.mark_processed(hdfs_file_path)  # 标记文件已处理
-                return True  # 成功处理后退出
-            except Exception as e:
-                retry_count += 1
-                print(f"处理文件 {hdfs_file_path} 时发生错误: {e}，正在重试 {retry_count}/{self.retries}")
-                await asyncio.sleep(2**retry_count)
-        print(f"处理文件 {hdfs_file_path} 失败，达到重试上限")
-        return False
-
-    def _process_file_wrapper(self, args):
-        """为多进程执行准备的同步包装函数"""
-        hdfs_file_path, process_func = args
-        if sys.platform == "win32":
-            loop = asyncio.get_event_loop()
-            return loop.run_until_complete(self._retry_process_file(hdfs_file_path, process_func))
-        else:
-            global _loop
-            if _loop is None:
-                _loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(_loop)
-            return _loop.run_until_complete(self._retry_process_file(hdfs_file_path, process_func))
-
-    async def _run_multi_process(self, gz_files, process_func):
-        """多进程并发运行文件处理任务"""
-        args_list = [(file_path, process_func) for file_path in gz_files]
-        with ProcessPoolExecutor(max_workers=self.max_processes) as executor:
-            # return executor.map(self._process_file_wrapper, args_list)
-            loop = asyncio.get_running_loop()
-            self._client = None  # 避免连接对象无法序列化导致卡死
-            tasks = [loop.run_in_executor(executor, self._process_file_wrapper, args) for args in args_list]
-            results = await asyncio.gather(*tasks)
-
-        if all(results):
-            # 处理完成后清理断点记录
-            self.tracker.clear_processed_items()
-            print(f"已清空断点记录: {self.tracker.db_path}")
-            return results
-        else:
-            raise Exception("部分或全部文件处理失败")
-
-    @abc.abstractmethod
-    async def map(self, process_func: Callable[[Any, Any], Awaitable[Any]]) -> None:
-        pass
+import abc
+import asyncio
+import sys
+from concurrent.futures import ProcessPoolExecutor
+import gzip
+
+import multiprocessing
+from pathlib import Path
+
+from io import BytesIO
+import time
+from typing import Awaitable, Callable, Any, Generator, List, Literal, Union
+
+from hdfs import InsecureClient
+
+
+from re_common.v2.baselibrary.tools.resume_tracker import ResumeTracker
+
+
+_pool = None
+_loop = None
+
+
+class HDFSBaseProcessor(abc.ABC):
+    def __init__(
+        self,
+        hdfs_dir: str,
+        hdfs_url: str = "http://VIP-DC-MASTER-2:9870",
+        hdfs_user: str = "root",
+        db_path: Union[str, Path] = "processed_files.db",
+        concurrency: int = 50,
+        batch_size: int = 50,
+        encoding: str = "utf-8",
+        read_mode: Literal["all", "stream"] = "all",
+        retries: int = 3,
+        pool_factory: Callable[[], Awaitable[Any]] = None,
+        max_processes: int = None,  # 添加多进程支持
+        result_dir: str = None,
+    ):
+        self.hdfs_dir = hdfs_dir
+        self.hdfs_url = hdfs_url
+        self.hdfs_user = hdfs_user
+        self.tracker = ResumeTracker(db_path)
+        self.concurrency = concurrency
+        self.batch_size = batch_size
+        self.encoding = encoding
+        self.read_mode = read_mode
+        self.retries = retries
+        self.result_dir = result_dir
+        self.pool_factory = pool_factory
+        self.max_processes = max_processes or max(multiprocessing.cpu_count() - 1, 1)  # 默认使用CPU核心数-1
+        self._client = None
+
+        self.tracker.init_db()
+
+    @property
+    def client(self):
+        if self._client is None:
+            self._client = InsecureClient(self.hdfs_url, user=self.hdfs_user)
+        return self._client
+
+    async def _get_pool(self):
+        if self.pool_factory is None:
+            return None
+        global _pool
+        if _pool is None:
+            _pool = await self.pool_factory()
+        return _pool
+
+    def _list_gz_files(self) -> List[str]:
+        """列出 HDFS 目录中的所有 gzip 文件"""
+        return [
+            f"{self.hdfs_dir}/{file[0]}"
+            for file in self.client.list(self.hdfs_dir, status=True)
+            if file[0].endswith(".gz")
+        ]
+
+    def _count_total_lines(self, gz_file_path: str) -> int:
+        with self.client.read(gz_file_path) as hdfs_file:
+            with gzip.GzipFile(fileobj=hdfs_file) as gz:
+                return sum(1 for _ in gz)
+
+    def _batch_read_gz_stream(self, gz_file_path: str) -> Generator[List[str], Any, None]:
+        """流式读取gz文件，分批yield返回"""
+        with self.client.read(gz_file_path) as hdfs_file:
+            with gzip.GzipFile(fileobj=hdfs_file) as gz:
+                while True:
+                    lines = []
+                    for _ in range(self.batch_size):
+                        try:
+                            line = next(gz)
+                            if line.strip():  # 移除空行
+                                lines.append(line.decode(self.encoding))  # 解码
+                        except StopIteration:  # 文件已读完
+                            break
+                    if not lines:
+                        break
+                    yield lines
+
+    def _batch_read_gz_all(self, gz_file_path: str) -> List[List[str]]:
+        """一次读取gz文件全部内容，二维数组批量返回"""
+        with self.client.read(gz_file_path) as reader:  # 以二进制模式读取
+            compressed_data = reader.read()  # 读取压缩数据
+            with gzip.GzipFile(fileobj=BytesIO(compressed_data)) as gz_file:  # 解压缩
+                content = gz_file.read().decode(self.encoding)  # 解码为字符串
+                print(f"文件读取成功: {gz_file_path}")
+                lines = [i for i in content.split("\n") if i.strip()]
+                batch_lines = [lines[i : i + self.batch_size] for i in range(0, len(lines), self.batch_size)]
+                return batch_lines
+
+    def _batch_read_gz(self, gz_file_path: str):
+        # 这里根据不同的配置选用不同的读取文件的方法
+        if self.read_mode == "stream":
+            return self._batch_read_gz_stream(gz_file_path)
+        else:
+            return self._batch_read_gz_all(gz_file_path)
+
+    def _generate_write_data(self, results):
+        for res in results:
+            yield str(res) + "\n"
+
+    def _print_progress(self, file_path, processed_lines, total_lines, start_time, processing_start_time):
+        elapsed_time = time.perf_counter() - start_time  # 已用时间
+        processing_time = time.perf_counter() - processing_start_time  # 本次处理时间
+        avg_processing_time = (
+            (elapsed_time * 1000) / processed_lines if processed_lines > 0 else float("inf")
+        )  # 平均每条数据的处理时间（毫秒）
+        # 估算剩余时间
+        remaining_time = (
+            ((avg_processing_time / 1000) * (total_lines - processed_lines)) if processed_lines > 0 else float("inf")
+        )
+        # 显示进度信息
+        print(
+            f"文件: {file_path} 总进度: {processed_lines}/{total_lines} 行 | "
+            f"已用时间: {elapsed_time:.2f}秒 | 本次处理时间: {processing_time:.2f}秒 | "
+            f"预估剩余时间: {remaining_time:.2f}秒 | 平均每条处理时间: {avg_processing_time:.2f}毫秒"
+        )
+
+    def _print_final_progress(self, file_path, processed_lines, total_lines, start_time):
+        final_elapsed_time = time.perf_counter() - start_time  # 最终已用时间
+        print(
+            f"文件: {file_path} 处理完成 | 总进度: {processed_lines}/{total_lines} 行 | "
+            f"总已用时间: {final_elapsed_time:.2f}秒 | "
+            f"平均每条处理时间: {(final_elapsed_time * 1000) / processed_lines:.2f}毫秒"
+            if processed_lines > 0
+            else "处理无数据"
+        )
+
+    @abc.abstractmethod
+    async def _process_file(self, hdfs_file_path, process_func):
+        pass
+
+    async def _retry_process_file(self, hdfs_file_path, process_func):
+        """带重试机制的文件处理"""
+        retry_count = 0
+        while retry_count < self.retries:
+            try:
+                if self.tracker.is_processed(hdfs_file_path):
+                    print(f"跳过已处理文件: {hdfs_file_path}")
+                    return True
+                await self._process_file(hdfs_file_path, process_func)
+                self.tracker.mark_processed(hdfs_file_path)  # 标记文件已处理
+                return True  # 成功处理后退出
+            except Exception as e:
+                retry_count += 1
+                print(f"处理文件 {hdfs_file_path} 时发生错误: {e}，正在重试 {retry_count}/{self.retries}")
+                await asyncio.sleep(2**retry_count)
+        print(f"处理文件 {hdfs_file_path} 失败，达到重试上限")
+        return False
+
+    def _process_file_wrapper(self, args):
+        """为多进程执行准备的同步包装函数"""
+        hdfs_file_path, process_func = args
+        if sys.platform == "win32":
+            loop = asyncio.get_event_loop()
+            return loop.run_until_complete(self._retry_process_file(hdfs_file_path, process_func))
+        else:
+            global _loop
+            if _loop is None:
+                _loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(_loop)
+            return _loop.run_until_complete(self._retry_process_file(hdfs_file_path, process_func))
+
+    async def _run_multi_process(self, gz_files, process_func):
+        """多进程并发运行文件处理任务"""
+        args_list = [(file_path, process_func) for file_path in gz_files]
+        with ProcessPoolExecutor(max_workers=self.max_processes) as executor:
+            # return executor.map(self._process_file_wrapper, args_list)
+            loop = asyncio.get_running_loop()
+            self._client = None  # 避免连接对象无法序列化导致卡死
+            tasks = [loop.run_in_executor(executor, self._process_file_wrapper, args) for args in args_list]
+            results = await asyncio.gather(*tasks)
+
+        if all(results):
+            # 处理完成后清理断点记录
+            self.tracker.clear_processed_items()
+            print(f"已清空断点记录: {self.tracker.db_path}")
+            return results
+        else:
+            raise Exception("部分或全部文件处理失败")
+
+    @abc.abstractmethod
+    async def map(self, process_func: Callable[[Any, Any], Awaitable[Any]]) -> None:
+        pass

re_common/v2/baselibrary/tools/hdfs_bulk_processor.py

@@ -1,67 +1,67 @@
-from pathlib import Path
-import time
-from typing import Any, Awaitable, Callable, List
-from re_common.v2.baselibrary.tools.concurrency import AsyncTaskPool
-from re_common.v2.baselibrary.tools.hdfs_base_processor import HDFSBaseProcessor
-
-
-class HDFSBulkProcessor(HDFSBaseProcessor):
-    def _flat_map(self, results):
-        # return itertools.chain.from_iterable(chunked_results)
-        for res in results:
-            if isinstance(res, list):
-                yield from res
-            else:
-                yield res
-
-    async def _process_file(self, hdfs_file_path, process_func):
-        start_time = time.perf_counter()
-        total_lines = self._count_total_lines(hdfs_file_path)
-        processed_lines = 0
-        pool = await self._get_pool()
-
-        tasks = []
-        for lines in self._batch_read_gz(hdfs_file_path):
-            # 处理读取到的批次数据
-            if lines:
-                tasks.append(process_func(lines, pool))  # 将批次数据传递给处理函数并收集任务
-                processed_lines += len(lines)  # 更新已处理行数
-        results = await AsyncTaskPool(self.concurrency).run(tasks)
-
-        if self.result_dir is not None:
-            self.client.write(
-                self.result_dir.rstrip("/") + f"/{Path(hdfs_file_path).stem}",
-                data=self._generate_write_data(self._flat_map(results)),
-                overwrite=True,
-                encoding=self.encoding,
-            )
-
-        # 最终进度显示
-        self._print_final_progress(hdfs_file_path, processed_lines, total_lines, start_time)
-
-    async def map(self, process_func: Callable[[List[str], Any], Awaitable[Any]]) -> None:
-        gz_files = self._list_gz_files()
-        await self._run_multi_process(gz_files, process_func)
-
-
-# async def test_func(lines: List[str], pool):
-#     pass
-
-
-# async def main():
-#     processor = HDFSBulkProcessor(
-#         "/xx/xx",
-#         db_path=Path(__file__).parent / "test_bulk.db",
-#         concurrency=200,
-#         batch_size=1000,
-#         pool_factory=get_pool,
-#         max_processes=2,
-#         result_dir="/xx/xx_res",
-#     )
-#     # processor.tracker.mark_many_processed(f"/xx/xx/part-{num:05d}.gz" for num in range(0, 6000))
-
-#     await processor.map(test_func)
-
-
-# if __name__ == "__main__":
-#     asyncio.run(main())
+from pathlib import Path
+import time
+from typing import Any, Awaitable, Callable, List
+from re_common.v2.baselibrary.tools.concurrency import AsyncTaskPool
+from re_common.v2.baselibrary.tools.hdfs_base_processor import HDFSBaseProcessor
+
+
+class HDFSBulkProcessor(HDFSBaseProcessor):
+    def _flat_map(self, results):
+        # return itertools.chain.from_iterable(chunked_results)
+        for res in results:
+            if isinstance(res, list):
+                yield from res
+            else:
+                yield res
+
+    async def _process_file(self, hdfs_file_path, process_func):
+        start_time = time.perf_counter()
+        total_lines = self._count_total_lines(hdfs_file_path)
+        processed_lines = 0
+        pool = await self._get_pool()
+
+        tasks = []
+        for lines in self._batch_read_gz(hdfs_file_path):
+            # 处理读取到的批次数据
+            if lines:
+                tasks.append(process_func(lines, pool))  # 将批次数据传递给处理函数并收集任务
+                processed_lines += len(lines)  # 更新已处理行数
+        results = await AsyncTaskPool(self.concurrency).run(tasks)
+
+        if self.result_dir is not None:
+            self.client.write(
+                self.result_dir.rstrip("/") + f"/{Path(hdfs_file_path).stem}",
+                data=self._generate_write_data(self._flat_map(results)),
+                overwrite=True,
+                encoding=self.encoding,
+            )
+
+        # 最终进度显示
+        self._print_final_progress(hdfs_file_path, processed_lines, total_lines, start_time)
+
+    async def map(self, process_func: Callable[[List[str], Any], Awaitable[Any]]) -> None:
+        gz_files = self._list_gz_files()
+        await self._run_multi_process(gz_files, process_func)
+
+
+# async def test_func(lines: List[str], pool):
+#     pass
+
+
+# async def main():
+#     processor = HDFSBulkProcessor(
+#         "/xx/xx",
+#         db_path=Path(__file__).parent / "test_bulk.db",
+#         concurrency=200,
+#         batch_size=1000,
+#         pool_factory=get_pool,
+#         max_processes=2,
+#         result_dir="/xx/xx_res",
+#     )
+#     # processor.tracker.mark_many_processed(f"/xx/xx/part-{num:05d}.gz" for num in range(0, 6000))
+
+#     await processor.map(test_func)
+
+
+# if __name__ == "__main__":
+#     asyncio.run(main())