Skip to content

Commit 8497e3e

Browse files
koperagenkoperagen
committed
[Compiler plugin] Support reorder + byName
1 parent 9f5e719 commit 8497e3e

File tree

6 files changed

+242
-0
lines changed

6 files changed

+242
-0
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public data class Reorder<T, C>(
2323
public fun <R> cast(): Reorder<T, R> = this as Reorder<T, R>
2424
}
2525

26+
@Interpretable("Reorder")
2627
public fun <T, C> DataFrame<T>.reorder(selector: ColumnsSelector<T, C>): Reorder<T, C> = Reorder(this, selector, false)
2728

2829
@AccessApiOverload
@@ -37,6 +38,8 @@ public fun <T> DataFrame<T>.reorder(vararg columns: String): Reorder<T, *> = reo
3738
public fun <T, C, V : Comparable<V>> Reorder<T, C>.by(expression: ColumnExpression<C, V>): DataFrame<T> =
3839
reorderImpl(false, expression)
3940

41+
@Refine
42+
@Interpretable("ByName")
4043
public fun <T, C> Reorder<T, C>.byName(desc: Boolean = false): DataFrame<T> =
4144
if (desc) byDesc { it.name } else by { it.name }
4245

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/reorder.kt

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22

3+
import org.jetbrains.kotlinx.dataframe.api.byName
4+
import org.jetbrains.kotlinx.dataframe.api.reorder
35
import org.jetbrains.kotlinx.dataframe.api.reorderColumnsByName
6+
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
47
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
58
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
69
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
@@ -18,3 +21,23 @@ class ReorderColumnsByName : AbstractSchemaModificationInterpreter() {
1821
return receiver.asDataFrame().reorderColumnsByName(atAnyDepth, desc).toPluginDataFrameSchema()
1922
}
2023
}
24+
25+
class Reorder : AbstractInterpreter<ReorderApproximation>() {
26+
val Arguments.receiver by dataFrame()
27+
val Arguments.selector: ColumnsResolver by arg()
28+
29+
override fun Arguments.interpret(): ReorderApproximation {
30+
return ReorderApproximation(receiver, selector)
31+
}
32+
}
33+
34+
class ReorderApproximation(val df: PluginDataFrameSchema, val selector: ColumnsResolver)
35+
36+
class ByName : AbstractSchemaModificationInterpreter() {
37+
val Arguments.receiver: ReorderApproximation by arg()
38+
val Arguments.desc: Boolean by arg(defaultValue = Present(false))
39+
40+
override fun Arguments.interpret(): PluginDataFrameSchema {
41+
return receiver.df.asDataFrame().reorder { receiver.selector }.byName(desc).toPluginDataFrameSchema()
42+
}
43+
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllFrom2
9191
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo0
9292
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo1
9393
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo2
94+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ByName
9495
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups0
9596
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups1
9697
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups2
@@ -176,6 +177,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
176177
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCounts
177178
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCase
178179
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCaseClause
180+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Reorder
179181
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReorderColumnsByName
180182
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single0
181183
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single1
@@ -425,6 +427,8 @@ internal inline fun <reified T> String.load(): T {
425427
"MergeBy0" -> MergeBy0()
426428
"MergeBy1" -> MergeBy1()
427429
"ReorderColumnsByName" -> ReorderColumnsByName()
430+
"Reorder" -> Reorder()
431+
"ByName" -> ByName()
428432
"GroupByCount0" -> GroupByCount0()
429433
"GroupByReducePredicate" -> GroupByReducePredicate()
430434
"GroupByReduceExpression" -> GroupByReduceExpression()

‎plugins/kotlin-dataframe/testData/box/reorderRange.fir.txt

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
FILE: reorderRange.kt
2+
public final fun box(): R|kotlin/String| {
3+
lval df: R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_09>| = R|org/jetbrains/kotlinx/dataframe/api/dataFrameOf|(vararg(String(name), String(age), String(city), String(weight))).R|kotlin/let|<R|org/jetbrains/kotlinx/dataframe/api/DataFrameBuilder|, R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_09>|>(<L> = fun <anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/DataFrameBuilder|): R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_09>| <inline=Inline, kind=EXACTLY_ONCE> {
4+
local abstract class Invoke_09I : R|kotlin/Any| {
5+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val city: R|kotlin/String?|
6+
public get(): R|kotlin/String?|
7+
8+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val age: R|kotlin/Int|
9+
public get(): R|kotlin/Int|
10+
11+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val name: R|kotlin/String|
12+
public get(): R|kotlin/String|
13+
14+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(3)) public abstract val weight: R|kotlin/Int?|
15+
public get(): R|kotlin/Int?|
16+
17+
public constructor(): R|<local>/Invoke_09I|
18+
19+
}
20+
21+
local final class Scope0 : R|kotlin/Any| {
22+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_09I>|.city: R|kotlin/String?|
23+
public get(): R|kotlin/String?|
24+
25+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_09I>|.city: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String?>|
26+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String?>|
27+
28+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_09I>|.age: R|kotlin/Int|
29+
public get(): R|kotlin/Int|
30+
31+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_09I>|.age: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
32+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
33+
34+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_09I>|.name: R|kotlin/String|
35+
public get(): R|kotlin/String|
36+
37+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_09I>|.name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
38+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
39+
40+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_09I>|.weight: R|kotlin/Int?|
41+
public get(): R|kotlin/Int?|
42+
43+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_09I>|.weight: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int?>|
44+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int?>|
45+
46+
public constructor(): R|<local>/Scope0|
47+
48+
}
49+
50+
local abstract class Invoke_09 : R|<local>/Invoke_09I| {
51+
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
52+
public get(): R|<local>/Scope0|
53+
54+
public constructor(): R|<local>/Invoke_09|
55+
56+
}
57+
58+
^ R|<local>/it|.R|org/jetbrains/kotlinx/dataframe/api/DataFrameBuilder.invoke|(vararg(String(Alice), Int(15), String(London), Int(54), String(Bob), Int(45), String(Dubai), Int(87), String(Charlie), Int(20), String(Moscow), Null(null), String(Charlie), Int(40), String(Milan), Null(null), String(Bob), Int(30), String(Tokyo), Int(68), String(Alice), Int(20), Null(null), Int(55), String(Charlie), Int(30), String(Moscow), Int(90)))
59+
}
60+
)
61+
lval df1: R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_83>| = R|<local>/df|.R|org/jetbrains/kotlinx/dataframe/api/reorder|<R|<local>/Invoke_09|, R|kotlin/Any?|>(<L> = reorder@fun R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_09>|.<anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_09>|): R|org/jetbrains/kotlinx/dataframe/columns/ColumnsResolver<kotlin/Any?>| <inline=NoInline> {
62+
^ (this@R|special/anonymous|, (this@R|/box|, this@R|special/anonymous|).R|<local>/Scope0.name|).R|org/jetbrains/kotlinx/dataframe/api/ColumnRangeColumnsSelectionDsl.rangeTo|((this@R|/box|, this@R|special/anonymous|).R|<local>/Scope0.age|)
63+
}
64+
).R|kotlin/let|<R|org/jetbrains/kotlinx/dataframe/api/Reorder<<local>/Invoke_09, kotlin/Any?>|, R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_83>|>(<L> = fun <anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/Reorder<<local>/Invoke_09, kotlin/Any?>|): R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_83>| <inline=Inline, kind=EXACTLY_ONCE> {
65+
local abstract class Invoke_83I : R|kotlin/Any| {
66+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val city: R|kotlin/String?|
67+
public get(): R|kotlin/String?|
68+
69+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val age: R|kotlin/Int|
70+
public get(): R|kotlin/Int|
71+
72+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val name: R|kotlin/String|
73+
public get(): R|kotlin/String|
74+
75+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(3)) public abstract val weight: R|kotlin/Int?|
76+
public get(): R|kotlin/Int?|
77+
78+
public constructor(): R|<local>/Invoke_83I|
79+
80+
}
81+
82+
local final class Scope0 : R|kotlin/Any| {
83+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_83I>|.city: R|kotlin/String?|
84+
public get(): R|kotlin/String?|
85+
86+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_83I>|.city: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String?>|
87+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String?>|
88+
89+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_83I>|.age: R|kotlin/Int|
90+
public get(): R|kotlin/Int|
91+
92+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_83I>|.age: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
93+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
94+
95+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_83I>|.name: R|kotlin/String|
96+
public get(): R|kotlin/String|
97+
98+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_83I>|.name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
99+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
100+
101+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_83I>|.weight: R|kotlin/Int?|
102+
public get(): R|kotlin/Int?|
103+
104+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_83I>|.weight: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int?>|
105+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int?>|
106+
107+
public constructor(): R|<local>/Scope0|
108+
109+
}
110+
111+
local abstract class Invoke_83 : R|<local>/Invoke_83I| {
112+
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
113+
public get(): R|<local>/Scope0|
114+
115+
public constructor(): R|<local>/Invoke_83|
116+
117+
}
118+
119+
^ R|<local>/it|.R|org/jetbrains/kotlinx/dataframe/api/byName|<R|<local>/Invoke_09|, R|kotlin/Any?|>()
120+
}
121+
)
122+
lval df2: R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_47>| = R|<local>/df|.R|org/jetbrains/kotlinx/dataframe/api/reorder|<R|<local>/Invoke_09|, R|kotlin/Any?|>(<L> = reorder@fun R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_09>|.<anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_09>|): R|org/jetbrains/kotlinx/dataframe/columns/ColumnsResolver<kotlin/Any?>| <inline=NoInline> {
123+
^ (this@R|special/anonymous|, (this@R|/box|, this@R|special/anonymous|).R|<local>/Scope0.name|).R|org/jetbrains/kotlinx/dataframe/api/ColumnRangeColumnsSelectionDsl.rangeTo|((this@R|/box|, this@R|special/anonymous|).R|<local>/Scope0.city|)
124+
}
125+
).R|kotlin/let|<R|org/jetbrains/kotlinx/dataframe/api/Reorder<<local>/Invoke_09, kotlin/Any?>|, R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_47>|>(<L> = fun <anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/Reorder<<local>/Invoke_09, kotlin/Any?>|): R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_47>| <inline=Inline, kind=EXACTLY_ONCE> {
126+
local abstract class Invoke_47I : R|kotlin/Any| {
127+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val city: R|kotlin/String?|
128+
public get(): R|kotlin/String?|
129+
130+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val age: R|kotlin/Int|
131+
public get(): R|kotlin/Int|
132+
133+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val name: R|kotlin/String|
134+
public get(): R|kotlin/String|
135+
136+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(3)) public abstract val weight: R|kotlin/Int?|
137+
public get(): R|kotlin/Int?|
138+
139+
public constructor(): R|<local>/Invoke_47I|
140+
141+
}
142+
143+
local final class Scope0 : R|kotlin/Any| {
144+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_47I>|.city: R|kotlin/String?|
145+
public get(): R|kotlin/String?|
146+
147+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_47I>|.city: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String?>|
148+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String?>|
149+
150+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_47I>|.age: R|kotlin/Int|
151+
public get(): R|kotlin/Int|
152+
153+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_47I>|.age: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
154+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
155+
156+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_47I>|.name: R|kotlin/String|
157+
public get(): R|kotlin/String|
158+
159+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_47I>|.name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
160+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
161+
162+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_47I>|.weight: R|kotlin/Int?|
163+
public get(): R|kotlin/Int?|
164+
165+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_47I>|.weight: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int?>|
166+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int?>|
167+
168+
public constructor(): R|<local>/Scope0|
169+
170+
}
171+
172+
local abstract class Invoke_47 : R|<local>/Invoke_47I| {
173+
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
174+
public get(): R|<local>/Scope0|
175+
176+
public constructor(): R|<local>/Invoke_47|
177+
178+
}
179+
180+
^ R|<local>/it|.R|org/jetbrains/kotlinx/dataframe/api/byName|<R|<local>/Invoke_09|, R|kotlin/Any?|>(Boolean(true))
181+
}
182+
)
183+
^box String(OK)
184+
}

