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

Add df.properties() function to improve visibility of generated API for typed column access #957

Merged
merged 1 commit into from
Dec 12, 2024
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
7 changes: 6 additions & 1 deletion core/api/core.api
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ public final class org/jetbrains/dataframe/keywords/SoftKeywords$Companion {
public final fun getVALUES ()Ljava/util/List;
}

public abstract interface class org/jetbrains/kotlinx/dataframe/ColumnsContainer {
public abstract interface class org/jetbrains/kotlinx/dataframe/ColumnsContainer : org/jetbrains/kotlinx/dataframe/ColumnsScope {
public abstract fun columns ()Ljava/util/List;
public abstract fun columnsCount ()I
public abstract fun containsColumn (Ljava/lang/String;)Z
Expand Down Expand Up @@ -234,6 +234,10 @@ public final class org/jetbrains/kotlinx/dataframe/ColumnsContainer$DefaultImpls
public static fun get (Lorg/jetbrains/kotlinx/dataframe/ColumnsContainer;Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;)Lorg/jetbrains/kotlinx/dataframe/columns/FrameColumn;
}

public abstract interface class org/jetbrains/kotlinx/dataframe/ColumnsScope {
public abstract fun get (Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
}

public abstract interface class org/jetbrains/kotlinx/dataframe/DataColumn : org/jetbrains/kotlinx/dataframe/columns/BaseColumn {
public static final field Companion Lorg/jetbrains/kotlinx/dataframe/DataColumn$Companion;
public abstract fun distinct ()Lorg/jetbrains/kotlinx/dataframe/DataColumn;
Expand Down Expand Up @@ -4206,6 +4210,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/DataFrameGetKt {
public static final fun getOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;I)Lorg/jetbrains/kotlinx/dataframe/DataRow;
public static final fun getRows (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Ljava/lang/Iterable;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
public static final fun getRows (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/ranges/IntRange;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
public static final fun properties (Lorg/jetbrains/kotlinx/dataframe/DataFrame;)Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;
public static final fun rows (Lorg/jetbrains/kotlinx/dataframe/DataFrame;)Ljava/lang/Iterable;
public static final fun rowsReversed (Lorg/jetbrains/kotlinx/dataframe/DataFrame;)Ljava/lang/Iterable;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,9 @@ import kotlin.reflect.KProperty
*
* Base interface for [DataFrame] and [ColumnSelectionDsl]
*
* @param T Schema marker. Used to generate extension properties for typed column access.
* @param T Schema marker. Used to resolve generated extension properties for typed column access.
*/
public interface ColumnsContainer<out T> {
public interface ColumnsContainer<out T> : ColumnsScope<T> {

// region columns

Expand Down Expand Up @@ -54,7 +54,7 @@ public interface ColumnsContainer<out T> {

// region get

public operator fun get(columnName: String): AnyCol = getColumn(columnName)
public override operator fun get(columnName: String): AnyCol = getColumn(columnName)

public operator fun get(columnPath: ColumnPath): AnyCol = getColumn(columnPath)

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
package org.jetbrains.kotlinx.dataframe

/**
* Provides minimal API required for generated column properties:
*
* `val ColumnsScope<Schema marker>.column: DataColumn<String> get() = this["column"] as DataColumn<String>`
*
* @param T Schema marker. Used to resolve generated extension properties for typed column access.
*/
public interface ColumnsScope<out T> {
public operator fun get(columnName: String): AnyCol
}
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.AnyCol
import org.jetbrains.kotlinx.dataframe.AnyColumnReference
import org.jetbrains.kotlinx.dataframe.ColumnSelector
import org.jetbrains.kotlinx.dataframe.ColumnsContainer
import org.jetbrains.kotlinx.dataframe.ColumnsScope
import org.jetbrains.kotlinx.dataframe.ColumnsSelector
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
Expand Down Expand Up @@ -58,6 +59,28 @@ public fun <T> ColumnsContainer<T>.getFrameColumn(columnName: String): FrameColu
public fun <T> ColumnsContainer<T>.getColumnGroup(columnPath: ColumnPath): ColumnGroup<*> =
get(columnPath).asColumnGroup()

/**
* Utility property to access scope with only dataframe column properties for code completion,
* filtering out DataFrame API.
*
* It's a quick way to check that code generation in notebooks or compiler plugin
* worked as expected or find columns you're interested in.
*
* In notebooks:
* ```
* val df = DataFrame.read("file.csv")
* ==== next code cell
* df. // column properties are mixed together with methods, not easy to find unless you already know names
* df.properties(). // easy to overview available columns
* ```
* In compiler plugin:
* ```
* val df = @Import DataFrame.read("file.csv")
* df.properties().
* ```
*/
public fun <T> DataFrame<T>.properties(): ColumnsScope<T> = this

// region getColumn

public fun <T> ColumnsContainer<T>.getColumn(name: String): AnyCol =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import org.jetbrains.dataframe.impl.codeGen.InterfaceGenerationMode.WithFields
import org.jetbrains.dataframe.keywords.HardKeywords
import org.jetbrains.dataframe.keywords.ModifierKeywords
import org.jetbrains.kotlinx.dataframe.ColumnsContainer
import org.jetbrains.kotlinx.dataframe.ColumnsScope
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.DataRow
Expand Down Expand Up @@ -166,7 +167,7 @@ internal object FullyQualifiedNames : TypeRenderingStrategy {
internal object ShortNames : TypeRenderingStrategy {

private val dataRow = DataRow::class.simpleName!!
private val columnsContainer = ColumnsContainer::class.simpleName!!
private val columnsContainer = ColumnsScope::class.simpleName!!
private val dataFrame = DataFrame::class.simpleName!!
private val dataColumn = DataColumn::class.simpleName!!
private val columnGroup = ColumnGroup::class.simpleName!!
Expand Down Expand Up @@ -565,7 +566,7 @@ public fun Code.toStandaloneSnippet(packageName: String, additionalImports: List
appendLine("package $packageName")
appendLine()
}
appendLine("import org.jetbrains.kotlinx.dataframe.ColumnsContainer")
appendLine("import org.jetbrains.kotlinx.dataframe.ColumnsScope")
appendLine("import org.jetbrains.kotlinx.dataframe.DataColumn")
appendLine("import org.jetbrains.kotlinx.dataframe.DataFrame")
appendLine("import org.jetbrains.kotlinx.dataframe.DataRow")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import org.jetbrains.dataframe.impl.codeGen.InterfaceGenerationMode
import org.jetbrains.dataframe.impl.codeGen.ReplCodeGenerator
import org.jetbrains.dataframe.impl.codeGen.generate
import org.jetbrains.kotlinx.dataframe.AnyRow
import org.jetbrains.kotlinx.dataframe.ColumnsContainer
import org.jetbrains.kotlinx.dataframe.ColumnsScope
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
Expand All @@ -30,7 +30,7 @@ class CodeGenerationTests : BaseTest() {

val personShortName = Person::class.simpleName!!

val dfName = (ColumnsContainer::class).simpleName!!
val dfName = (ColumnsScope::class).simpleName!!
val dfRowName = (DataRow::class).simpleName!!
val dataCol = (DataColumn::class).simpleName!!
val dataRow = (DataRow::class).simpleName!!
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import io.kotest.matchers.shouldBe
import io.kotest.matchers.string.shouldNotBeEmpty
import org.jetbrains.dataframe.impl.codeGen.ReplCodeGenerator
import org.jetbrains.dataframe.impl.codeGen.process
import org.jetbrains.kotlinx.dataframe.ColumnsContainer
import org.jetbrains.kotlinx.dataframe.ColumnsScope
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
Expand All @@ -20,7 +20,7 @@ import org.junit.Test
@Suppress("ktlint:standard:class-naming")
class ReplCodeGenTests : BaseTest() {

val dfName = (ColumnsContainer::class).simpleName!!
val dfName = (ColumnsScope::class).simpleName!!
val dfRowName = (DataRow::class).simpleName!!
val dataCol = (DataColumn::class).simpleName!!
val intName = Int::class.simpleName!!
Expand Down Expand Up @@ -192,7 +192,7 @@ class ReplCodeGenTests : BaseTest() {
repl.process<Test3.B>()
repl.process<Test3.D>()
val c = repl.process(Test3.df, Test3::df)
"""val .*ColumnsContainer<\w*>.x:""".toRegex().findAll(c.declarations).count() shouldBe 1
"""val .*ColumnsScope<\w*>.x:""".toRegex().findAll(c.declarations).count() shouldBe 1
}

object Test4 {
Expand All @@ -216,6 +216,6 @@ class ReplCodeGenTests : BaseTest() {
repl.process<Test4.A>()
repl.process<Test4.B>()
val c = repl.process(Test4.df, Test4::df)
"""val .*ColumnsContainer<\w*>.a:""".toRegex().findAll(c.declarations).count() shouldBe 1
"""val .*ColumnsScope<\w*>.a:""".toRegex().findAll(c.declarations).count() shouldBe 1
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,8 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
companion object {
val COLUMNS_CONTAINER_ID =
CallableId(ClassId(FqName("org.jetbrains.kotlinx.dataframe"), Name.identifier("ColumnsContainer")), Name.identifier("get"))
val COLUMNS_SCOPE_ID =
CallableId(ClassId(FqName("org.jetbrains.kotlinx.dataframe"), Name.identifier("ColumnsScope")), Name.identifier("get"))
val DATA_ROW_ID =
CallableId(ClassId(FqName("org.jetbrains.kotlinx.dataframe"), Name.identifier("DataRow")), Name.identifier("get"))
}
Expand Down Expand Up @@ -196,7 +198,7 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering

val get = if (isDataColumn) {
context
.referenceFunctions(COLUMNS_CONTAINER_ID)
.referenceFunctions(COLUMNS_SCOPE_ID)
.single {
it.owner.valueParameters.size == 1 && it.owner.valueParameters[0].type == context.irBuiltIns.stringType
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
val columnContainerExtension = generateExtensionProperty(
callableId = callableId,
receiverType = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.COLUMNS_CONTAINER_CLASS_ID),
ConeClassLikeLookupTagImpl(Names.COLUMNS_SCOPE_CLASS_ID),
typeArguments = arrayOf(schemaProperty.marker),
isNullable = false
),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,12 @@ object Names {
FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe")),
Name.identifier("ColumnsContainer")
)

val COLUMNS_SCOPE_CLASS_ID: ClassId
get() = ClassId(
FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe")),
Name.identifier("ColumnsScope")
)
val DATA_ROW_CLASS_ID: ClassId
get() = ClassId(FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe")), Name.identifier("DataRow"))
val DF_ANNOTATIONS_PACKAGE: Name
Expand Down
10 changes: 5 additions & 5 deletions plugins/kotlin-dataframe/testData/box/propertiesOrder.fir.txt
Original file line number Diff line number Diff line change
Expand Up @@ -25,31 +25,31 @@ FILE: propertiesOrder.kt
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.full_name: R|kotlin/String|
public get(): R|kotlin/String|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.full_name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Read_16I>|.full_name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.topics: R|kotlin/String|
public get(): R|kotlin/String|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.topics: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Read_16I>|.topics: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.watchers: R|kotlin/Int|
public get(): R|kotlin/Int|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.watchers: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Read_16I>|.watchers: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.stargazers_count: R|kotlin/Int|
public get(): R|kotlin/Int|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.stargazers_count: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Read_16I>|.stargazers_count: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.html_url: R|java/net/URL|
public get(): R|java/net/URL|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.html_url: R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Read_16I>|.html_url: R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|

public constructor(): R|<local>/Scope0|
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ FILE: selectDuringTyping.kt
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/ExplodeSchema_94I>|.timestamps: R|kotlin/collections/List<kotlin/Int>|
public get(): R|kotlin/collections/List<kotlin/Int>|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/ExplodeSchema_94I>|.timestamps: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/collections/List<kotlin/Int>>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/ExplodeSchema_94I>|.timestamps: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/collections/List<kotlin/Int>>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/collections/List<kotlin/Int>>|

public constructor(): R|<local>/Scope0|
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ FILE: test.kt
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/S_43I>|.javaRecord: R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/JavaRecord_771>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/JavaRecord_771>|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/S_43I>|.javaRecord: R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/JavaRecord_771>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/S_43I>|.javaRecord: R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/JavaRecord_771>|
public get(): R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/JavaRecord_771>|

public constructor(): R|<local>/Scope0|
Expand All @@ -47,19 +47,19 @@ FILE: test.kt
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/JavaRecord_771>|.aaa: R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Aaa_771>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Aaa_771>|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/JavaRecord_771>|.aaa: R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/Aaa_771>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/JavaRecord_771>|.aaa: R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/Aaa_771>|
public get(): R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/Aaa_771>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/JavaRecord_771>|.bean: R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Bean_771>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Bean_771>|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/JavaRecord_771>|.bean: R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/Bean_771>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/JavaRecord_771>|.bean: R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/Bean_771>|
public get(): R|org/jetbrains/kotlinx/dataframe/columns/ColumnGroup<<local>/Bean_771>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/JavaRecord_771>|.i: R|kotlin/Int|
public get(): R|kotlin/Int|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/JavaRecord_771>|.i: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/JavaRecord_771>|.i: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|

public constructor(): R|<local>/Scope1|
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ class ExtensionsGenerator(

private fun OutputStreamWriter.writeImports() {
appendLine("import org.jetbrains.kotlinx.dataframe.annotations.*")
appendLine("import org.jetbrains.kotlinx.dataframe.ColumnsContainer")
appendLine("import org.jetbrains.kotlinx.dataframe.ColumnsScope")
appendLine("import org.jetbrains.kotlinx.dataframe.DataColumn")
appendLine("import org.jetbrains.kotlinx.dataframe.DataFrame")
appendLine("import org.jetbrains.kotlinx.dataframe.DataRow")
Expand Down
Loading
Loading