Adds full support for the VARIANT type

Fixes conformance tests related to the lack of a VARIANT type.
partiql · Dec 13, 2024 · f70aa7a · f70aa7a
1 parent b7852b7
commit f70aa7a
Show file tree

Hide file tree

Showing 24 changed files with 203 additions and 174 deletions.
diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/helpers/ValueUtility.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/helpers/ValueUtility.kt
@@ -15,7 +15,13 @@ internal object ValueUtility {
      */
     @JvmStatic
     fun Datum.isTrue(): Boolean {
-        return this.type.code() == PType.BOOL && !this.isNull && !this.isMissing && this.boolean
+        if (this.isNull || this.isMissing) {
+            return false
+        }
+        if (this.type.code() == PType.VARIANT) {
+            return this.lower().isTrue()
+        }
+        return this.type.code() == PType.BOOL && this.boolean
     }
 
     /**
@@ -29,6 +35,9 @@ internal object ValueUtility {
         if (this.type == type) {
             return this
         }
+        if (this.type.code() == PType.VARIANT) {
+            return this.lower().check(type)
+        }
         if (!this.isNull) {
             throw TypeCheckException("Expected type $type but received ${this.type}.")
         }
@@ -43,29 +52,42 @@ internal object ValueUtility {
      */
     fun Datum.getText(): String {
         return when (this.type.code()) {
+            PType.VARIANT -> this.lower().getText()
             PType.STRING, PType.CHAR -> this.string
             else -> throw TypeCheckException("Expected text, but received ${this.type}.")
         }
     }
 
     /**
-     * Takes in a [Datum] that is any integer type ([PartiQLValueType.INT8], [PartiQLValueType.INT8],
-     * [PartiQLValueType.INT8], [PartiQLValueType.INT8], [PartiQLValueType.INT8]) and returns the [BigInteger] (potentially
-     * coerced) that represents the integer.
+     * Converts all number values to [BigInteger]. If the number is [PType.DECIMAL] or [PType.NUMERIC], this asserts that
+     * the scale is zero.
      *
      * INTERNAL NOTE: The PLANNER should be handling the coercion. This function should not be necessary.
      *
+     * TODO: This is used specifically for LIMIT and OFFSET. This makes the conformance tests pass by coercing values
+     *  of type [PType.NUMERIC] and [PType.DECIMAL], but this is unspecified. Do we allow for LIMIT 2.0? Or of
+     *  a value that is greater than [PType.BIGINT]'s MAX value by using a [PType.DECIMAL] with a high precision and scale
+     *  of zero? This hasn't been decided, however, as the conformance tests allow for this, this function coerces
+     *  the value to a [BigInteger] regardless of the number's type.
+     *
      * @throws NullPointerException if the value is null
      * @throws TypeCheckException if type is not an integer type
      */
     fun Datum.getBigIntCoerced(): BigInteger {
         return when (this.type.code()) {
+            PType.VARIANT -> this.lower().getBigIntCoerced()
             PType.TINYINT -> this.byte.toInt().toBigInteger()
             PType.SMALLINT -> this.short.toInt().toBigInteger()
             PType.INTEGER -> this.int.toBigInteger()
             PType.BIGINT -> this.long.toBigInteger()
-            PType.NUMERIC -> this.bigInteger
-            else -> throw TypeCheckException()
+            PType.NUMERIC, PType.DECIMAL -> {
+                val decimal = this.bigDecimal
+                if (decimal.scale() != 0) {
+                    throw TypeCheckException("Expected integer, but received decimal.")
+                }
+                return decimal.toBigInteger()
+            }
+            else -> throw TypeCheckException("Type: ${this.type}")
         }
     }
 
@@ -82,6 +104,7 @@ internal object ValueUtility {
      */
     fun Datum.getInt32Coerced(): Int {
         return when (this.type.code()) {
+            PType.VARIANT -> this.lower().getInt32Coerced()
             PType.TINYINT -> this.byte.toInt()
             PType.SMALLINT -> this.short.toInt()
             PType.INTEGER -> this.int

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpIterate.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpIterate.kt
@@ -18,7 +18,12 @@ internal class RelOpIterate(
     override fun open(env: Environment) {
         val r = expr.eval(env.push(Row()))
         index = 0
-        iterator = when (r.type.code()) {
+        iterator = records(r)
+    }
+
+    private fun records(r: Datum): Iterator<Datum> {
+        return when (r.type.code()) {
+            PType.VARIANT -> records(r.lower())
             PType.BAG -> {
                 close()
                 throw TypeCheckException()

diff --git a/...iql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpIteratePermissive.kt b/...iql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpIteratePermissive.kt
@@ -18,7 +18,12 @@ internal class RelOpIteratePermissive(
     override fun open(env: Environment) {
         val r = expr.eval(env.push(Row()))
         index = 0
-        iterator = when (r.type.code()) {
+        iterator = records(r)
+    }
+
+    private fun records(r: Datum): Iterator<Datum> {
+        return when (r.type.code()) {
+            PType.VARIANT -> records(r.lower())
             PType.BAG -> {
                 isIndexable = false
                 r.iterator()

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpScan.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpScan.kt
@@ -6,6 +6,7 @@ import org.partiql.eval.ExprRelation
 import org.partiql.eval.ExprValue
 import org.partiql.eval.Row
 import org.partiql.eval.internal.helpers.RecordValueIterator
+import org.partiql.spi.value.Datum
 import org.partiql.types.PType
 
 internal class RelOpScan(
@@ -16,11 +17,16 @@ internal class RelOpScan(
 
     override fun open(env: Environment) {
         val r = expr.eval(env.push(Row()))
-        records = when (r.type.code()) {
-            PType.ARRAY, PType.BAG -> RecordValueIterator(r.iterator())
+        records = r.records()
+    }
+
+    private fun Datum.records(): RecordValueIterator {
+        return when (this.type.code()) {
+            PType.VARIANT -> this.lower().records()
+            PType.ARRAY, PType.BAG -> RecordValueIterator(this.iterator())
             else -> {
                 close()
-                throw TypeCheckException("Unexpected type for scan: ${r.type}")
+                throw TypeCheckException("Unexpected type for scan: ${this.type}")
             }
         }
     }

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpScanPermissive.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpScanPermissive.kt
@@ -5,6 +5,7 @@ import org.partiql.eval.ExprRelation
 import org.partiql.eval.ExprValue
 import org.partiql.eval.Row
 import org.partiql.eval.internal.helpers.RecordValueIterator
+import org.partiql.spi.value.Datum
 import org.partiql.types.PType
 
 internal class RelOpScanPermissive(
@@ -15,7 +16,13 @@ internal class RelOpScanPermissive(
 
     override fun open(env: Environment) {
         val r = expr.eval(env.push(Row()))
-        records = when (r.type.code()) {
+        records = r.records()
+    }
+
+    private fun Datum.records(): Iterator<Row> {
+        val r = this
+        return when (type.code()) {
+            PType.VARIANT -> r.lower().records()
             PType.BAG, PType.ARRAY -> RecordValueIterator(r.iterator())
             else -> iterator { yield(Row(arrayOf(r))) }
         }

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpUnpivot.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelOpUnpivot.kt
@@ -58,7 +58,13 @@ internal sealed class RelOpUnpivot : ExprRelation {
     class Strict(private val expr: ExprValue) : RelOpUnpivot() {
 
         override fun struct(): Datum {
-            val v = expr.eval(env.push(Row()))
+            val v = expr.eval(env.push(Row())).let {
+                if (it.type.code() == PType.VARIANT) {
+                    it.lower()
+                } else {
+                    it
+                }
+            }
             if (v.type.code() != PType.STRUCT && v.type.code() != PType.ROW) {
                 throw TypeCheckException()
             }
@@ -78,7 +84,13 @@ internal sealed class RelOpUnpivot : ExprRelation {
     class Permissive(private val expr: ExprValue) : RelOpUnpivot() {
 
         override fun struct(): Datum {
-            val v = expr.eval(env.push(Row()))
+            val v = expr.eval(env.push(Row())).let {
+                if (it.type.code() == PType.VARIANT) {
+                    it.lower()
+                } else {
+                    it
+                }
+            }
             if (v.isMissing) {
                 return Datum.struct(emptyList())
             }

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/CastTable.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/CastTable.kt
@@ -29,6 +29,7 @@ import org.partiql.types.PType.TIMESTAMPZ
 import org.partiql.types.PType.TIMEZ
 import org.partiql.types.PType.TINYINT
 import org.partiql.types.PType.VARCHAR
+import org.partiql.types.PType.VARIANT
 import org.partiql.value.datetime.DateTimeValue
 import java.math.BigDecimal
 import java.math.BigInteger
@@ -70,6 +71,9 @@ internal object CastTable {
         if (target.code() == DYNAMIC) {
             return source
         }
+        if (source.type.code() == VARIANT) {
+            return cast(source.lower(), target)
+        }
         val cast = _table[source.type.code()][target.code()]
             ?: throw TypeCheckException("CAST(${source.type} AS $target) is not supported.")
         return try {

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCallDynamic.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCallDynamic.kt
@@ -45,7 +45,14 @@ internal class ExprCallDynamic(
     private val candidates: MutableMap<List<PType>, Candidate> = mutableMapOf()
 
     override fun eval(env: Environment): Datum {
-        val actualArgs = args.map { it.eval(env) }.toTypedArray()
+        val actualArgs = args.map {
+            val arg = it.eval(env)
+            if (arg.type.code() == PType.VARIANT) {
+                arg.lower()
+            } else {
+                arg
+            }
+        }.toTypedArray()
         val actualTypes = actualArgs.map { it.type }
         var candidate = candidates[actualTypes]
         if (candidate == null) {

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprStructPermissive.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprStructPermissive.kt
@@ -32,6 +32,7 @@ internal class ExprStructPermissive(private val fields: List<ExprStructField>) :
                 return null
             }
             return when (this.type.code()) {
+                PType.VARIANT -> this.lower().getTextOrNull()
                 PType.STRING, PType.CHAR -> this.string
                 else -> null
             }

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/FnResolver.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/FnResolver.kt
@@ -47,7 +47,7 @@ internal object FnResolver {
         }
 
         // 2. If there are DYNAMIC arguments, return all candidates
-        val isDynamic = args.any { it.code() == PType.DYNAMIC }
+        val isDynamic = args.any { it.code() == PType.DYNAMIC || it.code() == PType.VARIANT }
         if (isDynamic) {
             val orderedMatches = candidates.sortedWith(FnComparator)
             return FnMatch.Dynamic(orderedMatches)

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/casts/CastTable.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/casts/CastTable.kt
@@ -63,6 +63,11 @@ internal class CastTable private constructor(
                     cast(it)
                 }
             }
+            graph[PType.VARIANT] = relationships {
+                PType.codes().map {
+                    cast(it)
+                }
+            }
             graph[PType.DYNAMIC] = relationships {
                 cast(PType.DYNAMIC)
                 PType.codes().filterNot { it == PType.DYNAMIC }.forEach {

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RexConverter.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RexConverter.kt
@@ -89,7 +89,6 @@ import org.partiql.planner.internal.typer.PlanTyper.Companion.toCType
 import org.partiql.planner.internal.utils.DateTimeUtils
 import org.partiql.spi.catalog.Identifier
 import org.partiql.spi.value.Datum
-import org.partiql.spi.value.DatumReader
 import org.partiql.types.PType
 import org.partiql.value.datetime.DateTimeValue
 import java.math.BigDecimal
@@ -216,10 +215,7 @@ internal object RexConverter {
             if (node.encoding != "ion") {
                 throw IllegalArgumentException("unsupported encoding ${node.encoding}")
             }
-            // TODO: Does this result in a Datum of type variant?
-            val v = DatumReader.ion(node.value.byteInputStream())
-            val datum = v.next() ?: error("Expected a single value")
-            v.next()?.let { throw TypeCheckException("Expected a single value") }
+            val datum = Datum.ion(node.value)
             val type = CompilerType(datum.type)
             return rex(type, rexOpLit(datum))
         }

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt
@@ -224,7 +224,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
             val kType = PType.string()
 
             // Check Root (Dynamic)
-            if (rex.type.code() == PType.DYNAMIC) {
+            if (rex.type.code() == PType.DYNAMIC || rex.type.code() == PType.VARIANT) {
                 val type = ctx!!.copyWithSchema(listOf(kType, PType.dynamic()).toCType())
                 return rel(type, op)
             }
@@ -624,7 +624,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
             }
 
             // Check if Root is DYNAMIC
-            if (root.type.code() == PType.DYNAMIC) {
+            if (root.type.code() == PType.DYNAMIC || root.type.code() == PType.VARIANT) {
                 return Rex(CompilerType(PType.dynamic()), Rex.Op.Path.Index(root, key))
             }
 
@@ -653,7 +653,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
             }
 
             // Check if Root is DYNAMIC
-            if (root.type.code() == PType.DYNAMIC) {
+            if (root.type.code() == PType.DYNAMIC || root.type.code() == PType.VARIANT) {
                 return Rex(CompilerType(PType.dynamic()), Rex.Op.Path.Key(root, key))
             }
 
@@ -685,7 +685,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
             val root = visitRex(node.root, node.root.type)
 
             // Check if Root is DYNAMIC
-            if (root.type.code() == PType.DYNAMIC) {
+            if (root.type.code() == PType.DYNAMIC || root.type.code() == PType.VARIANT) {
                 return Rex(CompilerType(PType.dynamic()), Rex.Op.Path.Symbol(root, node.key))
             }
 
@@ -923,7 +923,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
          * Hence, we permit Static Type BOOL, Static Type NULL, Static Type Missing here.
          */
         private fun canBeBoolean(type: CompilerType): Boolean {
-            return type.code() == PType.DYNAMIC || type.code() == PType.BOOL
+            return type.code() == PType.DYNAMIC || type.code() == PType.VARIANT || type.code() == PType.BOOL
         }
 
         /**
@@ -1040,7 +1040,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
          * Calculate output type of a scalar subquery.
          */
         private fun visitRexOpSubqueryScalar(subquery: Rex.Op.Subquery, cons: CompilerType): Rex {
-            if (cons.code() == PType.DYNAMIC) {
+            if (cons.code() == PType.DYNAMIC || cons.code() == PType.VARIANT) {
                 return Rex(PType.dynamic().toCType(), subquery)
             }
             if (cons.code() != PType.ROW) {
@@ -1186,7 +1186,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
                 when (arg.code()) {
                     PType.ROW -> fields.addAll(arg.fields!!)
                     PType.STRUCT -> structIsOpen = true
-                    PType.DYNAMIC -> containsDynamic = true
+                    PType.DYNAMIC, PType.VARIANT -> containsDynamic = true
                     PType.UNKNOWN -> structIsOpen = true
                     else -> containsNonStruct = true
                 }

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/Scope.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/Scope.kt
@@ -144,6 +144,7 @@ internal data class Scope(
             PType.ROW -> this.fields.any { name.matches(it.name) }
             PType.STRUCT -> null
             PType.DYNAMIC -> null
+            PType.VARIANT -> null
             else -> false
         }
     }

diff --git a/partiql-spi/api/partiql-spi.api b/partiql-spi/api/partiql-spi.api
@@ -501,6 +501,7 @@ public abstract interface class org/partiql/spi/value/Datum : java/lang/Iterable
 	public fun getTimestamp ()Lorg/partiql/value/datetime/Timestamp;
 	public abstract fun getType ()Lorg/partiql/types/PType;
 	public static fun integer (I)Lorg/partiql/spi/value/Datum;
+	public static fun ion (Ljava/lang/String;)Lorg/partiql/spi/value/Datum;
 	public fun isMissing ()Z
 	public fun isNull ()Z
 	public fun iterator ()Ljava/util/Iterator;

diff --git a/partiql-spi/src/main/java/org/partiql/spi/value/Datum.java b/partiql-spi/src/main/java/org/partiql/spi/value/Datum.java
@@ -1,13 +1,18 @@
 package org.partiql.spi.value;
 
+import com.amazon.ionelement.api.AnyElement;
+import com.amazon.ionelement.api.ElementLoader;
+import com.amazon.ionelement.api.IonElementLoader;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.partiql.errors.DataException;
+import org.partiql.spi.value.ion.IonVariant;
 import org.partiql.types.PType;
 import org.partiql.value.datetime.Date;
 import org.partiql.value.datetime.Time;
 import org.partiql.value.datetime.Timestamp;
 
+import java.io.InputStream;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.math.MathContext;
@@ -198,7 +203,7 @@ default short getShort() {
      *                                       {@link #isNull()} returns false before attempting to invoke this method.
      */
     default int getInt() {
-        throw new UnsupportedOperationException();
+        throw new UnsupportedOperationException("Has type: " + getType() + " and class " + this.getClass().getName());
     }
 
     /**
@@ -572,6 +577,13 @@ static Datum struct(@NotNull Iterable<Field> values) {
         return new DatumStruct(values);
     }
 
+    @NotNull
+    static Datum ion(@NotNull String value) {
+        IonElementLoader loader = ElementLoader.createIonElementLoader();
+        AnyElement element = loader.loadSingleElement(value);
+        return new IonVariant(element);
+    }
+
     /**
      * Comparator for PartiQL values.
      * <p>