From 4b3706184aa3b09ca6ae04096af4069366cc0f75 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Wed, 2 Sep 2026 16:17:44 +0200 Subject: [PATCH] feat(engine): extend SelectRunner with aggregateWindowValue --- .../kotlin/kestrel/engine/SelectRunner.kt | 201 ++++++++++++++++-- 1 file changed, 182 insertions(+), 19 deletions(-) diff --git a/src/main/kotlin/kestrel/engine/SelectRunner.kt b/src/main/kotlin/kestrel/engine/SelectRunner.kt index defb5bf..b29c66c 100644 --- a/src/main/kotlin/kestrel/engine/SelectRunner.kt +++ b/src/main/kotlin/kestrel/engine/SelectRunner.kt @@ -1136,6 +1136,66 @@ class SelectRunner(private val catalog: Catalog) { } return ranges } + + private fun aggregateWindowValue(window: WindowExpression, rows: List>): Value = + AggregateFunctions.compute(window.name, window.argument, window.star, rows, ::evaluateRow) + + /** + * LAG (direction -1) / LEAD (direction +1): reads [WindowExpression.argument] from the row + * [WindowExpression.offset] positions before/after the current one in [ordered] (default 1, + * evaluated per row so it may itself be a computed expression, though in practice it is + * almost always a literal). Past either edge of the partition, falls back to + * [WindowExpression.defaultValue] (evaluated against the *current* row) or NULL. + */ + private fun assignOffsetWindowValues( + window: WindowExpression, + ordered: List, + combined: List>, + out: Array, + direction: Int + ) { + ordered.forEachIndexed { position, index -> + val offset = resolveOffset(window, combined[index]) + val targetPosition = position + direction * offset + out[index] = if (targetPosition in ordered.indices) { + evaluateRow(window.argument!!, combined[ordered[targetPosition]]) + } else { + window.defaultValue?.let { evaluateRow(it, combined[index]) } ?: Value.NullValue + } + } + } + + private fun resolveOffset(window: WindowExpression, row: List): Int { + val expression = window.offset ?: return 1 + val value = evaluateRow(expression, row) + val amount = (value as? Value.IntegerValue)?.value + ?: throw KestrelException("${window.name}() offset must be an integer") + if (amount < 0) throw KestrelException("${window.name}() offset must not be negative") + return amount.toInt() + } + + /** + * NTILE(n): splits the partition into `n` buckets as evenly as possible - when the partition + * size does not divide evenly, the first `size % n` buckets get one extra row, exactly like + * every other SQL engine's NTILE. `n` larger than the partition just leaves the highest + * bucket numbers unused, rather than being an error. + */ + private fun assignNtileWindowValues(window: WindowExpression, ordered: List, combined: List>, out: Array) { + val bucketValue = evaluateRow(window.argument!!, combined[ordered.first()]) + val buckets = (bucketValue as? Value.IntegerValue)?.value + ?: throw KestrelException("${window.name}() bucket count must be an integer") + if (buckets <= 0) throw KestrelException("${window.name}() bucket count must be positive") + val bucketCount = buckets.toInt() + val base = ordered.size / bucketCount + val remainder = ordered.size % bucketCount + var position = 0 + for (bucket in 1..bucketCount) { + val size = base + if (bucket <= remainder) 1 else 0 + repeat(size) { + out[ordered[position]] = Value.IntegerValue(bucket.toLong()) + position++ + } + } } private fun primaryRows(): List { @@ -1150,30 +1210,82 @@ class SelectRunner(private val catalog: Catalog) { private fun executeJoin(left: List>, join: JoinClause, rightIndex: Int): List> { val right = sources[rightIndex] val nullRight = List(right.schema.columns.size) { Value.NullValue } + // Width of everything already combined to the left of this join - exactly what a RIGHT/ + // FULL JOIN needs to pad an unmatched right row with, since `right.offset` is defined as + // the running column count of every source before it. + val nullLeft = List(right.offset) { Value.NullValue } + val outerRight = join.type == JoinType.Right || join.type == JoinType.Full + val outerLeft = join.type == JoinType.Left || join.type == JoinType.Full + val lateralSelect = lateralSelectOf(join.from) + if (lateralSelect != null) return executeLateralJoin(left, join, lateralSelect, nullRight) val hashKeys = hashJoinKeys(join, rightIndex) val result = mutableListOf>() if (hashKeys != null) { - val buckets = HashMap>() - for (rightRow in right.rows) { - val key = evaluateOnRight(hashKeys.buildExpression, right, rightRow) - if (key !is Value.NullValue) buckets.getOrPut(key) { mutableListOf() }.add(rightRow) + val buckets = HashMap>>() + for (indexedRow in right.rows.withIndex()) { + val key = foldHashJoinKey(evaluateOnRight(hashKeys.buildExpression, right, indexedRow.value), hashKeys.collation) + if (key !is Value.NullValue) buckets.getOrPut(key) { mutableListOf() }.add(indexedRow) } + val matchedRight = BooleanArray(right.rows.size) for (leftRow in left) { - val probe = evaluateRow(hashKeys.probeExpression, leftRow) + val probe = foldHashJoinKey(evaluateRow(hashKeys.probeExpression, leftRow), hashKeys.collation) val matches = if (probe is Value.NullValue) emptyList() else buckets[probe].orEmpty() if (matches.isNotEmpty()) { - matches.forEach { result.add(leftRow + it.values) } - } else if (join.type == JoinType.Left) { + matches.forEach { (index, row) -> result.add(leftRow + row.values); matchedRight[index] = true } + } else if (outerLeft) { result.add(leftRow + nullRight) } } + if (outerRight) { + right.rows.forEachIndexed { index, row -> if (!matchedRight[index]) result.add(nullLeft + row.values) } + } return result } + val condition = join.on + val matchedRight = BooleanArray(right.rows.size) for (leftRow in left) { var matched = false - for (rightRow in right.rows) { + right.rows.forEachIndexed { index, rightRow -> val candidate = leftRow + rightRow.values - if (isTrue(evaluateRow(join.on, candidate))) { + if (condition == null || isTrue(evaluateRow(condition, candidate))) { + result.add(candidate) + matched = true + matchedRight[index] = true + } + } + if (!matched && outerLeft) result.add(leftRow + nullRight) + } + if (outerRight) { + right.rows.forEachIndexed { index, row -> if (!matchedRight[index]) result.add(nullLeft + row.values) } + } + return result + } + + private fun lateralSelectOf(item: FromItem): QueryStatement? = + (item as? SubquerySource)?.takeIf { it.lateral }?.select + + /** + * Unlike every other join, a LATERAL derived table has no fixed row set: it is + * re-evaluated once per row already accumulated on the left, correlated to that row via + * [lateralOuterResolver] exactly like a correlated subquery in WHERE/SELECT (same + * [SubqueryRunner] plumbing, just invoked from FROM instead). Hash-joining is therefore + * meaningless here (there is no static right-hand row set to build buckets from) - always + * nested loop, matching how real engines execute LATERAL. + */ + private fun executeLateralJoin( + left: List>, + join: JoinClause, + lateralSelect: QueryStatement, + nullRight: List + ): List> { + val condition = join.on + val result = mutableListOf>() + for (leftRow in left) { + val rightRows = runQuery(lateralSelect, lateralOuterResolver(leftRow)).rows + var matched = false + for (rightRow in rightRows) { + val candidate = leftRow + rightRow + if (condition == null || isTrue(evaluateRow(condition, candidate))) { result.add(candidate) matched = true } @@ -1183,6 +1295,27 @@ class SelectRunner(private val catalog: Catalog) { return result } + /** + * Resolves a column against the sources already joined by the time a LATERAL subquery + * runs (everything to its left in the FROM clause) - [leftRow] holds exactly their values, + * so a resolved global index landing past its end means the expression names a source that + * has not been joined yet (or is the LATERAL item itself), which SQL forbids. + */ + private fun lateralOuterResolver(leftRow: List): kestrel.engine.ColumnResolver = object : kestrel.engine.ColumnResolver { + override fun resolve(expression: ColumnExpression): Value { + val index = resolveIndex(expression) + if (index >= leftRow.size) { + throw KestrelException( + "LATERAL subquery cannot reference ${describe(expression)}: not joined yet at this point in the FROM clause" + ) + } + return leftRow[index] + } + + override fun declaredCollation(expression: ColumnExpression): Collation? = + resolveIndexOrNull(expression)?.takeIf { it < leftRow.size }?.let { declaredCollationOrNull(expression) } + } + private fun evaluateOnRight(expression: Expression, right: Source, rightRow: Row): Value { val resolver = kestrel.engine.ColumnResolver { columnExpression -> rightRow.values[resolveIndex(columnExpression) - right.offset] @@ -1190,17 +1323,35 @@ class SelectRunner(private val catalog: Catalog) { return Evaluator(resolver).evaluate(expression) } + /** + * A [Value.TextValue] hash-join key under any collation but [Collation.Binary] must land in + * the same bucket as every other spelling the declared/explicit `COLLATE` treats as equal + * (e.g. `'Bob'`/`'bob'` under `NOCASE`) - otherwise the fast path silently disagrees with + * the nested-loop `Evaluator` path it is an optimization of, which already resolves this via + * [pairCollation] for every non-hashable equality condition. [Collation.Binary]'s + * [Collation.canonicalize] is the identity, so a key from an equality with no COLLATE + * anywhere in play keeps its exact original `equals`/`hashCode`. + */ + private fun foldHashJoinKey(value: Value, collation: Collation): Value = + if (value is Value.TextValue) Value.TextValue(collation.canonicalize(value.value)) else value + private fun hashJoinKeys(join: JoinClause, rightIndex: Int): HashKeys? { + if (lateralSelectOf(join.from) != null) return null val condition = join.on if (condition !is BinaryExpression || condition.operator != BinaryOperator.Equal) return null val built = (0 until rightIndex).toSet() val leftRefs = referencedSources(condition.left) val rightRefs = referencedSources(condition.right) + // Resolved from the ORIGINAL (unswapped) left/right of `a = b`, not buildExpression/ + // probeExpression below, so `pairCollation`'s left-first tie-break matches exactly what + // evaluating this same `=` through Evaluator would have decided regardless of which side + // ends up building the hash buckets. + val collation = pairCollation(condition.left, condition.right, ::declaredCollationOrNull) return when { rightRefs == setOf(rightIndex) && leftRefs.all { it in built } -> - HashKeys(buildExpression = condition.right, probeExpression = condition.left) + HashKeys(buildExpression = condition.right, probeExpression = condition.left, collation = collation) leftRefs == setOf(rightIndex) && rightRefs.all { it in built } -> - HashKeys(buildExpression = condition.left, probeExpression = condition.right) + HashKeys(buildExpression = condition.left, probeExpression = condition.right, collation = collation) else -> null } } @@ -1212,16 +1363,28 @@ class SelectRunner(private val catalog: Catalog) { } fun projectRows(combined: List>): SelectResult { + statement.limit?.let { if (it < 0) throw KestrelException("LIMIT must not be negative") } + statement.offset?.let { if (it < 0) throw KestrelException("OFFSET must not be negative") } var rows = combined - orderExpression?.let { expression -> - val comparator = Comparator> { first, second -> - ValueOrdering.compare(evaluateRow(expression, first), evaluateRow(expression, second)) - } - rows = rows.sortedWith(if (statement.orderBy!!.descending) comparator.reversed() else comparator) + if (orderKeys.isNotEmpty()) { + rows = rows.sortedWith(multiKeyComparator(orderKeys, ::evaluateRow, ::declaredCollationOrNull)) } - statement.limit?.let { - if (it < 0) throw KestrelException("LIMIT must not be negative") - rows = rows.take(it) + if (statement.distinctOn.isNotEmpty()) { + // Kept in the order already established above (scan order with no ORDER BY, sorted + // order otherwise) - `LinkedHashMap.putIfAbsent` keeps the first row per key and + // never disturbs where that first row sits relative to the others kept. + val kept = LinkedHashMap, List>() + for (row in rows) { + val key = statement.distinctOn.map { evaluateRow(it, row) } + kept.putIfAbsent(collationFoldedKey(key, statement.distinctOn, ::declaredCollationOrNull), row) + } + rows = kept.values.toList() + } + if (!statement.distinct) { + if (statement.offset != null || statement.limit != null) { + rows = rows.drop(statement.offset ?: 0) + statement.limit?.let { rows = rows.take(it) } + } } return when (val projection = statement.projection) { is StarProjection -> SelectResult(columnNames, rows)