Skip to content

Commit 9e433d5

Browse files
JolanrensenJolanrensen
committed
wip allExcept. Added helper function for ColumnPath to drop its start with respect to a certain parent.
1 parent 022e7b1 commit 9e433d5

File tree

6 files changed

+143
-93
lines changed

6 files changed

+143
-93
lines changed

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

Lines changed: 13 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.createColumnSet
2020
import org.jetbrains.kotlinx.dataframe.impl.columns.isMissingColumn
2121
import org.jetbrains.kotlinx.dataframe.impl.columns.transformSingle
2222
import org.jetbrains.kotlinx.dataframe.impl.getColumnsWithPaths
23+
import org.jetbrains.kotlinx.dataframe.io.renderToString
2324
import kotlin.reflect.KProperty
2425

2526
// region ColumnsSelectionDsl
@@ -237,11 +238,6 @@ public interface AllExceptColumnsSelectionDsl<out T> {
237238
allColsExcept(others.toColumnSet())
238239

239240
// reference and path
240-
public infix fun SingleColumn<DataRow<*>>.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
241-
allColsExceptInternal(other)
242-
243-
public fun SingleColumn<DataRow<*>>.allColsExcept(vararg other: ColumnReference<*>): ColumnSet<*> =
244-
allColsExceptInternal(*other)
245241

246242
// endregion
247243

@@ -268,12 +264,6 @@ public interface AllExceptColumnsSelectionDsl<out T> {
268264
public fun String.allColsExcept(vararg others: KProperty<*>): ColumnSet<*> =
269265
allColsExcept(others.toColumnSet())
270266

271-
public fun String.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
272-
columnGroup(this).allColsExcept(other)
273-
274-
public fun String.allColsExcept(vararg others: ColumnReference<*>): ColumnSet<*> =
275-
columnGroup(this).allColsExcept(*others)
276-
277267
// endregion
278268

279269
// region KProperty
@@ -299,12 +289,6 @@ public interface AllExceptColumnsSelectionDsl<out T> {
299289
public fun KProperty<*>.allColsExcept(vararg others: KProperty<*>): ColumnSet<*> =
300290
allColsExcept(others.toColumnSet())
301291

302-
public infix fun KProperty<*>.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
303-
columnGroup(this).allColsExcept(other)
304-
305-
public fun KProperty<*>.allColsExcept(vararg others: ColumnReference<*>): ColumnSet<*> =
306-
columnGroup(this).allColsExcept(*others)
307-
308292
// endregion
309293

310294
// region ColumnPath
@@ -330,29 +314,10 @@ public interface AllExceptColumnsSelectionDsl<out T> {
330314
public fun ColumnPath.allColsExcept(vararg others: KProperty<*>): ColumnSet<*> =
331315
allColsExcept(others.toColumnSet())
332316

333-
public infix fun ColumnPath.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
334-
columnGroup(this).allColsExcept(other)
335-
336-
public fun ColumnPath.allColsExcept(vararg others: ColumnReference<*>): ColumnSet<*> =
337-
columnGroup(this).allColsExcept(*others)
338-
339317
// endregion
340318

341319
// endregion
342320

343-
/**
344-
* streamlines column references such that both relative and absolute paths can be used
345-
*/
346-
// TODO remove this overload again
347-
private fun SingleColumn<DataRow<*>>.allColsExceptInternal(vararg others: ColumnReference<*>): ColumnSet<*> =
348-
allColsExceptInternal(others.toColumnSet())
349-
// transformSingleWithContext { col ->
350-
// val correctedOthers = others.map {
351-
// it.path().dropStartWrt(col.path)
352-
// }
353-
// allColsExceptInternal(correctedOthers.toColumnSet()).resolve(this)
354-
// }
355-
356321
private fun SingleColumn<DataRow<*>>.allColsExceptInternal(other: ColumnsResolver<*>) =
357322
createColumnSet { context ->
358323
val col = this.ensureIsColumnGroup().resolveSingle(context)
@@ -365,30 +330,36 @@ public interface AllExceptColumnsSelectionDsl<out T> {
365330
val parentCol = parentScope.ensureIsColumnGroup().resolveSingle(context)
366331
?: return@createColumnSet emptyList()
367332
val parentColGroup = parentCol.asColumnGroup()
368-
val parentPath = parentCol.path
369333

370334
val allCols = colGroup.getColumnsWithPaths { all() }
371335

372336
val colsToExceptRelativeToParent = parentColGroup
373337
.getColumnsWithPaths(UnresolvedColumnsPolicy.Skip) { other }
374338

375-
376339
val colsToExceptRelativeToCol = colGroup
377340
.getColumnsWithPaths(UnresolvedColumnsPolicy.Skip) { other }
378341

379342
// throw exceptions for columns that weren't in this or parent scope
380-
(colsToExceptRelativeToParent + colsToExceptRelativeToCol).groupBy { it.path }
343+
(colsToExceptRelativeToParent + colsToExceptRelativeToCol)
344+
.groupBy { it.path }
381345
.forEach { (path, cols) ->
382346
if (cols.all { it.data.isMissingColumn() }) {
347+
val columnTitles = parentColGroup.renderToString(0, 0, columnTypes = true).trim(' ', '\n', '.')
383348
throw IllegalArgumentException(
384-
"Column ${(colPath + path).joinToString()} and ${(parentPath + path).joinToString()} not found."
349+
"Column ${(colPath + path).joinToString()} and ${path.joinToString()} not found among columns: $columnTitles."
385350
)
386351
}
387352
}
388353

389354
val colsToExcept = colsToExceptRelativeToCol +
390-
colsToExceptRelativeToParent.map { // adjust the path to be relative to the current column
391-
it.changePath(it.path.dropFirst(colPath.size - parentPath.size))
355+
colsToExceptRelativeToParent.mapNotNull { // adjust the path to be relative to the current column
356+
for (i in colPath.indices) {
357+
if (colPath[i] != it.path.getOrNull(i)) {
358+
return@mapNotNull null
359+
}
360+
}
361+
val newPath = it.path.dropFirst(colPath.size)
362+
if (newPath.isEmpty()) null else it.changePath(newPath)
392363
}
393364

394365
allCols.allColumnsExceptKeepingStructure(

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

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

33
import io.kotest.assertions.throwables.shouldThrow
4+
import io.kotest.matchers.shouldBe
45
import org.intellij.lang.annotations.Language
56
import org.jetbrains.kotlinx.dataframe.DataFrame
67
import org.jetbrains.kotlinx.dataframe.DataRow
@@ -136,23 +137,37 @@ class AllExceptTests : ColumnsSelectionDslTests() {
136137
name.firstName allColsExcept pathOf("secondName")
137138
},
138139
).shouldAllBeEqual()
140+
141+
listOf(
142+
dfGroup.select { name.firstName { secondName and thirdName } },
143+
dfGroup.select { name { firstName allColsExcept "firstName" } }.alsoDebug(),
144+
dfGroup.select { name { firstName allColsExcept pathOf("firstName") } }.alsoDebug(),
145+
dfGroup.select { (name allColsExcept "firstName"["firstName"]).first().asColumnGroup().allCols() },
146+
dfGroup.remove { name.firstName.firstName }.select { name.firstName.allCols() },
147+
).shouldAllBeEqual()
139148
}
140149

141150
@Test
142151
fun temp() {
143-
df.alsoDebug()
152+
// df.alsoDebug()
144153
// df.select {
145154
// name allColsExcept "lastName"
146155
// }.alsoDebug()
147156
//
148-
// df.select {
149-
// name allColsExcept pathOf("lastName")
157+
158+
// dfGroup.select {
159+
// name {
160+
// firstName allColsExcept firstName.firstName
161+
// }
150162
// }.alsoDebug()
163+
TODO()
164+
dfGroup.select { name { firstName allColsExcept firstName.firstName } }.alsoDebug() // should work
165+
dfGroup.select { name { firstName allColsExcept firstName } }.alsoDebug() // should fail
151166

152167
shouldThrow<IllegalArgumentException> {
153168
dfGroup.select {
154169
name.firstName allColsExcept "firstName"["secondName"]
155-
}.alsoDebug()
170+
}
156171
}
157172

158173

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

Lines changed: 13 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.createColumnSet
2020
import org.jetbrains.kotlinx.dataframe.impl.columns.isMissingColumn
2121
import org.jetbrains.kotlinx.dataframe.impl.columns.transformSingle
2222
import org.jetbrains.kotlinx.dataframe.impl.getColumnsWithPaths
23+
import org.jetbrains.kotlinx.dataframe.io.renderToString
2324
import kotlin.reflect.KProperty
2425

2526
// region ColumnsSelectionDsl
@@ -191,11 +192,6 @@ public interface AllExceptColumnsSelectionDsl<out T> {
191192
allColsExcept(others.toColumnSet())
192193

193194
// reference and path
194-
public infix fun SingleColumn<DataRow<*>>.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
195-
allColsExceptInternal(other)
196-
197-
public fun SingleColumn<DataRow<*>>.allColsExcept(vararg other: ColumnReference<*>): ColumnSet<*> =
198-
allColsExceptInternal(*other)
199195

200196
// endregion
201197

@@ -222,12 +218,6 @@ public interface AllExceptColumnsSelectionDsl<out T> {
222218
public fun String.allColsExcept(vararg others: KProperty<*>): ColumnSet<*> =
223219
allColsExcept(others.toColumnSet())
224220

225-
public fun String.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
226-
columnGroup(this).allColsExcept(other)
227-
228-
public fun String.allColsExcept(vararg others: ColumnReference<*>): ColumnSet<*> =
229-
columnGroup(this).allColsExcept(*others)
230-
231221
// endregion
232222

233223
// region KProperty
@@ -253,12 +243,6 @@ public interface AllExceptColumnsSelectionDsl<out T> {
253243
public fun KProperty<*>.allColsExcept(vararg others: KProperty<*>): ColumnSet<*> =
254244
allColsExcept(others.toColumnSet())
255245

256-
public infix fun KProperty<*>.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
257-
columnGroup(this).allColsExcept(other)
258-
259-
public fun KProperty<*>.allColsExcept(vararg others: ColumnReference<*>): ColumnSet<*> =
260-
columnGroup(this).allColsExcept(*others)
261-
262246
// endregion
263247

264248
// region ColumnPath
@@ -284,29 +268,10 @@ public interface AllExceptColumnsSelectionDsl<out T> {
284268
public fun ColumnPath.allColsExcept(vararg others: KProperty<*>): ColumnSet<*> =
285269
allColsExcept(others.toColumnSet())
286270

287-
public infix fun ColumnPath.allColsExcept(other: ColumnReference<*>): ColumnSet<*> =
288-
columnGroup(this).allColsExcept(other)
289-
290-
public fun ColumnPath.allColsExcept(vararg others: ColumnReference<*>): ColumnSet<*> =
291-
columnGroup(this).allColsExcept(*others)
292-
293271
// endregion
294272

295273
// endregion
296274

297-
/**
298-
* streamlines column references such that both relative and absolute paths can be used
299-
*/
300-
// TODO remove this overload again
301-
private fun SingleColumn<DataRow<*>>.allColsExceptInternal(vararg others: ColumnReference<*>): ColumnSet<*> =
302-
allColsExceptInternal(others.toColumnSet())
303-
// transformSingleWithContext { col ->
304-
// val correctedOthers = others.map {
305-
// it.path().dropStartWrt(col.path)
306-
// }
307-
// allColsExceptInternal(correctedOthers.toColumnSet()).resolve(this)
308-
// }
309-
310275
private fun SingleColumn<DataRow<*>>.allColsExceptInternal(other: ColumnsResolver<*>) =
311276
createColumnSet { context ->
312277
val col = this.ensureIsColumnGroup().resolveSingle(context)
@@ -319,30 +284,36 @@ public interface AllExceptColumnsSelectionDsl<out T> {
319284
val parentCol = parentScope.ensureIsColumnGroup().resolveSingle(context)
320285
?: return@createColumnSet emptyList()
321286
val parentColGroup = parentCol.asColumnGroup()
322-
val parentPath = parentCol.path
323287

324288
val allCols = colGroup.getColumnsWithPaths { all() }
325289

326290
val colsToExceptRelativeToParent = parentColGroup
327291
.getColumnsWithPaths(UnresolvedColumnsPolicy.Skip) { other }
328292

329-
330293
val colsToExceptRelativeToCol = colGroup
331294
.getColumnsWithPaths(UnresolvedColumnsPolicy.Skip) { other }
332295

333296
// throw exceptions for columns that weren't in this or parent scope
334-
(colsToExceptRelativeToParent + colsToExceptRelativeToCol).groupBy { it.path }
297+
(colsToExceptRelativeToParent + colsToExceptRelativeToCol)
298+
.groupBy { it.path }
335299
.forEach { (path, cols) ->
336300
if (cols.all { it.data.isMissingColumn() }) {
301+
val columnTitles = parentColGroup.renderToString(0, 0, columnTypes = true).trim(' ', '\n', '.')
337302
throw IllegalArgumentException(
338-
"Column ${(colPath + path).joinToString()} and ${(parentPath + path).joinToString()} not found."
303+
"Column ${(colPath + path).joinToString()} and ${path.joinToString()} not found among columns: $columnTitles."
339304
)
340305
}
341306
}
342307

343308
val colsToExcept = colsToExceptRelativeToCol +
344-
colsToExceptRelativeToParent.map { // adjust the path to be relative to the current column
345-
it.changePath(it.path.dropFirst(colPath.size - parentPath.size))
309+
colsToExceptRelativeToParent.mapNotNull { // adjust the path to be relative to the current column
310+
for (i in colPath.indices) {
311+
if (colPath[i] != it.path.getOrNull(i)) {
312+
return@mapNotNull null
313+
}
314+
}
315+
val newPath = it.path.dropFirst(colPath.size)
316+
if (newPath.isEmpty()) null else it.changePath(newPath)
346317
}
347318

348319
allCols.allColumnsExceptKeepingStructure(

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@ public data class ColumnPath(val path: List<String>) : List<String> by path, Col
5050
*/
5151
public fun take(first: Int): ColumnPath = ColumnPath(path.take(first))
5252

53-
public fun replaceLast(name: String): ColumnPath = ColumnPath(if (path.size < 2) listOf(name) else path.dropLast(1) + name)
53+
public fun replaceLast(name: String): ColumnPath =
54+
ColumnPath(if (path.size < 2) listOf(name) else path.dropLast(1) + name)
5455

5556
/**
5657
* Returns a shortened [ColumnPath] containing just the last [last] elements.
@@ -80,3 +81,34 @@ public data class ColumnPath(val path: List<String>) : List<String> by path, Col
8081

8182
override fun <C> get(column: ColumnReference<C>): ColumnAccessor<C> = ColumnAccessorImpl(this + column.path())
8283
}
84+
85+
/**
86+
* Drops the overlapping start of the child path with respect to the parent path, and returns the resulting ColumnPath.
87+
*
88+
* For example:
89+
* ```kt
90+
* val parentPath = pathOf("a", "b", "c")
91+
* val childPath = pathOf("a", "b", "c", "d", "e")
92+
*
93+
* childPath.dropStartWrt(parentPath) // returns pathOf("d", "e")
94+
* ```
95+
*
96+
* @param otherPath The parent path to compare against.
97+
* @return The resulting ColumnPath after dropping the overlapping start.
98+
*/
99+
internal fun ColumnPath.dropStartWrt(otherPath: ColumnPath): ColumnPath {
100+
val first = dropOverlappingStartOfChild(parent = otherPath, child = this)
101+
return ColumnPath(first)
102+
}
103+
104+
internal fun <T> dropOverlappingStartOfChild(parent: List<T>, child: List<T>): List<T> {
105+
var indexToRemoveTill = 0
106+
for (i in child.indices) {
107+
val subFirst = child.subList(0, i + 1)
108+
val sufficientSubSecond = parent.subList(maxOf(0, parent.size - (i + 1)), parent.size)
109+
if (subFirst == sufficientSubSecond) {
110+
indexToRemoveTill = i + 1
111+
}
112+
}
113+
return child.subList(indexToRemoveTill, child.size)
114+
}

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package org.jetbrains.kotlinx.dataframe.api
2+
3+
import io.kotest.matchers.shouldBe
4+
import org.jetbrains.kotlinx.dataframe.columns.dropOverlappingStartOfChild
5+
import org.jetbrains.kotlinx.dataframe.columns.dropStartWrt
6+
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
7+
import org.junit.Test
8+
9+
class ColumnPathTests : TestBase() {
10+
11+
@Test
12+
fun `should trim overlapping start of first list from the end of second list`() {
13+
val parent = pathOf("something", "name", "firstName")
14+
val child = pathOf("name", "firstName", "secondName")
15+
16+
dropOverlappingStartOfChild(parent, child) shouldBe listOf("secondName")
17+
child.dropStartWrt(parent) shouldBe pathOf("secondName")
18+
}
19+
20+
@Test
21+
fun `should return first list as is when there is no overlap`() {
22+
val parent = pathOf("city", "country")
23+
val child = pathOf("name", "firstName", "secondName")
24+
25+
dropOverlappingStartOfChild(parent, child) shouldBe listOf("name", "firstName", "secondName")
26+
child.dropStartWrt(parent) shouldBe pathOf("name", "firstName", "secondName")
27+
}
28+
29+
@Test
30+
fun `should return empty list when first list is completely overlapped`() {
31+
val parent = pathOf("city", "name", "firstName")
32+
val child = pathOf("name", "firstName")
33+
34+
dropOverlappingStartOfChild(parent, child) shouldBe emptyList()
35+
child.dropStartWrt(parent) shouldBe pathOf()
36+
}
37+
38+
@Test
39+
fun `if parent is empty`() {
40+
val parent = pathOf()
41+
val child = pathOf("name", "firstName")
42+
43+
dropOverlappingStartOfChild(parent, child) shouldBe listOf("name", "firstName")
44+
child.dropStartWrt(parent) shouldBe pathOf("name", "firstName")
45+
}
46+
}

0 commit comments

Comments
 (0)