1use nodedb_types::columnar::ColumnType;
10
11use super::registry::ArgTypeSpec;
12
13const 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
38pub(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
75pub 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
88pub 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
101pub 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
117pub 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
125pub static RRF_SCORE_ARGS: &[ArgTypeSpec] = &[any("rank1"), any("rank2"), any("k1"), any("k2")];
130
131pub 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
142pub static GRAPH_SCORE_ARGS: &[ArgTypeSpec] =
150 &[any("node_id_col"), any("seed_id"), any_variadic("options")];
151
152pub 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
189pub 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
208pub 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
219pub static GROUPING_ARGS: &[ArgTypeSpec] = &[any_variadic("col")];
223
224pub 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
235pub static NDB_CHUNK_TEXT_ARGS: &[ArgTypeSpec] = &[
238 typed("text", TEXT),
239 typed("chunk_size", INT64_ONLY),
240 typed("overlap", INT64_ONLY),
241];
242
243pub 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
270pub 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
276pub 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
284pub 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
306pub 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)];