Skip to content

Commit 2dddcc0

Browse files
JolanrensenJolanrensen
committed
children
1 parent e3a115b commit 2dddcc0

File tree

6 files changed

+176
-46
lines changed

6 files changed

+176
-46
lines changed

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

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4496,23 +4496,26 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
44964496

44974497
// region groups
44984498

4499-
public fun SingleColumn<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
4500-
children { it.isColumnGroup() && filter(it.asColumnGroup()) } as ColumnSet<AnyRow>
4499+
public fun SingleColumn<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSetWithRecursively<AnyRow> =
4500+
children { it.isColumnGroup() && filter(it.asColumnGroup()) } as ColumnSetWithRecursively<AnyRow>
45014501

4502-
public fun String.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
4502+
public fun String.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSetWithRecursively<AnyRow> =
45034503
toColumnAccessor().groups(filter)
45044504

4505-
public fun KProperty<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
4505+
public fun KProperty<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSetWithRecursively<AnyRow> =
45064506
toColumnAccessor().groups(filter)
45074507

45084508
// endregion
45094509

45104510
// region children
45114511

4512-
public fun ColumnSet<*>.children(predicate: (ColumnWithPath<Any?>) -> Boolean = { true }): ColumnSet<Any?> =
4512+
// takes children of all columns in the column set
4513+
public fun ColumnSet<*>.children(predicate: ColumnFilter<Any?> = { true }): ColumnSetWithRecursively<Any?> =
45134514
transform { it.flatMap { it.children().filter { predicate(it) } } }
45144515

4515-
public fun ColumnGroupReference.children(): ColumnSet<Any?> = transformSingle { it.children() }
4516+
// same as cols
4517+
public fun SingleColumn<*>.children(predicate: ColumnFilter<Any?> = { true }): ColumnSetWithRecursively<Any?> =
4518+
(this as ColumnSet<*>).children(predicate)
45164519

45174520
// endregion
45184521

@@ -4796,9 +4799,7 @@ internal fun ColumnSet<*>.colsInternal(range: IntRange): ColumnSetWithRecursivel
47964799
it.single().children()
47974800
} else {
47984801
it
4799-
}.let { cols ->
4800-
cols.subList(range.first, range.last + 1)
4801-
}
4802+
}.subList(range.first, range.last + 1)
48024803
}
48034804

48044805
internal fun ColumnSet<*>.allInternal(): ColumnSetWithRecursively<*> =
@@ -4810,6 +4811,15 @@ internal fun ColumnSet<*>.allInternal(): ColumnSetWithRecursively<*> =
48104811
}
48114812
}
48124813

4814+
//internal fun ColumnSet<*>.childrenInternal(predicate: (ColumnWithPath<Any?>) -> Boolean): ColumnSetWithRecursively<*> =
4815+
// transform {
4816+
// if (isSingleColumnGroup(it)) {
4817+
// it.single().children()
4818+
// } else {
4819+
// it
4820+
// }
4821+
// }
4822+
48134823

48144824
@Deprecated("Replaced with recursively()")
48154825
internal fun ColumnSet<*>.dfsInternal(predicate: (ColumnWithPath<*>) -> Boolean) =

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/Utils.kt

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@ import org.jetbrains.kotlinx.dataframe.ColumnsContainer
66
import org.jetbrains.kotlinx.dataframe.DataColumn
77
import org.jetbrains.kotlinx.dataframe.DataFrame
88
import org.jetbrains.kotlinx.dataframe.DataRow
9-
import org.jetbrains.kotlinx.dataframe.api.*
9+
import org.jetbrains.kotlinx.dataframe.api.cast
10+
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
1011
import org.jetbrains.kotlinx.dataframe.api.name
12+
import org.jetbrains.kotlinx.dataframe.api.pathOf
1113
import org.jetbrains.kotlinx.dataframe.columns.*
1214
import org.jetbrains.kotlinx.dataframe.columns.values
1315
import org.jetbrains.kotlinx.dataframe.impl.DataFrameImpl
@@ -85,8 +87,35 @@ internal fun <A, B> SingleColumn<A>.transformSingle(
8587
.resolveSingle(context)
8688
?.let(converter)
8789
?: emptyList()
90+
91+
// override fun resolveRecursively(
92+
// context: ColumnResolutionContext,
93+
// includeGroups: Boolean,
94+
// includeTopLevel: Boolean,
95+
// ): List<ColumnWithPath<B>> =
96+
// this@transformSingle
97+
// .flattenRecursively(includeGroups, includeTopLevel)
98+
// .resolve(context) as List<ColumnWithPath<B>>
8899
}
89100

