Skip to content

Commit c932196

Browse files
JolanrensenJolanrensen
committed
updating from master
1 parent d9af0b9 commit c932196

File tree

3 files changed

+162
-19
lines changed

3 files changed

+162
-19
lines changed

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

Lines changed: 125 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,18 @@ public interface ColumnSelectionDsl<out T> : ColumnsContainer<T> {
4242

4343
public operator fun <C> ColumnPath.invoke(): DataColumn<C> = getColumn(this).cast()
4444

45+
public operator fun <T> KProperty<T>.invoke(): DataColumn<T> = this@ColumnSelectionDsl[this]
46+
47+
public operator fun <T> KProperty<DataRow<T>>.invoke(): ColumnGroup<T> = this@ColumnSelectionDsl[this]
48+
49+
public operator fun <T> KProperty<DataFrame<T>>.invoke(): FrameColumn<T> = this@ColumnSelectionDsl[this]
50+
51+
public operator fun <T, R> KProperty<DataRow<T>>.get(column: KProperty<R>): DataColumn<R> = invoke()[column]
52+
53+
public operator fun <T, R> KProperty<DataRow<T>>.get(column: KProperty<DataRow<R>>): ColumnGroup<R> = invoke()[column]
54+
55+
public operator fun <T, R> KProperty<DataRow<T>>.get(column: KProperty<DataFrame<R>>): FrameColumn<R> = invoke()[column]
56+
4557
public operator fun <C> String.invoke(): DataColumn<C> = getColumn(this).cast()
4658

4759
public operator fun String.get(column: String): ColumnPath = pathOf(this, column)
@@ -68,6 +80,8 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
6880
public operator fun String.rangeTo(endInclusive: String): ColumnSet<*> =
6981
toColumnAccessor().rangeTo(endInclusive.toColumnAccessor())
7082

83+
public operator fun KProperty<*>.rangeTo(endInclusive: KProperty<*>): ColumnSet<*> = toColumnAccessor().rangeTo(endInclusive.toColumnAccessor())
84+
7185
public operator fun AnyColumnReference.rangeTo(endInclusive: AnyColumnReference): ColumnSet<*> =
7286
object : ColumnSet<Any?> {
7387
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
@@ -102,19 +116,19 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
102116
transform { it.flatMap { col -> names.mapNotNull { col.getChild(it) } } }
103117
}
104118

119+
public fun <C> ColumnSet<*>.cols(firstCol: KProperty<C>, vararg otherCols: KProperty<C>): ColumnSet<C> =
120+
(listOf(firstCol) + otherCols).let { names ->
121+
transform { it.flatMap { col -> names.mapNotNull { col.getChild(it) } } }
122+
}
123+
105124
public fun ColumnSet<*>.cols(vararg indices: Int): ColumnSet<Any?> =
106125
transform { it.flatMap { it.children().let { children -> indices.map { children[it] } } } }
107126

108127
public fun ColumnSet<*>.cols(range: IntRange): ColumnSet<Any?> =
109-
transform { it.flatMap { it.children().subList(range.start, range.endInclusive + 1) } }
128+
transform { it.flatMap { it.children().subList(range.first, range.last + 1) } }
110129

111130
// region select
112131

113-
public fun <C> ColumnSet<DataRow<C>>.select(vararg columns: String): ColumnSet<*> = select { columns.toColumns() }
114-
115-
public fun <C, R> ColumnSet<DataRow<C>>.select(vararg columns: KProperty<R>): ColumnSet<R> =
116-
select { columns.toColumns() }
117-
118132
public fun <C, R> ColumnSet<DataRow<C>>.select(selector: ColumnsSelector<C, R>): ColumnSet<R> = createColumnSet {
119133
this@select.resolve(it).flatMap { group ->
120134
group.asColumnGroup().getColumnsWithPaths(selector).map {
@@ -123,6 +137,12 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
123137
}
124138
}
125139

140+
public fun <C> ColumnSet<DataRow<C>>.select(vararg columns: String): ColumnSet<*> = select { columns.toColumns() }
141+
142+
public fun <C, R> ColumnSet<DataRow<C>>.select(vararg columns: ColumnReference<R>): ColumnSet<R> = select { columns.toColumns() }
143+
144+
public fun <C, R> ColumnSet<DataRow<C>>.select(vararg columns: KProperty<R>): ColumnSet<R> = select { columns.toColumns() }
145+
126146
// endregion
127147

128148
// endregion
@@ -133,6 +153,8 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
133153

134154
public fun String.dfs(predicate: (ColumnWithPath<*>) -> Boolean): ColumnSet<*> = toColumnAccessor().dfs(predicate)
135155

156+
public fun <C> KProperty<C>.dfs(predicate: (ColumnWithPath<*>) -> Boolean): ColumnSet<*> = toColumnAccessor().dfs(predicate)
157+
136158
// endregion
137159

138160
// region all
@@ -141,13 +163,17 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
141163

142164
public fun String.all(): ColumnSet<*> = toColumnAccessor().transformSingle { it.children() }
143165

166+
public fun KProperty<*>.all(): ColumnSet<*> = toColumnAccessor().transformSingle { it.children() }
167+
144168
// region allDfs
145169

146170
public fun ColumnSet<*>.allDfs(includeGroups: Boolean = false): ColumnSet<Any?> =
147171
if (includeGroups) dfs { true } else dfs { !it.isColumnGroup() }
148172

149173
public fun String.allDfs(includeGroups: Boolean = false): ColumnSet<Any?> = toColumnAccessor().allDfs(includeGroups)
150174

175+
public fun KProperty<*>.allDfs(includeGroups: Boolean = false): ColumnSet<Any?> = toColumnAccessor().allDfs(includeGroups)
176+
151177
// endregion
152178

153179
// region allAfter
@@ -166,13 +192,24 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
166192

167193
public fun SingleColumn<*>.allAfter(colName: String): ColumnSet<Any?> = allAfter(pathOf(colName))
168194
public fun SingleColumn<*>.allAfter(column: AnyColumnReference): ColumnSet<Any?> = allAfter(column.path())
195+
public fun SingleColumn<*>.allAfter(column: KProperty<*>): ColumnSet<Any?> = allAfter(column.toColumnAccessor().path())
196+
197+
public fun String.allAfter(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allAfter(colPath)
198+
public fun String.allAfter(colName: String): ColumnSet<Any?> = toColumnAccessor().allAfter(colName)
199+
public fun String.allAfter(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allAfter(column)
200+
public fun String.allAfter(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allAfter(column)
201+
202+
public fun KProperty<*>.allAfter(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allAfter(colPath)
203+
public fun KProperty<*>.allAfter(colName: String): ColumnSet<Any?> = toColumnAccessor().allAfter(colName)
204+
public fun KProperty<*>.allAfter(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allAfter(column)
205+
public fun KProperty<*>.allAfter(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allAfter(column)
169206

170207
// endregion
171208

172-
// region allSince
209+
// region allFrom
173210

174211
// including current
175-
public fun SingleColumn<*>.allSince(colPath: ColumnPath): ColumnSet<Any?> {
212+
public fun SingleColumn<*>.allFrom(colPath: ColumnPath): ColumnSet<Any?> {
176213
var take = false
177214
return children {
178215
if (take) true
@@ -183,8 +220,19 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
183220
}
184221
}
185222

186-
public fun SingleColumn<*>.allSince(colName: String): ColumnSet<Any?> = allSince(pathOf(colName))
187-
public fun SingleColumn<*>.allSince(column: AnyColumnReference): ColumnSet<Any?> = allSince(column.path())
223+
public fun SingleColumn<*>.allFrom(colName: String): ColumnSet<Any?> = allFrom(pathOf(colName))
224+
public fun SingleColumn<*>.allFrom(column: AnyColumnReference): ColumnSet<Any?> = allFrom(column.path())
225+
public fun SingleColumn<*>.allFrom(column: KProperty<*>): ColumnSet<Any?> = allFrom(column.toColumnAccessor().path())
226+
227+
public fun String.allFrom(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allFrom(colPath)
228+
public fun String.allFrom(colName: String): ColumnSet<Any?> = toColumnAccessor().allFrom(colName)
229+
public fun String.allFrom(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allFrom(column)
230+
public fun String.allFrom(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allFrom(column)
231+
232+
public fun KProperty<*>.allFrom(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allFrom(colPath)
233+
public fun KProperty<*>.allFrom(colName: String): ColumnSet<Any?> = toColumnAccessor().allFrom(colName)
234+
public fun KProperty<*>.allFrom(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allFrom(column)
235+
public fun KProperty<*>.allFrom(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allFrom(column)
188236

189237
// endregion
190238

@@ -204,13 +252,24 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
204252

205253
public fun SingleColumn<*>.allBefore(colName: String): ColumnSet<Any?> = allBefore(pathOf(colName))
206254
public fun SingleColumn<*>.allBefore(column: AnyColumnReference): ColumnSet<Any?> = allBefore(column.path())
255+
public fun SingleColumn<*>.allBefore(column: KProperty<*>): ColumnSet<Any?> = allBefore(column.toColumnAccessor().path())
256+
257+
public fun String.allBefore(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allBefore(colPath)
258+
public fun String.allBefore(colName: String): ColumnSet<Any?> = toColumnAccessor().allBefore(colName)
259+
public fun String.allBefore(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allBefore(column)
260+
public fun String.allBefore(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allBefore(column)
261+
262+
public fun KProperty<*>.allBefore(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allBefore(colPath)
263+
public fun KProperty<*>.allBefore(colName: String): ColumnSet<Any?> = toColumnAccessor().allBefore(colName)
264+
public fun KProperty<*>.allBefore(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allBefore(column)
265+
public fun KProperty<*>.allBefore(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allBefore(column)
207266

208267
// endregion
209268

210-
// region allUntil
269+
// region allUpTo
211270

212271
// including current
213-
public fun SingleColumn<*>.allUntil(colPath: ColumnPath): ColumnSet<Any?> {
272+
public fun SingleColumn<*>.allUpTo(colPath: ColumnPath): ColumnSet<Any?> {
214273
var take = true
215274
return children {
216275
if (!take) false
@@ -221,21 +280,46 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
221280
}
222281
}
223282

224-
public fun SingleColumn<*>.allUntil(colName: String): ColumnSet<Any?> = allUntil(pathOf(colName))
225-
public fun SingleColumn<*>.allUntil(column: AnyColumnReference): ColumnSet<Any?> = allUntil(column.path())
283+
public fun SingleColumn<*>.allUpTo(colName: String): ColumnSet<Any?> = allUpTo(pathOf(colName))
284+
public fun SingleColumn<*>.allUpTo(column: AnyColumnReference): ColumnSet<Any?> = allUpTo(column.path())
285+
public fun SingleColumn<*>.allUpTo(column: KProperty<*>): ColumnSet<Any?> = allUpTo(column.toColumnAccessor().path())
286+
287+
public fun String.allUpTo(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allUpTo(colPath)
288+
public fun String.allUpTo(colName: String): ColumnSet<Any?> = toColumnAccessor().allUpTo(colName)
289+
public fun String.allUpTo(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allUpTo(column)
290+
public fun String.allUpTo(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allUpTo(column)
291+
292+
public fun KProperty<*>.allUpTo(colPath: ColumnPath): ColumnSet<Any?> = toColumnAccessor().allUpTo(colPath)
293+
public fun KProperty<*>.allUpTo(colName: String): ColumnSet<Any?> = toColumnAccessor().allUpTo(colName)
294+
public fun KProperty<*>.allUpTo(column: AnyColumnReference): ColumnSet<Any?> = toColumnAccessor().allUpTo(column)
295+
public fun KProperty<*>.allUpTo(column: KProperty<*>): ColumnSet<Any?> = toColumnAccessor().allUpTo(column)
226296

227297
// endregion
228298

229299
// endregion
230300

301+
// region groups
302+
231303
public fun SingleColumn<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
232304
children { it.isColumnGroup() && filter(it.asColumnGroup()) } as ColumnSet<AnyRow>
233305

234-
public fun <C> ColumnSet<C>.children(predicate: (ColumnWithPath<Any?>) -> Boolean = { true }): ColumnSet<Any?> =
306+
public fun String.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
307+
toColumnAccessor().groups(filter)
308+
309+
public fun KProperty<*>.groups(filter: (ColumnGroup<*>) -> Boolean = { true }): ColumnSet<AnyRow> =
310+
toColumnAccessor().groups(filter)
311+
312+
// endregion
313+
314+
// region children
315+
316+
public fun ColumnSet<*>.children(predicate: (ColumnWithPath<Any?>) -> Boolean = { true }): ColumnSet<Any?> =
235317
transform { it.flatMap { it.children().filter { predicate(it) } } }
236318

237319
public fun ColumnGroupReference.children(): ColumnSet<Any?> = transformSingle { it.children() }
238320

321+
// endregion
322+
239323
public operator fun <C> List<DataColumn<C>>.get(range: IntRange): ColumnSet<C> =
240324
ColumnsList(subList(range.first, range.last + 1))
241325

@@ -245,6 +329,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
245329
transform { it.mapNotNull { it.getChild(colName) } }
246330

247331
public operator fun <C> ColumnSet<*>.get(column: ColumnReference<C>): ColumnSet<C> = cols(column)
332+
public operator fun <C> ColumnSet<*>.get(column: KProperty<C>): ColumnSet<C> = cols(column)
248333

249334
public fun SingleColumn<AnyRow>.take(n: Int): ColumnSet<*> = transformSingle { it.children().take(n) }
250335
public fun SingleColumn<AnyRow>.takeLast(n: Int): ColumnSet<*> = transformSingle { it.children().takeLast(n) }
@@ -296,9 +381,18 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
296381
toColumnAccessor().into(column.columnName)
297382

298383
public infix fun <C> ColumnReference<C>.named(newName: String): ColumnReference<C> = renamedReference(newName)
299-
public infix fun <C> ColumnReference<C>.named(name: KProperty<*>): ColumnReference<C> = named(name.columnName)
384+
public infix fun <C> ColumnReference<C>.named(nameFrom: ColumnReference<*>): ColumnReference<C> = named(nameFrom.name)
385+
public infix fun <C> ColumnReference<C>.named(nameFrom: KProperty<*>): ColumnReference<C> = named(nameFrom.columnName)
300386

301387
public infix fun String.named(newName: String): ColumnReference<Any?> = toColumnAccessor().named(newName)
388+
public infix fun String.named(nameFrom: ColumnReference<*>): ColumnReference<Any?> = toColumnAccessor().named(nameFrom.name)
389+
public infix fun String.named(nameFrom: KProperty<*>): ColumnReference<Any?> = toColumnAccessor().named(nameFrom.columnName)
390+
391+
public infix fun <C> KProperty<C>.named(newName: String): ColumnReference<C> = toColumnAccessor().named(newName)
392+
393+
public infix fun <C> KProperty<C>.named(nameFrom: ColumnReference<*>): ColumnReference<C> = toColumnAccessor().named(nameFrom.name)
394+
395+
public infix fun <C> KProperty<C>.named(nameFrom: KProperty<*>): ColumnReference<C> = toColumnAccessor().named(nameFrom.columnName)
302396

303397
// region and
304398

@@ -339,6 +433,21 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
339433
// endregion
340434

341435
public fun <C> ColumnSet<C>.distinct(): ColumnSet<C> = DistinctColumnSet(this)
436+
437+
public fun <C> String.dfsOf(type: KType, predicate: (ColumnWithPath<C>) -> Boolean = { true }): ColumnSet<*> =
438+
toColumnAccessor().dfsOf(type, predicate)
439+
440+
public fun <C> KProperty<*>.dfsOf(type: KType, predicate: (ColumnWithPath<C>) -> Boolean = { true }): ColumnSet<*> =
441+
toColumnAccessor().dfsOf(type, predicate)
442+
443+
public fun String.colsOf(type: KType): ColumnSet<Any?> = toColumnAccessor().colsOf(type)
444+
public fun KProperty<*>.colsOf(type: KType): ColumnSet<Any?> = toColumnAccessor().colsOf(type)
445+
446+
public fun <C> String.colsOf(type: KType, filter: (DataColumn<C>) -> Boolean): ColumnSet<Any?> =
447+
toColumnAccessor().colsOf(type, filter)
448+
449+
public fun <C> KProperty<*>.colsOf(type: KType, filter: (DataColumn<C>) -> Boolean): ColumnSet<Any?> =
450+
toColumnAccessor().colsOf(type, filter)
342451
}
343452

344453
public inline fun <T, reified R> ColumnsSelectionDsl<T>.expr(
@@ -349,7 +458,6 @@ public inline fun <T, reified R> ColumnsSelectionDsl<T>.expr(
349458

350459
internal fun <T, C> ColumnsSelector<T, C>.filter(predicate: (ColumnWithPath<C>) -> Boolean): ColumnsSelector<T, C> =
351460
{ this@filter(it, it).filter(predicate) }
352-
// internal fun Columns<*>.filter(predicate: (AnyCol) -> Boolean) = transform { it.filter { predicate(it.data) } }
353461

354462
internal fun ColumnSet<*>.colsInternal(predicate: (AnyCol) -> Boolean) =
355463
transform { it.flatMap { it.children().filter { predicate(it.data) } } }

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -706,6 +706,41 @@ class Access : TestBase() {
706706
// SampleEnd
707707
}
708708

709+
@Test
710+
fun columnSelectors_kproperties() {
711+
// SampleStart
712+
// by column name
713+
df.select { it[Person::name] }
714+
df.select { (Person::name)() }
715+
df.select { col(Person::name) }
716+
717+
// by column path
718+
df.select { it[Person::name][Name::firstName] }
719+
df.select { Person::name[Name::firstName] }
720+
721+
// with a new name
722+
df.select { Person::name named "Full Name" }
723+
724+
// converted
725+
df.select { Person::name[Name::firstName].map { it.lowercase() } }
726+
727+
// column arithmetics
728+
df.select { 2021 - (Person::age)() }
729+
730+
// two columns
731+
df.select { Person::name and Person::age }
732+
733+
// range of columns
734+
df.select { Person::name..Person::age }
735+
736+
// all children of ColumnGroup
737+
df.select { Person::name.all() }
738+
739+
// dfs traversal of all children columns
740+
df.select { Person::name.allDfs() }
741+
// SampleEnd
742+
}
743+
709744
@Test
710745
fun columnSelectors_strings() {
711746
// SampleStart

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1899,9 +1899,9 @@ class DataFrameTests : BaseTest() {
18991899
@Test
19001900
fun `select all after`() {
19011901
typed.select { allAfter(age) } shouldBe typed.select { city and weight }
1902-
typed.select { allSince(age) } shouldBe typed.select { age and city and weight }
1902+
typed.select { allFrom(age) } shouldBe typed.select { age and city and weight }
19031903
typed.select { allBefore(age) } shouldBe typed.select { name }
1904-
typed.select { allUntil(age) } shouldBe typed.select { name and age }
1904+
typed.select { allUpTo(age) } shouldBe typed.select { name and age }
19051905
}
19061906

19071907
@Test

0 commit comments

Comments
 (0)