feat(sql): add containsAggregate and validateDistinctOn to DiskExecutor
This commit is contained in:
dualfroz committed
1 parent
bb57f5ac91
commit
de2cc8854d
1 file changed
+173
@@ -3187,3 +3187,176 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess {
|
|||||||
is ExpressionProjection -> projection.items.any { containsAggregate(it.expression) }
|
is ExpressionProjection -> projection.items.any { containsAggregate(it.expression) }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private fun containsAggregate(expression: Expression): Boolean = when (expression) {
|
||||||
|
is FunctionExpression -> true
|
||||||
|
is BinaryExpression -> containsAggregate(expression.left) || containsAggregate(expression.right)
|
||||||
|
is UnaryExpression -> containsAggregate(expression.operand)
|
||||||
|
is CallExpression -> expression.arguments.any { containsAggregate(it) }
|
||||||
|
is IsNullExpression -> containsAggregate(expression.operand)
|
||||||
|
is InExpression -> containsAggregate(expression.operand) || expression.values.any { containsAggregate(it) }
|
||||||
|
is BetweenExpression ->
|
||||||
|
containsAggregate(expression.operand) || containsAggregate(expression.low) || containsAggregate(expression.high)
|
||||||
|
is LikeExpression -> containsAggregate(expression.operand) || containsAggregate(expression.pattern)
|
||||||
|
is CaseExpression ->
|
||||||
|
(expression.operand?.let { containsAggregate(it) } ?: false) ||
|
||||||
|
expression.branches.any { containsAggregate(it.condition) || containsAggregate(it.result) } ||
|
||||||
|
(expression.elseExpression?.let { containsAggregate(it) } ?: false)
|
||||||
|
is CastExpression -> containsAggregate(expression.operand)
|
||||||
|
is CollateExpression -> containsAggregate(expression.operand)
|
||||||
|
else -> false
|
||||||
|
}
|
||||||
|
|
||||||
|
/** See [kestrel.engine.SelectRunner]'s `Plan.validateDistinctOn` for why both checks exist. */
|
||||||
|
private fun validateDistinctOn(statement: SelectStatement, orderKeys: List<OrderKey>, aggregated: Boolean) {
|
||||||
|
if (aggregated) {
|
||||||
|
throw DiskExecutorException("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 DiskExecutorException("SELECT DISTINCT ON expressions must match initial ORDER BY expressions")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun validateAggregation(statement: SelectStatement, orderKeys: List<OrderKey>, sources: List<JoinSource>) {
|
||||||
|
val groupColumns = statement.groupBy.filterIsInstance<ColumnExpression>()
|
||||||
|
.map { resolveIndex(sources, it) }.toSet()
|
||||||
|
if (statement.projection is ExpressionProjection) {
|
||||||
|
statement.projection.items.forEach { requireGrouped(it.expression, statement.groupBy, groupColumns, sources) }
|
||||||
|
}
|
||||||
|
statement.having?.let { requireGrouped(it, statement.groupBy, groupColumns, sources) }
|
||||||
|
orderKeys.forEach { requireGrouped(it.expression, statement.groupBy, groupColumns, sources) }
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun requireGrouped(expression: Expression, groupBy: List<Expression>, groupColumns: Set<Int>, sources: List<JoinSource>) {
|
||||||
|
if (expression in groupBy) return
|
||||||
|
when (expression) {
|
||||||
|
is FunctionExpression -> Unit
|
||||||
|
is ColumnExpression -> if (resolveIndex(sources, expression) !in groupColumns) {
|
||||||
|
throw DiskExecutorException(
|
||||||
|
"column ${expression.name} must appear in GROUP BY or be used in an aggregate function"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
is BinaryExpression -> {
|
||||||
|
requireGrouped(expression.left, groupBy, groupColumns, sources)
|
||||||
|
requireGrouped(expression.right, groupBy, groupColumns, sources)
|
||||||
|
}
|
||||||
|
is UnaryExpression -> requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
is CallExpression -> expression.arguments.forEach { requireGrouped(it, groupBy, groupColumns, sources) }
|
||||||
|
is IsNullExpression -> requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
is InExpression -> {
|
||||||
|
requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
expression.values.forEach { requireGrouped(it, groupBy, groupColumns, sources) }
|
||||||
|
}
|
||||||
|
is BetweenExpression -> {
|
||||||
|
requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
requireGrouped(expression.low, groupBy, groupColumns, sources)
|
||||||
|
requireGrouped(expression.high, groupBy, groupColumns, sources)
|
||||||
|
}
|
||||||
|
is LikeExpression -> {
|
||||||
|
requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
requireGrouped(expression.pattern, groupBy, groupColumns, sources)
|
||||||
|
}
|
||||||
|
is CaseExpression -> {
|
||||||
|
expression.operand?.let { requireGrouped(it, groupBy, groupColumns, sources) }
|
||||||
|
expression.branches.forEach {
|
||||||
|
requireGrouped(it.condition, groupBy, groupColumns, sources)
|
||||||
|
requireGrouped(it.result, groupBy, groupColumns, sources)
|
||||||
|
}
|
||||||
|
expression.elseExpression?.let { requireGrouped(it, groupBy, groupColumns, sources) }
|
||||||
|
}
|
||||||
|
is CastExpression -> requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
is CollateExpression -> requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
is ScalarSubqueryExpression -> Unit
|
||||||
|
is ExistsExpression -> Unit
|
||||||
|
is InSubqueryExpression -> requireGrouped(expression.operand, groupBy, groupColumns, sources)
|
||||||
|
else -> Unit
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Window functions are only legal in the SELECT list and ORDER BY - never WHERE/GROUP BY/HAVING/ON. */
|
||||||
|
private fun containsWindowFunction(expression: Expression): Boolean = when (expression) {
|
||||||
|
is WindowExpression -> true
|
||||||
|
is BinaryExpression -> containsWindowFunction(expression.left) || containsWindowFunction(expression.right)
|
||||||
|
is UnaryExpression -> containsWindowFunction(expression.operand)
|
||||||
|
is CallExpression -> expression.arguments.any { containsWindowFunction(it) }
|
||||||
|
is IsNullExpression -> containsWindowFunction(expression.operand)
|
||||||
|
is InExpression -> containsWindowFunction(expression.operand) || expression.values.any { containsWindowFunction(it) }
|
||||||
|
is BetweenExpression ->
|
||||||
|
containsWindowFunction(expression.operand) || containsWindowFunction(expression.low) || containsWindowFunction(expression.high)
|
||||||
|
is LikeExpression -> containsWindowFunction(expression.operand) || containsWindowFunction(expression.pattern)
|
||||||
|
is CaseExpression ->
|
||||||
|
(expression.operand?.let { containsWindowFunction(it) } ?: false) ||
|
||||||
|
expression.branches.any { containsWindowFunction(it.condition) || containsWindowFunction(it.result) } ||
|
||||||
|
(expression.elseExpression?.let { containsWindowFunction(it) } ?: false)
|
||||||
|
is CastExpression -> containsWindowFunction(expression.operand)
|
||||||
|
is CollateExpression -> containsWindowFunction(expression.operand)
|
||||||
|
else -> false
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun rejectWindowFunction(expression: Expression, clause: String) {
|
||||||
|
if (containsWindowFunction(expression)) throw DiskExecutorException("window functions are not allowed in $clause")
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Every distinct window-function call reachable from [expression] (nested inside arithmetic, CASE, ...). */
|
||||||
|
private fun collectWindowExpressions(expression: Expression): List<WindowExpression> = when (expression) {
|
||||||
|
is WindowExpression -> listOf(expression)
|
||||||
|
is BinaryExpression -> collectWindowExpressions(expression.left) + collectWindowExpressions(expression.right)
|
||||||
|
is UnaryExpression -> collectWindowExpressions(expression.operand)
|
||||||
|
is CallExpression -> expression.arguments.flatMap { collectWindowExpressions(it) }
|
||||||
|
is IsNullExpression -> collectWindowExpressions(expression.operand)
|
||||||
|
is InExpression -> collectWindowExpressions(expression.operand) + expression.values.flatMap { collectWindowExpressions(it) }
|
||||||
|
is BetweenExpression ->
|
||||||
|
collectWindowExpressions(expression.operand) + collectWindowExpressions(expression.low) + collectWindowExpressions(expression.high)
|
||||||
|
is LikeExpression -> collectWindowExpressions(expression.operand) + collectWindowExpressions(expression.pattern)
|
||||||
|
is CaseExpression ->
|
||||||
|
(expression.operand?.let { collectWindowExpressions(it) } ?: emptyList()) +
|
||||||
|
expression.branches.flatMap { collectWindowExpressions(it.condition) + collectWindowExpressions(it.result) } +
|
||||||
|
(expression.elseExpression?.let { collectWindowExpressions(it) } ?: emptyList())
|
||||||
|
is CastExpression -> collectWindowExpressions(expression.operand)
|
||||||
|
is CollateExpression -> collectWindowExpressions(expression.operand)
|
||||||
|
else -> emptyList()
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Every distinct window-function call in the SELECT list or (post alias-resolution) ORDER BY -
|
||||||
|
* the only two clauses where one is allowed, mirroring [kestrel.engine.SelectRunner]'s own
|
||||||
|
* `windowExpressions` property.
|
||||||
|
*/
|
||||||
|
private fun collectWindowExpressions(projection: Projection, orderKeys: List<OrderKey>): List<WindowExpression> {
|
||||||
|
val projectionWindows = (projection as? ExpressionProjection)?.items?.flatMap { collectWindowExpressions(it.expression) } ?: emptyList()
|
||||||
|
val orderWindows = orderKeys.flatMap { collectWindowExpressions(it.expression) }
|
||||||
|
return (projectionWindows + orderWindows).distinct()
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 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 executor
|
||||||
|
* does not implement (same scope cut as [kestrel.engine.SelectRunner]). RANK/DENSE_RANK need an
|
||||||
|
* explicit tie-break (unlike ROW_NUMBER, which falls back to partition arrival order when OVER
|
||||||
|
* has no ORDER BY) and nesting one window function inside another's argument/PARTITION BY/
|
||||||
|
* ORDER BY/frame offset is invalid SQL in the first place.
|
||||||
|
*/
|
||||||
|
private fun validateWindows(windowExpressions: List<WindowExpression>, aggregated: Boolean) {
|
||||||
|
if (aggregated) {
|
||||||
|
throw DiskExecutorException("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 DiskExecutorException("${window.name}() OVER (...) requires an ORDER BY clause")
|
||||||
|
}
|
||||||
|
if (window.frame != null && window.kind !in FRAME_CAPABLE_KINDS) {
|
||||||
|
throw DiskExecutorException("${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 DiskExecutorException("window functions cannot be nested")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Reference in new issue
Block a user