Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 59 additions & 0 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,65 @@ internal fun <T> InsertClause<T>.afterImpl(columnPath: ColumnPath): DataFrame<T>

// endregion

// region before

/**
* Inserts the new column previously specified with [insert]
* at the position immediately before the selected [column] (on the same level).
*
* For more information: {@include [DocumentationUrls.Insert]}
*
* See [Grammar][InsertDocs.Grammar] for more details.
*
* See also: [SelectingColumns.Dsl].
*
* ### Examples:
* ```kotlin
* // Insert a new column "age" before the "name" column
* df.insert(age).before { name }
*
* // Insert a new column "sum" before the nested "min" column (inside the "stats" column group)
* val dfWithSum = df.insert("sum") { a + b }.before { stats.min }
* ```
*
* @param [column] The [ColumnSelector] used to choose an existing column in this [DataFrame],
* before which the new column will be inserted.
* @return A new [DataFrame] with the inserted column placed before the selected column.
*/
@Refine
@Interpretable("InsertBefore0")
public fun <T> InsertClause<T>.before(column: ColumnSelector<T, *>): DataFrame<T> = beforeImpl(df.getColumnPath(column))

/**
* Inserts the new column previously specified with [insert]
* at the position immediately before the column with the given [name][column].
*
* For more information: {@include [DocumentationUrls.Insert]}
*
* See [Grammar][InsertDocs.Grammar] for more details.
*
* See also: [SelectingColumns.ColumnNames].
*
* ### Example
* ```kotlin
* // Insert a new column "age" before the "name" column
* df.insert(age).before("name")
* ```
*
* @param [column] The [String] name of the column in this [DataFrame]
* before which the new column will be inserted.
* @return A new [DataFrame] with the inserted column placed before the specified column.
*/
public fun <T> InsertClause<T>.before(column: String): DataFrame<T> =
df.add(this.column).move(this.column).before(column)

internal fun <T> InsertClause<T>.beforeImpl(columnPath: ColumnPath): DataFrame<T> {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's a DataFrame convention of putting the internal Impl-ementation of API functions under impl/api/<functionName>.kt. I see we missed it for afterImpl, but let's not make it worse with beforeImpl :) Could you move both beforeImpl and afterImpl to impl/api/insert.kt? :)
thanks in advance!

val dstPath = ColumnPath(columnPath.removeAt(columnPath.size - 1) + column.name())
return df.insertImpl(dstPath, column).move { dstPath }.before { columnPath }
}

// endregion

// region at

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import org.jetbrains.kotlinx.dataframe.api.asDataFrame
import org.jetbrains.kotlinx.dataframe.api.asFrame
import org.jetbrains.kotlinx.dataframe.api.asGroupBy
import org.jetbrains.kotlinx.dataframe.api.at
import org.jetbrains.kotlinx.dataframe.api.before
import org.jetbrains.kotlinx.dataframe.api.by
import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.colsOf
Expand Down Expand Up @@ -702,6 +703,22 @@ class DataFrameTreeTests : BaseTest() {
typed2.insert(colName) { nameAndCity.name.reversed() }.after { nameAndCity.name }.check()
}

@Test
fun `insert column before`() {
val colName = "reversed"

fun DataFrame<GroupedPerson>.check() {
nameAndCity.columnsCount() shouldBe 3
nameAndCity.columnNames() shouldBe listOf(
typed2.nameAndCity.name.name(),
colName,
typed2.nameAndCity.city.name(),
)
}

typed2.insert(colName) { nameAndCity.name.reversed() }.before { nameAndCity.city }.check()
}

@Test
fun append() {
val res = typed2.append(listOf("Bill", "San Francisco"), null, 66)
Expand Down
Loading