Skip to content

Commit 97d9879

Browse files
JolanrensenJolanrensen
committed
fixed first and last functions, added tests
1 parent dfe149d commit 97d9879

File tree

4 files changed

+285
-48
lines changed

4 files changed

+285
-48
lines changed

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

Lines changed: 95 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -259,13 +259,13 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
259259

260260
/**
261261
* ## First
262-
* Returns the first column in this [ColumnSet] that adheres to the given [condition\].
262+
* Returns the first column in this [ColumnSet] or [ColumnGroup] that adheres to the given [condition\].
263263
*
264264
* For example:
265265
*
266266
* {@includeArg [Examples]}
267267
*
268-
* @param [condition\] The [ColumnFilter] condition that the column must adhere to.
268+
* @param [condition\] The optional [ColumnFilter] condition that the column must adhere to.
269269
* @return A [SingleColumn] containing the first column that adheres to the given [condition\].
270270
* @see [last\]
271271
*/
@@ -277,61 +277,95 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
277277

278278
/**
279279
* ## First
280-
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that adheres to the given [condition].
280+
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
281+
*
282+
* For example:
283+
*
284+
* `df.`[select][select]` { `[colsOf][colsOf]`<`[String][String]`>().`[first][first]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
285+
*
286+
* `df.`[select][select]` { `[colsOf][colsOf]`<`[Int][Int]`>().`[first][first]`() }`
287+
*
288+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
289+
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the first column that adheres to the given [condition].
290+
* @see [last]
291+
*/
292+
public fun <C> ColumnSet<C>.first(condition: ColumnFilter<C> = { true }): SingleColumn<C> =
293+
transform { listOf(it.first(condition)) }.single()
294+
295+
/**
296+
* ## First
297+
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
281298
*
282299
* For example:
283300
*
284301
* `df.`[select][select]` { `[first][first]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
285302
*
303+
* `df.`[select][select]` { `[first][first]`() }`
304+
*
305+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
306+
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the first column that adheres to the given [condition].
307+
* @see [last]
308+
*/
309+
public fun first(condition: ColumnFilter<*> = { true }): SingleColumn<*> =
310+
all().first(condition)
311+
312+
/**
313+
* ## First
314+
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
315+
*
316+
* For example:
317+
*
318+
* `df.`[select][select]` { myColumnGroup.`[first][first]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
319+
*
286320
* `df.`[select][select]` { myColumnGroup.`[first][first]`() }`
287321
*
288-
* @param [condition] The [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
322+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
289323
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the first column that adheres to the given [condition].
290324
* @see [last]
291325
*/
292-
public fun <C> ColumnSet<C>.first(condition: ColumnFilter<C> = { true }): SingleColumn<C> =
293-
children { condition(it.cast()) }[0].cast()
326+
public fun ColumnGroupReference.first(condition: ColumnFilter<*> = { true }): SingleColumn<*> =
327+
all().first(condition)
294328

295329
/**
296330
* ## First
297-
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that adheres to the given [condition].
331+
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
298332
*
299333
* For example:
300334
*
301335
* `df.`[select][select]` { "myColumnGroup".`[first][first]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
302336
*
303-
* @param [condition] The [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
337+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
304338
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the first column that adheres to the given [condition].
305339
* @see [last]
306340
*/
307341
public fun String.first(condition: ColumnFilter<*> = { true }): SingleColumn<*> =
308-
toColumnAccessor().first(condition)
342+
getColumnGroup(this).first(condition)
309343

310344
/**
311345
* ## First
312-
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that adheres to the given [condition].
346+
* Returns the first column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
313347
*
314348
* For example:
315349
*
316350
* `df.`[select][select]` { Type::myColumnGroup.`[first][first]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
317351
*
318-
* @param [condition] The [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
352+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
319353
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the first column that adheres to the given [condition].
320354
* @see [last]
321355
*/
322-
public fun <C> KProperty<C>.first(condition: ColumnFilter<C>): SingleColumn<*> =
323-
toColumnAccessor().first(condition)
356+
public fun KProperty<*>.first(condition: ColumnFilter<*>): SingleColumn<*> =
357+
getColumnGroup(this).first(condition)
324358

325359

326360
/**
327361
* ## Last
328-
* Returns the last column in this [ColumnSet] that adheres to the given [condition\].
362+
* Returns the last column in this [ColumnSet] or [ColumnGroup] that adheres to the given [condition\].
329363
*
330364
* For example:
331365
*
332366
* {@includeArg [Examples]}
333367
*
334-
* @param [condition\] The [ColumnFilter] condition that the column must adhere to.
368+
* @param [condition\] The optional [ColumnFilter] condition that the column must adhere to.
335369
* @return A [SingleColumn] containing the last column that adheres to the given [condition\].
336370
* @see [first\]
337371
*/
@@ -343,53 +377,84 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
343377

