Use Iterable for PartiQLValue collections rather than Sequence (#1305)

partiql · Dec 14, 2023 · 79ee005 · 79ee005 · github-actions · Dec 15, 2023
1 parent 8076e6c
commit 79ee005
Show file tree

Hide file tree

Showing 14 changed files with 298 additions and 260 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -47,6 +47,7 @@ Thank you to all who have contributed!
   - The new plan is fully resolved and typed.
   - Operators will be converted to function call. 
 - Changes the return type of `filter_distinct` to a list if input collection is list
+- Changes the `PartiQLValue` collections to implement Iterable rather than Sequence, allowing for multiple consumption.
 
 ### Deprecated
 

diff --git a/partiql-cli/src/main/kotlin/org/partiql/cli/utils/ServiceLoaderUtil.kt b/partiql-cli/src/main/kotlin/org/partiql/cli/utils/ServiceLoaderUtil.kt
@@ -268,28 +268,42 @@ class ServiceLoaderUtil {
                 PartiQLValueType.INTERVAL -> TODO()
 
                 PartiQLValueType.BAG -> {
-                    (partiqlValue as? BagValue<*>)?.elements?.map { PartiQLtoExprValue(it) }?.let { newBag(it) }
-                        ?: ExprValue.nullValue
+                    if (partiqlValue.isNull) {
+                        ExprValue.nullValue
+                    } else {
+                        newBag((partiqlValue as? BagValue<*>)!!.map { PartiQLtoExprValue(it) })
+                    }
                 }
 
                 PartiQLValueType.LIST -> {
-                    (partiqlValue as? ListValue<*>)?.elements?.map { PartiQLtoExprValue(it) }?.let { newList(it) }
-                        ?: ExprValue.nullValue
+                    if (partiqlValue.isNull) {
+                        ExprValue.nullValue
+                    } else {
+                        newList((partiqlValue as? ListValue<*>)!!.map { PartiQLtoExprValue(it) })
+                    }
                 }
 
                 PartiQLValueType.SEXP -> {
-                    (partiqlValue as? SexpValue<*>)?.elements?.map { PartiQLtoExprValue(it) }?.let { newSexp(it) }
-                        ?: ExprValue.nullValue
+                    if (partiqlValue.isNull) {
+                        ExprValue.nullValue
+                    } else {
+                        newSexp((partiqlValue as? SexpValue<*>)!!.map { PartiQLtoExprValue(it) })
+                    }
                 }
 
                 PartiQLValueType.STRUCT -> {
-                    (partiqlValue as? StructValue<*>)?.fields?.map {
-                        PartiQLtoExprValue(it.second).namedValue(
-                            newString(
-                                it.first
+                    if (partiqlValue.isNull) {
+                        ExprValue.nullValue
+                    } else {
+                        val entries = (partiqlValue as? StructValue<*>)!!.entries
+                        entries.map {
+                            PartiQLtoExprValue(it.second).namedValue(
+                                newString(
+                                    it.first
+                                )
                             )
-                        )
-                    }?.let { newStruct(it, StructOrdering.ORDERED) } ?: ExprValue.nullValue
+                        }.let { newStruct(it, StructOrdering.ORDERED) }
+                    }
                 }
 
                 PartiQLValueType.DECIMAL -> TODO()
@@ -447,7 +461,7 @@ class ServiceLoaderUtil {
                 PartiQLValueType.INTERVAL -> TODO()
                 PartiQLValueType.BAG -> when (exprValue.type) {
                     ExprValueType.NULL -> bagValue(null)
-                    ExprValueType.BAG -> bagValue(exprValue.map { ExprToPartiQLValue(it, ExprToPartiQLValueType(it)) }.asSequence())
+                    ExprValueType.BAG -> bagValue(exprValue.map { ExprToPartiQLValue(it, ExprToPartiQLValueType(it)) })
                     else -> throw ExprToPartiQLValueTypeMismatchException(
                         PartiQLValueType.BAG, ExprToPartiQLValueType(exprValue)
                     )
@@ -459,7 +473,7 @@ class ServiceLoaderUtil {
                             ExprToPartiQLValue(
                                 it, ExprToPartiQLValueType(it)
                             )
-                        }.asSequence()
+                        }
                     )
                     else -> throw ExprToPartiQLValueTypeMismatchException(
                         PartiQLValueType.LIST, ExprToPartiQLValueType(exprValue)
@@ -472,7 +486,7 @@ class ServiceLoaderUtil {
                             ExprToPartiQLValue(
                                 it, ExprToPartiQLValueType(it)
                             )
-                        }.asSequence()
+                        }
                     )
                     else -> throw ExprToPartiQLValueTypeMismatchException(
                         PartiQLValueType.SEXP, ExprToPartiQLValueType(exprValue)
@@ -485,7 +499,7 @@ class ServiceLoaderUtil {
                             Pair(
                                 it.name?.stringValue() ?: "", ExprToPartiQLValue(it, ExprToPartiQLValueType(it))
                             )
-                        }.asSequence()
+                        }
                     )
                     else -> throw ExprToPartiQLValueTypeMismatchException(
                         PartiQLValueType.STRUCT, ExprToPartiQLValueType(exprValue)

diff --git a/partiql-types/src/main/kotlin/org/partiql/value/PartiQL.kt b/partiql-types/src/main/kotlin/org/partiql/value/PartiQL.kt
@@ -38,12 +38,12 @@ import org.partiql.value.impl.Int64ValueImpl
 import org.partiql.value.impl.Int8ValueImpl
 import org.partiql.value.impl.IntValueImpl
 import org.partiql.value.impl.IntervalValueImpl
+import org.partiql.value.impl.IterableStructValueImpl
 import org.partiql.value.impl.ListValueImpl
 import org.partiql.value.impl.MapStructValueImpl
 import org.partiql.value.impl.MissingValueImpl
 import org.partiql.value.impl.MultiMapStructValueImpl
 import org.partiql.value.impl.NullValueImpl
-import org.partiql.value.impl.SequenceStructValueImpl
 import org.partiql.value.impl.SexpValueImpl
 import org.partiql.value.impl.StringValueImpl
 import org.partiql.value.impl.SymbolValueImpl
@@ -344,10 +344,25 @@ public fun intervalValue(
 @JvmOverloads
 @PartiQLValueExperimental
 public fun <T : PartiQLValue> bagValue(
-    elements: Sequence<T>?,
+    elements: Iterable<T>?,
     annotations: Annotations = emptyList(),
 ): BagValue<T> = BagValueImpl(elements, annotations.toPersistentList())
 
+/**
+ * BAG type value.
+ *
+ * @param T
+ * @param elements
+ * @param annotations
+ * @return
+ */
+@JvmOverloads
+@PartiQLValueExperimental
+public fun <T : PartiQLValue> bagValue(
+    vararg elements: T,
+    annotations: Annotations = emptyList(),
+): BagValue<T> = BagValueImpl(elements.asIterable(), annotations.toPersistentList())
+
 /**
  * LIST type value.
  *
@@ -359,10 +374,25 @@ public fun <T : PartiQLValue> bagValue(
 @JvmOverloads
 @PartiQLValueExperimental
 public fun <T : PartiQLValue> listValue(
-    elements: Sequence<T>?,
+    elements: Iterable<T>?,
     annotations: Annotations = emptyList(),
 ): ListValue<T> = ListValueImpl(elements, annotations.toPersistentList())
 
+/**
+ * LIST type value.
+ *
+ * @param T
+ * @param elements
+ * @param annotations
+ * @return
+ */
+@JvmOverloads
+@PartiQLValueExperimental
+public fun <T : PartiQLValue> listValue(
+    vararg elements: T,
+    annotations: Annotations = emptyList(),
+): ListValue<T> = ListValueImpl(elements.asIterable(), annotations.toPersistentList())
+
 /**
  * SEXP type value.
  *
@@ -374,12 +404,42 @@ public fun <T : PartiQLValue> listValue(
 @JvmOverloads
 @PartiQLValueExperimental
 public fun <T : PartiQLValue> sexpValue(
-    elements: Sequence<T>?,
+    elements: Iterable<T>?,
     annotations: Annotations = emptyList(),
 ): SexpValue<T> = SexpValueImpl(elements, annotations.toPersistentList())
 
 /**
- * STRUCT type value.
+ * SEXP type value.
+ *
+ * @param T
+ * @param elements
+ * @param annotations
+ * @return
+ */
+@JvmOverloads
+@PartiQLValueExperimental
+public fun <T : PartiQLValue> sexpValue(
+    vararg elements: T,
+    annotations: Annotations = emptyList(),
+): SexpValue<T> = SexpValueImpl(elements.asIterable(), annotations.toPersistentList())
+
+/**
+ * Create a PartiQL struct value backed by an iterable of key-value field pairs.
+ *
+ * @param T
+ * @param fields
+ * @param annotations
+ * @return
+ */
+@JvmOverloads
+@PartiQLValueExperimental
+public fun <T : PartiQLValue> structValue(
+    fields: Iterable<Pair<String, T>>?,
+    annotations: Annotations = emptyList(),
+): StructValue<T> = IterableStructValueImpl(fields, annotations.toPersistentList())
+
+/**
+ * Create a PartiQL struct value backed by an iterable of key-value field pairs.
  *
  * @param T
  * @param fields
@@ -389,12 +449,13 @@ public fun <T : PartiQLValue> sexpValue(
 @JvmOverloads
 @PartiQLValueExperimental
 public fun <T : PartiQLValue> structValue(
-    fields: Sequence<Pair<String, T>>?,
+    vararg fields: Pair<String, T>,
     annotations: Annotations = emptyList(),
-): StructValue<T> = SequenceStructValueImpl(fields, annotations.toPersistentList())
+): StructValue<T> = IterableStructValueImpl(fields.toList(), annotations.toPersistentList())
 
 /**
- * STRUCT type value.
+ * Create a PartiQL struct value backed by a multimap of keys with a list of values. This supports having multiple
+ * values per key, while improving lookup performance compared to using an iterable.
  *
  * @param T
  * @param fields
@@ -403,13 +464,14 @@ public fun <T : PartiQLValue> structValue(
  */
 @JvmOverloads
 @PartiQLValueExperimental
-public fun <T : PartiQLValue> structValueWithDuplicates(
+public fun <T : PartiQLValue> structValueMultiMap(
     fields: Map<String, Iterable<T>>?,
     annotations: Annotations = emptyList(),
 ): StructValue<T> = MultiMapStructValueImpl(fields, annotations.toPersistentList())
 
 /**
- * STRUCT type value.
+ * Create a PartiQL struct value backed by a map of keys with a list of values. This does not support having multiple
+ * values per key, but uses a Java HashMap for quicker lookup than an iterable backed StructValue.
  *
  * @param T
  * @param fields
@@ -418,7 +480,7 @@ public fun <T : PartiQLValue> structValueWithDuplicates(
  */
 @JvmOverloads
 @PartiQLValueExperimental
-public fun <T : PartiQLValue> structValueNoDuplicates(
+public fun <T : PartiQLValue> structValueMap(
     fields: Map<String, T>?,
     annotations: Annotations = emptyList(),
 ): StructValue<T> = MapStructValueImpl(fields, annotations.toPersistentList())

diff --git a/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt b/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt
@@ -65,14 +65,11 @@ public sealed interface ScalarValue<T> : PartiQLValue {
 }
 
 @PartiQLValueExperimental
-public sealed interface CollectionValue<T : PartiQLValue> : PartiQLValue, Sequence<T> {
-
-    public val elements: Sequence<T>?
+public sealed interface CollectionValue<T : PartiQLValue> : PartiQLValue, Iterable<T> {
 
     override val isNull: Boolean
-        get() = elements == null
 
-    override fun iterator(): Iterator<T> = elements!!.iterator()
+    override fun iterator(): Iterator<T>
 
     override fun copy(annotations: Annotations): CollectionValue<T>
 
@@ -388,8 +385,8 @@ public abstract class BagValue<T : PartiQLValue> : CollectionValue<T> {
         if (this.isNull || other.isNull) return this.isNull == other.isNull
 
         // both not null, compare values
-        val lhs = this.elements!!.toList()
-        val rhs = other.elements!!.toList()
+        val lhs = this.toList()
+        val rhs = other.toList()
         // this is incorrect as it assumes ordered-ness, but we don't have a sort or hash yet
         return lhs == rhs
     }
@@ -421,8 +418,8 @@ public abstract class ListValue<T : PartiQLValue> : CollectionValue<T> {
         if (this.isNull || other.isNull) return this.isNull == other.isNull
 
         // both not null, compare values
-        val lhs = this.elements!!.toList()
-        val rhs = other.elements!!.toList()
+        val lhs = this.toList()
+        val rhs = other.toList()
         return lhs == rhs
     }
 
@@ -452,8 +449,8 @@ public abstract class SexpValue<T : PartiQLValue> : CollectionValue<T> {
         if (this.isNull || other.isNull) return this.isNull == other.isNull
 
         // both not null, compare values
-        val lhs = this.elements!!.toList()
-        val rhs = other.elements!!.toList()
+        val lhs = this.toList()
+        val rhs = other.toList()
         return lhs == rhs
     }
 
@@ -464,16 +461,15 @@ public abstract class SexpValue<T : PartiQLValue> : CollectionValue<T> {
 }
 
 @PartiQLValueExperimental
-public abstract class StructValue<T : PartiQLValue> : PartiQLValue, Sequence<Pair<String, T>> {
+public abstract class StructValue<T : PartiQLValue> : PartiQLValue {
 
     override val type: PartiQLValueType = PartiQLValueType.STRUCT
 
-    public abstract val fields: Sequence<Pair<String, T>>?
+    public abstract val fields: Iterable<String>
 
-    override val isNull: Boolean
-        get() = fields == null
+    public abstract val values: Iterable<T>
 
-    override fun iterator(): Iterator<Pair<String, T>> = fields!!.iterator()
+    public abstract val entries: Iterable<Pair<String, T>>
 
     public abstract operator fun get(key: String): T?
 
@@ -486,9 +482,7 @@ public abstract class StructValue<T : PartiQLValue> : PartiQLValue, Sequence<Pai
     abstract override fun withoutAnnotations(): StructValue<T>
 
     /**
-     * See equality of IonElement StructElementImpl
-     *
-     * https://github.com/amazon-ion/ion-element-kotlin/blob/master/src/com/amazon/ionelement/impl/StructElementImpl.kt
+     * Checks equality of struct entries, ignoring ordering.
      *
      * @param other
      * @return
@@ -502,15 +496,15 @@ public abstract class StructValue<T : PartiQLValue> : PartiQLValue, Sequence<Pai
         if (this.isNull || other.isNull) return this.isNull == other.isNull
 
         // both not null, compare fields
-        val lhs = this.fields!!.groupBy({ it.first }, { it.second })
-        val rhs = other.fields!!.groupBy({ it.first }, { it.second })
+        val lhs = this.entries.asIterable().groupBy({ it.first }, { it.second })
+        val rhs = other.entries.asIterable().groupBy({ it.first }, { it.second })
 
         // check size
         if (lhs.size != rhs.size) return false
         if (lhs.keys != rhs.keys) return false
 
         // check values
-        lhs.forEach { (key, values) ->
+        lhs.entries.forEach { (key, values) ->
             val lGroup: Map<PartiQLValue, Int> = values.groupingBy { it }.eachCount()
             val rGroup: Map<PartiQLValue, Int> = rhs[key]!!.groupingBy { it }.eachCount()
             if (lGroup != rGroup) return false

diff --git a/partiql-types/src/main/kotlin/org/partiql/value/helpers/ToIon.kt b/partiql-types/src/main/kotlin/org/partiql/value/helpers/ToIon.kt
@@ -258,31 +258,31 @@ internal object ToIon : PartiQLValueBaseVisitor<IonElement, Unit>() {
     override fun visitInterval(v: IntervalValue, ctx: Unit): IonElement = TODO("Not Yet supported")
 
     override fun visitBag(v: BagValue<*>, ctx: Unit): IonElement = v.annotate {
-        when (val elements = v.elements) {
-            null -> ionNull(ElementType.LIST)
-            else -> ionListOf(elements.map { it.accept(ToIon, Unit) }.toList())
+        when (v.isNull) {
+            true -> ionNull(ElementType.LIST)
+            else -> ionListOf(v.map { it.accept(ToIon, Unit) }.toList())
         }
     }.withAnnotations(BAG_ANNOTATION)
 
     override fun visitList(v: ListValue<*>, ctx: Unit): IonElement = v.annotate {
-        when (val elements = v.elements) {
-            null -> ionNull(ElementType.LIST)
-            else -> ionListOf(elements.map { it.accept(ToIon, Unit) }.toList())
+        when (v.isNull) {
+            true -> ionNull(ElementType.LIST)
+            else -> ionListOf(v.map { it.accept(ToIon, Unit) }.toList())
         }
     }
 
     override fun visitSexp(v: SexpValue<*>, ctx: Unit): IonElement = v.annotate {
-        when (val elements = v.elements) {
-            null -> ionNull(ElementType.SEXP)
-            else -> ionSexpOf(elements.map { it.accept(ToIon, Unit) }.toList())
+        when (v.isNull) {
+            true -> ionNull(ElementType.SEXP)
+            else -> ionSexpOf(v.map { it.accept(ToIon, Unit) }.toList())
         }
     }
 
     override fun visitStruct(v: StructValue<*>, ctx: Unit): IonElement = v.annotate {
-        when (val fields = v.fields) {
-            null -> ionNull(ElementType.STRUCT)
+        when (v.isNull) {
+            true -> ionNull(ElementType.STRUCT)
             else -> {
-                val ionFields = fields.map {
+                val ionFields = entries.map {
                     val fk = it.first
                     val fv = it.second.accept(ToIon, ctx)
                     field(fk, fv)