Skip to content

Commit b2aa8af

Browse files
committed
support applying column operations on row expression result in AddDsl
1 parent d9d5f92 commit b2aa8af

File tree

3 files changed

+16
-1
lines changed
  • core/src

3 files changed

+16
-1
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame
77
import org.jetbrains.kotlinx.dataframe.AnyRow
88
import org.jetbrains.kotlinx.dataframe.Column
99
import org.jetbrains.kotlinx.dataframe.ColumnsContainer
10+
import org.jetbrains.kotlinx.dataframe.DataColumn
1011
import org.jetbrains.kotlinx.dataframe.DataFrame
1112
import org.jetbrains.kotlinx.dataframe.DataRow
1213
import org.jetbrains.kotlinx.dataframe.RowExpression
@@ -161,6 +162,10 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
161162
noinline expression: RowExpression<T, R>
162163
): Boolean = add(df.mapToColumn(name, infer, expression))
163164

165+
public inline fun <reified R> expr(noinline expression: RowExpression<T, R>): DataColumn<R> {
166+
return df.mapToColumn("", Infer.Nulls, expression)
167+
}
168+
164169
public inline infix fun <reified R> String.from(noinline expression: RowExpression<T, R>): Boolean = add(this, Infer.Nulls, expression)
165170

166171
// TODO: use path instead of name

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import org.jetbrains.kotlinx.dataframe.api.SortDsl
2020
import org.jetbrains.kotlinx.dataframe.api.asDataColumn
2121
import org.jetbrains.kotlinx.dataframe.api.cast
2222
import org.jetbrains.kotlinx.dataframe.api.concat
23+
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
2324
import org.jetbrains.kotlinx.dataframe.api.indices
2425
import org.jetbrains.kotlinx.dataframe.api.toColumnAccessor
2526
import org.jetbrains.kotlinx.dataframe.api.toColumnOf
@@ -56,7 +57,8 @@ internal fun <T, R> ColumnsContainer<T>.newColumn(
5657
infer: Infer = Infer.Nulls,
5758
expression: AddExpression<T, R>,
5859
): DataColumn<R> {
59-
val (nullable, values) = computeValues(this as DataFrame<T>, expression)
60+
val df = this as? DataFrame<T> ?: dataFrameOf(columns()).cast()
61+
val (nullable, values) = computeValues(df, expression)
6062
return when (infer) {
6163
Infer.Nulls -> DataColumn.create(name, values, type.withNullability(nullable), Infer.None)
6264
Infer.Type -> DataColumn.createWithTypeInference(name, values, nullable)

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,12 @@ class MapTests {
1010
val frames by columnOf(dataFrameOf("a")(1), emptyDataFrame())
1111
frames.map { it.firstOrNull() }.size() shouldBe frames.size()
1212
}
13+
14+
@Test
15+
fun `map ColumnsContainer`() {
16+
val df = dataFrameOf("a")(1, 2).add {
17+
expr { "a"<Int>() + 1 }.cumSum() into "b"
18+
}
19+
df["b"][1] shouldBe 5
20+
}
1321
}

0 commit comments

Comments
 (0)