Skip to content

Commit da89e93

Browse files
JolanrensenJolanrensen
committed
first {}.recursively now seems to work. Why, I'm not sure. Needs more tests
1 parent 149e96f commit da89e93

File tree

10 files changed

+147
-86
lines changed

10 files changed

+147
-86
lines changed

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

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4224,7 +4224,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
42244224
*
42254225
* @see [cols]
42264226
*/
4227-
public fun ColumnSet<*>.all(): ColumnSet<*> = wrap()
4227+
public fun <C> ColumnSet<C>.all(): ColumnSet<C> = allInternal() as ColumnSet<C>
42284228

42294229
/**
42304230
* ## All
@@ -4250,7 +4250,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
42504250
*
42514251
* @see [cols]
42524252
*/
4253-
public fun SingleColumn<*>.all(): ColumnSet<*> = transformSingle { it.children() }
4253+
public fun SingleColumn<*>.all(): ColumnSet<*> = allInternal()
42544254

42554255
/**
42564256
* ## All
@@ -4330,6 +4330,8 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
43304330
toColumnAccessor().allDfs(includeGroups)
43314331

43324332
/**
4333+
* ## Recursively / Rec
4334+
*
43334335
* Modifies the previous call to run not only on the current column set, but also on all of its children.
43344336
*
43354337
* For example:
@@ -4341,10 +4343,16 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
43414343
*
43424344
* `df.`[select][DataFrame.select]` { `[colsOf][ColumnSet.colsOf]`<`[String][String]`>().`[recursively][ColumnSet.recursivelyImpl]`() }`
43434345
*
4344-
* This will return all columns of type [String] in lower levels (unless [includeTopLevel]` == true`).
4346+
* This will return all columns of type [String] in lower levels (unless [includeTopLevel\]` == true`).
43454347
*
43464348
* TODO
43474349
*/
4350+
private interface CommonRecursivelyDocs {
4351+
4352+
/** Example argument */
4353+
interface Examples
4354+
}
4355+
43484356
public fun <C> ColumnSet<C>.recursively(
43494357
includeGroups: Boolean = true,
43504358
includeTopLevel: Boolean = true,
@@ -4358,7 +4366,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
43584366
public fun SingleColumn<*>.recursively(
43594367
includeGroups: Boolean = true,
43604368
includeTopLevel: Boolean = true,
4361-
): ColumnSet<*> = all().recursivelyImpl(includeTopLevel = includeTopLevel, includeGroups = includeGroups)
4369+
): ColumnSet<*> = recursivelyImpl(includeTopLevel = includeTopLevel, includeGroups = includeGroups)
43624370

