dualfroz/kestrel
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(engine): extend SelectRunner with aggregateWindowValue

Browse Source
This commit is contained in:
dualfroz committed 2026-09-02 16:17:44 +02:00
1 parent 98138f80e9
commit 4b3706184a
1 file changed
+182 -19
Download Patch File Download Diff File Expand all files Collapse all files
src/main/kotlin/kestrel/engine/SelectRunner.kt
+182 -19
View File
@@ -1136,6 +1136,66 @@ class SelectRunner(private val catalog: Catalog) {
}
return ranges
}
private fun aggregateWindowValue(window: WindowExpression, rows: List<List<Value>>): 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<Int>,
combined: List<List<Value>>,
out: Array<Value?>,
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<Value>): 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<Int>, combined: List<List<Value>>, out: Array<Value?>) {
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<Row> {
@@ -1150,30 +1210,82 @@ class SelectRunner(private val catalog: Catalog) {
private fun executeJoin(left: List<List<Value>>, join: JoinClause, rightIndex: Int): List<List<Value>> {
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<List<Value>>()
if (hashKeys != null) {
val buckets = HashMap<Value, MutableList<Row>>()
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<Value, MutableList<IndexedValue<Row>>>()
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<List<Value>>,
join: JoinClause,
lateralSelect: QueryStatement,
nullRight: List<Value>
): List<List<Value>> {
val condition = join.on
val result = mutableListOf<List<Value>>()
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<Value>): 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<List<Value>>): 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<List<Value>> { 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<Value>, List<Value>>()
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)