From 5a9b63dc9d317a74b2ab89d17534c8c2c36ccdc7 Mon Sep 17 00:00:00 2001
From: John Ed Quinn <johqunn@amazon.com>
Date: Fri, 27 Dec 2024 12:51:10 -0800
Subject: [PATCH] Addresses PR comments #2

Updates Accumulator to check for variant on number types

Updates internal docs for clarification

Simplifies IonVariant conversion for BigIntegers
---
 .../partiql/eval/internal/helpers/DatumUtils.kt |  2 +-
 .../function/builtins/internal/Accumulator.kt   |  3 +++
 .../org/partiql/spi/value/ion/IonVariant.kt     | 17 +++++++++++------
 3 files changed, 15 insertions(+), 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
index 3e3cf110d..369f255e7 100644
--- 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
index 9e3ca87d3..70ced293f 100644
--- 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
index 23ee66330..90d3a586b 100644
--- 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)
                     }
                 }