43634371
public fun SingleColumn<*>.rec(
43644372
includeGroups: Boolean = true,
@@ -4795,13 +4803,23 @@ internal fun <T, C> ColumnsSelector<T, C>.filter(predicate: (ColumnWithPath<C>)
47954803
*/
47964804
internal fun ColumnSet<*>.colsInternal(predicate: ColumnFilter<*>): ColumnSet<*> =
47974805
transform {
4798-
if (isSingleColumn() && it.singleOrNull()?.isColumnGroup() == true) {
4806+
if (isSingleColumnGroup(it)) {
47994807
it.single().children()
48004808
} else {
48014809
it
48024810
}.filter(predicate)
48034811
}
48044812

4813+
internal fun ColumnSet<*>.allInternal(): ColumnSet<*> =
4814+
transform {
4815+
if (isSingleColumnGroup(it)) {
4816+
it.single().children()
4817+
} else {
4818+
it
4819+
}
4820+
}
4821+
4822+
48054823
@Deprecated("Replaced with recursively()")
48064824
internal fun ColumnSet<*>.dfsInternal(predicate: (ColumnWithPath<*>) -> Boolean) =
48074825
transform { it.filter { it.isColumnGroup() }.flatMap { it.children().flattenRecursively().filter(predicate) } }

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

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,11 @@ public interface ColumnSet<out C> {
3030
): List<ColumnWithPath<C>>
3131
}
3232

33-
public interface ColumnSetTransformer {
34-
35-
// @Deprecated("see if this can be removed")
36-
// public fun transformRemainingSingle(singleColumn: SingleColumn<*>): SingleColumn<*>
37-
38-
// public fun transformSingle(singleColumn: SingleColumn<*>): ColumnSet<*>
33+
public fun interface ColumnSetTransformer {
3934

4035
public fun transform(columnSet: ColumnSet<*>): ColumnSet<*>
36+
37+
public operator fun invoke(columnSet: ColumnSet<*>): ColumnSet<*> = transform(columnSet)
4138
}
4239

4340
public class ColumnResolutionContext internal constructor(

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

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package org.jetbrains.kotlinx.dataframe.columns
22

33
import org.jetbrains.kotlinx.dataframe.DataColumn
4+
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
5+
import kotlin.contracts.ExperimentalContracts
6+
import kotlin.contracts.contract
47

58
/**
69
* Entity that can be [resolved][resolveSingle] into [DataColumn].
@@ -17,11 +20,18 @@ public interface SingleColumn<out C> : ColumnSet<C> {
1720
context: ColumnResolutionContext,
1821
transformer: ColumnSetTransformer,
1922
): List<ColumnWithPath<C>> =
20-
throw UnsupportedOperationException(
21-
"SingleColumn.resolveAfterTransform is not supported because transformations can only be applied on normal ColumnSets. use '.wrap()' to wrap a SingleColumn into a ColumnSet"
22-
)
23+
transformer(this).resolve(context) as List<ColumnWithPath<C>>
2324

2425
public fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>?
2526
}
2627

27-
public fun ColumnSet<*>.isSingleColumn(): Boolean = this is SingleColumn<*>
28+
@OptIn(ExperimentalContracts::class)
29+
public fun ColumnSet<*>.isSingleColumn(): Boolean {
30+
contract {
31+
returns(true) implies (this@isSingleColumn is SingleColumn<*>)
32+
}
33+
return this is SingleColumn<*>
34+
}
35+
36+
public fun ColumnSet<*>.isSingleColumnGroup(cols: List<ColumnWithPath<*>>): Boolean =
37+
isSingleColumn() && cols.singleOrNull()?.isColumnGroup() == true

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

Lines changed: 16 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,10 @@ internal fun <A, B> SingleColumn<A>.transformSingle(
8989
override fun resolveAfterTransform(
9090
context: ColumnResolutionContext,
9191
transformer: ColumnSetTransformer,
92-
): List<ColumnWithPath<B>> = throw NotImplementedError("converter cannot be applied after transform.")
93-
// transformer.transformRemainingSingle(this@transformSingle).cast<A>()
94-
// .resolveSingle(context)
95-
// ?.let(converter)
96-
// ?: emptyList()
97-
}.wrap()
92+
): List<ColumnWithPath<B>> =
93+
this@transformSingle.resolveAfterTransform(context, transformer)
94+
.flatMap(converter)
95+
}
9896

9997
internal fun <A, B> SingleColumn<A>.transformRemainingSingle(
10098
converter: (ColumnWithPath<A>) -> ColumnWithPath<B>,
@@ -125,9 +123,9 @@ internal fun <C> ColumnSet<C>.recursivelyImpl(
125123

126124
val flattenTransformer = object : ColumnSetTransformer {
127125

128-
private fun flattenColumnWithPaths(list: List<ColumnWithPath<*>>, isSingleColumn: Boolean): List<ColumnWithPath<*>> {
126+
private fun flattenColumnWithPaths(list: List<ColumnWithPath<*>>, columnSet: ColumnSet<*>): List<ColumnWithPath<*>> {
129127
val cols =
130-
if (isSingleColumn && list.singleOrNull()?.isColumnGroup() == true) {
128+
if (columnSet.isSingleColumnGroup(list)) {
131129
list.single().children()
132130
} else {
133131
list
@@ -142,27 +140,8 @@ internal fun <C> ColumnSet<C>.recursivelyImpl(
142140
}.filter { includeGroups || !it.isColumnGroup() }
143141
}
144142

145-
// TODO remove/clean
146-
// override fun transformRemainingSingle(singleColumn: SingleColumn<*>): SingleColumn<*> =
147-
// singleColumn.transformRemainingSingle {
148-
// if (!it.isColumnGroup()) return@transformRemainingSingle it
149-
//
150-
// val list = it.asColumnGroup().getColumnsWithPaths { all() }
151-
// val res = flattenColumnWithPaths(list, true)
152-
//
153-
// DataColumn.createColumnGroup(it.name, res.toDataFrame()).addPath(it.path())
154-
// }
155-
156-
// TODO remove/clean
157-
// override fun transformSingle(singleColumn: SingleColumn<*>): ColumnSet<*> = singleColumn.transformSingle {
158-
// if (!it.isColumnGroup()) return@transformSingle listOf(it)
159-
//
160-
// val list = it.asColumnGroup().getColumnsWithPaths { all() }
161-
// flattenColumnWithPaths(list, true)
162-
// }
163-
164143
override fun transform(columnSet: ColumnSet<*>): ColumnSet<*> = columnSet.transform {
165-
flattenColumnWithPaths(it, columnSet.isSingleColumn())
144+
flattenColumnWithPaths(it, columnSet)
166145
}
167146
}
168147

@@ -217,6 +196,15 @@ internal fun <A, B> ColumnSet<A>.transformWithContext(
217196
internal fun <T> ColumnSet<T>.singleImpl() = object : SingleColumn<T> {
218197
override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<T>? =
219198
this@singleImpl.resolve(context).singleOrNull()
199+
200+
override fun resolveAfterTransform(
201+
context: ColumnResolutionContext,
202+
transformer: ColumnSetTransformer,
203+
): List<ColumnWithPath<T>> =
204+
this@singleImpl.resolveAfterTransform(
205+
context = context,
206+
transformer = transformer,
207+
)
220208
}
221209

222210
internal fun <T> ColumnSet<T>.getAt(index: Int) = object : SingleColumn<T> {

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.matchers.shouldBe
4+
import org.jetbrains.kotlinx.dataframe.alsoDebug
45
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
6+
import org.jetbrains.kotlinx.dataframe.impl.columns.recursivelyImpl
7+
import org.jetbrains.kotlinx.dataframe.impl.columns.singleImpl
8+
import org.jetbrains.kotlinx.dataframe.impl.columns.transform
59
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
610
import org.jetbrains.kotlinx.dataframe.samples.api.name
711
import org.junit.Test
@@ -29,6 +33,18 @@ class Recursively : TestBase() {
2933
private val recursivelyString = dfGroup.getColumnsWithPaths { dfsOf<String?>() }
3034
.sortedBy { it.name }
3135

36+
@Test
37+
fun first() {
38+
dfGroup.select {
39+
first { it.data.any { it == "Alice" } }
40+
.recursively()
41+
}.alsoDebug()
42+
43+
dfGroup.select {
44+
first { it.data.any { it == "London" } }.recursively()
45+
}.alsoDebug()
46+
}
47+
3248
@Test
3349
fun recursively() {
3450
dfGroup.getColumnsWithPaths { recursively() }.sortedBy { it.name } shouldBe recursivelyGoal

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

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1630,7 +1630,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
16301630
* NOTE: This is an identity call and can be omitted in most cases. However, it can still prove useful
16311631
* for readability or in combination with [recursively].
16321632
*/
1633-
public fun ColumnSet<*>.all(): ColumnSet<*> = wrap()
1633+
public fun <C> ColumnSet<C>.all(): ColumnSet<C> = allInternal() as ColumnSet<C>
16341634

16351635
/**
16361636
* @include [CommonAllDocs]
@@ -1642,7 +1642,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
16421642
*
16431643
* `df.`[select][select]` { "pathTo"["myGroup"].`[all][all]`() }`
16441644
*/
1645-
public fun SingleColumn<*>.all(): ColumnSet<*> = transformSingle { it.children() }
1645+
public fun SingleColumn<*>.all(): ColumnSet<*> = allInternal()
16461646

16471647
/**
16481648
* @include [CommonAllDocs]
@@ -1694,6 +1694,8 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
16941694
toColumnAccessor().allDfs(includeGroups)
16951695

16961696
/**
1697+
* ## Recursively / Rec
1698+
*
16971699
* Modifies the previous call to run not only on the current column set, but also on all of its children.
16981700
*
16991701
* For example:
@@ -1705,10 +1707,16 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
17051707
*
17061708
* `df.`[select][DataFrame.select]` { `[colsOf][ColumnSet.colsOf]`<`[String][String]`>().`[recursively][ColumnSet.recursivelyImpl]`() }`
17071709
*
1708-
* This will return all columns of type [String] in lower levels (unless [includeTopLevel]` == true`).
1710+
* This will return all columns of type [String] in lower levels (unless [includeTopLevel\]` == true`).
17091711
*
17101712
* TODO
17111713
*/
1714+
private interface CommonRecursivelyDocs {
1715+
1716+
/** Example argument */
1717+
interface Examples
1718+
}
1719+
17121720
public fun <C> ColumnSet<C>.recursively(
17131721
includeGroups: Boolean = true,
17141722
includeTopLevel: Boolean = true,
@@ -1722,7 +1730,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
17221730
public fun SingleColumn<*>.recursively(
17231731
includeGroups: Boolean = true,
17241732
includeTopLevel: Boolean = true,
1725-
): ColumnSet<*> = all().recursivelyImpl(includeTopLevel = includeTopLevel, includeGroups = includeGroups)
1733+
): ColumnSet<*> = recursivelyImpl(includeTopLevel = includeTopLevel, includeGroups = includeGroups)
17261734

17271735
public fun SingleColumn<*>.rec(
17281736
includeGroups: Boolean = true,
@@ -2107,13 +2115,23 @@ internal fun <T, C> ColumnsSelector<T, C>.filter(predicate: (ColumnWithPath<C>)
21072115
*/
21082116
internal fun ColumnSet<*>.colsInternal(predicate: ColumnFilter<*>): ColumnSet<*> =
21092117
transform {
2110-
if (isSingleColumn() && it.singleOrNull()?.isColumnGroup() == true) {
2118+
if (isSingleColumnGroup(it)) {
21112119
it.single().children()
21122120
} else {
21132121
it
21142122
}.filter(predicate)
21152123
}
21162124

2125+
internal fun ColumnSet<*>.allInternal(): ColumnSet<*> =
2126+
transform {
2127+
if (isSingleColumnGroup(it)) {
2128+
it.single().children()
2129+
} else {
2130+
it
2131+
}
2132+
}
2133+
2134+
21172135
@Deprecated("Replaced with recursively()")
21182136
internal fun ColumnSet<*>.dfsInternal(predicate: (ColumnWithPath<*>) -> Boolean) =
21192137
transform { it.filter { it.isColumnGroup() }.flatMap { it.children().flattenRecursively().filter(predicate) } }

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

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,11 @@ public interface ColumnSet<out C> {
3030
): List<ColumnWithPath<C>>
3131
}
3232

33-
public interface ColumnSetTransformer {
34-
35-
// @Deprecated("see if this can be removed")
36-
// public fun transformRemainingSingle(singleColumn: SingleColumn<*>): SingleColumn<*>
37-
38-
// public fun transformSingle(singleColumn: SingleColumn<*>): ColumnSet<*>
33+
public fun interface ColumnSetTransformer {
3934

4035
public fun transform(columnSet: ColumnSet<*>): ColumnSet<*>
36+
37+
public operator fun invoke(columnSet: ColumnSet<*>): ColumnSet<*> = transform(columnSet)
4138
}
4239

4340
public class ColumnResolutionContext internal constructor(

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

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package org.jetbrains.kotlinx.dataframe.columns
22

33
import org.jetbrains.kotlinx.dataframe.DataColumn
4+
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
5+
import kotlin.contracts.ExperimentalContracts
6+
import kotlin.contracts.contract
47

58
/**
69
* Entity that can be [resolved][resolveSingle] into [DataColumn].
@@ -17,11 +20,18 @@ public interface SingleColumn<out C> : ColumnSet<C> {
1720
context: ColumnResolutionContext,
1821
transformer: ColumnSetTransformer,
1922
): List<ColumnWithPath<C>> =
20-
throw UnsupportedOperationException(
21-
"SingleColumn.resolveAfterTransform is not supported because transformations can only be applied on normal ColumnSets. use '.wrap()' to wrap a SingleColumn into a ColumnSet"
22-
)
23+
transformer(this).resolve(context) as List<ColumnWithPath<C>>
2324

2425
public fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<C>?
2526
}
2627

27-
public fun ColumnSet<*>.isSingleColumn(): Boolean = this is SingleColumn<*>
28+
@OptIn(ExperimentalContracts::class)
29+
public fun ColumnSet<*>.isSingleColumn(): Boolean {
30+
contract {
31+
returns(true) implies (this@isSingleColumn is SingleColumn<*>)
32+
}
33+
return this is SingleColumn<*>
34+
}
35+
36+
public fun ColumnSet<*>.isSingleColumnGroup(cols: List<ColumnWithPath<*>>): Boolean =
37+
isSingleColumn() && cols.singleOrNull()?.isColumnGroup() == true

0 commit comments

Comments
 (0)