Skip to content

Commit ea12560

Browse files
AndreiKingsleyAndreiKingsley
committed
concatWithKeys impl
1 parent 5e9df02 commit ea12560

File tree

7 files changed

+172
-64
lines changed

7 files changed

+172
-64
lines changed

‎core/api/core.api

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1340,6 +1340,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/ConcatKt {
13401340
public static final fun concat (Lorg/jetbrains/kotlinx/dataframe/api/ReducedGroupBy;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
13411341
public static final fun concatRows (Ljava/lang/Iterable;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
13421342
public static final fun concatT (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Ljava/lang/Iterable;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
1343+
public static final fun concatWithKeys (Lorg/jetbrains/kotlinx/dataframe/api/GroupBy;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
13431344
}
13441345

13451346
public final class org/jetbrains/kotlinx/dataframe/api/ConstructorsKt {

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,12 @@ package org.jetbrains.kotlinx.dataframe.api
33
import org.jetbrains.kotlinx.dataframe.DataColumn
44
import org.jetbrains.kotlinx.dataframe.DataFrame
55
import org.jetbrains.kotlinx.dataframe.DataRow
6+
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
7+
import org.jetbrains.kotlinx.dataframe.annotations.Refine
68
import org.jetbrains.kotlinx.dataframe.columns.values
79
import org.jetbrains.kotlinx.dataframe.impl.api.concatImpl
810
import org.jetbrains.kotlinx.dataframe.impl.asList
11+
import org.jetbrains.kotlinx.dataframe.type
912

1013
// region DataColumn
1114

@@ -40,6 +43,64 @@ public fun <T> DataFrame<T>.concat(frames: Iterable<DataFrame<T>>): DataFrame<T>
4043

4144
public fun <T, G> GroupBy<T, G>.concat(): DataFrame<G> = groups.concat()
4245

46+
/**
47+
* Concatenates all groups in this [GroupBy] into a single [DataFrame],
48+
* preserving and including all grouping key columns that are not present in the group's columns.
49+
*
50+
* Doesn't affect key columns that have the same names as columns inside the groups (even if their content differs).
51+
*
52+
* This function is especially useful when grouping by expressions or renamed columns,
53+
* and you want the resulting [DataFrame] to include those keys as part of the output.
54+
*
55+
* ### Example
56+
*
57+
* ```kotlin
58+
* val df = dataFrameOf(
59+
* "value" to listOf(1, 2, 3, 3),
60+
* "type" to listOf("a", "b", "a", "b")
61+
* )
62+
*
63+
* val gb = df.groupBy { expr { "Category: \${type.uppercase()}" } named "category" }
64+
* ```
65+
*
66+
* A regular `concat()` will return a [DataFrame] similar to the original `df`
67+
* (with the same columns and rows but in the different orders):
68+
*
69+
* ```
70+
* gb.concat()
71+
* ```
72+
* | value | type |
73+
* | :---- | :--- |
74+
* | 1 | a |
75+
* | 3 | a |
76+
* | 2 | b |
77+
* | 3 | b |
78+
*
79+
* But `concatWithKeys()` will include the new "category" key column:
80+
*
81+
* ```
82+
* gb.concatWithKeys()
83+
* ```
84+
* | value | type | category |
85+
* | :---- | :--- | :------------ |
86+
* | 1 | a | Category: A |
87+
* | 3 | a | Category: A |
88+
* | 2 | b | Category: B |
89+
* | 3 | b | Category: B |
90+
*
91+
* @return A new [DataFrame] where all groups are combined and additional key columns are included in each row.
92+
*/
93+
@Refine
94+
@Interpretable("ConcatWithKeys")
95+
public fun <T, G> GroupBy<T, G>.concatWithKeys(): DataFrame<G> =
96+
mapToFrames {
97+
val rowsCount = group.rowsCount()
98+
val keyColumns = keys.columns().filter { it.name !in group.columnNames() }.map { keyColumn ->
99+
DataColumn.createByType(keyColumn.name, List(rowsCount) { key[keyColumn] }, keyColumn.type)
100+
}
101+
group.addAll(keyColumns)
102+
}.concat()
103+
43104
// endregion
44105

45106
// region ReducedGroupBy

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,16 @@ class ConcatTests {
1111
val b by columnOf(3.0, null)
1212
a.concat(b) shouldBe columnOf(1, 2, 3.0, null).named("a")
1313
}
14+
15+
@Test
16+
fun `concat with keys`() {
17+
val df = dataFrameOf(
18+
"value" to listOf(1, 2, 3, 3),
19+
"type" to listOf("a", "b", "a", "b"),
20+
)
21+
val gb = df.groupBy { expr { "Category: ${(this["type"] as String).uppercase()}" } named "category" }
22+
val dfWithCategory = gb.concatWithKeys()
23+
24+
dfWithCategory.columnNames() shouldBe listOf("value", "type", "category")
25+
}
1426
}

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
3636
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
3737
import org.jetbrains.kotlinx.dataframe.plugin.interpret
3838
import org.jetbrains.kotlinx.dataframe.plugin.loadInterpreter
39+
import kotlin.collections.plus
3940

4041
class GroupBy(val keys: PluginDataFrameSchema, val groups: PluginDataFrameSchema) {
4142
companion object {
@@ -420,6 +421,12 @@ private fun isIntraComparable(col: SimpleDataColumn, session: FirSession): Boole
420421
return col.type.type.isSubtypeOf(comparable, session)
421422
}
422423

424+
class ConcatWithKeys : AbstractSchemaModificationInterpreter() {
425+
val Arguments.receiver by groupBy()
423426

424-
425-
427+
override fun Arguments.interpret(): PluginDataFrameSchema {
428+
val originalColumns = receiver.groups.columns()
429+
return PluginDataFrameSchema(
430+
originalColumns + receiver.keys.columns().filter { it.name !in originalColumns.map { it.name } })
431+
}
432+
}

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

Lines changed: 64 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -6,49 +6,6 @@
66
package org.jetbrains.kotlinx.dataframe.plugin
77

88
import org.jetbrains.kotlin.fir.FirSession
9-
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
10-
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
11-
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.INTERPRETABLE_FQNAME
12-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Add
13-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddWithDsl
14-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.And0
15-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.And10
16-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert0
17-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert2
18-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert6
19-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameGroupBy
20-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropNulls0
21-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Exclude0
22-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Exclude1
23-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Explode0
24-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Expr0
25-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.From
26-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Group0
27-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateDslInto
28-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByToDataFrame
29-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert0
30-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert1
31-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert2
32-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert3
33-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Into
34-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Into0
35-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Join0
36-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Match0
37-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Preserve0
38-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Preserve1
39-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Properties0
40-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Remove0
41-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Rename
42-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameInto
43-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Select0
44-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.To0
45-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under0
46-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under1
47-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under2
48-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under3
49-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under4
50-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Ungroup0
51-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.With0
529
import org.jetbrains.kotlin.fir.declarations.findArgumentByName
5310
import org.jetbrains.kotlin.fir.expressions.FirClassReferenceExpression
5411
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
@@ -66,10 +23,15 @@ import org.jetbrains.kotlin.name.CallableId
6623
import org.jetbrains.kotlin.name.ClassId
6724
import org.jetbrains.kotlin.name.Name
6825
import org.jetbrains.kotlin.name.StandardClassIds
26+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
27+
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
28+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Add
6929
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslNamedGroup
7030
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
7131
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
32+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddWithDsl
7233
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate
34+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateDslInto
7335
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateRow
7436
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
7537
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All1
@@ -87,6 +49,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllFrom2
8749
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo0
8850
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo1
8951
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo2
52+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.And0
53+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.And10
9054
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ByName
9155
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups0
9256
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups1
@@ -99,7 +63,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
9963
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
10064
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf2
10165
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColumnRange
66+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ConcatWithKeys
67+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert0
68+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert2
69+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Convert6
10270
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0
71+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameGroupBy
10372
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0
10473
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf3
10574
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameXs
@@ -110,6 +79,11 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropLast0
11079
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropLast1
11180
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropLast2
11281
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropNa0
82+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropNulls0
83+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Exclude0
84+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Exclude1
85+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Explode0
86+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Expr0
11387
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0
11488
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.First0
11589
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.First1
@@ -119,6 +93,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault
11993
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
12094
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols1
12195
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols2
96+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.From
97+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Group0
12298
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByAdd
12399
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByCount0
124100
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByInto
@@ -137,29 +113,38 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByMinOf
137113
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceExpression
138114
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceInto
139115
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReducePredicate
140-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByXs
141-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAfter0
142-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last0
143-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last1
144-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last2
145116
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStd0
146117
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStd1
147118
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStdOf
148119
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum0
149120
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum1
150121
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySumOf
122+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByToDataFrame
123+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByXs
124+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert0
125+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert1
126+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert2
127+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert3
128+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAfter0
129+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Into
130+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Into0
131+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Join0
132+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last0
133+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last1
134+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last2
151135
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
136+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Match0
152137
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Merge0
153-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeId
154138
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeBy0
155139
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeBy1
140+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeId
156141
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeInto0
157142
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
158143
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveAfter0
159144
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveInto0
145+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToEnd0
160146
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart0
161147
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart1
162-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToEnd0
163148
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveUnder0
164149
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveUnder1
165150
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameContains0
@@ -175,38 +160,54 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Named0
175160
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairConstructor
176161
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairToConstructor
177162
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PerRowCol
163+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Preserve0
164+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Preserve1
165+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Properties0
166+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Remove0
167+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Rename
168+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameInto
178169
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameMapping
179-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
180-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
181-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
182-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
183-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDslStringInvoke
184-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
185-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
186-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
187-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Update0
188-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
189-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCounts
190170
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCase
191171
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCaseClause
192172
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Reorder
193173
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReorderColumnsByName
174+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Select0
194175
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single0
195176
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single1
196177
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Single2
197-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols2
198178
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Take0
199179
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Take1
200180
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Take2
201181
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TakeLast0
202182
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TakeLast1
203183
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TakeLast2
184+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.To0
185+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
186+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
187+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
188+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
189+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDslStringInvoke
190+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
191+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
192+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
193+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under0
194+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under1
195+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under2
196+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under3
197+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Under4
198+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Ungroup0
199+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Update0
200+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
204201
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols0
205202
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols1
203+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCols2
204+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCounts
205+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.With0
206206
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.WithoutNulls0
207207
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.WithoutNulls1
208208
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.WithoutNulls2
209209
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
210+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.INTERPRETABLE_FQNAME
210211

211212
internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
212213
val interpreter = Stdlib.interpreter(this)
@@ -459,6 +460,7 @@ internal inline fun <reified T> String.load(): T {
459460
"GroupByStdOf" -> GroupByStdOf()
460461
"DataFrameXs" -> DataFrameXs()
461462
"GroupByXs" -> GroupByXs()
463+
"ConcatWithKeys" -> ConcatWithKeys()
462464
else -> error("$this")
463465
} as T
464466
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
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(
8+
"value" to listOf(1, 2, 3, 3),
9+
"type" to listOf("a", "b", "a", "b")
10+
)
11+
val gb = df.groupBy { expr { "Category: ${type.uppercase()}" } named "category" }
12+
val categoryKey = gb.keys.category
13+
14+
val dfWithCategory = gb.concatWithKeys()
15+
16+
val category: DataColumn<String> = dfWithCategory.category
17+
18+
return "OK"
19+
}

0 commit comments

Comments
 (0)