Module Talon.Kit

head n q is the first n rows of q:

Query.slice ~offset:0 ~length:n q

tail n q is the last n rows of q:

Query.slice ~offset:(-n) ~length:n q

top_k k keys q is the first k rows of q in the order of keys:

Query.(q |> sort keys |> slice ~offset:0 ~length:k)

A slice from the start of a sort runs as a selection of its first rows, which sorts only the rows that precede or tie with the kth on the first key.

distinct q is the first of each set of q's rows that are the same on every column, by key identity (null is one key, and NaN is one), in order:

match List.map fst (Schema.columns (Query.schema q)) with
| [] -> head 1 q
| names -> Query.aggregate ~by:names [] q

A query without columns has one distinct row if it has rows.

count_by ks q is one row per group of q's rows that have the same keys in the columns ks, in order of first appearance: the columns ks, then "count", the group's number of rows as int64:

Query.aggregate ~by:ks Expr.[ "count" := rows ] q

value_counts c q is each distinct value of q's column c, null included, with its number of rows, most frequent first and ties in order of first appearance:

count_by [ c ] q |> Query.sort [ Order.desc "count" ]

describe q summarizes each integer or float column of q, in order, as one row of the columns "column", its name; "count", its number of values that are not null; "nulls", its number of nulls; "mean"; "std"; "min"; "q25", "median" and "q75", its quartiles; and "max". "count" and "nulls" are int64 and the others float64. The statistics are Expr's reductions, which skip nulls:

let stats name x =
  Query.aggregate ~by:[]
    Expr.
      [
        "column" := string name;
        "count" := count x;
        "nulls" := rows - count x;
        "mean" := mean x;
        "std" := std x;
        "min" := cast Type.float64 (min x);
        "q25" := quantile 0.25 x;
        "median" := median x;
        "q75" := quantile 0.75 x;
        "max" := cast Type.float64 (max x);
      ]
    q
in
s0 |> Query.append s1 |> … |> Query.append sk

where s0 to sk are stats n (Col.int n) or stats n (Col.float n) for each integer or float column n. Without such a column, describe q is head 0 (stats "" Expr.(store Type.float64 null)), which has no rows. It reads q once per integer or float column.

null_count q is one row with, for each column of q, in order and under its name, its number of nulls as int64:

Query.aggregate ~by:[]
  Expr.[ each Sel.all { column = (fun n x -> n := rows - count x) } ]
  q

drop sel q is q without the columns that sel selects:

Query.select Expr.[ keep Sel.(all - sel) ] q

rename pairs q is q with each column old of pairs renamed to the name it pairs with, in place:

let name n = Option.value ~default:n (List.assoc_opt n pairs) in
Query.select
  Expr.
    [
      each
        Sel.(all + names (List.map fst pairs))
        { column = (fun n x -> name n := x) };
    ]
  q

An old that q lacks is a problem of the select, through Sel.names, and so is a name that two columns would take.

Raises Invalid_argument if pairs renames a column twice.

complete ks q is q with a row for each combination of the values of the columns ks that q lacks, its other columns null. Each column's values come in order of first appearance, combinations vary the last column fastest, and each combination is followed by its rows of q, in order:

let values k = Query.select Expr.[ keep Sel.(names [ k ]) ] q |> distinct in
let combos =
  List.fold_left
    (fun acc k -> acc |> Query.join ~on:Join.all (values k))
    (values k0) ks'                               (* ks = k0 :: ks' *)
in
combos
|> Query.join ~kind:Left ~on:(Join.keys ks) q
|> Query.select Expr.[ keep (Sel.names (columns of q)) ]

It has q's schema. Its key columns hold the combinations' values, so where a key holds floats, a row of q holds the first of its equal values to appear in q: -0. or 0..

Raises Invalid_argument if ks is empty or names a column twice.

longer ~name ~value sel q is q with the columns of sel replaced by two columns: name, the categorical whose dictionary is the selected names in order, and value, their cells at the type where their types meet (Type.common), one row per cell, nulls included. The other columns come first, repeated, in q's order, then name and value, each row's cells in column order. It is:

