flatMapIndexed

Author: vincentlauvlwj

ktorm-core/src/main/kotlin/me/liuwj/ktorm/dsl/Query.kt

@@ -550,6 +550,30 @@ public inline fun <R, C : MutableCollection<in R>> Query.flatMapTo(
     return destination
 }
 
+/**
+ * Return a single list of all elements yielded from results of [transform] function being invoked on each row
+ * and its index of the query.
+ *
+ * @since 3.1.0
+ */
+public inline fun <R> Query.flatMapIndexed(transform: (index: Int, row: QueryRowSet) -> Iterable<R>): List<R> {
+    return flatMapIndexedTo(ArrayList(), transform)
+}
+
+/**
+ * Append all elements yielded from results of [transform] function being invoked on each row and its index
+ * of the query, to the given [destination].
+ *
+ * @since 3.1.0
+ */
+public inline fun <R, C : MutableCollection<in R>> Query.flatMapIndexedTo(
+    destination: C,
+    transform: (index: Int, row: QueryRowSet) -> Iterable<R>
+): C {
+    var index = 0
+    return flatMapTo(destination) { transform(index++, it) }
+}
+
 /**
  * Return a [Map] containing key-value pairs provided by [transform] function applied to rows of the query.
  *

ktorm-core/src/main/kotlin/me/liuwj/ktorm/entity/EntitySequence.kt

@@ -460,6 +460,34 @@ public inline fun <E : Any, R, C : MutableCollection<in R>> EntitySequence<E, *>
     return destination
 }
 
+/**
+ * Return a single list of all elements yielded from results of [transform] function being invoked
+ * on each element and its index of original sequence.
+ *
+ * The operation is terminal.
+ *
+ * @since 3.1.0
+ */
+public inline fun <E : Any, R> EntitySequence<E, *>.flatMapIndexed(transform: (index: Int, E) -> Iterable<R>): List<R> {
+    return flatMapIndexedTo(ArrayList(), transform)
+}
+
+/**
+ * Append all elements yielded from results of [transform] function being invoked on each element
+ * and its index of original sequence, to the given [destination].
+ *
+ * The operation is terminal.
+ *
+ * @since 3.1.0
+ */
+public inline fun <E : Any, R, C : MutableCollection<in R>> EntitySequence<E, *>.flatMapIndexedTo(
+    destination: C,
+    transform: (index: Int, E) -> Iterable<R>
+): C {
+    var index = 0
+    return flatMapTo(destination) { transform(index++, it) }
+}
+
 /**
  * Customize the selected columns of the internal query by the given [columnSelector] function, and return a [List]
  * containing the query results.

ktorm-core/src/test/kotlin/me/liuwj/ktorm/dsl/QueryTest.kt

@@ -285,4 +285,18 @@ class QueryTest : BaseTest() {
             }
         }
     }
+
+    @Test
+    fun testFlatMap() {
+        val names = database
+            .from(Employees)
+            .select(Employees.name)
+            .where { Employees.departmentId eq 1 }
+            .orderBy(Employees.salary.desc())
+            .flatMapIndexed { index, row -> listOf("$index:${row.getString(1)}") }
+
+        assert(names.size == 2)
+        assert(names[0] == "0:vince")
+        assert(names[1] == "1:marry")
+    }
 }

ktorm-core/src/test/kotlin/me/liuwj/ktorm/entity/EntitySequenceTest.kt

@@ -196,4 +196,15 @@ class EntitySequenceTest : BaseTest() {
         assert(names.size == 4)
         assert(names[0] == "vince")
     }
+
+    @Test
+    fun testFlatMap() {
+        val names = database.employees
+            .sortedBy { it.id }
+            .flatMapIndexed { index, employee -> listOf("$index:${employee.name}") }
+
+        println(names)
+        assert(names.size == 4)
+        assert(names[0] == "0:vince")
+    }
 }