squirreling 0.12.14 → 0.12.16

package/README.md

@@ -154,13 +154,14 @@ Squirreling mostly follows the SQL standard. The following features are supporte
 
 ### Functions
 
-- Aggregate: `COUNT`, `SUM`, `AVG`, `MIN`, `MAX`, `MEDIAN`, `PERCENTILE_CONT`, `APPROX_QUANTILE`, `STDDEV_POP`, `STDDEV_SAMP`, `ARRAY_AGG`, `JSON_ARRAYAGG`, `STRING_AGG`
+- Aggregate: `COUNT`, `COUNTIF`, `SUM`, `AVG`, `MIN`, `MAX`, `MEDIAN`, `PERCENTILE_CONT`, `APPROX_QUANTILE`, `STDDEV_POP`, `STDDEV_SAMP`, `ARRAY_AGG`, `JSON_ARRAYAGG`, `STRING_AGG`
+- Window: `ROW_NUMBER`, `LAG`, `LEAD`
 - String: `CONCAT`, `SUBSTRING`, `REPLACE`, `LENGTH`, `UPPER`, `LOWER`, `TRIM`, `LEFT`, `RIGHT`, `INSTR`, `POSITION`, `STRPOS`
 - Math: `ABS`, `SIGN`, `CEIL`, `FLOOR`, `ROUND`, `MOD`, `RAND`, `RANDOM`, `LN`, `LOG10`, `EXP`, `POWER`, `SQRT`
 - Trig: `SIN`, `COS`, `TAN`, `COT`, `ASIN`, `ACOS`, `ATAN`, `ATAN2`, `DEGREES`, `RADIANS`, `PI`
-- Date: `CURRENT_DATE`, `CURRENT_TIME`, `CURRENT_TIMESTAMP`, `DATE_PART`, `DATE_TRUNC`, `EXTRACT`, `INTERVAL`
+- Date: `CURRENT_DATE`, `CURRENT_TIME`, `CURRENT_TIMESTAMP`, `DATE_DIFF`, `DATEDIFF`, `DATE_PART`, `DATE_TRUNC`, `EPOCH`, `EXTRACT`, `INTERVAL`
 - Json: `JSON_VALUE`, `JSON_QUERY`, `JSON_EXTRACT`, `JSON_OBJECT`, `JSON_ARRAY_LENGTH`, `JSON_VALID`, `JSON_TYPE`
-- Array: `ARRAY_LENGTH`, `ARRAY_POSITION`, `ARRAY_SORT`, `CARDINALITY`
+- Array: `ARRAY_LENGTH`, `ARRAY_POSITION`, `ARRAY_CONTAINS`, `ARRAY_SORT`, `CARDINALITY`, `SIZE`
 - Table functions: `UNNEST`, `EXPLODE`, `JSON_EACH`
 - Regex: `REGEXP_SUBSTR`, `REGEXP_EXTRACT`, `REGEXP_REPLACE`, `REGEXP_MATCHES`
 - Spatial: `ST_GeomFromText`, `ST_MakeEnvelope`, `ST_AsText`, `ST_Intersects`, `ST_Contains`, `ST_ContainsProperly`, `ST_Within`, `ST_Overlaps`, `ST_Touches`, `ST_Equals`, `ST_Crosses`, `ST_Covers`, `ST_CoveredBy`, `ST_DWithin`

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.12.14",
+  "version": "0.12.16",
   "description": "Squirreling Async SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",

package/src/execute/window.js

