feat(engine): extend GroupAggregation with LogicalGroup and aggregate
This commit is contained in:
dualfroz committed
1 parent
be312cda54
commit
4653c74cf9
1 file changed
+158
-24
@@ -1,54 +1,188 @@
|
||||
package kestrel.engine
|
||||
|
||||
import kestrel.KestrelException
|
||||
import kestrel.sql.ast.BetweenExpression
|
||||
import kestrel.sql.ast.BinaryExpression
|
||||
import kestrel.sql.ast.CallExpression
|
||||
import kestrel.sql.ast.CaseBranch
|
||||
import kestrel.sql.ast.CaseExpression
|
||||
import kestrel.sql.ast.CastExpression
|
||||
import kestrel.sql.ast.Expression
|
||||
import kestrel.sql.ast.ExpressionProjection
|
||||
import kestrel.sql.ast.FunctionExpression
|
||||
import kestrel.sql.ast.InExpression
|
||||
import kestrel.sql.ast.IsNullExpression
|
||||
import kestrel.sql.ast.LikeExpression
|
||||
import kestrel.sql.ast.LiteralExpression
|
||||
import kestrel.sql.ast.OrderKey
|
||||
import kestrel.sql.ast.SelectItem
|
||||
import kestrel.sql.ast.SelectStatement
|
||||
import kestrel.sql.ast.UnaryExpression
|
||||
import kestrel.storage.Collation
|
||||
import kestrel.storage.Value
|
||||
import kestrel.storage.ValueOrdering
|
||||
import kestrel.storage.typeName
|
||||
|
||||
class GroupAggregation(
|
||||
private val resolveIndex: (kestrel.sql.ast.ColumnExpression) -> Int,
|
||||
private val evaluateRow: (Expression, List<Value>) -> Value,
|
||||
private val label: (SelectItem, Int) -> String
|
||||
private val label: (SelectItem, Int) -> String,
|
||||
private val subqueries: SubqueryRunner? = null,
|
||||
/** A bare grouping column's declared `COLLATE` (see [kestrel.catalog.Column.collation]) - defaults to reporting none, matching every [ColumnResolver] built before this existed. */
|
||||
private val declaredCollationOf: (kestrel.sql.ast.ColumnExpression) -> Collation? = { null }
|
||||
) {
|
||||
fun aggregate(statement: SelectStatement, orderExpression: Expression?, combined: List<List<Value>>): SelectResult {
|
||||
val groups = LinkedHashMap<List<Value>, MutableList<List<Value>>>()
|
||||
for (row in combined) {
|
||||
val key = statement.groupBy.map { evaluateRow(it, row) }
|
||||
groups.getOrPut(key) { mutableListOf() }.add(row)
|
||||
}
|
||||
var groupRows: List<List<List<Value>>> =
|
||||
if (statement.groupBy.isEmpty() && groups.isEmpty()) listOf(emptyList()) else groups.values.toList()
|
||||
/**
|
||||
* One row of the eventual result, before projection: the underlying source rows that fell into
|
||||
* it, and the grouping-set expressions that were "active" (present, not rolled up away) for it.
|
||||
* For an ordinary `GROUP BY` (no [kestrel.sql.ast.GroupingSpec]) there is exactly one grouping set
|
||||
* per query - [activeExpressions] is always the full `GROUP BY` list, so every column reads
|
||||
* normally and `GROUPING(...)` (see [substituteGroupingAwareCalls]) always answers 0, matching
|
||||
* plain SQL. [rows] is empty only for the grand-total group of a set with zero expressions over a
|
||||
* zero-row input - the same "ungrouped aggregate over nothing" case a bare `SELECT COUNT(*)` with
|
||||
* no rows at all already has to handle.
|
||||
*/
|
||||
private data class LogicalGroup(val rows: List<List<Value>>, val activeExpressions: List<Expression>)
|
||||
|
||||
fun aggregate(statement: SelectStatement, orderKeys: List<OrderKey>, combined: List<List<Value>>): SelectResult {
|
||||
val allGroupingColumns = statement.groupBy
|
||||
val sets = GroupingSetExpansion.setsFor(statement.groupingSpec, statement.groupBy)
|
||||
var logicalGroups = sets.flatMap { groupsForSet(it, combined) }
|
||||
statement.having?.let { having ->
|
||||
groupRows = groupRows.filter { isTrue(evaluateInGroup(having, it)) }
|
||||
logicalGroups = logicalGroups.filter { isTrue(evaluateInGroup(having, it, allGroupingColumns)) }
|
||||
}
|
||||
orderExpression?.let { expression ->
|
||||
val comparator = Comparator<List<List<Value>>> { first, second ->
|
||||
ValueOrdering.compare(evaluateInGroup(expression, first), evaluateInGroup(expression, second))
|
||||
if (orderKeys.isNotEmpty()) {
|
||||
val comparator = Comparator<LogicalGroup> { first, second ->
|
||||
var result = 0
|
||||
for (key in orderKeys) {
|
||||
val comparison = compareOrderKeyValues(
|
||||
evaluateInGroup(key.expression, first, allGroupingColumns),
|
||||
evaluateInGroup(key.expression, second, allGroupingColumns),
|
||||
key,
|
||||
declaredCollationOf
|
||||
)
|
||||
if (comparison != 0) {
|
||||
result = comparison
|
||||
break
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
groupRows = groupRows.sortedWith(if (statement.orderBy!!.descending) comparator.reversed() else comparator)
|
||||
logicalGroups = logicalGroups.sortedWith(comparator)
|
||||
}
|
||||
statement.limit?.let {
|
||||
if (it < 0) throw KestrelException("LIMIT must not be negative")
|
||||
groupRows = groupRows.take(it)
|
||||
statement.offset?.let { if (it < 0) throw KestrelException("OFFSET must not be negative") }
|
||||
statement.limit?.let { if (it < 0) throw KestrelException("LIMIT must not be negative") }
|
||||
// DISTINCT can still collapse two separate groups into the same output row (e.g.
|
||||
// `SELECT DISTINCT dept FROM t GROUP BY dept, sub_dept`) - LIMIT/OFFSET are deferred past
|
||||
// dedup below so they count distinct output rows, not pre-dedup groups.
|
||||
if (!statement.distinct && (statement.offset != null || statement.limit != null)) {
|
||||
logicalGroups = logicalGroups.drop(statement.offset ?: 0)
|
||||
statement.limit?.let { logicalGroups = logicalGroups.take(it) }
|
||||
}
|
||||
val projection = statement.projection as ExpressionProjection
|
||||
val labels = projection.items.mapIndexed { index, item -> label(item, index) }
|
||||
val projected = groupRows.map { group -> projection.items.map { evaluateInGroup(it.expression, group) } }
|
||||
var projected = logicalGroups.map { group -> projection.items.map { evaluateInGroup(it.expression, group, allGroupingColumns) } }
|
||||
if (statement.distinct) {
|
||||
val itemExpressions = projection.items.map { it.expression }
|
||||
projected = projected.distinctBy { row -> collationFoldedKey(row, itemExpressions, declaredCollationOf) }
|
||||
if (statement.offset != null || statement.limit != null) {
|
||||
projected = projected.drop(statement.offset ?: 0)
|
||||
statement.limit?.let { projected = projected.take(it) }
|
||||
}
|
||||
}
|
||||
return SelectResult(labels, projected)
|
||||
}
|
||||
|
||||
private fun evaluateInGroup(expression: Expression, group: List<List<Value>>): Value {
|
||||
val substituted = substituteAggregates(expression, group)
|
||||
val resolver = ColumnResolver { columnExpression -> group.first()[resolveIndex(columnExpression)] }
|
||||
return Evaluator(resolver).evaluate(substituted)
|
||||
/**
|
||||
* Groups [combined] by [setExpressions] alone (not the query's full grouping-column union) -
|
||||
* exactly the single-set logic this class always had, just parameterized per set instead of
|
||||
* always using `statement.groupBy`. A set with zero expressions over zero source rows still
|
||||
* produces the one grand-total group (`combined.isEmpty()` means the inner loop never ran, so
|
||||
* `buckets` is genuinely empty rather than holding a single `emptyList()`-keyed bucket) - every
|
||||
* *other* set, or this same set over a non-empty input, groups normally.
|
||||
*/
|
||||
private fun groupsForSet(setExpressions: List<Expression>, combined: List<List<Value>>): List<LogicalGroup> {
|
||||
val buckets = LinkedHashMap<List<Value>, MutableList<List<Value>>>()
|
||||
for (row in combined) {
|
||||
val key = setExpressions.map { evaluateRow(it, row) }
|
||||
val bucketKey = collationFoldedKey(key, setExpressions, declaredCollationOf)
|
||||
buckets.getOrPut(bucketKey) { mutableListOf() }.add(row)
|
||||
}
|
||||
return if (setExpressions.isEmpty() && buckets.isEmpty()) {
|
||||
listOf(LogicalGroup(emptyList(), setExpressions))
|
||||
} else {
|
||||
buckets.values.map { LogicalGroup(it, setExpressions) }
|
||||
}
|
||||
}
|
||||
|
||||
private fun evaluateInGroup(expression: Expression, group: LogicalGroup, allGroupingColumns: List<Expression>): Value {
|
||||
val nulled = substituteInactiveGroupingColumns(expression, group.activeExpressions, allGroupingColumns)
|
||||
val substituted = substituteAggregates(nulled, group.rows)
|
||||
// Only reached for a bare reference to one of `group.activeExpressions` (anything else was
|
||||
// just replaced with NULL above), and a group with any active expression at all is never
|
||||
// empty - it was built from at least one matching source row - so `.first()` cannot throw.
|
||||
val resolver = object : ColumnResolver {
|
||||
override fun resolve(columnExpression: kestrel.sql.ast.ColumnExpression): Value =
|
||||
group.rows.first()[resolveIndex(columnExpression)]
|
||||
|
||||
override fun declaredCollation(columnExpression: kestrel.sql.ast.ColumnExpression): Collation? =
|
||||
declaredCollationOf(columnExpression)
|
||||
}
|
||||
return Evaluator(resolver, subqueries).evaluate(substituted)
|
||||
}
|
||||
|
||||
/**
|
||||
* Replaces every grouping-column expression not in [activeExpressions] with a NULL literal (SQL's
|
||||
* rule for a ROLLUP/CUBE/GROUPING SETS super-aggregate row: a collapsed column reads as NULL
|
||||
* wherever it is used outside an aggregate's own argument) and every `GROUPING(...)` call with its
|
||||
* literal bitmask. Stops at a [FunctionExpression] without recursing into its argument - an
|
||||
* aggregate always sees real per-row values regardless of which columns this particular row
|
||||
* collapsed, which is exactly what makes `SUM(amount)` correct on the grand-total row too.
|
||||
*/
|
||||
private fun substituteInactiveGroupingColumns(
|
||||
expression: Expression,
|
||||
activeExpressions: List<Expression>,
|
||||
allGroupingColumns: List<Expression>
|
||||
): Expression {
|
||||
if (expression in allGroupingColumns && expression !in activeExpressions) {
|
||||
return LiteralExpression(Value.NullValue)
|
||||
}
|
||||
fun recurse(child: Expression) = substituteInactiveGroupingColumns(child, activeExpressions, allGroupingColumns)
|
||||
return when (expression) {
|
||||
is BinaryExpression -> BinaryExpression(expression.operator, recurse(expression.left), recurse(expression.right))
|
||||
is UnaryExpression -> UnaryExpression(expression.operator, recurse(expression.operand))
|
||||
is CallExpression ->
|
||||
if (expression.name.equals("GROUPING", ignoreCase = true)) {
|
||||
LiteralExpression(computeGrouping(expression, activeExpressions))
|
||||
} else {
|
||||
CallExpression(expression.name, expression.arguments.map(::recurse))
|
||||
}
|
||||
is IsNullExpression -> IsNullExpression(recurse(expression.operand), expression.negated)
|
||||
is InExpression -> InExpression(recurse(expression.operand), expression.values.map(::recurse), expression.negated)
|
||||
is BetweenExpression -> BetweenExpression(
|
||||
recurse(expression.operand), recurse(expression.low), recurse(expression.high), expression.negated
|
||||
)
|
||||
is LikeExpression -> LikeExpression(recurse(expression.operand), recurse(expression.pattern), expression.negated)
|
||||
is CaseExpression -> CaseExpression(
|
||||
expression.operand?.let(::recurse),
|
||||
expression.branches.map { CaseBranch(recurse(it.condition), recurse(it.result)) },
|
||||
expression.elseExpression?.let(::recurse)
|
||||
)
|
||||
is CastExpression -> CastExpression(recurse(expression.operand), expression.targetType, expression.precision, expression.scale)
|
||||
else -> expression
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* `GROUPING(e1, ..., en)` per the SQL standard: an integer whose bits, most significant first, are
|
||||
* 1 for each `ei` collapsed away in this row's grouping set and 0 for each one still present -
|
||||
* `computeAggregate`'s own arguments never reach this ([FunctionExpression] is not
|
||||
* [CallExpression]), so there is no ambiguity with a real aggregate named similarly.
|
||||
*/
|
||||
private fun computeGrouping(call: CallExpression, activeExpressions: List<Expression>): Value {
|
||||
if (call.arguments.isEmpty()) throw KestrelException("GROUPING requires at least one argument")
|
||||
var bits = 0L
|
||||
for (argument in call.arguments) {
|
||||
bits = (bits shl 1) or if (argument in activeExpressions) 0L else 1L
|
||||
}
|
||||
return Value.IntegerValue(bits)
|
||||
}
|
||||
|
||||
private fun substituteAggregates(expression: Expression, group: List<List<Value>>): Expression = when (expression) {
|
||||
|
||||
Reference in new issue
Block a user