@query-farm/vgi 0.6.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (206) hide show
  1. package/LICENSE +134 -0
  2. package/README.md +494 -0
  3. package/dist/__dbg.test-d.d.ts +1 -0
  4. package/dist/arguments/argument-spec.d.ts +51 -0
  5. package/dist/arguments/arguments.d.ts +13 -0
  6. package/dist/arrow/codec/branded.d.ts +83 -0
  7. package/dist/arrow/codec/registry.d.ts +29 -0
  8. package/dist/arrow/codec/repr.d.ts +23 -0
  9. package/dist/arrow/codec/type-descriptors.d.ts +109 -0
  10. package/dist/arrow/impl-arrowjs/build.d.ts +19 -0
  11. package/dist/arrow/impl-arrowjs/canonical.d.ts +15 -0
  12. package/dist/arrow/impl-arrowjs/empty.d.ts +7 -0
  13. package/dist/arrow/impl-arrowjs/filter.d.ts +11 -0
  14. package/dist/arrow/impl-arrowjs/index.d.ts +11 -0
  15. package/dist/arrow/impl-arrowjs/ipc.d.ts +24 -0
  16. package/dist/arrow/impl-arrowjs/iterate.d.ts +39 -0
  17. package/dist/arrow/impl-arrowjs/project.d.ts +10 -0
  18. package/dist/arrow/impl-arrowjs/schema.d.ts +62 -0
  19. package/dist/arrow/impl-arrowjs/statistics.d.ts +13 -0
  20. package/dist/arrow/impl-flechette/build.d.ts +14 -0
  21. package/dist/arrow/impl-flechette/canonical.d.ts +18 -0
  22. package/dist/arrow/impl-flechette/empty.d.ts +9 -0
  23. package/dist/arrow/impl-flechette/filter.d.ts +10 -0
  24. package/dist/arrow/impl-flechette/index.d.ts +11 -0
  25. package/dist/arrow/impl-flechette/ipc.d.ts +33 -0
  26. package/dist/arrow/impl-flechette/iterate.d.ts +29 -0
  27. package/dist/arrow/impl-flechette/message-meta.d.ts +10 -0
  28. package/dist/arrow/impl-flechette/normalize-type.d.ts +11 -0
  29. package/dist/arrow/impl-flechette/project.d.ts +3 -0
  30. package/dist/arrow/impl-flechette/schema.d.ts +47 -0
  31. package/dist/arrow/impl-flechette/statistics.d.ts +13 -0
  32. package/dist/arrow/index.d.ts +11 -0
  33. package/dist/arrow/predicates.d.ts +47 -0
  34. package/dist/arrow/repr.test-d.d.ts +52 -0
  35. package/dist/arrow/schema-types.d.ts +71 -0
  36. package/dist/arrow/types.d.ts +57 -0
  37. package/dist/catalog/attach-option.d.ts +33 -0
  38. package/dist/catalog/attach-options.d.ts +30 -0
  39. package/dist/catalog/composite.d.ts +33 -0
  40. package/dist/catalog/descriptors.d.ts +154 -0
  41. package/dist/catalog/interface.d.ts +239 -0
  42. package/dist/catalog/read-only.d.ts +33 -0
  43. package/dist/client/client.d.ts +189 -0
  44. package/dist/client/errors.d.ts +26 -0
  45. package/dist/client/helpers.d.ts +12 -0
  46. package/dist/client/protocol.d.ts +16 -0
  47. package/dist/client/types.d.ts +125 -0
  48. package/dist/client-entry.d.ts +11 -0
  49. package/dist/client-entry.js +48 -0
  50. package/dist/client-entry.js.map +9 -0
  51. package/dist/codec/asd.d.ts +25 -0
  52. package/dist/crypto.d.ts +36 -0
  53. package/dist/errors.d.ts +24 -0
  54. package/dist/filter-pushdown/collector.d.ts +32 -0
  55. package/dist/filter-pushdown/deserialize.d.ts +16 -0
  56. package/dist/filter-pushdown/evaluate.d.ts +51 -0
  57. package/dist/filter-pushdown/index.d.ts +4 -0
  58. package/dist/filter-pushdown/types.d.ts +84 -0
  59. package/dist/functions/aggregate.d.ts +129 -0
  60. package/dist/functions/copy-from.d.ts +81 -0
  61. package/dist/functions/copy-to.d.ts +121 -0
  62. package/dist/functions/registry.d.ts +16 -0
  63. package/dist/functions/scalar.d.ts +98 -0
  64. package/dist/functions/storage-cf-do.d.ts +47 -0
  65. package/dist/functions/storage.d.ts +86 -0
  66. package/dist/functions/table-buffering.d.ts +99 -0
  67. package/dist/functions/table-in-out.d.ts +64 -0
  68. package/dist/functions/table.d.ts +118 -0
  69. package/dist/functions/types.d.ts +157 -0
  70. package/dist/generated/vgi-client.d.ts +466 -0
  71. package/dist/generated/vgi-protocol-schemas.d.ts +115 -0
  72. package/dist/index.core.d.ts +42 -0
  73. package/dist/index.d.ts +2 -0
  74. package/dist/index.js +10888 -0
  75. package/dist/index.js.map +85 -0
  76. package/dist/metadata/resolve.d.ts +3 -0
  77. package/dist/metadata/serialize.d.ts +12 -0
  78. package/dist/metadata/types.d.ts +47 -0
  79. package/dist/protocol/dispatch.d.ts +23 -0
  80. package/dist/protocol/handlers/aggregate.d.ts +3 -0
  81. package/dist/protocol/handlers/catalog/admin.d.ts +3 -0
  82. package/dist/protocol/handlers/catalog/index.d.ts +3 -0
  83. package/dist/protocol/handlers/catalog/index_methods.d.ts +3 -0
  84. package/dist/protocol/handlers/catalog/macro.d.ts +3 -0
  85. package/dist/protocol/handlers/catalog/shared.d.ts +70 -0
  86. package/dist/protocol/handlers/catalog/table.d.ts +3 -0
  87. package/dist/protocol/handlers/catalog/view.d.ts +3 -0
  88. package/dist/protocol/handlers/function.d.ts +8 -0
  89. package/dist/protocol/handlers/shared.d.ts +31 -0
  90. package/dist/protocol/handlers/table-buffering.d.ts +3 -0
  91. package/dist/protocol/serialize.d.ts +4 -0
  92. package/dist/protocol/serializers/arguments.d.ts +3 -0
  93. package/dist/protocol/serializers/bind.d.ts +9 -0
  94. package/dist/protocol/serializers/cardinality.d.ts +7 -0
  95. package/dist/protocol/serializers/init.d.ts +9 -0
  96. package/dist/protocol/serializers/shared.d.ts +12 -0
  97. package/dist/protocol/state-serializer.d.ts +23 -0
  98. package/dist/protocol/types.d.ts +129 -0
  99. package/dist/secrets/helpers.d.ts +15 -0
  100. package/dist/types.d.ts +52 -0
  101. package/dist/util/arrow/index.d.ts +1 -0
  102. package/dist/util/bytes.d.ts +2 -0
  103. package/dist/util/statistics.d.ts +7 -0
  104. package/dist/worker-cf-entry.d.ts +30 -0
  105. package/dist/worker-cf-entry.js +8767 -0
  106. package/dist/worker-cf-entry.js.map +74 -0
  107. package/dist/worker.d.ts +36 -0
  108. package/package.json +96 -0
  109. package/src/arguments/argument-spec.ts +200 -0
  110. package/src/arguments/arguments.ts +81 -0
  111. package/src/arrow/codec/branded.ts +206 -0
  112. package/src/arrow/codec/registry.ts +674 -0
  113. package/src/arrow/codec/repr.ts +181 -0
  114. package/src/arrow/codec/type-descriptors.ts +137 -0
  115. package/src/arrow/impl-arrowjs/build.ts +88 -0
  116. package/src/arrow/impl-arrowjs/canonical.ts +509 -0
  117. package/src/arrow/impl-arrowjs/empty.ts +72 -0
  118. package/src/arrow/impl-arrowjs/filter.ts +53 -0
  119. package/src/arrow/impl-arrowjs/index.ts +54 -0
  120. package/src/arrow/impl-arrowjs/ipc.ts +67 -0
  121. package/src/arrow/impl-arrowjs/iterate.ts +159 -0
  122. package/src/arrow/impl-arrowjs/project.ts +73 -0
  123. package/src/arrow/impl-arrowjs/schema.ts +256 -0
  124. package/src/arrow/impl-arrowjs/statistics.ts +185 -0
  125. package/src/arrow/impl-flechette/build.ts +65 -0
  126. package/src/arrow/impl-flechette/canonical.ts +289 -0
  127. package/src/arrow/impl-flechette/empty.ts +20 -0
  128. package/src/arrow/impl-flechette/filter.ts +42 -0
  129. package/src/arrow/impl-flechette/index.ts +54 -0
  130. package/src/arrow/impl-flechette/ipc.ts +110 -0
  131. package/src/arrow/impl-flechette/iterate.ts +121 -0
  132. package/src/arrow/impl-flechette/message-meta.ts +175 -0
  133. package/src/arrow/impl-flechette/normalize-type.ts +99 -0
  134. package/src/arrow/impl-flechette/project.ts +64 -0
  135. package/src/arrow/impl-flechette/schema.ts +202 -0
  136. package/src/arrow/impl-flechette/statistics.ts +198 -0
  137. package/src/arrow/index.ts +129 -0
  138. package/src/arrow/predicates.ts +66 -0
  139. package/src/arrow/repr.test-d.ts +189 -0
  140. package/src/arrow/schema-types.ts +213 -0
  141. package/src/arrow/types.ts +74 -0
  142. package/src/catalog/attach-option.ts +107 -0
  143. package/src/catalog/attach-options.ts +79 -0
  144. package/src/catalog/composite.ts +180 -0
  145. package/src/catalog/descriptors.ts +164 -0
  146. package/src/catalog/interface.ts +721 -0
  147. package/src/catalog/read-only.ts +909 -0
  148. package/src/client/client.ts +1348 -0
  149. package/src/client/errors.ts +69 -0
  150. package/src/client/helpers.ts +56 -0
  151. package/src/client/protocol.ts +33 -0
  152. package/src/client/types.ts +140 -0
  153. package/src/client-entry.ts +112 -0
  154. package/src/codec/asd.ts +147 -0
  155. package/src/crypto.ts +143 -0
  156. package/src/errors.ts +61 -0
  157. package/src/filter-pushdown/collector.ts +147 -0
  158. package/src/filter-pushdown/deserialize.ts +253 -0
  159. package/src/filter-pushdown/evaluate.ts +591 -0
  160. package/src/filter-pushdown/index.ts +28 -0
  161. package/src/filter-pushdown/types.ts +91 -0
  162. package/src/functions/aggregate.ts +384 -0
  163. package/src/functions/copy-from.ts +282 -0
  164. package/src/functions/copy-to.ts +319 -0
  165. package/src/functions/registry.ts +258 -0
  166. package/src/functions/scalar.ts +356 -0
  167. package/src/functions/storage-cf-do.ts +303 -0
  168. package/src/functions/storage.ts +521 -0
  169. package/src/functions/table-buffering.ts +562 -0
  170. package/src/functions/table-in-out.ts +373 -0
  171. package/src/functions/table.ts +521 -0
  172. package/src/functions/types.ts +192 -0
  173. package/src/generated/vgi-client.ts +560 -0
  174. package/src/generated/vgi-protocol-schemas.ts +864 -0
  175. package/src/index.core.ts +267 -0
  176. package/src/index.ts +11 -0
  177. package/src/metadata/resolve.ts +109 -0
  178. package/src/metadata/serialize.ts +242 -0
  179. package/src/metadata/types.ts +60 -0
  180. package/src/protocol/dispatch.ts +52 -0
  181. package/src/protocol/handlers/aggregate.ts +355 -0
  182. package/src/protocol/handlers/catalog/admin.ts +368 -0
  183. package/src/protocol/handlers/catalog/index.ts +25 -0
  184. package/src/protocol/handlers/catalog/index_methods.ts +65 -0
  185. package/src/protocol/handlers/catalog/macro.ts +114 -0
  186. package/src/protocol/handlers/catalog/shared.ts +291 -0
  187. package/src/protocol/handlers/catalog/table.ts +401 -0
  188. package/src/protocol/handlers/catalog/view.ts +119 -0
  189. package/src/protocol/handlers/function.ts +302 -0
  190. package/src/protocol/handlers/shared.ts +80 -0
  191. package/src/protocol/handlers/table-buffering.ts +207 -0
  192. package/src/protocol/serialize.ts +23 -0
  193. package/src/protocol/serializers/arguments.ts +160 -0
  194. package/src/protocol/serializers/bind.ts +212 -0
  195. package/src/protocol/serializers/cardinality.ts +45 -0
  196. package/src/protocol/serializers/init.ts +247 -0
  197. package/src/protocol/serializers/shared.ts +46 -0
  198. package/src/protocol/state-serializer.ts +181 -0
  199. package/src/protocol/types.ts +153 -0
  200. package/src/secrets/helpers.ts +70 -0
  201. package/src/types.ts +71 -0
  202. package/src/util/arrow/index.ts +23 -0
  203. package/src/util/bytes.ts +14 -0
  204. package/src/util/statistics.ts +40 -0
  205. package/src/worker-cf-entry.ts +172 -0
  206. package/src/worker.ts +212 -0