‎plugins/kotlin-dataframe/testData/box/reorderRange.kt

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
// FIR_DUMP
2+
3+
import org.jetbrains.kotlinx.dataframe.*
4+
import org.jetbrains.kotlinx.dataframe.annotations.*
5+
import org.jetbrains.kotlinx.dataframe.api.*
6+
import org.jetbrains.kotlinx.dataframe.io.*
7+
8+
fun box(): String {
9+
val df = dataFrameOf("name", "age", "city", "weight")(
10+
"Alice", 15, "London", 54,
11+
"Bob", 45, "Dubai", 87,
12+
"Charlie", 20, "Moscow", null,
13+
"Charlie", 40, "Milan", null,
14+
"Bob", 30, "Tokyo", 68,
15+
"Alice", 20, null, 55,
16+
"Charlie", 30, "Moscow", 90,
17+
)
18+
19+
val df1 = df.reorder { name..age }.byName()
20+
val df2 = df.reorder { name..city }.byName(desc = true)
21+
return "OK"
22+
}

‎plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -514,6 +514,12 @@ public void testReorder() {
514514
runTest("testData/box/reorder.kt");
515515
}
516516

517+
@Test
518+
@TestMetadata("reorderRange.kt")
519+
public void testReorderRange() {
520+
runTest("testData/box/reorderRange.kt");
521+
}
522+
517523
@Test
518524
@TestMetadata("Schema.kt")
519525
public void testSchema() {

0 commit comments

Comments
 (0)