Skip to content

Commit a8d3e1d

Browse files
JolanrensenJolanrensen
committed
tests for new col functions
1 parent 4e76c25 commit a8d3e1d

File tree

2 files changed

+192
-2
lines changed

2 files changed

+192
-2
lines changed

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

Lines changed: 96 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.matchers.shouldBe
4-
import org.jetbrains.kotlinx.dataframe.impl.columns.singleImpl
4+
import org.jetbrains.kotlinx.dataframe.DataFrame
5+
import org.jetbrains.kotlinx.dataframe.DataRow
6+
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
7+
import org.jetbrains.kotlinx.dataframe.impl.columns.asValueColumn
58
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
69
import org.jetbrains.kotlinx.dataframe.samples.api.age
710
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
811
import org.jetbrains.kotlinx.dataframe.samples.api.isHappy
12+
import org.jetbrains.kotlinx.dataframe.samples.api.lastName
913
import org.jetbrains.kotlinx.dataframe.samples.api.name
1014
import org.junit.Test
1115

@@ -110,4 +114,95 @@ class ColumnsSelectionDslTests : TestBase() {
110114
}
111115
} shouldBe df.select { name.firstName }
112116
}
117+
118+
@Test
119+
fun col() {
120+
df.select { col("age") } shouldBe df.select { age }
121+
df.select { col<Int>("age") } shouldBe df.select { age }
122+
df.select { col(pathOf("age")) } shouldBe df.select { age }
123+
df.select { col<Int>(pathOf("age")) } shouldBe df.select { age }
124+
df.select { col(Person::age) } shouldBe df.select { age }
125+
126+
df.select { colGroup("name").col("firstName") } shouldBe df.select { name.firstName }
127+
df.select { colGroup("name").col<String>("firstName") } shouldBe df.select { name.firstName }
128+
df.select { colGroup("name").col(pathOf("firstName")) } shouldBe df.select { name.firstName }
129+
df.select { colGroup("name").col<String>(pathOf("firstName")) } shouldBe df.select { name.firstName }
130+
df.select { colGroup("name").col(Name::firstName) } shouldBe df.select { name.firstName }
131+
}
132+
133+
@DataSchema
134+
interface FirstNames {
135+
val firstName: String
136+
val secondName: String?
137+
val thirdName: String?
138+
}
139+
140+
@DataSchema
141+
interface MyName : Name {
142+
val firstNames: FirstNames
143+
}
144+
145+
@Test
146+
fun colGroup() {
147+
val firstNames by columnGroup<FirstNames>()
148+
val dfGroup = df.convert { name.firstName }.to {
149+
val firstName by it
150+
val secondName by it.map<_, String?> { null }.asValueColumn()
151+
val thirdName by it.map<_, String?> { null }.asValueColumn()
152+
153+
dataFrameOf(firstName, secondName, thirdName)
154+
.cast<FirstNames>(verify = true)
155+
.asColumnGroup(firstNames)
156+
}
157+
158+
dfGroup.print(columnTypes = true, title = true)
159+
160+
dfGroup.select { colGroup("name") } shouldBe dfGroup.select { name }
161+
dfGroup.select { colGroup<String>("name") } shouldBe dfGroup.select { name }
162+
dfGroup.select { colGroup(pathOf("name")) } shouldBe dfGroup.select { name }
163+
dfGroup.select { colGroup<String>(pathOf("name")) } shouldBe dfGroup.select { name }
164+
dfGroup.select { colGroup(Person::name) } shouldBe dfGroup.select { name }
165+
166+
dfGroup.select { colGroup("name").colGroup("firstNames") } shouldBe dfGroup.select { name[firstNames] }
167+
dfGroup.select { colGroup("name").colGroup<String>("firstNames") } shouldBe dfGroup.select { name[firstNames] }
168+
dfGroup.select { colGroup("name").colGroup(pathOf("firstNames")) } shouldBe dfGroup.select { name[firstNames] }
169+
dfGroup.select { colGroup("name").colGroup<String>(pathOf("firstNames")) } shouldBe dfGroup.select { name[firstNames] }
170+
dfGroup.select { colGroup("name").colGroup(MyName::firstNames) } shouldBe dfGroup.select { name[firstNames] }
171+
}
172+
173+
@DataSchema
174+
interface PersonWithFrame : Person {
175+
val frameCol: DataFrame<Person>
176+
}
177+
178+
@Test
179+
fun frameCol() {
180+
val frameCol by frameColumn<Person>()
181+
182+
val dfWithFrames = df
183+
.add {
184+
expr { df } into frameCol
185+
}
186+
.convert { name }.to {
187+
val firstName by it.asColumnGroup().firstName
188+
val lastName by it.asColumnGroup().lastName
189+
190+
@Suppress("NAME_SHADOWING")
191+
val frameCol by it.map { df }.asFrameColumn()
192+
193+
dataFrameOf(firstName, lastName, frameCol).asColumnGroup("name")
194+
}
195+
196+
dfWithFrames.select { frameCol("frameCol") } shouldBe dfWithFrames.select { frameCol }
197+
dfWithFrames.select { frameCol<Person>("frameCol") } shouldBe dfWithFrames.select { frameCol }
198+
dfWithFrames.select { frameCol(pathOf("frameCol")) } shouldBe dfWithFrames.select { frameCol }
199+
dfWithFrames.select { frameCol<Person>(pathOf("frameCol")) } shouldBe dfWithFrames.select { frameCol }
200+
dfWithFrames.select { frameCol(PersonWithFrame::frameCol) } shouldBe dfWithFrames.select { frameCol }
201+
202+
dfWithFrames.select { colGroup("name").frameCol("frameCol") } shouldBe dfWithFrames.select { name[frameCol] }
203+
dfWithFrames.select { colGroup("name").frameCol<Person>("frameCol") } shouldBe dfWithFrames.select { name[frameCol] }
204+
dfWithFrames.select { colGroup("name").frameCol(pathOf("frameCol")) } shouldBe dfWithFrames.select { name[frameCol] }
205+
dfWithFrames.select { colGroup("name").frameCol<Person>(pathOf("frameCol")) } shouldBe dfWithFrames.select { name[frameCol] }
206+
dfWithFrames.select { colGroup("name").frameCol(PersonWithFrame::frameCol) } shouldBe dfWithFrames.select { name[frameCol] }
207+
}
113208
}

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