101+
internal fun <A, B> SingleColumn<A>.transformRemainingSingle(
102+
converter: (ColumnWithPath<A>?) -> ColumnWithPath<B>?,
103+
): SingleColumnWithRecursively<B> = object : SingleColumnWithRecursively<B> {
104+
override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<B>? =
105+
this@transformRemainingSingle
106+
.resolveSingle(context)
107+
?.let(converter)
108+
109+
override fun resolveSingleRecursively(
110+
context: ColumnResolutionContext,
111+
includeGroups: Boolean,
112+
includeTopLevel: Boolean,
113+
): ColumnWithPath<B>? =
114+
this@transformRemainingSingle
115+
.flattenRecursively(includeGroups, includeTopLevel)
116+
.resolveSingle(context)
117+
.let { converter(it as ColumnWithPath<A>?) }
118+
}
90119

91120
internal fun ColumnSet<*>.flattenRecursively(
92121
includeGroups: Boolean = true,
@@ -108,6 +137,27 @@ internal fun ColumnSet<*>.flattenRecursively(
108137
}.filter { includeGroups || !it.isColumnGroup() }
109138
}
110139

140+
internal fun SingleColumn<*>.flattenRecursively(
141+
includeGroups: Boolean = true,
142+
includeTopLevel: Boolean = true,
143+
): SingleColumn<*> = transformRemainingSingle {
144+
val cols = it?.let {
145+
if (isSingleColumnGroup(listOf(it))) {
146+
it.children()
147+
} else {
148+
listOf(it)
149+
}
150+
} ?: emptyList()
151+
152+
if (includeTopLevel) {
153+
cols.flattenRecursively()
154+
} else {
155+
cols
156+
.filter { it.isColumnGroup() }
157+
.flatMap { it.children().flattenRecursively() }
158+
}.singleOrNull { includeGroups || !it.isColumnGroup() }
159+
}
160+
111161
internal fun <A, B> ColumnSet<A>.transform(
112162
converter: (List<ColumnWithPath<A>>) -> List<ColumnWithPath<B>>,
113163
): ColumnSetWithRecursively<B> = object : ColumnSetWithRecursively<B> {

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/recursively.kt

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,7 @@ package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.matchers.shouldBe
44
import io.kotest.matchers.shouldNotBe
5-
import org.jetbrains.kotlinx.dataframe.alsoDebug
65
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
7-
import org.jetbrains.kotlinx.dataframe.impl.columns.singleImpl
8-
import org.jetbrains.kotlinx.dataframe.impl.columns.transform
96
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
107
import org.jetbrains.kotlinx.dataframe.samples.api.city
118
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
@@ -44,24 +41,32 @@ class Recursively : TestBase() {
4441
listOf(
4542
dfGroup.select { name.firstName.firstName },
4643

47-
dfGroup.select { first { it.data.any { it == "Alice" } }.recursively() },
48-
dfGroup.select { last { it.data.any { it == "Alice" } }.recursively() },
49-
dfGroup.select { single { it.data.any { it == "Alice" } }.recursively() },
44+
dfGroup.select { first { col -> col.any { it == "Alice" } }.recursively() },
45+
dfGroup.select { last { col -> col.any { it == "Alice" } }.recursively() },
46+
dfGroup.select { single { col -> col.any { it == "Alice" } }.recursively() },
5047
).shouldAllBeEqual()
5148

5249
listOf(
5350
dfGroup.select { city },
5451

55-
dfGroup.select { first { it.data.any { it == "London" } }.recursively() },
56-
dfGroup.select { last { it.data.any { it == "London" } }.recursively() },
57-
dfGroup.select { single { it.data.any { it == "London" } }.recursively() },
52+
dfGroup.select { first { col -> col.any { it == "London" } }.recursively() },
53+
dfGroup.select { last { col -> col.any { it == "London" } }.recursively() },
54+
dfGroup.select { single { col -> col.any { it == "London" } }.recursively() },
5855
).shouldAllBeEqual()
5956
}
6057

61-
// @Test
62-
// fun `get at`() {
63-
// dfGroup.getColumnsWithPaths { it[0].recursively() }.print()
64-
// }
58+
@Test
59+
fun `get at`() {
60+
// dfGroup.getColumnsWithPaths { it[0..3].recursively() }.print()
61+
dfGroup.getColumnsWithPaths { all().recursively()[0..3] }.print()
62+
}
63+
64+
@Test
65+
fun `children`() {
66+
dfGroup.getColumnsWithPaths { children().recursively() }.print()
67+
dfGroup.getColumnsWithPaths { name.children() }.print()
68+
}
69+
6570

6671
// @Test
6772
// fun `combination`() {

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

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1860,23 +1860,26 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
18601860

18611861
// region groups
18621862

1863-
public fun SingleColumn<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
1864-
children { it.isColumnGroup() && filter(it.asColumnGroup()) } as ColumnSet<AnyRow>
1863+
public fun SingleColumn<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSetWithRecursively<AnyRow> =
1864+
children { it.isColumnGroup() && filter(it.asColumnGroup()) } as ColumnSetWithRecursively<AnyRow>
18651865

1866-
public fun String.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
1866+
public fun String.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSetWithRecursively<AnyRow> =
18671867
toColumnAccessor().groups(filter)
18681868

1869-
public fun KProperty<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
1869+
public fun KProperty<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSetWithRecursively<AnyRow> =
18701870
toColumnAccessor().groups(filter)
18711871

18721872
// endregion
18731873

18741874
// region children
18751875

1876-
public fun ColumnSet<*>.children(predicate: (ColumnWithPath<Any?>) -> Boolean = { true }): ColumnSet<Any?> =
1876+
// takes children of all columns in the column set
1877+
public fun ColumnSet<*>.children(predicate: ColumnFilter<Any?> = { true }): ColumnSetWithRecursively<Any?> =
18771878
transform { it.flatMap { it.children().filter { predicate(it) } } }
18781879

1879-
public fun ColumnGroupReference.children(): ColumnSet<Any?> = transformSingle { it.children() }
1880+
// same as cols
1881+
public fun SingleColumn<*>.children(predicate: ColumnFilter<Any?> = { true }): ColumnSetWithRecursively<Any?> =
1882+
(this as ColumnSet<*>).children(predicate)
18801883

18811884
// endregion
18821885

@@ -2108,9 +2111,7 @@ internal fun ColumnSet<*>.colsInternal(range: IntRange): ColumnSetWithRecursivel
21082111
it.single().children()
21092112
} else {
21102113
it
2111-
}.let { cols ->
2112-
cols.subList(range.first, range.last + 1)
2113-
}
2114+
}.subList(range.first, range.last + 1)
21142115
}
21152116

