Addresses PR comments #2

Updates Accumulator to check for variant on number types Updates internal docs for clarification Simplifies IonVariant conversion for BigIntegers
partiql · Dec 27, 2024 · 5a9b63d · 5a9b63d
1 parent 0f565af
commit 5a9b63d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/helpers/DatumUtils.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/helpers/DatumUtils.kt
@@ -9,7 +9,7 @@ internal object DatumUtils {
      * Calls [Datum.lower] if the datum is a variant, otherwise returns the datum. If you don't know whether the value
      * is of type [PType.VARIANT], you should use [Datum.lowerSafe] before invoking whatever methods you intend to use.
      * This is essentially a workaround for the fact that we currently don't know whether a particular expression will be
-     * [PType.VARIANT] or not. This can eventually be optimized to accommodate this.
+     * [PType.VARIANT] or not. The planner/plan can eventually be optimized to accommodate this.
      */
     internal fun Datum.lowerSafe(): Datum {
         return when (this.type.code()) {

diff --git a/partiql-spi/src/main/kotlin/org/partiql/spi/function/builtins/internal/Accumulator.kt b/partiql-spi/src/main/kotlin/org/partiql/spi/function/builtins/internal/Accumulator.kt
@@ -45,6 +45,9 @@ internal fun comparisonAccumulator(comparator: Comparator<Datum>): (Datum?, Datu
     }
 
 internal fun checkIsNumberType(funcName: String, value: Datum) {
+    if (value.type.code() == PType.VARIANT) {
+        return checkIsNumberType(funcName, value.lower())
+    }
     if (!value.type.isNumber()) {
         throw TypeCheckException("Expected NUMBER but received ${value.type}.")
     }

diff --git a/partiql-spi/src/main/kotlin/org/partiql/spi/value/ion/IonVariant.kt b/partiql-spi/src/main/kotlin/org/partiql/spi/value/ion/IonVariant.kt
@@ -16,6 +16,7 @@ import com.amazon.ionelement.api.ElementType.STRING
 import com.amazon.ionelement.api.ElementType.STRUCT
 import com.amazon.ionelement.api.ElementType.SYMBOL
 import com.amazon.ionelement.api.ElementType.TIMESTAMP
+import com.amazon.ionelement.api.IntElementSize
 import org.partiql.spi.value.Datum
 import org.partiql.spi.value.Field
 import org.partiql.types.PType
@@ -63,12 +64,16 @@ internal class IonVariant(private var value: AnyElement) : Datum {
             BLOB -> Datum.blob(value.blobValue.copyOfBytes())
             TIMESTAMP -> Datum.timestamp(DateTimeValue.timestamp(value.timestampValue))
             INT -> {
-                val bigInteger = value.bigIntegerValue
-                when {
-                    bigInteger < Int.MAX_VALUE.toBigInteger() && bigInteger > Int.MIN_VALUE.toBigInteger() -> Datum.integer(bigInteger.toInt())
-                    bigInteger < Long.MAX_VALUE.toBigInteger() && bigInteger > Long.MIN_VALUE.toBigInteger() -> Datum.bigint(bigInteger.toLong())
-                    else -> {
-                        val dec = bigInteger.toBigDecimal()
+                when (value.integerSize) {
+                    IntElementSize.LONG -> {
+                        val long = value.longValue
+                        when (long < Int.MAX_VALUE && long > Int.MIN_VALUE) {
+                            true -> Datum.integer(long.toInt())
+                            false -> Datum.bigint(long)
+                        }
+                    }
+                    IntElementSize.BIG_INTEGER -> {
+                        val dec = value.bigIntegerValue.toBigDecimal()
                         Datum.decimal(dec, dec.precision(), 0)
                     }
                 }