Lines changed: 96 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.matchers.shouldBe
4-
import org.jetbrains.kotlinx.dataframe.impl.columns.singleImpl
4+
import org.jetbrains.kotlinx.dataframe.DataFrame
5+
import org.jetbrains.kotlinx.dataframe.DataRow
6+
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
7+
import org.jetbrains.kotlinx.dataframe.impl.columns.asValueColumn
58
import org.jetbrains.kotlinx.dataframe.samples.api.TestBase
69
import org.jetbrains.kotlinx.dataframe.samples.api.age
710
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
811
import org.jetbrains.kotlinx.dataframe.samples.api.isHappy
12+
import org.jetbrains.kotlinx.dataframe.samples.api.lastName
913
import org.jetbrains.kotlinx.dataframe.samples.api.name
1014
import org.junit.Test
1115

@@ -110,4 +114,95 @@ class ColumnsSelectionDslTests : TestBase() {
110114
}
111115
} shouldBe df.select { name.firstName }
112116
}
117+
118+
@Test
119+
fun col() {
120+
df.select { col("age") } shouldBe df.select { age }
121+
df.select { col<Int>("age") } shouldBe df.select { age }
122+
df.select { col(pathOf("age")) } shouldBe df.select { age }
123+
df.select { col<Int>(pathOf("age")) } shouldBe df.select { age }
124+
df.select { col(Person::age) } shouldBe df.select { age }
125+
126+
df.select { colGroup("name").col("firstName") } shouldBe df.select { name.firstName }
127+
df.select { colGroup("name").col<String>("firstName") } shouldBe df.select { name.firstName }
128+
df.select { colGroup("name").col(pathOf("firstName")) } shouldBe df.select { name.firstName }
129+
df.select { colGroup("name").col<String>(pathOf("firstName")) } shouldBe df.select { name.firstName }
130+
df.select { colGroup("name").col(Name::firstName) } shouldBe df.select { name.firstName }
131+
}
132+
133+
@DataSchema
134+
interface FirstNames {
135+
val firstName: String
136+
val secondName: String?
137+
val thirdName: String?
138+
}
139+
140+
@DataSchema
141+
interface MyName : Name {
142+
val firstNames: FirstNames
143+
}
144+
145+
@Test
146+
fun colGroup() {
147+
val firstNames by columnGroup<FirstNames>()
148+
val dfGroup = df.convert { name.firstName }.to {
149+
val firstName by it
150+
val secondName by it.map<_, String?> { null }.asValueColumn()
151+
val thirdName by it.map<_, String?> { null }.asValueColumn()
152+
153+
dataFrameOf(firstName, secondName, thirdName)
154+
.cast<FirstNames>(verify = true)
155+
.asColumnGroup(firstNames)
156+
}
157+
158+
dfGroup.print(columnTypes = true, title = true)
159+
160+
dfGroup.select { colGroup("name") } shouldBe dfGroup.select { name }
161+
dfGroup.select { colGroup<String>("name") } shouldBe dfGroup.select { name }
162+
dfGroup.select { colGroup(pathOf("name")) } shouldBe dfGroup.select { name }
163+
dfGroup.select { colGroup<String>(pathOf("name")) } shouldBe dfGroup.select { name }
164+
dfGroup.select { colGroup(Person::name) } shouldBe dfGroup.select { name }
165+
166+
dfGroup.select { colGroup("name").colGroup("firstNames") } shouldBe dfGroup.select { name[firstNames] }
167+
dfGroup.select { colGroup("name").colGroup<String>("firstNames") } shouldBe dfGroup.select { name[firstNames] }
168+
dfGroup.select { colGroup("name").colGroup(pathOf("firstNames")) } shouldBe dfGroup.select { name[firstNames] }
169+
dfGroup.select { colGroup("name").colGroup<String>(pathOf("firstNames")) } shouldBe dfGroup.select { name[firstNames] }
170+
dfGroup.select { colGroup("name").colGroup(MyName::firstNames) } shouldBe dfGroup.select { name[firstNames] }
171+
}
172+
173+
@DataSchema
174+
interface PersonWithFrame : Person {
175+
val frameCol: DataFrame<Person>
176+
}
177+
178+
@Test
179+
fun frameCol() {
180+
val frameCol by frameColumn<Person>()
181+
182+
val dfWithFrames = df
183+
.add {
184+
expr { df } into frameCol
185+
}
186+
.convert { name }.to {
187+
val firstName by it.asColumnGroup().firstName
188+
val lastName by it.asColumnGroup().lastName
189+
190+
@Suppress("NAME_SHADOWING")
191+
val frameCol by it.map { df }.asFrameColumn()
192+
193+
dataFrameOf(firstName, lastName, frameCol).asColumnGroup("name")
194+
}
195+
196+
dfWithFrames.select { frameCol("frameCol") } shouldBe dfWithFrames.select { frameCol }
197+
dfWithFrames.select { frameCol<Person>("frameCol") } shouldBe dfWithFrames.select { frameCol }
198+
dfWithFrames.select { frameCol(pathOf("frameCol")) } shouldBe dfWithFrames.select { frameCol }
199+
dfWithFrames.select { frameCol<Person>(pathOf("frameCol")) } shouldBe dfWithFrames.select { frameCol }
200+
dfWithFrames.select { frameCol(PersonWithFrame::frameCol) } shouldBe dfWithFrames.select { frameCol }
201+
202+
dfWithFrames.select { colGroup("name").frameCol("frameCol") } shouldBe dfWithFrames.select { name[frameCol] }
203+
dfWithFrames.select { colGroup("name").frameCol<Person>("frameCol") } shouldBe dfWithFrames.select { name[frameCol] }
204+
dfWithFrames.select { colGroup("name").frameCol(pathOf("frameCol")) } shouldBe dfWithFrames.select { name[frameCol] }
205+
dfWithFrames.select { colGroup("name").frameCol<Person>(pathOf("frameCol")) } shouldBe dfWithFrames.select { name[frameCol] }
206+
dfWithFrames.select { colGroup("name").frameCol(PersonWithFrame::frameCol) } shouldBe dfWithFrames.select { name[frameCol] }
207+
}
113208
}

0 commit comments

Comments
 (0)