Skip to main content

nodedb_sql/functions/
arg_types.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Static argument type specifications for all built-in functions.
4//!
5//! Each constant defines the accepted argument types per position.
6//! An empty `accepted` slice means any type is accepted (wildcard).
7//! `variadic: true` on the last entry means the argument repeats.
8
9use nodedb_types::columnar::ColumnType;
10
11use super::registry::ArgTypeSpec;
12
13// ── Shorthands ──────────────────────────────────────────────────────────────
14
15const ANY: &[ColumnType] = &[];
16
17const NUMERIC: &[ColumnType] = &[
18    ColumnType::Int64,
19    ColumnType::Float64,
20    ColumnType::Decimal(None),
21];
22
23const TEXT: &[ColumnType] = &[ColumnType::String];
24
25const FLOAT64_ONLY: &[ColumnType] = &[ColumnType::Float64];
26
27const INT64_ONLY: &[ColumnType] = &[ColumnType::Int64];
28
29const VECTOR_ONLY: &[ColumnType] = &[ColumnType::Vector(0)];
30
31const TIMESTAMP_TYPES: &[ColumnType] = &[ColumnType::Timestamp, ColumnType::Timestamptz];
32
33const ARRAY_ONLY: &[ColumnType] = &[ColumnType::Array];
34
35// ── Helper constructors ──────────────────────────────────────────────────────
36
37pub(super) const fn any(name: &'static str) -> ArgTypeSpec {
38    ArgTypeSpec {
39        name,
40        accepted: ANY,
41        variadic: false,
42    }
43}
44
45pub(super) const fn any_variadic(name: &'static str) -> ArgTypeSpec {
46    ArgTypeSpec {
47        name,
48        accepted: ANY,
49        variadic: true,
50    }
51}
52
53pub(super) const fn typed(name: &'static str, accepted: &'static [ColumnType]) -> ArgTypeSpec {
54    ArgTypeSpec {
55        name,
56        accepted,
57        variadic: false,
58    }
59}
60
61pub(super) const fn typed_variadic(
62    name: &'static str,
63    accepted: &'static [ColumnType],
64) -> ArgTypeSpec {
65    ArgTypeSpec {
66        name,
67        accepted,
68        variadic: true,
69    }
70}
71
72// ── Standard aggregates ──────────────────────────────────────────────────────
73
74/// count(*) / count(expr)
75pub static COUNT_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
76
77pub static SUM_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC)];
78
79pub static AVG_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC)];
80
81/// `stddev` / `variance` and their `_samp` / `_pop` forms.
82pub static DISPERSION_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC)];
83
84pub static MIN_ARGS: &[ArgTypeSpec] = &[any("expr")];
85
86pub static MAX_ARGS: &[ArgTypeSpec] = &[any("expr")];
87
88// ── Standard window ──────────────────────────────────────────────────────────
89
90pub static NO_ARGS: &[ArgTypeSpec] = &[];
91
92pub static LAG_LEAD_ARGS: &[ArgTypeSpec] =
93    &[any("expr"), typed("offset", INT64_ONLY), any("default")];
94
95pub static FIRST_LAST_VALUE_ARGS: &[ArgTypeSpec] = &[any("expr")];
96
97pub static NTH_VALUE_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("n", INT64_ONLY)];
98
99pub static NTILE_ARGS: &[ArgTypeSpec] = &[typed("buckets", INT64_ONLY)];
100
101// ── Vector search ────────────────────────────────────────────────────────────
102
103pub static VECTOR_DISTANCE_ARGS: &[ArgTypeSpec] = &[
104    typed("column", VECTOR_ONLY),
105    typed("query", VECTOR_ONLY),
106    any("metric"),
107];
108
109pub static MULTI_VECTOR_SEARCH_ARGS: &[ArgTypeSpec] = &[any("query"), any("options")];
110
111pub static MULTI_VECTOR_SCORE_ARGS: &[ArgTypeSpec] =
112    &[any("col1"), any("col2"), typed("query", VECTOR_ONLY)];
113
114pub static SPARSE_SCORE_ARGS: &[ArgTypeSpec] =
115    &[any("col"), any("query"), typed("boost", FLOAT64_ONLY)];
116
117// ── Text search ───────────────────────────────────────────────────────────────
118
119pub static BM25_SCORE_ARGS: &[ArgTypeSpec] = &[any("column"), typed("query", TEXT)];
120
121pub static SEARCH_SCORE_ARGS: &[ArgTypeSpec] = &[any("column"), typed("query", TEXT)];
122
123pub static TEXT_MATCH_ARGS: &[ArgTypeSpec] = &[any("column"), typed("query", TEXT)];
124
125// ── Hybrid search ─────────────────────────────────────────────────────────────
126
127/// Two-source `rrf_score(rank1, rank2, k1?, k2?)` — vector + text.
128/// Accepts 2–4 arguments; the planner validates arity at plan time.
129pub static RRF_SCORE_ARGS: &[ArgTypeSpec] = &[any("rank1"), any("rank2"), any("k1"), any("k2")];
130
131/// Three-source `rrf_score(rank1, rank2, rank3, k1?, k2?, k3?)` — vector + text + graph.
132/// Shares the same function name; arity dispatch happens in the planner.
133pub static RRF_SCORE_TRIPLE_ARGS: &[ArgTypeSpec] = &[
134    any("rank1"),
135    any("rank2"),
136    any("rank3"),
137    any("k1"),
138    any("k2"),
139    any("k3"),
140];
141
142/// `graph_score(node_id_col, seed_id, depth => N, label => 'edge_label')`.
143///
144/// This is a planner-intercepted marker function — it is never evaluated
145/// per-row. The hybrid planner recognises it as the graph-distance source
146/// in a three-source `rrf_score(...)` call and lowers it to a graph BFS
147/// spec attached to the physical plan. Arguments beyond the first two are
148/// named (`depth =>`, `label =>`).
149pub static GRAPH_SCORE_ARGS: &[ArgTypeSpec] =
150    &[any("node_id_col"), any("seed_id"), any_variadic("options")];
151
152// ── Timeseries ────────────────────────────────────────────────────────────────
153
154pub static TIME_BUCKET_ARGS: &[ArgTypeSpec] = &[any("interval"), typed("ts", TIMESTAMP_TYPES)];
155
156pub static TS_PERCENTILE_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("percentile", FLOAT64_ONLY)];
157
158pub static TS_STDDEV_ARGS: &[ArgTypeSpec] = &[any("expr")];
159
160pub static TS_CORRELATE_ARGS: &[ArgTypeSpec] = &[any("col1"), any("col2")];
161
162pub static TS_WINDOW_1_ARGS: &[ArgTypeSpec] = &[any("expr")];
163
164pub static TS_MOVING_AVG_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("window", INT64_ONLY)];
165
166pub static TS_EMA_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("alpha", FLOAT64_ONLY)];
167
168pub static TS_LAG_LEAD_ARGS: &[ArgTypeSpec] =
169    &[any("expr"), typed("offset", INT64_ONLY), any("default")];
170
171// ── Approximate aggregates ────────────────────────────────────────────────────
172
173pub static APPROX_COUNT_DISTINCT_ARGS: &[ArgTypeSpec] = &[any("expr")];
174
175pub static APPROX_PERCENTILE_ARGS: &[ArgTypeSpec] =
176    &[any("expr"), typed("percentile", FLOAT64_ONLY)];
177
178pub static APPROX_TOPK_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("k", INT64_ONLY)];
179
180pub static APPROX_COUNT_ARGS: &[ArgTypeSpec] = &[any("expr")];
181
182// ── Grouping set helpers ──────────────────────────────────────────────────────
183
184/// `GROUPING(col [, col2, ...])` — accepts 1 or more column references.
185pub static GROUPING_ARGS: &[ArgTypeSpec] = &[any_variadic("col")];
186
187// ── Document helpers ──────────────────────────────────────────────────────────
188
189pub static DOC_GET_ARGS: &[ArgTypeSpec] = &[any("doc"), typed("path", TEXT), any("default")];
190
191pub static DOC_EXISTS_ARGS: &[ArgTypeSpec] = &[any("doc"), typed("path", TEXT)];
192
193pub static DOC_ARRAY_CONTAINS_ARGS: &[ArgTypeSpec] =
194    &[any("doc"), typed("path", TEXT), any("value")];
195
196pub static NAV_ARGS: &[ArgTypeSpec] = &[any("doc"), typed("path", TEXT)];
197
198// ── Utility ───────────────────────────────────────────────────────────────────
199
200pub static NDB_CHUNK_TEXT_ARGS: &[ArgTypeSpec] = &[
201    typed("text", TEXT),
202    typed("chunk_size", INT64_ONLY),
203    typed("overlap", INT64_ONLY),
204];
205
206// ── Standard scalar ───────────────────────────────────────────────────────────
207
208pub static COALESCE_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
209
210pub static NULLIF_ARGS: &[ArgTypeSpec] = &[any("expr1"), any("expr2")];
211
212pub static MATH_1_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC)];
213
214/// `mod(a, b)` / `power(base, exp)` / `pow(base, exp)` — two required
215/// numeric arguments.
216pub static MATH_2_ARGS: &[ArgTypeSpec] = &[typed("left", NUMERIC), typed("right", NUMERIC)];
217
218pub static ROUND_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC), typed("scale", INT64_ONLY)];
219
220/// `greatest(expr, ...)` / `least(expr, ...)` — one or more arguments of any
221/// (comparable) type.
222pub static GREATEST_LEAST_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
223
224/// `typeof(expr)` / `type_of(expr)` — accepts any value, including NULL.
225pub static TYPEOF_ARGS: &[ArgTypeSpec] = &[any("expr")];
226
227pub static STRING_1_ARGS: &[ArgTypeSpec] = &[typed("expr", TEXT)];
228
229pub static LENGTH_ARGS: &[ArgTypeSpec] = &[typed("expr", TEXT)];
230
231pub static SUBSTRING_ARGS: &[ArgTypeSpec] = &[
232    typed("expr", TEXT),
233    typed("start", INT64_ONLY),
234    typed("length", INT64_ONLY),
235];
236
237pub static CONCAT_ARGS: &[ArgTypeSpec] = &[typed_variadic("expr", TEXT)];
238
239pub static REPLACE_ARGS: &[ArgTypeSpec] =
240    &[typed("expr", TEXT), typed("from", TEXT), typed("to", TEXT)];
241
242pub static MAKE_ARRAY_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
243
244pub static PG_CATALOG_2_ARGS: &[ArgTypeSpec] = &[any("value"), any("context")];
245
246// ── PostgreSQL JSON operators ─────────────────────────────────────────────────
247
248pub static PG_JSON_2_ARGS: &[ArgTypeSpec] = &[any("json_col"), any("key")];
249
250pub static PG_JSON_BOOL_2_ARGS: &[ArgTypeSpec] = &[any("json_col"), any("operand")];
251
252// ── SQL/JSON standard functions ───────────────────────────────────────────────
253
254pub static JSON_VALUE_ARGS: &[ArgTypeSpec] = &[any("json_col"), typed("path", TEXT)];
255
256pub static JSON_QUERY_ARGS: &[ArgTypeSpec] = &[any("json_col"), typed("path", TEXT)];
257
258pub static JSON_EXISTS_ARGS: &[ArgTypeSpec] = &[any("json_col"), typed("path", TEXT)];
259
260// ── PostgreSQL FTS operators ──────────────────────────────────────────────────
261
262pub static PG_FTS_MATCH_ARGS: &[ArgTypeSpec] = &[any("tsvector"), any("tsquery")];
263
264pub static PG_TO_TSVECTOR_ARGS: &[ArgTypeSpec] = &[typed("config", TEXT), typed("document", TEXT)];
265
266pub static PG_TO_TSQUERY_ARGS: &[ArgTypeSpec] = &[typed("config", TEXT), typed("query", TEXT)];
267
268pub static PG_TS_RANK_ARGS: &[ArgTypeSpec] = &[
269    any("tsvector"),
270    any("tsquery"),
271    typed("weights", FLOAT64_ONLY),
272    typed("normalization", INT64_ONLY),
273];
274
275pub static PG_TS_HEADLINE_ARGS: &[ArgTypeSpec] = &[
276    typed("config", TEXT),
277    any("document"),
278    any("tsquery"),
279    typed("options", TEXT),
280];
281
282// ── Array engine ──────────────────────────────────────────────────────────────
283
284pub static ARRAY_SLICE_ARGS: &[ArgTypeSpec] = &[
285    typed("name", TEXT),
286    any("slice_obj"),
287    any("attrs"),
288    typed("limit", INT64_ONLY),
289];
290
291pub static ARRAY_PROJECT_ARGS: &[ArgTypeSpec] = &[typed("name", TEXT), any("attrs")];
292
293pub static ARRAY_AGG_ARGS: &[ArgTypeSpec] = &[
294    typed("name", TEXT),
295    typed("attr", TEXT),
296    typed("reducer", TEXT),
297    typed("group_by_dim", INT64_ONLY),
298];
299
300pub static ARRAY_ELEMENTWISE_ARGS: &[ArgTypeSpec] = &[
301    typed("left", TEXT),
302    typed("right", TEXT),
303    typed("op", TEXT),
304    typed("attr", TEXT),
305];
306
307pub static ARRAY_MAINT_ARGS: &[ArgTypeSpec] = &[typed("name", TEXT)];
308
309// ── ID generation / detection (nodedb-query's `functions/id.rs`) ──────────────
310
311/// `is_uuid(value)` / `is_ulid(value)` / `is_cuid2(value)` / `is_nanoid(value)`
312/// / `id_type(value)` / `uuid_version(value)` / `ulid_timestamp(value)` — a
313/// single text value to classify. Shared shape across all detector functions.
314pub static ID_CHECK_ARGS: &[ArgTypeSpec] = &[typed("value", TEXT)];
315
316/// `nanoid(length?)` — optional custom length, defaults to the standard
317/// 21-character id when omitted.
318pub static NANOID_ARGS: &[ArgTypeSpec] = &[typed("length", INT64_ONLY)];
319
320// ── Array element functions (nodedb-query's `functions/array.rs`) ─────────────
321//
322// Distinct from the "Array engine" table-valued functions above
323// (`array_slice`/`array_project`/...): these operate on an in-row `Array`
324// value, mirroring the match arms in `nodedb_query::functions::array::try_eval`.
325
326/// `array_length(arr)` / `cardinality(arr)` / `array_distinct(arr)` /
327/// `array_reverse(arr)` — a single array argument.
328pub static ARRAY_1_ARGS: &[ArgTypeSpec] = &[typed("arr", ARRAY_ONLY)];
329
330/// `array_append(arr, value)` / `array_remove(arr, value)` /
331/// `array_contains(arr, value)` / `array_position(arr, value)` — an array
332/// plus an element of any type.
333pub static ARRAY_ELEM_ARGS: &[ArgTypeSpec] = &[typed("arr", ARRAY_ONLY), any("value")];
334
335/// `array_prepend(value, arr)` — value-first argument order, matching
336/// `nodedb_query::functions::array::try_eval`'s `"array_prepend"` arm.
337pub static ARRAY_PREPEND_ARGS: &[ArgTypeSpec] = &[any("value"), typed("arr", ARRAY_ONLY)];
338
339/// `array_concat(arr1, arr2)` / `array_cat(arr1, arr2)`.
340pub static ARRAY_CAT_ARGS: &[ArgTypeSpec] = &[typed("arr1", ARRAY_ONLY), typed("arr2", ARRAY_ONLY)];
341
342// ── DateTime / duration (nodedb-query's `functions/datetime.rs`) ──────────────
343
344/// `datetime(value)` / `to_datetime(value)` / `unix_secs(value)` /
345/// `epoch_secs(value)` / `unix_millis(value)` / `epoch_millis(value)` /
346/// `duration(value)` / `to_duration(value)` / `decimal(value)` /
347/// `to_decimal(value)` — a single value whose accepted shape (string,
348/// integer micros, or an existing timestamp/duration) varies by function, so
349/// the position is left untyped (`any`) rather than over-constrained.
350pub static DATETIME_1_ARGS: &[ArgTypeSpec] = &[any("value")];
351
352/// `extract(part, ts)` / `date_part(part, ts)` / `date_trunc(part, ts)` /
353/// `datetrunc(part, ts)` — a text field-name selector plus the
354/// timestamp-like value to read it from.
355pub static DATE_PART_ARGS: &[ArgTypeSpec] = &[typed("part", TEXT), any("ts")];
356
357/// `date_add(ts, duration)` / `datetime_add(ts, duration)` /
358/// `date_sub(ts, duration)` / `datetime_sub(ts, duration)`.
359pub static DATE_ARITH_ARGS: &[ArgTypeSpec] = &[any("ts"), typed("duration", TEXT)];
360
361/// `date_diff(ts1, ts2)` / `datediff(ts1, ts2)`.
362pub static DATE_DIFF_ARGS: &[ArgTypeSpec] = &[any("ts1"), any("ts2")];
363
364// ── Legacy JSON manipulation (nodedb-query's `functions/json/legacy.rs`) ──────
365//
366// Distinct from the PostgreSQL JSON operators and SQL/JSON standard
367// functions above: these are the pre-existing `json_*` document-helper
368// functions.
369
370/// `json_extract(json, path)` / `json_get(json, path)`.
371pub static JSON_EXTRACT_ARGS: &[ArgTypeSpec] = &[any("json"), typed("path", TEXT)];
372
373/// `json_set(json, key, value)`.
374pub static JSON_SET_ARGS: &[ArgTypeSpec] = &[any("json"), typed("key", TEXT), any("value")];
375
376/// `json_remove(json, key)`.
377pub static JSON_REMOVE_ARGS: &[ArgTypeSpec] = &[any("json"), typed("key", TEXT)];
378
379/// `json_keys(json)` / `json_values(json)` / `json_length(json)` /
380/// `json_len(json)` / `json_type(json)` — a single JSON-shaped value.
381pub static JSON_1_ARGS: &[ArgTypeSpec] = &[any("json")];
382
383/// `json_array(value, ...)` — zero or more values of any type.
384pub static JSON_ARRAY_ARGS: &[ArgTypeSpec] = &[any_variadic("value")];
385
386/// `json_object(key, value, ...)` — a flat, variadic key/value list.
387pub static JSON_OBJECT_ARGS: &[ArgTypeSpec] = &[any_variadic("kv")];
388
389/// `json_contains(container, needle)`.
390pub static JSON_CONTAINS_ARGS: &[ArgTypeSpec] = &[any("container"), any("needle")];
391
392/// `json_merge(base, overlay)` / `json_patch(base, overlay)`.
393pub static JSON_MERGE_ARGS: &[ArgTypeSpec] = &[any("base"), any("overlay")];
394
395/// `to_jsonb(value)` — any value, or `*` for the whole row.
396pub static TO_JSONB_ARGS: &[ArgTypeSpec] = &[any("value")];
397
398// ── Sequence accessors ───────────────────────────────────────────────────────
399
400/// `nextval(name)` / `currval(name)` — one sequence name.
401pub static SEQUENCE_NAME_ARGS: &[ArgTypeSpec] = &[typed("sequence", TEXT)];
402
403/// `setval(name, value)` — a sequence name and the value it must take.
404pub static SETVAL_ARGS: &[ArgTypeSpec] = &[typed("sequence", TEXT), typed("value", INT64_ONLY)];