Rename DataRow.near function to relative and change its return type to DataFrame.

Author: nikitinas

docs/StardustDocs/topics/DataRow.md

@@ -17,8 +17,8 @@
 * `namedValues(): List<NameValuePair<Any?>>` — list of name-value pairs where `name` is a column name and `value` is cell value
 * `namedValuesOf<T>(): List<NameValuePair<T>>` — list of name-value pairs where value has given type 
 * `getRow(Int): DataRow` — row from `DataFrame` by row index
-* `near(Iterable<Int>): Sequence<DataRow>` — sequence of the nearest rows by relative index: `near(-1..1)` will return previous, current and next row. Requested indices will be coerced to valid range and invalid indices will be skipped
-* `rows(Iterable<Int>): Sequence<DataRow>` — sequence of the rows by absolute index. Requested indices are not coerced to valid boundaries and you should care about it
+* `getRows(Iterable<Int>): DataFrame` — dataframe with subset of rows selected by absolute row index. 
+* `relative(Iterable<Int>): DataFrame` — dataframe with subset of rows selected by relative row index: `relative(-1..1)` will return previous, current and next row. Requested indices will be coerced to the valid range and invalid indices will be skipped
 * `get(column): T` — cell value by this row and given `column`
 * `df()` — `DataFrame` that current row belongs to
 

examples/jupyter-notebooks/puzzles/40 puzzles.ipynb

@@ -2408,7 +2408,7 @@
    "source": [
     "df.add(\"id\"){ index() }\n",
     "    .groupBy { groups }.add(\"res\") {\n",
-    "        near(-2..0).map { it.value }.filter { !it.isNaN() }.average()\n",
+    "        relative(-2..0).value.filter { !it.isNaN() }.mean()\n",
     "    }.concat()\n",
     "    .sortBy(\"id\")\n",
     "    .remove(\"id\")"
@@ -3118,4 +3118,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
+}

src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/DataRowApi.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.impl.columnName
 import org.jetbrains.kotlinx.dataframe.impl.owner
 import org.jetbrains.kotlinx.dataframe.index
+import org.jetbrains.kotlinx.dataframe.indices
 import org.jetbrains.kotlinx.dataframe.ncol
 import org.jetbrains.kotlinx.dataframe.nrow
 import org.jetbrains.kotlinx.dataframe.type
@@ -55,6 +56,9 @@ public fun AnyRow.columnTypes(): List<KType> = df().columnTypes()
 
 public fun <T> DataRow<T>.getRow(index: Int): DataRow<T> = getRowOrNull(index)!!
 
+public fun <T> DataRow<T>.getRows(indices: Iterable<Int>): DataFrame<T> = df().getRows(indices)
+public fun <T> DataRow<T>.getRows(indices: IntRange): DataFrame<T> = df().getRows(indices)
+
 public fun <T> DataRow<T>.getRowOrNull(index: Int): DataRow<T>? {
     val df = df()
     return if (index >= 0 && index < df.nrow) df[index] else null
@@ -71,11 +75,11 @@ public fun <T> DataRow<T>.next(): DataRow<T>? {
     return if (index < df.nrow - 1) df[index + 1] else null
 }
 
-public fun <T> DataRow<T>.rows(absoluteIndices: Iterable<Int>): Sequence<DataRow<T>> =
-    absoluteIndices.asSequence().map { getRow(index + it) }
+public fun <T> DataRow<T>.relative(relativeIndices: Iterable<Int>): DataFrame<T> =
+    getRows(relativeIndices.mapNotNull { (index + it).let { if (it >= 0 && it < df().rowsCount()) it else null } })
 
-public fun <T> DataRow<T>.near(relativeIndices: Iterable<Int>): Sequence<DataRow<T>> =
-    relativeIndices.asSequence().mapNotNull { getRowOrNull(index + it) }
+public fun <T> DataRow<T>.relative(relativeIndices: IntRange): DataFrame<T> =
+    getRows((relativeIndices.start + index).coerceIn(df().indices)..(relativeIndices.endInclusive + index).coerceIn(df().indices))
 
 internal fun <T> DataRow<T>.movingAverage(k: Int, expression: RowExpression<T, Number>): Double {
     var count = 0

src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/HardTests.kt

@@ -115,7 +115,7 @@ class HardTests {
 
         df.add("id") { index() }
             .groupBy { groups }.add("res") {
-                round(near(-2..0).map { it[value] }.filter { !it.isNaN() }.average())
+                round(relative(-2..0)[value].filter { !it.isNaN() }.mean())
             }.concat()
             .sortBy("id")
             .remove("id") shouldBe expected

src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt

@@ -102,7 +102,6 @@ import org.jetbrains.kotlinx.dataframe.api.moveToLeft
 import org.jetbrains.kotlinx.dataframe.api.moveToRight
 import org.jetbrains.kotlinx.dataframe.api.name
 import org.jetbrains.kotlinx.dataframe.api.named
-import org.jetbrains.kotlinx.dataframe.api.near
 import org.jetbrains.kotlinx.dataframe.api.notNull
 import org.jetbrains.kotlinx.dataframe.api.nullable
 import org.jetbrains.kotlinx.dataframe.api.parse
@@ -1876,11 +1875,6 @@ class DataFrameTests : BaseTest() {
         stringCols.columnNames() shouldBe listOf("name", "city")
     }
 
-    @Test
-    fun neighbours() {
-        typed[2].near(-1..1).toList() shouldBe listOf(typed[1], typed[2], typed[3])
-    }
-
     @Test
     fun `get row value by expression`() {
         val expression: RowExpression<Person, Int> = { it.age * 2 }

src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataRowTests.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlinx.dataframe.api.namedValues
 import org.jetbrains.kotlinx.dataframe.api.namedValuesOf
 import org.jetbrains.kotlinx.dataframe.api.next
 import org.jetbrains.kotlinx.dataframe.api.prev
+import org.jetbrains.kotlinx.dataframe.api.relative
 import org.jetbrains.kotlinx.dataframe.api.rowMean
 import org.jetbrains.kotlinx.dataframe.api.rowStd
 import org.jetbrains.kotlinx.dataframe.api.rowSum
@@ -85,4 +86,11 @@ class DataRowTests : BaseTest() {
             if (firstNullIndex == -1) null else columnNames()[firstNullIndex]
         }
     }
+
+    @Test
+    fun relativeTest() {
+        typed[1].relative(0..0) shouldBe typed[1..1]
+        typed[1].relative(-2..2) shouldBe typed[0..3]
+        typed[1].relative(listOf(2, -1, -3, 0)) shouldBe typed[3, 0, 1]
+    }
 }