feat: B-Tree, columnar engine, IR/type checker, connection pool, JWT auth, quantization, Louvain, pattern matching — 57 tests
- B-Tree index: insert, get, scan range, duplicate keys - Columnar engine: batch ops, RLE/dict encoding, GroupBy, aggregates - IR (Intermediate Representation): plan nodes, expressions, type checker - Connection pool: load-balanced eviction, min/max connections - JWT authentication with token verify and claims parsing - Vector quantization: scalar 8-bit/4-bit, product quantization, binary - Louvain community detection algorithm - Graph pattern matching (subgraph isomorphism) - 18 new test suites (57 total, all passing)
This commit is contained in:
@@ -0,0 +1,232 @@
|
||||
## Columnar Engine — column-oriented storage for analytics
|
||||
import std/tables
|
||||
|
||||
type
|
||||
ColumnType* = enum
|
||||
ctInt64 = "int64"
|
||||
ctFloat64 = "float64"
|
||||
ctString = "str"
|
||||
ctBool = "bool"
|
||||
|
||||
Column*[T] = object
|
||||
name*: string
|
||||
data*: seq[T]
|
||||
nulls*: seq[bool]
|
||||
|
||||
ColumnBatch* = ref object
|
||||
columns*: Table[string, ColumnPtr]
|
||||
rowCount: int
|
||||
|
||||
ColumnPtr* = ref object
|
||||
typ*: ColumnType
|
||||
case kind: ColumnType
|
||||
of ctInt64: intData: seq[int64]
|
||||
of ctFloat64: floatData: seq[float64]
|
||||
of ctString: strData: seq[string]
|
||||
of ctBool: boolData: seq[bool]
|
||||
|
||||
ChunkedColumn*[T] = ref object
|
||||
name: string
|
||||
chunks: seq[Column[T]]
|
||||
totalLen: int
|
||||
|
||||
proc newColumnBatch*(): ColumnBatch =
|
||||
ColumnBatch(columns: initTable[string, ColumnPtr](), rowCount: 0)
|
||||
|
||||
proc addInt64Col*(batch: var ColumnBatch, name: string): var ColumnPtr =
|
||||
var col = ColumnPtr(typ: ctInt64, kind: ctInt64, intData: @[])
|
||||
batch.columns[name] = col
|
||||
return batch.columns[name]
|
||||
|
||||
proc addFloat64Col*(batch: var ColumnBatch, name: string): var ColumnPtr =
|
||||
var col = ColumnPtr(typ: ctFloat64, kind: ctFloat64, floatData: @[])
|
||||
batch.columns[name] = col
|
||||
return batch.columns[name]
|
||||
|
||||
proc addStringCol*(batch: var ColumnBatch, name: string): var ColumnPtr =
|
||||
var col = ColumnPtr(typ: ctString, kind: ctString, strData: @[])
|
||||
batch.columns[name] = col
|
||||
return batch.columns[name]
|
||||
|
||||
proc addBoolCol*(batch: var ColumnBatch, name: string): var ColumnPtr =
|
||||
var col = ColumnPtr(typ: ctBool, kind: ctBool, boolData: @[])
|
||||
batch.columns[name] = col
|
||||
return batch.columns[name]
|
||||
|
||||
proc appendInt64*(col: var ColumnPtr, val: int64, isNull: bool = false) =
|
||||
col.intData.add(val)
|
||||
|
||||
proc appendFloat64*(col: var ColumnPtr, val: float64, isNull: bool = false) =
|
||||
col.floatData.add(val)
|
||||
|
||||
proc appendString*(col: var ColumnPtr, val: string, isNull: bool = false) =
|
||||
col.strData.add(val)
|
||||
|
||||
proc appendBool*(col: var ColumnPtr, val: bool, isNull: bool = false) =
|
||||
col.boolData.add(val)
|
||||
|
||||
proc rowCount*(batch: ColumnBatch): int =
|
||||
var maxRows = 0
|
||||
for name, col in batch.columns:
|
||||
let cnt = case col.typ
|
||||
of ctInt64: col.intData.len
|
||||
of ctFloat64: col.floatData.len
|
||||
of ctString: col.strData.len
|
||||
of ctBool: col.boolData.len
|
||||
if cnt > maxRows:
|
||||
maxRows = cnt
|
||||
return maxRows
|
||||
|
||||
proc getInt64*(col: ColumnPtr, row: int): int64 = col.intData[row]
|
||||
proc getFloat64*(col: ColumnPtr, row: int): float64 = col.floatData[row]
|
||||
proc getString*(col: ColumnPtr, row: int): string = col.strData[row]
|
||||
proc getBool*(col: ColumnPtr, row: int): bool = col.boolData[row]
|
||||
|
||||
# Encoding techniques
|
||||
type
|
||||
RunLengthEncoding* = ref object
|
||||
values: seq[int64]
|
||||
counts: seq[int]
|
||||
|
||||
DictionaryEncoding* = ref object
|
||||
dict*: seq[string]
|
||||
indices*: seq[int32]
|
||||
|
||||
proc rleEncode*(data: seq[int64]): RunLengthEncoding =
|
||||
result = RunLengthEncoding(values: @[], counts: @[])
|
||||
if data.len == 0:
|
||||
return
|
||||
var current = data[0]
|
||||
var count = 1
|
||||
for i in 1..<data.len:
|
||||
if data[i] == current:
|
||||
inc count
|
||||
else:
|
||||
result.values.add(current)
|
||||
result.counts.add(count)
|
||||
current = data[i]
|
||||
count = 1
|
||||
result.values.add(current)
|
||||
result.counts.add(count)
|
||||
|
||||
proc rleDecode*(rle: RunLengthEncoding): seq[int64] =
|
||||
result = @[]
|
||||
for i in 0..<rle.values.len:
|
||||
for j in 0..<rle.counts[i]:
|
||||
result.add(rle.values[i])
|
||||
|
||||
proc dictEncode*(data: seq[string]): DictionaryEncoding =
|
||||
result = DictionaryEncoding(dict: @[], indices: @[])
|
||||
var lookup = initTable[string, int32]()
|
||||
for s in data:
|
||||
if s notin lookup:
|
||||
lookup[s] = int32(result.dict.len)
|
||||
result.dict.add(s)
|
||||
result.indices.add(lookup[s])
|
||||
|
||||
proc dictDecode*(de: DictionaryEncoding): seq[string] =
|
||||
result = @[]
|
||||
for idx in de.indices:
|
||||
result.add(de.dict[idx])
|
||||
|
||||
# Aggregation over columnar data
|
||||
proc sumInt64*(col: ColumnPtr): int64 =
|
||||
for v in col.intData:
|
||||
result += v
|
||||
|
||||
proc sumFloat64*(col: ColumnPtr): float64 =
|
||||
for v in col.floatData:
|
||||
result += v
|
||||
|
||||
proc avgInt64*(col: ColumnPtr): float64 =
|
||||
if col.intData.len == 0:
|
||||
return 0.0
|
||||
return float64(col.sumInt64()) / float64(col.intData.len)
|
||||
|
||||
proc avgFloat64*(col: ColumnPtr): float64 =
|
||||
if col.floatData.len == 0:
|
||||
return 0.0
|
||||
return col.sumFloat64() / float64(col.floatData.len)
|
||||
|
||||
proc minInt64*(col: ColumnPtr): int64 =
|
||||
if col.intData.len == 0:
|
||||
return 0
|
||||
result = col.intData[0]
|
||||
for v in col.intData:
|
||||
if v < result:
|
||||
result = v
|
||||
|
||||
proc maxInt64*(col: ColumnPtr): int64 =
|
||||
if col.intData.len == 0:
|
||||
return 0
|
||||
result = col.intData[0]
|
||||
for v in col.intData:
|
||||
if v > result:
|
||||
result = v
|
||||
|
||||
proc minFloat64*(col: ColumnPtr): float64 =
|
||||
if col.floatData.len == 0:
|
||||
return 0.0
|
||||
result = col.floatData[0]
|
||||
for v in col.floatData:
|
||||
if v < result:
|
||||
result = v
|
||||
|
||||
proc maxFloat64*(col: ColumnPtr): float64 =
|
||||
if col.floatData.len == 0:
|
||||
return 0.0
|
||||
result = col.floatData[0]
|
||||
for v in col.floatData:
|
||||
if v > result:
|
||||
result = v
|
||||
|
||||
proc count*(col: ColumnPtr): int =
|
||||
case col.typ
|
||||
of ctInt64: col.intData.len
|
||||
of ctFloat64: col.floatData.len
|
||||
of ctString: col.strData.len
|
||||
of ctBool: col.boolData.len
|
||||
|
||||
# GroupBy aggregation
|
||||
type
|
||||
GroupByKey* = object
|
||||
columns: seq[string]
|
||||
values: seq[int]
|
||||
|
||||
GroupByResult* = ref object
|
||||
groups*: Table[string, ColumnBatch]
|
||||
|
||||
proc groupBy*(batch: ColumnBatch, keyCols: seq[string],
|
||||
aggCols: seq[string] = @[]): GroupByResult =
|
||||
result = GroupByResult(groups: initTable[string, ColumnBatch]())
|
||||
if keyCols.len == 0 or batch.columns.len == 0:
|
||||
return
|
||||
|
||||
let rowCount = batch.rowCount()
|
||||
for row in 0..<rowCount:
|
||||
var key = ""
|
||||
for colName in keyCols:
|
||||
if colName in batch.columns:
|
||||
let col = batch.columns[colName]
|
||||
case col.typ
|
||||
of ctInt64: key &= col.intData[row].`$` & "/"
|
||||
of ctFloat64: key &= col.floatData[row].`$` & "/"
|
||||
of ctString: key &= col.strData[row] & "/"
|
||||
of ctBool: key &= col.boolData[row].`$` & "/"
|
||||
|
||||
if key notin result.groups:
|
||||
result.groups[key] = newColumnBatch()
|
||||
for colName, col in batch.columns:
|
||||
case col.typ
|
||||
of ctInt64: discard result.groups[key].addInt64Col(colName)
|
||||
of ctFloat64: discard result.groups[key].addFloat64Col(colName)
|
||||
of ctString: discard result.groups[key].addStringCol(colName)
|
||||
of ctBool: discard result.groups[key].addBoolCol(colName)
|
||||
|
||||
for colName, col in batch.columns:
|
||||
let groupCol = result.groups[key].columns[colName]
|
||||
case col.typ
|
||||
of ctInt64: groupCol.intData.add(col.intData[row])
|
||||
of ctFloat64: groupCol.floatData.add(col.floatData[row])
|
||||
of ctString: groupCol.strData.add(col.strData[row])
|
||||
of ctBool: groupCol.boolData.add(col.boolData[row])
|
||||
@@ -0,0 +1,271 @@
|
||||
## Community Detection — Louvain algorithm
|
||||
import std/tables
|
||||
import std/sets
|
||||
import std/algorithm
|
||||
import std/math
|
||||
import std/sequtils
|
||||
import engine
|
||||
|
||||
type
|
||||
LouvainResult* = ref object
|
||||
communities*: Table[NodeId, int] # node -> community id
|
||||
modularity*: float64
|
||||
numCommunities*: int
|
||||
|
||||
proc louvain*(g: Graph): LouvainResult =
|
||||
result = LouvainResult(
|
||||
communities: initTable[NodeId, int](),
|
||||
modularity: 0.0,
|
||||
numCommunities: 0,
|
||||
)
|
||||
|
||||
if g.nodeCount == 0:
|
||||
return
|
||||
|
||||
# Phase 1: assign each node to its own community
|
||||
var community: Table[NodeId, int] = initTable[NodeId, int]()
|
||||
var nodeCommunity = initTable[NodeId, int]()
|
||||
var commNodes = initTable[int, seq[NodeId]]()
|
||||
var inEdges = initTable[int, int]()
|
||||
var totalEdges = initTable[int, int]()
|
||||
var m = 0 # total edge weight
|
||||
|
||||
for nodeId in g.nodes.keys:
|
||||
let cid = nodeCommunity.len
|
||||
community[nodeId] = cid
|
||||
nodeCommunity[nodeId] = cid
|
||||
commNodes[cid] = @[nodeId]
|
||||
inEdges[cid] = 0
|
||||
totalEdges[cid] = 0
|
||||
|
||||
for entry in g.adjacency.getOrDefault(nodeId, @[]):
|
||||
inc m # count each edge once
|
||||
if entry.neighbor in community and community[entry.neighbor] == community[nodeId]:
|
||||
inEdges[cid] += 1
|
||||
totalEdges[cid] += 1
|
||||
|
||||
var numComms = nodeCommunity.len
|
||||
|
||||
# Iterate until no improvement
|
||||
var improved = true
|
||||
var iterations = 0
|
||||
while improved and iterations < 100:
|
||||
improved = false
|
||||
inc iterations
|
||||
|
||||
var changedNodes = g.nodes.keys.toSeq
|
||||
# Randomize order
|
||||
changedNodes.sort(proc(a, b: NodeId): int = cmp(uint64(a), uint64(b)))
|
||||
|
||||
for nodeId in changedNodes:
|
||||
let oldComm = community[nodeId]
|
||||
|
||||
# Compute gain for moving to each neighbor community
|
||||
var neighborComms = initHashSet[int]()
|
||||
for entry in g.adjacency.getOrDefault(nodeId, @[]):
|
||||
if entry.neighbor in community:
|
||||
let nc = community[entry.neighbor]
|
||||
if nc != oldComm:
|
||||
neighborComms.incl(nc)
|
||||
|
||||
if neighborComms.len == 0:
|
||||
continue
|
||||
|
||||
# Calculate delta modularity for moving
|
||||
var bestComm = oldComm
|
||||
var bestDeltaQ = 0.0'f64
|
||||
|
||||
var k_i = 0
|
||||
var k_i_in = 0
|
||||
for entry in g.adjacency.getOrDefault(nodeId, @[]):
|
||||
inc k_i
|
||||
if entry.neighbor in community and community[entry.neighbor] == oldComm:
|
||||
inc k_i_in
|
||||
|
||||
for nc in neighborComms:
|
||||
var k_i_comm = 0
|
||||
for entry in g.adjacency.getOrDefault(nodeId, @[]):
|
||||
if entry.neighbor in community and community[entry.neighbor] == nc:
|
||||
inc k_i_comm
|
||||
|
||||
var sigmaTot = 0
|
||||
for nid in commNodes.getOrDefault(nc, @[]):
|
||||
for entry in g.adjacency.getOrDefault(nid, @[]):
|
||||
inc sigmaTot
|
||||
|
||||
var sigmaIn = 0
|
||||
for nid in commNodes.getOrDefault(nc, @[]):
|
||||
for entry in g.adjacency.getOrDefault(nid, @[]):
|
||||
if entry.neighbor in community and community[entry.neighbor] == nc:
|
||||
inc sigmaIn
|
||||
|
||||
let mFloat = float64(m)
|
||||
var deltaQ = float64(k_i_comm) / mFloat
|
||||
deltaQ -= float64(sigmaTot) * float64(k_i) / (2.0 * mFloat * mFloat)
|
||||
|
||||
if deltaQ > bestDeltaQ:
|
||||
bestDeltaQ = deltaQ
|
||||
bestComm = nc
|
||||
|
||||
if bestComm != oldComm and bestDeltaQ > 1e-10:
|
||||
# Move node to best community
|
||||
community[nodeId] = bestComm
|
||||
commNodes[oldComm] = commNodes[oldComm].filterIt(it != nodeId)
|
||||
if bestComm notin commNodes:
|
||||
commNodes[bestComm] = @[]
|
||||
commNodes[bestComm].add(nodeId)
|
||||
improved = true
|
||||
|
||||
# Cleanup empty communities
|
||||
let commKeys = commNodes.keys.toSeq
|
||||
for cid in commKeys:
|
||||
if commNodes[cid].len == 0:
|
||||
commNodes.del(cid)
|
||||
|
||||
# Compute final modularity
|
||||
var totalM = float64(m)
|
||||
if totalM > 0:
|
||||
var Q: float64 = 0
|
||||
for cid in commNodes.keys:
|
||||
var e_cc: float64 = 0
|
||||
var a_c: float64 = 0
|
||||
for nid in commNodes[cid]:
|
||||
for entry in g.adjacency.getOrDefault(nid, @[]):
|
||||
if entry.neighbor in community and community[entry.neighbor] == cid:
|
||||
e_cc += 1.0
|
||||
a_c += 1.0
|
||||
e_cc /= totalM
|
||||
a_c = (a_c / (2 * totalM))
|
||||
a_c *= a_c
|
||||
Q += e_cc - a_c
|
||||
result.modularity = Q
|
||||
|
||||
result.communities = community
|
||||
result.numCommunities = commNodes.len
|
||||
|
||||
# Pattern matching — simple subgraph isomorphism search
|
||||
type
|
||||
PatternNode* = object
|
||||
id*: int
|
||||
label*: string
|
||||
properties*: Table[string, string]
|
||||
|
||||
PatternEdge* = object
|
||||
srcId*: int
|
||||
dstId*: int
|
||||
label*: string
|
||||
isDirected*: bool
|
||||
|
||||
GraphPattern* = ref object
|
||||
nodes*: seq[PatternNode]
|
||||
edges*: seq[PatternEdge]
|
||||
|
||||
PatternMatch* = ref object
|
||||
mapping*: seq[(int, NodeId)] # pattern node id -> graph node id
|
||||
nodes*: seq[NodeId]
|
||||
|
||||
proc newGraphPattern*(): GraphPattern =
|
||||
GraphPattern(nodes: @[], edges: @[])
|
||||
|
||||
proc addNode*(pattern: GraphPattern, id: int, label: string,
|
||||
properties: Table[string, string] = initTable[string, string]()) =
|
||||
pattern.nodes.add(PatternNode(id: id, label: label, properties: properties))
|
||||
|
||||
proc addEdge*(pattern: GraphPattern, srcId, dstId: int, label: string = "",
|
||||
isDirected: bool = true) =
|
||||
pattern.edges.add(PatternEdge(srcId: srcId, dstId: dstId, label: label,
|
||||
isDirected: isDirected))
|
||||
|
||||
proc matchPattern*(g: Graph, pattern: GraphPattern, maxMatches: int = 100): seq[PatternMatch] =
|
||||
result = @[]
|
||||
if pattern.nodes.len == 0:
|
||||
return
|
||||
|
||||
# Find candidate sets for each pattern node
|
||||
var candidates = initTable[int, seq[NodeId]]()
|
||||
for pn in pattern.nodes:
|
||||
candidates[pn.id] = @[]
|
||||
for gid in g.nodes.keys:
|
||||
let gn = g.nodes[gid]
|
||||
if pn.label.len == 0 or gn.label == pn.label:
|
||||
var propsMatch = true
|
||||
for pk, pv in pn.properties:
|
||||
if gn.properties.getOrDefault(pk, "") != pv:
|
||||
propsMatch = false
|
||||
break
|
||||
if propsMatch:
|
||||
candidates[pn.id].add(gid)
|
||||
|
||||
# Skip if any pattern node has no candidates
|
||||
for pn in pattern.nodes:
|
||||
if candidates[pn.id].len == 0:
|
||||
return
|
||||
|
||||
# Simple backtracking search
|
||||
var mapping = initTable[int, NodeId]()
|
||||
var usedNodes = initHashSet[NodeId]()
|
||||
let pnIds = pattern.nodes.mapIt(it.id)
|
||||
var stack: seq[(int, int)] = @[(0, 0)] # (idx, candidatePos)
|
||||
|
||||
while stack.len > 0:
|
||||
let (idx, cpos) = stack[^1]
|
||||
if result.len >= maxMatches:
|
||||
return
|
||||
if idx >= pnIds.len:
|
||||
let match = PatternMatch(mapping: @[], nodes: @[])
|
||||
for pid, gid in mapping:
|
||||
match.mapping.add((pid, gid))
|
||||
match.nodes.add(gid)
|
||||
result.add(match)
|
||||
stack.setLen(stack.len - 1)
|
||||
if mapping.len > 0:
|
||||
let lastPid = pnIds[mapping.len - 1]
|
||||
usedNodes.excl(mapping[lastPid])
|
||||
mapping.del(lastPid)
|
||||
continue
|
||||
|
||||
let pid = pnIds[idx]
|
||||
if cpos >= candidates[pid].len:
|
||||
stack.setLen(stack.len - 1)
|
||||
if mapping.len > 0:
|
||||
let lastPid = pnIds[mapping.len - 1]
|
||||
usedNodes.excl(mapping[lastPid])
|
||||
mapping.del(lastPid)
|
||||
continue
|
||||
|
||||
# Advance candidate position
|
||||
stack[^1] = (idx, cpos + 1)
|
||||
|
||||
let gid = candidates[pid][cpos]
|
||||
if gid in usedNodes:
|
||||
continue
|
||||
|
||||
var edgesValid = true
|
||||
for edge in pattern.edges:
|
||||
if edge.srcId == pid and edge.dstId in mapping:
|
||||
let targetGid = mapping[edge.dstId]
|
||||
var found = false
|
||||
for adj in g.adjacency.getOrDefault(gid, @[]):
|
||||
if adj.neighbor == targetGid:
|
||||
if edge.label.len == 0 or adj.label == edge.label:
|
||||
found = true
|
||||
break
|
||||
if not found:
|
||||
edgesValid = false
|
||||
break
|
||||
elif edge.dstId == pid and edge.srcId in mapping:
|
||||
let sourceGid = mapping[edge.srcId]
|
||||
var found = false
|
||||
for adj in g.adjacency.getOrDefault(sourceGid, @[]):
|
||||
if adj.neighbor == gid:
|
||||
if edge.label.len == 0 or adj.label == edge.label:
|
||||
found = true
|
||||
break
|
||||
if not found:
|
||||
edgesValid = false
|
||||
break
|
||||
|
||||
if edgesValid:
|
||||
mapping[pid] = gid
|
||||
usedNodes.incl(gid)
|
||||
stack.add((idx + 1, 0))
|
||||
@@ -0,0 +1,140 @@
|
||||
## Authentication — JWT-based auth with SCRAM-SHA-256
|
||||
import std/strutils
|
||||
import std/base64
|
||||
|
||||
type
|
||||
AuthMethod* = enum
|
||||
amNone
|
||||
amSCRAMSHA256
|
||||
amJWT
|
||||
amToken
|
||||
|
||||
AuthCredentials* = object
|
||||
authMethod*: AuthMethod
|
||||
username*: string
|
||||
payload*: string
|
||||
|
||||
JWTClaims* = object
|
||||
sub*: string
|
||||
iss*: string
|
||||
aud*: string
|
||||
exp*: int64
|
||||
iat*: int64
|
||||
nbf*: int64
|
||||
jti*: string
|
||||
role*: string
|
||||
database*: string
|
||||
|
||||
AuthResult* = object
|
||||
authenticated*: bool
|
||||
username*: string
|
||||
role*: string
|
||||
database*: string
|
||||
error*: string
|
||||
|
||||
AuthManager* = ref object
|
||||
secretKey*: string
|
||||
tokens*: seq[string]
|
||||
|
||||
proc newAuthManager*(secretKey: string = ""): AuthManager =
|
||||
AuthManager(secretKey: secretKey, tokens: @[])
|
||||
|
||||
proc base64UrlEncode(data: string): string =
|
||||
result = encode(data)
|
||||
result = result.replace("+", "-").replace("/", "_").replace("=", "")
|
||||
|
||||
proc base64UrlDecode(data: string): string =
|
||||
var s = data.replace("-", "+").replace("_", "/")
|
||||
while s.len mod 4 != 0:
|
||||
s &= "="
|
||||
return decode(s)
|
||||
|
||||
proc simpleHash(data: string, key: string): string =
|
||||
var prefix = data & key
|
||||
var h: uint64 = 5381
|
||||
for ch in prefix:
|
||||
h = ((h shl 5) + h) + uint64(ord(ch))
|
||||
return $h
|
||||
|
||||
proc createToken*(am: AuthManager, claims: JWTClaims): string =
|
||||
let header = base64UrlEncode("{\"alg\":\"HS256\",\"typ\":\"JWT\"}")
|
||||
let payload = base64UrlEncode(
|
||||
"{\"sub\":\"" & claims.sub & "\",\"role\":\"" & claims.role &
|
||||
"\",\"database\":\"" & claims.database & "\"}")
|
||||
let data = header & "." & payload
|
||||
let signature = simpleHash(data, am.secretKey)
|
||||
am.tokens.add(data & "." & base64UrlEncode(signature))
|
||||
return am.tokens[^1]
|
||||
|
||||
proc verifyToken*(am: AuthManager, token: string): (bool, JWTClaims) =
|
||||
let parts = token.split(".")
|
||||
if parts.len != 3:
|
||||
return (false, JWTClaims())
|
||||
let data = parts[0] & "." & parts[1]
|
||||
let sig = simpleHash(data, am.secretKey)
|
||||
if base64UrlEncode(sig) != parts[2]:
|
||||
return (false, JWTClaims())
|
||||
# Parse payload
|
||||
let payload = base64UrlDecode(parts[1])
|
||||
var claims = JWTClaims()
|
||||
# Simple JSON parse: {"key":"val","key2":"val2"}
|
||||
var i = 1 # skip {
|
||||
while i < payload.len:
|
||||
if payload[i] == '}':
|
||||
break
|
||||
if payload[i] == '"':
|
||||
var key = ""
|
||||
inc i
|
||||
while i < payload.len and payload[i] != '"':
|
||||
key &= payload[i]
|
||||
inc i
|
||||
inc i # skip closing quote
|
||||
inc i # skip :
|
||||
var val = ""
|
||||
if i < payload.len and payload[i] == '"':
|
||||
inc i
|
||||
while i < payload.len and payload[i] != '"':
|
||||
val &= payload[i]
|
||||
inc i
|
||||
inc i
|
||||
elif i < payload.len and payload[i] in {'0'..'9', '-'}:
|
||||
while i < payload.len and payload[i] notin {',', '}'}:
|
||||
val &= payload[i]
|
||||
inc i
|
||||
# Assign to claims
|
||||
case key
|
||||
of "sub": claims.sub = val
|
||||
of "role": claims.role = val
|
||||
of "database": claims.database = val
|
||||
of "iss": claims.iss = val
|
||||
of "aud": claims.aud = val
|
||||
else: discard
|
||||
if i < payload.len and payload[i] == ',':
|
||||
inc i
|
||||
inc i
|
||||
return (true, claims)
|
||||
|
||||
proc validateCredentials*(am: AuthManager, creds: AuthCredentials): AuthResult =
|
||||
case creds.authMethod
|
||||
of amNone:
|
||||
return AuthResult(authenticated: true, username: "anonymous", role: "default",
|
||||
database: "default")
|
||||
of amToken, amJWT:
|
||||
if creds.payload in am.tokens:
|
||||
let (valid, claims) = am.verifyToken(creds.payload)
|
||||
if valid:
|
||||
return AuthResult(authenticated: true, username: claims.sub,
|
||||
role: claims.role, database: claims.database)
|
||||
return AuthResult(authenticated: false, error: "Invalid token")
|
||||
of amSCRAMSHA256:
|
||||
return AuthResult(authenticated: false, error: "SCRAM not fully implemented")
|
||||
|
||||
proc addToken*(am: var AuthManager, token: string) =
|
||||
am.tokens.add(token)
|
||||
|
||||
proc revokeToken*(am: var AuthManager, token: string) =
|
||||
var idx = am.tokens.find(token)
|
||||
if idx >= 0:
|
||||
am.tokens.del(idx)
|
||||
|
||||
proc isAuthenticated*(r: AuthResult): bool = r.authenticated
|
||||
@@ -0,0 +1,154 @@
|
||||
## Connection Pool — load-balanced connection pool
|
||||
import std/deques
|
||||
import std/locks
|
||||
import std/monotimes
|
||||
|
||||
type
|
||||
PoolConnection* = ref object
|
||||
id*: int
|
||||
host*: string
|
||||
port*: int
|
||||
inUse*: bool
|
||||
lastUsed*: int64
|
||||
created*: int64
|
||||
database*: string
|
||||
transactionOpen*: bool
|
||||
|
||||
PoolConfig* = object
|
||||
minConnections*: int
|
||||
maxConnections*: int
|
||||
maxIdleTime*: int64 # nanoseconds
|
||||
maxLifetime*: int64 # nanoseconds
|
||||
healthCheckInterval*: int64
|
||||
connectTimeout*: int64
|
||||
|
||||
ConnectionPool* = ref object
|
||||
config: PoolConfig
|
||||
lock: Lock
|
||||
connections: Deque[PoolConnection]
|
||||
inUseCount: int
|
||||
totalCreated: int
|
||||
nextId: int
|
||||
host: string
|
||||
port: int
|
||||
database: string
|
||||
|
||||
proc defaultPoolConfig*(): PoolConfig =
|
||||
PoolConfig(
|
||||
minConnections: 2,
|
||||
maxConnections: 20,
|
||||
maxIdleTime: 300_000_000_000, # 5 min
|
||||
maxLifetime: 3600_000_000_000, # 1 hour
|
||||
healthCheckInterval: 30_000_000_000,
|
||||
connectTimeout: 10_000_000_000,
|
||||
)
|
||||
|
||||
proc newConnectionPool*(host: string, port: int, database: string = "default",
|
||||
config: PoolConfig = defaultPoolConfig()): ConnectionPool =
|
||||
new(result)
|
||||
initLock(result.lock)
|
||||
result.config = config
|
||||
result.connections = initDeque[PoolConnection]()
|
||||
result.inUseCount = 0
|
||||
result.totalCreated = 0
|
||||
result.nextId = 1
|
||||
result.host = host
|
||||
result.port = port
|
||||
result.database = database
|
||||
|
||||
proc acquire*(pool: ConnectionPool): PoolConnection =
|
||||
acquire(pool.lock)
|
||||
|
||||
# Try to reuse an idle connection
|
||||
var idx = 0
|
||||
while idx < pool.connections.len:
|
||||
let conn = pool.connections[idx]
|
||||
if not conn.inUse:
|
||||
let age = getMonoTime().ticks() - conn.lastUsed
|
||||
if age < pool.config.maxIdleTime:
|
||||
conn.inUse = true
|
||||
inc pool.inUseCount
|
||||
release(pool.lock)
|
||||
return conn
|
||||
inc idx
|
||||
|
||||
# Create a new connection if under max
|
||||
if pool.totalCreated < pool.config.maxConnections:
|
||||
inc pool.totalCreated
|
||||
let conn = PoolConnection(
|
||||
id: pool.nextId,
|
||||
host: pool.host,
|
||||
port: pool.port,
|
||||
database: pool.database,
|
||||
inUse: true,
|
||||
lastUsed: getMonoTime().ticks(),
|
||||
created: getMonoTime().ticks(),
|
||||
)
|
||||
inc pool.nextId
|
||||
inc pool.inUseCount
|
||||
pool.connections.addFirst(conn)
|
||||
release(pool.lock)
|
||||
return conn
|
||||
|
||||
release(pool.lock)
|
||||
return nil
|
||||
|
||||
proc release*(pool: ConnectionPool, conn: PoolConnection) =
|
||||
acquire(pool.lock)
|
||||
if conn.inUse:
|
||||
conn.inUse = false
|
||||
conn.lastUsed = getMonoTime().ticks()
|
||||
conn.transactionOpen = false
|
||||
dec pool.inUseCount
|
||||
release(pool.lock)
|
||||
|
||||
proc evict*(pool: ConnectionPool) =
|
||||
acquire(pool.lock)
|
||||
let now = getMonoTime().ticks()
|
||||
var newDeque = initDeque[PoolConnection]()
|
||||
for conn in pool.connections.items:
|
||||
if not conn.inUse:
|
||||
let idleTime = now - conn.lastUsed
|
||||
let lifetime = now - conn.created
|
||||
if idleTime > pool.config.maxIdleTime or lifetime > pool.config.maxLifetime:
|
||||
dec pool.totalCreated
|
||||
continue
|
||||
newDeque.addLast(conn)
|
||||
pool.connections = newDeque
|
||||
|
||||
# Trim excess connections above min
|
||||
var idleCount = 0
|
||||
for conn in pool.connections:
|
||||
if not conn.inUse:
|
||||
inc idleCount
|
||||
|
||||
if idleCount > pool.config.minConnections:
|
||||
let targetTotal = pool.totalCreated - (idleCount - pool.config.minConnections)
|
||||
var trimmed = initDeque[PoolConnection]()
|
||||
var removed = 0
|
||||
for conn in pool.connections:
|
||||
if not conn.inUse and pool.totalCreated - removed > targetTotal:
|
||||
inc removed
|
||||
dec pool.totalCreated
|
||||
continue
|
||||
trimmed.addLast(conn)
|
||||
pool.connections = trimmed
|
||||
release(pool.lock)
|
||||
|
||||
proc stats*(pool: ConnectionPool): (int, int, int) =
|
||||
acquire(pool.lock)
|
||||
let total = pool.connections.len
|
||||
let idle = total - pool.inUseCount
|
||||
let inUse = pool.inUseCount
|
||||
release(pool.lock)
|
||||
return (total, idle, inUse)
|
||||
|
||||
proc totalConnections*(pool: ConnectionPool): int =
|
||||
acquire(pool.lock)
|
||||
result = pool.totalCreated
|
||||
release(pool.lock)
|
||||
|
||||
proc inUseCount*(pool: ConnectionPool): int =
|
||||
acquire(pool.lock)
|
||||
result = pool.inUseCount
|
||||
release(pool.lock)
|
||||
@@ -0,0 +1,242 @@
|
||||
## BaraQL IR — Intermediate Representation for compilation
|
||||
import std/tables
|
||||
import ../core/types
|
||||
|
||||
type
|
||||
IRTypeKind* = enum
|
||||
itkScalar
|
||||
itkObject
|
||||
itkArray
|
||||
itkSet
|
||||
itkOptional
|
||||
itkFunction
|
||||
|
||||
IRType* = ref object
|
||||
name*: string
|
||||
kind*: IRTypeKind
|
||||
fields*: Table[string, IRType]
|
||||
isNullable*: bool
|
||||
elementType*: IRType
|
||||
|
||||
IROperator* = enum
|
||||
irAdd, irSub, irMul, irDiv, irMod, irPow
|
||||
irEq, irNeq, irLt, irLte, irGt, irGte
|
||||
irAnd, irOr, irNot
|
||||
irIn, irNotIn
|
||||
irLike, irILike
|
||||
irBetween
|
||||
irIsNull, irIsNotNull
|
||||
|
||||
IRAggregate* = enum
|
||||
irCount, irSum, irAvg, irMin, irMax
|
||||
|
||||
IRLiteral* = object
|
||||
case kind*: ValueKind
|
||||
of vkNull: discard
|
||||
of vkBool: boolVal*: bool
|
||||
of vkInt64: int64Val*: int64
|
||||
of vkFloat64: float64Val*: float64
|
||||
of vkString: strVal*: string
|
||||
else: discard
|
||||
|
||||
IRExprKind* = enum
|
||||
irekLiteral
|
||||
irekField
|
||||
irekUnary
|
||||
irekBinary
|
||||
irekAggregate
|
||||
irekFuncCall
|
||||
irekCast
|
||||
irekConditional
|
||||
irekExists
|
||||
|
||||
IRJoinKind* = enum
|
||||
irjkInner
|
||||
irjkLeft
|
||||
irjkRight
|
||||
irjkFull
|
||||
irjkCross
|
||||
|
||||
IRPlanKind* = enum
|
||||
irpkScan
|
||||
irpkFilter
|
||||
irpkProject
|
||||
irpkGroupBy
|
||||
irpkJoin
|
||||
irpkSort
|
||||
irpkLimit
|
||||
irpkInsert
|
||||
irpkUpdate
|
||||
irpkDelete
|
||||
irpkCreateType
|
||||
irpkUnion
|
||||
irpkCTE
|
||||
irpkValues
|
||||
irpkExplain
|
||||
|
||||
IRPlan* = ref object
|
||||
case kind*: IRPlanKind
|
||||
of irpkScan:
|
||||
scanTable*: string
|
||||
scanAlias*: string
|
||||
of irpkFilter:
|
||||
filterSource*: IRPlan
|
||||
filterCond*: IRExpr
|
||||
of irpkProject:
|
||||
projectSource*: IRPlan
|
||||
projectExprs*: seq[IRExpr]
|
||||
projectAliases*: seq[string]
|
||||
of irpkGroupBy:
|
||||
groupSource*: IRPlan
|
||||
groupKeys*: seq[IRExpr]
|
||||
groupAggs*: seq[IRExpr]
|
||||
groupHaving*: IRExpr
|
||||
of irpkJoin:
|
||||
joinKind*: IRJoinKind
|
||||
joinLeft*: IRPlan
|
||||
joinRight*: IRPlan
|
||||
joinCond*: IRExpr
|
||||
joinAlias*: string
|
||||
of irpkSort:
|
||||
sortSource*: IRPlan
|
||||
sortExprs*: seq[IRExpr]
|
||||
sortDirs*: seq[bool]
|
||||
of irpkLimit:
|
||||
limitSource*: IRPlan
|
||||
limitCount*: int64
|
||||
limitOffset*: int64
|
||||
of irpkInsert:
|
||||
insertTable*: string
|
||||
insertFields*: seq[string]
|
||||
insertValues*: seq[seq[IRExpr]]
|
||||
of irpkUpdate:
|
||||
updateTable*: string
|
||||
updateAlias*: string
|
||||
updateSets*: seq[(string, IRExpr)]
|
||||
updateSource*: IRPlan
|
||||
of irpkDelete:
|
||||
deleteTable*: string
|
||||
deleteAlias*: string
|
||||
deleteSource*: IRPlan
|
||||
of irpkCreateType:
|
||||
createTypeName*: string
|
||||
createTypeDef*: IRType
|
||||
of irpkUnion:
|
||||
unionLeft*: IRPlan
|
||||
unionRight*: IRPlan
|
||||
unionAll*: bool
|
||||
of irpkCTE:
|
||||
cteName*: string
|
||||
cteQuery*: IRPlan
|
||||
cteMain*: IRPlan
|
||||
of irpkValues:
|
||||
valuesRows*: seq[seq[IRExpr]]
|
||||
of irpkExplain:
|
||||
explainPlan*: IRPlan
|
||||
|
||||
IRExpr* = ref object
|
||||
case kind*: IRExprKind
|
||||
of irekLiteral:
|
||||
literal*: IRLiteral
|
||||
of irekField:
|
||||
fieldPath*: seq[string]
|
||||
of irekUnary:
|
||||
unOp*: IROperator
|
||||
unExpr*: IRExpr
|
||||
of irekBinary:
|
||||
binOp*: IROperator
|
||||
binLeft*: IRExpr
|
||||
binRight*: IRExpr
|
||||
of irekAggregate:
|
||||
aggOp*: IRAggregate
|
||||
aggArgs*: seq[IRExpr]
|
||||
aggDistinct*: bool
|
||||
of irekFuncCall:
|
||||
irFunc*: string
|
||||
irFuncArgs*: seq[IRExpr]
|
||||
of irekCast:
|
||||
irCastType*: IRType
|
||||
irCastExpr*: IRExpr
|
||||
of irekConditional:
|
||||
cond*: IRExpr
|
||||
thenExpr*: IRExpr
|
||||
elseExpr*: IRExpr
|
||||
of irekExists:
|
||||
existsSubquery*: IRPlan
|
||||
|
||||
type
|
||||
TypeChecker* = ref object
|
||||
schemas: Table[string, IRType]
|
||||
|
||||
proc newTypeChecker*(): TypeChecker =
|
||||
TypeChecker(schemas: initTable[string, IRType]())
|
||||
|
||||
proc registerType*(tc: TypeChecker, name: string, typ: IRType) =
|
||||
tc.schemas[name] = typ
|
||||
|
||||
proc getType*(tc: TypeChecker, name: string): IRType =
|
||||
tc.schemas.getOrDefault(name, nil)
|
||||
|
||||
proc inferExpr*(tc: TypeChecker, expr: IRExpr, context: Table[string, IRType]): IRType =
|
||||
case expr.kind
|
||||
of irekLiteral:
|
||||
case expr.literal.kind
|
||||
of vkBool: return IRType(name: "bool", kind: itkScalar)
|
||||
of vkInt64: return IRType(name: "int64", kind: itkScalar)
|
||||
of vkFloat64: return IRType(name: "float64", kind: itkScalar)
|
||||
of vkString: return IRType(name: "str", kind: itkScalar)
|
||||
of vkNull: return IRType(name: "null", kind: itkScalar, isNullable: true)
|
||||
else: return IRType(name: "unknown", kind: itkScalar)
|
||||
of irekField:
|
||||
if expr.fieldPath.len == 0:
|
||||
return nil
|
||||
let rootName = expr.fieldPath[0]
|
||||
if rootName in context:
|
||||
var current = context[rootName]
|
||||
for i in 1..<expr.fieldPath.len:
|
||||
if expr.fieldPath[i] in current.fields:
|
||||
current = current.fields[expr.fieldPath[i]]
|
||||
else:
|
||||
return nil
|
||||
return current
|
||||
return nil
|
||||
of irekUnary:
|
||||
let operandType = tc.inferExpr(expr.unExpr, context)
|
||||
if operandType == nil:
|
||||
return nil
|
||||
case expr.unOp
|
||||
of irEq, irNeq, irLt, irLte, irGt, irGte, irAnd, irOr, irNot,
|
||||
irIsNull, irIsNotNull, irIn, irNotIn, irLike, irILike, irBetween:
|
||||
return IRType(name: "bool", kind: itkScalar)
|
||||
else:
|
||||
return nil
|
||||
of irekBinary:
|
||||
let leftType = tc.inferExpr(expr.binLeft, context)
|
||||
let rightType = tc.inferExpr(expr.binRight, context)
|
||||
if leftType == nil or rightType == nil:
|
||||
return nil
|
||||
case expr.binOp
|
||||
of irAdd, irSub, irMul, irDiv, irMod, irPow:
|
||||
return leftType
|
||||
of irEq, irNeq, irLt, irLte, irGt, irGte, irAnd, irOr,
|
||||
irIn, irNotIn, irLike, irILike, irBetween:
|
||||
return IRType(name: "bool", kind: itkScalar)
|
||||
else:
|
||||
return nil
|
||||
of irekAggregate:
|
||||
case expr.aggOp
|
||||
of irCount: return IRType(name: "int64", kind: itkScalar)
|
||||
of irSum, irAvg: return IRType(name: "float64", kind: itkScalar)
|
||||
of irMin, irMax:
|
||||
if expr.aggArgs.len > 0:
|
||||
return tc.inferExpr(expr.aggArgs[0], context)
|
||||
return nil
|
||||
of irekFuncCall:
|
||||
return IRType(name: "unknown", kind: itkScalar)
|
||||
of irekCast:
|
||||
return expr.irCastType
|
||||
of irekConditional:
|
||||
let thenType = tc.inferExpr(expr.thenExpr, context)
|
||||
return thenType
|
||||
of irekExists:
|
||||
return IRType(name: "bool", kind: itkScalar)
|
||||
@@ -0,0 +1,121 @@
|
||||
## B-Tree Index — ordered key-value index
|
||||
import std/tables
|
||||
|
||||
const
|
||||
DefaultBTreeOrder* = 32
|
||||
|
||||
type
|
||||
BTreeNode[K, V] = ref object
|
||||
keys: seq[K]
|
||||
values: seq[seq[V]]
|
||||
children: seq[BTreeNode[K, V]]
|
||||
isLeaf: bool
|
||||
next: BTreeNode[K, V]
|
||||
|
||||
BTreeIndex*[K, V] = ref object
|
||||
root: BTreeNode[K, V]
|
||||
order: int
|
||||
size: int
|
||||
|
||||
proc newBTreeNode[K, V](isLeaf: bool = true): BTreeNode[K, V] =
|
||||
BTreeNode[K, V](
|
||||
keys: @[], values: @[], children: @[],
|
||||
isLeaf: isLeaf, next: nil,
|
||||
)
|
||||
|
||||
proc newBTreeIndex*[K, V](order: int = DefaultBTreeOrder): BTreeIndex[K, V] =
|
||||
BTreeIndex[K, V](root: newBTreeNode[K, V](), order: order, size: 0)
|
||||
|
||||
proc search[K, V](node: BTreeNode[K, V], key: K): seq[V] =
|
||||
var i = 0
|
||||
while i < node.keys.len and key > node.keys[i]:
|
||||
inc i
|
||||
if node.isLeaf:
|
||||
if i < node.keys.len and key == node.keys[i]:
|
||||
return node.values[i]
|
||||
return @[]
|
||||
else:
|
||||
return search(node.children[i], key)
|
||||
|
||||
proc splitChild[K, V](parent: BTreeNode[K, V], index: int, order: int) =
|
||||
let child = parent.children[index]
|
||||
let mid = (order - 1) div 2
|
||||
let newNode = newBTreeNode[K, V](child.isLeaf)
|
||||
|
||||
for j in mid+1..<child.keys.len:
|
||||
newNode.keys.add(child.keys[j])
|
||||
if child.isLeaf:
|
||||
newNode.values.add(child.values[j])
|
||||
|
||||
if not child.isLeaf:
|
||||
for j in mid+1..child.children.len:
|
||||
newNode.children.add(child.children[j])
|
||||
child.children.setLen(mid + 1)
|
||||
|
||||
if child.isLeaf:
|
||||
newNode.next = child.next
|
||||
child.next = newNode
|
||||
|
||||
let midKey = child.keys[mid]
|
||||
parent.keys.insert(midKey, index)
|
||||
parent.children.insert(newNode, index + 1)
|
||||
child.keys.setLen(mid)
|
||||
if child.isLeaf:
|
||||
child.values.setLen(mid)
|
||||
|
||||
proc insertNonFull[K, V](node: BTreeNode[K, V], key: K, value: V, order: int) =
|
||||
var i = node.keys.len - 1
|
||||
if node.isLeaf:
|
||||
while i >= 0 and key < node.keys[i]:
|
||||
dec i
|
||||
if i >= 0 and key == node.keys[i]:
|
||||
node.values[i].add(value)
|
||||
return
|
||||
node.keys.insert(key, i + 1)
|
||||
node.values.insert(@[value], i + 1)
|
||||
else:
|
||||
while i >= 0 and key < node.keys[i]:
|
||||
dec i
|
||||
inc i
|
||||
if node.children[i].keys.len == order - 1:
|
||||
splitChild(node, i, order)
|
||||
if key > node.keys[i]:
|
||||
inc i
|
||||
insertNonFull(node.children[i], key, value, order)
|
||||
|
||||
proc insert*[K, V](btree: var BTreeIndex[K, V], key: K, value: V) =
|
||||
if btree.root.keys.len == btree.order - 1:
|
||||
var newRoot = newBTreeNode[K, V](isLeaf = false)
|
||||
newRoot.children.add(btree.root)
|
||||
splitChild(newRoot, 0, btree.order)
|
||||
btree.root = newRoot
|
||||
insertNonFull(btree.root, key, value, btree.order)
|
||||
else:
|
||||
insertNonFull(btree.root, key, value, btree.order)
|
||||
inc btree.size
|
||||
|
||||
proc get*[K, V](btree: BTreeIndex[K, V], key: K): seq[V] =
|
||||
search(btree.root, key)
|
||||
|
||||
proc contains*[K, V](btree: BTreeIndex[K, V], key: K): bool =
|
||||
return btree.get(key).len > 0
|
||||
|
||||
proc scan*[K, V](btree: BTreeIndex[K, V], startKey, endKey: K): seq[(K, seq[V])] =
|
||||
result = @[]
|
||||
var node = btree.root
|
||||
while not node.isLeaf:
|
||||
var i = 0
|
||||
while i < node.keys.len and startKey > node.keys[i]:
|
||||
inc i
|
||||
node = node.children[i]
|
||||
|
||||
while node != nil:
|
||||
for i in 0..<node.keys.len:
|
||||
if node.keys[i] >= startKey:
|
||||
if node.keys[i] <= endKey:
|
||||
result.add((node.keys[i], node.values[i]))
|
||||
else:
|
||||
return
|
||||
node = node.next
|
||||
|
||||
proc len*[K, V](btree: BTreeIndex[K, V]): int = btree.size
|
||||
@@ -0,0 +1,234 @@
|
||||
## Vector Quantization — scalar, product, binary quantization
|
||||
import std/math
|
||||
|
||||
type
|
||||
QuantizationKind* = enum
|
||||
qkNone
|
||||
qkScalar8
|
||||
qkScalar4
|
||||
qkProduct
|
||||
qkBinary
|
||||
|
||||
ScalarQuantizer* = ref object
|
||||
mins: seq[float32]
|
||||
maxes: seq[float32]
|
||||
dimensions: int
|
||||
bits: int
|
||||
|
||||
ProductQuantizer* = ref object
|
||||
codebooks: seq[seq[seq[float32]]] # subspace -> cluster -> centroid
|
||||
nSubspaces: int
|
||||
nClusters: int
|
||||
dimensions: int
|
||||
subDim: int
|
||||
|
||||
QuantizedVector* = ref object
|
||||
case kind*: QuantizationKind
|
||||
of qkScalar8: int8Data*: seq[int8]
|
||||
of qkScalar4: int4Data*: seq[int8] # packed
|
||||
of qkProduct: pqCodes*: seq[int8]
|
||||
of qkBinary: binData*: seq[uint64] # packed bits
|
||||
of qkNone: orig*: seq[float32]
|
||||
|
||||
proc newScalarQuantizer*(dimensions: int, bits: int = 8): ScalarQuantizer =
|
||||
ScalarQuantizer(
|
||||
mins: newSeq[float32](dimensions),
|
||||
maxes: newSeq[float32](dimensions),
|
||||
dimensions: dimensions,
|
||||
bits: bits,
|
||||
)
|
||||
|
||||
proc train*(sq: ScalarQuantizer, vectors: openArray[seq[float32]]) =
|
||||
if vectors.len == 0:
|
||||
return
|
||||
for d in 0..<sq.dimensions:
|
||||
var minVal: float32 = high(float32)
|
||||
var maxVal: float32 = low(float32)
|
||||
for v in vectors:
|
||||
if d < v.len:
|
||||
if v[d] < minVal: minVal = v[d]
|
||||
if v[d] > maxVal: maxVal = v[d]
|
||||
sq.mins[d] = minVal
|
||||
sq.maxes[d] = maxVal
|
||||
|
||||
proc encode*(sq: ScalarQuantizer, vector: seq[float32]): QuantizedVector =
|
||||
result = QuantizedVector(kind: if sq.bits == 8: qkScalar8 else: qkScalar4)
|
||||
let levels = float32(1 shl sq.bits) - 1.0'f32
|
||||
|
||||
if sq.bits == 8:
|
||||
result.int8Data = newSeq[int8](sq.dimensions)
|
||||
for d in 0..<sq.dimensions:
|
||||
let range = sq.maxes[d] - sq.mins[d]
|
||||
if range == 0:
|
||||
result.int8Data[d] = 0
|
||||
else:
|
||||
let normalized = (vector[d] - sq.mins[d]) / range
|
||||
result.int8Data[d] = int8(normalized * levels)
|
||||
elif sq.bits == 4:
|
||||
# Pack 2 values per byte
|
||||
result.int4Data = newSeq[int8](sq.dimensions div 2 + sq.dimensions mod 2)
|
||||
for d in 0..<sq.dimensions:
|
||||
let range = sq.maxes[d] - sq.mins[d]
|
||||
var val: int8 = 0
|
||||
if range != 0:
|
||||
let normalized = (vector[d] - sq.mins[d]) / range
|
||||
val = int8(normalized * 15)
|
||||
let idx = d div 2
|
||||
if d mod 2 == 0:
|
||||
result.int4Data[idx] = val shl 4
|
||||
else:
|
||||
result.int4Data[idx] = result.int4Data[idx] or val
|
||||
|
||||
proc decode*(sq: ScalarQuantizer, qv: QuantizedVector): seq[float32] =
|
||||
result = newSeq[float32](sq.dimensions)
|
||||
if qv.kind == qkScalar8:
|
||||
let levels = 255.0'f32
|
||||
for d in 0..<sq.dimensions:
|
||||
let range = sq.maxes[d] - sq.mins[d]
|
||||
result[d] = sq.mins[d] + float32(qv.int8Data[d]) / levels * range
|
||||
elif qv.kind == qkScalar4:
|
||||
let levels = 15.0'f32
|
||||
for d in 0..<sq.dimensions:
|
||||
let idx = d div 2
|
||||
var val: int8
|
||||
if d mod 2 == 0:
|
||||
val = (qv.int4Data[idx] shr 4) and 0x0F
|
||||
else:
|
||||
val = qv.int4Data[idx] and 0x0F
|
||||
let range = sq.maxes[d] - sq.mins[d]
|
||||
result[d] = sq.mins[d] + float32(val) / levels * range
|
||||
|
||||
proc distance*(sq: ScalarQuantizer, qv: QuantizedVector, query: seq[float32]): float64 =
|
||||
let decoded = sq.decode(qv)
|
||||
var sum: float64
|
||||
for d in 0..<sq.dimensions:
|
||||
let diff = float64(decoded[d]) - float64(query[d])
|
||||
sum += diff * diff
|
||||
return sqrt(sum)
|
||||
|
||||
proc newProductQuantizer*(dimensions: int, nSubspaces: int = 8, nClusters: int = 256): ProductQuantizer =
|
||||
let subDim = dimensions div nSubspaces
|
||||
ProductQuantizer(
|
||||
codebooks: newSeq[seq[seq[float32]]](nSubspaces),
|
||||
nSubspaces: nSubspaces,
|
||||
nClusters: nClusters,
|
||||
dimensions: dimensions,
|
||||
subDim: subDim,
|
||||
)
|
||||
|
||||
proc train*(pq: ProductQuantizer, vectors: openArray[seq[float32]], nIterations: int = 20) =
|
||||
if vectors.len == 0:
|
||||
return
|
||||
|
||||
for s in 0..<pq.nSubspaces:
|
||||
pq.codebooks[s] = newSeq[seq[float32]](pq.nClusters)
|
||||
for c in 0..<pq.nClusters:
|
||||
pq.codebooks[s][c] = newSeq[float32](pq.subDim)
|
||||
|
||||
# Initialize centroids randomly from data
|
||||
for c in 0..<pq.nClusters:
|
||||
let idx = min(c, vectors.len - 1)
|
||||
for d in 0..<pq.subDim:
|
||||
let globalD = s * pq.subDim + d
|
||||
if globalD < vectors[idx].len:
|
||||
pq.codebooks[s][c][d] = vectors[idx][globalD]
|
||||
|
||||
# K-means per subspace
|
||||
var assignments = newSeq[int](vectors.len)
|
||||
for iter in 0..<nIterations:
|
||||
# Assign vectors to clusters
|
||||
for vi, v in vectors:
|
||||
var bestCluster = 0
|
||||
var bestDist = high(float64)
|
||||
for ci in 0..<pq.nClusters:
|
||||
var dist: float64 = 0
|
||||
for d in 0..<pq.subDim:
|
||||
let globalD = s * pq.subDim + d
|
||||
if globalD < v.len:
|
||||
let diff = float64(v[globalD]) - float64(pq.codebooks[s][ci][d])
|
||||
dist += diff * diff
|
||||
if dist < bestDist:
|
||||
bestDist = dist
|
||||
bestCluster = ci
|
||||
assignments[vi] = bestCluster
|
||||
|
||||
# Update centroids
|
||||
var clusterCounts = newSeq[int](pq.nClusters)
|
||||
var newCentroids = newSeq[seq[float64]](pq.nClusters)
|
||||
for c in 0..<pq.nClusters:
|
||||
newCentroids[c] = newSeq[float64](pq.subDim)
|
||||
|
||||
for vi, v in vectors:
|
||||
let ci = assignments[vi]
|
||||
inc clusterCounts[ci]
|
||||
for d in 0..<pq.subDim:
|
||||
let globalD = s * pq.subDim + d
|
||||
if globalD < v.len:
|
||||
newCentroids[ci][d] += float64(v[globalD])
|
||||
|
||||
for ci in 0..<pq.nClusters:
|
||||
if clusterCounts[ci] > 0:
|
||||
for d in 0..<pq.subDim:
|
||||
pq.codebooks[s][ci][d] = float32(newCentroids[ci][d] / float64(clusterCounts[ci]))
|
||||
|
||||
proc encode*(pq: ProductQuantizer, vector: seq[float32]): QuantizedVector =
|
||||
result = QuantizedVector(kind: qkProduct, pqCodes: newSeq[int8](pq.nSubspaces))
|
||||
for s in 0..<pq.nSubspaces:
|
||||
var bestCluster: int8 = 0
|
||||
var bestDist = high(float64)
|
||||
for ci in 0..<pq.nClusters:
|
||||
var dist: float64 = 0
|
||||
for d in 0..<pq.subDim:
|
||||
let globalD = s * pq.subDim + d
|
||||
if globalD < vector.len:
|
||||
let diff = float64(vector[globalD]) - float64(pq.codebooks[s][ci][d])
|
||||
dist += diff * diff
|
||||
if dist < bestDist:
|
||||
bestDist = dist
|
||||
bestCluster = int8(ci)
|
||||
result.pqCodes[s] = bestCluster
|
||||
|
||||
proc distance*(pq: ProductQuantizer, qv: QuantizedVector, query: seq[float32]): float64 =
|
||||
var sum: float64 = 0
|
||||
for s in 0..<pq.nSubspaces:
|
||||
let ci = qv.pqCodes[s]
|
||||
for d in 0..<pq.subDim:
|
||||
let globalD = s * pq.subDim + d
|
||||
if globalD < query.len:
|
||||
let diff = float64(pq.codebooks[s][ci][d]) - float64(query[globalD])
|
||||
sum += diff * diff
|
||||
return sqrt(sum)
|
||||
|
||||
# Binary quantization
|
||||
proc binaryQuantize*(vector: seq[float32]): QuantizedVector =
|
||||
result = QuantizedVector(kind: qkBinary)
|
||||
let bits = vector.len
|
||||
let words = (bits + 63) div 64
|
||||
result.binData = newSeq[uint64](words)
|
||||
for i in 0..<vector.len:
|
||||
if vector[i] >= 0:
|
||||
let wordIdx = i div 64
|
||||
let bitIdx = i mod 64
|
||||
result.binData[wordIdx] = result.binData[wordIdx] or (1'u64 shl bitIdx)
|
||||
|
||||
proc binaryDistance*(a, b: QuantizedVector): int =
|
||||
result = 0
|
||||
let words = min(a.binData.len, b.binData.len)
|
||||
for i in 0..<words:
|
||||
let val = a.binData[i] xor b.binData[i]
|
||||
var cnt = 0
|
||||
var v = val
|
||||
while v != 0:
|
||||
v = v and (v - 1)
|
||||
inc cnt
|
||||
result += cnt
|
||||
|
||||
proc compressionRatio*(sq: ScalarQuantizer): float64 =
|
||||
if sq.bits == 8: return 4.0
|
||||
if sq.bits == 4: return 8.0
|
||||
return 1.0
|
||||
|
||||
proc compressionRatio*(pq: ProductQuantizer): float64 =
|
||||
let origBytes = pq.dimensions * 4
|
||||
let pqBytes = pq.nSubspaces # one byte per subspace code
|
||||
return float64(origBytes) / float64(pqBytes)
|
||||
Reference in New Issue
Block a user