Skip to content

[Compiler plugin] Support move to and insert at #1130

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 22, 2025
Merged

[Compiler plugin] Support move to and insert at #1130

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ public fun <T> InsertClause<T>.after(columnPath: ColumnPath): DataFrame<T> {

// region at

@Refine
@Interpretable("InsertAt")
public fun <T> InsertClause<T>.at(position: Int): DataFrame<T> = df.add(column).move(column).to(position)

// endregion
Expand Down
2 changes: 2 additions & 0 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,6 +520,8 @@ public fun <T, C> MoveClause<T, C>.under(
* @param [columnIndex] The index specifying the position in the [DataFrame] columns
* * where the selected columns will be moved.
*/
@Refine
@Interpretable("MoveTo")
public fun <T, C> MoveClause<T, C>.to(columnIndex: Int): DataFrame<T> = moveTo(columnIndex)

/**
Expand Down
12 changes: 12 additions & 0 deletions core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2340,6 +2340,18 @@ class DataFrameTests : BaseTest() {
.move { city }.to(1)
}

@Test
fun `move beyond range of existing column indices`() {
val res = typed.move { city }.to(11)
res.columnNames() shouldBe listOf("name", "age", "weight", "city")
}

@Test
fun `move multiple columns beyond range of existing column indices`() {
val res = typed.move { city and name }.to(11)
res.columnNames() shouldBe listOf("age", "weight", "city", "name")
}

@Test
fun splitIntoThisAndNewColumn() {
val split = typed.split { name }.by { listOf(it.dropLast(1), it.last()) }.into("name", "lastChar")
Expand Down
12 changes: 12 additions & 0 deletions plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/insert.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api

import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.api.after
import org.jetbrains.kotlinx.dataframe.api.at
import org.jetbrains.kotlinx.dataframe.api.insert
import org.jetbrains.kotlinx.dataframe.api.pathOf
import org.jetbrains.kotlinx.dataframe.api.under
Expand Down Expand Up @@ -86,3 +87,14 @@ internal class InsertAfter0 : AbstractInterpreter<PluginDataFrameSchema>() {
.toPluginDataFrameSchema()
}
}

internal class InsertAt : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.receiver: InsertClauseApproximation by arg()
val Arguments.position: Int by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame()
.insert(receiver.column.asDataColumn()).at(position)
Copy link
Collaborator

@zaleslaw zaleslaw Apr 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we call somehow to calculate the schema and other functions? Just a new idea for me

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For operations that only operate on column structure or names: move, rename, flatten, remove, group / ungroup, it's possible. insert here is a bit unique because this way you can only insert this special column converted from SimpleCol.

.toPluginDataFrameSchema()
}
}
9 changes: 9 additions & 0 deletions plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/move.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.api.after
import org.jetbrains.kotlinx.dataframe.api.into
import org.jetbrains.kotlinx.dataframe.api.move
import org.jetbrains.kotlinx.dataframe.api.moveToStart
import org.jetbrains.kotlinx.dataframe.api.to
import org.jetbrains.kotlinx.dataframe.api.toStart
import org.jetbrains.kotlinx.dataframe.api.toEnd
import org.jetbrains.kotlinx.dataframe.api.toTop
Expand Down Expand Up @@ -104,5 +105,13 @@ class MoveAfter0 : AbstractSchemaModificationInterpreter() {
}
}

class MoveTo : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()
val Arguments.columnIndex: Int by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame().move { receiver.columns }.to(columnIndex).toPluginDataFrameSchema()
}
}

class MoveClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)
4 changes: 4 additions & 0 deletions plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStdOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum1
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySumOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAt
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Merge0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeId
Expand All @@ -157,6 +158,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeInto0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveAfter0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveInto0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveTo
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart1
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToEnd0
Expand Down Expand Up @@ -307,6 +309,7 @@ internal inline fun <reified T> String.load(): T {
"Under1" -> Under1()
"Under4" -> Under4()
"InsertAfter0" -> InsertAfter0()
"InsertAt" -> InsertAt()
"Join0" -> Join0()
"Match0" -> Match0()
"Rename" -> Rename()
Expand Down Expand Up @@ -431,6 +434,7 @@ internal inline fun <reified T> String.load(): T {
"MoveToStart1" -> MoveToStart1()
"MoveToEnd0" -> MoveToEnd0()
"MoveAfter0" -> MoveAfter0()
"MoveTo" -> MoveTo()
"GroupByAdd" -> GroupByAdd()
"Merge0" -> Merge0()
"MergeInto0" -> MergeInto0()
Expand Down
15 changes: 15 additions & 0 deletions plugins/kotlin-dataframe/testData/box/insertAt.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import org.jetbrains.kotlinx.dataframe.*
import org.jetbrains.kotlinx.dataframe.annotations.*
import org.jetbrains.kotlinx.dataframe.api.*
import org.jetbrains.kotlinx.dataframe.io.*

fun box(): String {
val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6)
val dCol: DataColumn<String> = df.insert("d") { (b * c).toString() }.at(2)
// weird way to check order, but ok?..
.select { drop(2) }
.select { take(1) }
.select { d }.d

return "OK"
}
11 changes: 11 additions & 0 deletions plugins/kotlin-dataframe/testData/box/moveTo.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import org.jetbrains.kotlinx.dataframe.*
import org.jetbrains.kotlinx.dataframe.annotations.*
import org.jetbrains.kotlinx.dataframe.api.*
import org.jetbrains.kotlinx.dataframe.io.*

fun box(): String {
val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6)
val res = df.move { c and a }.to(3).select { drop(1) }
res.compareSchemas(strict = true)
return "OK"
}
12 changes: 12 additions & 0 deletions ...e/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,12 @@ public void testInsertAfterNested() {
runTest("testData/box/insertAfterNested.kt");
}

@Test
@TestMetadata("insertAt.kt")
public void testInsertAt() {
runTest("testData/box/insertAt.kt");
}

@Test
@TestMetadata("inventNamesForLocalClasses.kt")
public void testInventNamesForLocalClasses() {
Expand Down Expand Up @@ -424,6 +430,12 @@ public void testMoveInto() {
runTest("testData/box/moveInto.kt");
}

@Test
@TestMetadata("moveTo.kt")
public void testMoveTo() {
runTest("testData/box/moveTo.kt");
}

@Test
@TestMetadata("moveToEnd.kt")
public void testMoveToEnd() {
Expand Down
Loading