21162117
internal fun ColumnSet<*>.allInternal(): ColumnSetWithRecursively<*> =
@@ -2122,6 +2123,15 @@ internal fun ColumnSet<*>.allInternal(): ColumnSetWithRecursively<*> =
21222123
}
21232124
}
21242125

2126+
//internal fun ColumnSet<*>.childrenInternal(predicate: (ColumnWithPath<Any?>) -> Boolean): ColumnSetWithRecursively<*> =
2127+
// transform {
2128+
// if (isSingleColumnGroup(it)) {
2129+
// it.single().children()
2130+
// } else {
2131+
// it
2132+
// }
2133+
// }
2134+
21252135

21262136
@Deprecated("Replaced with recursively()")
21272137
internal fun ColumnSet<*>.dfsInternal(predicate: (ColumnWithPath<*>) -> Boolean) =

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/Utils.kt

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@ import org.jetbrains.kotlinx.dataframe.ColumnsContainer
66
import org.jetbrains.kotlinx.dataframe.DataColumn
77
import org.jetbrains.kotlinx.dataframe.DataFrame
88
import org.jetbrains.kotlinx.dataframe.DataRow
9-
import org.jetbrains.kotlinx.dataframe.api.*
9+
import org.jetbrains.kotlinx.dataframe.api.cast
10+
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
1011
import org.jetbrains.kotlinx.dataframe.api.name
12+
import org.jetbrains.kotlinx.dataframe.api.pathOf
1113
import org.jetbrains.kotlinx.dataframe.columns.*
1214
import org.jetbrains.kotlinx.dataframe.columns.values
1315
import org.jetbrains.kotlinx.dataframe.impl.DataFrameImpl
@@ -85,8 +87,35 @@ internal fun <A, B> SingleColumn<A>.transformSingle(
8587
.resolveSingle(context)
8688
?.let(converter)
8789
?: emptyList()
90+
91+
// override fun resolveRecursively(
92+
// context: ColumnResolutionContext,
93+
// includeGroups: Boolean,
94+
// includeTopLevel: Boolean,
95+
// ): List<ColumnWithPath<B>> =
96+
// this@transformSingle
97+
// .flattenRecursively(includeGroups, includeTopLevel)
98+
// .resolve(context) as List<ColumnWithPath<B>>
8899
}
89100