@@ -20,10 +20,13 @@ export function executeWindow(plan, context) {
   const child = executePlan({ plan: plan.child, context })
   const extraColumns = plan.windows.map(w => w.alias)
 
-  // Streaming fast path: every window is OVER () with no partition/order, so
-  // each row's output depends only on its position in the input stream. Avoids
-  // buffering — critical for large scans (e.g. parquet).
-  const streamable = plan.windows.every(w => w.partitionBy.length === 0 && w.orderBy.length === 0)
+  // Streaming fast path: every window is a positional function (e.g.
+  // ROW_NUMBER) with OVER () — no partition/order — so each row's output
+  // depends only on its index in the input stream. Avoids buffering, which
+  // matters for large scans (e.g. parquet).
+  const streamable = plan.windows.every(w =>
+    w.funcName === 'ROW_NUMBER' && w.partitionBy.length === 0 && w.orderBy.length === 0
+  )
 
   if (streamable) {
     return {
@@ -37,7 +40,7 @@ export function executeWindow(plan, context) {
           i++
           const cells = { ...row.cells }
           for (const w of plan.windows) {
-            const value = assignRowNumber(w.funcName, i - 1)
+            const value = i
             cells[w.alias] = () => Promise.resolve(value)
           }
           yield {
@@ -119,6 +122,8 @@ async function computeWindow(spec, rows, output, context) {
     if (context.signal?.aborted) return
 
     // Order within the partition. Empty ORDER BY → input order.
+    /** @type {number[]} */
+    let ordered
     if (spec.orderBy.length) {
       const orderValues = await Promise.all(bucket.map(idx =>
         Promise.all(spec.orderBy.map(term => evaluateExpr({ node: term.expr, row: rows[idx], context })))
@@ -132,23 +137,51 @@ async function computeWindow(spec, rows, output, context) {
         }
         return a.pos - b.pos
       })
-      for (let k = 0; k < entries.length; k++) {
-        output[entries[k].idx] = assignRowNumber(spec.funcName, k)
-      }
+      ordered = entries.map(e => e.idx)
     } else {
-      for (let k = 0; k < bucket.length; k++) {
-        output[bucket[k]] = assignRowNumber(spec.funcName, k)
-      }
+      ordered = bucket
     }
+
+    await applyWindowFunction(spec, ordered, rows, output, context)
   }
 }
 
 /**
- * @param {string} funcName
- * @param {number} rank - 0-based rank within the partition
- * @returns {SqlPrimitive}
+ * Computes window function values for a single partition's rows in order.
+ *
+ * @param {WindowSpec} spec
+ * @param {number[]} ordered - row indices in window order
+ * @param {AsyncRow[]} rows
+ * @param {SqlPrimitive[]} output
+ * @param {ExecuteContext} context
  */
-function assignRowNumber(funcName, rank) {
-  if (funcName === 'ROW_NUMBER') return rank + 1
-  throw new Error(`Unsupported window function: ${funcName}`)
+async function applyWindowFunction(spec, ordered, rows, output, context) {
+  if (spec.funcName === 'ROW_NUMBER') {
+    for (let k = 0; k < ordered.length; k++) {
+      output[ordered[k]] = k + 1
+    }
+    return
+  }
+  if (spec.funcName === 'LAG' || spec.funcName === 'LEAD') {
+    const direction = spec.funcName === 'LAG' ? -1 : 1
+    const [valueExpr, offsetExpr, defaultExpr] = spec.args
+    for (let k = 0; k < ordered.length; k++) {
+      if (context.signal?.aborted) return
+      const idx = ordered[k]
+      const row = rows[idx]
+      const offset = offsetExpr
+        ? Number(await evaluateExpr({ node: offsetExpr, row, context }))
+        : 1
+      const target = k + direction * offset
+      if (target >= 0 && target < ordered.length) {
+        output[idx] = await evaluateExpr({ node: valueExpr, row: rows[ordered[target]], context })
+      } else if (defaultExpr) {
+        output[idx] = await evaluateExpr({ node: defaultExpr, row, context })
+      } else {
+        output[idx] = null
+      }
+    }
+    return
+  }
+  throw new Error(`Unsupported window function: ${spec.funcName}`)
 }

package/src/expression/date.js

@@ -100,6 +100,32 @@ export function extractField(field, dateVal) {
   return null
 }
 
+/**
+ * Compute the number of unit boundaries between two dates (end - start).
+ * @param {SqlPrimitive} unit
+ * @param {SqlPrimitive} startVal
+ * @param {SqlPrimitive} endVal
+ * @returns {number | null}
+ */
+export function dateDiff(unit, startVal, endVal) {
+  if (unit == null || startVal == null || endVal == null) return null
+  const start = toDate(startVal)
+  const end = toDate(endVal)
+  if (start == null || end == null) return null
+
+  const u = String(unit).toUpperCase()
+  if (u === 'YEAR') return end.getUTCFullYear() - start.getUTCFullYear()
+  if (u === 'MONTH') {
+    return (end.getUTCFullYear() - start.getUTCFullYear()) * 12 + (end.getUTCMonth() - start.getUTCMonth())
+  }
+  const ms = end.getTime() - start.getTime()
+  if (u === 'DAY') return Math.trunc(ms / 86400000)
+  if (u === 'HOUR') return Math.trunc(ms / 3600000)
+  if (u === 'MINUTE') return Math.trunc(ms / 60000)
+  if (u === 'SECOND') return Math.trunc(ms / 1000)
+  return null
+}
+
 /**
  * @param {SqlPrimitive} val
  * @returns {Date | null}

package/src/expression/evaluate.js

@@ -6,7 +6,7 @@ import { UnknownFunctionError } from '../validation/parseErrors.js'
 import { ColumnNotFoundError } from '../validation/tables.js'
 import { derivedAlias } from './alias.js'
 import { applyBinaryOp } from './binary.js'
-import { applyIntervalToDate, dateTrunc, extractField } from './date.js'
+import { applyIntervalToDate, dateDiff, dateTrunc, extractField } from './date.js'
 import { evaluateMathFunc } from './math.js'
 import { evaluateRegexpFunc } from './regexp.js'
 import { evaluateSpatialFunc } from '../spatial/spatial.js'
@@ -208,6 +208,17 @@ export async function evaluateExpr({ node, row, rowIndex, rows, context }) {
         return count
       }
 
+      if (funcName === 'COUNTIF') {
+        const values = await Promise.all(filteredRows.map(row =>
+          evaluateExpr({ node: argNode, row, context })
+        ))
+        let count = 0
+        for (const v of values) {
+          if (v) count++
+        }
+        return count
+      }
+
       if (funcName === 'SUM' || funcName === 'AVG' || funcName === 'MIN' || funcName === 'MAX') {
         const rawValues = await Promise.all(filteredRows.map(row =>
           evaluateExpr({ node: argNode, row, context })
@@ -406,6 +417,14 @@ export async function evaluateExpr({ node, row, rowIndex, rows, context }) {
       return extractField(args[0], args[1])
     }
 
+    if (funcName === 'EPOCH') {
+      return extractField('EPOCH', args[0])
+    }
+
+    if (funcName === 'DATE_DIFF' || funcName === 'DATEDIFF') {
+      return dateDiff(args[0], args[1], args[2])
+    }
+
     if (funcName === 'CURRENT_DATE') {
       return new Date().toISOString().split('T')[0]
     }
@@ -497,7 +516,7 @@ export async function evaluateExpr({ node, row, rowIndex, rows, context }) {
       return arr.length
     }
 
-    if (funcName === 'ARRAY_LENGTH' || funcName === 'CARDINALITY') {
+    if (funcName === 'ARRAY_LENGTH' || funcName === 'CARDINALITY' || funcName === 'SIZE') {
       const arr = args[0]
       if (!Array.isArray(arr)) return null
       if (funcName === 'ARRAY_LENGTH' && args.length === 2) {
@@ -527,6 +546,12 @@ export async function evaluateExpr({ node, row, rowIndex, rows, context }) {
       return index === -1 ? null : index + 1
     }
 
+    if (funcName === 'ARRAY_CONTAINS') {
+      const [arr, target] = args
+      if (!Array.isArray(arr)) return null
+      return arr.includes(target)
+    }
+
     if (funcName === 'ARRAY_SORT') {
       const arr = args[0]
       if (!Array.isArray(arr)) return null
@@ -612,8 +637,8 @@ export async function evaluateExpr({ node, row, rowIndex, rows, context }) {
       if (typeof val === 'object') return stringify(val)
       return String(val)
     }
-    // Can only cast primitives to other primitive types
-    if (typeof val === 'object') {
+    // Can only cast primitives (and Dates) to other primitive types
+    if (typeof val === 'object' && !(val instanceof Date)) {
       throw new ExecutionError({ message: `Cannot CAST object to ${toType}`, rowIndex, ...node })
     }
     if (toType === 'INTEGER' || toType === 'INT') {

package/src/types.d.ts

@@ -129,7 +129,7 @@ export interface UserDefinedFunction {
   arguments: FunctionSignature
 }
 
-export type AggregateFunc = 'COUNT' | 'SUM' | 'AVG' | 'MIN' | 'MAX' | 'ARRAY_AGG' | 'JSON_ARRAYAGG' | 'STDDEV_SAMP' | 'STDDEV_POP' | 'MEDIAN' | 'PERCENTILE_CONT' | 'APPROX_QUANTILE' | 'STRING_AGG'
+export type AggregateFunc = 'COUNT' | 'COUNTIF' | 'SUM' | 'AVG' | 'MIN' | 'MAX' | 'ARRAY_AGG' | 'JSON_ARRAYAGG' | 'STDDEV_SAMP' | 'STDDEV_POP' | 'MEDIAN' | 'PERCENTILE_CONT' | 'APPROX_QUANTILE' | 'STRING_AGG'
 
 export type RegExpFunction = 'REGEXP_SUBSTR' | 'REGEXP_EXTRACT' | 'REGEXP_REPLACE' | 'REGEXP_MATCHES'
 

package/src/validation/functions.js

@@ -11,7 +11,7 @@ export const niladicFuncs = ['CURRENT_DATE', 'CURRENT_TIME', 'CURRENT_TIMESTAMP'
  * @returns {name is AggregateFunc}
  */
 export function isAggregateFunc(name) {
-  return ['COUNT', 'SUM', 'AVG', 'MIN', 'MAX', 'ARRAY_AGG', 'JSON_ARRAYAGG', 'STDDEV_SAMP', 'STDDEV_POP', 'MEDIAN', 'PERCENTILE_CONT', 'APPROX_QUANTILE', 'STRING_AGG'].includes(name)
+  return ['COUNT', 'COUNTIF', 'SUM', 'AVG', 'MIN', 'MAX', 'ARRAY_AGG', 'JSON_ARRAYAGG', 'STDDEV_SAMP', 'STDDEV_POP', 'MEDIAN', 'PERCENTILE_CONT', 'APPROX_QUANTILE', 'STRING_AGG'].includes(name)
 }
 
 /**
@@ -31,7 +31,7 @@ export function isMathFunc(name) {
  * @returns {boolean}
  */
 export function isWindowFunc(name) {
-  return ['ROW_NUMBER'].includes(name)
+  return ['ROW_NUMBER', 'LAG', 'LEAD'].includes(name)
 }
 
 /**
@@ -140,7 +140,10 @@ export const FUNCTION_SIGNATURES = {
   CURRENT_TIMESTAMP: { min: 0, max: 0, signature: '' },
   DATE_TRUNC: { min: 2, max: 2, signature: 'unit, date' },
   DATE_PART: { min: 2, max: 2, signature: 'field, date' },
+  DATE_DIFF: { min: 3, max: 3, signature: 'unit, start, end' },
+  DATEDIFF: { min: 3, max: 3, signature: 'unit, start, end' },
   EXTRACT: { min: 2, max: 2, signature: 'field FROM date' },
+  EPOCH: { min: 1, max: 1, signature: 'date' },
 
   // Math functions
   FLOOR: { min: 1, max: 1, signature: 'number' },
@@ -182,8 +185,10 @@ export const FUNCTION_SIGNATURES = {
   // Array functions
   ARRAY_LENGTH: { min: 1, max: 2, signature: 'array[, dimension]' },
   ARRAY_POSITION: { min: 2, max: 2, signature: 'array, element' },
+  ARRAY_CONTAINS: { min: 2, max: 2, signature: 'array, element' },
   ARRAY_SORT: { min: 1, max: 1, signature: 'array' },
   CARDINALITY: { min: 1, max: 1, signature: 'array' },
+  SIZE: { min: 1, max: 1, signature: 'array' },
 
   // Table functions (used in FROM clause)
   UNNEST: { min: 1, max: 1, signature: 'array' },
@@ -198,6 +203,7 @@ export const FUNCTION_SIGNATURES = {
 
   // Aggregate functions
   COUNT: { min: 1, max: 1, signature: 'expression' },
+  COUNTIF: { min: 1, max: 1, signature: 'condition' },
   SUM: { min: 1, max: 1, signature: 'expression' },
   AVG: { min: 1, max: 1, signature: 'expression' },
   MIN: { min: 1, max: 1, signature: 'expression' },
@@ -211,6 +217,8 @@ export const FUNCTION_SIGNATURES = {
 
   // Window functions
   ROW_NUMBER: { min: 0, max: 0, signature: '' },
+  LAG: { min: 1, max: 3, signature: 'value[, offset[, default]]' },
+  LEAD: { min: 1, max: 3, signature: 'value[, offset[, default]]' },
 
   // Spatial functions
   ST_INTERSECTS: { min: 2, max: 2, signature: 'geometry, geometry' },