feat(engine): extend SelectRunner with lateralSchema

src/main/kotlin/kestrel/engine/SelectRunner.kt

@@ -406,9 +406,93 @@ class SelectRunner(private val catalog: Catalog) {
         return sources
     }
 
-    private inner class Plan(val statement: SelectStatement, val sources: List<Source>) {
+    private fun lateralSchema(item: SubquerySource): TableSchema {
+        val labels = structuralColumnNames(item.select)
+        val columns = labels.map { Column(it, ColumnType.Text, notNull = false, primaryKey = false) }
+        return TableSchema(item.alias, columns)
+    }
+
+    /**
+     * The output column names of [query], derived purely from its AST (no execution, no catalog
+     * row data) - needed to give a LATERAL derived table a stable schema before any row of the
+     * source it correlates with is known. Column names never depend on data (aliases and bare
+     * column references are syntactic, and a table/view/CTE's shape is fixed), so this is safe
+     * even though the real rows can only be produced later, per outer row.
+     */
+    private fun structuralColumnNames(query: QueryStatement, localCtes: Map<String, List<String>> = emptyMap()): List<String> =
+        when (query) {
+            is SelectStatement -> when (val projection = query.projection) {
+                is StarProjection ->
+                    structuralSourceColumnNames(query.from, localCtes) +
+                        query.joins.flatMap { structuralSourceColumnNames(it.from, localCtes) }
+                is ExpressionProjection -> projection.items.mapIndexed { index, item -> columnLabel(item, index) }
+            }
+            is SetOperationStatement -> structuralColumnNames(query.left, localCtes)
+            is WithQuery -> {
+                val extended = localCtes.toMutableMap()
+                for (cte in query.ctes) {
+                    extended[cte.name.lowercase()] = cte.columns ?: structuralColumnNames(cte.query, extended)
+                }
+                structuralColumnNames(query.body, extended)
+            }
+        }
+
+    private fun structuralSourceColumnNames(item: FromItem, localCtes: Map<String, List<String>>): List<String> =
+        when (item) {
+            is TableSource -> {
+                val key = item.name.lowercase()
+                localCtes[key]
+                    ?: activeCtes.lastOrNull { it.name.equals(item.name, ignoreCase = true) }?.schema?.columnNames
+                    ?: catalog.tableOrNull(item.name)?.schema?.columnNames
+                    ?: catalog.viewOrNull(item.name)?.let { structuralColumnNames(it.query, localCtes) }
+                    ?: throw KestrelException("no such table: ${item.name}")
+            }
+            is SubquerySource -> structuralColumnNames(item.select, localCtes)
+        }
+
+    private fun sourceLabel(item: FromItem): String = when (item) {
+        is TableSource -> item.alias ?: item.name
+        is SubquerySource -> item.alias
+    }
+
+    private fun resolveSource(item: FromItem): Pair<TableSchema, List<Row>> = when (item) {
+        is TableSource -> resolveNamedSource(item.name)
+        is SubquerySource -> {
+            val result = runQuery(item.select)
+            derivedSchema(item.alias, result) to result.rows.map { Row(it) }
+        }
+    }
+
+    private fun resolveNamedSource(name: String): Pair<TableSchema, List<Row>> {
+        activeCtes.lastOrNull { it.name.equals(name, ignoreCase = true) }?.let { return it.schema to it.rows }
+        catalog.tableOrNull(name)?.let { return it.schema to it.rows }
+        val view = catalog.viewOrNull(name) ?: throw KestrelException("no such table: $name")
+        val key = name.lowercase()
+        if (!expandingViews.add(key)) {
+            throw KestrelException("recursive view definition: $name")
+        }
+        try {
+            val result = runQuery(view.query)
+            return derivedSchema(view.name, result) to result.rows.map { Row(it) }
+        } finally {
+            expandingViews.remove(key)
+        }
+    }
+
+    private fun derivedSchema(name: String, result: SelectResult, labels: List<String> = result.columns): TableSchema {
+        val columns = labels.mapIndexed { index, columnName ->
+            val type = result.rows.asSequence().map { it[index] }.firstOrNull { !it.isNull }?.type ?: ColumnType.Text
+            Column(columnName, type, notNull = false, primaryKey = false)
+        }
+        return TableSchema(name, columns)
+    }
+
+    private class CteScope(val name: String, val schema: TableSchema, var rows: List<Row>)
+
+    private inner class Plan(val statement: SelectStatement, val sources: List<Source>, val outer: ColumnResolver?) {
         val aggregated: Boolean =
             statement.groupBy.isNotEmpty() ||
+                statement.groupingSpec != null ||
                 containsAggregate(statement.projection) ||
                 (statement.having != null && containsAggregate(statement.having))
 
@@ -419,7 +503,8 @@ class SelectRunner(private val catalog: Catalog) {
                 ?.mapNotNull { item -> item.alias?.let { it.lowercase() to item.expression } }
                 ?.toMap() ?: emptyMap()
 
-        private val orderExpression: Expression? = statement.orderBy?.let { resolveOrderAlias(it.expression) }
+        private val orderKeys: List<OrderKey> =
+            statement.orderBy?.keys.orEmpty().map { OrderKey(resolveOrderAlias(it.expression), it.descending, it.nullsFirst) }
 
         private fun resolveOrderAlias(expression: Expression): Expression {
             if (expression is ColumnExpression && expression.qualifier == null) {
@@ -428,7 +513,19 @@ class SelectRunner(private val catalog: Catalog) {
             return expression
         }
 
-        fun resolveIndex(expression: ColumnExpression): Int {
+        /**
+         * Every distinct window-function call in the SELECT list or (post alias-resolution) ORDER
+         * BY - the only two clauses where one is allowed. Computed once per [Plan] and reused both
+         * to append the values ([appendWindowColumns]) and to substitute them back in ([evaluateRow]).
+         */
+        private val windowExpressions: List<WindowExpression> = run {
+            val projectionWindows = (statement.projection as? ExpressionProjection)?.items
+                ?.flatMap { collectWindowExpressions(it.expression) } ?: emptyList()
+            val orderWindows = orderKeys.flatMap { collectWindowExpressions(it.expression) }
+            (projectionWindows + orderWindows).distinct()
+        }
+
+        fun resolveIndexOrNull(expression: ColumnExpression): Int? {
             val qualifier = expression.qualifier
             val matches = mutableListOf<Int>()
             for (source in sources) {
@@ -437,12 +534,32 @@ class SelectRunner(private val catalog: Catalog) {
                 matches.add(source.offset + local)
             }
             return when {
-                matches.isEmpty() -> throw KestrelException("no such column: ${describe(expression)}")
+                matches.isEmpty() -> null
                 matches.size > 1 -> throw KestrelException("ambiguous column reference: ${describe(expression)}")
                 else -> matches.single()
             }
         }
 
+        fun resolveIndex(expression: ColumnExpression): Int =
+            resolveIndexOrNull(expression) ?: throw KestrelException("no such column: ${describe(expression)}")
+
+        /**
+         * [expression]'s declared `COLLATE` (its [kestrel.catalog.Column.collation]) from whichever
+         * [sources] entry it resolves against - `null` when it names no column here (ambiguous
+         * across sources, same as [resolveIndexOrNull], resolves to nothing at all, or the matching
+         * column simply declared no COLLATE of its own).
+         */
+        fun declaredCollationOrNull(expression: ColumnExpression): Collation? {
+            val qualifier = expression.qualifier
+            val matches = mutableListOf<Collation>()
+            for (source in sources) {
+                if (qualifier != null && !source.alias.equals(qualifier, ignoreCase = true)) continue
+                val column = source.schema.column(expression.name) ?: continue
+                matches.add(column.collation)
+            }
+            return matches.singleOrNull()
+        }
+
         private fun sourceOfGlobal(global: Int): Int =
             sources.indexOfLast { global >= it.offset }
 
@@ -450,25 +567,89 @@ class SelectRunner(private val catalog: Catalog) {
             collectColumns(expression).map { sourceOfGlobal(resolveIndex(it)) }.toSet()
 
         fun validate() {
-            statement.where?.let { validateColumns(it) }
-            orderExpression?.let { validateColumns(it) }
-            statement.groupBy.forEach { validateColumns(it) }
-            statement.having?.let { validateColumns(it) }
+            statement.where?.let { validateColumns(it); rejectWindowFunction(it, "WHERE") }
+            orderKeys.forEach { validateColumns(it.expression) }
+            statement.groupBy.forEach { validateColumns(it); rejectWindowFunction(it, "GROUP BY") }
+            statement.having?.let { validateColumns(it); rejectWindowFunction(it, "HAVING") }
             if (statement.projection is ExpressionProjection) {
                 statement.projection.items.forEach { validateColumns(it.expression) }
             }
             statement.joins.forEachIndexed { position, join ->
-                validateColumns(join.on)
+                // CROSS JOIN has no ON condition at all (its `on` is null) - nothing to validate.
+                val on = join.on ?: return@forEachIndexed
+                validateColumns(on)
+                rejectWindowFunction(on, "an ON condition")
                 val allowed = position + 1
-                if (referencedSources(join.on).any { it > allowed }) {
-                    throw KestrelException("ON condition for ${join.tableName} references a table joined later")
+                if (referencedSources(on).any { it > allowed }) {
+                    throw KestrelException("ON condition for ${sourceLabel(join.from)} references a table joined later")
                 }
             }
             if (aggregated) validateAggregation()
+            if (windowExpressions.isNotEmpty()) validateWindows()
+            if (statement.distinctOn.isNotEmpty()) validateDistinctOn()
+        }
+
+        /**
+         * `DISTINCT ON (...)` has no unsurprising meaning once rows are already collapsed into
+         * aggregate groups - rejected outright rather than guessing which pre-aggregate row a group
+         * should be "first" for. The leading-`ORDER BY` requirement mirrors Postgres: `DISTINCT ON`
+         * exists to pick a deterministic row per group, which only means something if `ORDER BY`
+         * (when present at all) agrees on what "first" means for those same expressions.
+         */
+        private fun validateDistinctOn() {
+            statement.distinctOn.forEach { validateColumns(it); rejectWindowFunction(it, "DISTINCT ON") }
+            if (aggregated) {
+                throw KestrelException("DISTINCT ON cannot be combined with GROUP BY or aggregate functions")
+            }
+            if (statement.orderBy != null) {
+                val leading = orderKeys.take(statement.distinctOn.size).map { it.expression }
+                if (leading != statement.distinctOn) {
+                    throw KestrelException("SELECT DISTINCT ON expressions must match initial ORDER BY expressions")
+                }
+            }
+        }
+
+        private fun rejectWindowFunction(expression: Expression, clause: String) {
+            if (containsWindowFunction(expression)) {
+                throw KestrelException("window functions are not allowed in $clause")
+            }
+        }
+
+        /**
+         * Window functions are rejected outright alongside GROUP BY/aggregates ([aggregated]) -
+         * windowing over an already-grouped result is valid SQL but requires running GroupAggregation
+         * first and then a second pass over its output rows, which this engine does not implement.
+         * RANK/DENSE_RANK need an explicit tie-break (unlike ROW_NUMBER, which falls back to this
+         * partition's row arrival order when OVER has no ORDER BY) and nesting one window function
+         * inside another's argument/PARTITION BY/ORDER BY is invalid SQL in the first place.
+         */
+        private fun validateWindows() {
+            if (aggregated) {
+                throw KestrelException("window functions cannot be combined with GROUP BY or aggregate functions in the same query")
+            }
+            for (window in windowExpressions) {
+                if ((window.kind == WindowFunctionKind.Rank || window.kind == WindowFunctionKind.DenseRank) && window.orderBy.isEmpty()) {
+                    throw KestrelException("${window.name}() OVER (...) requires an ORDER BY clause")
+                }
+                if (window.frame != null && window.kind !in FRAME_CAPABLE_KINDS) {
+                    throw KestrelException("${window.name}() cannot be used with a ROWS/RANGE/GROUPS frame clause")
+                }
+                val nested = listOfNotNull(
+                    window.argument, window.offset, window.defaultValue,
+                    window.frame?.start?.offset, window.frame?.end?.offset
+                ) + window.partitionBy + window.orderBy.map { it.expression }
+                if (nested.any { containsWindowFunction(it) }) {
+                    throw KestrelException("window functions cannot be nested")
+                }
+            }
         }
 
         private fun validateColumns(expression: Expression) {
-            collectColumns(expression).forEach { resolveIndex(it) }
+            collectColumns(expression).forEach { column ->
+                if (resolveIndexOrNull(column) == null && outer == null) {
+                    throw KestrelException("no such column: ${describe(column)}")
+                }
+            }
         }
 
         private fun validateAggregation() {