101+
internal fun <A, B> SingleColumn<A>.transformRemainingSingle(
102+
converter: (ColumnWithPath<A>?) -> ColumnWithPath<B>?,
103+
): SingleColumnWithRecursively<B> = object : SingleColumnWithRecursively<B> {
104+
override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<B>? =
105+
this@transformRemainingSingle
106+
.resolveSingle(context)
107+
?.let(converter)
108+
109+
override fun resolveSingleRecursively(
110+
context: ColumnResolutionContext,
111+
includeGroups: Boolean,
112+
includeTopLevel: Boolean,
113+
): ColumnWithPath<B>? =
114+
this@transformRemainingSingle
115+
.flattenRecursively(includeGroups, includeTopLevel)
116+
.resolveSingle(context)
117+
.let { converter(it as ColumnWithPath<A>?) }
118+
}
90119

91120
internal fun ColumnSet<*>.flattenRecursively(
92121
includeGroups: Boolean = true,
@@ -108,6 +137,27 @@ internal fun ColumnSet<*>.flattenRecursively(
108137
}.filter { includeGroups || !it.isColumnGroup() }
109138
}
110139

140+
internal fun SingleColumn<*>.flattenRecursively(
141+
includeGroups: Boolean = true,
142+
includeTopLevel: Boolean = true,
143+
): SingleColumn<*> = transformRemainingSingle {
144+
val cols = it?.let {
145+
if (isSingleColumnGroup(listOf(it))) {
146+
it.children()
147+
} else {
148+
listOf(it)
149+
}
150+
} ?: emptyList()
151+
152+
if (includeTopLevel) {
153+
cols.flattenRecursively()
154+
} else {
155+
cols
156+
.filter { it.isColumnGroup() }
157+
.flatMap { it.children().flattenRecursively() }
158+
}.singleOrNull { includeGroups || !it.isColumnGroup() }
159+
}
160+
111161
internal fun <A, B> ColumnSet<A>.transform(
112162
converter: (List<ColumnWithPath<A>>) -> List<ColumnWithPath<B>>,
113163
): ColumnSetWithRecursively<B> = object : ColumnSetWithRecursively<B> {

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

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,7 @@ package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.matchers.shouldBe
44
import io.kotest.matchers.shouldNotBe
5-
import org.jetbrains.kotlinx.dataframe.alsoDebug
65
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
7-
import org.jetbrains.kotlinx.dataframe.impl.columns.singleImpl
8-
import org.jetbrains.kotlinx.dataframe.impl.columns.transform
96
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
107
import org.jetbrains.kotlinx.dataframe.samples.api.city
118
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
@@ -44,24 +41,32 @@ class Recursively : TestBase() {
4441
listOf(
4542
dfGroup.select { name.firstName.firstName },
4643

47-
dfGroup.select { first { it.data.any { it == "Alice" } }.recursively() },
48-
dfGroup.select { last { it.data.any { it == "Alice" } }.recursively() },
49-
dfGroup.select { single { it.data.any { it == "Alice" } }.recursively() },
44+
dfGroup.select { first { col -> col.any { it == "Alice" } }.recursively() },
45+
dfGroup.select { last { col -> col.any { it == "Alice" } }.recursively() },
46+
dfGroup.select { single { col -> col.any { it == "Alice" } }.recursively() },
5047
).shouldAllBeEqual()
5148

5249
listOf(
5350
dfGroup.select { city },
5451

55-
dfGroup.select { first { it.data.any { it == "London" } }.recursively() },
56-
dfGroup.select { last { it.data.any { it == "London" } }.recursively() },
57-
dfGroup.select { single { it.data.any { it == "London" } }.recursively() },
52+
dfGroup.select { first { col -> col.any { it == "London" } }.recursively() },
53+
dfGroup.select { last { col -> col.any { it == "London" } }.recursively() },
54+
dfGroup.select { single { col -> col.any { it == "London" } }.recursively() },
5855
).shouldAllBeEqual()
5956
}
6057

61-
// @Test
62-
// fun `get at`() {
63-
// dfGroup.getColumnsWithPaths { it[0].recursively() }.print()
64-
// }
58+
@Test
59+
fun `get at`() {
60+
// dfGroup.getColumnsWithPaths { it[0..3].recursively() }.print()
61+
dfGroup.getColumnsWithPaths { all().recursively()[0..3] }.print()
62+
}
63+
64+
@Test
65+
fun `children`() {
66+
dfGroup.getColumnsWithPaths { children().recursively() }.print()
67+
dfGroup.getColumnsWithPaths { name.children() }.print()
68+
}
69+
6570

6671
// @Test
6772
// fun `combination`() {

0 commit comments

Comments
 (0)