Skip to content

Commit 6b892d1

Browse files
koperagenkoperagen
committed
[Compiler plugin] Support dataFrameOf(header)(values)
1 parent 4caf141 commit 6b892d1

File tree

8 files changed

+94
-21
lines changed

8 files changed

+94
-21
lines changed

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

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
99
import org.jetbrains.kotlinx.dataframe.DataFrame
1010
import org.jetbrains.kotlinx.dataframe.DataRow
1111
import org.jetbrains.kotlinx.dataframe.RowExpression
12+
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
13+
import org.jetbrains.kotlinx.dataframe.annotations.Refine
1214
import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
1315
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1416
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -18,6 +20,7 @@ import org.jetbrains.kotlinx.dataframe.exceptions.UnequalColumnSizesException
1820
import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
1921
import org.jetbrains.kotlinx.dataframe.impl.DataFrameImpl
2022
import org.jetbrains.kotlinx.dataframe.impl.UNNAMED_COLUMN_PREFIX
23+
import org.jetbrains.kotlinx.dataframe.impl.api.withValuesImpl
2124
import org.jetbrains.kotlinx.dataframe.impl.asList
2225
import org.jetbrains.kotlinx.dataframe.impl.columnName
2326
import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnAccessorImpl
@@ -268,6 +271,7 @@ public fun dataFrameOf(vararg header: ColumnReference<*>): DataFrameBuilder = Da
268271

269272
public fun dataFrameOf(vararg columns: AnyBaseCol): AnyFrame = dataFrameOf(columns.asIterable())
270273

274+
@Interpretable("DataFrameOf0")
271275
public fun dataFrameOf(vararg header: String): DataFrameBuilder = dataFrameOf(header.toList())
272276

273277
public inline fun <reified C> dataFrameOf(vararg header: String, fill: (String) -> Iterable<C>): AnyFrame =
@@ -302,27 +306,15 @@ public class DataFrameBuilder(private val header: List<String>) {
302306
}.toDataFrame()
303307
}
304308

309+
@Refine
310+
@Interpretable("DataFrameBuilderInvoke0")
305311
public operator fun invoke(vararg values: Any?): AnyFrame = withValues(values.asIterable())
306312

