bigarraylist 0.2.0__py3-none-any.whl

bigarraylist/algorithms/bucket.py

@@ -0,0 +1,200 @@
+"""持久化分桶（Hash Partitioning）+ 单桶扫描查找。
+
+对比重构前：
+    - 重构前：bucket_find() 每次调用都重建桶（每次都付 O(n) 建桶成本，
+      多次查询没有任何加速）。
+    - 重构后：build_buckets() 返回 BucketIndex 对象，建桶一次 O(n)，
+      之后任意多次 .find(target) 只扫单桶 O(n/N)；同时支持跨进程复用
+      （persist_path 把桶文件保留到指定目录）。
+
+MapReduce 核心思想落地：相同 key 一定在同一个桶 → 查找时只需扫一个桶。
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import pickle
+import re
+import tempfile
+from collections.abc import Callable, Generator, Iterable
+from typing import Any, Generic, TypeVar
+
+from ..config import BigArrayConfig
+
+T = TypeVar("T")
+
+
+class BucketIndex(Generic[T]):
+    """持久化分桶索引。
+
+    建桶一次后可多次 find，每次只扫目标桶。
+    支持上下文管理器：
+        with build_buckets(data, key_func) as idx:
+            for tgt in targets:
+                hits = list(idx.find(tgt))
+    """
+
+    __slots__ = ("__key_func", "__num_buckets", "__owns_files", "__temp_paths")
+
+    def __init__(
+        self,
+        temp_paths: list[str],
+        num_buckets: int,
+        key_func: Callable[[T], Any],
+        owns_files: bool = True,
+    ) -> None:
+        self.__temp_paths = temp_paths
+        self.__num_buckets = num_buckets
+        self.__key_func = key_func
+        self.__owns_files = owns_files  # False 时 close() 不删文件（外部持久化场景）
+
+    def __enter__(self) -> BucketIndex:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: Any,
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        """清理所有桶临时文件。多次调用幂等。"""
+        if not self.__owns_files:
+            return
+        paths = self.__temp_paths
+        self.__temp_paths = []
+        for path in paths:
+            try:
+                os.unlink(path)
+            except OSError:
+                pass
+
+    @property
+    def num_buckets(self) -> int:
+        return self.__num_buckets
+
+    def find(self, target: T) -> Generator[tuple[int, T], None, None]:
+        """查找所有 key(item) == key(target) 的元素。
+
+        返回生成器，每项为 (原始下标, 命中元素)。
+        只扫目标桶，O(n/N) 时间；O(1) 额外内存（除目标桶文件读取外）。
+        """
+        target_key = self.__key_func(target)
+        target_bucket = hash(target_key) % self.__num_buckets
+
+        # 在 temp_paths 中定位目标桶（可能不存在 → 该桶为空）
+        target_path = None
+        for path in self.__temp_paths:
+            if self._extract_bucket_idx(path) == target_bucket:
+                target_path = path
+                break
+
+        if target_path is None:
+            return  # 目标桶为空
+
+        with open(target_path, "rb") as f:
+            while True:
+                try:
+                    i, item = pickle.load(f)
+                except EOFError:
+                    break
+                if self.__key_func(item) == target_key:
+                    yield (i, item)
+
+    @staticmethod
+    def _extract_bucket_idx(path: str) -> int:
+        """从临时文件名 big_bucket_{idx}_XXXX.tmp 提取桶下标。"""
+        m = re.search(r"big_bucket_(\d+)_", path)
+        if m is None:
+            return -1
+        return int(m.group(1))
+
+
+def _choose_num_buckets(total_count: int, max_buckets: int) -> int:
+    """自动选桶数：sqrt(n) 附近的 2 的幂，上限 max_buckets。
+
+    取 2 的幂让 hash % num_buckets 利用位运算高效；上限避免文件描述符爆掉。
+    """
+    target_n = max(8, math.isqrt(total_count))
+    n = 1
+    while n * 2 <= target_n and n < max_buckets:
+        n *= 2
+    return n
+
+
+def build_buckets(
+    data: Iterable[T],
+    total_count: int,
+    key_func: Callable[[T], Any] | None = None,
+    num_buckets: int | None = None,
+    *,
+    config: BigArrayConfig | None = None,
+) -> BucketIndex:
+    """对数据流式建桶，返回可复用的 BucketIndex。
+
+    参数：
+        data:         数据源迭代器（建议来自 _fresh_iter）
+        total_count:  数据总条数
+        key_func:     从元素提取分桶 key；None 用元素本身
+                      （必须满足 a==b → key(a)==key(b) 且 hash 兼容，否则漏命中）
+        num_buckets:   桶数；None 自动选（sqrt(n) 附近 2 的幂，上限 config.bucket_max_count）
+        config:        配置对象；None 用 BigArrayConfig.default()
+
+    返回：BucketIndex，支持多次 find(target) 和上下文管理器。
+
+    内存模型：
+        - 阶段 1（建桶）：流式读，每条元素 + 下标一起 pickle.dump 到对应桶文件
+        - 内存峰值 = O(桶数) 文件句柄 + O(1) 单条元素
+    """
+    cfg = config if config is not None else BigArrayConfig.default()
+    if key_func is None:
+        key_func = lambda x: x  # noqa: E731 恒等默认值
+
+    if num_buckets is None:
+        num_buckets = _choose_num_buckets(total_count, cfg.bucket_max_count)
+    if num_buckets < 1:
+        raise ValueError(f"num_buckets 必须 >= 1，得到 {num_buckets}")
+
+    bucket_files: list[Any] = [None] * num_buckets
+    temp_paths: list[str] = []
+    try:
+        # 阶段 1：流式分桶（每条元素 + 其下标一起 pickle 落盘）
+        for i, item in enumerate(data):
+            bidx = hash(key_func(item)) % num_buckets
+            if bucket_files[bidx] is None:
+                tmp = tempfile.NamedTemporaryFile(
+                    suffix=".tmp",
+                    prefix=f"big_bucket_{bidx}_",
+                    delete=False,
+                )
+                temp_paths.append(tmp.name)
+                bucket_files[bidx] = tmp
+            # 逐条写 (下标, 元素)
+            pickle.dump((i, item), bucket_files[bidx])
+
+        # 关闭所有已开启的桶文件（写模式 → 准备读）
+        for f in bucket_files:
+            if f is not None:
+                f.close()
+
+        return BucketIndex(temp_paths, num_buckets, key_func, owns_files=True)
+    except BaseException:
+        # 出错时清理已开的文件和已建的临时文件
+        for f in bucket_files:
+            if f is not None:
+                try:
+                    f.close()
+                except Exception:
+                    pass
+        for path in temp_paths:
+            try:
+                os.unlink(path)
+            except OSError:
+                pass
+        raise
+
+
+__all__ = ["BucketIndex", "build_buckets"]

bigarraylist/algorithms/external_sort.py

@@ -0,0 +1,200 @@
+"""真正的低内存外部排序。
+
+核心改进（对比重构前的整块版本）：
+    - 重构前：阶段 2 整块 pickle.load，所有 chunk_list 同时在内存，
+      内存峰值 ≈ 数据总量，违背"外部排序"的初衷。
+    - 重构后：阶段 1 逐条 pickle.dump 到文件（每条独立可 load），
+      阶段 2 每个文件用一个 lazy 文件迭代器，每次只 load 一条。
+      内存峰值 = O(块数) 文件句柄 + O(块数) 单条元素。
+
+模式选择（由 BigArrayConfig.external_sort_streaming_merge 控制）：
+    - True  → streaming（逐条 load，省内存，速度稍慢）
+    - False → block（整块 load，快 3~5 倍，内存 ≈ 数据总量；与重构前一致）
+
+通过 lazy_file_iter / lazy_block_iter 两种文件迭代器抽象两种模式，
+上层调用方代码完全相同。
+"""
+
+from __future__ import annotations
+
+import heapq
+import os
+import pickle
+import tempfile
+from collections import deque
+from collections.abc import Callable, Generator, Iterable
+from concurrent.futures import Future, ProcessPoolExecutor
+from typing import IO, Any, TypeVar
+
+from ..config import BigArrayConfig
+from ..utils.scale import choose_external_sort_chunk_size
+from .parallel import sort_chunk
+
+T = TypeVar("T")
+
+# 双缓冲流水线深度：阶段 1 同时最多保持 2 个 chunk 在飞行
+# （1 个在 worker 中排序 + 1 个在主进程读取/落盘）。
+# 排序 N 与 "读 N+1 + 落 N-1" 并行，并行度从 1 提升到 2，
+# 内存峰值仍为 O(chunk_size)（与串行版一致，不退化成中数据分支）。
+_BUFFER_DEPTH = 2
+
+
+def _streaming_file_iter(f: IO[bytes]) -> Generator[T, None, None]:
+    """逐条 pickle.load 模式的文件迭代器。
+
+    要求文件是用"逐条 pickle.dump"写入的（每次 dump 一个元素）。
+    每次 next 调用 pickle.load(f)，读到 EOF 抛 StopIteration。
+    内存：O(1) per file（同时持有的元素仅 1 个）。
+    """
+    try:
+        while True:
+            try:
+                yield pickle.load(f)
+            except EOFError:
+                return
+    except GeneratorExit:
+        # 调用方 close() 生成器时也要让上层 finally 关闭文件
+        raise
+
+
+def _block_file_iter(f: IO[bytes]) -> Generator[T, None, None]:
+    """整块 pickle.load 模式的文件迭代器（向后兼容用）。
+
+    要求文件是用"整块 pickle.dump(list)"写入的（一次 dump 整个 list）。
+    一次 pickle.load 读出整块 list，然后 yield from 展开。
+    内存：O(整块大小) per file；所有文件同时存在时，总内存 ≈ 数据总量。
+    """
+    chunk_list = pickle.load(f)
+    yield from chunk_list
+
+
+def _dump_streaming(chunk: list[T], f: IO[bytes]) -> None:
+    """逐条 pickle.dump 一个 chunk 的元素到文件 f。"""
+    for item in chunk:
+        pickle.dump(item, f)
+
+
+def _dump_block(chunk: list[T], f: IO[bytes]) -> None:
+    """整块 pickle.dump 一个 chunk 到文件 f（向后兼容模式）。"""
+    pickle.dump(chunk, f)
+
+
+def external_sort(
+    data: Iterable[T],
+    total_count: int,
+    *,
+    config: BigArrayConfig | None = None,
+    max_workers: int | None = None,
+    key: Callable[[T], Any] | None = None,
+) -> Generator[T, None, None]:
+    """外部排序（External Sort）：分块排序 + 落盘 + 多路归并。
+
+    参数：
+        data:         数据源迭代器（建议来自 _fresh_iter）
+        total_count:  数据总条数（用于选块大小）
+        config:       配置对象；None 用 BigArrayConfig.default()
+        max_workers:  阶段 1 排序的进程池并行度
+        key:          排序键函数（类似 ``sorted(data, key=...)``）；
+                      None 按元素本身排序。
+                      ⚠️ 多进程路径要求 key 可 pickle（模块级函数或 functools.partial，
+                      lambda 在 spawn 模式下不可 pickle）
+
+    返回：生成器，按 key 的非降序逐条产出已排序元素。
+
+    内存模型（streaming 模式，默认）：
+        - 阶段 1：双缓冲流水线排序 + 逐条落盘。submit 一块后立即继续读下一块，
+          飞行队列满（_BUFFER_DEPTH=2 块）时才 .result() 落盘最老的一块；
+          排序 N 与 "读 N+1 + 落 N-1" 并行。内存峰值 ≈ 2*chunk_size。
+        - 阶段 2：所有块各持一个文件句柄 + lazy 迭代器，每次 yield 一条。
+          内存峰值 = O(块数)（≈ 32）文件句柄 + O(块数) 临时元素。
+        适合"数据量远超内存"的场景。
+
+    内存模型（block 模式，external_sort_streaming_merge=False）：
+        - 阶段 1：双缓冲流水线排序 + 整块落盘。比逐条快 3~5 倍。
+        - 阶段 2：所有块整块 pickle.load 到内存，heapq.merge 同时驱动所有 list。
+          内存峰值 ≈ 数据总量。仅当数据总量放得下内存时使用。
+
+    约束：
+        - 元素必须可 pickle（落盘要求）
+        - 文件描述符上限：块数受 OS 限制（Linux 默认 1024）
+          块数由 choose_external_sort_chunk_size 控制（目标 32）
+    """
+    cfg = config if config is not None else BigArrayConfig.default()
+    streaming = cfg.external_sort_streaming_merge
+    chunk_size = choose_external_sort_chunk_size(
+        total_count,
+        target=cfg.external_sort_chunk_target,
+        min_chunk=cfg.external_sort_min_chunk,
+        max_chunk=cfg.external_sort_max_chunk,
+    )
+    _dump = _dump_streaming if streaming else _dump_block
+    _file_iter_cls = _streaming_file_iter if streaming else _block_file_iter
+
+    temp_paths: list[str] = []
+
+    def _dump_chunk(sorted_chunk: list) -> None:
+        tmp = tempfile.NamedTemporaryFile(suffix=".tmp", prefix="big_sort_", delete=False)
+        try:
+            _dump(sorted_chunk, tmp)
+            tmp.close()
+        except Exception:
+            tmp.close()
+            try:
+                os.unlink(tmp.name)
+            except OSError:
+                pass
+            raise
+        temp_paths.append(tmp.name)
+
+    try:
+        # 阶段 1：双缓冲流水线排序 + 落盘
+        # 改进点（对比串行版）：原版每块 submit().result() 串行，排序与 I/O 完全串行。
+        # 双缓冲：submit 一块后立即继续读下一块；飞行队列满（_BUFFER_DEPTH 块）时
+        # 才 .result() 落盘最老的一块。排序 N 与 "读 N+1 + 落 N-1" 并行，并行度 1→2。
+        # 内存：飞行队列最多 2 个 future；每个 future 的 chunk 在 worker pickled 后即可
+        # 释放，主进程峰值 ≈ 2*chunk_size（与串行版一致，不退化成中数据分支）。
+        in_flight: deque[Future[list[T]]] = deque()
+        with ProcessPoolExecutor(max_workers=max_workers) as pool:
+            chunk: list = []
+            try:
+                for item in data:
+                    chunk.append(item)
+                    if len(chunk) >= chunk_size:
+                        in_flight.append(pool.submit(sort_chunk, (chunk, key)))
+                        chunk = []
+                        if len(in_flight) >= _BUFFER_DEPTH:
+                            _dump_chunk(in_flight.popleft().result())
+                if chunk:  # 末尾不足一块的也要排
+                    in_flight.append(pool.submit(sort_chunk, (chunk, key)))
+                # 排空飞行队列
+                while in_flight:
+                    _dump_chunk(in_flight.popleft().result())
+            except BaseException:
+                # 异常时取消未启动的 future；已启动的让其自然结束（pool.__exit__ 会 join）
+                for fut in in_flight:
+                    if not fut.running():
+                        fut.cancel()
+                raise
+
+        # 阶段 2：多路归并
+        # heapq.merge 接受多个有序迭代器，逐条比较产出
+        files: list[IO[bytes]] = []
+        try:
+            files.extend(open(path, "rb") for path in temp_paths)
+            file_iters: list[Generator[T, None, None]] = [_file_iter_cls(f) for f in files]
+            yield from heapq.merge(*file_iters, key=key)  # type: ignore[arg-type]
+        finally:
+            for f in files:
+                try:
+                    f.close()
+                except Exception:
+                    pass
+    finally:
+        for path in temp_paths:
+            try:
+                os.unlink(path)
+            except OSError:
+                pass
+
+
+__all__ = ["external_sort"]

bigarraylist/algorithms/parallel.py

@@ -0,0 +1,68 @@
+"""模块级可 pickle 函数：供 ProcessPoolExecutor 子进程调用。
+
+⚠️ 这些函数必须放在模块级（不能是嵌套函数/lambda），否则子进程
+pickle 时找不到定义。本文件按职责拆分原 core.py 顶部的四个函数。
+
+子进程通过 `from bigarraylist.algorithms.parallel import _sort_chunk` 路径找到定义。
+"""
+
+from __future__ import annotations
+
+import bisect as _bisect_mod
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+
+def linear_find_in_chunk(args: tuple[list[T], T, int]) -> list[tuple[int, T]]:
+    """在单个 chunk 内查找所有等于 target 的元素，返回 (全局下标, 元素) 列表。"""
+    chunk, target, start_index = args
+    return [(i, item) for i, item in enumerate(chunk, start=start_index) if item == target]
+
+
+def sort_chunk(args: tuple[Any, ...]) -> list[T]:
+    """对单个 chunk 排序，返回新列表（不修改原 chunk）。
+
+    args 形如 ``(chunk, key)``；``key=None`` 表示按元素本身排序。
+    向后兼容：也接受旧式 ``(chunk,)`` 一元组（等价于 ``key=None``）。
+    """
+    if len(args) == 2:
+        chunk, key = args
+    else:
+        # 旧式 (chunk,) 调用（_sort_chunk 别名的遗留调用方）
+        (chunk,) = args
+        key = None
+    return sorted(chunk, key=key)
+
+
+def chunk_is_sorted(args: tuple[list[T]]) -> bool:
+    """判断单个 chunk 是否升序（模块级，供子进程 pickle 调用）。"""
+    (chunk,) = args
+    return all(chunk[i] <= chunk[i + 1] for i in range(len(chunk) - 1))  # type: ignore[operator]
+
+
+def bisect_find_in_chunk(args: tuple[list[T], T, int]) -> list[tuple[int, T]]:
+    """在**已排序**的单个 chunk 内用 bisect 找所有等于 target 的元素。
+
+    返回 (全局下标, 元素) 列表。
+    """
+    chunk, target, start_index = args
+    left = _bisect_mod.bisect_left(chunk, target)  # type: ignore[call-overload]
+    right = _bisect_mod.bisect_right(chunk, target)  # type: ignore[call-overload]
+    return [(start_index + i, chunk[i]) for i in range(left, right)]
+
+
+# ==================== 向后兼容别名 ====================
+# 重构前这些函数在 bigarraylist.core 模块下，名字带下划线前缀。
+# 现有测试 / 外部代码若直接 import，保留别名避免破坏。
+_linear_find_in_chunk: Any = linear_find_in_chunk
+_sort_chunk: Any = sort_chunk
+_chunk_is_sorted: Any = chunk_is_sorted
+_bisect_find_in_chunk: Any = bisect_find_in_chunk
+
+__all__ = [
+    "bisect_find_in_chunk",
+    "chunk_is_sorted",
+    "linear_find_in_chunk",
+    "sort_chunk",
+]

bigarraylist/algorithms/sample.py

@@ -0,0 +1,51 @@
+"""蓄水池抽样（Reservoir Sampling）：从总量未知的流里等概率抽 k 个。
+
+数学保证：每个元素被选中的概率都是 k/N（N 为实际总量，无需预先知道）。
+内存 O(k)，无需预知 N。大数据分析里做随机采样的必备算法。
+"""
+
+from __future__ import annotations
+
+import random
+from collections.abc import Generator, Iterable
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def reservoir_sample(
+    stream: Iterable[T],
+    k: int,
+    seed: int | None = None,
+) -> Generator[T, None, None]:
+    """蓄水池抽样算法 R。
+
+    参数：
+        stream: 任意可迭代对象（不要求预知长度）
+        k:      要抽取的样本数（必须 >= 1，调用方校验）
+        seed:   随机种子，None 表示系统随机；传 int 可复现
+
+    返回：生成器，产出 k 个元素（若流总量 < k，产出全部元素）
+
+    算法：
+        - 先存前 k 个到 reservoir
+        - 第 i 个（i >= k，从 0 计数即 i+1 个元素）以 k/(i+1) 概率
+          随机替换池中某一个
+        - 一遍流式扫描结束，池中 k 个即为等概率样本
+    """
+    rng = random.Random(seed)
+    reservoir: list[T] = []
+    for i, item in enumerate(stream):
+        if i < k:
+            reservoir.append(item)
+        else:
+            # 第 i+1 个元素（索引 i >= k），以 k/(i+1) 概率替换池中随机一个
+            # randint(0, i) 返回 [0, i] 闭区间，共 i+1 个值
+            # 落在 [0, k-1] 的概率 = k/(i+1)，正是我们要的
+            j = rng.randint(0, i)
+            if j < k:
+                reservoir[j] = item
+    yield from reservoir
+
+
+__all__ = ["reservoir_sample"]

bigarraylist/config.py

@@ -0,0 +1,86 @@
+"""配置对象：所有阈值集中可配。
+
+向后兼容：所有字段都有默认值，等于重构前的硬编码常量。
+不传 config 时，行为与重构前完全一致。
+
+用法：
+    cfg = BigArrayConfig(linear_parallel_threshold=50_000)
+    arr = BigArrayList(data, config=cfg)
+
+也可类级别覆盖默认值，影响所有后续实例：
+    BigArrayConfig.set_default(external_sort_chunk_target=64)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class BigArrayConfig:
+    """全部阈值配置。
+
+    所有字段单位均为"元素条数"。各字段的工程意义见注释。
+    """
+
+    # ===== 并行分流阈值 =====
+    # 数据量 >= 该值才走多进程路径；小数据单进程更快
+    linear_parallel_threshold: int = 100_000
+
+    # ===== 内存排序分块 =====
+    # 内存多进程排序的单块大小映射表（数据量 -> 块大小）
+    # 100k→100/块、1M→300/块、10M→800/块、100M→1000/块
+    chunk_size_map: dict[int, int] | None = None  # None → 用默认 _DEFAULT_CHUNK_MAP
+
+    # ===== 外部排序 =====
+    # 目标块数：32 路归并，heapq.merge 效率最高
+    external_sort_chunk_target: int = 32
+    # 每块至少 10 万条（再小进程调度不划算）
+    external_sort_min_chunk: int = 100_000
+    # 每块最多 500 万条（sorted() 单块内存安全）
+    external_sort_max_chunk: int = 5_000_000
+    # 数据量 >= 该值才走外部排序（否则走内存多进程分块）
+    external_sort_threshold: int = 1_000_000
+    # 真正的低内存归并：逐条 load 而非整块 load。
+    # True  → 内存峰值 O(块数)，速度较慢但能处理远超内存的数据
+    # False → 整块 load，归并内存峰值 ≈ 数据总量，速度快 3~5 倍
+    #  True（符合"外部排序"语义）；要恢复旧行为显式设 False
+    external_sort_streaming_merge: bool = True
+
+    # ===== 分桶查找 =====
+    # 桶数上限（避免 OS 文件描述符耗尽，Linux 默认 1024）
+    bucket_max_count: int = 1024
+
+    # ===== 默认配置管理 =====
+    _default: BigArrayConfig | None = None
+
+    @classmethod
+    def default(cls) -> BigArrayConfig:
+        """当前生效的默认配置（实例级覆盖优先于类级默认）。"""
+        if cls._default is not None:
+            return cls._default
+        return cls()
+
+    @classmethod
+    def set_default(cls, **kwargs: Any) -> BigArrayConfig:
+        """类级别覆盖默认值；影响后续所有未显式传 config 的实例。
+
+        kwargs 不全传时，缺失字段保留原默认值。
+        """
+        cls._default = cls(**kwargs)
+        return cls._default
+
+    @classmethod
+    def reset_default(cls) -> None:
+        """恢复类级默认值为内置默认。"""
+        cls._default = None
+
+
+# 内存排序分块大小默认映射（不变于重构前）
+_DEFAULT_CHUNK_MAP: dict[int, int] = {
+    100_000: 100,  # 十万级 → 100 条/块
+    1_000_000: 300,  # 百万级 → 300 条/块
+    10_000_000: 800,  # 千万级 → 800 条/块
+    100_000_000: 1_000,  # 亿级   → 1000 条/块
+}

bigarraylist/core/__init__.py

@@ -0,0 +1,5 @@
+"""核心层：迭代协议 + _fresh_iter（通过 mixin 组合）。"""
+
+from .base import BaseBigArrayList
+
+__all__ = ["BaseBigArrayList"]