344378
/**
345379
* ## Last
346-
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that adheres to the given [condition].
380+
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
347381
*
348382
* For example:
349383
*
350-
* `df.`[select][select]` { `[last][last]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
384+
* `df.`[select][select]` { `[colsOf][colsOf]`<`[String][String]`>().`[last][last]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
351385
*
352-
* `df.`[select][select]` { myColumnGroup.`[first][last]`() }`
386+
* `df.`[select][select]` { `[colsOf][colsOf]`<`[Int][Int]`>().`[first][last]`() }`
353387
*
354-
* @param [condition] The [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
388+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
355389
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the last column that adheres to the given [condition].
356390
* @see [first]
357391
*/
358392
public fun <C> ColumnSet<C>.last(condition: ColumnFilter<C> = { true }): SingleColumn<C> =
359-
children { condition(it.cast()) }
360-
.transform { listOf(it.last()) }
361-
.single()
362-
.cast()
393+
transform { listOf(it.last(condition)) }.single()
394+
395+
/**
396+
* ## Last
397+
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
398+
*
399+
* For example:
400+
*
401+
* `df.`[select][select]` { `[last][last]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
402+
*
403+
* `df.`[select][select]` { `[first][last]`() }`
404+
*
405+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
406+
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the last column that adheres to the given [condition].
407+
* @see [first]
408+
*/
409+
public fun last(condition: ColumnFilter<*> = { true }): SingleColumn<*> =
410+
all().last(condition)
411+
412+
/**
413+
* ## Last
414+
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
415+
*
416+
* For example:
417+
*
418+
* `df.`[select][select]` { myColumnGroup.`[last][last]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
419+
*
420+
* `df.`[select][select]` { myColumnGroup.`[last][last]`() }`
421+
*
422+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
423+
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the last column that adheres to the given [condition].
424+
* @see [first]
425+
*/
426+
public fun ColumnGroupReference.last(condition: ColumnFilter<*> = { true }): SingleColumn<*> =
427+
all().last(condition)
363428

364429
/**
365430
* ## Last
366-
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that adheres to the given [condition].
431+
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
367432
*
368433
* For example:
369434
*
370435
* `df.`[select][select]` { "myColumnGroup".`[last][last]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
371436
*
372-
* @param [condition] The [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
437+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
373438
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the last column that adheres to the given [condition].
374439
* @see [first]
375440
*/
376441
public fun String.last(condition: ColumnFilter<*> = { true }): SingleColumn<*> =
377-
toColumnAccessor().last(condition)
442+
getColumnGroup(this).last(condition)
378443

379444
/**
380445
* ## Last
381-
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that adheres to the given [condition].
446+
* Returns the last column in this [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] or [ColumnGroup][org.jetbrains.kotlinx.dataframe.columns.ColumnGroup] that adheres to the given [condition].
382447
*
383448
* For example:
384449
*
385450
* `df.`[select][select]` { Type::myColumnGroup.`[last][last]` { it.`[name][ColumnReference.name]`().`[startsWith][String.startsWith]`("year") } }`
386451
*
387-
* @param [condition] The [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
452+
* @param [condition] The optional [ColumnFilter][org.jetbrains.kotlinx.dataframe.ColumnFilter] condition that the column must adhere to.
388453
* @return A [SingleColumn][org.jetbrains.kotlinx.dataframe.columns.SingleColumn] containing the last column that adheres to the given [condition].
389454
* @see [first]
390455
*/
391-
public fun <C> KProperty<C>.last(condition: ColumnFilter<C>): SingleColumn<*> =
392-
toColumnAccessor().last(condition)
456+
public fun KProperty<*>.last(condition: ColumnFilter<*>): SingleColumn<*> =
457+
getColumnGroup(this).last(condition)
393458

394459
public fun <C> ColumnSet<C>.single(condition: ColumnFilter<C>): SingleColumn<C> =
395460
transform { listOf(it.single(condition)) }.single()

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

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package org.jetbrains.kotlinx.dataframe.api
2+
3+
import io.kotest.matchers.shouldBe
4+
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
5+
import org.jetbrains.kotlinx.dataframe.samples.api.age
6+
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
7+
import org.jetbrains.kotlinx.dataframe.samples.api.isHappy
8+
import org.jetbrains.kotlinx.dataframe.samples.api.lastName
9+
import org.jetbrains.kotlinx.dataframe.samples.api.name
10+
import org.junit.Test
11+
12+
class ColumnsSelectionDslTests : TestBase() {
13+
14+
@Test
15+
fun first() {
16+
df.select { all().first() } shouldBe df.select { first() }
17+
18+
df.select { all().first() } shouldBe df.select { name }
19+
20+
df.select { first() } shouldBe df.select { name }
21+
22+
df.select { first { it.name().startsWith("a") } } shouldBe df.select { age }
23+
24+
df.select {
25+
name.first {
26+
it.any { it == "Alice" }
27+
}
28+
} shouldBe df.select {
29+
name.colsOf<String>().first {
30+
it.any { it == "Alice" }
31+
}
32+
}
33+
34+
df.select {
35+
"name".first {
36+
it.any { it == "Alice" }
37+
}
38+
} shouldBe df.select { name.firstName }
39+
40+
df.select {
41+
Person::name.first {
42+
it.any { it == "Alice" }
43+
}
44+
} shouldBe df.select { name.firstName }
45+
}
46+
47+
@Test
48+
fun last() {
49+
df.select { all().last() } shouldBe df.select { last() }
50+
51+
df.select { all().last() } shouldBe df.select { isHappy }
52+
53+
df.select { last() } shouldBe df.select { isHappy }
54+
55+
df.select { last { it.name().startsWith("a") } } shouldBe df.select { age }
56+
57+
df.select {
58+
name.last {
59+
it.any { it == "Alice" }
60+
}
61+
} shouldBe df.select {
62+
name.colsOf<String>().last {
63+
it.any { it == "Alice" }
64+
}
65+
}
66+
}
67+
68+
}

0 commit comments

Comments
 (0)