Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

More variations of "move", AddDsl.group, CreateDataFrameDsl String.invoke #1035

Merged
merged 3 commits into from
Jan 29, 2025
Merged
Show file tree
Hide file tree
Changes from all 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
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,7 @@ public class AddDsl<T>(
@AccessApiOverload
public fun group(column: AnyColumnGroupAccessor, body: AddDsl<T>.() -> Unit): Unit = group(column.name(), body)

@Interpretable("AddDslNamedGroup")
public fun group(name: String, body: AddDsl<T>.() -> Unit) {
val dsl = AddDsl(df)
body(dsl)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T>

public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)

@AccessApiOverload
public fun <T> InsertClause<T>.after(column: ColumnAccessor<*>): DataFrame<T> = after(column.path())

@AccessApiOverload
Expand Down
27 changes: 27 additions & 0 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,13 @@ public fun <T> DataFrame<T>.moveTo(newColumnIndex: Int, vararg columns: KPropert

// region moveToLeft

@Refine
@Interpretable("MoveToLeft1")
public fun <T> DataFrame<T>.moveToLeft(columns: ColumnsSelector<T, *>): DataFrame<T> = move(columns).toLeft()

public fun <T> DataFrame<T>.moveToLeft(vararg columns: String): DataFrame<T> = moveToLeft { columns.toColumnSet() }

@AccessApiOverload
public fun <T> DataFrame<T>.moveToLeft(vararg columns: AnyColumnReference): DataFrame<T> =
moveToLeft { columns.toColumnSet() }

Expand All @@ -68,13 +71,17 @@ public fun <T> DataFrame<T>.moveToLeft(vararg columns: KProperty<*>): DataFrame<

// region moveToRight

@Refine
@Interpretable("MoveToRight1")
public fun <T> DataFrame<T>.moveToRight(columns: ColumnsSelector<T, *>): DataFrame<T> = move(columns).toRight()

public fun <T> DataFrame<T>.moveToRight(vararg columns: String): DataFrame<T> = moveToRight { columns.toColumnSet() }

@AccessApiOverload
public fun <T> DataFrame<T>.moveToRight(vararg columns: AnyColumnReference): DataFrame<T> =
moveToRight { columns.toColumnSet() }

@AccessApiOverload
public fun <T> DataFrame<T>.moveToRight(vararg columns: KProperty<*>): DataFrame<T> =
moveToRight { columns.toColumnSet() }

Expand All @@ -94,6 +101,11 @@ public fun <T, C> MoveClause<T, C>.into(
newPathExpression = column,
)

/**
* Move a single selected column to top level with a new name
*/
@Refine
@Interpretable("MoveInto0")
public fun <T, C> MoveClause<T, C>.into(column: String): DataFrame<T> = pathOf(column).let { path -> into { path } }

public fun <T, C> MoveClause<T, C>.intoIndexed(
Expand All @@ -109,11 +121,19 @@ public fun <T, C> MoveClause<T, C>.intoIndexed(

// region under

@Refine
@Interpretable("MoveUnder0")
public fun <T, C> MoveClause<T, C>.under(column: String): DataFrame<T> = pathOf(column).let { path -> under { path } }

@AccessApiOverload
public fun <T, C> MoveClause<T, C>.under(column: AnyColumnGroupAccessor): DataFrame<T> =
column.path().let { path -> under { path } }

/**
* Move selected columns under existing column group
*/
@Refine
@Interpretable("MoveUnder1")
public fun <T, C> MoveClause<T, C>.under(
column: ColumnsSelectionDsl<T>.(ColumnWithPath<C>) -> AnyColumnReference,
): DataFrame<T> =
Expand All @@ -138,10 +158,13 @@ public fun <T, C> MoveClause<T, C>.toTop(

// region after

@Refine
@Interpretable("MoveAfter0")
public fun <T, C> MoveClause<T, C>.after(column: ColumnSelector<T, *>): DataFrame<T> = afterOrBefore(column, true)

public fun <T, C> MoveClause<T, C>.after(column: String): DataFrame<T> = after { column.toColumnAccessor() }

@AccessApiOverload
public fun <T, C> MoveClause<T, C>.after(column: AnyColumnReference): DataFrame<T> = after { column }

@AccessApiOverload
Expand All @@ -157,8 +180,12 @@ fun <T, C> MoveColsClause<T, C>.before(column: String) = before { column.toColum
fun <T, C> MoveColsClause<T, C>.before(column: ColumnSelector<T, *>) = afterOrBefore(column, false)
*/

@Refine
@Interpretable("MoveToLeft0")
public fun <T, C> MoveClause<T, C>.toLeft(): DataFrame<T> = to(0)

@Refine
@Interpretable("MoveToRight0")
public fun <T, C> MoveClause<T, C>.toRight(): DataFrame<T> = to(df.ncol)

public class MoveClause<T, C>(internal val df: DataFrame<T>, internal val columns: ColumnsSelector<T, C>) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,7 @@ public abstract class CreateDataFrameDsl<T> : TraversePropertiesDsl {

public inline fun <reified R> inferType(noinline expression: (T) -> R): InferType<T, R> = InferType(expression)

@Interpretable("ToDataFrameDslStringInvoke")
public abstract operator fun String.invoke(builder: CreateDataFrameDsl<T>.() -> Unit)
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -795,4 +795,10 @@ class DataFrameTreeTests : BaseTest() {
val df = typed2.convert { nameAndCity }.asFrame { it.remove { city } }
df.nameAndCity.columns() shouldBe typed2.nameAndCity.remove { city }.columns()
}

@Test
fun `move under existing group`() {
val df = typed2.move { age }.under { nameAndCity }
df.nameAndCity.columnNames() shouldBe listOf("name", "city", "age")
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -68,3 +68,15 @@ class AddDslStringInvoke : AbstractInterpreter<Unit>() {
dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
}
}

class AddDslNamedGroup : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.name: String by arg()
val Arguments.body by dsl()

override fun Arguments.interpret() {
val addDsl = AddDslApproximation(mutableListOf())
body(addDsl, emptyMap())
dsl.columns.add(SimpleColumnGroup(name, addDsl.columns))
}
}
Original file line number Diff line number Diff line change
@@ -1,8 +1,13 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api

import org.jetbrains.kotlinx.dataframe.api.after
import org.jetbrains.kotlinx.dataframe.api.into
import org.jetbrains.kotlinx.dataframe.api.move
import org.jetbrains.kotlinx.dataframe.api.pathOf
import org.jetbrains.kotlinx.dataframe.api.moveToLeft
import org.jetbrains.kotlinx.dataframe.api.toLeft
import org.jetbrains.kotlinx.dataframe.api.toRight
import org.jetbrains.kotlinx.dataframe.api.toTop
import org.jetbrains.kotlinx.dataframe.api.under
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
Expand Down Expand Up @@ -30,4 +35,74 @@ class ToTop : AbstractSchemaModificationInterpreter() {
}
}

class MoveUnder0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()
val Arguments.column: String by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.df).map { it.path }
return receiver.df.asDataFrame().move { columns.toColumnSet() }.under(column).toPluginDataFrameSchema()
}
}

class MoveUnder1 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()
val Arguments.column: SingleColumnApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.df).map { it.path }
return receiver.df.asDataFrame().move { columns.toColumnSet() }.under { column.col.path }.toPluginDataFrameSchema()
}
}

