Line data Source code
1 : export type FilterOp = "=" | "<" | "<=" | ">" | ">=";
2 :
3 : export type AggFunc = "min" | "max" | "avg" | "count" | "count_star";
4 :
5 : export interface OrderByConstraint {
6 : field: string;
7 : ascending?: boolean;
8 : nullsFirst?: boolean;
9 : }
10 :
11 : export interface BinningConfig {
12 : fieldName: string;
13 : statsTable: string;
14 : statsMinField: string;
15 : statsMaxField: string;
16 : binCount: number;
17 : minBinWidth?: number;
18 : outputAlias: string;
19 : toNumericFn?: string;
20 : }
21 :
22 : export interface GroupByKey {
23 : fieldName: string;
24 : outputAlias: string;
25 : binning?: GroupByKeyBinning;
26 : }
27 :
28 : export interface GroupByKeyBinning {
29 : preBinnedFieldName?: string;
30 : binCount: number;
31 : minBinWidth?: number;
32 : statsTable: string;
33 : statsMinField: string;
34 : statsMaxField: string;
35 : outputBinLbAlias: string;
36 : outputBinUbAlias: string;
37 : outputBinWidthAlias: string;
38 : includeNullBin?: boolean;
39 : toNumericFn?: string;
40 : }
41 :
42 : export interface GroupByAggregate {
43 : fieldName?: string;
44 : outputAlias: string;
45 : func: AggFunc;
46 : distinct?: boolean;
47 : }
48 :
49 : export interface GroupByConfig {
50 : keys: GroupByKey[];
51 : aggregates: GroupByAggregate[];
52 : }
53 :
54 : export interface ScalarFilter {
55 : fieldName: string;
56 : op: FilterOp;
57 : value: number;
58 : }
59 :
60 : export interface SemiJoinFilter {
61 : fieldName: string;
62 : tableName: string;
63 : joinFieldName: string;
64 : }
65 :
66 : interface ValueIdentifierEntry {
67 : fieldName: string;
68 : outputAlias: string;
69 : }
70 :
71 1 : export function quoteIdent(name: string): string {
72 1 : return `"${name.replace(/"/g, '""')}"`;
73 1 : }
74 :
75 1 : export function quoteStringLiteral(value: string): string {
76 1 : return `'${value.replace(/'/g, "''")}'`;
77 1 : }
78 :
79 1 : function formatLiteral(value: number): string {
80 1 : return value.toString();
81 1 : }
82 :
83 1 : function binWidthExpression(range: string, binCount: number, minBinWidth?: number): string {
84 1 : const minimum = minBinWidth ?? 1e-15;
85 1 : return `GREATEST(${range} / ${binCount}, ${minimum})`;
86 1 : }
87 :
88 1 : function toNumeric(expr: string, fn?: string): string {
89 1 : switch (fn) {
90 1 : case "EPOCH":
91 1 : return `EXTRACT(EPOCH FROM ${expr})`;
92 1 : case "TIME_EPOCH":
93 1 : return `EXTRACT(EPOCH FROM CAST(${expr} AS INTERVAL))`;
94 1 : case "BOOLEAN":
95 1 : return `CASE WHEN ${expr} THEN 1.0 WHEN NOT ${expr} THEN 0.0 END`;
96 1 : default:
97 1 : return fn
98 0 : ? `${fn}(${expr})`
99 1 : : `CAST(${expr} AS DOUBLE PRECISION)`;
100 1 : }
101 1 : }
102 :
103 1 : function formatAggExpr(agg: GroupByAggregate): string {
104 1 : switch (agg.func) {
105 1 : case "min":
106 1 : return `MIN(${quoteIdent(agg.fieldName!)})`;
107 1 : case "max":
108 1 : return `MAX(${quoteIdent(agg.fieldName!)})`;
109 1 : case "avg":
110 1 : return `AVG(${quoteIdent(agg.fieldName!)})`;
111 1 : case "count":
112 1 : if (agg.distinct) {
113 1 : return `COUNT(DISTINCT ${quoteIdent(agg.fieldName!)})`;
114 : }
115 1 : return `COUNT(${quoteIdent(agg.fieldName!)})`;
116 1 : case "count_star":
117 1 : return `COUNT(*)`;
118 1 : }
119 1 : }
120 :
121 1 : function formatOrderByItem(c: OrderByConstraint): string {
122 1 : const dir = (c.ascending ?? true) ? "ASC" : "DESC";
123 1 : const nulls = (c.nullsFirst ?? false) ? "NULLS FIRST" : "NULLS LAST";
124 1 : return `${quoteIdent(c.field)} ${dir} ${nulls}`;
125 1 : }
126 :
127 1 : export class SQLFrame {
128 1 : private _inputTable: string;
129 1 : private _valueIdentifiers: ValueIdentifierEntry[] = [];
130 1 : private _binnings: BinningConfig[] = [];
131 1 : private _scalarFilters: ScalarFilter[] = [];
132 1 : private _semiJoinFilters: SemiJoinFilter[] = [];
133 1 : private _groupBy: GroupByConfig | null = null;
134 1 : private _orderBy: OrderByConstraint[] | null = null;
135 1 : private _limit: number | undefined = undefined;
136 1 : private _projection: string[] | null = null;
137 :
138 1 : private constructor(inputTable: string) {
139 1 : this._inputTable = inputTable;
140 1 : }
141 :
142 1 : static from(inputTable: string): SQLFrame {
143 1 : return new SQLFrame(inputTable);
144 1 : }
145 :
146 1 : valueIdentifier(fieldName: string, outputAlias: string): SQLFrame {
147 1 : this._valueIdentifiers.push({ fieldName, outputAlias });
148 1 : return this;
149 1 : }
150 :
151 1 : binning(config: BinningConfig): SQLFrame {
152 1 : this._binnings.push(config);
153 1 : return this;
154 1 : }
155 :
156 1 : filter(fieldName: string, op: FilterOp, value: number): SQLFrame {
157 1 : this._scalarFilters.push({ fieldName, op, value });
158 1 : return this;
159 1 : }
160 :
161 1 : semiJoinFilter(fieldName: string, tableName: string, joinFieldName: string): SQLFrame {
162 1 : this._semiJoinFilters.push({ fieldName, tableName, joinFieldName });
163 1 : return this;
164 1 : }
165 :
166 1 : groupBy(config: GroupByConfig): SQLFrame {
167 1 : this._groupBy = config;
168 1 : return this;
169 1 : }
170 :
171 1 : orderBy(constraints: OrderByConstraint[], limit?: number): SQLFrame {
172 1 : this._orderBy = constraints;
173 1 : this._limit = limit;
174 1 : return this;
175 1 : }
176 :
177 1 : project(fields: string[]): SQLFrame {
178 1 : this._projection = fields;
179 1 : return this;
180 1 : }
181 :
182 1 : toSQL(): string {
183 1 : const ctes: string[] = [];
184 1 : let prev = quoteIdent(this._inputTable);
185 :
186 : // 1. ValueIdentifiers
187 1 : if (this._valueIdentifiers.length > 0) {
188 1 : const cteName = "__value_ids";
189 1 : const windowExprs = this._valueIdentifiers.map(vi =>
190 1 : `dense_rank() OVER (ORDER BY ${quoteIdent(vi.fieldName)} ASC NULLS LAST) AS ${quoteIdent(vi.outputAlias)}`
191 : );
192 1 : const body = `SELECT *, ${windowExprs.join(", ")} FROM ${prev}`;
193 1 : ctes.push(`${cteName} AS (\n ${body}\n )`);
194 1 : prev = cteName;
195 1 : }
196 :
197 : // 2. Binning
198 1 : if (this._binnings.length > 0) {
199 1 : const cteName = "__binning";
200 1 : const binExprs: string[] = [];
201 1 : const statsJoins: string[] = [];
202 :
203 1 : for (let idx = 0; idx < this._binnings.length; idx++) {
204 1 : const b = this._binnings[idx];
205 1 : const fn = b.toNumericFn;
206 1 : const alias = `__bs${idx}`;
207 1 : const sub =
208 1 : `SELECT ` +
209 1 : `${toNumeric(quoteIdent(b.statsMinField), fn)} AS __min, ` +
210 1 : `${binWidthExpression(`ABS(${toNumeric(quoteIdent(b.statsMaxField), fn)} - ${toNumeric(quoteIdent(b.statsMinField), fn)})`, b.binCount, b.minBinWidth)} AS __bin_width ` +
211 1 : `FROM ${quoteIdent(b.statsTable)}`;
212 1 : statsJoins.push(`CROSS JOIN (${sub}) ${alias}`);
213 1 : binExprs.push(
214 1 : `(${toNumeric(`t.${quoteIdent(b.fieldName)}`, fn)} - ${alias}.__min) / ${alias}.__bin_width AS ${quoteIdent(b.outputAlias)}`
215 1 : );
216 1 : }
217 :
218 1 : const body = `SELECT t.*, ${binExprs.join(", ")} FROM ${prev} t ${statsJoins.join(" ")}`;
219 1 : ctes.push(`${cteName} AS (\n ${body}\n )`);
220 1 : prev = cteName;
221 1 : }
222 :
223 : // 3. Filters (scalar)
224 1 : if (this._scalarFilters.length > 0) {
225 1 : const cteName = "__filtered";
226 1 : const conditions = this._scalarFilters.map(f =>
227 1 : `${quoteIdent(f.fieldName)} ${f.op} ${formatLiteral(f.value)}`
228 1 : );
229 1 : const body = `SELECT * FROM ${prev} WHERE ${conditions.join(" AND ")}`;
230 1 : ctes.push(`${cteName} AS (\n ${body}\n )`);
231 1 : prev = cteName;
232 1 : }
233 :
234 : // 3b. Filters (semi-join)
235 1 : for (let i = 0; i < this._semiJoinFilters.length; i++) {
236 1 : const sj = this._semiJoinFilters[i];
237 1 : const cteName = `__semi_${i}`;
238 1 : const body = `SELECT * FROM ${prev} WHERE ${quoteIdent(sj.fieldName)} IN (SELECT ${quoteIdent(sj.joinFieldName)} FROM ${quoteIdent(sj.tableName)})`;
239 1 : ctes.push(`${cteName} AS (\n ${body}\n )`);
240 1 : prev = cteName;
241 1 : }
242 :
243 : // 4. GroupBy
244 1 : if (this._groupBy !== null) {
245 1 : const gb = this._groupBy;
246 1 : let binnedKey: { key: GroupByKey; binning: GroupByKeyBinning } | null = null;
247 :
248 1 : for (const key of gb.keys) {
249 1 : if (key.binning) {
250 1 : binnedKey = { key, binning: key.binning };
251 1 : break;
252 1 : }
253 : }
254 :
255 1 : if (binnedKey !== null) {
256 1 : prev = this._buildBinnedGroupBy(ctes, prev, gb, binnedKey.key, binnedKey.binning);
257 : } else {
258 1 : prev = this._buildSimpleGroupBy(ctes, prev, gb);
259 : }
260 1 : }
261 :
262 : // Final SELECT
263 1 : const selectCols = this._projection
264 1 : ? this._projection.map(f => quoteIdent(f)).join(", ")
265 1 : : "*";
266 :
267 1 : let finalSelect = `SELECT ${selectCols} FROM ${prev}`;
268 :
269 : // 5. OrderBy
270 1 : if (this._orderBy !== null && this._orderBy.length > 0) {
271 1 : finalSelect += `\nORDER BY ${this._orderBy.map(formatOrderByItem).join(", ")}`;
272 : }
273 :
274 : // 5b. Limit
275 1 : if (this._limit !== undefined) {
276 1 : finalSelect += `\nLIMIT ${this._limit}`;
277 : }
278 :
279 1 : if (ctes.length === 0) {
280 1 : return finalSelect;
281 : }
282 1 : return `WITH\n ${ctes.join(",\n ")}\n${finalSelect}`;
283 1 : }
284 :
285 1 : private _buildSimpleGroupBy(ctes: string[], prev: string, gb: GroupByConfig): string {
286 1 : const cteName = "__grouped";
287 1 : const keyExprs = gb.keys.map(k => `${quoteIdent(k.fieldName)} AS ${quoteIdent(k.outputAlias)}`);
288 1 : const aggExprs = gb.aggregates.map(a => `${formatAggExpr(a)} AS ${quoteIdent(a.outputAlias)}`);
289 1 : const selectList = [...keyExprs, ...aggExprs].join(", ");
290 1 : const groupByPositions = gb.keys.map((_, i) => i + 1).join(", ");
291 1 : const body = gb.keys.length > 0
292 1 : ? `SELECT ${selectList} FROM ${prev} GROUP BY ${groupByPositions}`
293 1 : : `SELECT ${selectList} FROM ${prev}`;
294 1 : ctes.push(`${cteName} AS (\n ${body}\n )`);
295 1 : return cteName;
296 1 : }
297 :
298 1 : private _buildBinnedGroupBy(
299 1 : ctes: string[],
300 1 : prev: string,
301 1 : gb: GroupByConfig,
302 1 : binnedKeyDef: GroupByKey,
303 1 : binning: GroupByKeyBinning,
304 1 : ): string {
305 1 : const binCount = binning.binCount;
306 1 : const includeNullBin = binning.includeNullBin ?? false;
307 1 : const totalBins = includeNullBin ? binCount + 1 : binCount;
308 1 : const binAlias = binnedKeyDef.outputAlias;
309 1 : const fn = binning.toNumericFn;
310 :
311 : // Temporal bounds are kept in the original type for date formatting below.
312 : // Other converted values, such as booleans, use a numeric __min.
313 1 : const statsCteName = "__grp_stats";
314 1 : const minField = quoteIdent(binning.statsMinField);
315 1 : const maxField = quoteIdent(binning.statsMaxField);
316 1 : const isTemporal = fn === "EPOCH" || fn === "TIME_EPOCH";
317 1 : const statsSql = isTemporal
318 1 : ? `SELECT ` +
319 1 : `${minField} AS __min, ` +
320 1 : `${binWidthExpression(`ABS(${toNumeric(maxField, fn)} - ${toNumeric(minField, fn)})`, binCount, binning.minBinWidth)} AS __bin_width ` +
321 1 : `FROM ${quoteIdent(binning.statsTable)}`
322 1 : : `SELECT ` +
323 1 : `${toNumeric(minField, fn)} AS __min, ` +
324 1 : `${binWidthExpression(`ABS(${toNumeric(maxField, fn)} - ${toNumeric(minField, fn)})`, binCount, binning.minBinWidth)} AS __bin_width ` +
325 1 : `FROM ${quoteIdent(binning.statsTable)}`;
326 1 : ctes.push(`${statsCteName} AS (\n ${statsSql}\n )`);
327 :
328 : // Bin expression
329 1 : let binExpr: string;
330 1 : if (binning.preBinnedFieldName) {
331 1 : binExpr = `t.${quoteIdent(binning.preBinnedFieldName)}`;
332 : } else {
333 1 : binExpr = isTemporal
334 1 : ? `(${toNumeric(`t.${quoteIdent(binnedKeyDef.fieldName)}`, fn)} - ${toNumeric('s.__min', fn)}) / s.__bin_width`
335 1 : : `(${toNumeric(`t.${quoteIdent(binnedKeyDef.fieldName)}`, fn)} - s.__min) / s.__bin_width`;
336 : }
337 :
338 : // Clamped bin. The floating-point bin index must be clamped to
339 : // [0, binCount-1] *before* casting to INTEGER: for non-finite values
340 : // (NaN / +-Inf, common in REAL/Float32 columns) or near-constant columns
341 : // where __bin_width is clamped to 1e-15, FLOOR(binExpr) can be NaN or far
342 : // outside the INT32 range, and a bare CAST(... AS INTEGER) would raise a
343 : // conversion / out-of-range error. We therefore clamp in the DOUBLE domain
344 : // (NaN falls through every comparison, so it is mapped to bin 0) and only
345 : // cast the already-bounded value.
346 : const clampedFloat =
347 : `CASE ` +
348 1 : `WHEN FLOOR(${binExpr}) >= ${binCount} THEN ${binCount - 1} ` +
349 1 : `WHEN FLOOR(${binExpr}) >= 0 THEN FLOOR(${binExpr}) ` +
350 : `ELSE 0 ` +
351 : `END`;
352 1 : const flooredBin = `CAST(${clampedFloat} AS INTEGER)`;
353 1 : let binCaseExpr: string;
354 1 : if (includeNullBin) {
355 1 : binCaseExpr =
356 1 : `CASE ` +
357 1 : `WHEN t.${quoteIdent(binnedKeyDef.fieldName)} IS NULL THEN ${binCount} ` +
358 1 : `ELSE ${flooredBin} ` +
359 : `END`;
360 : } else {
361 1 : binCaseExpr = flooredBin;
362 : }
363 :
364 : // Build key exprs and agg exprs
365 1 : const keyExprs: string[] = [];
366 1 : for (const key of gb.keys) {
367 1 : if (key === binnedKeyDef) {
368 1 : keyExprs.push(`${binCaseExpr} AS ${quoteIdent(binAlias)}`);
369 : } else {
370 0 : keyExprs.push(`${quoteIdent(key.fieldName)} AS ${quoteIdent(key.outputAlias)}`);
371 : }
372 : }
373 1 : const aggExprs = gb.aggregates.map(a => `${formatAggExpr(a)} AS ${quoteIdent(a.outputAlias)}`);
374 1 : const selectList = [...keyExprs, ...aggExprs].join(", ");
375 1 : const groupByPositions = gb.keys.map((_, i) => i + 1).join(", ");
376 :
377 1 : const groupedCteName = "__grouped";
378 1 : const groupedSql = `SELECT ${selectList} FROM ${prev} t CROSS JOIN ${statsCteName} s GROUP BY ${groupByPositions}`;
379 1 : ctes.push(`${groupedCteName} AS (\n ${groupedSql}\n )`);
380 :
381 : // All bins CTE
382 1 : const allBinsCteName = "__all_bins";
383 1 : const allBinsSql = `SELECT * FROM generate_series(0, ${totalBins - 1}) AS bins(${quoteIdent(binAlias)})`;
384 1 : ctes.push(`${allBinsCteName} AS (\n ${allBinsSql}\n )`);
385 :
386 : // Join missing bins
387 1 : const otherCols: string[] = [];
388 1 : for (const key of gb.keys) {
389 1 : if (key !== binnedKeyDef) {
390 0 : otherCols.push(`g.${quoteIdent(key.outputAlias)}`);
391 : }
392 : }
393 1 : for (const agg of gb.aggregates) {
394 1 : otherCols.push(`g.${quoteIdent(agg.outputAlias)}`);
395 : }
396 :
397 1 : const withBinsCteName = "__with_bins";
398 1 : const withBinsSelect = [`ab.${quoteIdent(binAlias)}`, ...otherCols].join(", ");
399 1 : const withBinsSql = `SELECT ${withBinsSelect} FROM ${allBinsCteName} ab LEFT JOIN ${groupedCteName} g ON ab.${quoteIdent(binAlias)} = g.${quoteIdent(binAlias)}`;
400 1 : ctes.push(`${withBinsCteName} AS (\n ${withBinsSql}\n )`);
401 :
402 : // Temporal bounds are returned as epoch milliseconds for Arrow formatting.
403 1 : const metaCteName = "__bin_meta";
404 1 : const widthAlias = binning.outputBinWidthAlias;
405 1 : const lbAlias = binning.outputBinLbAlias;
406 1 : const ubAlias = binning.outputBinUbAlias;
407 1 : const binRef = quoteIdent(binAlias);
408 :
409 1 : const offsetExpr = (offset: string) => isTemporal
410 1 : ? `${toNumeric('s.__min', fn)} * 1000.0 + (${offset}) * 1000.0`
411 1 : : `s.__min + (${offset})`;
412 :
413 1 : let widthExpr: string;
414 1 : let lbExpr: string;
415 1 : let ubExpr: string;
416 :
417 1 : if (includeNullBin) {
418 1 : widthExpr = `CASE WHEN ${binRef} = ${binCount} THEN NULL ELSE s.__bin_width END`;
419 1 : lbExpr = `CASE WHEN ${binRef} = ${binCount} THEN NULL ELSE ${offsetExpr(`${binRef} * s.__bin_width`)} END`;
420 1 : ubExpr = `CASE WHEN ${binRef} = ${binCount} THEN NULL ELSE ${offsetExpr(`(${binRef} + 1) * s.__bin_width`)} END`;
421 1 : } else {
422 1 : widthExpr = `s.__bin_width`;
423 1 : lbExpr = offsetExpr(`${binRef} * s.__bin_width`);
424 1 : ubExpr = offsetExpr(`(${binRef} + 1) * s.__bin_width`);
425 1 : }
426 :
427 1 : const metaSql =
428 1 : `SELECT t.*, ` +
429 1 : `${widthExpr} AS ${quoteIdent(widthAlias)}, ` +
430 1 : `${lbExpr} AS ${quoteIdent(lbAlias)}, ` +
431 1 : `${ubExpr} AS ${quoteIdent(ubAlias)} ` +
432 1 : `FROM ${withBinsCteName} t CROSS JOIN ${statsCteName} s`;
433 1 : ctes.push(`${metaCteName} AS (\n ${metaSql}\n )`);
434 :
435 1 : return metaCteName;
436 1 : }
437 1 : }
|
