Skip to content

Commit

Permalink
Fix LIMIT clause execution order (#300)
Browse files Browse the repository at this point in the history
  • Loading branch information
@alancai98
alancai98 committed Oct 19, 2020
1 parent df9d8cb commit c1dedfc
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 47 deletions.
116 changes: 69 additions & 47 deletions lang/src/org/partiql/lang/eval/EvaluatingCompiler.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -889,6 +889,47 @@ internal class EvaluatingCompiler(
}
}

private fun evalLimit(limitThunk: ThunkEnv, env: Environment, limitLocationMeta: SourceLocationMeta?): Long {
val limitExprValue = limitThunk(env)

if(limitExprValue.type != ExprValueType.INT) {
err("LIMIT value was not an integer",
ErrorCode.EVALUATOR_NON_INT_LIMIT_VALUE,
errorContextFrom(limitLocationMeta).also {
it[Property.ACTUAL_TYPE] = limitExprValue.type.toString()
},
internal = false)
}

// `Number.toLong()` (used below) does *not* cause an overflow exception if the underlying [Number]
// implementation (i.e. Decimal or BigInteger) exceeds the range that can be represented by Longs.
// This can cause very confusing behavior if the user specifies a LIMIT value that exceeds
// Long.MAX_VALUE, because no results will be returned from their query. That no overflow exception
// is thrown is not a problem as long as PartiQL's restriction of integer values to +/- 2^63 remains.
// We throw an exception here if the value exceeds the supported range (say if we change that
// restriction or if a custom [ExprValue] is provided which exceeds that value).
val limitIonValue = limitExprValue.ionValue as IonInt
if(limitIonValue.integerSize == IntegerSize.BIG_INTEGER) {
err("IntegerSize.BIG_INTEGER not supported for LIMIT values",
errorContextFrom(limitLocationMeta),
internal = true)
}

val limitValue = limitExprValue.numberValue().toLong()

if (limitValue < 0) {
err("negative LIMIT",
ErrorCode.EVALUATOR_NEGATIVE_LIMIT,
errorContextFrom(limitLocationMeta),
internal = false)
}

// we can't use the Kotlin's Sequence<T>.take(n) for this since it accepts only an integer.
// this references [Sequence<T>.take(count: Long): Sequence<T>] defined in [org.partiql.util].
return limitValue
}