@@ -0,0 +1,384 @@
1
+ // Copyright 2025, 2026 Query Farm LLC - https://query.farm
2
+ // Aggregate function framework — mirrors vgi-python's vgi/aggregate_function.py
3
+ // protocol at the wire level. Per-execution state lives in an in-memory store
4
+ // keyed by (execution_id, group_id); fine for subprocess workers and for
5
+ // single-instance HTTP. Multi-host HTTP deployments would need a pluggable
6
+ // state backend (same tradeoff Python's FunctionStorage wraps).
7
+ //
8
+ // DuckDB drives the lifecycle via five unary RPCs:
9
+ //
10
+ // aggregate_bind → resolve output schema + allocate execution_id
11
+ // aggregate_update → fold rows into per-group state
12
+ // aggregate_combine → merge parallel-build states
13
+ // aggregate_finalize → emit one result row per group_id
14
+ // aggregate_destructor → free states (best-effort; see Python worker.py)
15
+
16
+ import { type VgiSchema, schema, type VgiField, field, type VgiBatch, type VgiDataType, nullType, isNull } from "../arrow/index.js";
17
+ import type { Arguments } from "../arguments/arguments.js";
18
+ import type { ArgumentSpec } from "../arguments/argument-spec.js";
19
+ import type {
20
+ FunctionMeta,
21
+ VgiFunction,
22
+ } from "./types.js";
23
+ import { DEFAULT_MAX_WORKERS } from "../types.js";
24
+ import type { FunctionExample } from "./types.js";
25
+ // Aggregate state persistence uses node:fs/path/os to spill per-group state to
26
+ // disk under VGI_AGG_STATE_DIR. Workers (Cloudflare workerd / browsers)
27
+ // don't define aggregate functions in the first place, so we resolve the
28
+ // node:* modules through indirect string variables that esbuild can't trace.
29
+ // The bundle stays clean even without nodejs_compat polyfills.
30
+ const _NODE_FS_MOD = "node:fs";
31
+ const _NODE_PATH_MOD = "node:path";
32
+ const _NODE_OS_MOD = "node:os";
33
+ function _aggReq(): any {
34
+ const req: any = (import.meta as any).require ?? (globalThis as any).require ?? null;
35
+ if (!req) {
36
+ throw new Error(
37
+ "Aggregate state persistence requires Node.js or Bun (node:fs/path/os).",
38
+ );
39
+ }
40
+ return req;
41
+ }
42
+ const fs: any = new Proxy({} as any, {
43
+ get: (_t, prop) => (_aggReq()(_NODE_FS_MOD) as any)[prop],
44
+ });
45
+ const path: any = new Proxy({} as any, {
46
+ get: (_t, prop) => (_aggReq()(_NODE_PATH_MOD) as any)[prop],
47
+ });
48
+ const os: any = new Proxy({} as any, {
49
+ get: (_t, prop) => (_aggReq()(_NODE_OS_MOD) as any)[prop],
50
+ });
51
+
52
+ // Column name DuckDB prepends to every aggregate_update batch. One entry per
53
+ // input row; the value identifies which group the row belongs to.
54
+ export const GROUP_COLUMN_NAME = "__vgi_group_id";
55
+
56
+ export interface AggregateBindParams<TArgs = Record<string, any>> {
57
+ args: TArgs;
58
+ arguments: Arguments;
59
+ inputSchema: VgiSchema | null;
60
+ settings: Record<string, any>;
61
+ secrets: Record<string, Record<string, any>>;
62
+ }
63
+
64
+ export interface AggregateUpdateParams<TArgs = Record<string, any>, TState = any> {
65
+ states: Map<bigint, TState>;
66
+ groupIds: bigint[];
67
+ columns: any[]; // arrow Vectors per declared param
68
+ args: TArgs;
69
+ /**
70
+ * Get-or-initialize accessor for per-group state. Call this inside the loop
71
+ * only when the row actually contributes to the aggregate — skipping it
72
+ * (e.g. on NULL input under NullHandling.DEFAULT) leaves `states` without
73
+ * an entry for the group, so finalize() can return SQL NULL for that
74
+ * group (matches vgi-python's "state absent = NULL" semantics for
75
+ * SUM/AVG/MIN/MAX over zero non-null rows).
76
+ */
77
+ ensureState: (gid: bigint) => TState;
78
+ }
79
+
80
+ export interface AggregateFinalizeParams<TArgs = Record<string, any>, TState = any> {
81
+ groupIds: bigint[];
82
+ states: Map<bigint, TState | null>;
83
+ outputSchema: VgiSchema;
84
+ args: TArgs;
85
+ }
86
+
87
+ export interface AggregateFunctionConfig<TArgs = Record<string, any>, TState = any> {
88
+ name: string;
89
+ description?: string;
90
+ /**
91
+ * Map of positional argument name → Arrow type. Passed through to DuckDB
92
+ * function registration. Use `null` (untyped) for varargs placeholders.
93
+ */
94
+ args?: Record<string, VgiDataType>;
95
+ /**
96
+ * Names of args that accept a variable number of column arguments. DuckDB
97
+ * treats these as varargs of the declared Arrow type; at call time the
98
+ * function may receive 1..N columns of that type. update() sees all of
99
+ * them as consecutive entries in `columns`.
100
+ */
101
+ varargs?: string[];
102
+ /**
103
+ * Names of args whose value is constant (known at bind time) and must be
104
+ * folded away by DuckDB before reaching update(). The value arrives in
105
+ * `bindParams.args[name]`; it is NOT passed as an input column in the
106
+ * update batch. Use for aggregates parameterized by a literal (e.g.
107
+ * percentile=0.5 in vgi_percentile).
108
+ */
109
+ constParams?: string[];
110
+ /**
111
+ * Arrow type of the aggregate's output (one value per group). DuckDB uses
112
+ * this to build the result column type at bind time. May be overridden
113
+ * at bind time by `onBind` (returning a different Arrow type) for
114
+ * aggregates whose return type depends on the input column type
115
+ * (e.g. vgi_generic_sum: BIGINT input → BIGINT output, DOUBLE → DOUBLE).
116
+ */
117
+ outputType: VgiDataType;
118
+ /**
119
+ * Optional bind-time hook for dynamic output types. Receives the bind
120
+ * parameters (including input_schema) and returns the Arrow type to
121
+ * advertise for this invocation. Defaults to returning `config.outputType`.
122
+ */
123
+ onBind?: (params: AggregateBindParams<TArgs>) => VgiDataType | Promise<VgiDataType>;
124
+ /** Optional per-arg default values (positional only here). */
125
+ argDefaults?: Record<string, any>;
126
+
127
+ initialState: (params: AggregateBindParams<TArgs>) => TState;
128
+ /**
129
+ * Fold a batch of input rows into the provided per-group state map.
130
+ * Implementations mutate `states` in place. `groupIds[i]` identifies which
131
+ * group row i belongs to; `columns[k][i]` is the value for the kth declared
132
+ * column argument at row i.
133
+ */
134
+ update: (params: AggregateUpdateParams<TArgs, TState>) => void;
135
+ combine: (source: TState, target: TState, params: AggregateBindParams<TArgs>) => TState;
136
+ /** Must return a single-column RecordBatch with `groupIds.length` rows. */
137
+ finalize: (params: AggregateFinalizeParams<TArgs, TState>) => VgiBatch;
138
+
139
+ examples?: FunctionExample[];
140
+ categories?: string[];
141
+ /** Arbitrary function-level tags surfaced into duckdb_functions().tags. */
142
+ tags?: Record<string, string>;
143
+ nullHandling?: "DEFAULT" | "SPECIAL";
144
+ /**
145
+ * DuckDB secret types this aggregate needs. Advertised on the catalog's
146
+ * `required_secrets` so the C++ extension pre-resolves matching secrets and
147
+ * delivers their VALUES on the aggregate_bind request (keyed by secret name).
148
+ * The secret is read at bind time only — update/combine/finalize receive an
149
+ * empty ResolvedSecrets. Mirrors vgi-python's `Secret()` annotation on an
150
+ * aggregate `on_bind`.
151
+ */
152
+ requiredSecrets?: string[];
153
+ }
154
+
155
+ // Per-execution state registry. Value = (gid → state) + bind context.
156
+ //
157
+ // Under DuckDB's parallel aggregate with worker pooling, aggregate_* RPCs for
158
+ // the same execution_id can land on different worker processes — so state
159
+ // needs to live in a place all workers can reach. We use a filesystem-backed
160
+ // store (one directory per execution_id, one JSON file per group_id) with
161
+ // atomic rename-on-write. Fine for subprocess transport and for
162
+ // single-instance HTTP; multi-host HTTP would need a shared path or a
163
+ // swap-in network backend. Override with VGI_AGG_STATE_DIR.
164
+ //
165
+ // In-memory cache sits in front of the filesystem to keep single-process
166
+ // workloads fast — writes go to disk, reads prefer cache, and getExecutionState
167
+ // lazily reconstructs from disk when a different worker's exec is requested.
168
+
169
+ type StateMap = Map<bigint, any>;
170
+
171
+ interface ExecEntry {
172
+ states: StateMap;
173
+ bindArgs: Arguments;
174
+ bindParams: AggregateBindParams<any>;
175
+ loaded: Set<bigint>; // gids fetched from disk (so we don't refetch in same call)
176
+ }
177
+
178
+ const stateStore = new Map<string, ExecEntry>();
179
+
180
+ // Custom JSON codec: BigInt ↔ {"__bigint__":"…"}. Keep the encoder/decoder
181
+ // centralized so aggregate state classes with int64 fields (CountState,
182
+ // SumState, etc.) round-trip without extra wiring.
183
+ function jsonReplacer(_: string, v: any): any {
184
+ if (typeof v === "bigint") return { __bigint__: v.toString() };
185
+ return v;
186
+ }
187
+ function jsonReviver(_: string, v: any): any {
188
+ if (v && typeof v === "object" && typeof (v as any).__bigint__ === "string") {
189
+ return BigInt((v as any).__bigint__);
190
+ }
191
+ return v;
192
+ }
193
+
194
+ function stateDirRoot(): string {
195
+ return process.env.VGI_AGG_STATE_DIR ?? path.join(os.tmpdir(), "vgi-agg");
196
+ }
197
+
198
+ function execDir(execHex: string): string {
199
+ return path.join(stateDirRoot(), execHex);
200
+ }
201
+
202
+ function execKey(executionId: Uint8Array): string {
203
+ let s = "";
204
+ for (const b of executionId) s += b.toString(16).padStart(2, "0");
205
+ return s;
206
+ }
207
+
208
+ // Atomic write: write to tmp then rename. Avoids partial reads from another
209
+ // worker mid-update. Node's fs.renameSync is atomic on POSIX for same FS.
210
+ function atomicWriteJson(filePath: string, value: any): void {
211
+ const tmp = filePath + ".tmp." + process.pid + "." + Math.random().toString(36).slice(2);
212
+ const parent = path.dirname(filePath);
213
+ fs.mkdirSync(parent, { recursive: true });
214
+ fs.writeFileSync(tmp, JSON.stringify(value, jsonReplacer));
215
+ fs.renameSync(tmp, filePath);
216
+ }
217
+
218
+ function tryReadJson(filePath: string): any | undefined {
219
+ try {
220
+ const text = fs.readFileSync(filePath, "utf8");
221
+ return JSON.parse(text, jsonReviver);
222
+ } catch (e: any) {
223
+ if (e?.code === "ENOENT") return undefined;
224
+ throw e;
225
+ }
226
+ }
227
+
228
+ /**
229
+ * Fetch the execution context — bind args/params + states for requested gids.
230
+ * On cache miss, tries to reconstruct from on-disk files written by a
231
+ * sibling worker. Returns undefined if the exec never existed anywhere.
232
+ * `requestedGids` (optional) tells us which group states to pre-load.
233
+ */
234
+ export function getExecutionState(
235
+ executionId: Uint8Array,
236
+ requestedGids?: Iterable<bigint>,
237
+ ): ExecEntry | undefined {
238
+ const key = execKey(executionId);
239
+ let entry = stateStore.get(key);
240
+ if (!entry) {
241
+ // Try to reconstruct bind context from disk.
242
+ const ctx = tryReadJson(path.join(execDir(key), "context.json"));
243
+ if (!ctx) return undefined;
244
+ // Restore bindParams + args from the serialized context. args is kept
245
+ // only as a plain dict; the Arguments instance isn't needed after bind.
246
+ entry = {
247
+ states: new Map(),
248
+ bindArgs: ctx.bindArgs, // placeholder — Arguments instance is only needed during bind
249
+ bindParams: ctx.bindParams,
250
+ loaded: new Set(),
251
+ };
252
+ stateStore.set(key, entry);
253
+ }
254
+ if (requestedGids) {
255
+ // Always re-read from disk each RPC. A sibling worker may have just
256
+ // written a newer value for this gid; relying on our in-memory cache
257
+ // would overwrite their update on our next persist (read-modify-write
258
+ // race). File I/O per-RPC is cheap compared to the RPC overhead.
259
+ for (const gid of requestedGids) {
260
+ const stateFile = path.join(execDir(key), gid.toString() + ".json");
261
+ const s = tryReadJson(stateFile);
262
+ if (s !== undefined) {
263
+ entry.states.set(gid, s);
264
+ } else {
265
+ entry.states.delete(gid);
266
+ }
267
+ entry.loaded.add(gid);
268
+ }
269
+ }
270
+ return entry;
271
+ }
272
+
273
+ export function setExecutionState(
274
+ executionId: Uint8Array,
275
+ entry: Omit<ExecEntry, "loaded">,
276
+ ): void {
277
+ const key = execKey(executionId);
278
+ stateStore.set(key, { ...entry, loaded: new Set() });
279
+ // Persist bind context so sibling workers can reconstruct it on demand.
280
+ atomicWriteJson(path.join(execDir(key), "context.json"), {
281
+ bindParams: entry.bindParams,
282
+ // bindArgs is a live Arguments instance; stash its args map separately
283
+ // if needed for future state-dependent behavior. Workers reading the
284
+ // context receive `bindParams.args` which is what update/finalize use.
285
+ });
286
+ }
287
+
288
+ /**
289
+ * Persist the given group states to disk so sibling workers can pick them up.
290
+ * Deletes the on-disk entry when the in-memory state was removed (gid no
291
+ * longer present in states map).
292
+ */
293
+ export function persistGroupStates(
294
+ executionId: Uint8Array,
295
+ gids: Iterable<bigint>,
296
+ ): void {
297
+ const key = execKey(executionId);
298
+ const entry = stateStore.get(key);
299
+ if (!entry) return;
300
+ for (const gid of gids) {
301
+ const stateFile = path.join(execDir(key), gid.toString() + ".json");
302
+ if (entry.states.has(gid)) {
303
+ atomicWriteJson(stateFile, entry.states.get(gid));
304
+ } else {
305
+ try { fs.unlinkSync(stateFile); } catch { /* already gone */ }
306
+ }
307
+ entry.loaded.add(gid);
308
+ }
309
+ }
310
+
311
+ export function deleteExecutionState(executionId: Uint8Array): void {
312
+ const key = execKey(executionId);
313
+ stateStore.delete(key);
314
+ try { fs.rmSync(execDir(key), { recursive: true, force: true }); } catch { /* best-effort */ }
315
+ }
316
+
317
+ /**
318
+ * Declare an aggregate function. The returned VgiFunction has `kind:
319
+ * "aggregate"` plus a `config` handle the dispatch layer uses to drive the
320
+ * aggregate_* RPCs; the scalar/table-function code paths ignore it.
321
+ */
322
+ export function defineAggregate<TArgs = Record<string, any>, TState = any>(
323
+ config: AggregateFunctionConfig<TArgs, TState>,
324
+ ): VgiFunction & { aggregateConfig: AggregateFunctionConfig<TArgs, TState> } {
325
+ const specs: ArgumentSpec[] = [];
326
+ let posIdx = 0;
327
+ const varargsSet = new Set(config.varargs ?? []);
328
+ const constSet = new Set(config.constParams ?? []);
329
+ if (config.args) {
330
+ for (const [name, type] of Object.entries(config.args)) {
331
+ const hasDefault = config.argDefaults?.[name] !== undefined;
332
+ const isConst = constSet.has(name);
333
+ const isAnyType = isNull(type);
334
+ // Const params are always positional at the SQL call-site — a default
335
+ // value just means "the user may omit it and we'll use the default",
336
+ // not "rename to a named kwarg". Non-const args with defaults become
337
+ // named (matches table-function behavior).
338
+ specs.push({
339
+ name,
340
+ position: isConst || !hasDefault ? posIdx++ : name,
341
+ arrowType: isAnyType ? nullType() : type,
342
+ isAnyType,
343
+ isVarargs: varargsSet.has(name),
344
+ isConst,
345
+ });
346
+ }
347
+ }
348
+
349
+ const meta: FunctionMeta = {
350
+ name: config.name,
351
+ description: config.description,
352
+ examples: config.examples,
353
+ categories: config.categories,
354
+ tags: config.tags,
355
+ nullHandling: config.nullHandling as any,
356
+ requiredSecrets: config.requiredSecrets,
357
+ };
358
+
359
+ // The normal VgiFunction methods (bind/globalInit/createStreamHandlers) are
360
+ // never invoked for aggregates — the dispatch layer routes aggregate_* RPCs
361
+ // through the separate aggregate_* handler. Provide stubs that throw so
362
+ // misrouted calls surface loudly rather than silently hanging.
363
+ // Catalog output schema: single "result" column. Advertise config.outputType
364
+ // when it's a real type (anything other than Null) — onBind may still override
365
+ // at bind time, but for the catalog listing we prefer the static declaration
366
+ // so duckdb_functions() shows a meaningful return type. Aggregates whose
367
+ // return type genuinely depends on input set outputType: nullType() to
368
+ // advertise ANY explicitly.
369
+ const isAnyType = isNull(config.outputType);
370
+ const defaultOutputSchema = isAnyType
371
+ ? schema([field("result", nullType(), true, new Map([["vgi:any", "true"]]))])
372
+ : schema([field("result", config.outputType, true)]);
373
+
374
+ return {
375
+ kind: "aggregate" as any,
376
+ meta,
377
+ argumentSpecs: specs,
378
+ defaultOutputSchema,
379
+ bind: () => { throw new Error(`aggregate '${config.name}' received bind() — should be routed through aggregate_bind`); },
380
+ globalInit: () => { throw new Error(`aggregate '${config.name}' received globalInit() — should be routed through aggregate_bind`); },
381
+ createStreamHandlers: () => { throw new Error(`aggregate '${config.name}' is unary — no stream handlers`); },
382
+ aggregateConfig: config,
383
+ } as VgiFunction & { aggregateConfig: AggregateFunctionConfig<TArgs, TState> };
384
+ }
@@ -0,0 +1,282 @@
1
+ // Copyright 2025, 2026 Query Farm LLC - https://query.farm
2
+ // Custom `COPY ... FROM` format reader.
3
+ //
4
+ // A CopyFromFunction lets a VGI catalog act as a remote file-format reader: the
5
+ // user runs `COPY target FROM 'path' (FORMAT <name>, opt val, ...)` and the
6
+ // worker parses the source and streams Arrow batches that DuckDB inserts into
7
+ // the local `target` table. FROM only (no COPY TO).
8
+ //
9
+ // Mechanically this is an ordinary producer-mode table function (it reuses the
10
+ // table-function bind/init/scan path on both sides). What makes it a COPY format
11
+ // is twofold:
12
+ // * it sets `meta.copyFromFormat` to the SQL FORMAT identifier, and
13
+ // * the catalog advertises it via `copyFromFormats()` (the
14
+ // `catalog_copy_from_formats` RPC), so the VGI DuckDB extension registers a
15
+ // DuckDB CopyFunction for it.
16
+ //
17
+ // The COPY statement's file path and the target table's schema arrive on the
18
+ // bind through `BindRequest.copy_from` (CopyFromContext). The COPY options
19
+ // arrive as the function's normal named arguments. Mirrors vgi-python's
20
+ // `vgi.copy_from_function.CopyFromFunction`.
21
+
22
+ import { type VgiSchema, type VgiDataType } from "../arrow/index.js";
23
+ import type { OutputCollector } from "@query-farm/vgi-rpc";
24
+ import { DEFAULT_MAX_WORKERS } from "../types.js";
25
+ import type {
26
+ BindRequest,
27
+ BindResponse,
28
+ InitRequest,
29
+ GlobalInitResponse,
30
+ } from "../protocol/types.js";
31
+ import type { VgiFunction, FunctionMeta, StreamHandlers, FunctionExample } from "./types.js";
32
+ import type { CopySecretLookup } from "./copy-to.js";
33
+ import type { ArgumentSpec } from "../arguments/argument-spec.js";
34
+ import { batchToScalarDict, batchToSecretDict, safeNumber } from "../util/arrow/index.js";
35
+ import { BoundStorage, storage as globalStorage } from "./storage.js";
36
+ import type { TableProcessParams } from "./table.js";
37
+
38
+ /** Declaration of a single COPY option (a named function argument). */
39
+ export interface CopyFromOption {
40
+ /** Arrow type of the option value (e.g. `utf8()`, `int64()`). */
41
+ type: VgiDataType;
42
+ /** Per-option description, surfaced by `vgi_copy_formats()`. */
43
+ doc?: string;
44
+ /**
45
+ * Default value. When omitted the option is REQUIRED — the worker throws a
46
+ * clear error at COPY bind if the user does not supply it. Mirrors
47
+ * vgi-python's `Arg(..., default=...)` (no default => required).
48
+ */
49
+ default?: unknown;
50
+ /**
51
+ * Optional allowed value set, validated worker-side at bind. Mirrors
52
+ * vgi-python's `Arg(..., choices=[...])`.
53
+ */
54
+ choices?: unknown[];
55
+ }
56
+
57
+ export interface CopyFromReadParams<TArgs = Record<string, unknown>> {
58
+ /** Source path from the `COPY ... FROM 'path'` statement. */
59
+ path: string;
60
+ /** Parsed COPY options (defaults applied). */
61
+ options: TArgs;
62
+ /**
63
+ * The COPY target's schema. Every emitted batch must have this exact schema
64
+ * (names + types, in order) — DuckDB inserts no cast before the INSERT.
65
+ */
66
+ expectedSchema: VgiSchema;
67
+ /** Full process parameters (settings, secrets, storage). */
68
+ processParams: TableProcessParams<TArgs>;
69
+ /** Collector to emit batches / log. `finish()` is called for you. */
70
+ out: OutputCollector;
71
+ }
72
+
73
+ export interface CopyFromFunctionConfig<TArgs = Record<string, unknown>> {
74
+ /** Handler name (the function's registered name; `Meta.name`). */
75
+ name: string;
76
+ /** SQL `FORMAT` identifier users type, e.g. `COPY t FROM 'x' (FORMAT myfmt)`. */
77
+ format: string;
78
+ /** Function description (intrinsic documentation; `Meta.description`). */
79
+ description?: string;
80
+ /** Optional free-text comment surfaced by `vgi_copy_formats()`. */
81
+ comment?: string | null;
82
+ /** Reserved for a future `COPY ... TO`; only `"from"` is supported today. */
83
+ direction?: string;
84
+ /** COPY options, keyed by option name. The `file_path` is NOT an option. */
85
+ options?: Record<string, CopyFromOption>;
86
+ categories?: string[];
87
+ tags?: Record<string, string>;
88
+ examples?: FunctionExample[];
89
+ requiredSettings?: string[];
90
+ requiredSecrets?: string[];
91
+ /**
92
+ * Optional secret-bind hook: forward CREATE SECRET credentials for
93
+ * secret-backed cloud sources (S3/GCS/HTTP/…). Called during bind (only on the
94
+ * first pass); return the secrets to resolve — typically scoped by the source
95
+ * `path`. The framework's two-phase secret bind resolves each lookup from the
96
+ * caller's SecretManager and surfaces the resolved values on
97
+ * `processParams.secrets` at `read` time. Mirrors vgi-python's
98
+ * `CopyFromFunction.on_secrets`.
99
+ */
100
+ onSecrets?: (params: {
101
+ options: TArgs;
102
+ path: string;
103
+ bindCall: BindRequest;
104
+ }) => CopySecretLookup[] | void;
105
+ /** Parse `path` and emit Arrow batches matching `expectedSchema`. */
106
+ read: (params: CopyFromReadParams<TArgs>) => void | Promise<void>;
107
+ }
108
+
109
+ /**
110
+ * Define a custom `COPY ... FROM` format reader, returned as a `VgiFunction`
111
+ * (kind `"table"`) carrying the `copyFromFormat` metadata marker. Register it in
112
+ * the catalog's function list like any table function; the catalog's
113
+ * `copyFromFormats()` introspection picks it up and advertises the format.
114
+ */
115
+ export function defineCopyFromFunction<TArgs = Record<string, unknown>>(
116
+ config: CopyFromFunctionConfig<TArgs>,
117
+ ): VgiFunction {
118
+ // Build named-argument specs from the options. COPY passes options by name,
119
+ // so every option is a named argument (string position).
120
+ const specs: ArgumentSpec[] = [];
121
+ for (const [name, opt] of Object.entries(config.options ?? {})) {
122
+ specs.push({
123
+ name,
124
+ position: name, // named argument
125
+ arrowType: opt.type,
126
+ doc: opt.doc,
127
+ });
128
+ }
129
+
130
+ const options = config.options ?? {};
131
+ const requiredNames = Object.entries(options)
132
+ .filter(([, opt]) => !("default" in opt))
133
+ .map(([name]) => name);
134
+
135
+ const meta: FunctionMeta = {
136
+ name: config.name,
137
+ description: config.description,
138
+ categories: config.categories,
139
+ tags: config.tags,
140
+ examples: config.examples,
141
+ requiredSettings: config.requiredSettings,
142
+ requiredSecrets: config.requiredSecrets,
143
+ copyFromFormat: config.format,
144
+ copyFromDirection: config.direction ?? "from",
145
+ copyFromComment: config.comment ?? null,
146
+ };
147
+
148
+ // Extract + validate the COPY options from the bind/init arguments. Options
149
+ // arrive as named arguments; required ones (no default) must be present.
150
+ function extractOptions(request: BindRequest): TArgs {
151
+ const args: Record<string, unknown> = {};
152
+ for (const [name, opt] of Object.entries(options)) {
153
+ const hasDefault = "default" in opt;
154
+ let val: unknown;
155
+ try {
156
+ val = request.arguments.get(name, hasDefault ? opt.default : undefined);
157
+ } catch {
158
+ val = hasDefault ? opt.default : undefined;
159
+ }
160
+ if (val === undefined || val === null) {
161
+ if (!hasDefault) {
162
+ throw new Error(
163
+ `COPY (FORMAT ${config.format}): missing required option '${name}'`,
164
+ );
165
+ }
166
+ val = opt.default;
167
+ }
168
+ if (typeof val === "bigint") val = safeNumber(val);
169
+ if (opt.choices && !opt.choices.includes(val)) {
170
+ throw new Error(
171
+ `COPY (FORMAT ${config.format}): option '${name}' value ${JSON.stringify(val)} ` +
172
+ `is not one of ${JSON.stringify(opt.choices)}`,
173
+ );
174
+ }
175
+ args[name] = val;
176
+ }
177
+ return args as TArgs;
178
+ }
179
+
180
+ function requireCopyFrom(request: BindRequest) {
181
+ const cf = request.copy_from;
182
+ if (cf == null) {
183
+ throw new Error(
184
+ `${config.name} is a COPY FROM format reader; invoke it via ` +
185
+ `COPY <table> FROM '<path>' (FORMAT ${config.format}), not as a table function.`,
186
+ );
187
+ }
188
+ return cf;
189
+ }
190
+
191
+ return {
192
+ kind: "table",
193
+ meta,
194
+ argumentSpecs: specs,
195
+ // Table functions determine output at bind time; the COPY target schema is
196
+ // not known until bind, so advertise an empty default output schema.
197
+
198
+ bind(request: BindRequest): BindResponse {
199
+ const cf = requireCopyFrom(request);
200
+ // Validate options eagerly so a missing/invalid option fails at COPY bind.
201
+ const opts = extractOptions(request);
202
+ // Forward any requested secrets on the first bind pass so the two-phase
203
+ // secret bind resolves them (the resolved values reach read via
204
+ // processParams.secrets).
205
+ if (config.onSecrets && !request.resolved_secrets_provided) {
206
+ const lookups =
207
+ config.onSecrets({ options: opts, path: cf.file_path, bindCall: request }) ?? [];
208
+ if (lookups.length > 0) {
209
+ return {
210
+ output_schema: cf.expected_schema,
211
+ opaque_data: null,
212
+ lookup_secret_types: lookups.map((l) => l.secretType),
213
+ lookup_scopes: lookups.map((l) => l.scope ?? ""),
214
+ lookup_names: lookups.map((l) => l.name ?? ""),
215
+ };
216
+ }
217
+ }
218
+ // DuckDB forces the scan's output to the target table's columns, so a
219
+ // COPY-FROM reader must produce exactly expected_schema.
220
+ return {
221
+ output_schema: cf.expected_schema,
222
+ opaque_data: null,
223
+ };
224
+ },
225
+
226
+ globalInit(request: InitRequest): GlobalInitResponse {
227
+ const executionId = new Uint8Array(16);
228
+ crypto.getRandomValues(executionId);
229
+ if (request.execution_id) {
230
+ return { max_workers: 1, execution_id: request.execution_id, opaque_data: null };
231
+ }
232
+ return { max_workers: 1, execution_id: executionId, opaque_data: null };
233
+ },
234
+
235
+ createStreamHandlers(request: InitRequest, response: GlobalInitResponse): StreamHandlers {
236
+ const cf = requireCopyFrom(request.bind_call);
237
+ const args = extractOptions(request.bind_call);
238
+ const settings = batchToScalarDict(request.bind_call.settings);
239
+ const secrets = batchToSecretDict(request.bind_call.secrets);
240
+ const outputSchema = request.output_schema;
241
+ const boundStorage = new BoundStorage(globalStorage, response.execution_id);
242
+
243
+ const processParams: TableProcessParams<TArgs> = {
244
+ args,
245
+ initCall: request,
246
+ initResponse: response,
247
+ outputSchema,
248
+ settings,
249
+ secrets,
250
+ storage: boundStorage,
251
+ atUnit: request.bind_call.at_unit ?? undefined,
252
+ atValue: request.bind_call.at_value ?? undefined,
253
+ };
254
+
255
+ return {
256
+ outputSchema,
257
+ // Single-shot: the whole source is read on the first producer tick. The
258
+ // `done` guard lives under `.state` so it round-trips across HTTP
259
+ // exchanges (the framework serializes handlerState.state).
260
+ producerInit: () => ({ state: { done: false } }),
261
+ producerFn: async (
262
+ pState: { state: { done: boolean } },
263
+ out: OutputCollector,
264
+ ) => {
265
+ if (pState.state.done) {
266
+ out.finish();
267
+ return;
268
+ }
269
+ await config.read({
270
+ path: cf.file_path,
271
+ options: args,
272
+ expectedSchema: outputSchema,
273
+ processParams,
274
+ out,
275
+ });
276
+ pState.state.done = true;
277
+ out.finish();
278
+ },
279
+ };
280
+ },
281
+ };
282
+ }