Skip to content

Commit 6b04891

Browse files
Automated commit of generated code
1 parent 390cd30 commit 6b04891

File tree

3 files changed

+36
-7
lines changed

3 files changed

+36
-7
lines changed

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/typeConversions.kt

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -394,12 +394,16 @@ public fun <T> DataFrame<T>.asColumnGroup(column: ColumnGroupAccessor<T>): Colum
394394

395395
// region as GroupedDataFrame
396396

397-
public fun <T> DataFrame<T>.asGroupBy(groupedColumnName: String): GroupBy<T, T> =
398-
GroupByImpl(this, getFrameColumn(groupedColumnName).castFrameColumn()) { none() }
397+
public fun <T> DataFrame<T>.asGroupBy(groupedColumnName: String): GroupBy<T, T> {
398+
val groups = getFrameColumn(groupedColumnName)
399+
return asGroupBy { groups.cast() }
400+
}
399401

400402
@AccessApiOverload
401-
public fun <T, G> DataFrame<T>.asGroupBy(groupedColumn: ColumnReference<DataFrame<G>>): GroupBy<T, G> =
402-
GroupByImpl(this, getFrameColumn(groupedColumn.name()).castFrameColumn()) { none() }
403+
public fun <T, G> DataFrame<T>.asGroupBy(groupedColumn: ColumnReference<DataFrame<G>>): GroupBy<T, G> {
404+
val groups = getFrameColumn(groupedColumn.name()).castFrameColumn<G>()
405+
return asGroupBy { groups }
406+
}
403407

404408
public fun <T> DataFrame<T>.asGroupBy(): GroupBy<T, T> {
405409
val groupCol = columns().single { it.isFrameColumn() }.asAnyFrameColumn().castFrameColumn<T>()
@@ -408,7 +412,7 @@ public fun <T> DataFrame<T>.asGroupBy(): GroupBy<T, T> {
408412

409413
public fun <T, G> DataFrame<T>.asGroupBy(selector: ColumnSelector<T, DataFrame<G>>): GroupBy<T, G> {
410414
val column = getColumn(selector).asFrameColumn()
411-
return GroupByImpl(this, column) { none() }
415+
return GroupByImpl(this.move { column }.toEnd(), column) { none() }
412416
}
413417

414418
// endregion

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/PlaylistJsonTest.kt

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.jetbrains.kotlinx.dataframe.io
22

3+
import io.kotest.assertions.withClue
34
import io.kotest.matchers.shouldBe
45
import org.jetbrains.kotlinx.dataframe.DataFrame
56
import org.jetbrains.kotlinx.dataframe.DataRow
@@ -220,8 +221,10 @@ class PlaylistJsonTest {
220221
minBy { snippet.publishedAt }.snippet into "earliest"
221222
}
222223

223-
res.columnsCount() shouldBe typed.columnsCount() + 1
224-
res.getColumnIndex("earliest") shouldBe typed.getColumnIndex("items") + 1
224+
withClue(res.columnNames() to typed.columnNames()) {
225+
res.columnsCount() shouldBe typed.columnsCount() + 1
226+
res.getColumnIndex("earliest") shouldBe res.columns().indices.last
227+
}
225228

226229
val expected = typed.items.map { it.snippet.minBy { publishedAt } }.toList()
227230
res["earliest"].toList() shouldBe expected

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,28 @@ class DataFrameTreeTests : BaseTest() {
138138
res shouldBe typed
139139
}
140140

141+
@Test
142+
fun asGroupByOverloads() {
143+
val rowsColumn by columnOf(typed[0..3], typed[4..5], typed[6..6])
144+
val df = dataFrameOf(rowsColumn)
145+
val res = df.asGroupBy { rowsColumn }.max()
146+
df.asGroupBy("rowsColumn").max() shouldBe res
147+
df.asGroupBy(rowsColumn).max() shouldBe res
148+
}
149+
150+
@Test
151+
fun moveGroupedColumn() {
152+
val df = dataFrameOf(
153+
"group" to listOf(typed[0..3], typed[4..5], typed[6..6]),
154+
"col" to listOf(1, 2, 3),
155+
)
156+
157+
// We need to match the order of columns in the runtime and in compiler plugin
158+
// GroupBy with the same schema should give the same result after aggregation, no matter how it was created
159+
// We cannot track position of group in original df, so we align `asGroupBy` with `groupBy` and move `group` column to end
160+
df.asGroupBy("group").toDataFrame().columnNames() shouldBe listOf("col", "group")
161+
}
162+
141163
@Test
142164
fun createFrameColumn2() {
143165
val id by column(typed.indices())

0 commit comments

Comments
 (0)