let ns = (* the names [sel] selects, in order *) in
let others = (* [q]'s columns but [ns] *) in
let cells = Array.of_list ns in
Query.unnest
  Expr.[
    name := store Type.(list (categorical cells)) (const cells);
    value := list (List.map (Col.v k) ns);
  ]
  q
|> Query.select Expr.[ keep (Sel.names (others @ [ name; value ])) ]

with k the kind of the first selected column. It streams in O(n k), and building and optimizing its plan take O(k²) time. Filtering with Expr.is_null after it drops null cells. Selected columns of different kinds are a problem of the handle; widen them first with Expr.cast.

Raises Invalid_argument if sel selects no column, if name or value is a column sel does not select, or if the first selected column holds an extension, which no Col handle reads.

wider ~by ~key ~value q is one row per group of q's rows that have the same by, in order of first appearance, holding by, then one column per category of key's categorical type, in dictionary order, holding value at the row of that group and category, null where there is none (Expr.pivot). With by empty it is one row, even from no rows. It is:

Query.aggregate ~by
  Expr.[
    each (Sel.names [ value ])
      { column = (fun n x -> n := pivot (Col.string key) x) };
  ]
  q
|> Query.filter Expr.(not (is_null (Col.v Record.kind value)))
|> Query.select
     Expr.[
       keep (Sel.names by);
       within (Col.v Record.kind value) (keep Sel.all);
     ]

The columns come from key's type, so the plan's schema is known before any data is read. Cast a key to the categorical type whose dictionary names the columns, or read them from the data with a run:

let open Result.Syntax in
let* q = Kit.categorize [ "element" ] q in
Ok (Kit.wider ~by:[ "station"; "date" ] ~key:"element" ~value:"value" q)

An unused category is a column of nulls; categorize a filtered query to read only its categories. Expr.pivot states the failures: duplicates in a cell, a null key holding a value, and filtering before it. After longer ~name ~value sel, wider ~by ~key:name ~value restores the input when the selected columns have one type, by is a key of the input, lists its other columns in their order, and is non-empty or the input has a row, and those columns precede the selected ones. Equality there is by key identity (Type.compare_value).

It blocks in O(n + groups k), and building and optimizing its plan take O(k²) time in its number of columns. Its problems are those of its aggregate and select: a key that is not categorical, a category named as a column of by.

one_hot c q is q with its categorical column c replaced, in place, by one bool column per category, in dictionary order, named c ^ "_" ^ category: true where c is that category, false where it is another, and null where c is null:

Query.select
  Expr.(
    [ keep (Sel.names before) ]
    @ List.map
        (fun cat -> c ^ "_" ^ cat := Col.string c = string cat)
        categories
    @ [ keep (Sel.names after) ])
  q

The categories are in c's type, so the columns are known before any data is read; a string column is cast to a categorical type first.

If q lacks c, it is Query.select Expr.[ keep Sel.(names [ c ]) ] q's problem, reported with the nearest names.

Raises Invalid_argument if c is not categorical.

union rest q is q's rows, then rest's, written q |> Kit.union rest, with q's columns, then those that only rest has, in order; a column is null on the side that lacks it:

let nulls from into =
  List.filter_map
    (fun (n, Type.Any t) ->
      match Schema.find into n with
      | Some _ -> None
      | None -> Some Expr.(n := store t null))
    (Schema.columns from)
in
let pad os q = if os = [] then q else Query.derive os q in
let s = Query.schema q and r = Query.schema rest in
pad (nulls r s) q |> Query.append (pad (nulls s r) rest)

A column that both have, of two types, takes the type where they meet, and one whose types do not meet is a problem of the append.

union_all qs is the rows of qs, in order, with the columns of the first query, then each column that a later one adds, in order, null on the queries that lack it. Each column has the type where all of its types meet (Type.common), which no order of qs can change: columns of int8, uint8 and int16 make int16, where a fold of union in that order fails. Every query is cast to those types and padded before any is appended, in a balanced tree of appends, so its plan takes work linear in the length of qs, beyond a comparison of each leaf with those of its name and shape. A filter on a column stored at its common type reaches every such query, whatever the others' types; a query whose column was cast is read and filtered after the cast. A column whose types do not meet is a problem of the appends. union_all [] is a query of no column and no row, the identity of union: q |> union (union_all []) has q's rows and columns.

cumulative r is r over the frame's rows up to and including each:

Expr.window r

fill_forward x is x, or where x is null its last earlier value in the frame:

Expr.(window (last x))

fill_backward x is x, or where x is null its next value in the frame:

Expr.(window ~before:(int 1) ~after:(int max_int) (first x))

interpolate ~on x is x where it holds a value. Elsewhere it is the value at the row's key on the line through the last value at or before that key and the first value after it, in float64. Where a row of the same key holds a value, it is that value. It is null before the first value and after the last. Keys are on's values cast to float64, an integer or a float, or without on the rows' places. Inside Expr.( … ) it is:

let line ~at ~x0 ~k0 ~x1 ~k1 =
  coalesce
    [ x; if_ (at = k0) x0 (x0 +. (x1 -. x0) *. (at -. k0) /. (k1 -. k0)) ]
in
match on with
| Some on ->
    let k = cast Type.float64 on in
    let held v = if_ (is_null x) null v in
    let upto r = window ~on:k r in
    let past r =
      window ~on:k ~before:(float 0.) ~after:(float infinity) r
    in
    line ~at:k ~x0:(upto (last x)) ~k0:(upto (last (held k)))
      ~x1:(past (first x)) ~k1:(past (first (held k)))
| None ->
    (* [c] values reach the row; the [c]th starts the run of rows of
       count [c], which ends before the next value. *)
    let place r = cast Type.float64 r in
    let c = window (count x) in
    let k0 = place (window ~on:c ~exclude_key:true rows + int 1) in
    let k1 = k0 +. place (window ~on:c ~before:(int 1) rows) in
    line ~at:(place (window rows)) ~x0:(window (last x)) ~k0
      ~x1:(window ~before:(int 0) ~after:(int max_int) (first x)) ~k1

The type for the directions of an as-of join.

join_asof ~by ~direction ~strict ~on:(l, r) right left is each row of left, in order, with its pick in right (As-of joins), written left |> Kit.join_asof ~on right. direction defaults to Backward and strict to false. Its columns are left's, then right's without by, null where a row has no pick. Neither side need be sorted. A tolerance is applied after the join (Tolerances). Definition shows its definition.

A name on both sides, r included when it is l, is a problem of the join.

Raises Invalid_argument if by names a column twice, or if direction is Nearest and the keys meet at another type.

merge_intervals ~by (s, e) q is q's sets in normal form (Normal form), so periods that overlap merge, and so do periods that touch, where one's stop is the next one's start. It is the sweep (Sweep) without rest and with holds a _ = a > 0. The union of two sets is the merge of their Query.append.

intersect_intervals ~by (s, e) rest q is the values that both q's and rest's set of a key hold, written q |> Kit.intersect_intervals (s, e) rest, in normal form (Normal form). It is the sweep (Sweep) with holds a b = a > 0 && b > 0.

subtract_intervals ~by (s, e) rest q is the values that q's set of a key holds and rest's does not, written q |> Kit.subtract_intervals (s, e) rest, in normal form (Normal form). It is the sweep (Sweep) with holds a b = a > 0 && b = 0.

shuffle ~seed q is q's rows in the order of the hash of each row's place in q, counted across its batches:

Query.derive Expr.[ key := hash ~seed (window rows) ] q
|> Query.sort [ Order.asc key ]
|> drop (Sel.names [ key ])

It reads no value, so equal rows take independent places, and it blocks.

sample ~seed size q draws rows of q without replacement. `Rows n is the first n rows of shuffle's order, or all of them where q has fewer. `Fraction f keeps, in q's order, the rows whose place hashes below f * 2{^ 62}, so at one seed a smaller fraction keeps a subset of a larger one:

Query.filter Expr.(hash ~seed (window rows) < int cutoff) q

It blocks. Bootstrapping with replacement is Talon.take of Nx.Rng.randint indices.

Raises Invalid_argument if n is negative or f is not in [0;1].

folds ~seed ~by k q is q with a last column "fold", an int64 in [0;k - 1], that gives every row of a group of by one fold, by key identity (null is one group and NaN one). Groups rank by the hash of the record of by's columns, in that order, and the rth of g groups takes fold ⌊k (r - 1) / g⌋, so fold sizes in groups differ by at most one. It blocks, and reads every key before the first fold is known.

Raises Invalid_argument if by is empty or names a column twice, if k is below 1 or if q already has a column "fold"; a run raises it if q has fewer than k groups.

time_range ty ~start ~stop step name is the table of one column name of ty, the values offset (k * step) start below stop, in order, for k from 0, so that adjacent ranges meet at stop. A day past a month's end becomes that month's last day, and each step counts from start, so 2024-01-31 by one month gives 2024-02-29, then 2024-03-31. It is empty when start >= stop.

It is a table where the rest of Kit is a plan, since the kth value scales a step by k, which no verb computes.

Raises Invalid_argument if step is not positive, if Expr.Temporal.offset does not move ty by step, or if the range holds 2{^ 62} steps or more.

categorize cs q is q with each text column of cs cast, in place, to the categorical type whose dictionary is the column's distinct non-null values in byte order, the order of Order.asc on text, which does not depend on the order or the batches of q's rows. It runs one query per column, after building them all:

let words c =
  let x = Col.string c in
  Query.select Expr.[ c := cast Type.string x ] q
  |> Query.filter Expr.(not (is_null x))
  |> distinct
  |> Query.sort [ Order.asc c ]
  |> Query.values x
in
Query.derive
  (List.map
     (fun (c, ws) ->
       Expr.(c := cast (Type.categorical ws) (Col.string c)))
     dictionaries)
  (* the [(c, words c)] of each [c] *)
  q

cs empty is Ok q. A column of q that is missing or is not text is a problem of the select. A categorical column takes the dictionary of the values it holds, in byte order, since it is read as text.

Error e if a run fails, or if a column holds more than 2³¹ - 1 distinct strings, the most a dictionary holds.

Raises Invalid_argument if cs names a column twice.