private fun compileSelect(selectExpr: Select): ThunkEnv {
// Get all the FROM source aliases and LET bindings for binding error checks
val fold = object : PartiqlAst.VisitorFold<Set<String>>() {
Expand All @@ -913,14 +954,15 @@ internal class EvaluatingCompiler(
.union(pigGeneratedAst.fromLet?.let { fold.walkLet(pigGeneratedAst.fromLet, emptySet()) } ?: emptySet())

return nestCompilationContext(ExpressionContext.NORMAL, emptySet()) {
val (setQuantifier, projection, from, fromLet, _, groupBy, having, _, metas: MetaContainer) = selectExpr
val (setQuantifier, projection, from, fromLet, _, groupBy, having, limit, metas: MetaContainer) = selectExpr

val fromSourceThunks = compileFromSources(from)

val letSourceThunks = fromLet?.let { compileLetSources(it) }

val sourceThunks = compileQueryWithoutProjection(selectExpr, fromSourceThunks, letSourceThunks)

val limitThunk = limit?.let { compileExprNode(limit) }
val limitLocationMeta = limit?.metas?.sourceLocationMeta

// Returns a thunk that invokes [sourceThunks], and invokes [projectionThunk] to perform the projection.
fun getQueryThunk(selectProjectionThunk: ThunkEnvValue<List<ExprValue>>): ThunkEnv {
val (_, groupByItems, groupAsName) = groupBy ?: GroupBy(GroupingStrategy.FULL, listOf())
Expand All @@ -941,6 +983,11 @@ internal class EvaluatingCompiler(
// wrap the ExprValue to use ExprValue.equals as the equality
SetQuantifier.DISTINCT -> projectedRows.filter(createUniqueExprValueFilter())
SetQuantifier.ALL -> projectedRows
}.let { rows ->
when (limitThunk) {
null -> rows
else -> rows.take(evalLimit(limitThunk, env, limitLocationMeta))
}
}

valueFactory.newBag(quantifiedRows.map {
Expand Down Expand Up @@ -983,7 +1030,12 @@ internal class EvaluatingCompiler(
// Create a closure that groups all the rows in the FROM source into a single group.
thunkFactory.thunkEnv(metas) { env ->
// Evaluate the FROM clause
val fromProductions: Sequence<FromProduction> = sourceThunks(env)
val fromProductions: Sequence<FromProduction> = sourceThunks(env).let { rows ->
when (limitThunk) {
null -> rows
else -> rows.take(evalLimit(limitThunk, env, limitLocationMeta))
}
}
val registers = createRegisterBank()

// note: the group key can be anything here because we only ever have a single
Expand Down Expand Up @@ -1066,9 +1118,14 @@ internal class EvaluatingCompiler(
val projectedRows = env.groups.mapNotNull { g ->
val groupByEnv = getGroupEnv(env, g.value)
filterHavingAndProject(groupByEnv, g.value)
}.asSequence().let { rows ->
when (limitThunk) {
null -> rows
else -> rows.take(evalLimit(limitThunk, env, limitLocationMeta))
}
}

valueFactory.newBag(projectedRows.asSequence())
valueFactory.newBag(projectedRows)
}
}
}
Expand All @@ -1094,9 +1151,14 @@ internal class EvaluatingCompiler(
val asThunk = compileExprNode(asExpr)
val atThunk = compileExprNode(atExpr)
thunkFactory.thunkEnv(metas) { env ->
val sourceValue = sourceThunks(env)
val sourceValue = sourceThunks(env).asSequence().let { rows ->
when (limitThunk) {
null -> rows
else -> rows.take(evalLimit(limitThunk, env, limitLocationMeta))
}
}

val seq = sourceValue
.asSequence()
.map { (_, env) -> Pair(asThunk(env), atThunk(env)) }
.filter { (name, _) -> name.type.isText }
.map { (name, value) -> value.namedValue(name) }
Expand Down Expand Up @@ -1385,8 +1447,6 @@ internal class EvaluatingCompiler(

val localsBinder = compiledSources.map { it.alias }.localsBinder(valueFactory.missingValue)
val whereThunk = ast.where?.let { compileExprNode(it) }
val limitThunk = ast.limit?.let { compileExprNode(it) }
val limitLocationMeta = ast.limit?.metas?.sourceLocationMeta

return { rootEnv ->
val fromEnv = rootEnv.flipToGlobalsFirst()
Expand Down Expand Up @@ -1481,44 +1541,6 @@ internal class EvaluatingCompiler(
}
}
}
if (limitThunk != null) {
val limitExprValue = limitThunk(rootEnv)
if(limitExprValue.type != ExprValueType.INT) {
err("LIMIT value was not an integer",
ErrorCode.EVALUATOR_NON_INT_LIMIT_VALUE,
errorContextFrom(limitLocationMeta).also {
it[Property.ACTUAL_TYPE] = limitExprValue.type.toString()
},
internal = false)
}

// `Number.toLong()` (used below) does *not* cause an overflow exception if the underlying [Number]
// implementation (i.e. Decimal or BigInteger) exceeds the range that can be represented by Longs.
// This can cause very confusing behavior if the user specifies a LIMIT value that exceeds
// Long.MAX_VALUE, because no results will be returned from their query. That no overflow exception
// is thrown is not a problem as long as PartiQL's restriction of integer values to +/- 2^63 remains.
// We throw an exception here if the value exceeds the supported range (say if we change that
// restriction or if a custom [ExprValue] is provided which exceeds that value).
val limitIonValue = limitExprValue.ionValue as IonInt
if(limitIonValue.integerSize == IntegerSize.BIG_INTEGER) {
err("IntegerSize.BIG_INTEGER not supported for LIMIT values",
errorContextFrom(limitLocationMeta),
internal = true)
}

val limitValue = limitExprValue.numberValue().toLong()

if (limitValue < 0) {
err("negative LIMIT",
ErrorCode.EVALUATOR_NEGATIVE_LIMIT,
errorContextFrom(limitLocationMeta),
internal = false)
}

// we can't use the Kotlin's Sequence<T>.take(n) for this since it accepts only an integer.
// this references [Sequence<T>.take(count: Long): Sequence<T>] defined in [org.partiql.util].
seq = seq.take(limitValue)
}
seq
}
}
Expand Down
7 changes: 7 additions & 0 deletions lang/test/org/partiql/lang/eval/EvaluatingCompilerLimitTests.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,4 +61,11 @@ class EvaluatingCompilerLimitTests : EvaluatorTestBase() {
""" select * from <<1>> limit 'this won''t work' """,
ErrorCode.EVALUATOR_NON_INT_LIMIT_VALUE,
sourceLocationProperties(1, 28) + mapOf(Property.ACTUAL_TYPE to "STRING"))

@Test
fun `LIMIT applied after GROUP BY`() =
assertEval(
"SELECT g FROM `[{foo: 1, bar: 10}, {foo: 1, bar: 11}]` AS f GROUP BY f.foo GROUP AS g LIMIT 1",
"""[ { 'g': [ { 'f': { 'foo': 1, 'bar': 10 } }, { 'f': { 'foo': 1, 'bar': 11 } } ] } ]"""
)
}

0 comments on commit c1dedfc

Please sign in to comment.