Add support for OFFSET

docs/dev/RewriterToVisitorTransformGuide.md

@@ -78,11 +78,11 @@ In the rewriter code, you may encounter [`innerRewriteSelect`](https://github.co
 
 However, currently there are some slight differences between the order of transforms in `innerRewriteSelect` and `transformExprSelect`. `innerRewriteSelect` follows the traversal order (SQL semantic order) that is
 
-FROM → (FROM LET) → (WHERE) → (GROUP BY) → (HAVING) → PROJECTION → (LIMIT)
+FROM → (FROM LET) → (WHERE) → (GROUP BY) → (HAVING) → PROJECTION → (OFFSET) → (LIMIT) 
 
 `transformExprSelect` transforms the clauses in the order they are written for a PartiQL query, that is
 
-PROJECTION → FROM → (FROM LET) → (WHERE) → (GROUP BY) → (HAVING) → (LIMIT)
+PROJECTION → FROM → (FROM LET) → (WHERE) → (GROUP BY) → (HAVING) → (LIMIT) → (OFFSET)
 
 This slight difference can lead to some different behaviors when translating `innerRewriteSelect` (e.g. StaticTypeRewriter). So to completely solve this issue, you can have your VisitorTransform implement [`VisitorTransformBase`](https://github.com/partiql/partiql-lang-kotlin/blob/master/lang/src/org/partiql/lang/eval/visitors/VisitorTransformBase.kt) and call `transformExprSelectEvaluationOrder` rather than `transformExprSelect` to get the same behavior.
 
@@ -97,6 +97,7 @@ fun transformExprSelectEvaluationOrder(node: PartiqlAst.Expr.Select): PartiqlAst
     val having = transformExprSelect_having(node)
     val setq = transformExprSelect_setq(node)
     val project = transformExprSelect_project(node)
+    val offset = transformExprSelect_offset(node)
     val limit = transformExprSelect_limit(node)
     val metas = transformExprSelect_metas(node)
     return PartiqlAst.build {
@@ -109,6 +110,7 @@ fun transformExprSelectEvaluationOrder(node: PartiqlAst.Expr.Select): PartiqlAst
             group = group,
             having = having,
             limit = limit,
+            offset = offset,
             metas = metas)
     }
 }
@@ -158,6 +160,7 @@ private fun copyProjectionToSelect(node: PartiqlAst.Expr.Select, newProjection:
             group = node.group,
             having = node.having,
             limit = node.limit,
+            offset = node.offset,
             metas = node.metas)
     }
 }

lang/resources/org/partiql/type-domains/partiql.ion

@@ -121,7 +121,8 @@
                 (group (? group_by))
                 (having (? expr))
                 (order (? order_by))
-                (limit (? expr))))
+                (limit (? expr))
+                (offset (? expr))))
         // end of sum expr
 
         // Time

lang/src/org/partiql/lang/ast/AstSerialization.kt