class MoveInto0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()
val Arguments.column: String by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.df).map { it.path }
return receiver.df.asDataFrame().move { columns.toColumnSet() }.into(column).toPluginDataFrameSchema()
}
}

class MoveToLeft0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.df).map { it.path }
return receiver.df.asDataFrame().move { columns.toColumnSet() }.toLeft().toPluginDataFrameSchema()
}
}

class MoveToLeft1 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: ColumnsResolver by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = columns.resolve(receiver).map { it.path }
return receiver.asDataFrame().moveToLeft { columns.toColumnSet() }.toPluginDataFrameSchema()
}
}


class MoveToRight0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.df).map { it.path }
return receiver.df.asDataFrame().move { columns.toColumnSet() }.toRight().toPluginDataFrameSchema()
}
}

class MoveAfter0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: MoveClauseApproximation by arg()
val Arguments.column: SingleColumnApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.df).map { it.path }
return receiver.df.asDataFrame().move { columns.toColumnSet() }.after { column.col.path }.toPluginDataFrameSchema()
}
}


class MoveClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,18 @@ class Properties0 : AbstractInterpreter<Unit>() {
}
}

class ToDataFrameDslStringInvoke : AbstractInterpreter<Unit>() {
val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
val Arguments.receiver: String by arg()
val Arguments.builder by dsl()

override fun Arguments.interpret() {
val addDsl = CreateDataFrameDslImplApproximation()
builder(addDsl, emptyMap())
dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
}
}

class CreateDataFrameConfiguration {
var maxDepth = DEFAULT_MAX_DEPTH
var traverseConfiguration: TraverseConfiguration = TraverseConfiguration()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.name.StandardClassIds
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslNamedGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate
Expand All @@ -88,6 +89,13 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveAfter0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveInto0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToLeft0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToLeft1
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToRight0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveUnder0
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveUnder1
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairConstructor
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairToConstructor
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
Expand All @@ -96,6 +104,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDslStringInvoke
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
Expand Down Expand Up @@ -235,6 +244,7 @@ internal inline fun <reified T> String.load(): T {
"toDataFrameDsl" -> ToDataFrameDsl()
"toDataFrame" -> ToDataFrame()
"toDataFrameDefault" -> ToDataFrameDefault()
"ToDataFrameDslStringInvoke" -> ToDataFrameDslStringInvoke()
"DataFrameOf0" -> DataFrameOf0()
"DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
"ToDataFrameColumn" -> ToDataFrameColumn()
Expand All @@ -246,6 +256,7 @@ internal inline fun <reified T> String.load(): T {
"FlattenDefault" -> FlattenDefault()
"AddId" -> AddId()
"AddDslStringInvoke" -> AddDslStringInvoke()
"AddDslNamedGroup" -> AddDslNamedGroup()
"MapToFrame" -> MapToFrame()
"Move0" -> Move0()
"ToTop" -> ToTop()
Expand All @@ -255,6 +266,13 @@ internal inline fun <reified T> String.load(): T {
"ValueCounts" -> ValueCounts()
"RenameToCamelCase" -> RenameToCamelCase()
"RenameToCamelCaseClause" -> RenameToCamelCaseClause()
"MoveUnder0" -> MoveUnder0()
"MoveUnder1" -> MoveUnder1()
"MoveInto0" -> MoveInto0()
"MoveToLeft0" -> MoveToLeft0()
"MoveToLeft1" -> MoveToLeft1()
"MoveToRight0" -> MoveToRight0()
"MoveAfter0" -> MoveAfter0()
else -> error("$this")
} as T
}
4 changes: 4 additions & 0 deletions plugins/kotlin-dataframe/testData/box/addDsl.kt
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,12 @@ fun box(): String {
"group" {
"a" from { it }
}
group("group1") {
"b" from { it }
}
}

df.group.a
df.group1.b
return "OK"
}
Loading
Loading