Merge pull request #1303 from Kotlin/AddExpression-in-add-dsl

Makes `prev()?.newValue()` work in Add DSL

Author: Jolanrensen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlinx.dataframe.ColumnsContainer
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
-import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.Selector
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
 import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
@@ -259,27 +258,27 @@ public class AddDsl<T>(
     internal inline fun <reified R> add(
         name: String,
         infer: Infer = Infer.Nulls,
-        noinline expression: RowExpression<T, R>,
+        noinline expression: AddExpression<T, R>,
     ): Boolean = add(df.mapToColumn(name, infer, expression))
 
     public inline fun <reified R> expr(
         infer: Infer = Infer.Nulls,
-        noinline expression: RowExpression<T, R>,
+        noinline expression: AddExpression<T, R>,
     ): DataColumn<R> = df.mapToColumn("", infer, expression)
 
     @Interpretable("From")
-    public inline infix fun <reified R> String.from(noinline expression: RowExpression<T, R>): Boolean =
+    public inline infix fun <reified R> String.from(noinline expression: AddExpression<T, R>): Boolean =
         add(this, Infer.Nulls, expression)
 
     // TODO: use path instead of name
     @Deprecated(DEPRECATED_ACCESS_API)
     @AccessApiOverload
-    public inline infix fun <reified R> ColumnAccessor<R>.from(noinline expression: RowExpression<T, R>): Boolean =
+    public inline infix fun <reified R> ColumnAccessor<R>.from(noinline expression: AddExpression<T, R>): Boolean =
         name().from(expression)
 
     @Deprecated(DEPRECATED_ACCESS_API)
     @AccessApiOverload
-    public inline infix fun <reified R> KProperty<R>.from(noinline expression: RowExpression<T, R>): Boolean =
+    public inline infix fun <reified R> KProperty<R>.from(noinline expression: AddExpression<T, R>): Boolean =
         add(name, Infer.Nulls, expression)
 
     public infix fun String.from(column: AnyColumnReference): Boolean = add(column.rename(this))
@@ -417,15 +416,15 @@ public fun <T> DataFrame<T>.add(body: AddDsl<T>.() -> Unit): DataFrame<T> {
 public inline fun <reified R, T, G> GroupBy<T, G>.add(
     name: String,
     infer: Infer = Infer.Nulls,
-    noinline expression: RowExpression<G, R>,
+    noinline expression: AddExpression<G, R>,
 ): GroupBy<T, G> = updateGroups { add(name, infer, expression) }
 
 @Deprecated(DEPRECATED_ACCESS_API)
 @AccessApiOverload
 public inline fun <reified R, T, G> GroupBy<T, G>.add(
     column: ColumnAccessor<G>,
     infer: Infer = Infer.Nulls,
-    noinline expression: RowExpression<G, R>,
+    noinline expression: AddExpression<G, R>,
 ): GroupBy<T, G> = add(column.name(), infer, expression)
 
 public class AddGroup<T>(internal val body: AddDsl<T>.() -> Unit)

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt

@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.api
 import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.junit.Test
 import kotlin.reflect.typeOf
 
@@ -33,4 +34,28 @@ class AddTests {
         val df = dataFrameOf("a")(1).addValue(2)
         df["value"].type() shouldBe typeOf<List<Any?>>()
     }
+
+    @Test
+    fun `fibonacci simple add`() {
+        val df = DataFrame.empty(10).add("fibonacci") {
+            if (index() < 2) 1 else prev()!!.newValue<Int>() + prev()!!.prev()!!.newValue<Int>()
+        }
+
+        df["fibonacci"].toList() shouldBe listOf(1, 1, 2, 3, 5, 8, 13, 21, 34, 55)
+    }
+
+    @Test
+    fun `fibonacci add dsl`() {
+        val df = DataFrame.empty(10).add {
+            "fibonacci1" from {
+                if (index() < 2) 1 else prev()!!.newValue<Int>() + prev()!!.prev()!!.newValue<Int>()
+            }
+            expr {
+                if (index() < 2) 1 else prev()!!.newValue<Int>() + prev()!!.prev()!!.newValue<Int>()
+            } into "fibonacci2"
+        }
+
+        df["fibonacci1"].toList() shouldBe listOf(1, 1, 2, 3, 5, 8, 13, 21, 34, 55)
+        df["fibonacci2"].toList() shouldBe listOf(1, 1, 2, 3, 5, 8, 13, 21, 34, 55)
+    }
 }