Merge pull request #196 from Kotlin/add-dsl

support applying column operations on row expression result in AddDsl

Author: koperagen

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

@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.Column
 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
@@ -161,6 +162,10 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
         noinline expression: RowExpression<T, R>
     ): Boolean = add(df.mapToColumn(name, infer, expression))
 
+    public inline fun <reified R> expr(noinline expression: RowExpression<T, R>): DataColumn<R> {
+        return df.mapToColumn("", Infer.Nulls, expression)
+    }
+
     public inline infix fun <reified R> String.from(noinline expression: RowExpression<T, R>): Boolean = add(this, Infer.Nulls, expression)
 
     // TODO: use path instead of name

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/constructors.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlinx.dataframe.api.SortDsl
 import org.jetbrains.kotlinx.dataframe.api.asDataColumn
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.concat
+import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.indices
 import org.jetbrains.kotlinx.dataframe.api.toColumnAccessor
 import org.jetbrains.kotlinx.dataframe.api.toColumnOf
@@ -57,7 +58,8 @@ internal fun <T, R> ColumnsContainer<T>.newColumn(
     infer: Infer = Infer.Nulls,
     expression: AddExpression<T, R>,
 ): DataColumn<R> {
-    val (nullable, values) = computeValues(this as DataFrame<T>, expression)
+    val df = this as? DataFrame<T> ?: dataFrameOf(columns()).cast()
+    val (nullable, values) = computeValues(df, expression)
     return when (infer) {
         Infer.Nulls -> DataColumn.create(
             name = name,

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

@@ -10,4 +10,12 @@ class MapTests {
         val frames by columnOf(dataFrameOf("a")(1), emptyDataFrame())
         frames.map { it.firstOrNull() }.size() shouldBe frames.size()
     }
+
+    @Test
+    fun `map ColumnsContainer`() {
+        val df = dataFrameOf("a")(1, 2).add {
+            expr { "a"<Int>() + 1 }.cumSum() into "b"
+        }
+        df["b"][1] shouldBe 5
+    }
 }