Skip to content

Commit 845e6cd

Browse files
JolanrensenJolanrensen
committed
col indices and ranges
1 parent 9acddcd commit 845e6cd

File tree

4 files changed

+374
-8
lines changed

4 files changed

+374
-8
lines changed

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

Lines changed: 99 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1510,9 +1510,11 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
15101510
// endregion
15111511

15121512
// region cols
1513+
15131514
// TODO treat same as first/single etc.
15141515

15151516
// region predicate
1517+
15161518
public fun <C> ColumnSet<C>.cols(
15171519
predicate: ColumnFilter<C> = { true },
15181520
): ColumnSet<C> = colsInternal(predicate as ColumnFilter<*>).cast()
@@ -1760,7 +1762,7 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
17601762
vararg otherCols: ColumnPath,
17611763
): ColumnSet<*> = cols(firstCol, *otherCols)
17621764

1763-
// endRegion
1765+
// endregion
17641766

17651767
// region properties
17661768

@@ -1828,16 +1830,109 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
18281830

18291831
// region indices
18301832

1831-
public fun ColumnSet<*>.cols(
1833+
public fun <C> ColumnSet<C>.cols(
1834+
firstIndex: Int,
1835+
vararg otherIndices: Int,
1836+
): ColumnSet<C> = object : ColumnSet<C> {
1837+
1838+
@Suppress("UNCHECKED_CAST")
1839+
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>> =
1840+
this@cols.resolve(context)
1841+
.let(::dataFrameOf)
1842+
.asColumnGroup()
1843+
.cols(firstIndex, *otherIndices)
1844+
.resolve(context) as List<ColumnWithPath<C>>
1845+
}
1846+
1847+
public operator fun <C> ColumnSet<C>.get(
18321848
firstIndex: Int,
18331849
vararg otherIndices: Int,
1834-
): ColumnSet<Any?> = headPlusArray(firstIndex, otherIndices).let { indices ->
1850+
): ColumnSet<C> = cols(firstIndex, *otherIndices)
1851+
1852+
public fun SingleColumn<AnyRow>.cols(
1853+
firstIndex: Int,
1854+
vararg otherIndices: Int,
1855+
): ColumnSet<*> = headPlusArray(firstIndex, otherIndices).let { indices ->
18351856
transform { it.flatMap { it.children().let { children -> indices.map { children[it] } } } }
18361857
}
18371858

1838-
public fun ColumnSet<*>.cols(range: IntRange): ColumnSet<Any?> =
1859+
/**
1860+
*
1861+
*/
1862+
public operator fun SingleColumn<AnyRow>.get(
1863+
firstIndex: Int,
1864+
vararg otherIndices: Int,
1865+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1866+
1867+
public fun String.cols(
1868+
firstIndex: Int,
1869+
vararg otherIndices: Int,
1870+
): ColumnSet<*> = getColumnGroup(this).cols(firstIndex, *otherIndices)
1871+
1872+
public operator fun String.get(
1873+
firstIndex: Int,
1874+
vararg otherIndices: Int,
1875+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1876+
1877+
public fun ColumnPath.cols(
1878+
firstIndex: Int,
1879+
vararg otherIndices: Int,
1880+
): ColumnSet<*> = getColumnGroup(this).cols(firstIndex, *otherIndices)
1881+
1882+
public operator fun ColumnPath.get(
1883+
firstIndex: Int,
1884+
vararg otherIndices: Int,
1885+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1886+
1887+
public fun KProperty<*>.cols(
1888+
firstIndex: Int,
1889+
vararg otherIndices: Int,
1890+
): ColumnSet<*> = getColumnGroup(this).cols(firstIndex, *otherIndices)
1891+
1892+
public operator fun KProperty<*>.get(
1893+
firstIndex: Int,
1894+
vararg otherIndices: Int,
1895+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1896+
1897+
// endregion
1898+
1899+
// region ranges
1900+
1901+
public fun <C> ColumnSet<C>.cols(range: IntRange): ColumnSet<C> =
1902+
object : ColumnSet<C> {
1903+
1904+
@Suppress("UNCHECKED_CAST")
1905+
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>> =
1906+
this@cols.resolve(context)
1907+
.let(::dataFrameOf)
1908+
.asColumnGroup()
1909+
.cols(range)
1910+
.resolve(context) as List<ColumnWithPath<C>>
1911+
}
1912+
1913+
public operator fun <C> ColumnSet<C>.get(range: IntRange): ColumnSet<C> = cols(range)
1914+
1915+
1916+
public fun SingleColumn<AnyRow>.cols(range: IntRange): ColumnSet<*> =
18391917
transform { it.flatMap { it.children().subList(range.first, range.last + 1) } }
18401918

1919+
/**
1920+
*
1921+
*/
1922+
public operator fun SingleColumn<AnyRow>.get(range: IntRange): ColumnSet<*> = cols(range)
1923+
1924+
public fun String.cols(range: IntRange): ColumnSet<*> = getColumnGroup(this).cols(range)
1925+
1926+
public operator fun String.get(range: IntRange): ColumnSet<*> = cols(range)
1927+
1928+
public fun ColumnPath.cols(range: IntRange): ColumnSet<*> = getColumnGroup(this).cols(range)
1929+
1930+
public operator fun ColumnPath.get(range: IntRange): ColumnSet<*> = cols(range)
1931+
1932+
public fun KProperty<*>.cols(range: IntRange): ColumnSet<*> = getColumnGroup(this).cols(range)
1933+
1934+
public operator fun KProperty<*>.get(range: IntRange): ColumnSet<*> = cols(range)
1935+
18411936
// endregion
18421937

18431938
// endregion

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,22 @@ class ColumnsSelectionDslTests : TestBase() {
329329

330330
@Test
331331
fun `cols and get with column paths`() {
332+
listOf(
333+
df.select {
334+
all().cols(pathOf("name", "firstName"))
335+
},
336+
df.select {
337+
cols(pathOf("name", "firstName"))
338+
},
339+
df.select {
340+
pathOf("name", "firstName")
341+
}
342+
).reduce { acc, dataFrame ->
343+
acc shouldBe dataFrame
344+
dataFrame
345+
}
346+
347+
332348
df.select { all().cols(pathOf("name"), pathOf("age")) } shouldBe df.select {
333349
cols(
334350
pathOf("name"),
@@ -404,5 +420,79 @@ class ColumnsSelectionDslTests : TestBase() {
404420
pathOf("name")[Name::firstName, Name::lastName]
405421
}
406422
}
423+
424+
@Test
425+
fun `cols and get with indices`() {
426+
df.select { all().cols(0, 1) } shouldBe df.select { cols(0, 1) }
427+
df.select { all()[0, 1] } shouldBe df.select { this[0, 1] }
428+
429+
df.select {
430+
name.cols(0, 1)
431+
} shouldBe df.select {
432+
name.colsOf<String>().cols(0, 1)
433+
}
434+
435+
df.select {
436+
// name[0, 1]
437+
name.cols(0, 1)
438+
} shouldBe df.select {
439+
name.colsOf<String>()[0, 1]
440+
}
441+
442+
df.select {
443+
"name".cols(0, 1)
444+
} shouldBe df.select {
445+
Person::name.cols(0, 1)
446+
}
447+
448+
df.select {
449+
"name"[0, 1]
450+
} shouldBe df.select {
451+
Person::name[0, 1]
452+
}
453+
454+
df.select {
455+
pathOf("name").cols(0, 1)
456+
} shouldBe df.select {
457+
pathOf("name")[0, 1]
458+
}
459+
}
460+
461+
@Test
462+
fun `cols and get with range`() {
463+
df.select { all().cols(0..1) } shouldBe df.select { cols(0..1) }
464+
df.select { all()[0..1] } shouldBe df.select { this[0..1] }
465+
466+
df.select {
467+
name.cols(0..1)
468+
} shouldBe df.select {
469+
name.colsOf<String>().cols(0..1)
470+
}
471+
472+
df.select {
473+
// name[0..1]
474+
name.cols(0..1)
475+
} shouldBe df.select {
476+
name.colsOf<String>()[0..1]
477+
}
478+
479+
df.select {
480+
"name".cols(0..1)
481+
} shouldBe df.select {
482+
Person::name.cols(0..1)
483+
}
484+
485+
df.select {
486+
"name"[0..1]
487+
} shouldBe df.select {
488+
Person::name[0..1]
489+
}
490+
491+
df.select {
492+
pathOf("name").cols(0..1)
493+
} shouldBe df.select {
494+
pathOf("name")[0..1]
495+
}
496+
}
407497
}
408498

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

Lines changed: 95 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1125,17 +1125,108 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
11251125

11261126
// region indices
11271127

1128-
public fun ColumnSet<*>.cols(
1128+
public fun <C> ColumnSet<C>.cols(
1129+
firstIndex: Int,
1130+
vararg otherIndices: Int,
1131+
): ColumnSet<C> = object : ColumnSet<C> {
1132+
1133+
@Suppress("UNCHECKED_CAST")
1134+
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>> =
1135+
this@cols.resolve(context)
1136+
.let(::dataFrameOf)
1137+
.asColumnGroup()
1138+
.cols(firstIndex, *otherIndices)
1139+
.resolve(context) as List<ColumnWithPath<C>>
1140+
}
1141+
1142+
public operator fun <C> ColumnSet<C>.get(
1143+
firstIndex: Int,
1144+
vararg otherIndices: Int,
1145+
): ColumnSet<C> = cols(firstIndex, *otherIndices)
1146+
1147+
public fun SingleColumn<AnyRow>.cols(
11291148
firstIndex: Int,
11301149
vararg otherIndices: Int,
1131-
): ColumnSet<Any?> = headPlusArray(firstIndex, otherIndices).let { indices ->
1150+
): ColumnSet<*> = headPlusArray(firstIndex, otherIndices).let { indices ->
11321151
transform { it.flatMap { it.children().let { children -> indices.map { children[it] } } } }
11331152
}
11341153

1135-
public fun ColumnSet<*>.cols(range: IntRange): ColumnSet<Any?> =
1154+
/**
1155+
* {@comment this function is shadowed by [ColumnGroup.get] for accessors}
1156+
*/
1157+
public operator fun SingleColumn<AnyRow>.get(
1158+
firstIndex: Int,
1159+
vararg otherIndices: Int,
1160+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1161+
1162+
public fun String.cols(
1163+
firstIndex: Int,
1164+
vararg otherIndices: Int,
1165+
): ColumnSet<*> = getColumnGroup(this).cols(firstIndex, *otherIndices)
1166+
1167+
public operator fun String.get(
1168+
firstIndex: Int,
1169+
vararg otherIndices: Int,
1170+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1171+
1172+
public fun ColumnPath.cols(
1173+
firstIndex: Int,
1174+
vararg otherIndices: Int,
1175+
): ColumnSet<*> = getColumnGroup(this).cols(firstIndex, *otherIndices)
1176+
1177+
public operator fun ColumnPath.get(
1178+
firstIndex: Int,
1179+
vararg otherIndices: Int,
1180+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1181+
1182+
public fun KProperty<*>.cols(
1183+
firstIndex: Int,
1184+
vararg otherIndices: Int,
1185+
): ColumnSet<*> = getColumnGroup(this).cols(firstIndex, *otherIndices)
1186+
1187+
public operator fun KProperty<*>.get(
1188+
firstIndex: Int,
1189+
vararg otherIndices: Int,
1190+
): ColumnSet<*> = cols(firstIndex, *otherIndices)
1191+
1192+
// endregion
1193+
1194+
// region ranges
1195+
1196+
public fun <C> ColumnSet<C>.cols(range: IntRange): ColumnSet<C> =
1197+
object : ColumnSet<C> {
1198+
1199+
@Suppress("UNCHECKED_CAST")
1200+
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<C>> =
1201+
this@cols.resolve(context)
1202+
.let(::dataFrameOf)
1203+
.asColumnGroup()
1204+
.cols(range)
1205+
.resolve(context) as List<ColumnWithPath<C>>
1206+
}
1207+
1208+
public operator fun <C> ColumnSet<C>.get(range: IntRange): ColumnSet<C> = cols(range)
1209+
1210+
1211+
public fun SingleColumn<AnyRow>.cols(range: IntRange): ColumnSet<*> =
11361212
transform { it.flatMap { it.children().subList(range.first, range.last + 1) } }
11371213

1138-
// TODO
1214+
/**
1215+
* {@comment this function is shadowed by [ColumnGroup.get] for accessors}
1216+
*/
1217+
public operator fun SingleColumn<AnyRow>.get(range: IntRange): ColumnSet<*> = cols(range)
1218+
1219+
public fun String.cols(range: IntRange): ColumnSet<*> = getColumnGroup(this).cols(range)
1220+
1221+
public operator fun String.get(range: IntRange): ColumnSet<*> = cols(range)
1222+
1223+
public fun ColumnPath.cols(range: IntRange): ColumnSet<*> = getColumnGroup(this).cols(range)
1224+
1225+
public operator fun ColumnPath.get(range: IntRange): ColumnSet<*> = cols(range)
1226+
1227+
public fun KProperty<*>.cols(range: IntRange): ColumnSet<*> = getColumnGroup(this).cols(range)
1228+
1229+
public operator fun KProperty<*>.get(range: IntRange): ColumnSet<*> = cols(range)
11391230

11401231
// endregion
11411232

0 commit comments

Comments
 (0)