@@ -207,7 +207,12 @@ private class AstSerializerImpl(val astVersion: AstVersion, val ion: IonSystem):
     }
 
     private fun IonWriterContext.writeSelect(expr: Select) {
-        val (setQuantifier, projection, from, fromLet, where, groupBy, having, orderBy, limit, _: MetaContainer) = expr
+        val (setQuantifier, projection, from, fromLet, where, groupBy, having, orderBy, limit, offset, _: MetaContainer) = expr
+
+        if (offset != null){
+            throw UnsupportedOperationException("OFFSET clause is not supported by the V0 AST")
+        }
+
         if (orderBy != null) {
             throw UnsupportedOperationException("ORDER BY clause is not supported by the V0 AST")
         }

lang/src/org/partiql/lang/ast/ExprNodeToStatement.kt

@@ -160,6 +160,7 @@ fun ExprNode.toAstExpr(): PartiqlAst.Expr {
                     group = node.groupBy?.toAstGroupSpec(),
                     having = node.having?.toAstExpr(),
                     limit = node.limit?.toAstExpr(),
+                    offset = node.offset?.toAstExpr(),
                     metas = metas)
             is Struct -> struct(node.fields.map { exprPair(it.name.toAstExpr(), it.expr.toAstExpr()) }, metas)
             is Seq ->

lang/src/org/partiql/lang/ast/StatementToExprNode.kt

@@ -190,6 +190,7 @@ private class StatementTransformer(val ion: IonSystem) {
                     having = having?.toExprNode(),
                     orderBy = order?.toOrderBy(),
                     limit = limit?.toExprNode(),
+                    offset = offset?.toExprNode(),
                     metas = metas
             )
             is Expr.Date ->

lang/src/org/partiql/lang/ast/ast.kt

@@ -440,9 +440,10 @@ data class Select(
     val having: ExprNode? = null,
     val orderBy: OrderBy? = null,
     val limit: ExprNode? = null,
+    val offset: ExprNode? = null,
     override val metas: MetaContainer
 ) : ExprNode() {
-    override val children: List<AstNode> = listOfNotNull(projection, from, fromLet, where, groupBy, having, orderBy, limit)
+    override val children: List<AstNode> = listOfNotNull(projection, from, fromLet, where, groupBy, having, orderBy, limit, offset)
 }
 
 //********************************

lang/src/org/partiql/lang/ast/passes/AstRewriterBase.kt

@@ -154,7 +154,8 @@ open class AstRewriterBase : AstRewriter {
      * 5. `HAVING`
      * 6. *projection*
      * 7. `ORDER BY` (to be implemented)
-     * 8. `LIMIT`
+     * 8. `OFFSET`
+     * 9. `LIMIT`
      */
     protected open fun innerRewriteSelect(selectExpr: Select): Select {
         val from = rewriteFromSource(selectExpr.from)
@@ -164,6 +165,7 @@ open class AstRewriterBase : AstRewriter {
         val having = selectExpr.having?.let { rewriteSelectHaving(it) }
         val projection = rewriteSelectProjection(selectExpr.projection)
         val orderBy = selectExpr.orderBy?.let { rewriteOrderBy(it) }
+        val offset = selectExpr.offset?.let { rewriteSelectOffset(it) }
         val limit = selectExpr.limit?.let { rewriteSelectLimit(it) }
         val metas = rewriteSelectMetas(selectExpr)
 
@@ -177,6 +179,7 @@ open class AstRewriterBase : AstRewriter {
             having = having,
             orderBy = orderBy,
             limit = limit,
+            offset = offset,
             metas = metas)
     }
 
@@ -186,6 +189,8 @@ open class AstRewriterBase : AstRewriter {
 
     open fun rewriteSelectLimit(node: ExprNode): ExprNode = rewriteExprNode(node)
 
+    open fun rewriteSelectOffset(node: ExprNode): ExprNode = rewriteExprNode(node)
+
     open fun rewriteSelectMetas(selectExpr: Select): MetaContainer = rewriteMetas(selectExpr)
 
     open fun rewriteSelectProjection(projection: SelectProjection): SelectProjection =

lang/src/org/partiql/lang/ast/passes/AstWalker.kt

@@ -88,7 +88,7 @@ open class AstWalker(private val visitor: AstVisitor) {
                     }
                 }
                 is Select       -> case {
-                    val (_, projection, from, fromLet, where, groupBy, having, orderBy, limit, _: MetaContainer) = expr
+                    val (_, projection, from, fromLet, where, groupBy, having, orderBy, limit, offset, _: MetaContainer) = expr
                     walkSelectProjection(projection)
                     walkFromSource(from)
                     walkExprNode(where)
@@ -107,6 +107,7 @@ open class AstWalker(private val visitor: AstVisitor) {
                         }
                     }
                     walkExprNode(limit)
+                    walkExprNode(offset)
                 }
                 is DataManipulation -> case {
                     val (dmlOperation, from, where, returning, _: MetaContainer) = expr

lang/src/org/partiql/lang/errors/ErrorCode.kt

@@ -654,6 +654,19 @@ enum class ErrorCode(private val category: ErrorCategory,
         LOCATION,
         "LIMIT must not be negative"),
 
+    EVALUATOR_NON_INT_OFFSET_VALUE (
+        ErrorCategory.EVALUATOR,
+        LOCATION + setOf(Property.ACTUAL_TYPE),
+        "") {
+        override fun getErrorMessage(errorContext: PropertyValueMap?): String =
+            "OFFSET value must be an integer but found ${errorContext.getProperty(Property.ACTUAL_TYPE)}"
+    },
+
+    EVALUATOR_NEGATIVE_OFFSET(
+        ErrorCategory.EVALUATOR,
+        LOCATION,
+        "OFFSET must not be negative"),
+
     EVALUATOR_DIVIDE_BY_ZERO(
         ErrorCategory.EVALUATOR,
         LOCATION,

lang/src/org/partiql/lang/eval/EvaluatingCompiler.kt

@@ -910,7 +910,7 @@ internal class EvaluatingCompiler(
     private fun evalLimit(limitThunk: ThunkEnv, env: Environment, limitLocationMeta: SourceLocationMeta?): Long {
         val limitExprValue = limitThunk(env)
 
-        if(limitExprValue.type != ExprValueType.INT) {
+        if (limitExprValue.type != ExprValueType.INT) {
             err("LIMIT value was not an integer",
                 ErrorCode.EVALUATOR_NON_INT_LIMIT_VALUE,
                 errorContextFrom(limitLocationMeta).also {
@@ -927,7 +927,7 @@ internal class EvaluatingCompiler(
         // 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) {
+        if (limitIonValue.integerSize == IntegerSize.BIG_INTEGER) {
             err("IntegerSize.BIG_INTEGER not supported for LIMIT values",
                 errorContextFrom(limitLocationMeta),
                 internal = true)
@@ -947,6 +947,43 @@ internal class EvaluatingCompiler(
         return limitValue
     }
 
+    private fun evalOffset(offsetThunk: ThunkEnv, env: Environment, offsetLocationMeta: SourceLocationMeta?): Long {
+        val offsetExprValue = offsetThunk(env)
+
+        if (offsetExprValue.type != ExprValueType.INT) {
+            err("OFFSET value was not an integer",
+                ErrorCode.EVALUATOR_NON_INT_OFFSET_VALUE,
+                errorContextFrom(offsetLocationMeta).also {
+                    it[Property.ACTUAL_TYPE] = offsetExprValue.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 OFFSET 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 offsetIonValue = offsetExprValue.ionValue as IonInt
+        if (offsetIonValue.integerSize == IntegerSize.BIG_INTEGER) {
+            err("IntegerSize.BIG_INTEGER not supported for OFFSET values",
+                errorContextFrom(offsetLocationMeta),
+                internal = true)
+        }
+
+        val offsetValue = offsetExprValue.numberValue().toLong()
+
+        if (offsetValue < 0) {
+            err("negative OFFSET",
+                ErrorCode.EVALUATOR_NEGATIVE_OFFSET,
+                errorContextFrom(offsetLocationMeta),
+                internal = false)
+        }
+
+        return offsetValue
+    }
 
     private fun compileSelect(selectExpr: Select): ThunkEnv {
         selectExpr.orderBy?.let {
@@ -980,15 +1017,28 @@ 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, _, limit, metas: MetaContainer) = selectExpr
+            val (setQuantifier, projection, from, fromLet, _, groupBy, having, _, limit, offset, 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 offsetThunk = offset?.let { compileExprNode(it) }
+            val offsetLocationMeta = offset?.metas?.sourceLocationMeta
+            val limitThunk = limit?.let { compileExprNode(it) }
             val limitLocationMeta = limit?.metas?.sourceLocationMeta
 
+            fun <T> rowsWithOffsetAndLimit (rows: Sequence<T>, env: Environment): Sequence<T> {
+                val rowsWithOffset = when (offsetThunk){
+                    null -> rows
+                    else -> rows.drop(evalOffset(offsetThunk, env, offsetLocationMeta))
+                }
+                return when (limitThunk) {
+                    null -> rowsWithOffset
+                    else -> rowsWithOffset.take(evalLimit(limitThunk, env, limitLocationMeta))
+                }
+            }
+
             // 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())
@@ -1009,12 +1059,7 @@ 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))
-                                }
-                            }
+                            }.let { rowsWithOffsetAndLimit(it, env) }
 
                             valueFactory.newBag(quantifiedRows.map {
                                 // TODO make this expose the ordinal for ordered sequences
@@ -1056,12 +1101,7 @@ 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).let { rows ->
-                                        when (limitThunk) {
-                                            null -> rows
-                                            else -> rows.take(evalLimit(limitThunk, env, limitLocationMeta))
-                                        }
-                                    }
+                                    val fromProductions: Sequence<FromProduction> = rowsWithOffsetAndLimit(sourceThunks(env), env)
                                     val registers = createRegisterBank()
 
                                     // note: the group key can be anything here because we only ever have a single
@@ -1144,12 +1184,7 @@ 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))
-                                        }
-                                    }
+                                    }.asSequence().let { rowsWithOffsetAndLimit(it, env) }
 
                                     valueFactory.newBag(projectedRows)
                                 }
@@ -1177,13 +1212,7 @@ internal class EvaluatingCompiler(
                         val asThunk = compileExprNode(asExpr)
                         val atThunk = compileExprNode(atExpr)
                         thunkFactory.thunkEnv(metas) { env ->
-                            val sourceValue = sourceThunks(env).asSequence().let { rows ->
-                                when (limitThunk) {
-                                    null -> rows
-                                    else -> rows.take(evalLimit(limitThunk, env, limitLocationMeta))
-                                }
-                            }
-
+                            val sourceValue = rowsWithOffsetAndLimit(sourceThunks(env).asSequence(), env)
                             val seq = sourceValue
                                     .map { (_, env) -> Pair(asThunk(env), atThunk(env)) }
                                     .filter { (name, _) -> name.type.isText }

lang/src/org/partiql/lang/eval/visitors/GroupByPathExpressionVisitorTransform.kt

@@ -103,6 +103,7 @@ class GroupByPathExpressionVisitorTransform(
         val groupBy = unshadowedTransformer.transformExprSelect_group(node)
         val having = currentAndUnshadowedTransformer.transformExprSelect_having(node)
         val order = currentAndUnshadowedTransformer.transformExprSelect_order(node)
+        val offset = unshadowedTransformer.transformExprSelect_offset(node)
         val limit = unshadowedTransformer.transformExprSelect_limit(node)
         val metas = unshadowedTransformer.transformExprSelect_metas(node)
 
@@ -116,6 +117,7 @@ class GroupByPathExpressionVisitorTransform(
                 group = groupBy,
                 having = having,
                 order = order,
+                offset = offset,
                 limit = limit,
                 metas = metas)
         }

lang/src/org/partiql/lang/eval/visitors/SelectStarVisitorTransform.kt

@@ -24,6 +24,7 @@ class SelectStarVisitorTransform : VisitorTransformBase() {
                 having = node.having,
                 order = node.order,
                 limit = node.limit,
+                offset = node.offset,
                 metas = node.metas)
         }
     }

lang/src/org/partiql/lang/eval/visitors/VisitorTransformBase.kt

@@ -27,8 +27,9 @@ abstract class VisitorTransformBase : PartiqlAst.VisitorTransform() {
      * 4. `GROUP BY`
      * 5. `HAVING`
      * 6. *projection*
-     * 7. `LIMIT`
-     * 8. The metas.
+     * 7. `OFFSET`
+     * 8. `LIMIT`
+     * 9. The metas.
      *
      * This differs from [transformExprSelect], which executes following the written order of clauses.
      */
@@ -41,6 +42,7 @@ abstract class VisitorTransformBase : PartiqlAst.VisitorTransform() {
         val setq = transformExprSelect_setq(node)
         val project = transformExprSelect_project(node)
         val order = transformExprSelect_order(node)
+        val offset = transformExprSelect_offset(node)
         val limit = transformExprSelect_limit(node)
         val metas = transformExprSelect_metas(node)
         return PartiqlAst.build {
@@ -53,6 +55,7 @@ abstract class VisitorTransformBase : PartiqlAst.VisitorTransform() {
                 group = group,
                 having = having,
                 order = order,
+                offset = offset,
                 limit = limit,
                 metas = metas)
         }

lang/src/org/partiql/lang/syntax/LexerConstants.kt

@@ -258,6 +258,7 @@ internal val DATE_PART_KEYWORDS: Set<String> = DatePart.values()
     "pivot",
     "unpivot",
     "limit",
+    "offset",
     "tuple",
     "remove",
     "index",