307313
@JvmName("invoke1")
308-
internal fun withValues(values: Iterable<Any?>): AnyFrame {
309-
val list = values.asList()
310-
311-
val ncol = header.size
312-
313-
require(header.isNotEmpty() && list.size.rem(ncol) == 0) {
314-
"Number of values ${list.size} is not divisible by number of columns $ncol"
315-
}
316-
317-
val nrow = list.size / ncol
318-
319-
return (0 until ncol).map { col ->
320-
val colValues = (0 until nrow).map { row ->
321-
list[row * ncol + col]
322-
}
323-
DataColumn.createWithTypeInference(header[col], colValues)
314+
internal fun withValues(values: Iterable<Any?>): AnyFrame =
315+
withValuesImpl(header, values.asList()).map { (name, values) ->
316+
DataColumn.createWithTypeInference(name, values)
324317
}.toDataFrame()
325-
}
326318

327319
public operator fun invoke(args: Sequence<Any?>): AnyFrame = invoke(*args.toList().toTypedArray())
328320

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package org.jetbrains.kotlinx.dataframe.impl.api
2+
3+
internal fun <T> withValuesImpl(header: List<String>, values: List<T>): List<Pair<String, List<T>>> {
4+
val ncol = header.size
5+
6+
require(header.isNotEmpty() && values.size.rem(ncol) == 0) {
7+
"Number of values ${values.size} is not divisible by number of columns $ncol"
8+
}
9+
10+
val nrow = values.size / ncol
11+
12+
return (0 until ncol).map { col ->
13+
val colValues = (0 until nrow).map { row ->
14+
values[row * ncol + col]
15+
}
16+
header[col] to colValues
17+
}
18+
}

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -364,8 +364,13 @@ class FunctionCallTransformer(
364364
resolvedSymbol = parameterSymbol
365365
}
366366
}
367+
if (callDispatchReceiver != null) {
368+
call.replaceDispatchReceiver(itPropertyAccess)
369+
}
367370
call.replaceExplicitReceiver(itPropertyAccess)
368-
call.replaceExtensionReceiver(itPropertyAccess)
371+
if (callExtensionReceiver != null) {
372+
call.replaceExtensionReceiver(itPropertyAccess)
373+
}
369374
result = call
370375
this.target = target
371376
}
@@ -403,9 +408,9 @@ class FunctionCallTransformer(
403408
}
404409
variance = Variance.INVARIANT
405410
}
406-
dispatchReceiver = callDispatchReceiver
411+
dispatchReceiver = null
407412
this.explicitReceiver = callExplicitReceiver
408-
extensionReceiver = callExtensionReceiver
413+
extensionReceiver = callExtensionReceiver ?: callDispatchReceiver
409414
argumentList = buildResolvedArgumentList(original = null, linkedMapOf(argument to parameter.fir))
410415
calleeReference = buildResolvedNamedReference {
411416
source = call.calleeReference.source

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
2+
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
3+
4+
import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
5+
import org.jetbrains.kotlin.fir.types.commonSuperTypeOrNull
6+
import org.jetbrains.kotlin.fir.types.resolvedType
7+
import org.jetbrains.kotlin.fir.types.typeContext
8+
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
9+
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
10+
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
11+
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
12+
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
13+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
14+
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
15+
import org.jetbrains.kotlinx.dataframe.impl.api.withValuesImpl
16+
17+
class DataFrameOf0 : AbstractInterpreter<DataFrameBuilderApproximation>() {
18+
val Arguments.header: List<String> by varargString()
19+
20+
override fun Arguments.interpret(): DataFrameBuilderApproximation {
21+
return DataFrameBuilderApproximation(header)
22+
}
23+
}
24+
25+
class DataFrameBuilderApproximation(val header: List<String>)
26+
27+
class DataFrameBuilderInvoke0 : AbstractSchemaModificationInterpreter() {
28+
val Arguments.receiver: DataFrameBuilderApproximation by arg()
29+
val Arguments.values: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
30+
31+
override fun Arguments.interpret(): PluginDataFrameSchema {
32+
val columns = withValuesImpl(receiver.header, values.arguments).map { (name, values) ->
33+
val type = session.typeContext.commonSuperTypeOrNull(values.map { it.resolvedType }) ?: error("$name $values")
34+
simpleColumnOf(name, type)
35+
}
36+
return PluginDataFrameSchema(columns)
37+
}
38+
}

‎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
@@ -71,6 +71,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
7171
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0
7272
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
7373
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
74+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0
75+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0
7476
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
7577
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
7678
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
@@ -177,6 +179,8 @@ internal inline fun <reified T> String.load(): T {
177179
"toDataFrameDsl" -> ToDataFrameDsl()
178180
"toDataFrame" -> ToDataFrame()
179181
"toDataFrameDefault" -> ToDataFrameDefault()
182+
"DataFrameOf0" -> DataFrameOf0()
183+
"DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
180184
else -> error("$this")
181185
} as T
182186
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
fun box(): String {
7+
val df = dataFrameOf("a")(1, 2, 3)
8+
val i: Int = df.a[0]
9+
return "OK"
10+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ interface Schema {
1010
}
1111

1212
fun main(args: Array<String>) {
13-
val df = dataFrameOf("i")(1, 2, 3).cast<Schema>()
13+
val df = dataFrameOf("i", "fff")(1, "321").cast<Schema>()
1414
println(df.i)
1515

1616
val df1 = df.add("ca") { 423 }

‎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
@@ -64,6 +64,12 @@ public void testDataFrameOf() {
6464
runTest("testData/box/dataFrameOf.kt");
6565
}
6666

67+
@Test
68+
@TestMetadata("dataFrameOf_vararg.kt")
69+
public void testDataFrameOf_vararg() {
70+
runTest("testData/box/dataFrameOf_vararg.kt");
71+
}
72+
6773
@Test
6874
@TestMetadata("dataRowSchemaApi.kt")
6975
public void testDataRowSchemaApi() {

0 commit comments

Comments
 (0)