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 {
21        precision: 38,
22        scale: 10,
23    },
24];
25
26const TEXT: &[ColumnType] = &[ColumnType::String];
27
28const FLOAT64_ONLY: &[ColumnType] = &[ColumnType::Float64];
29
30const INT64_ONLY: &[ColumnType] = &[ColumnType::Int64];
31
32const VECTOR_ONLY: &[ColumnType] = &[ColumnType::Vector(0)];
33
34const GEOMETRY_ONLY: &[ColumnType] = &[ColumnType::Geometry];
35
36const TIMESTAMP_TYPES: &[ColumnType] = &[ColumnType::Timestamp, ColumnType::Timestamptz];
37
38// ── Helper constructors ──────────────────────────────────────────────────────
39
40pub(super) const fn any(name: &'static str) -> ArgTypeSpec {
41    ArgTypeSpec {
42        name,
43        accepted: ANY,
44        variadic: false,
45    }
46}
47
48pub(super) const fn any_variadic(name: &'static str) -> ArgTypeSpec {
49    ArgTypeSpec {
50        name,
51        accepted: ANY,
52        variadic: true,
53    }
54}
55
56pub(super) const fn typed(name: &'static str, accepted: &'static [ColumnType]) -> ArgTypeSpec {
57    ArgTypeSpec {
58        name,
59        accepted,
60        variadic: false,
61    }
62}
63
64pub(super) const fn typed_variadic(
65    name: &'static str,
66    accepted: &'static [ColumnType],
67) -> ArgTypeSpec {
68    ArgTypeSpec {
69        name,
70        accepted,
71        variadic: true,
72    }
73}
74
75// ── Standard aggregates ──────────────────────────────────────────────────────
76
77/// count(*) / count(expr)
78pub static COUNT_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
79
80pub static SUM_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC)];
81
82pub static AVG_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), any("options")];
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// ── Spatial ───────────────────────────────────────────────────────────────────
153
154pub static SPATIAL_3_ARGS: &[ArgTypeSpec] = &[
155    typed("geom1", GEOMETRY_ONLY),
156    typed("geom2", GEOMETRY_ONLY),
157    typed("distance", FLOAT64_ONLY),
158];
159
160pub static SPATIAL_2_ARGS: &[ArgTypeSpec] =
161    &[typed("geom1", GEOMETRY_ONLY), typed("geom2", GEOMETRY_ONLY)];
162
163pub static ST_BUFFER_ARGS: &[ArgTypeSpec] = &[
164    typed("geom", GEOMETRY_ONLY),
165    typed("distance", FLOAT64_ONLY),
166    typed("segments", INT64_ONLY),
167];
168
169pub static ST_ENVELOPE_ARGS: &[ArgTypeSpec] = &[typed("geom", GEOMETRY_ONLY)];
170
171pub static ST_POINT_ARGS: &[ArgTypeSpec] = &[typed("x", FLOAT64_ONLY), typed("y", FLOAT64_ONLY)];
172
173pub static ST_GEOHASH_ARGS: &[ArgTypeSpec] = &[
174    typed("lng", FLOAT64_ONLY),
175    typed("lat", FLOAT64_ONLY),
176    typed("precision", INT64_ONLY),
177];
178
179pub static ST_GEOHASHDECODE_ARGS: &[ArgTypeSpec] = &[typed("geohash", TEXT)];
180
181pub static H3_LATLNGTOCELL_ARGS: &[ArgTypeSpec] = &[
182    typed("lat", FLOAT64_ONLY),
183    typed("lng", FLOAT64_ONLY),
184    typed("resolution", INT64_ONLY),
185];
186
187pub static H3_CELLTOLATLNG_ARGS: &[ArgTypeSpec] = &[typed("h3_index", TEXT)];
188
189// ── Timeseries ────────────────────────────────────────────────────────────────
190
191pub static TIME_BUCKET_ARGS: &[ArgTypeSpec] = &[any("interval"), typed("ts", TIMESTAMP_TYPES)];
192
193pub static TS_PERCENTILE_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("percentile", FLOAT64_ONLY)];
194
195pub static TS_STDDEV_ARGS: &[ArgTypeSpec] = &[any("expr")];
196
197pub static TS_CORRELATE_ARGS: &[ArgTypeSpec] = &[any("col1"), any("col2")];
198
199pub static TS_WINDOW_1_ARGS: &[ArgTypeSpec] = &[any("expr")];
200
201pub static TS_MOVING_AVG_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("window", INT64_ONLY)];
202
203pub static TS_EMA_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("alpha", FLOAT64_ONLY)];
204
205pub static TS_LAG_LEAD_ARGS: &[ArgTypeSpec] =
206    &[any("expr"), typed("offset", INT64_ONLY), any("default")];
207
208// ── Approximate aggregates ────────────────────────────────────────────────────
209
210pub static APPROX_COUNT_DISTINCT_ARGS: &[ArgTypeSpec] = &[any("expr")];
211
212pub static APPROX_PERCENTILE_ARGS: &[ArgTypeSpec] =
213    &[any("expr"), typed("percentile", FLOAT64_ONLY)];
214
215pub static APPROX_TOPK_ARGS: &[ArgTypeSpec] = &[any("expr"), typed("k", INT64_ONLY)];
216
217pub static APPROX_COUNT_ARGS: &[ArgTypeSpec] = &[any("expr")];
218
219// ── Grouping set helpers ──────────────────────────────────────────────────────
220
221/// `GROUPING(col [, col2, ...])` — accepts 1 or more column references.
222pub static GROUPING_ARGS: &[ArgTypeSpec] = &[any_variadic("col")];
223
224// ── Document helpers ──────────────────────────────────────────────────────────
225
226pub static DOC_GET_ARGS: &[ArgTypeSpec] = &[any("doc"), typed("path", TEXT), any("default")];
227
228pub static DOC_EXISTS_ARGS: &[ArgTypeSpec] = &[any("doc"), typed("path", TEXT)];
229
230pub static DOC_ARRAY_CONTAINS_ARGS: &[ArgTypeSpec] =
231    &[any("doc"), typed("path", TEXT), any("value")];
232
233pub static NAV_ARGS: &[ArgTypeSpec] = &[any("doc"), typed("path", TEXT)];
234
235// ── Utility ───────────────────────────────────────────────────────────────────
236
237pub static NDB_CHUNK_TEXT_ARGS: &[ArgTypeSpec] = &[
238    typed("text", TEXT),
239    typed("chunk_size", INT64_ONLY),
240    typed("overlap", INT64_ONLY),
241];
242
243// ── Standard scalar ───────────────────────────────────────────────────────────
244
245pub static COALESCE_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
246
247pub static NULLIF_ARGS: &[ArgTypeSpec] = &[any("expr1"), any("expr2")];
248
249pub static MATH_1_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC)];
250
251pub static ROUND_ARGS: &[ArgTypeSpec] = &[typed("expr", NUMERIC), typed("scale", INT64_ONLY)];
252
253pub static STRING_1_ARGS: &[ArgTypeSpec] = &[typed("expr", TEXT)];
254
255pub static LENGTH_ARGS: &[ArgTypeSpec] = &[typed("expr", TEXT)];
256
257pub static SUBSTRING_ARGS: &[ArgTypeSpec] = &[
258    typed("expr", TEXT),
259    typed("start", INT64_ONLY),
260    typed("length", INT64_ONLY),
261];
262
263pub static CONCAT_ARGS: &[ArgTypeSpec] = &[typed_variadic("expr", TEXT)];
264
265pub static REPLACE_ARGS: &[ArgTypeSpec] =
266    &[typed("expr", TEXT), typed("from", TEXT), typed("to", TEXT)];
267
268pub static MAKE_ARRAY_ARGS: &[ArgTypeSpec] = &[any_variadic("expr")];
269
270// ── PostgreSQL JSON operators ─────────────────────────────────────────────────
271
272pub static PG_JSON_2_ARGS: &[ArgTypeSpec] = &[any("json_col"), any("key")];
273
274pub static PG_JSON_BOOL_2_ARGS: &[ArgTypeSpec] = &[any("json_col"), any("operand")];
275
276// ── SQL/JSON standard functions ───────────────────────────────────────────────
277
278pub static JSON_VALUE_ARGS: &[ArgTypeSpec] = &[any("json_col"), typed("path", TEXT)];
279
280pub static JSON_QUERY_ARGS: &[ArgTypeSpec] = &[any("json_col"), typed("path", TEXT)];
281
282pub static JSON_EXISTS_ARGS: &[ArgTypeSpec] = &[any("json_col"), typed("path", TEXT)];
283
284// ── PostgreSQL FTS operators ──────────────────────────────────────────────────
285
286pub static PG_FTS_MATCH_ARGS: &[ArgTypeSpec] = &[any("tsvector"), any("tsquery")];
287
288pub static PG_TO_TSVECTOR_ARGS: &[ArgTypeSpec] = &[typed("config", TEXT), typed("document", TEXT)];
289
290pub static PG_TO_TSQUERY_ARGS: &[ArgTypeSpec] = &[typed("config", TEXT), typed("query", TEXT)];
291
292pub static PG_TS_RANK_ARGS: &[ArgTypeSpec] = &[
293    any("tsvector"),
294    any("tsquery"),
295    typed("weights", FLOAT64_ONLY),
296    typed("normalization", INT64_ONLY),
297];
298
299pub static PG_TS_HEADLINE_ARGS: &[ArgTypeSpec] = &[
300    typed("config", TEXT),
301    any("document"),
302    any("tsquery"),
303    typed("options", TEXT),
304];
305
306// ── Array engine ──────────────────────────────────────────────────────────────
307
308pub static ARRAY_SLICE_ARGS: &[ArgTypeSpec] = &[
309    typed("name", TEXT),
310    any("slice_obj"),
311    any("attrs"),
312    typed("limit", INT64_ONLY),
313];
314
315pub static ARRAY_PROJECT_ARGS: &[ArgTypeSpec] = &[typed("name", TEXT), any("attrs")];
316
317pub static ARRAY_AGG_ARGS: &[ArgTypeSpec] = &[
318    typed("name", TEXT),
319    typed("attr", TEXT),
320    typed("reducer", TEXT),
321    typed("group_by_dim", INT64_ONLY),
322];
323
324pub static ARRAY_ELEMENTWISE_ARGS: &[ArgTypeSpec] = &[
325    typed("left", TEXT),
326    typed("right", TEXT),
327    typed("op", TEXT),
328    typed("attr", TEXT),
329];
330
331pub static ARRAY_MAINT_ARGS: &[ArgTypeSpec] = &[typed("name", TEXT)];