Spaces:
Running
Running
File size: 4,757 Bytes
a1dd5ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """BPT1 — immutable, bulk-loaded B+ tree (numpy twin of the C++20 reader in
src/streetdex/index/bptree.{hpp,cpp}; identical bytes on disk).
Why a B+ tree here at all: zone maps prune brilliantly on `ts` because files
are time-sorted, and prune *nothing* on an unsorted column (every row group
spans nearly the whole domain). A secondary index re-sorts (key → location)
once, so point/range predicates on ANY numeric column become a descent plus
a contiguous leaf scan instead of a full-table read. Immutability makes the
classic hard parts vanish: bulk load at 100% fill, no splits, no
rebalancing, rebuilt per table version like every other derived artifact.
Values are opaque u64s; the store packs (file_idx << 40) | row_in_file, so a
leaf hit maps straight to one Parquet row group — the index prunes I/O, not
just rows.
"""
from __future__ import annotations
import struct
from pathlib import Path
import numpy as np
MAGIC = b"BPT1"
DEFAULT_ORDER = 256 # 256 × 8B keys = one 4 KiB page per node touch
def encode_key(values: np.ndarray) -> np.ndarray:
"""Order-preserving i64 encoding: ints pass through; IEEE-754 doubles use
the sign-flip trick (identical to bpt_encode_f64 in C++)."""
if np.issubdtype(values.dtype, np.integer):
return values.astype("<i8")
if np.issubdtype(values.dtype, np.floating):
u = values.astype("<f8").view("<u8").copy()
neg = (u & 0x8000000000000000) != 0
u[neg] = ~u[neg]
u[~neg] |= 0x8000000000000000
return (u ^ 0x8000000000000000).view("<i8")
raise TypeError(f"unindexable dtype {values.dtype}")
def encode_scalar(v) -> int:
return int(encode_key(np.array([v]))[0])
def build(keys: np.ndarray, values: np.ndarray,
order: int = DEFAULT_ORDER) -> bytes:
order_idx = np.argsort(keys, kind="stable")
k = np.ascontiguousarray(keys[order_idx], dtype="<i8")
v = np.ascontiguousarray(values[order_idx], dtype="<u8")
n = len(k)
levels = [] # bottom-up internal levels: first key of each child node
below = k
while len(below) > order:
firsts = below[::order].copy()
levels.append(firsts)
below = firsts
levels.reverse() # root first, as the C++ builder writes them
out = bytearray()
out += MAGIC
out += struct.pack("<IQII", order, n, len(levels), 0)
dir_pos = len(out)
out += b"\0" * (16 * (len(levels) + 1))
directory = []
for lv in levels:
while len(out) % 8:
out += b"\0"
directory.append((len(out), len(lv)))
out += lv.astype("<i8").tobytes()
while len(out) % 8:
out += b"\0"
directory.append((len(out), n))
leaf = np.empty(n, dtype=[("k", "<i8"), ("v", "<u8")])
leaf["k"] = k
leaf["v"] = v
out += leaf.tobytes()
for i, (off, cnt) in enumerate(directory):
struct.pack_into("<QQ", out, dir_pos + 16 * i, off, cnt)
return bytes(out)
class Reader:
def __init__(self, data: bytes | np.memmap):
buf = np.frombuffer(data, dtype=np.uint8) if isinstance(data, bytes) \
else data
if bytes(buf[:4]) != MAGIC:
raise ValueError("not a BPT1 index")
self.order, self.n, nlevels, _ = struct.unpack_from(
"<IQII", buf.tobytes()[:24] if isinstance(buf, np.memmap)
else data, 4)
raw = buf.tobytes() if isinstance(buf, np.memmap) else data
dirs = [struct.unpack_from("<QQ", raw, 24 + 16 * i)
for i in range(nlevels + 1)]
self.levels = [np.frombuffer(raw, "<i8", cnt, off)
for off, cnt in dirs[:-1]]
loff, lcnt = dirs[-1]
leaf = np.frombuffer(raw, dtype=[("k", "<i8"), ("v", "<u8")],
count=lcnt, offset=loff)
self.keys = leaf["k"]
self.vals = leaf["v"]
@classmethod
def open(cls, path: str | Path):
return cls(np.memmap(path, dtype=np.uint8, mode="r"))
def lower_bound(self, k: int) -> int:
node = 0
for lv in self.levels:
begin = node * self.order
end = min(begin + self.order, len(lv))
# lower_bound + step-back (see C++ twin): duplicates of k may
# start in the child before the first child whose first-key == k
pos = int(np.searchsorted(lv[begin:end], k, side="left")) + begin
node = pos - 1 if pos > begin else begin
begin = node * self.order
end = min(begin + self.order, self.n)
return int(np.searchsorted(self.keys[begin:end], k, side="left")) + begin
def range(self, lo: int, hi: int) -> np.ndarray:
a = self.lower_bound(lo)
b = int(np.searchsorted(self.keys, hi, side="right"))
return self.vals[a:b]
|