Something went wrong. Try again.
My own corner of monopam
Something went wrong. Try again.
17 kB · 449 lines
OCaml
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450(*--------------------------------------------------------------------------- Copyright (c) 2026 Thomas Gazagnaire. All rights reserved. SPDX-License-Identifier: ISC ---------------------------------------------------------------------------*)
module Ir = Catalog.Irmodule Spi = Catalog.Spimodule V = Catalog.Value
(* One window function as this pass runs it: the call, and the position in the projected input row of each expression it reads. *)type bound = | Unbounded_preceding | Preceding of int | Current_row | Following of int | Unbounded_following
type frame = { rows : bool; start : bound; end_ : bound }
type win = { name : string; star : bool; (* a star [count]: it counts rows and reads no argument *) args : int list; coll : Catalog.Collation.t option; (* the sequence the fold's own comparison runs under, read only by a fold that compares its arguments *) partition : (int * Catalog.Collation.t option) list; order : (int * Ast.dir * Catalog.Collation.t option) list; frame : frame option;}
(* The projected input columns, in the order they are registered. An expression appearing twice (the same partition key under two windows, say) is projected once. *)let registry () = let acc = ref [] and n = ref 0 in let slot e = let rec go = function | (e', i) :: _ when Ast.equal_expr e' e -> i | _ :: t -> go t | [] -> let i = !n in acc := !acc @ [ (e, i) ]; incr n; i in go !acc in (slot, fun () -> List.map fst !acc)
(* The call under an [OVER]: its name and its arguments. DISTINCT is not one of them -- sqlite3WindowAttach (src/window.c) refuses it as it attaches the window, and the parser refuses it at the same point -- so nothing reaching here folds distinct arguments. *)let call = function | Ast.Func (n, args) -> (n, args) | _ -> failwith "window: not a function call"
let bare = Collation.striplet collation = Collation.sequence
let spec slot (e : Ast.expr) = let f, w = match e with Ast.Window (f, w) -> (f, w) | _ -> assert false in let name, args = call f in let key e = slot (bare e) in (* A star [count] counts the rows and has no argument to project; every other call projects each of its arguments. *) let star = name = "count" && match args with [ a ] -> a = Ast.Star | _ -> true in let bound : Ast.frame_bound -> bound = function | Unbounded_preceding -> Unbounded_preceding | Preceding e -> Preceding (key e) | Current_row -> Current_row | Following e -> Following (key e) | Unbounded_following -> Unbounded_following in { name; star; args = (if star then [] else List.map key args); (* A fold that compares its arguments -- [min], [max] -- reads a sequence of its own off the argument list, exactly as the same fold over a GROUP BY group does: one {!Collation.call_sequence} for both. Every other fold orders nothing and carries none. *) coll = (if Collation.fold_compares name then Collation.call_sequence args else None); partition = List.map (fun e -> (key e, collation e)) w.partition; order = List.map (fun (e, d) -> (key e, d, collation e)) w.Ast.order; frame = Option.map (fun (fr : Ast.frame) -> { rows = fr.rows; start = bound fr.start; end_ = bound fr.end_ }) w.frame; }
(* ---- Running an expression projection on the register VM ---------------- *)
let project ~params ?affinity ~membership layout (s : Ast.select) rows = let q, subs = Lower.query ~params ?affinity ~membership layout s in if Lower.needs_runner subs then raise Ir.Unsupported; let builtins = { Lower.builtins with subquery = Lower.member subs } in Ir.run ~builtins ~params ~collate:Lower.collate (Ir.rebase q) [| Spi.list_cursor rows |]
(* ---- Partitioning, ordering, framing ------------------------------------ *)
let order_cmp keys (a : V.t array) (b : V.t array) = let rec go = function | [] -> 0 | (col, dir, coll) :: rest -> let c = match coll with | Some seq -> Lower.collate seq a.(col) b.(col) | None -> V.compare a.(col) b.(col) in let c = match (dir : Ast.dir) with Asc -> c | Desc -> -c in if c <> 0 then c else go rest in go keys
let frame_bounds ~pos ~k ~peer_start ~peer_end ~off w = let lo, hi = match w.frame with | None -> (0, if w.order = [] then k - 1 else peer_end.(pos)) | Some f -> let resolve at_start = function | Unbounded_preceding -> 0 | Unbounded_following -> k - 1 | Current_row -> if f.rows then pos else if at_start then peer_start.(pos) else peer_end.(pos) | Preceding c -> if f.rows then pos - off c else peer_start.(pos) | Following c -> if f.rows then pos + off c else peer_end.(pos) in (resolve true f.start, resolve false f.end_) in (max 0 lo, min (k - 1) hi)
(* One row's value for a navigation or aggregate window function. [at i c] is input column [c] of the [i]-th ordered row, [cur c] the same at [pos]. *)let lag_lead ~name ~args ~k ~pos ~at ~cur = let off = match List.nth_opt args 1 with | Some c -> Int64.to_int (Catalog.Value_ops.intify (cur c)) | None -> 1 in let dflt = match List.nth_opt args 2 with Some c -> cur c | None -> V.Null in let t = if name = "lag" then pos - off else pos + off in if t >= 0 && t < k then at t (List.hd args) else dflt
let nth_value ~args ~lo ~hi ~at ~cur = let n = Int64.to_int (Catalog.Value_ops.intify (cur (List.nth args 1))) in let t = lo + (n - 1) in if n >= 1 && t <= hi then at t (List.hd args) else V.Null
let ntile ~args ~k ~pos ~cur = let nb = max 1 (Int64.to_int (Catalog.Value_ops.intify (cur (List.hd args)))) in let base = k / nb and rem = k mod nb in let big = rem * (base + 1) in let b = if pos < big then pos / (base + 1) else rem + ((pos - big) / base) in V.Int (Int64.of_int (b + 1))
(* Fold an aggregate over the frame's rows through the shared {!Catalog.Value_ops.agg_acc}, the same fold the register VM's GROUP BY path uses, so one policy decides what an aggregate means wherever it is folded -- the sequence a comparing fold runs under included. No dedup: a window call carries no DISTINCT to fold under ({!call}). *)let agg_over w ~inputs ~oa ~lo ~hi = if w.star then ( (* a star [count] reads no argument, so it compares nothing *) let acc = Catalog.Value_ops.agg_acc w.name ~nargs:0 ~coll:None in for _ = lo to hi do acc.step [||] 0 done; acc.finish ()) else if w.args = [] then Fmt.failwith "wrong number of arguments to function %s()" w.name else let tuples = if lo > hi then [] else List.init (hi - lo + 1) (fun r -> let row = inputs.(oa.(lo + r)) in Array.of_list (List.map (fun c -> row.(c)) w.args)) in let acc = Catalog.Value_ops.agg_acc w.name ~nargs:(List.length w.args) ~coll:w.coll in List.iter (fun t -> acc.step t 0) tuples; acc.finish ()
let percent_rank ~peer_start ~k ~pos = if k > 1 then V.Float (float_of_int peer_start.(pos) /. float_of_int (k - 1)) else V.Float 0.
let cume_dist ~peer_end ~k ~pos = V.Float (float_of_int (peer_end.(pos) + 1) /. float_of_int k)
let value w ~inputs ~oa ~peer_start ~peer_end ~k ~dense ~pos ~lo ~hi = let at i c = inputs.(oa.(i)).(c) in let cur c = at pos c in let args = w.args in match w.name with | "row_number" -> V.Int (Int64.of_int (pos + 1)) | "rank" -> V.Int (Int64.of_int (peer_start.(pos) + 1)) | "dense_rank" -> V.Int (Int64.of_int dense) | "percent_rank" -> percent_rank ~peer_start ~k ~pos | "cume_dist" -> cume_dist ~peer_end ~k ~pos | "first_value" -> if lo <= hi then at lo (List.hd args) else V.Null | "last_value" -> if lo <= hi then at hi (List.hd args) else V.Null | "lag" | "lead" -> lag_lead ~name:w.name ~args ~k ~pos ~at ~cur | "nth_value" -> nth_value ~args ~lo ~hi ~at ~cur | "ntile" -> ntile ~args ~k ~pos ~cur | _ -> agg_over w ~inputs ~oa ~lo ~hi
let tie_break ~grouped w = let desc (_, d, _) = d = Ast.Desc in if grouped && w.order <> [] && List.for_all desc w.order then fun i j -> compare j i else fun _ _ -> 0
(* The row indices [idxs] in the order the window sorts them: by its ORDER BY, ties broken as {!tie_break} says. *)let ordered ~grouped w inputs idxs = let tie = tie_break ~grouped w in List.stable_sort (fun i j -> match order_cmp w.order inputs.(i) inputs.(j) with 0 -> tie i j | c -> c) idxs
(* Compute one partition's values into [out]: order its row indices by the window's ORDER BY, find the peer-group bounds (peers share a rank and a frame), then write each row's value. *)let fill_partition ~grouped w inputs out idxs = let oa = Array.of_list (ordered ~grouped w inputs idxs) in let k = Array.length oa in let peer_start = Array.make k 0 and peer_end = Array.make k 0 in let p = ref 0 in while !p < k do let q = ref !p in while !q + 1 < k && order_cmp w.order inputs.(oa.(!p)) inputs.(oa.(!q + 1)) = 0 do incr q done; for r = !p to !q do peer_start.(r) <- !p; peer_end.(r) <- !q done; p := !q + 1 done; let dense = ref 0 and prev_start = ref (-1) in for pos = 0 to k - 1 do if peer_start.(pos) <> !prev_start then begin incr dense; prev_start := peer_start.(pos) end; let off c = Int64.to_int (Catalog.Value_ops.intify inputs.(oa.(pos)).(c)) in let lo, hi = frame_bounds ~pos ~k ~peer_start ~peer_end ~off w in out.(oa.(pos)) <- value w ~inputs ~oa ~peer_start ~peer_end ~k ~dense:!dense ~pos ~lo ~hi done
(* Two rows' PARTITION BY keys compared, most significant first, each under its own collating sequence. sqlite3 gives the window's sorter one KeyInfo entry per key with that key's collation (window.c), so a partition's sequence and the window's ORDER BY sequences are independent of one another. *)let partition_cmp w inputs i j = let rec go = function | [] -> 0 | (c, coll) :: rest -> let d = match coll with | Some seq -> Lower.collate seq inputs.(i).(c) inputs.(j).(c) | None -> V.compare inputs.(i).(c) inputs.(j).(c) in if d <> 0 then d else go rest in go w.partition
let partitions w inputs = let cmp = partition_cmp w inputs in let sorted = List.stable_sort cmp (List.init (Array.length inputs) Fun.id) in let rec runs = function | [] -> [] | i :: rest -> let rec take acc = function | j :: t when cmp i j = 0 -> take (j :: acc) t | t -> (List.rev acc, t) in let group, tail = take [ i ] rest in group :: runs tail in runs sorted
(* [w]'s value for every input row, aligned to [inputs]: partition by the key tuple, then fill each partition. *)let evaluate ~grouped w inputs = let out = Array.make (Array.length inputs) V.Null in List.iter (fill_partition ~grouped w inputs out) (partitions w inputs); out
let row_order ~grouped w inputs = List.concat_map (ordered ~grouped w inputs) (partitions w inputs)
(* ---- The pass ----------------------------------------------------------- *)
(* The synthetic column a window's value is carried in for the second projection. The NUL byte keeps the name unspellable, so no user column collides with it. *)let win_column i = Fmt.str "\000win%d" i
(* [*] and [t.*] listed column by column, over the sources the rows came from. Expanded here rather than by the lowering, whose layout also carries the window columns, which a bare [*] must not reach. *)let expand_star segments (rc : Ast.result_col) : Ast.result_col list = let cols_of (alias, cols) = let qual = if alias = "" then None else Some alias in List.filter_map (fun c : Ast.result_col option -> if String.equal c Lower.rowid_sentinel then None else Some { expr = Ast.Col (qual, c); alias = None }) cols in match rc.Ast.expr with | Star -> List.concat_map cols_of segments | Qualified_star q -> ( let same (al, _) = String.uppercase_ascii al = String.uppercase_ascii q in match List.find_opt same segments with | Some seg -> cols_of seg | None -> [ rc ]) | _ -> [ rc ]
(* Replace each [Window] node by a reference to the column its value was written to. Windows do not nest, so a [Window] is not descended into. *)let rec substitute at (e : Ast.expr) : Ast.expr = match e with | Window _ -> Ast.Col (None, at e) (* A subquery computes its own windows over its own rows, so it is handed back as written; every other part of the expression is projected here and reads the column this pass wrote. *) | e -> Ast.map_expr ~select:Fun.id (substitute at) e
(* [rows], each widened by one column per window, in the order the first window sorts them -- as sqlite3 emits them: the sort a window needs is what the caller sees when nothing else orders the result. *)let widen ~grouped specs values inputs rows = let widened = List.mapi (fun i row -> row @ List.map (fun v -> v.(i)) values) rows in match specs with | [] -> widened | w :: _ -> let a = Array.of_list widened in List.map (fun i -> a.(i)) (row_order ~grouped w inputs)
(* The select list, ORDER BY, DISTINCT and LIMIT read back over the widened rows, each window replaced by the column [at] names for it. *)let second_projection at cols (s : Ast.select) = { (Lower.post_select (List.map (fun (rc : Ast.result_col) -> { rc with Ast.expr = substitute at rc.expr }) cols)) with Ast.distinct = s.distinct; order_by = List.map (fun (e, d) -> (substitute at e, d)) s.order_by; limit = s.limit; offset = s.offset; }
let scan_select (layout : Lower.layout) (s : Ast.select) = let rowid (alias, cols) = if List.mem Lower.rowid_sentinel cols then let qual = if alias = "" then None else Some alias in Some (alias, { Ast.expr = Col (qual, Lower.rowid_sentinel); alias = None }) else None in let rowids = List.filter_map rowid layout in let visible = List.map (fun (a, c) -> (a, List.filter (fun n -> n <> Lower.rowid_sentinel) c)) layout in ( { s with Ast.cols = { Ast.expr = Star; alias = None } :: List.map snd rowids; group_by = []; having = None; order_by = []; limit = None; offset = None; distinct = false; }, visible @ List.map (fun (a, _) -> (a, [ Lower.rowid_sentinel ])) rowids )
(* A rowid read through the segments {!scan_select} appends is an INTEGER, as it is on the source it came from. *)let rowid_affinity segments affinity (i, j) = match List.nth_opt segments i with | Some (_, cols) when List.nth_opt cols j = Some Lower.rowid_sentinel -> Some `Integer | _ -> affinity (i, j)
let scan ~params ?affinity ?(membership = fun _ -> Query.opaque_membership) ?(grouped = false) (segments : Lower.layout) (s : Ast.select) rows = let affinity = Option.map (rowid_affinity segments) affinity in let cols = List.concat_map (expand_star segments) s.Ast.cols in (* the select list and the ORDER BY both name windows this pass computes; expanding a star adds none, so the two agree on which *) let windows = Query.select_windows s in let slot, projected = registry () in let specs = List.map (spec slot) windows in let result_col e : Ast.result_col = { expr = e; alias = None } in let inputs = Array.of_list (List.map Array.of_list (project ~params ?affinity ~membership segments (Lower.post_select (List.map result_col (projected ()))) rows)) in let values = List.map (fun w -> evaluate ~grouped w inputs) specs in let names = List.mapi (fun i _ -> win_column i) windows in (* Every window the select names has a column here, the ORDER BY's included, so a node reaching this without one is a window the collection missed rather than a shape to answer differently. *) let at = let table = List.combine windows names in fun e -> match List.find_opt (fun (w, _) -> Ast.equal_expr w e) table with | Some (_, n) -> n | None -> raise Ir.Unsupported in let out = second_projection at cols s in project ~params ?affinity ~membership (segments @ [ ("", names) ]) out (widen ~grouped specs values inputs rows)